From aef4e9b3b1423be9d5e39ca463fb085b85c221f6 Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Sat, 12 Sep 2026 16:49:36 -0700 Subject: [PATCH 001/151] [kernel] VSA kernel: one sm_100a / sm_103a image per listed arch; un-gate backward on sm_103a (#1833) Co-authored-by: Claude Fable 5.1 --- .github/workflows/publish-kernel.yml | 4 +-- fastvideo-kernel/CMakeLists.txt | 36 +++++++++++++------ .../block_sparse_bwd_kernel_sm100a.cuh | 16 ++++++--- fastvideo-kernel/csrc/common_extension.cpp | 2 +- .../block_sparse_attn_bwd_sm100a.py | 6 ++-- tests/test_block_sparse_bwd_sm100a.py | 4 +-- tests/test_block_sparse_sm100a_dispatch.py | 5 +-- 7 files changed, 48 insertions(+), 25 deletions(-) diff --git a/.github/workflows/publish-kernel.yml b/.github/workflows/publish-kernel.yml index 5e89ecf2c..6c4fea1ca 100644 --- a/.github/workflows/publish-kernel.yml +++ b/.github/workflows/publish-kernel.yml @@ -176,11 +176,11 @@ jobs: # the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps # Ninja so heavy CUTLASS/TK template TUs don't OOM the 16 GB runner (exit 143). if [ "${{ matrix.platform.arch }}" = "aarch64" ]; then - export TORCH_CUDA_ARCH_LIST="10.0a;12.0a" + export TORCH_CUDA_ARCH_LIST="10.0a;10.3a;12.0a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON" export CMAKE_BUILD_PARALLEL_LEVEL=1 elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then - export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;12.0a" + export TORCH_CUDA_ARCH_LIST="9.0a;10.0a;10.3a;12.0a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a" # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize. export CMAKE_BUILD_PARALLEL_LEVEL=1 diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index a7f968242..405d7b751 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -396,27 +396,41 @@ if(BUILD_CXX_KERNELS) # exported (build.sh, and `pip install` with it set) the branch above only prints it -- # the cmake variable stays empty, so testing that alone silently skips the kernel and # leaves a build that succeeds with the op missing. - set(ENABLE_VSA_SM100_FAMILY OFF) + set(ENABLE_VSA_SM100A OFF) + set(ENABLE_VSA_SM103A OFF) set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST}) set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") endif() - if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.(0|3)a|10(0|3)a|sm_10(0|3)a)([; ,]|$)") - set(ENABLE_VSA_SM100_FAMILY ON) + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)") + set(ENABLE_VSA_SM100A ON) endif() - if(ENABLE_VSA_SM100_FAMILY) - message(STATUS "fastvideo-kernel: building block_sparse_sm100a (sm_100a/sm_103a, 64- and 128-token blocks)") + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.3a|103a|sm_103a)([; ,]|$)") + set(ENABLE_VSA_SM103A ON) + endif() + set(_VSA_FWD_GENCODE "") + set(_VSA_FWD_ARCHS "") + if(ENABLE_VSA_SM100A) + list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_100a,code=sm_100a") + list(APPEND _VSA_FWD_ARCHS "sm_100a") + endif() + if(ENABLE_VSA_SM103A) + list(APPEND _VSA_FWD_GENCODE "-gencode;arch=compute_103a,code=sm_103a") + list(APPEND _VSA_FWD_ARCHS "sm_103a") + endif() + if(_VSA_FWD_GENCODE) + message(STATUS "fastvideo-kernel: building block_sparse_sm100a (${_VSA_FWD_ARCHS}; 64- and 128-token blocks)") list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu) set_source_files_properties(csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES - COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true") - # VSA block-sparse attention BACKWARD, 64-token blocks. sm_100a only for now: validated - # on GB200, not yet on B300/GB300, so no sm_103a image is built and the Python side keeps - # the Triton backward for sm_103a devices. + COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") + endif() + # VSA block-sparse attention BACKWARD, 64-token blocks: the same images as the forward. + if(_VSA_FWD_GENCODE) list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu) set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES - COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true") + COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI @@ -434,7 +448,7 @@ if(BUILD_CXX_KERNELS) # Build compile definitions list set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) - if(ENABLE_VSA_SM100_FAMILY) + if(ENABLE_VSA_SM100A OR ENABLE_VSA_SM103A) list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) endif() if(ENABLE_TK_KERNELS) diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh index 86d738cac..3d3014ff0 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh @@ -123,7 +123,9 @@ __global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa const int* __restrict__ k2q_num, const int* __restrict__ workitem_remap, const int* __restrict__ variable_block_sizes, int max_q_blocks, int num_samples, int num_heads, int seqlen, float scale_log2, float sm_scale) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) using DQ = DQConfig; const int num_kv_blocks_per_seq = seqlen / BLOCK; @@ -863,7 +865,9 @@ __global__ void __launch_bounds__(ORDER_THREADS, 1) vsa_bwd_order_kernel(const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num, int max_q_blocks, int num_kv_blocks_per_seq, int order_bin, bool snake, int* __restrict__ order_out) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) extern __shared__ int order_smem[]; int* sbin = order_smem; int* smid = order_smem + num_kv_blocks_per_seq; @@ -908,7 +912,9 @@ __global__ void __launch_bounds__(256, 1) __nv_bfloat16* __restrict__ dk, __nv_bfloat16* __restrict__ dv, const int* __restrict__ k2q_num, int num_samples, int num_heads, int seqlen) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) __shared__ __align__(128) __nv_bfloat16 tile[PRE_TOKENS][SUB_COLS_BF16 + 4]; const int num_kv_blocks_per_seq = seqlen / BLOCK; const int token_block_id = (int)blockIdx.x; @@ -1038,7 +1044,9 @@ template __global__ void __launch_bounds__(128, 1) vsa_bwd_postprocess_kernel(const DQ_DTYPE* __restrict__ dqaccum, __nv_bfloat16* __restrict__ dq, int num_heads, int seqlen, float sm_scale) { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) const int q_block_id = (int)blockIdx.x; const int batch_head = (int)blockIdx.y; const int batch = batch_head / num_heads, head = batch_head % num_heads; diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index f58a39711..45d5744bb 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -60,7 +60,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)"); m.def("block_sparse_sm100a_bwd", torch::wrap_pybind_function(block_sparse_sm100a_bwd), - "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a)"); + "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)"); #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py index 76d338cd6..ea073d868 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""sm_100a (Blackwell) CUDA block-sparse VSA backward. +"""sm_100a/sm_103a (data-center Blackwell) CUDA block-sparse VSA backward. Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and @@ -23,7 +23,7 @@ except ImportError: # pragma: no cover - extension not built _BWD = None _HAS_VSA_BWD_SM100A = False -_SM100 = (10, 0) +_SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 BLOCK = 64 # Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true). @@ -58,7 +58,7 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: """ if not _HAS_VSA_BWD_SM100A or not q.is_cuda: return False - if torch.cuda.get_device_capability(q.device) != _SM100: + if torch.cuda.get_device_capability(q.device) not in _SUPPORTED_COMPUTE_CAPABILITIES: return False if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM: return False diff --git a/tests/test_block_sparse_bwd_sm100a.py b/tests/test_block_sparse_bwd_sm100a.py index a2a7c1351..b9daea240 100644 --- a/tests/test_block_sparse_bwd_sm100a.py +++ b/tests/test_block_sparse_bwd_sm100a.py @@ -31,9 +31,9 @@ REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref| MEAN_ABS_TOL = 1e-3 # mean|got - ref| pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0) + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} or not bwd._HAS_VSA_BWD_SM100A, - reason="requires a compute-capability (10, 0) GPU (sm_100a) and a fastvideo_kernel " + reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel " "extension built with block_sparse_sm100a_bwd", ) diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py index 9bfdc719c..122aec2ea 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -175,8 +175,9 @@ def test_backward_uses_sm100a_kernel_when_built(monkeypatch): """Guards against a silent Triton fallback: with the op built on an sm_100a device the CUDA backward must be the one that runs.""" from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd - if not vsa_bwd._HAS_VSA_BWD_SM100A or torch.cuda.get_device_capability() != (10, 0): - pytest.skip("sm_100a backward not built for this device") + if (not vsa_bwd._HAS_VSA_BWD_SM100A + or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}): + pytest.skip("sm_100a/sm_103a backward not built for this device") _, _, sm100a_backward = _grads_sm100a_route_vs_triton(monkeypatch) assert sm100a_backward From 3a3ad3d20941a41826a41a72e2f3befaa003ca86 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:31:59 -0700 Subject: [PATCH 002/151] [feat]: NVFP4 text encoder for MiniMax H3 (#1837) --- fastvideo/layers/quantization/base_config.py | 9 + .../encoders/minimax_h3_checkpoint_nvfp4.py | 425 +++++++++++++++ .../models/encoders/minimax_h3_qwen3_vl.py | 21 +- fastvideo/models/loader/component_loader.py | 20 +- ...st_minimax_h3_qwen3_vl_checkpoint_nvfp4.py | 483 ++++++++++++++++++ 5 files changed, 946 insertions(+), 12 deletions(-) create mode 100644 fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py create mode 100644 fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py diff --git a/fastvideo/layers/quantization/base_config.py b/fastvideo/layers/quantization/base_config.py index c7dc2754a..c9ef3f6e4 100644 --- a/fastvideo/layers/quantization/base_config.py +++ b/fastvideo/layers/quantization/base_config.py @@ -139,5 +139,14 @@ class QuantizationConfig(ABC): """ raise NotImplementedError + def validate_runtime(self, device: torch.device) -> None: + """Fail fast on a device the scheme cannot execute on. + + Called once by loaders that select a serialized checkpoint scheme, + before any weight is read. The default accepts every device; schemes + with kernel or capability requirements override it. + """ + return None + def get_cache_scale(self, name: str) -> str | None: return None diff --git a/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py new file mode 100644 index 000000000..0c96e636c --- /dev/null +++ b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialized NVFP4 execution for the MiniMax-H3 Qwen3-VL encoder. + +The bf16 conditioner is 48 GB resident and sets the single-GB10 peak once the +DiT runs in FP8. A checkpoint written by +``scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py`` +stores every ``language_model.layers.*`` linear as three tensors and no +``weight``:: + + .weight_packed uint8 [out, in // 2] two E2M1 values per byte + .weight_scale uint8 [out, in // 16] one E4M3 scale per 16 values, + FlashInfer 128x4 swizzled layout + .weight_global_scale float32 [1] (448 * 6) / amax(|W|) + +The bytes are exactly what ``flashinfer.nvfp4_quantize(W, global_scale, +sfLayout=SfLayout.layout_128x4)`` returns, so loading them reproduces the state +``nvfp4_config.convert_model_to_nvfp4`` builds at runtime without ever +materializing the bf16 weight. Activations are quantized per call with a unit +global scale and multiplied with ``flashinfer.mm_fp4``. + +The checkpoint selects this path through ``config.json``; every field is +required so a checkpoint from another exporter cannot pass by omission:: + + "quantization_config": {"quant_method": "nvfp4", "activation_scheme": "dynamic", + "fmt": "e2m1", "group_size": 16, "scale_fmt": "e4m3", + "scale_layout": "128x4", + "modules_to_not_convert": ["model.visual", "lm_head"]} + +Whole projection kinds may stay bf16 by listing their suffix in +``modules_to_not_convert`` (for example ``"mlp.down_proj"``); those linears +keep a plain ``weight`` and the unquantized method. Only the seven language +projection names are accepted there, so a typo fails at config time. + +Single GPU only: packed columns and swizzled scale rows cannot be narrowed per +tensor-parallel rank without repacking. Missing tensors are reported by the +loader's strict check before the post-load hook runs; the hook adds only the +content checks a copied tensor can still fail. +""" + +from typing import Any + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +from fastvideo.distributed import get_tp_world_size +from fastvideo.layers.linear import LinearBase, LinearMethodBase +from fastvideo.layers.quantization.base_config import QuantizationConfig +from fastvideo.layers.quantization.nvfp4_config import ( + _coerce_fp4_input_dtype, + _mm_fp4, + _nvfp4_quantize, + _register_ops_once, + _require_flashinfer, +) +from fastvideo.models.utils import set_weight_attrs + +NVFP4_GROUP_SIZE = 16 +NVFP4_SCALE_LAYOUT = "128x4" +# FlashInfer's 128x4 layout tiles scales in 128 rows by 4 scale columns. Keeping +# every weight a whole number of tiles means the serialized scale tensor is +# exactly [out, in // 16] with no padding to describe. Every Qwen3-VL language +# linear satisfies this (out in {1024, 5120, 8192, 25600}, in in {5120, 8192, 25600}). +NVFP4_ROW_TILE = 128 +NVFP4_COLUMN_MULTIPLE = 4 * NVFP4_GROUP_SIZE +# The linears of one Qwen3-VL language layer, as the checkpoint and the model name them. +LANGUAGE_PROJECTIONS = ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", +) +# The three tensors a quantized linear stores, as parameter names on the layer +# and as checkpoint key suffixes. The converter imports these so both sides +# spell them once. +NVFP4_TENSOR_SUFFIXES = ("weight_packed", "weight_scale", "weight_global_scale") +_REQUIRED_METADATA_KEYS = ("activation_scheme", "fmt", "group_size", "scale_fmt", "scale_layout", + "modules_to_not_convert") +_E2M1_MAX = 6.0 +_E4M3_MAX = 448.0 +# An E4M3 byte of 0xFF is NaN, so a scale tile that is still all 0xFF after +# loading is one the checkpoint never filled. The loader copies whole tensors +# or asserts, so the first 128x4 tile speaks for the tensor. +_UNLOADED_SCALE_BYTE = 0xFF +_SCALE_TILE_BYTES = NVFP4_ROW_TILE * 4 + + +def validate_nvfp4_geometry(output_size: int, input_size: int) -> None: + if output_size % NVFP4_ROW_TILE: + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires output_size divisible by {NVFP4_ROW_TILE}, " + f"got {output_size}") + if input_size % NVFP4_COLUMN_MULTIPLE: + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires input_size divisible by {NVFP4_COLUMN_MULTIPLE}, " + f"got {input_size}") + + +def nvfp4_packed_weight_shape(output_size: int, input_size: int) -> tuple[int, int]: + return output_size, input_size // 2 + + +def nvfp4_scale_shape(output_size: int, input_size: int) -> tuple[int, int]: + return output_size, input_size // NVFP4_GROUP_SIZE + + +def nvfp4_weight_global_scale(weight: torch.Tensor) -> torch.Tensor: + """The per-tensor scale ``convert_model_to_nvfp4`` uses: E4M3 max times E2M1 max over amax. + + Unlike the runtime converter this refuses NaN or infinite weights instead of + mapping them to zero, because a checkpoint written from them would be wrong + forever. + """ + amax = weight.float().abs().max() + if not torch.isfinite(amax) or amax <= 0: + raise ValueError("MiniMax-H3 NVFP4 global scale needs a finite, non-zero weight amax") + scale = ((_E4M3_MAX * _E2M1_MAX) / amax).to(torch.float32) + if not torch.isfinite(scale): + raise ValueError(f"MiniMax-H3 NVFP4 global scale overflowed float32 for weight amax {amax.item():.3e}") + return scale + + +def serialized_nvfp4_quantization_config( + *, + keep_bf16: tuple[str, ...] | list[str] = (), + producer: dict[str, Any] | None = None, +) -> dict[str, Any]: + """The ``config.json`` ``quantization_config`` the converter writes and ``from_config`` accepts. + + ``keep_bf16`` lists projection kinds from ``LANGUAGE_PROJECTIONS`` that stay + bf16 in every language layer; they are appended to ``modules_to_not_convert``. + """ + unknown = sorted(set(keep_bf16) - set(LANGUAGE_PROJECTIONS)) + if unknown: + raise ValueError(f"keep_bf16 names unknown projection kinds {unknown}; choose from {LANGUAGE_PROJECTIONS}") + config: dict[str, Any] = { + "quant_method": "nvfp4", + "activation_scheme": "dynamic", + "fmt": "e2m1", + "group_size": NVFP4_GROUP_SIZE, + "scale_fmt": "e4m3", + "scale_layout": NVFP4_SCALE_LAYOUT, + "modules_to_not_convert": ["model.visual", "lm_head", *keep_bf16], + } + if producer: + config["producer"] = dict(producer) + return config + + +def _quantize_activation_nvfp4(x_2d: torch.Tensor, global_scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize one bf16 activation the way ``NVFP4QuantizeMethod`` does: unit global scale, 128x4 layout. + + The scale tensor is returned as uint8 bytes whatever dtype FlashInfer labels + it with, so it always matches the uint8 weight scales the checkpoint stores. + """ + sf_layout, _, _ = _require_flashinfer() + x_fp4, x_scale = _nvfp4_quantize(x_2d, global_scale, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + if x_scale.dtype != torch.uint8: + x_scale = x_scale.view(torch.uint8) + return x_fp4, x_scale + + +def _nvfp4_linear( + x_fp4: torch.Tensor, + x_scale: torch.Tensor, + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, +) -> torch.Tensor: + """``x @ W.T`` on packed operands. Mirrors ``NVFP4QuantizeMethod.apply``: the weight is + handed to ``mm_fp4`` transposed, the output is bf16, the backend is FlashInfer's choice.""" + return _mm_fp4( + x_fp4, + weight_packed.t(), + x_scale, + weight_scale.t(), + alpha, + torch.bfloat16, + None, + backend="auto", + ) + + +class MiniMaxH3SerializedNVFP4Config(QuantizationConfig): + """Serialized 16-group NVFP4 contract for the H3 text encoder. + + The group size and scale layout are fixed by the loader's parameter shapes, + so the only state is which projection kinds the checkpoint kept in bf16. + """ + + def __init__(self, bf16_projections: tuple[str, ...] = ()) -> None: + super().__init__() + unknown = sorted(set(bf16_projections) - set(LANGUAGE_PROJECTIONS)) + if unknown: + raise ValueError(f"MiniMax-H3 serialized NVFP4 cannot keep unknown projection kinds {unknown} in bf16; " + f"choose from {LANGUAGE_PROJECTIONS}") + # Projection kinds the checkpoint kept in bf16 in every language layer, + # e.g. ("mlp.down_proj",). Those linears load a plain weight. + self.bf16_projections = tuple(bf16_projections) + + @classmethod + def get_name(cls) -> str: + return "nvfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @staticmethod + def get_config_filenames() -> list[str]: + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "MiniMaxH3SerializedNVFP4Config": + quant_method = str(config.get("quant_method", "")).lower() + if quant_method != "nvfp4": + raise ValueError(f"MiniMax-H3 serialized NVFP4 config got quant_method {quant_method!r}") + missing = [key for key in _REQUIRED_METADATA_KEYS if key not in config] + if missing: + raise ValueError("MiniMax-H3 serialized NVFP4 requires explicit quantization_config fields; " + f"missing {missing}") + if str(config["activation_scheme"]).lower() != "dynamic": + raise ValueError("MiniMax-H3 serialized NVFP4 requires dynamic activation quantization") + if str(config["fmt"]).lower() not in ("e2m1", "float4_e2m1fn", "nvfp4"): + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires E2M1 weights, got {config['fmt']!r}") + if str(config["scale_fmt"]).lower() not in ("e4m3", "float8_e4m3fn"): + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires E4M3 block scales, got {config['scale_fmt']!r}") + group_size = config["group_size"] + if isinstance(group_size, bool) or group_size != NVFP4_GROUP_SIZE: + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires group_size={NVFP4_GROUP_SIZE}, got {group_size!r}") + if config["scale_layout"] != NVFP4_SCALE_LAYOUT: + raise ValueError(f"MiniMax-H3 serialized NVFP4 requires scale_layout={NVFP4_SCALE_LAYOUT!r}, " + f"got {config['scale_layout']!r}") + ignored_layers = config["modules_to_not_convert"] + if not isinstance(ignored_layers, list | tuple): + raise ValueError("MiniMax-H3 serialized NVFP4 modules_to_not_convert must be a sequence") + bf16_projections: list[str] = [] + saw_visual = False + for name in ignored_layers: + if not isinstance(name, str) or not name: + raise ValueError("MiniMax-H3 serialized NVFP4 modules_to_not_convert entries must be non-empty " + f"strings, got {name!r}") + if "visual" in name: + saw_visual = True + elif name == "lm_head" or name.endswith(".lm_head"): + continue + elif name in LANGUAGE_PROJECTIONS: + bf16_projections.append(name) + else: + # A whole projection kind may stay bf16 across every language + # layer; a single layer, a typo, or an HF-style full path is not + # a contract this loader supports. + raise ValueError("MiniMax-H3 serialized NVFP4 keeps bf16 per projection kind only; " + f"modules_to_not_convert entry {name!r} is not one of {LANGUAGE_PROJECTIONS}") + if not saw_visual: + raise ValueError("MiniMax-H3 serialized NVFP4 requires the vision stack to be listed in " + "modules_to_not_convert") + return cls(tuple(bf16_projections)) + + def validate_runtime(self, device: torch.device) -> None: + if device.type != "cuda": + raise RuntimeError(f"MiniMax-H3 serialized NVFP4 requires a CUDA device; got {device.type!r}") + capability = torch.cuda.get_device_capability(device) + capability_number = capability[0] * 10 + capability[1] + if capability_number < self.get_min_capability(): + raise RuntimeError("MiniMax-H3 serialized NVFP4 requires GPU capability " + f"sm{self.get_min_capability()} or newer, got sm{capability_number}") + if capability[0] not in (10, 12): + raise RuntimeError("MiniMax-H3 serialized NVFP4 runs FlashInfer's Blackwell FP4 GEMM; " + f"got unsupported sm{capability_number}") + if get_tp_world_size() > 1: + raise NotImplementedError("MiniMax-H3 serialized NVFP4 supports a single GPU: packed FP4 columns and " + "128x4 swizzled scale rows cannot be narrowed per tensor-parallel rank") + sf_layout, _, _ = _require_flashinfer() + if not hasattr(sf_layout, "layout_128x4"): + raise RuntimeError("The installed flashinfer has no SfLayout.layout_128x4; MiniMax-H3 serialized NVFP4 " + "checkpoints store their scales in that layout") + _register_ops_once() + + def is_kept_bf16(self, prefix: str) -> bool: + return any(prefix == name or prefix.endswith("." + name) for name in self.bf16_projections) + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + if not isinstance(layer, LinearBase) or ".language_model.layers." not in prefix: + return None + if self.is_kept_bf16(prefix): + return None + return MiniMaxH3SerializedNVFP4LinearMethod() + + +def _strict_dtype_loader(base_loader): + """Wrap a layer's weight loader so a checkpoint tensor must carry the parameter's exact dtype.""" + if base_loader is None: + raise ValueError("MiniMax-H3 serialized NVFP4 linears need the layer's weight_loader") + + def load(param: torch.Tensor, loaded_weight: torch.Tensor, *args: Any, **kwargs: Any): + if loaded_weight.dtype != param.dtype: + raise ValueError("Serialized MiniMax-H3 NVFP4 tensors must be stored as their declared dtype; " + f"got {loaded_weight.dtype} for a {param.dtype} parameter of shape {tuple(param.shape)}") + return base_loader(param, loaded_weight, *args, **kwargs) + + return load + + +class MiniMaxH3SerializedNVFP4LinearMethod(LinearMethodBase): + """Execute serialized NVFP4 weights without re-quantizing them.""" + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + validate_nvfp4_geometry(output_size_per_partition, input_size_per_partition) + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # No input_dim/output_dim: these tensors are not shardable, and without + # the attributes the weight loaders copy them whole. The copy is a numeric + # cast, so a checkpoint that stored scales as float8 bytes would be + # silently converted value by value; refuse any dtype but the declared one. + loader_attrs = {"weight_loader": _strict_dtype_loader(extra_weight_attrs.get("weight_loader"))} + packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES + weight_packed = Parameter( + torch.zeros(nvfp4_packed_weight_shape(output_size_per_partition, input_size_per_partition), + dtype=torch.uint8), + requires_grad=False, + ) + set_weight_attrs(weight_packed, loader_attrs) + layer.register_parameter(packed_name, weight_packed) + + weight_scale = Parameter( + torch.full(nvfp4_scale_shape(output_size_per_partition, input_size_per_partition), + _UNLOADED_SCALE_BYTE, + dtype=torch.uint8), + requires_grad=False, + ) + set_weight_attrs(weight_scale, loader_attrs) + layer.register_parameter(scale_name, weight_scale) + + weight_global_scale = Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False) + set_weight_attrs(weight_global_scale, loader_attrs) + layer.register_parameter(global_scale_name, weight_global_scale) + # No bf16 weight ever exists on this layer; ``None`` keeps ``layer.weight`` + # readable for code that inspects it, matching the purged NVFP4 path. + layer.register_parameter("weight", None) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Reject a layer the checkpoint did not fill and derive the GEMM multiplier. + + Shapes and dtypes are fixed by ``create_weights`` and enforced by the + weight loader's copy, and the loader reports missing tensors by name + before this runs; what remains is content a copied tensor can still get + wrong: a non-positive global scale, or a scale tile left at its 0xFF + initializer by a direct caller that bypassed the loader. + """ + weight_scale = layer.weight_scale + global_scale = float(layer.weight_global_scale.item()) + if not (global_scale > 0) or global_scale == float("inf"): + raise ValueError("Serialized MiniMax-H3 NVFP4 weight_global_scale was not loaded: " + f"it must be a finite positive value, got {global_scale}") + if bool((weight_scale.view(-1)[:_SCALE_TILE_BYTES] == _UNLOADED_SCALE_BYTE).all()): + raise ValueError("Serialized MiniMax-H3 NVFP4 weight_scale was not loaded: its first tile is still 0xFF") + # ``mm_fp4`` folds both global scales into one multiplier. Activations use a + # unit global scale, so the multiplier is the inverse weight global scale. + device = weight_scale.device + layer.register_buffer("_nvfp4_alpha", torch.tensor(1.0 / global_scale, dtype=torch.float32, device=device), + persistent=False) + layer.register_buffer("_nvfp4_x_global_scale", torch.ones((), dtype=torch.float32, device=device), + persistent=False) + + @staticmethod + def _apply_finalized(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + x = _coerce_fp4_input_dtype(x) + original_shape = x.shape + if x.numel() == 0: + # An empty prompt has nothing to quantize; the FP4 kernels are not defined for zero rows. + return x.new_zeros(*original_shape[:-1], layer.output_size_per_partition, dtype=torch.bfloat16) + x_fp4, x_scale = _quantize_activation_nvfp4(x.reshape(-1, original_shape[-1]), layer._nvfp4_x_global_scale) + output = _nvfp4_linear(x_fp4, x_scale, layer.weight_packed, layer.weight_scale, layer._nvfp4_alpha) + if bias is not None: + # The GEMM emits bf16; keep it that way whatever dtype the bias was built in. + output = output + bias.to(output.dtype) + return output.view(*original_shape[:-1], output.shape[-1]) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if x.device.type != "cuda": + raise RuntimeError("MiniMax-H3 serialized NVFP4 execution requires CUDA") + if getattr(layer, "_nvfp4_alpha", None) is None: + raise RuntimeError("MiniMax-H3 serialized NVFP4 linear was not finalized: " + "process_weights_after_loading has not run") + return self._apply_finalized(layer, x, bias) + + +__all__ = [ + "LANGUAGE_PROJECTIONS", + "NVFP4_TENSOR_SUFFIXES", + "MiniMaxH3SerializedNVFP4Config", + "MiniMaxH3SerializedNVFP4LinearMethod", + "NVFP4_GROUP_SIZE", + "NVFP4_SCALE_LAYOUT", + "nvfp4_packed_weight_shape", + "nvfp4_scale_shape", + "nvfp4_weight_global_scale", + "serialized_nvfp4_quantization_config", + "validate_nvfp4_geometry", +] diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index a67273195..9f4b5e8ea 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -12,9 +12,11 @@ from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3 from fastvideo.distributed import get_tp_world_size from fastvideo.layers.layernorm import RMSNorm from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear +from fastvideo.layers.quantization.base_config import QuantizationConfig from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding from fastvideo.models.encoders.base import TextEncoder from fastvideo.models.encoders.minimax_h3_checkpoint_fp8 import MiniMaxH3SerializedFP8Config +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config from fastvideo.models.loader.weight_utils import default_weight_loader @@ -503,14 +505,26 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]): """H3 conditioner returning the unnormalized layer-50 hidden tensor.""" supports_hf_from_pretrained = False - supported_checkpoint_quantization_methods = frozenset({"fp8"}) + # Serialized quantization contracts this conditioner can load, keyed by the + # ``quant_method`` a checkpoint declares. The loader gates on the key set + # before calling the factory, so both stay in step by construction. + _checkpoint_quantization_configs: dict[str, type[QuantizationConfig]] = { + MiniMaxH3SerializedFP8Config.get_name(): MiniMaxH3SerializedFP8Config, + MiniMaxH3SerializedNVFP4Config.get_name(): MiniMaxH3SerializedNVFP4Config, + } + supported_checkpoint_quantization_methods = frozenset(_checkpoint_quantization_configs) @classmethod def checkpoint_quantization_config_from_metadata( cls, metadata: dict[str, Any], - ) -> MiniMaxH3SerializedFP8Config: - return MiniMaxH3SerializedFP8Config.from_config(metadata) + ) -> QuantizationConfig: + quant_method = str(metadata.get("quant_method", "")).lower() + config_cls = cls._checkpoint_quantization_configs.get(quant_method) + if config_cls is None: + raise ValueError(f"MiniMax-H3 has no serialized {quant_method!r} text-encoder contract; " + f"supported: {sorted(cls._checkpoint_quantization_configs)}") + return config_cls.from_config(metadata) def __init__(self, config: MiniMaxH3Qwen3VLConfig) -> None: super().__init__(config) @@ -756,4 +770,5 @@ EntryClass = MiniMaxH3Qwen3VLConditioner __all__ = [ "MiniMaxH3Qwen3VLConditioner", "MiniMaxH3SerializedFP8Config", + "MiniMaxH3SerializedNVFP4Config", ] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 7f850744e..148a6d1bb 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -450,9 +450,19 @@ class TextEncoderLoader(ComponentLoader): self.counter_after_loading_weights - self.counter_before_loading_weights, ) + # Strict check for unquantized models and for serialized checkpoint + # schemes, which track every loaded tensor. It runs before the + # post-load hooks so a missing tensor is reported by name here + # instead of as one anonymous "not loaded" failure inside a hook. + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded and (model_config.quant_config is None or checkpoint_quant_config is not None): + raise ValueError("Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}") + if checkpoint_quant_config is not None: processed_linears = _process_quantized_text_encoder_weights(model, runtime_device) - logger.info("Validated %d serialized blockwise FP8 text-encoder linears", processed_linears) + logger.info("Validated %d serialized %s text-encoder linears", processed_linears, + checkpoint_quant_config.get_name()) # Explicitly move model to target device after loading weights model = model.to(target_device) @@ -492,14 +502,6 @@ class TextEncoderLoader(ComponentLoader): fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory, ) - # We only enable strict check for non-quantized models - # that have loaded weights tracking currently. - # if loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded and (model_config.quant_config is None or checkpoint_quant_config is not None): - raise ValueError("Following weights were not initialized from " - f"checkpoint: {weights_not_loaded}") - return model.eval() diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py new file mode 100644 index 000000000..74be1a9e8 --- /dev/null +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py @@ -0,0 +1,483 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import json +import os + +import pytest +import torch + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29515") + +import fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 as h3_nvfp4 +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import ( + MiniMaxH3Qwen3VLArchConfig, + MiniMaxH3Qwen3VLConfig, +) +from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear, UnquantizedLinearMethod +from fastvideo.layers.vocab_parallel_embedding import UnquantizedEmbeddingMethod, VocabParallelEmbedding +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.encoders.minimax_h3_checkpoint_fp8 import MiniMaxH3SerializedFP8Config +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import ( + LANGUAGE_PROJECTIONS, + NVFP4_TENSOR_SUFFIXES, + MiniMaxH3SerializedNVFP4Config, + MiniMaxH3SerializedNVFP4LinearMethod, + serialized_nvfp4_quantization_config, +) +from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner +from fastvideo.models.loader.text_encoder_quantization import ( + _configure_text_encoder_quantization, + _process_quantized_text_encoder_weights, + _read_text_encoder_checkpoint_quantization_config, +) + +LAYER_PREFIX = "minimax_h3_qwen3_vl.language_model.layers.0" +E4M3_ONE = 0x38 + + +def _checkpoint_quantization_config(**overrides) -> dict: + config = serialized_nvfp4_quantization_config() + config.update(overrides) + return config + + +def _language_linear(config: MiniMaxH3SerializedNVFP4Config, + input_size: int = 128, + output_size: int = 256, + proj: str = "self_attn.q_proj", + bias: bool = False) -> ColumnParallelLinear: + return ColumnParallelLinear( + input_size=input_size, + output_size=output_size, + bias=bias, + quant_config=config, + prefix=f"{LAYER_PREFIX}.{proj}", + ) + + +def _fill_loaded(layer: torch.nn.Module, global_scale: float = 2.0) -> None: + layer.weight_packed.data.fill_(0x11) + layer.weight_scale.data.fill_(E4M3_ONE) + layer.weight_global_scale.data.fill_(global_scale) + + +def test_h3_accepts_only_the_serialized_nvfp4_contract() -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + assert config.bf16_projections == () + assert config.get_name() == "nvfp4" + assert config.get_supported_act_dtypes() == [torch.bfloat16] + + with pytest.raises(ValueError, match="group_size=16"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(group_size=32)) + with pytest.raises(ValueError, match="scale_layout='128x4'"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(scale_layout="linear")) + with pytest.raises(ValueError, match="dynamic activation"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(activation_scheme="static")) + with pytest.raises(ValueError, match="E2M1"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(fmt="e4m3")) + with pytest.raises(ValueError, match="vision stack"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(modules_to_not_convert=["lm_head"])) + with pytest.raises(ValueError, match="quant_method 'fp8'"): + MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config(quant_method="fp8")) + + +def test_h3_requires_every_layout_field_to_be_explicit() -> None: + bare = {"quant_method": "nvfp4", "modules_to_not_convert": ["model.visual", "lm_head"]} + with pytest.raises(ValueError, match="explicit quantization_config fields"): + MiniMaxH3SerializedNVFP4Config.from_config(bare) + + +def test_h3_keeps_bf16_per_projection_kind_only() -> None: + kept = MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", "mlp.down_proj"])) + assert kept.bf16_projections == ("mlp.down_proj", ) + assert kept.is_kept_bf16(f"{LAYER_PREFIX}.mlp.down_proj") + assert not kept.is_kept_bf16(f"{LAYER_PREFIX}.mlp.up_proj") + + with pytest.raises(ValueError, match="per projection kind only"): + MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", "language_model.layers.3"])) + with pytest.raises(ValueError, match="mlp.dwn_proj"): + MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", "mlp.dwn_proj"])) + with pytest.raises(ValueError, match="non-empty strings"): + MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", ""])) + + +def test_converter_metadata_round_trips_and_tolerates_producer_notes() -> None: + metadata = serialized_nvfp4_quantization_config(producer={"converter": "x.py", "flashinfer": "0.6.13rc2"}) + assert MiniMaxH3SerializedNVFP4Config.from_config(metadata).bf16_projections == () + assert metadata["producer"]["flashinfer"] == "0.6.13rc2" + + mixed = serialized_nvfp4_quantization_config(keep_bf16=("mlp.down_proj", )) + assert mixed["modules_to_not_convert"] == ["model.visual", "lm_head", "mlp.down_proj"] + assert MiniMaxH3SerializedNVFP4Config.from_config(mixed).bf16_projections == ("mlp.down_proj", ) + with pytest.raises(ValueError, match="unknown projection kinds"): + serialized_nvfp4_quantization_config(keep_bf16=("mlp.dwn_proj", )) + assert set(LANGUAGE_PROJECTIONS) == { + "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", + "mlp.up_proj", "mlp.down_proj" + } + + +def test_serialized_nvfp4_allocates_packed_weight_and_scales_without_a_bf16_weight(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + layer = _language_linear(config) + + assert isinstance(layer.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) + assert layer.weight is None + assert layer.weight_packed.dtype == torch.uint8 + assert layer.weight_packed.shape == (256, 64) + assert layer.weight_scale.dtype == torch.uint8 + assert layer.weight_scale.shape == (256, 8) + assert layer.weight_global_scale.dtype == torch.float32 + assert layer.weight_global_scale.shape == (1, ) + for name in ("weight_packed", "weight_scale", "weight_global_scale"): + param = getattr(layer, name) + assert callable(getattr(param, "weight_loader", None)) + assert not hasattr(param, "output_dim") and not hasattr(param, "input_dim") + + _fill_loaded(layer, global_scale=4.0) + packed_pointer = layer.weight_packed.data_ptr() + scale_pointer = layer.weight_scale.data_ptr() + layer.quant_method.process_weights_after_loading(layer) + + assert layer.weight_packed.data_ptr() == packed_pointer + assert layer.weight_scale.data_ptr() == scale_pointer + assert layer.weight is None + assert layer._nvfp4_alpha.dtype == torch.float32 + assert layer._nvfp4_alpha.item() == pytest.approx(0.25) + assert layer._nvfp4_x_global_scale.item() == 1.0 + + +def test_serialized_nvfp4_finalization_rejects_tensors_the_checkpoint_never_filled(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + layer = _language_linear(config) + + with pytest.raises(ValueError, match="weight_global_scale was not loaded"): + layer.quant_method.process_weights_after_loading(layer) + + layer.weight_global_scale.data.fill_(2.0) + with pytest.raises(ValueError, match="weight_scale was not loaded"): + layer.quant_method.process_weights_after_loading(layer) + + layer.weight_scale.data.fill_(E4M3_ONE) + layer.quant_method.process_weights_after_loading(layer) + assert layer._nvfp4_alpha.item() == pytest.approx(0.5) + + +def test_serialized_nvfp4_tensors_arrive_through_the_layer_weight_loader(distributed_setup) -> None: + """Drive the three tensors the way ``load_weights`` does, through the layer's own loader.""" + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + column = _language_linear(config, proj="self_attn.q_proj") + row = RowParallelLinear( + input_size=128, + output_size=256, + bias=False, + quant_config=config, + prefix=f"{LAYER_PREFIX}.self_attn.o_proj", + ) + for layer in (column, row): + assert isinstance(layer.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) + packed = torch.full((256, 64), 0x22, dtype=torch.uint8) + scale = torch.full((256, 8), E4M3_ONE, dtype=torch.uint8) + layer.weight_packed.weight_loader(layer.weight_packed, packed) + layer.weight_scale.weight_loader(layer.weight_scale, scale) + layer.weight_global_scale.weight_loader(layer.weight_global_scale, torch.tensor(8.0)) + layer.quant_method.process_weights_after_loading(layer) + assert torch.equal(layer.weight_packed.data, packed) + assert layer._nvfp4_alpha.item() == pytest.approx(0.125) + + with pytest.raises(AssertionError): + column.weight_packed.weight_loader(column.weight_packed, torch.zeros(256, 128, dtype=torch.uint8)) + # A float8 scale tensor would be value-cast into uint8 by a plain copy; the loader refuses it. + with pytest.raises(ValueError, match="declared dtype"): + column.weight_scale.weight_loader(column.weight_scale, torch.zeros(256, 8, dtype=torch.float8_e4m3fn)) + with pytest.raises(ValueError, match="declared dtype"): + column.weight_global_scale.weight_loader(column.weight_global_scale, torch.tensor(8.0, dtype=torch.float16)) + + +def test_serialized_nvfp4_quantizes_only_language_linears(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + visual_linear = ColumnParallelLinear( + input_size=128, + output_size=128, + bias=False, + quant_config=config, + prefix="minimax_h3_qwen3_vl.visual.blocks.0.attn.proj", + ) + embedding = VocabParallelEmbedding( + num_embeddings=128, + embedding_dim=128, + org_num_embeddings=128, + quant_config=config, + prefix="minimax_h3_qwen3_vl.language_model.embed_tokens", + ) + + assert isinstance(visual_linear.quant_method, UnquantizedLinearMethod) + assert visual_linear.weight.dtype == torch.get_default_dtype() + assert isinstance(embedding.quant_method, UnquantizedEmbeddingMethod) + assert embedding.weight.dtype == torch.get_default_dtype() + + +def test_kept_bf16_projection_builds_a_plain_linear(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", "mlp.down_proj"])) + down = _language_linear(config, input_size=128, output_size=128, proj="mlp.down_proj") + up = _language_linear(config, input_size=128, output_size=128, proj="mlp.up_proj") + + assert isinstance(down.quant_method, UnquantizedLinearMethod) + assert down.weight is not None and down.weight.shape == (128, 128) + assert not hasattr(down, "weight_packed") + assert isinstance(up.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) + assert up.weight is None + + +def test_serialized_nvfp4_rejects_geometry_outside_flashinfer_tiles(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + with pytest.raises(ValueError, match="input_size divisible by 64"): + _language_linear(config, input_size=96, output_size=256) + with pytest.raises(ValueError, match="output_size divisible by 128"): + _language_linear(config, input_size=128, output_size=200) + + +def test_serialized_nvfp4_cpu_execution_fails_closed(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + layer = _language_linear(config, input_size=128, output_size=128) + _fill_loaded(layer) + layer.quant_method.process_weights_after_loading(layer) + + with pytest.raises(RuntimeError, match="requires CUDA"): + layer(torch.zeros(2, 128, dtype=torch.bfloat16)) + + +def test_runtime_preflight_gates_device_capability_parallelism_and_flashinfer(monkeypatch) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + with pytest.raises(RuntimeError, match="requires a CUDA device"): + config.validate_runtime(torch.device("cpu")) + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) + with pytest.raises(RuntimeError, match="sm100 or newer"): + config.validate_runtime(torch.device("cuda")) + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (12, 1)) + monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 2) + with pytest.raises(NotImplementedError, match="single GPU"): + config.validate_runtime(torch.device("cuda")) + + monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 1) + + def missing_flashinfer(): + raise ImportError("NVFP4 quantization requires flashinfer") + + monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", missing_flashinfer) + with pytest.raises(ImportError, match="requires flashinfer"): + config.validate_runtime(torch.device("cuda")) + + class LayoutWithoutSwizzle: + pass + + monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", lambda: (LayoutWithoutSwizzle, None, None)) + with pytest.raises(RuntimeError, match="layout_128x4"): + config.validate_runtime(torch.device("cuda")) + + class Layout: + layout_128x4 = 1 + + registered = [] + monkeypatch.setattr(h3_nvfp4, "_require_flashinfer", lambda: (Layout, None, None)) + monkeypatch.setattr(h3_nvfp4, "_register_ops_once", lambda: registered.append(True)) + config.validate_runtime(torch.device("cuda")) + assert registered == [True] + + +def test_loader_selects_nvfp4_from_checkpoint_metadata(tmp_path) -> None: + checkpoint_config = _checkpoint_quantization_config() + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": checkpoint_config}), encoding="utf-8") + + assert _read_text_encoder_checkpoint_quantization_config(str(tmp_path)) == checkpoint_config + model_config = MiniMaxH3Qwen3VLConfig() + quant_config = _configure_text_encoder_quantization(model_config, MiniMaxH3Qwen3VLConditioner, str(tmp_path)) + assert isinstance(quant_config, MiniMaxH3SerializedNVFP4Config) + assert model_config.quant_config is quant_config + + with pytest.raises(ValueError, match="does not support serialized 'nvfp4'"): + _configure_text_encoder_quantization(MiniMaxH3Qwen3VLConfig(), TextEncoder, str(tmp_path)) + + +def test_conditioner_routes_each_serialized_scheme_to_its_config() -> None: + fp8_metadata = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "fmt": "e4m3", + "weight_block_size": [128, 128], + "modules_to_not_convert": ["model.visual", "lm_head"], + } + assert MiniMaxH3Qwen3VLConditioner.supported_checkpoint_quantization_methods == frozenset({"fp8", "nvfp4"}) + assert isinstance(MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata(fp8_metadata), + MiniMaxH3SerializedFP8Config) + assert isinstance( + MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata(_checkpoint_quantization_config()), + MiniMaxH3SerializedNVFP4Config) + with pytest.raises(ValueError, match="no serialized 'int4'"): + MiniMaxH3Qwen3VLConditioner.checkpoint_quantization_config_from_metadata({"quant_method": "int4"}) + + +def test_post_load_processing_visits_only_serialized_nvfp4_linears(distributed_setup) -> None: + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + quantized = _language_linear(config, input_size=128, output_size=128) + plain = ColumnParallelLinear(input_size=128, output_size=128, bias=False, prefix="plain") + _fill_loaded(quantized) + model = torch.nn.ModuleList([quantized, plain]) + + assert _process_quantized_text_encoder_weights(model, torch.device("cpu")) == 1 + assert quantized.weight_packed.device.type == "cpu" + assert quantized._nvfp4_alpha.device.type == "cpu" + assert plain.weight.device.type == "cpu" + + +def test_nvfp4_linear_hands_mm_fp4_transposed_operands_and_a_bf16_output(monkeypatch) -> None: + x_fp4 = torch.zeros(4, 64, dtype=torch.uint8) + x_scale = torch.zeros(128, 8, dtype=torch.uint8) + weight_packed = torch.zeros(256, 64, dtype=torch.uint8) + weight_scale = torch.zeros(256, 8, dtype=torch.uint8) + alpha = torch.tensor(0.5, dtype=torch.float32) + receipt: dict[str, object] = {} + + def fake_mm_fp4(a, b, a_scale, b_scale, alpha_arg, out_dtype, out, **kwargs): + receipt.update(a=a, b=b, a_scale=a_scale, b_scale=b_scale, alpha=alpha_arg, out_dtype=out_dtype, out=out, + kwargs=kwargs) + return torch.zeros(a.shape[0], b.shape[1], dtype=out_dtype) + + monkeypatch.setattr(h3_nvfp4, "_mm_fp4", fake_mm_fp4) + output = h3_nvfp4._nvfp4_linear(x_fp4, x_scale, weight_packed, weight_scale, alpha) + + assert output.shape == (4, 256) + assert output.dtype == torch.bfloat16 + assert receipt["a"] is x_fp4 + assert receipt["a_scale"] is x_scale + assert receipt["b"].shape == (64, 256) + assert receipt["b"].data_ptr() == weight_packed.data_ptr() + assert receipt["b_scale"].shape == (8, 256) + assert receipt["b_scale"].data_ptr() == weight_scale.data_ptr() + assert receipt["alpha"] is alpha + assert receipt["out_dtype"] == torch.bfloat16 + assert receipt["out"] is None + assert receipt["kwargs"] == {"backend": "auto"} + + +def test_apply_flattens_quantizes_and_restores_the_leading_dims(distributed_setup, monkeypatch) -> None: + """The finalized forward on a 3-D fp32 input: coerce to bf16, quantize the 2-D view with the + unit global scale, GEMM with the layer's alpha, add the bias, restore [batch, seq, out].""" + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + layer = _language_linear(config, bias=True) + _fill_loaded(layer, global_scale=4.0) + layer.quant_method.process_weights_after_loading(layer) + layer.bias.data.fill_(3.0) + receipt: dict[str, object] = {} + + def fake_quantize(x_2d, global_scale): + receipt.update(x_2d=x_2d, global_scale=global_scale) + return torch.zeros(x_2d.shape[0], 64, dtype=torch.uint8), torch.zeros(128, 8, dtype=torch.uint8) + + def fake_linear(x_fp4, x_scale, weight_packed, weight_scale, alpha): + receipt.update(alpha=alpha, weight_packed=weight_packed, weight_scale=weight_scale) + return torch.zeros(x_fp4.shape[0], weight_packed.shape[0], dtype=torch.bfloat16) + + monkeypatch.setattr(h3_nvfp4, "_quantize_activation_nvfp4", fake_quantize) + monkeypatch.setattr(h3_nvfp4, "_nvfp4_linear", fake_linear) + output = MiniMaxH3SerializedNVFP4LinearMethod._apply_finalized(layer, torch.ones(1, 5, 128), layer.bias) + + assert receipt["x_2d"].shape == (5, 128) + assert receipt["x_2d"].dtype == torch.bfloat16 + assert receipt["global_scale"] is layer._nvfp4_x_global_scale + assert receipt["alpha"] is layer._nvfp4_alpha + assert receipt["weight_packed"] is layer.weight_packed + assert output.shape == (1, 5, 256) + assert output.dtype == torch.bfloat16 + assert torch.equal(output, torch.full((1, 5, 256), 3.0, dtype=torch.bfloat16)) + + +def _tiny_conditioner_config(keep_bf16: tuple[str, ...] = ("mlp.down_proj", )) -> MiniMaxH3Qwen3VLConfig: + """One language layer whose linears satisfy the FlashInfer tile geometry, and a small vision tower.""" + config = MiniMaxH3Qwen3VLConfig() + config.arch_config = MiniMaxH3Qwen3VLArchConfig( + vocab_size=64, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + output_hidden_state_index=1, + num_hidden_layers_override=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + rope_scaling={ + "mrope_interleaved": True, + "mrope_section": [32, 16, 16], + "rope_type": "default", + }, + vision_depth=1, + vision_hidden_size=64, + vision_intermediate_size=128, + vision_num_heads=1, + vision_deepstack_visual_indexes=(), + vision_out_hidden_size=128, + ) + config.quant_config = MiniMaxH3SerializedNVFP4Config.from_config( + _checkpoint_quantization_config(modules_to_not_convert=["model.visual", "lm_head", *keep_bf16])) + return config + + +def test_conditioner_loads_converter_named_tensors_end_to_end(distributed_setup) -> None: + """The real chain: a conditioner built with the NVFP4 config, checkpoint keys spelled the way the + converter writes them, ``load_weights``, the strict missing-tensor check, then the post-load hook.""" + model = MiniMaxH3Qwen3VLConditioner(_tiny_conditioner_config()) + layer = model.language_model.layers[0] + assert isinstance(layer.self_attn.q_proj.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) + assert isinstance(layer.self_attn.o_proj.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) + assert isinstance(layer.mlp.down_proj.quant_method, UnquantizedLinearMethod) + assert layer.self_attn.q_proj.weight is None and layer.mlp.down_proj.weight is not None + + packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES + checkpoint: dict[str, torch.Tensor] = {} + for name, param in model.named_parameters(): + if name.endswith("." + packed_name): + tensor = torch.full_like(param, 0x11) + elif name.endswith("." + scale_name): + tensor = torch.full_like(param, E4M3_ONE) + elif name.endswith("." + global_scale_name): + tensor = torch.full_like(param, 2.0) + else: + tensor = torch.zeros_like(param) + checkpoint["model." + name] = tensor + expected = {name for name, _ in model.named_parameters()} + assert sum(name.endswith("." + packed_name) for name in expected) == 6 + + loaded = model.load_weights(iter(checkpoint.items())) + assert loaded == expected + assert _process_quantized_text_encoder_weights(model, torch.device("cpu")) == 6 + assert layer.self_attn.q_proj._nvfp4_alpha.item() == pytest.approx(0.5) + assert layer.mlp.up_proj._nvfp4_alpha.item() == pytest.approx(0.5) + + # A tensor the checkpoint lacks is reported by name through the loaded set, which is what the + # loader's strict check subtracts, and the hook then refuses the unfilled layer. + missing = f"model.language_model.layers.0.mlp.up_proj.{scale_name}" + layer.mlp.up_proj.weight_scale.data.fill_(0xFF) + partial = {key: value for key, value in checkpoint.items() if key != missing} + loaded = model.load_weights(iter(partial.items())) + assert expected - loaded == {missing[len("model."):]} + with pytest.raises(ValueError, match="weight_scale was not loaded"): + _process_quantized_text_encoder_weights(model, torch.device("cpu")) + + wrong_dtype = dict(checkpoint) + wrong_dtype[missing] = wrong_dtype[missing].to(torch.float8_e4m3fn) + with pytest.raises(ValueError, match="declared dtype"): + model.load_weights(iter(wrong_dtype.items())) + + stray = dict(checkpoint) + stray["model.language_model.layers.0.self_attn.q_proj.weight"] = torch.zeros(128, 128) + with pytest.raises(ValueError, match="Unexpected"): + model.load_weights(iter(stray.items())) From bfc9c017977d1f428d43b6e580a0ff20503d442b Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:48:07 -0700 Subject: [PATCH 003/151] [feat]: convert the MiniMax H3 text encoder to NVFP4 (#1838) --- ...minimax_h3_text_encoder_nvfp4_converter.py | 125 +++++ .../convert_minimax_h3_text_encoder_nvfp4.py | 435 ++++++++++++++++++ tests/local_tests/minimax_h3/README.md | 4 + ...st_minimax_h3_text_encoder_nvfp4_parity.py | 149 ++++++ 4 files changed, 713 insertions(+) create mode 100644 fastvideo/tests/encoders/test_minimax_h3_text_encoder_nvfp4_converter.py create mode 100644 scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py create mode 100644 tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py diff --git a/fastvideo/tests/encoders/test_minimax_h3_text_encoder_nvfp4_converter.py b/fastvideo/tests/encoders/test_minimax_h3_text_encoder_nvfp4_converter.py new file mode 100644 index 000000000..8992b2d3e --- /dev/null +++ b/fastvideo/tests/encoders/test_minimax_h3_text_encoder_nvfp4_converter.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU tests for the pure-Python parts of the NVFP4 text-encoder converter.""" +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import LANGUAGE_PROJECTIONS + +SCRIPT = Path(__file__).resolve().parents[3] / "scripts" / "checkpoint_conversion" / \ + "convert_minimax_h3_text_encoder_nvfp4.py" + + +@pytest.fixture(scope="module") +def converter(): + spec = importlib.util.spec_from_file_location("convert_minimax_h3_text_encoder_nvfp4", SCRIPT) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_language_linear_regex_matches_only_the_seven_projections(converter) -> None: + for proj in LANGUAGE_PROJECTIONS: + match = converter.LANGUAGE_LINEAR.match(f"model.language_model.layers.3.{proj}.weight") + assert match is not None and match["proj"] == proj and match["layer"] == "3" + for key in ( + "model.language_model.layers.3.self_attn.q_proj.bias", + "model.language_model.layers.3.self_attn.q_norm.weight", + "model.language_model.embed_tokens.weight", + "model.visual.blocks.0.attn.proj.weight", + "model.layers.3.self_attn.q_proj.weight", + ): + assert converter.LANGUAGE_LINEAR.match(key) is None, key + assert converter.LANGUAGE_LAYER.match("model.language_model.layers.12.mlp.down_proj.weight")["layer"] == "12" + + +def test_layer_plan_follows_the_conditioner_build_depth(converter) -> None: + kept, total, tap = converter.layer_plan({"text_config": {}}) + assert (kept, total, tap) == (50, 64, 50) + kept, total, tap = converter.layer_plan({"text_config": {"num_hidden_layers": 60}}) + assert (kept, total, tap) == (50, 60, 50) + with pytest.raises(ValueError, match="output_hidden_state_index"): + converter.layer_plan({"text_config": {"num_hidden_layers": 40}}) + + +def test_scan_source_counts_only_language_linears_below_the_kept_depth(converter, tmp_path) -> None: + shard = tmp_path / "model-00001-of-00001.safetensors" + tensors = { + "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), + "model.language_model.layers.1.self_attn.q_proj.weight": torch.zeros(2, 2), + "model.language_model.layers.1.mlp.down_proj.weight": torch.zeros(2, 2), + "model.language_model.layers.1.self_attn.q_norm.weight": torch.zeros(2), + "model.visual.blocks.0.attn.proj.weight": torch.zeros(2, 2), + "lm_head.weight": torch.zeros(2, 2), + } + save_file(tensors, str(shard), metadata={"format": "pt"}) + + counts, highest_layer, layers_seen = converter.scan_source([str(shard)], kept=2) + assert counts == {"self_attn.q_proj": 2, "mlp.down_proj": 1} + assert highest_layer == 1 + assert layers_seen["self_attn.q_proj"] == {0, 1} + assert layers_seen["mlp.down_proj"] == {1} + assert layers_seen["mlp.up_proj"] == set() + + counts, highest_layer, _ = converter.scan_source([str(shard)], kept=1) + assert counts == {"self_attn.q_proj": 1} + assert highest_layer == 1 + + +def test_shard_writer_uses_hub_names_writes_an_index_and_can_abort(converter, tmp_path) -> None: + writer = converter.ShardWriter(tmp_path, shard_bytes=8) + writer.add("a", torch.zeros(4, dtype=torch.uint8)) + writer.add("b", torch.zeros(4, dtype=torch.uint8)) + writer.add("c", torch.zeros(4, dtype=torch.uint8)) + writer.finish() + + names = sorted(path.name for path in tmp_path.glob("*.safetensors")) + assert names == ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"] + index = json.loads((tmp_path / converter.SAFE_WEIGHTS_INDEX_NAME).read_text()) + assert index["metadata"]["total_size"] == 12 + assert index["weight_map"] == { + "a": "model-00001-of-00002.safetensors", + "b": "model-00001-of-00002.safetensors", + "c": "model-00002-of-00002.safetensors", + } + + writer.abort() + assert not list(tmp_path.glob("*.safetensors")) + + dry = converter.ShardWriter(None, shard_bytes=8) + dry.add("a", torch.zeros(4, dtype=torch.uint8)) + dry.finish() + assert dry.total_bytes == 4 and not list(tmp_path.glob("*.safetensors")) + + +def test_argument_validation_rejects_unusable_requests(converter, monkeypatch, tmp_path) -> None: + base = [str(SCRIPT), "--src", str(tmp_path), "--dst", str(tmp_path / "out")] + + monkeypatch.setattr(sys, "argv", base + ["--keep-bf16", "mlp.dwn_proj"]) + with pytest.raises(SystemExit): + converter.parse_args() + + monkeypatch.setattr(sys, "argv", base + ["--keep-bf16", ",".join(LANGUAGE_PROJECTIONS)]) + with pytest.raises(SystemExit): + converter.parse_args() + + monkeypatch.setattr(sys, "argv", base + ["--probe-rows", "-1"]) + with pytest.raises(SystemExit): + converter.parse_args() + + monkeypatch.setattr(sys, "argv", [str(SCRIPT), "--src", str(tmp_path)]) + with pytest.raises(SystemExit): + converter.parse_args() + + monkeypatch.setattr(sys, "argv", base + ["--keep-bf16", "mlp.down_proj,mlp.down_proj"]) + args = converter.parse_args() + assert args.keep_bf16 == ("mlp.down_proj", ) + assert args.max_probe_error == 0.5 diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py b/scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py new file mode 100644 index 000000000..4c47d0dbd --- /dev/null +++ b/scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialize the MiniMax-H3 Qwen3-VL text encoder as NVFP4. + +H3 conditions on hidden state 50 of a 64-layer Qwen3-VL and never predicts a +token, so the conditioner builds 50 decoder layers, no ``lm_head``, and reads +nothing above (#1711). Every linear in those layers is quantized with +``flashinfer.nvfp4_quantize`` exactly as ``convert_model_to_nvfp4`` would at +runtime and stored as ``weight_packed`` / ``weight_scale`` / +``weight_global_scale``; the byte layout is documented in +``fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py``. Everything else +(token embedding, norms, the vision tower) is copied unchanged. ``config.json`` +gains the ``quantization_config`` block that makes the loader select the +serialized NVFP4 path and a ``num_hidden_layers_override`` equal to the kept +layer count, so the loader builds exactly the layers the file holds. + +For the FastH3 conditioner: 50 layers x 487.6M values = 24.4B values, 12.2 GB +packed plus 1.5 GB of scales, next to about 3 GB of unquantized tensors, +against 48.8 GB of bf16 language layers (63 GB on disk with the 14 dropped +layers and ``lm_head``). + +Needs a Blackwell GPU with FlashInfer: the quantizer is a CUDA kernel, and the +error report runs the same ``mm_fp4`` path the loader executes. + +Usage:: + + python scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py \ + --src /path/to/FastH3/text_encoder \ + --dst /path/to/FastH3-nvfp4/text_encoder --keep-bf16 mlp.down_proj + + # use it with any FastH3 model directory + ComponentConfig(text_encoder_weights="/path/to/FastH3-nvfp4/text_encoder") + +``--keep-bf16 mlp.down_proj`` leaves one projection kind in bf16 in every +layer and records it in ``modules_to_not_convert`` so the loader builds those +linears unquantized. The SwiGLU product feeding ``down_proj`` carries the +widest activation outliers in the stack; measured on one GB10 it is the +projection whose 4-bit form moves the layer-50 output most. + +Every quantized linear is probed: random bf16 rows go through the same +``mm_fp4`` path the loader runs and the relative error against the bf16 +product must stay under ``--max-probe-error``, otherwise the conversion stops +and removes what it wrote. The genuine W4A4 noise on random inputs is about +0.13; a wrong scale layout reads as 1.0. ``--report-only`` runs the probe and +writes nothing. The kept layer count is the one the conditioner builds, 50 +for H3, and is also written into ``config.json`` for a directory used as the +``text_encoder`` component; when the checkpoint is supplied through +``text_encoder_weights`` the model structure comes from the model directory, +which builds the same 50 layers. +""" + +from __future__ import annotations + +import argparse +import json +import re +import shutil +import time +from collections import Counter +from importlib import metadata as importlib_metadata +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig +from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import ( + LANGUAGE_PROJECTIONS, + NVFP4_TENSOR_SUFFIXES, + _nvfp4_linear, + _quantize_activation_nvfp4, + nvfp4_packed_weight_shape, + nvfp4_scale_shape, + nvfp4_weight_global_scale, + serialized_nvfp4_quantization_config, + validate_nvfp4_geometry, +) +from fastvideo.models.loader.weight_utils import SAFE_WEIGHTS_INDEX_NAME, resolve_safetensors_files + +LANGUAGE_LINEAR = re.compile(r"^model\.language_model\.layers\.(?P\d+)\.(?P" + + "|".join(re.escape(name) for name in LANGUAGE_PROJECTIONS) + r")\.weight$") +LANGUAGE_LAYER = re.compile(r"^model\.language_model\.layers\.(?P\d+)\.") + +# Keys the converter drops on purpose: id -> (pattern, reason). +SKIPPED = { + "lm_head": ("lm_head.weight", "the conditioner never predicts tokens; H3 reads a hidden state"), + "upper_layers": ("model.language_model.layers.{N >= kept}.*", + "layers above the one H3 reads are never built (#1711)"), +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--src", type=Path, required=True, help="text_encoder directory of the bf16 checkpoint") + parser.add_argument("--dst", type=Path, help="text_encoder directory to write (required unless --report-only)") + parser.add_argument("--keep-bf16", action="append", default=[], metavar="PROJ", + help="projection kind to leave in bf16 in every layer, e.g. mlp.down_proj; " + "repeat or comma-separate for several") + parser.add_argument("--device", default="cuda", help="CUDA device that runs the quantizer") + parser.add_argument("--shard-size-gb", type=float, default=4.0, help="safetensors shard size") + parser.add_argument("--probe-rows", type=int, default=512, + help="rows of random activations per linear for the error probe, 0 disables it") + parser.add_argument("--max-probe-error", type=float, default=0.5, + help="stop when a linear's FP4 GEMM relative error exceeds this; W4A4 noise is about 0.13") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--report-only", action="store_true", help="quantize and report errors, write nothing") + args = parser.parse_args() + if not args.report_only and args.dst is None: + parser.error("--dst is required unless --report-only is given") + if args.probe_rows < 0: + parser.error("--probe-rows must be 0 or positive") + if args.shard_size_gb <= 0: + parser.error("--shard-size-gb must be positive") + if args.max_probe_error <= 0: + parser.error("--max-probe-error must be positive") + keep = [name.strip() for entry in args.keep_bf16 for name in entry.split(",") if name.strip()] + unknown = sorted(set(keep) - set(LANGUAGE_PROJECTIONS)) + if unknown: + parser.error(f"--keep-bf16 got unknown projection(s) {unknown}; choose from {list(LANGUAGE_PROJECTIONS)}") + args.keep_bf16 = tuple(dict.fromkeys(keep)) + if set(args.keep_bf16) == set(LANGUAGE_PROJECTIONS): + parser.error("--keep-bf16 names every projection kind, so nothing would be quantized; " + "use the bf16 checkpoint as is") + return args + + +def layer_plan(source_config: dict) -> tuple[int, int, int]: + """Return (kept, total, tap): the language layers the conditioner builds from + this config.json, how many the source has, and the hidden state H3 reads, + derived the way the model derives them.""" + model_config = MiniMaxH3Qwen3VLConfig() + model_config.update_model_arch(source_config) + arch = model_config.arch_config + total = int(arch.num_hidden_layers) + tap = int(arch.output_hidden_state_index) + override = getattr(arch, "num_hidden_layers_override", None) + kept = min(total, int(override)) if override else total + return kept, total, tap + + +def scan_source(shards: list[str], kept: int) -> tuple[Counter, int, dict[str, set[int]]]: + """Count the language linears the conversion will touch before any work starts. + + Returns per-projection counts below ``kept``, the highest language layer index + seen, and the layer indexes present per projection for diagnostics. + """ + counts: Counter = Counter() + layers_seen: dict[str, set[int]] = {proj: set() for proj in LANGUAGE_PROJECTIONS} + highest_layer = -1 + for shard in shards: + with safe_open(shard, framework="pt", device="cpu") as handle: + for key in handle.keys(): + layer_match = LANGUAGE_LAYER.match(key) + if layer_match is not None: + highest_layer = max(highest_layer, int(layer_match["layer"])) + linear_match = LANGUAGE_LINEAR.match(key) + if linear_match is not None and int(linear_match["layer"]) < kept: + counts[linear_match["proj"]] += 1 + layers_seen[linear_match["proj"]].add(int(linear_match["layer"])) + return counts, highest_layer, layers_seen + + +class ShardWriter: + """Accumulate tensors and flush them as fixed-size safetensors shards plus an index. + + Shards are renamed to the ``model-00001-of-0000N.safetensors`` convention + once the count is known, so the output looks like any Hub checkpoint. + """ + + def __init__(self, dst: Path | None, shard_bytes: int) -> None: + self.dst = dst + self.shard_bytes = shard_bytes + self.pending: dict[str, torch.Tensor] = {} + self.pending_bytes = 0 + self.weight_map: dict[str, str] = {} + self.total_bytes = 0 + self.shard_index = 0 + self.written: list[Path] = [] + + def add(self, name: str, tensor: torch.Tensor) -> None: + nbytes = tensor.numel() * tensor.element_size() + if self.pending and self.pending_bytes + nbytes > self.shard_bytes: + self.flush() + self.pending[name] = tensor + self.pending_bytes += nbytes + self.total_bytes += nbytes + + def flush(self) -> None: + if not self.pending: + return + self.shard_index += 1 + filename = f"model-{self.shard_index:05d}.partial.safetensors" + if self.dst is not None: + save_file(self.pending, str(self.dst / filename), metadata={"format": "pt"}) + self.written.append(self.dst / filename) + for name in self.pending: + self.weight_map[name] = filename + self.pending = {} + self.pending_bytes = 0 + + def finish(self) -> None: + self.flush() + if self.dst is None: + return + final_names = {} + for index in range(1, self.shard_index + 1): + partial = f"model-{index:05d}.partial.safetensors" + final = f"model-{index:05d}-of-{self.shard_index:05d}.safetensors" + (self.dst / partial).rename(self.dst / final) + final_names[partial] = final + self.written = [self.dst / final for final in final_names.values()] + weight_map = {name: final_names[partial] for name, partial in self.weight_map.items()} + index_payload = {"metadata": {"total_size": self.total_bytes}, "weight_map": weight_map} + (self.dst / SAFE_WEIGHTS_INDEX_NAME).write_text(json.dumps(index_payload, indent=2, sort_keys=True), + encoding="utf-8") + + def abort(self) -> None: + """Remove every shard this writer produced so a failed conversion leaves no half checkpoint.""" + for path in self.written: + path.unlink(missing_ok=True) + self.written = [] + + +def probe_relative_error( + weight: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + global_scale: torch.Tensor, + rows: int, + generator: torch.Generator, +) -> float: + """||fp4(x) @ fp4(W).T - x @ W.T|| / ||x @ W.T|| on random bf16 rows, through the loader's own GEMM path. + + The reference is the bf16 tensor-core product with fp32 accumulation; its + rounding sits two orders of magnitude below the FP4 error being measured. + """ + x = torch.randn(rows, weight.shape[1], generator=generator, device=weight.device, + dtype=torch.float32).to(torch.bfloat16) + reference = (x @ weight.t()).float() + x_fp4, x_scale = _quantize_activation_nvfp4(x, torch.ones((), dtype=torch.float32, device=weight.device)) + output = _nvfp4_linear(x_fp4, x_scale, packed, scale, (1.0 / global_scale).reshape(())) + return ((output.float() - reference).norm() / reference.norm().clamp_min(1e-12)).item() + + +def quantize_language_linear( + weight: torch.Tensor, + device: torch.device, + sf_layout: object, + probe_rows: int, + generator: torch.Generator | None, + max_probe_error: float, + key: str = "", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float | None]: + if weight.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"Expected a bf16, fp16 or fp32 source weight, got {weight.dtype}; " + "the source text_encoder must be unquantized") + output_size, input_size = weight.shape + validate_nvfp4_geometry(output_size, input_size) + weight_device = weight.to(device=device, dtype=torch.bfloat16) + global_scale = nvfp4_weight_global_scale(weight_device) + packed, scale = _nvfp4_quantize(weight_device, global_scale, sfLayout=sf_layout.layout_128x4, do_shuffle=False) + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + if scale.dtype != torch.uint8: + scale = scale.view(torch.uint8) + expected_packed = nvfp4_packed_weight_shape(output_size, input_size) + expected_scale = nvfp4_scale_shape(output_size, input_size) + if tuple(packed.shape) != expected_packed or tuple(scale.shape) != expected_scale: + raise RuntimeError("flashinfer.nvfp4_quantize returned an unexpected layout for a " + f"{output_size}x{input_size} weight: packed {tuple(packed.shape)} " + f"(expected {expected_packed}), scale {tuple(scale.shape)} (expected {expected_scale}). " + "The loader allocates exactly these shapes; check the FlashInfer version.") + error = None + if probe_rows and generator is not None: + error = probe_relative_error(weight_device, packed, scale, global_scale, probe_rows, generator) + if not error <= max_probe_error: + raise SystemExit(f"FP4 GEMM relative error {error:.3e} on {key or 'a language linear'} exceeds " + f"--max-probe-error {max_probe_error}; the packed layout or scales do not reproduce " + "the bf16 product, nothing was written") + return packed.cpu().contiguous(), scale.cpu().contiguous(), global_scale.reshape(1).cpu(), error + + +def flashinfer_version() -> str: + try: + return importlib_metadata.version("flashinfer-python") + except importlib_metadata.PackageNotFoundError: + import flashinfer + return str(getattr(flashinfer, "__version__", "unknown")) + + +def main() -> None: + args = parse_args() + src: Path = args.src + config_path = src / "config.json" + if not config_path.is_file(): + raise SystemExit(f"{src} has no config.json; point --src at the text_encoder directory") + source_config = json.loads(config_path.read_text(encoding="utf-8")) + if source_config.get("quantization_config"): + raise SystemExit(f"{src} already carries a quantization_config " + f"({source_config['quantization_config'].get('quant_method')!r}); this converter " + "quantizes the bf16 text_encoder, not a serialized one") + kept, total, tap = layer_plan(source_config) + + shards = resolve_safetensors_files(str(src)) + counts, highest_layer, layers_seen = scan_source(shards, kept) + if not counts: + raise SystemExit(f"No language linear in {src} matches model.language_model.layers.N..weight; " + "this converter expects the Qwen3-VL checkpoint layout") + if highest_layer + 1 < kept: + raise SystemExit(f"{src} holds {highest_layer + 1} language layers but the conditioner builds {kept}") + short = {} + for proj in LANGUAGE_PROJECTIONS: + missing = sorted(set(range(kept)) - layers_seen[proj]) + if missing: + short[proj] = f"{len(missing)} missing, e.g. model.language_model.layers.{missing[0]}.{proj}.weight" + if short: + raise SystemExit(f"Expected {kept} of every language projection below layer {kept}: {short}") + + device = torch.device(args.device) + if device.type != "cuda": + raise SystemExit("The NVFP4 quantizer is a CUDA kernel; pass --device cuda") + sf_layout, _, _ = _require_flashinfer() + + dst: Path | None = None + if not args.report_only: + dst = args.dst + dst.mkdir(parents=True, exist_ok=True) + existing = sorted(path.name for path in dst.glob("*.safetensors")) + \ + [name for name in ("config.json", SAFE_WEIGHTS_INDEX_NAME) if (dst / name).exists()] + if existing: + raise SystemExit(f"{dst} already holds {existing[:4]}{' and more' if len(existing) > 4 else ''}; " + "refusing to overwrite or mix outputs") + writer = ShardWriter(dst, int(args.shard_size_gb * (1 << 30))) + generator = torch.Generator(device=device).manual_seed(args.seed) if args.probe_rows else None + + quantized = 0 + kept_bf16 = 0 + kept_bf16_bytes = 0 + copied = 0 + skipped: Counter = Counter() + source_language_bytes = 0 + written_language_bytes = 0 + errors: dict[str, list[float]] = {} + started = time.perf_counter() + print(f"keeping {kept} of {total} language layers, H3 reads hidden state {tap}; " + f"{sum(counts.values())} language linears to visit", flush=True) + packed_name, scale_name, global_scale_name = NVFP4_TENSOR_SUFFIXES + + try: + for shard in shards: + with safe_open(shard, framework="pt", device="cpu") as handle: + for key in sorted(handle.keys()): + layer_match = LANGUAGE_LAYER.match(key) + if key == "lm_head.weight": + skipped["lm_head"] += 1 + continue + if layer_match is not None and int(layer_match["layer"]) >= kept: + skipped["upper_layers"] += 1 + continue + tensor = handle.get_tensor(key) + linear_match = LANGUAGE_LINEAR.match(key) + if linear_match is None: + writer.add(key, tensor.contiguous()) + copied += 1 + continue + if linear_match["proj"] in args.keep_bf16: + # Kept linears are stored as bf16 whatever the source dtype, matching the + # precision the conditioner runs them in. + writer.add(key, tensor.to(torch.bfloat16).contiguous()) + kept_bf16 += 1 + kept_bf16_bytes += tensor.numel() * 2 + continue + packed, scale, global_scale, error = quantize_language_linear( + tensor, device, sf_layout, args.probe_rows, generator, args.max_probe_error, key) + prefix = key[:-len(".weight")] + writer.add(f"{prefix}.{packed_name}", packed) + writer.add(f"{prefix}.{scale_name}", scale) + writer.add(f"{prefix}.{global_scale_name}", global_scale) + source_language_bytes += tensor.numel() * tensor.element_size() + written_language_bytes += packed.numel() + scale.numel() + 4 + quantized += 1 + if error is not None: + errors.setdefault(linear_match["proj"], []).append(error) + if quantized % 35 == 0: + print(f" quantized {quantized} language linears, {time.perf_counter() - started:.0f}s", + flush=True) + writer.finish() + except BaseException: + writer.abort() + raise + + if dst is not None: + config = dict(source_config) + config["num_hidden_layers_override"] = kept + config["quantization_config"] = serialized_nvfp4_quantization_config( + keep_bf16=args.keep_bf16, + producer={ + "converter": Path(__file__).name, + "flashinfer": flashinfer_version(), + "kept_language_layers": kept, + }) + (dst / "config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") + for extra in src.iterdir(): + if extra.is_file() and extra.name not in {SAFE_WEIGHTS_INDEX_NAME, "config.json"} \ + and not extra.name.endswith(".safetensors"): + shutil.copy2(extra, dst / extra.name) + + print(f"language linears quantized: {quantized}") + if args.keep_bf16: + print(f"language linears kept bf16: {kept_bf16} ({', '.join(args.keep_bf16)}; " + f"{kept_bf16_bytes / 1e9:.2f} GB)") + print(f"tensors copied unchanged: {copied}") + for skip_id, count in sorted(skipped.items()): + pattern, reason = SKIPPED[skip_id] + print(f"skipped {count:>4} keys: {pattern}: {reason}") + print(f"language linear bytes: {source_language_bytes / 1e9:.2f} GB bf16 -> " + f"{written_language_bytes / 1e9:.2f} GB NVFP4 (packed values + E4M3 scales)") + print(f"artifact total: {writer.total_bytes / 1e9:.2f} GB in {writer.shard_index} shard(s)" + f"{'' if dst is not None else ' (not written, --report-only)'}") + if errors: + print(f"FP4 GEMM relative error vs bf16, {args.probe_rows} random rows per linear, " + f"every linear under --max-probe-error {args.max_probe_error}:") + print(f" {'projection':<20}{'linears':>8}{'max':>12}{'mean':>12}") + for proj in LANGUAGE_PROJECTIONS: + values = errors.get(proj) + if values: + print(f" {proj:<20}{len(values):>8}{max(values):>12.3e}{sum(values) / len(values):>12.3e}") + elif quantized: + print("FP4 GEMM error probe skipped because --probe-rows is 0; nothing verified the written layout") + print(f"done in {time.perf_counter() - started:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/minimax_h3/README.md b/tests/local_tests/minimax_h3/README.md index f1bd8b7e2..ee596739a 100644 --- a/tests/local_tests/minimax_h3/README.md +++ b/tests/local_tests/minimax_h3/README.md @@ -46,6 +46,10 @@ PYTHONPATH="$MINIMAX_H3_OFFICIAL_REF_DIR/src:$PWD" \ MINIMAX_H3_RUN_ENCODER_PARITY=1 \ pytest tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py -v -s +MINIMAX_H3_RUN_NVFP4_PARITY=1 MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ +MINIMAX_H3_NVFP4_TEXT_ENCODER=/path/to/FastH3-text-encoder-nvfp4 \ +pytest tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py -s + PYTHONPATH="$MINIMAX_H3_OFFICIAL_REF_DIR/src:$PWD" \ MINIMAX_H3_RUN_DIT_PARITY=1 \ MINIMAX_H3_RUN_VIDEO_VAE_PARITY=1 \ diff --git a/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py b/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py new file mode 100644 index 000000000..b47132898 --- /dev/null +++ b/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke test for a converted NVFP4 MiniMax-H3 text encoder. + +Loads the bf16 conditioner and a checkpoint written by +``scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py`` +through the production ``TextEncoderLoader`` and compares the layer-50 hidden +states on a few prompts. A layout, transposition or scale regression in the +serialized NVFP4 path shows up here as a cosine near zero; genuine 4-bit +noise stays above the thresholds. + +Run on a Blackwell GPU:: + + MINIMAX_H3_RUN_NVFP4_PARITY=1 \ + MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ + MINIMAX_H3_NVFP4_TEXT_ENCODER=/path/to/FastH3-nvfp4/text_encoder \ + pytest tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py -s +""" + +from __future__ import annotations + +import gc +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig +from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4LinearMethod +from fastvideo.models.loader.component_loader import TextEncoderLoader + +PROMPTS = ( + "an alpine lake at sunrise, gentle wind over the water", + "a red fox crosses fresh snow at sunrise", + "a busy night market in Taipei, neon signs reflecting on wet pavement, handheld camera", + # 170 tokens with the Qwen3-VL tokenizer, so the activation path crosses a 128-row FlashInfer tile. + # Natural text on purpose: a clause repeated three times scores 0.975 mean cosine because attention + # between identical copies is a tie that any noise can flip, which says nothing about real prompts. + "A slow cinematic drone shot glides over a coastal town at golden hour. Terracotta rooftops catch the " + "last light while laundry lines sway between balconies. A fishing boat returns to the harbor with gulls " + "trailing behind it, and children chase a ball down a cobblestone lane past a bakery whose windows are " + "fogged with steam. An old man reads on a bench under a lemon tree as waves break softly against the " + "pier and a church bell rings twice. Warm haze softens the distant hills, a tram crosses a stone bridge, " + "and a street musician tunes a guitar near a fountain where pigeons drink. The camera finally settles " + "on a cafe terrace where a waiter pours coffee, a couple shares a newspaper, and a cat stretches in the " + "last patch of sunlight before the shadows reach the water and the harbor lights flicker on one by one.", +) +MIN_MEAN_COSINE = float(os.environ.get("MINIMAX_H3_NVFP4_MIN_COSINE", "0.99")) +# Individual low-content tokens such as an article or a word fragment rotate easily under 4-bit noise +# while carrying almost no signal, so the per-token criterion is a share, not a minimum. +LOW_TOKEN_COSINE = 0.90 +MAX_LOW_TOKEN_SHARE = float(os.environ.get("MINIMAX_H3_NVFP4_MAX_LOW_TOKEN_SHARE", "0.02")) +MAX_RELATIVE_ERROR = float(os.environ.get("MINIMAX_H3_NVFP4_MAX_REL_ERR", "0.10")) + + +def _require_assets() -> tuple[torch.device, Path, Path]: + if os.environ.get("MINIMAX_H3_RUN_NVFP4_PARITY") != "1": + pytest.skip("set MINIMAX_H3_RUN_NVFP4_PARITY=1 on a Blackwell GPU node") + if not torch.cuda.is_available(): + pytest.fail("MiniMax-H3 NVFP4 parity requires a CUDA GPU", pytrace=False) + root = os.environ.get("MINIMAX_H3_MODEL_ROOT") + converted = os.environ.get("MINIMAX_H3_NVFP4_TEXT_ENCODER") + if not root or not converted: + pytest.fail("set MINIMAX_H3_MODEL_ROOT and MINIMAX_H3_NVFP4_TEXT_ENCODER", pytrace=False) + root_path, converted_path = Path(root), Path(converted) + missing = [str(p) for p in (root_path / "text_encoder", root_path / "tokenizer", converted_path) if not p.is_dir()] + if missing: + pytest.fail(f"directories are missing: {missing}", pytrace=False) + if not (converted_path / "config.json").is_file(): + pytest.fail(f"{converted_path} has no config.json; point at the converted text_encoder directory", + pytrace=False) + return torch.device("cuda"), root_path, converted_path + + +@pytest.fixture(scope="module", autouse=True) +def _distributed_runtime(): + if os.environ.get("MINIMAX_H3_RUN_NVFP4_PARITY") != "1": + yield + return + for name, value in { + "MASTER_ADDR": "127.0.0.1", + "MASTER_PORT": "29624", + "RANK": "0", + "WORLD_SIZE": "1", + "LOCAL_RANK": "0" + }.items(): + os.environ.setdefault(name, value) + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() + + +def _loader_args() -> SimpleNamespace: + return SimpleNamespace( + pipeline_config=SimpleNamespace( + text_encoder_configs=(MiniMaxH3Qwen3VLConfig(), ), + text_encoder_precisions=("bf16", ), + ), + text_encoder_cpu_offload=False, + override_text_encoder_quant=None, + override_text_encoder_safetensors=None, + pin_cpu_memory=False, + disable_offload_on_unified_memory=lambda device_id=0, *, offload_flag=None: True, + ) + + +def _encode(text_encoder_dir: Path, tokenizer, device: torch.device, expect_nvfp4: bool, + prompts: tuple[str, ...]) -> list[torch.Tensor]: + model = TextEncoderLoader().load(str(text_encoder_dir), _loader_args()) + probe = model.language_model.layers[0].self_attn.q_proj + is_nvfp4 = isinstance(probe.quant_method, MiniMaxH3SerializedNVFP4LinearMethod) and probe.weight is None + assert is_nvfp4 == expect_nvfp4, (f"{text_encoder_dir} loaded {'through' if is_nvfp4 else 'without'} the " + "serialized NVFP4 path") + outputs = [] + for prompt in prompts: + ids = torch.tensor(tokenizer(prompt, add_special_tokens=False)["input_ids"], dtype=torch.long, device=device) + with torch.inference_mode(): + outputs.append(model(input_ids=ids).float().cpu()) + del model + gc.collect() + torch.cuda.empty_cache() + return outputs + + +def test_minimax_h3_text_encoder_nvfp4_parity() -> None: + from transformers import AutoTokenizer + + device, root, converted = _require_assets() + tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer", local_files_only=True) + prompts = PROMPTS + reference = _encode(root / "text_encoder", tokenizer, device, expect_nvfp4=False, prompts=prompts) + quantized = _encode(converted, tokenizer, device, expect_nvfp4=True, prompts=prompts) + assert reference[-1].shape[0] > 128, "the long prompt must cross a 128-row tile" + + print(f"\n{'prompt':<60}{'tokens':>7}{'cos mean':>10}{'cos min':>9}{'low share':>10}{'rel err':>10}", + flush=True) + for prompt, expected, actual in zip(prompts, reference, quantized, strict=True): + assert actual.shape == expected.shape + cosine = torch.nn.functional.cosine_similarity(expected, actual, dim=-1) + low_share = (cosine < LOW_TOKEN_COSINE).float().mean().item() + relative_error = ((expected - actual).norm() / expected.norm()).item() + print(f"{prompt[:58]:<60}{expected.shape[0]:>7}{cosine.mean().item():>10.4f}{cosine.min().item():>9.4f}" + f"{low_share:>10.3f}{relative_error:>10.3e}", flush=True) + assert cosine.mean().item() >= MIN_MEAN_COSINE, f"{prompt[:40]!r}: mean cosine {cosine.mean().item():.4f}" + assert low_share <= MAX_LOW_TOKEN_SHARE, (f"{prompt[:40]!r}: {low_share:.1%} of tokens below " + f"{LOW_TOKEN_COSINE} cosine") + assert relative_error <= MAX_RELATIVE_ERROR, f"{prompt[:40]!r}: relative error {relative_error:.3e}" From 9a3c45779cd4f869e619913076037b4519a5332b Mon Sep 17 00:00:00 2001 From: Yaegaki1Erika <70182590+Yaegaki1Erika@users.noreply.github.com> Date: Tue, 15 Sep 2026 06:58:13 +0800 Subject: [PATCH 004/151] [bugfix]: fix Wan I2V DMD conditioning under sequence parallelism (#1844) The pipeline pre-sharded only the image conditioning along the temporal axis while the noise input stayed full-length, so concatenation could not match for sp_world_size > 1. WanTransformer3DModel shards the flattened token sequence after patch embedding, so pass full-length mask+latent conditioning and let the transformer shard. Also reads temporal_compression_ratio from the VAE config, simplifies the conditioning mask, concatenates in the transformer's (bs, c, t, h, w) layout, and adds unit coverage. Co-authored-by: Yaegaki1Erika <70182590+Yaegaki1Erika@users.noreply.github.com> --- .../distill/test_wan_i2v_distillation.py | 77 +++++++++++++++++++ .../training/wan_i2v_distillation_pipeline.py | 77 +++++++++++++------ 2 files changed, 130 insertions(+), 24 deletions(-) create mode 100644 fastvideo/tests/training/distill/test_wan_i2v_distillation.py diff --git a/fastvideo/tests/training/distill/test_wan_i2v_distillation.py b/fastvideo/tests/training/distill/test_wan_i2v_distillation.py new file mode 100644 index 000000000..279ff1249 --- /dev/null +++ b/fastvideo/tests/training/distill/test_wan_i2v_distillation.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.training import wan_i2v_distillation_pipeline +from fastvideo.training.distillation_pipeline import DistillationPipeline +from fastvideo.training.wan_i2v_distillation_pipeline import ( + WanI2VDistillationPipeline, ) + + +def test_i2v_prepare_inputs_does_not_pre_shard_conditioning(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = WanI2VDistillationPipeline.__new__(WanI2VDistillationPipeline) + arch_config = SimpleNamespace(temporal_compression_ratio=4) + pipeline.training_args = SimpleNamespace( + num_latent_t=4, + pipeline_config=SimpleNamespace(vae_config=SimpleNamespace(arch_config=arch_config)), + ) + pipeline.sp_world_size = 8 + pipeline.rank_in_sp_group = 0 + batch = SimpleNamespace(image_latents=torch.zeros(1, 16, 4, 2, 3)) + + monkeypatch.setattr(DistillationPipeline, "_prepare_dit_inputs", lambda self, training_batch: training_batch) + monkeypatch.setattr(wan_i2v_distillation_pipeline, "get_local_torch_device", lambda: torch.device("cpu")) + + result = pipeline._prepare_dit_inputs(batch) + + assert result.image_latents.shape == (1, 20, 4, 2, 3) + + +def test_i2v_conditioning_preserves_full_temporal_length_for_sequence_parallelism() -> None: + image_latents = torch.zeros(1, 16, 4, 2, 3) + + conditioning = WanI2VDistillationPipeline._build_image_conditioning( + image_latents, + num_latent_t=4, + ) + + assert conditioning.shape == (1, 20, 4, 2, 3) + torch.testing.assert_close(conditioning[:, :4, 0], torch.ones_like(conditioning[:, :4, 0])) + torch.testing.assert_close(conditioning[:, :4, 1:], torch.zeros_like(conditioning[:, :4, 1:])) + + +def test_i2v_conditioning_rejects_temporal_length_mismatch() -> None: + image_latents = torch.zeros(1, 16, 3, 2, 3) + + with pytest.raises(ValueError, match="temporal length"): + WanI2VDistillationPipeline._build_image_conditioning( + image_latents, + num_latent_t=4, + ) + + +def test_i2v_distill_input_kwargs_concatenates_conditioning_in_transformer_layout( + monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = WanI2VDistillationPipeline.__new__(WanI2VDistillationPipeline) + monkeypatch.setattr(wan_i2v_distillation_pipeline, "get_local_torch_device", lambda: torch.device("cpu")) + + noise_input = torch.zeros(1, 4, 16, 2, 3) # bs, t, c, h, w + image_latents = torch.ones(1, 20, 4, 2, 3) # bs, c, t, h, w + batch = SimpleNamespace(image_embeds=torch.zeros(1, 4), image_latents=image_latents) + text_dict = { + "encoder_hidden_states": torch.zeros(1, 4), + "encoder_attention_mask": torch.zeros(1, 4), + } + + result = pipeline._build_distill_input_kwargs(noise_input, torch.zeros(1), text_dict, batch) + + hidden_states = result.input_kwargs["hidden_states"] + assert hidden_states.shape == (1, 36, 4, 2, 3) + torch.testing.assert_close(hidden_states[:, :16], torch.zeros_like(hidden_states[:, :16])) + torch.testing.assert_close(hidden_states[:, 16:], torch.ones_like(hidden_states[:, 16:])) + assert result.noise_latents is noise_input diff --git a/fastvideo/training/wan_i2v_distillation_pipeline.py b/fastvideo/training/wan_i2v_distillation_pipeline.py index 5b1c834c4..821742663 100644 --- a/fastvideo/training/wan_i2v_distillation_pipeline.py +++ b/fastvideo/training/wan_i2v_distillation_pipeline.py @@ -4,7 +4,6 @@ from copy import deepcopy from typing import Any import torch -from einops import rearrange from fastvideo.api.sampling_param import SamplingParam from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v @@ -129,30 +128,49 @@ class WanI2VDistillationPipeline(DistillationPipeline): assert isinstance(training_batch.image_latents, torch.Tensor) image_latents = training_batch.image_latents.to(get_local_torch_device(), dtype=torch.bfloat16) - - temporal_compression_ratio = 4 - num_frames = (self.training_args.num_latent_t - 1) * temporal_compression_ratio + 1 - batch_size, num_channels, _, latent_height, latent_width = image_latents.shape - mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width) - mask_lat_size[:, :, 1:] = 0 - - first_frame_mask = mask_lat_size[:, :, :1] - first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=temporal_compression_ratio) - mask_lat_size = torch.cat([first_frame_mask, mask_lat_size[:, :, 1:]], dim=2) - mask_lat_size = mask_lat_size.view(batch_size, -1, temporal_compression_ratio, latent_height, latent_width) - mask_lat_size = mask_lat_size.transpose(1, 2) - mask_lat_size = mask_lat_size.to(image_latents.device).to(dtype=torch.bfloat16) - - image_latents = torch.cat([mask_lat_size, image_latents], dim=1) - training_batch.image_latents = image_latents - - if self.sp_world_size > 1: - image_latents = rearrange(image_latents, "b c (n t) h w -> b c n t h w", n=self.sp_world_size).contiguous() - image_latents = image_latents[:, :, self.rank_in_sp_group, :, :, :] - training_batch.image_latents = image_latents + vae_config = self.training_args.pipeline_config.vae_config.arch_config + training_batch.image_latents = self._build_image_conditioning( + image_latents, + num_latent_t=self.training_args.num_latent_t, + temporal_compression_ratio=vae_config.temporal_compression_ratio, + ) return training_batch + @staticmethod + def _build_image_conditioning( + image_latents: torch.Tensor, + *, + num_latent_t: int, + temporal_compression_ratio: int = 4, + ) -> torch.Tensor: + """Build full-length Wan I2V mask and latent conditioning. + + WanTransformer3DModel applies sequence-parallel sharding after patch + embedding. Pre-sharding only the image conditioning here would make its + temporal length differ from the full noise input before concatenation. + """ + batch_size, _, latent_t, latent_height, latent_width = image_latents.shape + if latent_t != num_latent_t: + raise ValueError("Wan I2V image conditioning temporal length must match num_latent_t: " + f"got image_latents.shape[2]={latent_t}, num_latent_t={num_latent_t}") + + # The Wan I2V conditioning mask is one on the first latent frame and zero + # elsewhere across the temporal-compression channels; the repeat/view/ + # transpose form used by the inference path reduces to this. + mask_lat_size = torch.zeros( + batch_size, + temporal_compression_ratio, + num_latent_t, + latent_height, + latent_width, + device=image_latents.device, + dtype=image_latents.dtype, + ) + mask_lat_size[:, :, 0] = 1 + + return torch.cat([mask_lat_size, image_latents], dim=1) + def _build_distill_input_kwargs(self, noise_input: torch.Tensor, timestep: torch.Tensor, text_dict: dict[str, torch.Tensor], training_batch: TrainingBatch) -> TrainingBatch: # Image Embeds for conditioning @@ -160,10 +178,21 @@ class WanI2VDistillationPipeline(DistillationPipeline): assert torch.isnan(image_embeds).sum() == 0 image_embeds = image_embeds.to(get_local_torch_device(), dtype=torch.bfloat16) - noisy_model_input = torch.cat([noise_input, training_batch.image_latents.permute(0, 2, 1, 3, 4)], dim=2) + image_conditioning = training_batch.image_latents # bs, c, t, h, w + noise_latents = noise_input.permute(0, 2, 1, 3, 4) # bs, c, t, h, w + matching_batch_and_time = (noise_latents.shape[0] == image_conditioning.shape[0] + and noise_latents.shape[2] == image_conditioning.shape[2]) + matching_spatial_shape = noise_latents.shape[3:] == image_conditioning.shape[3:] + if not matching_batch_and_time or not matching_spatial_shape: + raise ValueError("Wan I2V noise and image conditioning must have matching batch, temporal, and spatial " + f"dimensions, got noise={tuple(noise_latents.shape)} and " + f"conditioning={tuple(image_conditioning.shape)}") + # Concatenate in the transformer's (bs, c, t, h, w) layout so the patch + # embedding receives a contiguous tensor. + hidden_states = torch.cat([noise_latents, image_conditioning], dim=1) training_batch.input_kwargs = { - "hidden_states": noisy_model_input.permute(0, 2, 1, 3, 4), # bs, c, t, h, w + "hidden_states": hidden_states, "encoder_hidden_states": text_dict["encoder_hidden_states"], "encoder_attention_mask": text_dict["encoder_attention_mask"], "timestep": timestep, From 1c14afd5599a0033d2ffe5db75a4b0d1106b5c58 Mon Sep 17 00:00:00 2001 From: William Lin Date: Mon, 14 Sep 2026 16:06:53 -0700 Subject: [PATCH 005/151] [docs]: refresh AGENTS.md maps and add Wan SP/I2V and CI test notes (#1845) --- AGENTS.md | 4 ++++ fastvideo/AGENTS.md | 16 +++++++++++++-- fastvideo/attention/AGENTS.md | 10 ++++++++-- fastvideo/models/AGENTS.md | 6 ++++-- fastvideo/models/wan/AGENTS.md | 6 ++++++ fastvideo/pipelines/AGENTS.md | 9 +++++---- fastvideo/tests/AGENTS.md | 36 ++++++++++++++++++++++++++++++---- fastvideo/train/AGENTS.md | 13 ++++++++++-- fastvideo/training/AGENTS.md | 24 +++++++++++++++++++++-- 9 files changed, 106 insertions(+), 18 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ff18d09d5..f79037dcc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,14 +66,18 @@ Local guidance lives next to the code. Read the in-scope file before editing: | `fastvideo/AGENTS.md` | Core package map, public API, registry-driven model dispatch | | `fastvideo/configs/AGENTS.md` | Arch + pipeline config dataclasses, `param_names_mapping` | | `fastvideo/models/AGENTS.md` | DiT / VAE / encoder / scheduler / loader layout (pre-commit excluded) | +| `fastvideo/models/wan/AGENTS.md` | Wan family-local transformers, VAE, configs, and the SP sharding invariant | | `fastvideo/layers/AGENTS.md` | Tensor-parallel linear/attention layer rules for ports | | `fastvideo/attention/AGENTS.md` | Backend registry + env-var override | | `fastvideo/pipelines/AGENTS.md` | Stage ABC, `basic//`, `preprocess/`, presets | +| `fastvideo/pipelines/basic/wan/AGENTS.md` | Wan sampling stages, first-frame conditioning, DMD/causal boundaries | +| `fastvideo/pipelines/basic/magi_human/AGENTS.md` | MagiHuman umbrella repo, lazy-loaded components, packing invariants | | `fastvideo/training/AGENTS.md` | Legacy monolithic pipelines (frozen for existing models) | | `fastvideo/train/AGENTS.md` | New modular trainer (methods × models × callbacks, YAML) | | `fastvideo/tests/AGENTS.md` | Test taxonomy, conftest, pre-commit-excluded path | | `fastvideo/tests/ssim/AGENTS.md` | GPU SSIM regression authoring + reference video sync | | `scripts/checkpoint_conversion/AGENTS.md` | Adding a converter for a new HF/official checkpoint | +| `apps/dreamverse/AGENTS.md` | DreamVerse app structure and conventions | ## Critical: Two Training Stacks Coexist diff --git a/fastvideo/AGENTS.md b/fastvideo/AGENTS.md index ed20c552f..cac7ffb16 100644 --- a/fastvideo/AGENTS.md +++ b/fastvideo/AGENTS.md @@ -1,6 +1,6 @@ # `fastvideo/` — Core Package -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 Inference + training framework for video DiTs. Public API entry: `from fastvideo import VideoGenerator, PipelineConfig, SamplingParam`. @@ -20,15 +20,22 @@ CLI entry: `fastvideo` script → `entrypoints/cli/main.py` (subcommands: `gener fastvideo/ ├── api/ # Schema + presets for the OpenAI-compatible serving layer ├── attention/ # Backends + selector (FlashAttn / SageAttn / SDPA / VSA / VMoBA / SLA) +├── benchmarks/ # Benchmark suites (MLX FastWan, MetalFX RIFE) ├── configs/ # Per-model arch configs + per-pipeline configs (registry-driven) ├── dataset/ # Dataloaders (pre-commit excluded — minimal lint surface) ├── distributed/ # SP/TP groups, device communicators, init helpers ├── entrypoints/ # cli/, openai/, streaming/, video_generator.py +├── eval/ # Evaluation harnesses, datasets, and reference scores ├── hooks/ # Runtime hook system for pipelines ├── layers/ # Tensor-parallel linears + attention wrappers (port targets) +├── logging_utils/ # Log formatters shared by the CLI and serving paths +├── mlx_runtime/ # Apple MLX runtime (FastWan, checkpoint compat, upsampling) ├── models/ # Components + Wan family-local package (mostly pre-commit excluded) +├── performance/ # Performance metric policy + HF artifact store +├── performance_dashboard/ # Metrics service and dashboard API ├── pipelines/ # basic//, preprocess/, stages/, training/ ├── platforms/ # CUDA/ROCm capability + AttentionBackendEnum +├── tests/ # Package-level tests (see tests/AGENTS.md) ├── third_party/ # Vendored externals (lint excluded; do not reformat) ├── train/ # NEW modular trainer — methods × models × callbacks ├── training/ # LEGACY monolithic *_training/distillation_pipeline.py @@ -37,7 +44,12 @@ fastvideo/ ├── registry.py # Pipeline-config + model-class lookup (canonical) ├── envs.py # Env-var declarations ├── fastvideo_args.py# Runtime arg dataclass passed through pipelines -└── utils.py # FlexibleArgumentParser, qualname resolver, etc. +├── forward_context.py # Forward-time context (timestep, attention metadata) +├── image_processor.py # Image preprocessing for conditioning +├── logger.py # init_logger / logging setup +├── profiler.py # Profiler regions and trace helpers +├── utils.py # FlexibleArgumentParser, qualname resolver, etc. +└── version.py # Package version ``` ## Where to Look diff --git a/fastvideo/attention/AGENTS.md b/fastvideo/attention/AGENTS.md index c46005a23..9a33843a4 100644 --- a/fastvideo/attention/AGENTS.md +++ b/fastvideo/attention/AGENTS.md @@ -1,6 +1,6 @@ # `fastvideo/attention/` — Attention Backends -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 Backend registry + selector wrapping FlashAttn / SageAttn / SageAttn3 / SDPA / VSA / VMoBA / SLA / BSA. @@ -18,11 +18,17 @@ attention/ │ ├── sage_attn3.py # SageAttention v3 │ ├── sdpa.py # torch SDPA fallback │ ├── video_sparse_attn.py # VSA (paper: Video Sparse Attention) +│ ├── video_sparse_attn_h3.py # VSA for MiniMax H3 packed mixed-modality attention +│ ├── video_sparse_attn_h3_probe.py # Attention-mass probe for VSA-H3 selection │ ├── vmoba.py # Video-MoBA │ ├── sla.py # Sliding-window (STA) -│ └── bsa_attn.py # Block-sparse +│ ├── bsa_attn.py # Block-sparse +│ ├── nabla.py # NABLA block-sparse flex attention (Kandinsky5) +│ ├── attn_qat_train.py # QAT training attention path +│ └── attn_qat_infer.py # QAT inference attention path └── utils/ ├── flash_attn_cute.py + ├── flash_attn_default.py └── flash_attn_no_pad.py ``` diff --git a/fastvideo/models/AGENTS.md b/fastvideo/models/AGENTS.md index 9a5aa009c..635f41eee 100644 --- a/fastvideo/models/AGENTS.md +++ b/fastvideo/models/AGENTS.md @@ -1,6 +1,6 @@ # `fastvideo/models/` — Model Implementations -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 DiT / VAE / encoder / scheduler / upsampler / audio model classes. **Pre-commit excludes this directory**, except Wan's configs, definition, and `__init__.py`. Match neighboring file style manually for excluded files. @@ -13,7 +13,9 @@ models/ │ ├── .py # Single-file DiT (ltx2, hunyuanvideo, cosmos, ...); wanvideo is a shim │ ├── hyworld/ # Multi-file DiT family │ ├── lingbotworld/ # ditto -│ └── matrixgame2/ # ditto +│ ├── lingbotworld2/ # ditto +│ ├── matrixgame2/ # ditto +│ └── minimax_h3_fusions/ # Fused H3 modulation / QK-norm-RoPE / SwiGLU kernels ├── vaes/ # AutoencoderKL variants + shared utilities; wanvae is a shim ├── encoders/ # T5, CLIP, Llama, Qwen2.5, Gemma, SigLIP, Reason1, audio conditioner ├── schedulers/ # FlowMatch / EulerDiscrete / DPM custom schedulers diff --git a/fastvideo/models/wan/AGENTS.md b/fastvideo/models/wan/AGENTS.md index 13f6ab795..46f478bba 100644 --- a/fastvideo/models/wan/AGENTS.md +++ b/fastvideo/models/wan/AGENTS.md @@ -42,6 +42,12 @@ explicit compatibility exports. `__init__.py` are pre-commit checked; `transformer.py` and `vae.py` retain the existing model-code exclusion. Avoid unrelated reformatting. +- `WanTransformer3DModel.forward` shards the flattened spatiotemporal token + sequence after patch embedding (`sequence_model_parallel_shard`, with + padding). Callers pass full-length latents and conditioning; a pipeline must + not pre-shard along the temporal axis. Pre-sharding only part of the model + input (for example I2V image conditioning) cannot match the full-length noise + and double-shards the rest. ## Focused checks diff --git a/fastvideo/pipelines/AGENTS.md b/fastvideo/pipelines/AGENTS.md index 0ec613925..991b47867 100644 --- a/fastvideo/pipelines/AGENTS.md +++ b/fastvideo/pipelines/AGENTS.md @@ -1,6 +1,6 @@ # `fastvideo/pipelines/` — Pipeline Composition -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 Diffusion pipelines are **compositions of `PipelineStage` objects**. Each stage owns one verb (validate / encode / schedule / denoise / decode). Adding a model means assembling stages, not subclassing a megapipeline. @@ -25,9 +25,10 @@ pipelines/ │ ├── gamecraft_denoising.py # GameCraft-specific │ └── matrixgame2_denoising.py # Matrix-Game 2.0-specific ├── basic/ # Per-model end-to-end pipelines -│ ├── hunyuan/, hunyuan15/, hyworld/, gamecraft/, gen3c/, cosmos/ -│ ├── wan/, longcat/, ltx2/, lingbotworld/, magi_human/, matrixgame2/ -│ ├── sd35/, stable_audio/, turbodiffusion/ +│ ├── cosmos/, dreamx_world/, flux/, flux_2/, gamecraft/, gen3c/, glm_image/ +│ ├── hunyuan/, hunyuan15/, hyworld/, kandinsky5/, lingbot_video/, lingbotworld/ +│ ├── lingbotworld2/, longcat/, ltx2/, magi_human/, matrixgame2/, matrixgame3/ +│ ├── minimax_h3/, mmaudio/, sd35/, stable_audio/, turbodiffusion/, wan/, zimage/ │ └── /{_pipeline.py, presets.py, __init__.py} ├── preprocess/ # Data preprocessing pipelines (ltx2, wan, matrixgame2) └── training/ # Training-time pipeline glue diff --git a/fastvideo/tests/AGENTS.md b/fastvideo/tests/AGENTS.md index 49dd687ad..3a713e212 100644 --- a/fastvideo/tests/AGENTS.md +++ b/fastvideo/tests/AGENTS.md @@ -1,6 +1,6 @@ # `fastvideo/tests/` — Package-Level Test Suite -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 > **Pre-commit excludes `fastvideo/tests/`.** Lint/format hooks do not run on > files here. Match style of neighboring tests manually. @@ -11,26 +11,37 @@ tests/ ├── conftest.py # distributed_setup fixture (1×1 SP/TP init+cleanup) ├── utils.py # Shared test helpers +├── test_partial_hf_download.py # Partial HF snapshot download behavior ├── api/ # api/ schema + presets ├── attention/ # Backend selector + layer parity +├── audio/ # Audio encoder/decoder tests (LTX-2 BWE) +├── contract/ # CI planner + cross-module shape contracts ├── dataset/ # Dataloader smoke tests ├── distributed/ # SP / TP collectives ├── encoders/ # Per-encoder parity (t5, clip, llama, qwen, ...) ├── entrypoints/ # CLI + streaming + OpenAI-compatible server │ └── streaming/ # Streaming-specific tests +├── eval/ # Evaluation metric tests + reference scores +├── golden_gate/ # Model-family golden tests ├── hooks/ # Runtime hook system ├── inference/ # End-to-end inference smoke +├── layers/ # Tensor-parallel layer unit tests +├── loader/ # Component loader / FSDP / LoRA patch tests ├── lora_extraction/ # LoRA merge/extract round-trip +├── mlx/ # Apple MLX runtime tests ├── modal/ # Dormant manual Modal rollback helpers + their unit tests ├── nightly/ # Long-running suites (gated) ├── ops/ # Custom kernel / op tests ├── performance/ # Throughput / memory benchmarks (informational) +├── pipelines/ # Pipeline composition tests +├── platforms/ # Platform capability / unified-memory tests ├── ssim/ # GPU SSIM regressions — see ssim/AGENTS.md ├── stages/ # Per-stage unit tests ├── train/ # New modular trainer tests ├── training/ # Legacy training pipeline tests ├── transformers/ # transformers-shim tests ├── vaes/ # Per-VAE encode/decode parity +├── worker/ # Worker / executor tests └── workflow/ # Preprocessing workflow tests ``` @@ -59,17 +70,34 @@ working tree but are not part of the package suite. - `nightly/` and `performance/` tests should be guarded by an environment marker so the default `pytest fastvideo/tests/` stays fast. +## CI Wiring + +Adding a test file does not add it to CI. `.buildkite/scripts/unit_test.sh` is +an explicit file list, and each domain lane script (for example +`.buildkite/scripts/lanes/distillation_dmd.sh`) names its files. A new test +under an existing directory will not run in Fastcheck or the merge gate until +it is added to a lane script or a new lane is registered; see the +[testing guide](../../docs/contributing/testing.md). + +CPU-only tests still import GPU backends transitively +(`attention/backends/vmoba.py` and `sla.py` import `fastvideo_kernel`), which +fails on a GPU-less host with `RuntimeError: 0 active drivers`. Stub +`fastvideo_kernel` (`moba_attn_varlen`, `process_moba_input`, +`process_moba_output`, and `triton_kernels.sla_triton._attention`) or run +inside the CI container. + ## Slurm CI -`tests/ssim/ci_runner.py` is the active CI orchestrator. It auto-discovers +`tests/ssim/ci_runner.py` orchestrates the SSIM lane. It auto-discovers `test_*.py` files under `ssim/` and schedules one subprocess per `*_MODEL_TO_PARAMS` key across the four GPUs granted by the Slinky Slurm lane. The change-aware merge planner may pass repeated `--test-file` basenames for focused model-family coverage; direct and scheduled full runs omit the filter. New SSIM tests need no pipeline wiring — just declare `REQUIRED_GPUS = N`. -The modules in `tests/modal/` are retained for dormant manual rollback and -must not be referenced by active Buildkite or slash-command routes. +The rollback helper modules in `tests/modal/` are retained for dormant manual +rollback and must not be invoked by active Buildkite or slash-command routes; +their unit tests are still collected by `unit_test.sh`. ## Anti-Patterns diff --git a/fastvideo/train/AGENTS.md b/fastvideo/train/AGENTS.md index 29f470d16..045e36d14 100644 --- a/fastvideo/train/AGENTS.md +++ b/fastvideo/train/AGENTS.md @@ -8,6 +8,7 @@ YAML-driven trainer composed from interchangeable **methods × models × callbac train/ ├── trainer.py # Core training loop coordinator ├── README.md # User-facing overview (legacy → new diff) +├── attn_qat/ # Attn-QAT integration notes (see docs/training/attn_qat.md) ├── entrypoint/ # train.py + dcp_to_diffusers.py CLI entrypoints ├── methods/ │ ├── base.py # TrainingMethod ABC @@ -18,15 +19,23 @@ train/ │ └── rl/ # RL methods, sampling helpers, rewards ├── models/ │ ├── base.py # ModelBase / CausalModelBase wrappers -│ ├── wan/, hunyuan/, cosmos/ # Per-family training wrappers +│ ├── wan/, hunyuan/, cosmos/, kandinsky5/, longcat/, ltx2/, matrixgame2/, minimax_h3/ # Per-family wrappers ├── callbacks/ # callback.py base + ema, grad_clip, validation └── utils/ ├── training_config.py # Hierarchical YAML config dataclasses + ├── config.py # _target_-based run config ├── builder.py # Build trainer from config + ├── instantiate.py # _target_-based instantiation utilities ├── checkpoint.py # DCP save/load + ├── activation_checkpoint.py # Activation checkpointing policies ├── optimizer.py # AdamW / fused optimizer factory ├── tracking.py # build_tracker (W&B / TensorBoard) - └── dataloader.py # StatefulDataLoader wiring + ├── dataloader.py # StatefulDataLoader wiring + ├── lora.py # Training-side LoRA utilities + ├── module_state.py # Train/eval + requires_grad role helpers + ├── moduleloader.py # Module loading helpers + ├── negative_prompt.py # Per-rank negative-prompt encoding + └── validation_media.py # Validation frame/audio encoding ``` ## Composition Model diff --git a/fastvideo/training/AGENTS.md b/fastvideo/training/AGENTS.md index b60c36605..a74a411a0 100644 --- a/fastvideo/training/AGENTS.md +++ b/fastvideo/training/AGENTS.md @@ -1,9 +1,9 @@ # `fastvideo/training/` — Legacy Monolithic Training Pipelines -**Generated:** 2026-05-02 +**Generated:** 2026-09-14 > **Status:** maintenance mode. New training work should go in `fastvideo/train/`. -> Existing shipped recipes (Wan, LTX-2, Matrix-Game 2.0 distillation) still live here. +> Existing shipped recipes (Wan, LTX-2, Matrix-Game 2.0, Cosmos 2.5) still live here. ## What's Here @@ -19,6 +19,10 @@ training/ ├── wan_self_forcing_distillation_pipeline.py ├── ltx2_training_pipeline.py # LTX-2 training ├── matrixgame2_training_pipeline.py # Matrix-Game 2.0 training +├── matrixgame2_self_forcing_distillation_pipeline.py # Matrix-Game 2.0 self-forcing distill +├── matrixgame2_ode_causal_pipeline.py # Matrix-Game 2.0 ODE-causal +├── matrixgame2_ar_diffusion_pipeline.py # Matrix-Game 2.0 AR diffusion +├── cosmos2_5_training_pipeline.py # Cosmos 2.5 training ├── ode_causal_pipeline.py # ODE-causal pipeline ├── checkpointing_utils.py # save/load helpers ├── activation_checkpoint.py # AC wrapping helpers @@ -36,6 +40,22 @@ training/ - Configs are flat argparse flags wired in `fastvideo_args.py` and per-pipeline argparse groups. No YAML. +## Wan I2V Conditioning + +The I2V conditioning tensor is `[mask(temporal_compression_ratio channels), +image_latents(16)]` concatenated with the 16-channel noise on the channel dim. +The mask is one on the first latent frame across all temporal-compression +channels and zero elsewhere. Read `temporal_compression_ratio` from +`training_args.pipeline_config.vae_config.arch_config` (inference: +`self.vae.temporal_compression_ratio`); never hardcode 4. The construction is +duplicated in `wan_i2v_training_pipeline.py`, `wan_i2v_distillation_pipeline.py`, +`matrixgame2_self_forcing_distillation_pipeline.py`, and +`pipelines/stages/image_encoding.py` — keep them consistent. + +`WanTransformer3DModel` shards the flattened token sequence after patch +embedding, so these pipelines pass full-length conditioning and never pre-shard +it; see `fastvideo/models/wan/AGENTS.md`. + ## When to Touch a File Here - Bugfix or behavior tweak to a shipped recipe (Wan, LTX-2, Matrix-Game 2.0). From 614b59543c40f56990dde2e55930635f9125c31d Mon Sep 17 00:00:00 2001 From: Yaegaki1Erika <70182590+Yaegaki1Erika@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:34:25 +0800 Subject: [PATCH 006/151] [bugfix]: respect serialized tensor dtype in parquet dataloader (#1843) --- fastvideo/dataset/utils.py | 112 ++++++++++++----- .../preprocess_minimax_h3_overfit.py | 2 +- .../dataset/test_parquet_dtype_roundtrip.py | 119 ++++++++++++++++++ 3 files changed, 203 insertions(+), 30 deletions(-) create mode 100644 fastvideo/tests/dataset/test_parquet_dtype_roundtrip.py diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index efe8d8e2b..a20567aba 100644 --- a/fastvideo/dataset/utils.py +++ b/fastvideo/dataset/utils.py @@ -1,9 +1,69 @@ +import functools +import math import random from typing import Any, cast import numpy as np import torch +_DTYPE_ALIASES = { + "float": "float32", + "double": "float64", + "half": "float16", + "long": "int64", + "int": "int32", +} + + +@functools.lru_cache(maxsize=None) +def _normalize_tensor_dtype(dtype_value: Any) -> str: + """Normalize dtype strings written by NumPy or PyTorch record creators.""" + dtype_name = str(dtype_value).strip().lower() + for prefix in ("torch.", "numpy.", "np."): + if dtype_name.startswith(prefix): + dtype_name = dtype_name.removeprefix(prefix) + return _DTYPE_ALIASES.get(dtype_name, dtype_name) + + +def _decode_tensor_bytes( + bytes_data: bytes, + shape: list[int] | tuple[int, ...], + dtype_value: Any, + *, + zero: bool = False, +) -> torch.Tensor: + """Decode one tensor using the dtype persisted beside its byte buffer. + + ``zero=True`` still validates the byte length but returns a correctly typed + zero tensor instead of copying the payload, for CFG-dropped embeddings. + """ + dtype_name = _normalize_tensor_dtype(dtype_value) + is_bfloat16 = dtype_name == "bfloat16" + if is_bfloat16: + # NumPy has no portable bfloat16 dtype. Read the raw 16-bit storage, + # then reinterpret it as torch.bfloat16 without changing the bits. + numpy_dtype = np.dtype(np.uint16) + else: + try: + numpy_dtype = np.dtype(dtype_name) + except (TypeError, ValueError) as exc: + raise ValueError(f"Unsupported serialized tensor dtype: {dtype_value!r}") from exc + + expected_nbytes = math.prod(shape) * numpy_dtype.itemsize + if len(bytes_data) != expected_nbytes: + raise ValueError( + "Serialized tensor byte length does not match its shape and dtype: " + f"shape={tuple(shape)}, dtype={dtype_name}, expected={expected_nbytes}, actual={len(bytes_data)}") + + if zero: + tensor = torch.from_numpy(np.zeros(shape, dtype=numpy_dtype)) + else: + array = np.frombuffer(bytes_data, dtype=numpy_dtype).reshape(shape).copy() + tensor = torch.from_numpy(array) + if is_bfloat16: + tensor = tensor.view(torch.bfloat16) + return tensor + def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]: """ @@ -26,32 +86,28 @@ def get_torch_tensors_from_row_dict(row_dict, keys, cfg_rate, rng=None) -> dict[ """ return_dict = {} for key in keys: - shape, bytes = None, None if isinstance(key, tuple): + output_key = key[0] + serialized_key = None for k in key: - try: - shape = row_dict[f"{k}_shape"] - bytes = row_dict[f"{k}_bytes"] - except KeyError: - continue - key = key[0] - if shape is None or bytes is None: - raise ValueError(f"Key {key} not found in row_dict") + if f"{k}_shape" in row_dict and f"{k}_bytes" in row_dict: + serialized_key = k + if serialized_key is None: + raise ValueError(f"Key {output_key} not found in row_dict") else: - shape = row_dict[f"{key}_shape"] - bytes = row_dict[f"{key}_bytes"] + output_key = serialized_key = key + + shape = row_dict[f"{serialized_key}_shape"] + bytes_data = row_dict[f"{serialized_key}_bytes"] + dtype_value = row_dict.get(f"{serialized_key}_dtype", "float32") + drop = output_key == 'text_embedding' and (rng.random() if rng else random.random()) < cfg_rate + data = _decode_tensor_bytes(bytes_data, shape, dtype_value, zero=drop) - # TODO (peiyuan): read precision - if key == 'text_embedding' and (rng.random() if rng else random.random()) < cfg_rate: - data = np.zeros(shape, dtype=np.float32) - else: - data = np.frombuffer(bytes, dtype=np.float32).reshape(shape).copy() - data = torch.from_numpy(data) if len(data.shape) == 3: B, L, D = data.shape assert B == 1, "Batch size must be 1" data = data.squeeze(0) - return_dict[key] = data + return_dict[output_key] = data return return_dict @@ -131,12 +187,12 @@ def collate_rows_from_parquet_schema(rows, # Process each tensor field for tensor_name in tensor_fields: tensor_list = [] + shape_key = f"{tensor_name}_shape" + bytes_key = f"{tensor_name}_bytes" + dtype_key = f"{tensor_name}_dtype" for row in rows: # Get tensor data from row using the existing helper function pattern - shape_key = f"{tensor_name}_shape" - bytes_key = f"{tensor_name}_bytes" - if shape_key in row and bytes_key in row: shape = row[shape_key] bytes_data = row[bytes_key] @@ -153,14 +209,12 @@ def collate_rows_from_parquet_schema(rows, drop = (random.Random(seed ^ sample_idx).random() < cfg_rate) else: drop = ((rng.random() if rng else random.random()) < cfg_rate) - if drop: - data = np.zeros(shape, dtype=np.float32) - else: - data = np.frombuffer( - bytes_data, - dtype=np.float32, - ).reshape(shape).copy() - tensor = torch.from_numpy(data) + tensor = _decode_tensor_bytes( + bytes_data, + shape, + row.get(dtype_key, "float32"), + zero=drop, + ) # if len(data.shape) == 3: # B, L, D = tensor.shape # assert B == 1, "Batch size must be 1" diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py index 3a8f706da..33a26ce99 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py @@ -238,7 +238,7 @@ def build_parquet_record( """Serialize one synchronized H3 sample for the Parquet schema collator. ``collate_rows_from_parquet_schema`` reconstructs every tensor from a - bytes/shape/dtype triplet and reads the byte payload as float32, so this + bytes/shape/dtype triplet and decodes it with the serialized dtype, so this boundary stores contiguous float32 tensors for lossless reconstruction. """ tensors = { diff --git a/fastvideo/tests/dataset/test_parquet_dtype_roundtrip.py b/fastvideo/tests/dataset/test_parquet_dtype_roundtrip.py new file mode 100644 index 000000000..54c7b8751 --- /dev/null +++ b/fastvideo/tests/dataset/test_parquet_dtype_roundtrip.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v +from fastvideo.dataset.utils import ( + _decode_tensor_bytes, + collate_rows_from_parquet_schema, + get_torch_tensors_from_row_dict, +) + + +def _serialized_fields(name: str, value: np.ndarray) -> dict[str, object]: + return { + f"{name}_bytes": value.tobytes(), + f"{name}_shape": list(value.shape), + f"{name}_dtype": str(value.dtype), + } + + +def test_i2v_collate_respects_each_serialized_tensor_dtype() -> None: + vae_latent = np.arange(16 * 2 * 2 * 2, dtype=np.float32).reshape(16, 2, 2, 2) + text_embedding = np.arange(3 * 4, dtype=np.float32).reshape(3, 4) + clip_feature = np.arange(2 * 4, dtype=np.float16).reshape(2, 4) + first_frame_latent = np.arange(16 * 2 * 2 * 2, dtype=np.float32).reshape(16, 2, 2, 2) + pil_image = np.arange(3 * 2 * 4, dtype=np.uint8).reshape(3, 2, 4) + + row: dict[str, object] = { + "id": "sample", + "file_name": "sample.mp4", + "caption": "sample caption", + "media_type": "video", + "width": 4, + "height": 2, + "num_frames": 5, + "duration_sec": 1.0, + "fps": 5.0, + } + for name, value in ( + ("vae_latent", vae_latent), + ("text_embedding", text_embedding), + ("clip_feature", clip_feature), + ("first_frame_latent", first_frame_latent), + ("pil_image", pil_image), + ): + row.update(_serialized_fields(name, value)) + + batch = collate_rows_from_parquet_schema( + [row], + pyarrow_schema_i2v, + text_padding_length=4, + ) + + assert batch["vae_latent"].dtype == torch.float32 + assert batch["clip_feature"].dtype == torch.float16 + assert batch["first_frame_latent"].dtype == torch.float32 + assert batch["pil_image"].dtype == torch.uint8 + torch.testing.assert_close(batch["clip_feature"][0], torch.from_numpy(clip_feature)) + torch.testing.assert_close(batch["pil_image"][0], torch.from_numpy(pil_image)) + + +def test_decode_tensor_bytes_accepts_torch_bfloat16_dtype() -> None: + expected = torch.tensor([[1.0, -2.5], [3.25, 0.0]], dtype=torch.bfloat16) + storage = expected.view(torch.uint16).numpy() + + actual = _decode_tensor_bytes(storage.tobytes(), list(expected.shape), "torch.bfloat16") + + assert actual.dtype == torch.bfloat16 + torch.testing.assert_close(actual, expected) + + +def test_decode_tensor_bytes_zero_returns_typed_zeros() -> None: + expected = torch.tensor([[1.0, -2.5], [3.25, 0.0]], dtype=torch.bfloat16) + storage = expected.view(torch.uint16).numpy() + + actual = _decode_tensor_bytes(storage.tobytes(), list(expected.shape), "torch.bfloat16", zero=True) + + assert actual.dtype == torch.bfloat16 + assert actual.shape == expected.shape + assert actual.count_nonzero() == 0 + + +def test_collate_cfg_drop_zeros_text_embedding_with_serialized_dtype() -> None: + text_embedding = np.arange(3 * 4, dtype=np.float16).reshape(3, 4) + row = _serialized_fields("text_embedding", text_embedding) + row["caption"] = "sample caption" + + batch = collate_rows_from_parquet_schema( + [row], + pyarrow_schema_i2v, + text_padding_length=4, + cfg_rate=1.0, + ) + + assert batch["text_embedding"].dtype == torch.float16 + assert batch["text_embedding"].count_nonzero() == 0 + + +def test_decode_tensor_bytes_rejects_mismatched_byte_length() -> None: + with pytest.raises(ValueError, match="byte length"): + _decode_tensor_bytes(b"\x00\x01", [2], "float32") + + +def test_legacy_collate_helper_respects_dtype_for_alias_keys() -> None: + latent = np.arange(4, dtype=np.float16).reshape(2, 2) + row = _serialized_fields("latent", latent) + + tensors = get_torch_tensors_from_row_dict( + row, + [("vae_latent", "latent")], + cfg_rate=0.0, + ) + + assert tensors["vae_latent"].dtype == torch.float16 + torch.testing.assert_close(tensors["vae_latent"], torch.from_numpy(latent)) From 316f3876c2ec13ebd25d435d814a94b6354e938e Mon Sep 17 00:00:00 2001 From: Lele Date: Tue, 15 Sep 2026 07:54:41 +0800 Subject: [PATCH 007/151] [bugfix]: allow MiniMax H3 frame padding at the 15-second limit Accept the causal-VAE-aligned 362-frame bucket (15.083 s) for 15-second H3 requests. - Hoist MINIMAX_H3_MIN/MAX_ALIGNED_FRAMES next to align_num_frames and reuse them in the CUDA stage, the MLX runtime, and the LoRA example so the bound is consistent across entry points. - State the accepted frame range in the error messages. - Cover seconds="15" -> 362 at OpenAI admission and the MLX resolve_geometry bound. - Update the Spark/openai cookbook frame caps from 345 to 362. Known gaps: no GPU-lane coverage for the 362 bucket (SSIM runs 124 frames; the golden gate is a fixed-geometry fingerprint), and explicit num_frames=360 still hits the pre-existing grid gate. --- docs/cookbook/minimax-h3.md | 2 +- docs/cookbook/openai-api.md | 2 +- .../installation/spark_pair.md | 8 +++---- .../lora/minimax_h3_lora_inference.py | 14 +++++++---- fastvideo/mlx_runtime/minimax_h3.py | 9 ++++++++ fastvideo/mlx_runtime/minimax_h3_pipeline.py | 11 ++++++--- .../pipelines/basic/minimax_h3/packing.py | 7 ++++++ .../stages/minimax_h3_input_preparation.py | 8 ++++--- .../entrypoints/test_openai_serving_engine.py | 7 +++--- .../mlx/test_mlx_minimax_h3_fast_mode.py | 9 ++++++++ .../test_minimax_h3_input_preparation.py | 23 +++++++++++++++++++ 11 files changed, 81 insertions(+), 19 deletions(-) create mode 100644 fastvideo/tests/stages/test_minimax_h3_input_preparation.py diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index a97b8351e..e28456cae 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -269,7 +269,7 @@ cd FastVideo
  • The MLX source runtime supports T2VA, optional temporal --fast, optional spatial --fast-spatial, and opt-in VSA on --include-vsa checkpoints. FL2VA, Ref2VA, and two-pass refinement are not wired.
  • GPU count and VAE decode backend are configurable in the builder above for FastH3 CUDA recipes. Only the value shown by default has a recorded run; other supported values are unmeasured here.
  • DGX Spark is a runtime on FastH3 Preview, not a separate family card. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. The CUDA GPU-count knob does not apply to Spark.
    • -
  • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 345 (15 s). A 345-frame request on one Spark can OOM.
    • +
  • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM.
  • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
    • diff --git a/docs/cookbook/openai-api.md b/docs/cookbook/openai-api.md index 0aae48625..c3fc952f6 100644 --- a/docs/cookbook/openai-api.md +++ b/docs/cookbook/openai-api.md @@ -58,7 +58,7 @@ FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ The configuration loads `FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` on one GB10 and advertises it as `fasth3`. Lazy module load still reloads Qwen3-VL and the DiT between phases of each request. Legal `num_frames` values -are `17n+5`, capped at 345 (15 s); a 345-frame request on one Spark can OOM. +are `17n+5`, capped at 362 (15.08 s); a 345-frame request on one Spark can OOM. There is no cookbook server for two Sparks; use the generate YAML after [pairing two Sparks](../getting_started/installation/spark_pair.md). diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index 616b11a73..4eb86cc9e 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/docs/getting_started/installation/spark_pair.md @@ -109,7 +109,7 @@ Run the driver on the **head**, same venv, same QSFP IP. `basic_fasth3.py` defaults target a four-GPU GB200 profile: 768×1344, `sm100a` VSA, FA4, four GPUs. On Sparks you must override the kernel flags. Height, width, frames, steps, seed, and prompt are yours. Change them. Legal -`num_frames` values are `17n+5`, capped at 345. +`num_frames` values are `17n+5`, capped at 362. GB10 has no FA4 / sm_100a VSA kernel, so `--vsa-kernel triton --no-fa4` stays required on this box. `--execution-backend ray` is optional when `RAY_ADDRESS` @@ -152,8 +152,8 @@ Stop the cluster when you are done: `ray stop` on both nodes. ## FastH3 frame counts H3 is 24 fps. Legal `num_frames` values are `17n+5`. The pipeline rejects -clips longer than **15 s**. The longest legal length is **345 frames** -(14.375 s). 360 frames aligns to 362 and fails the duration check. +clips longer than **15 s**. The longest legal length is **362 frames** +(15.083 s). 360 frames aligns to 362 and is accepted. ## Measured on two GB10s (2026-08-31) @@ -179,7 +179,7 @@ VAE, same 4-step schedule: | Two Sparks, SP=2 | 2 | 124 | **215.2 s** | 72.4 s | Those medians used `--height` / `--width` / `--num-frames` as CLI flags. Swap -them. Native 480p on this model is 480×832, 124 frames. The 15 s cap is 345 +them. Native 480p on this model is 480×832, 124 frames. The 15 s cap is 362 frames. The first VAE decode still pays `torch.compile`. Later `generate()` calls in diff --git a/examples/inference/lora/minimax_h3_lora_inference.py b/examples/inference/lora/minimax_h3_lora_inference.py index b74b8ce3c..9a0c6c78a 100644 --- a/examples/inference/lora/minimax_h3_lora_inference.py +++ b/examples/inference/lora/minimax_h3_lora_inference.py @@ -37,7 +37,8 @@ from collections.abc import Sequence from pathlib import Path # MiniMax-H3 generates 5-15 s at 24 fps, on a latent grid that only admits frame counts -# of the form 17n + 5. 124 is the 5-second point the FastH3 profile is measured at. +# of the form 17n + 5. 124 is the 5-second point the FastH3 profile is measured at; +# the 15-second cap pads to 362 frames. FRAMES_PER_CHUNK = 17 LATENTS_PER_CHUNK = 5 FPS = 24 @@ -52,6 +53,10 @@ def align_num_frames(num_frames: int) -> int: return LATENTS_PER_CHUNK + chunks * FRAMES_PER_CHUNK +MIN_ALIGNED_FRAMES = align_num_frames(int(MIN_DURATION * FPS)) +MAX_ALIGNED_FRAMES = align_num_frames(int(MAX_DURATION * FPS)) + + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--model-path", default="MiniMaxAI/MiniMax-H3", help="the BASE checkpoint the adapter targets") @@ -86,9 +91,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: # community adapters built against the ComfyUI layout carry no gate. Checked in # main(), once the path has been resolved. aligned = align_num_frames(args.num_frames) - if not MIN_DURATION <= aligned / FPS <= MAX_DURATION: - parser.error(f"MiniMax-H3 generates {MIN_DURATION}-{MAX_DURATION}s at {FPS} fps; " - f"aligned num_frames={aligned} is {aligned / FPS:.1f}s") + if not MIN_ALIGNED_FRAMES <= aligned <= MAX_ALIGNED_FRAMES: + parser.error(f"MiniMax-H3 generates {MIN_DURATION:g}-{MAX_DURATION:g}s at {FPS} fps; " + f"aligned num_frames={aligned} is {aligned / FPS:.1f}s, outside the accepted " + f"{MIN_ALIGNED_FRAMES}-{MAX_ALIGNED_FRAMES} frame range") args.num_frames = aligned if not math.isfinite(args.lora_strength): parser.error("--lora-strength must be finite") diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index f26b5afeb..5ddd335f6 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -97,6 +97,8 @@ MINIMAX_H3_AUDIO_TAG = 2 MINIMAX_H3_MODALITY_NUM = 3 MINIMAX_H3_FPS = 24 +MINIMAX_H3_MIN_DURATION = 5.0 +MINIMAX_H3_MAX_DURATION = 15.0 MINIMAX_H3_FRAMES_PER_CHUNK = 17 MINIMAX_H3_LATENTS_PER_CHUNK = 5 MINIMAX_H3_AUDIO_LATENTS_PER_SECOND = 40 @@ -161,6 +163,13 @@ def align_num_frames(num_frames: int) -> int: return num_frames +# Mirrors fastvideo/pipelines/basic/minimax_h3/packing.py: duration limits are +# seconds, validated against the aligned buckets they map to. The 15-second +# target (360 frames) pads to 362 on the causal-VAE grid. +MINIMAX_H3_MIN_ALIGNED_FRAMES = align_num_frames(int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)) +MINIMAX_H3_MAX_ALIGNED_FRAMES = align_num_frames(int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)) + + def video_latent_num_frames(num_frames: int) -> int: if num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: raise ValueError(f"`num_frames` must be of the form 17 * n + 5, got {num_frames}.") diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index 5f517855f..28c19dc49 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -43,6 +43,10 @@ from fastvideo.mlx_runtime.minimax_h3 import ( MINIMAX_H3_AUDIO_SHIFT, MINIMAX_H3_FPS, MINIMAX_H3_KEYFRAME_NOISE_AUG, + MINIMAX_H3_MAX_ALIGNED_FRAMES, + MINIMAX_H3_MAX_DURATION, + MINIMAX_H3_MIN_ALIGNED_FRAMES, + MINIMAX_H3_MIN_DURATION, MINIMAX_H3_VIDEO_SHIFT, MiniMaxH3SchedulerState, align_num_frames, @@ -380,9 +384,10 @@ class MiniMaxH3MLXPipeline: aligned_frames = align_num_frames(num_frames) latent_frames = video_latent_num_frames(aligned_frames) duration = aligned_frames / MINIMAX_H3_FPS - if enforce_duration and not 5.0 <= duration <= 15.0: - raise ValueError(f"H3 generates 5-15 s at {MINIMAX_H3_FPS} fps; {aligned_frames} frames " - f"is {duration:.2f} s.") + if enforce_duration and not MINIMAX_H3_MIN_ALIGNED_FRAMES <= aligned_frames <= MINIMAX_H3_MAX_ALIGNED_FRAMES: + raise ValueError(f"H3 generates {MINIMAX_H3_MIN_DURATION:g}-{MINIMAX_H3_MAX_DURATION:g} s at " + f"{MINIMAX_H3_FPS} fps; {aligned_frames} frames is {duration:.2f} s, outside the " + f"accepted {MINIMAX_H3_MIN_ALIGNED_FRAMES}-{MINIMAX_H3_MAX_ALIGNED_FRAMES} frame range.") return { "height": height, "width": width, diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index db5d0538e..9d273e48a 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -118,6 +118,13 @@ def align_num_frames(num_frames: int) -> int: return num_frames +# Duration limits are seconds; requests are validated against the aligned buckets +# they map to. The 15-second target (360 frames) pads to 362 on the causal-VAE +# grid, so 362 is the largest accepted bucket. +MINIMAX_H3_MIN_ALIGNED_FRAMES = align_num_frames(int(MINIMAX_H3_MIN_DURATION * MINIMAX_H3_FPS)) +MINIMAX_H3_MAX_ALIGNED_FRAMES = align_num_frames(int(MINIMAX_H3_MAX_DURATION * MINIMAX_H3_FPS)) + + def video_latent_num_frames(num_frames: int) -> int: if num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: raise ValueError(f"`num_frames` must be of the form 17 * n + 5, got {num_frames}.") diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py index b5712362b..4bccc2284 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py @@ -13,7 +13,9 @@ from fastvideo.models.vision_utils import load_image from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_CANVAS_MULTIPLE, MINIMAX_H3_FPS, + MINIMAX_H3_MAX_ALIGNED_FRAMES, MINIMAX_H3_MAX_DURATION, + MINIMAX_H3_MIN_ALIGNED_FRAMES, MINIMAX_H3_MIN_DURATION, align_num_frames, prepare_keyframe_image, @@ -94,10 +96,10 @@ def resolve_target_num_frames(num_frames: object) -> int: if not isinstance(num_frames, int): raise TypeError("MiniMax-H3 `num_frames` must be an integer.") aligned = align_num_frames(num_frames) - duration = aligned / MINIMAX_H3_FPS - if not MINIMAX_H3_MIN_DURATION <= duration <= MINIMAX_H3_MAX_DURATION: + if not MINIMAX_H3_MIN_ALIGNED_FRAMES <= aligned <= MINIMAX_H3_MAX_ALIGNED_FRAMES: raise ValueError(f"MiniMax-H3 generates {MINIMAX_H3_MIN_DURATION:g}-{MINIMAX_H3_MAX_DURATION:g} seconds at " - f"{MINIMAX_H3_FPS} fps; aligned num_frames={aligned}.") + f"{MINIMAX_H3_FPS} fps; aligned num_frames={aligned} is outside the accepted " + f"{MINIMAX_H3_MIN_ALIGNED_FRAMES}-{MINIMAX_H3_MAX_ALIGNED_FRAMES} frame range.") return aligned diff --git a/fastvideo/tests/entrypoints/test_openai_serving_engine.py b/fastvideo/tests/entrypoints/test_openai_serving_engine.py index 0a2c06f0a..f833f78a7 100644 --- a/fastvideo/tests/entrypoints/test_openai_serving_engine.py +++ b/fastvideo/tests/entrypoints/test_openai_serving_engine.py @@ -253,8 +253,9 @@ def test_fasth3_invalid_geometry_fails_at_admission(tmp_path: Path, overrides, m ) -def test_fasth3_seconds_align_to_causal_vae_grid(tmp_path: Path) -> None: - request = VideoGenerationRequest(prompt="a fox", aspect_ratio="16:9", seconds="5") +@pytest.mark.parametrize(("seconds", "expected_frames"), [("5", 124), ("15", 362)]) +def test_fasth3_seconds_align_to_causal_vae_grid(tmp_path: Path, seconds: str, expected_frames: int) -> None: + request = VideoGenerationRequest(prompt="a fox", aspect_ratio="16:9", seconds=seconds) adapted = build_generation_request( "video_gen_test", request, @@ -262,7 +263,7 @@ def test_fasth3_seconds_align_to_causal_vae_grid(tmp_path: Path) -> None: served_model_name="fasth3", output_dir=str(tmp_path), ) - assert adapted.sampling.num_frames == 124 + assert adapted.sampling.num_frames == expected_frames def test_fasth3_ref2va_limits_references_at_admission(tmp_path: Path) -> None: diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py index 5cc04cc29..e00950554 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py @@ -41,6 +41,15 @@ def test_fast_plan_keeps_full_audio_and_reduces_only_video() -> None: assert plan.video_temporal_scale > 1.0 +def test_resolve_geometry_accepts_the_aligned_duration_limit() -> None: + geometry = MiniMaxH3MLXPipeline.resolve_geometry(768, 1344, 360) + + assert geometry["num_frames"] == 362 + assert geometry["latent_frame_count"] == 107 + with pytest.raises(ValueError, match="H3 generates"): + MiniMaxH3MLXPipeline.resolve_geometry(768, 1344, 363) + + def test_fast_layout_stretches_video_positions_without_changing_audio() -> None: baseline = build_packed_layout(8, 22, 46, 80, 207) fast = build_packed_layout(8, 22, 46, 80, 207, video_temporal_scale=1.7) diff --git a/fastvideo/tests/stages/test_minimax_h3_input_preparation.py b/fastvideo/tests/stages/test_minimax_h3_input_preparation.py new file mode 100644 index 000000000..78227fb7b --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_input_preparation.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import resolve_target_num_frames + + +@pytest.mark.parametrize("requested, expected", [(120, 124), (345, 345), (346, 362), (360, 362), (362, 362)]) +def test_target_frames_allow_padding_at_duration_limit(requested: int, expected: int) -> None: + assert resolve_target_num_frames(requested) == expected + assert resolve_target_num_frames(expected) == expected + + +@pytest.mark.parametrize("requested", [107, 363, 379]) +def test_target_frames_reject_buckets_outside_duration_limits(requested: int) -> None: + with pytest.raises(ValueError, match="MiniMax-H3 generates"): + resolve_target_num_frames(requested) + + +@pytest.mark.parametrize("requested", [360.0, "360", None]) +def test_target_frames_require_integer(requested: object) -> None: + with pytest.raises(TypeError, match="must be an integer"): + resolve_target_num_frames(requested) From 61b91220c05d5c0835662d7061107cf0abe7643d Mon Sep 17 00:00:00 2001 From: sudhirpol522 <47249555+sudhirpol522@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:01:40 -0600 Subject: [PATCH 008/151] [bugfix] Validate image response format before generation (#1841) Co-authored-by: SolitaryThinker --- fastvideo/entrypoints/openai/image_api.py | 20 +++- .../test_image_response_format_validation.py | 101 ++++++++++++++++++ 2 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 fastvideo/tests/entrypoints/test_image_response_format_validation.py diff --git a/fastvideo/entrypoints/openai/image_api.py b/fastvideo/entrypoints/openai/image_api.py index 6b624b9e2..b52725fca 100644 --- a/fastvideo/entrypoints/openai/image_api.py +++ b/fastvideo/entrypoints/openai/image_api.py @@ -32,6 +32,16 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) router = APIRouter(prefix="/v1/images", tags=["images"]) +_SUPPORTED_RESPONSE_FORMATS = frozenset({"b64_json", "url"}) + + +def _normalize_response_format(value: str | None) -> str: + """Lowercase and validate an OpenAI image response_format value.""" + fmt = (value or "b64_json").lower() + if fmt not in _SUPPORTED_RESPONSE_FORMATS: + raise HTTPException(status_code=400, detail=f"response_format={fmt} is not supported") + return fmt + def _build_generation_kwargs( request_id: str, @@ -87,6 +97,8 @@ def _build_generation_kwargs( @router.post("", response_model=ImageResponse) async def generations(request: ImageGenerationsRequest): + resp_format = _normalize_response_format(request.response_format) + request_id = generate_request_id() engine = get_serving_engine() @@ -115,14 +127,13 @@ async def generations(request: ImageGenerationsRequest): save_file_path = gen_kwargs["output_path"] - resp_format = (request.response_format or "b64_json").lower() if resp_format == "b64_json": if not os.path.exists(save_file_path): raise HTTPException(status_code=500, detail="Image was not saved to disk") async with aiofiles.open(save_file_path, "rb") as f: b64_data = base64.b64encode(await f.read()).decode("utf-8") data = [ImageResponseData(b64_json=b64_data, revised_prompt=request.prompt)] - elif resp_format == "url": + else: data = [ ImageResponseData( url=f"/v1/images/{request_id}/content", @@ -130,8 +141,6 @@ async def generations(request: ImageGenerationsRequest): file_path=os.path.abspath(save_file_path), ) ] - else: - raise HTTPException(status_code=400, detail=f"response_format={resp_format} is not supported") await IMAGE_STORE.upsert( request_id, @@ -170,6 +179,8 @@ async def edits( num_inference_steps: int | None = Form(None), enable_teacache: bool | None = Form(False), ): + resp_format = _normalize_response_format(response_format) + request_id = generate_request_id() engine = get_serving_engine() @@ -218,7 +229,6 @@ async def edits( save_file_path = gen_kwargs["output_path"] - resp_format = (response_format or "b64_json").lower() if resp_format == "b64_json": async with aiofiles.open(save_file_path, "rb") as f: b64_data = base64.b64encode(await f.read()).decode("utf-8") diff --git a/fastvideo/tests/entrypoints/test_image_response_format_validation.py b/fastvideo/tests/entrypoints/test_image_response_format_validation.py new file mode 100644 index 000000000..d52c4f86f --- /dev/null +++ b/fastvideo/tests/entrypoints/test_image_response_format_validation.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate image response formats before invoking the generator.""" + +import base64 +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fastvideo.entrypoints.openai import image_api +from fastvideo.entrypoints.openai.stores import AsyncDictStore + + +@pytest.fixture +def image_client(monkeypatch, tmp_path): + generate = Mock(side_effect=lambda **kwargs: Path(kwargs["output_path"]).write_bytes(b"test-image-content")) + + async def run_serialized(fn, **kwargs): + return fn(**kwargs) + + engine = SimpleNamespace( + generator=SimpleNamespace(generate_video=generate), + run_serialized=AsyncMock(side_effect=run_serialized), + ) + monkeypatch.setattr(image_api, "get_serving_engine", lambda: engine) + monkeypatch.setattr(image_api, "get_output_dir", lambda: str(tmp_path)) + monkeypatch.setattr(image_api, "IMAGE_STORE", AsyncDictStore()) + app = FastAPI() + app.include_router(image_api.router) + with TestClient(app) as client: + yield client, engine + + +@pytest.mark.parametrize("response_format", ["invalid", "INVALID"]) +def test_invalid_response_format_does_not_generate(image_client, response_format, tmp_path): + client, engine = image_client + response = client.post("/v1/images", json={"prompt": "a cat", "response_format": response_format}) + + assert response.status_code == 400 + assert response.json()["detail"] == "response_format=invalid is not supported" + engine.run_serialized.assert_not_awaited() + engine.generator.generate_video.assert_not_called() + assert not (tmp_path / "images").exists() + + +def test_edits_invalid_response_format_does_not_generate(image_client, tmp_path): + client, engine = image_client + response = client.post( + "/v1/images/edits", + data={"prompt": "a cat", "response_format": "invalid"}, + files={"image": ("cat.png", b"fake-image", "image/png")}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "response_format=invalid is not supported" + engine.run_serialized.assert_not_awaited() + engine.generator.generate_video.assert_not_called() + assert not (tmp_path / "uploads").exists() + + +@pytest.mark.parametrize( + "payload,expected", + [ + pytest.param({"prompt": "a cat"}, "url", id="omitted-uses-protocol-default"), + pytest.param({"prompt": "a cat", "response_format": None}, "b64_json", id="null-falls-back-to-b64-json"), + pytest.param({"prompt": "a cat", "response_format": ""}, "b64_json", id="empty-falls-back-to-b64-json"), + pytest.param({"prompt": "a cat", "response_format": "b64_json"}, "b64_json", id="b64-json"), + pytest.param({"prompt": "a cat", "response_format": "B64_JSON"}, "b64_json", id="b64-json-uppercase"), + pytest.param({"prompt": "a cat", "response_format": "url"}, "url", id="url"), + pytest.param({"prompt": "a cat", "response_format": "URL"}, "url", id="url-uppercase"), + ], +) +def test_supported_response_formats_still_generate(image_client, payload, expected): + client, engine = image_client + response = client.post("/v1/images", json=payload) + + assert response.status_code == 200, response.text + body = response.json() + if expected == "url": + assert body["data"][0]["url"] == f"/v1/images/{body['id']}/content" + else: + assert base64.b64decode(body["data"][0]["b64_json"]) == b"test-image-content" + engine.run_serialized.assert_awaited_once() + engine.generator.generate_video.assert_called_once() + + +def test_edits_supported_response_format_still_generates(image_client): + client, engine = image_client + response = client.post( + "/v1/images/edits", + data={"prompt": "a cat", "response_format": "b64_json"}, + files={"image": ("cat.png", b"fake-image", "image/png")}, + ) + + assert response.status_code == 200, response.text + assert base64.b64decode(response.json()["data"][0]["b64_json"]) == b"test-image-content" + engine.run_serialized.assert_awaited_once() + engine.generator.generate_video.assert_called_once() From 8839ba8d4d155a50f7c8432f7c338e235bd98910 Mon Sep 17 00:00:00 2001 From: sudhirpol522 <47249555+sudhirpol522@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:55:56 -0600 Subject: [PATCH 009/151] [bugfix] Add OpenAI-compatible image generation endpoint (#1840) --- docs/design/server_contracts/openai.md | 4 +- fastvideo/entrypoints/openai/image_api.py | 1 + ...st_image_generation_route_compatibility.py | 58 +++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py diff --git a/docs/design/server_contracts/openai.md b/docs/design/server_contracts/openai.md index 3834ef48d..350705ee8 100644 --- a/docs/design/server_contracts/openai.md +++ b/docs/design/server_contracts/openai.md @@ -25,11 +25,13 @@ requests. HTTP handling and job polling remain asynchronous. | `GET` | `/v1/videos/{id}/content` | Download a completed MP4 | | `DELETE` | `/v1/videos/{id}` | Delete a job and its completed artifact | | `POST` | `/v1/images` | Generate an image | +| `POST` | `/v1/images/generations` | OpenAI-compatible alias for image generation | | `POST` | `/v1/images/edits` | Generate an image from image references | | `GET` | `/v1/images/{id}/content` | Download a generated image | | `GET` | `/health` | Liveness probe | -`POST /v1/videos/generations` remains an alias for older FastVideo clients. +`POST /v1/videos/generations` remains an alias for older FastVideo clients, and +`POST /v1/images/generations` is the OpenAI Python client's image generation path. The OpenAI Python and JavaScript clients can create, retrieve, list, download, and delete video jobs. Use the [H3 server cookbook](../../cookbook/openai-api.md) for pinned client versions and executable examples. Download variants other diff --git a/fastvideo/entrypoints/openai/image_api.py b/fastvideo/entrypoints/openai/image_api.py index b52725fca..350b90cb1 100644 --- a/fastvideo/entrypoints/openai/image_api.py +++ b/fastvideo/entrypoints/openai/image_api.py @@ -95,6 +95,7 @@ def _build_generation_kwargs( return kwargs +@router.post("/generations", response_model=ImageResponse) @router.post("", response_model=ImageResponse) async def generations(request: ImageGenerationsRequest): resp_format = _normalize_response_format(request.response_format) diff --git a/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py b/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py new file mode 100644 index 000000000..1cd3eb431 --- /dev/null +++ b/fastvideo/tests/entrypoints/test_image_generation_route_compatibility.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Image generation route compatibility with a mocked generator.""" + +import base64 +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fastvideo.entrypoints.openai import image_api +from fastvideo.entrypoints.openai.stores import AsyncDictStore + + +@pytest.mark.parametrize("path", [ + pytest.param("/v1/images/generations", id="openai-route"), + pytest.param("/v1/images", id="legacy-route"), +]) +@pytest.mark.parametrize("response_format", [ + pytest.param("b64_json", id="b64-json"), + pytest.param("url", id="url"), +]) +def test_image_generation_routes(path, response_format, monkeypatch, tmp_path): + image_bytes = b"test-image-content" + generate = Mock(side_effect=lambda **kwargs: Path(kwargs["output_path"]).write_bytes(image_bytes)) + + async def run_serialized(fn, **kwargs): + return fn(**kwargs) + + engine = SimpleNamespace( + generator=SimpleNamespace(generate_video=generate), + run_serialized=AsyncMock(side_effect=run_serialized), + ) + monkeypatch.setattr(image_api, "get_serving_engine", lambda: engine) + monkeypatch.setattr(image_api, "get_output_dir", lambda: str(tmp_path)) + monkeypatch.setattr(image_api, "IMAGE_STORE", AsyncDictStore()) + app = FastAPI() + app.include_router(image_api.router) + + with TestClient(app) as client: + response = client.post(path, json={"prompt": "a cat", "response_format": response_format}) + assert response.status_code == 200, response.text + body = response.json() + assert body["data"][0]["revised_prompt"] == "a cat" + if response_format == "b64_json": + assert base64.b64decode(body["data"][0]["b64_json"]) == image_bytes + else: + assert body["data"][0]["url"] == f"/v1/images/{body['id']}/content" + content = client.get(body["data"][0]["url"]) + assert content.status_code == 200 + assert content.content == image_bytes + + generate.assert_called_once() + assert generate.call_args.kwargs["prompt"] == "a cat" + assert generate.call_args.kwargs["num_frames"] == 1 + engine.run_serialized.assert_awaited_once() From 37d06a832f5e8cef3470f81ffb7026e89a405ff0 Mon Sep 17 00:00:00 2001 From: Ishan <34205320+Ishxn20@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:23:35 -0700 Subject: [PATCH 010/151] [feat] Add fastvideo serve configs for Wan CUDA models (#1801) Co-authored-by: SolitaryThinker --- examples/serving/README.md | 34 +++++++++++++-- examples/serving/openai_fastwan21_1_3b.yaml | 44 +++++++++++++++++++ examples/serving/openai_wan21_i2v_14b.yaml | 48 +++++++++++++++++++++ examples/serving/openai_wan22_t2v_a14b.yaml | 29 +++++++++++++ examples/serving/openai_wan22_ti2v_5b.yaml | 27 ++++++++++++ fastvideo/tests/api/test_serve_examples.py | 21 +++++++++ 6 files changed, 200 insertions(+), 3 deletions(-) create mode 100644 examples/serving/openai_fastwan21_1_3b.yaml create mode 100644 examples/serving/openai_wan21_i2v_14b.yaml create mode 100644 examples/serving/openai_wan22_t2v_a14b.yaml create mode 100644 examples/serving/openai_wan22_ti2v_5b.yaml create mode 100644 fastvideo/tests/api/test_serve_examples.py diff --git a/examples/serving/README.md b/examples/serving/README.md index 9424127a4..86e5d2cdb 100644 --- a/examples/serving/README.md +++ b/examples/serving/README.md @@ -2,8 +2,10 @@ The REST serving engine is model-agnostic. Any model supported by `VideoGenerator` can use the same `/v1/models`, `/v1/videos`, and `/v1/images` -surface. The two FastH3 configs here are source-backed examples, not recorded -serving benchmarks. Both configure four CUDA GPUs without a GPU model or VRAM claim. +surface. The two CUDA FastH3 configs here are source-backed examples, not +recorded serving benchmarks. Both configure four CUDA GPUs without a GPU model +or VRAM claim. The Wan configs are copied from the working generate config or +example for each checkpoint. Use the [H3 server cookbook](https://haoailab.com/FastVideo/cookbook/openai-api/) for the guided install, server, and client workflow. @@ -44,6 +46,32 @@ The same playground and clients work with MLX. Its pipeline still loads and releases components between phases to limit unified-memory use. See the cookbook for setup and the supported text-to-video/audio request fields. +## Wan configs + +Four Wan CUDA configs mirror the working generate config or example for each +checkpoint: + +| Config | Model | GPUs | Served alias | +| --- | --- | --- | --- | +| `openai_fastwan21_1_3b.yaml` | FastWan2.1 T2V 1.3B (DMD, VSA) | 1 | `fastwan21-1.3b` | +| `openai_wan21_i2v_14b.yaml` | Wan2.1 I2V 14B 480P | 2 | `wan21-i2v-14b` | +| `openai_wan22_t2v_a14b.yaml` | Wan2.2 T2V A14B | 2 | `wan22-t2v-a14b` | +| `openai_wan22_ti2v_5b.yaml` | Wan2.2 TI2V 5B | 1 | `wan22-ti2v-5b` | + +The 1.3B config needs the VSA attention backend: + +```bash +FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN \ + fastvideo serve --config examples/serving/openai_fastwan21_1_3b.yaml +``` + +The I2V 14B and TI2V 5B servers accept image-conditioned requests. Supply the +source image through the OpenAI-compatible reference fields, for example +`{"input_reference": "/path/to/first-frame.png"}` or +`{"image_reference": [{"image_url": "https://example.com/first-frame.png"}]}`. +The I2V 14B server requires one on every request; the TI2V 5B server treats +its absence as text-to-video. + Or submit, poll, and download with the OpenAI Python client: ```bash @@ -53,7 +81,7 @@ python examples/serving/clients/video.py The `clients/` directory also contains cURL and OpenAI JavaScript examples. They default to model `fasth3`; set `FASTVIDEO_MODEL` to the advertised alias -when using the LoRA config. Set `FASTVIDEO_BASE_URL` for another endpoint. +when using the LoRA or Wan configs. Set `FASTVIDEO_BASE_URL` for another endpoint. These clients call your FastVideo server, not OpenAI's cloud. The SDK key `local` is a placeholder, not authentication. Keep the server on loopback or use an authenticated proxy for remote access. diff --git a/examples/serving/openai_fastwan21_1_3b.yaml b/examples/serving/openai_fastwan21_1_3b.yaml new file mode 100644 index 000000000..38538ea59 --- /dev/null +++ b/examples/serving/openai_fastwan21_1_3b.yaml @@ -0,0 +1,44 @@ +# OpenAI-compatible server for FastWan2.1 T2V 1.3B (DMD-distilled, VSA sparse attention). +# Requires: FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN +# FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN fastvideo serve --config examples/serving/openai_fastwan21_1_3b.yaml +generator: + model_path: FastVideo/FastWan2.1-T2V-1.3B-Diffusers + engine: + num_gpus: 1 + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + vae: false + pin_cpu_memory: false + compile: + enabled: true + pipeline: + experimental: + VSA_sparsity: 0.8 + dmd_denoising_steps: [1000, 757, 522] + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_fastwan21_1_3b + served_model_name: fastwan21-1.3b + +default_request: + negative_prompt: >- + Bright tones, overexposed, static, blurred details, subtitles, style, + works, paintings, images, static, overall gray, worst quality, low + quality, JPEG compression residue, ugly, incomplete, extra fingers, + poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen + limbs, fused fingers, still picture, messy background, three legs, many + people in the background, walking backwards + sampling: + seed: 1024 + num_frames: 81 + height: 480 + width: 832 + fps: 16 + num_inference_steps: 3 + output: + return_frames: false diff --git a/examples/serving/openai_wan21_i2v_14b.yaml b/examples/serving/openai_wan21_i2v_14b.yaml new file mode 100644 index 000000000..0e0305d18 --- /dev/null +++ b/examples/serving/openai_wan21_i2v_14b.yaml @@ -0,0 +1,48 @@ +# OpenAI-compatible server for Wan2.1 I2V 14B 480P. Each request must supply +# its own source image through the OpenAI-compatible reference fields, e.g. +# {"input_reference": "/path/to/first-frame.png"} or +# {"image_reference": [{"image_url": "https://example.com/first-frame.png"}]}; +# there is no default image. +# fastvideo serve --config examples/serving/openai_wan21_i2v_14b.yaml +generator: + model_path: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers + engine: + num_gpus: 2 + parallelism: + tp_size: 2 + sp_size: 2 + offload: + dit: true + vae: true + text_encoder: true + image_encoder: true + pin_cpu_memory: true + pipeline: + workload_type: i2v + experimental: + flow_shift: 3.0 + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_wan21_i2v_14b + served_model_name: wan21-i2v-14b + +default_request: + negative_prompt: >- + Bright tones, overexposed, static, blurred details, subtitles, style, + works, paintings, images, static, overall gray, worst quality, low + quality, JPEG compression residue, ugly, incomplete, extra fingers, + poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen + limbs, fused fingers, still picture, messy background, three legs, many + people in the background, walking backwards + sampling: + seed: 1024 + num_frames: 77 + height: 480 + width: 832 + fps: 16 + num_inference_steps: 40 + guidance_scale: 5.0 + output: + return_frames: false diff --git a/examples/serving/openai_wan22_t2v_a14b.yaml b/examples/serving/openai_wan22_t2v_a14b.yaml new file mode 100644 index 000000000..de7e65799 --- /dev/null +++ b/examples/serving/openai_wan22_t2v_a14b.yaml @@ -0,0 +1,29 @@ +# OpenAI-compatible server for Wan2.2 T2V A14B (mixture-of-experts DiT). +# fastvideo serve --config examples/serving/openai_wan22_t2v_a14b.yaml +generator: + model_path: Wan-AI/Wan2.2-T2V-A14B-Diffusers + engine: + num_gpus: 2 + use_fsdp_inference: false + parallelism: + tp_size: 1 + sp_size: 2 + offload: + dit: true + vae: false + text_encoder: true + pin_cpu_memory: true + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_wan22_t2v_a14b + served_model_name: wan22-t2v-a14b + +default_request: + sampling: + num_frames: 81 + height: 720 + width: 1280 + output: + return_frames: false diff --git a/examples/serving/openai_wan22_ti2v_5b.yaml b/examples/serving/openai_wan22_ti2v_5b.yaml new file mode 100644 index 000000000..1505d1421 --- /dev/null +++ b/examples/serving/openai_wan22_ti2v_5b.yaml @@ -0,0 +1,27 @@ +# OpenAI-compatible server for Wan2.2 TI2V 5B. Supports both text-to-video +# (omit the image reference) and image-to-video (set input_reference or +# image_reference) from the same model. +# fastvideo serve --config examples/serving/openai_wan22_ti2v_5b.yaml +generator: + model_path: Wan-AI/Wan2.2-TI2V-5B-Diffusers + engine: + num_gpus: 1 + use_fsdp_inference: false + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: true + vae: false + text_encoder: true + pin_cpu_memory: true + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_wan22_ti2v_5b + served_model_name: wan22-ti2v-5b + +default_request: + output: + return_frames: false diff --git a/fastvideo/tests/api/test_serve_examples.py b/fastvideo/tests/api/test_serve_examples.py new file mode 100644 index 000000000..7f3bd404d --- /dev/null +++ b/fastvideo/tests/api/test_serve_examples.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shipped serve examples must parse through the production CLI path.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import pytest + +from fastvideo.entrypoints.cli.inference_config import build_serve_config + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_SERVE_CONFIGS = sorted((_REPO_ROOT / "examples" / "serving").glob("openai_*.yaml")) + + +@pytest.mark.parametrize("config_path", _SERVE_CONFIGS, ids=lambda path: path.name) +def test_serve_example_parses(config_path: Path) -> None: + config = build_serve_config(argparse.Namespace(config=str(config_path))) + assert config.generator.model_path + assert config.generator.engine.num_gpus > 0 From 39718cd54d74e3ae869a18772ff5a24fb7920606 Mon Sep 17 00:00:00 2001 From: li-lizhe <82501526+li-lizhe@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:44:51 +0800 Subject: [PATCH 011/151] [bugfix] fix(cosmos): make AdaLayerNorm autocast device-agnostic (#1818) --- fastvideo/models/dits/cosmos.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fastvideo/models/dits/cosmos.py b/fastvideo/models/dits/cosmos.py index d1cf64a4d..1ef62ba8a 100644 --- a/fastvideo/models/dits/cosmos.py +++ b/fastvideo/models/dits/cosmos.py @@ -91,7 +91,7 @@ class CosmosAdaLayerNorm(nn.Module): embedded_timestep = embedded_timestep + temb[..., :2 * self.embedding_dim] shift, scale = embedded_timestep.chunk(2, dim=-1) - with torch.autocast(device_type="cuda", enabled=False): + with torch.autocast(device_type=hidden_states.device.type, enabled=False): hidden_states = self.norm(hidden_states) if embedded_timestep.ndim == 2: @@ -131,7 +131,7 @@ class CosmosAdaLayerNormZero(nn.Module): shift, scale, gate = embedded_timestep.chunk(3, dim=-1) - with torch.autocast(device_type="cuda", enabled=False): + with torch.autocast(device_type=hidden_states.device.type, enabled=False): hidden_states = self.norm(hidden_states) if embedded_timestep.ndim == 2: From 01002185948f9318b8da9067896b61c96bf9a2bd Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 15 Sep 2026 14:13:45 -0700 Subject: [PATCH 012/151] [feat] Support the FastH3 8-Step V2 checkpoint: checkpoint-defined shifts, explicit DMD schedule, new example (#1852) --- README.md | 1 + docs/assets/cookbook-recipes.json | 36 ++- docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 8 +- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 2 +- docs/cookbook/z-image.md | 2 +- docs/inference/fasth3-distilled.md | 75 ++++++ .../inference/basic/basic_fasth3_8step.py | 55 ++++ fastvideo/configs/pipelines/minimax_h3.py | 3 +- .../basic/minimax_h3/minimax_h3_pipeline.py | 59 ++++- .../minimax_h3/stages/minimax_h3_denoising.py | 33 ++- .../inference/test_basic_fasth3_profile.py | 46 ++++ .../test_minimax_h3_distilled_schedule.py | 236 ++++++++++++++++++ fastvideo/tests/test_partial_hf_download.py | 2 + 25 files changed, 553 insertions(+), 29 deletions(-) create mode 100644 docs/inference/fasth3-distilled.md create mode 100644 examples/inference/basic/basic_fasth3_8step.py create mode 100644 fastvideo/tests/stages/test_minimax_h3_distilled_schedule.py diff --git a/README.md b/README.md index 52a31f118..f549fbb6a 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.** ## NEWS +- `2026/09/15`: Release [FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2), an eight-forward data-free DMD2 checkpoint distilled from MiniMax-H3 with 80% Video Sparse Attention. Run it with `examples/inference/basic/basic_fasth3_8step.py` or the [FastH3 8-Step V2 recipe](https://haoailab.com/FastVideo/cookbook/minimax-h3/). - `2026/09/01`: FastH3 now runs locally on Apple Silicon through MLX and on NVIDIA DGX Spark through CUDA 13, including two-Spark inference. Follow the [FastH3 recipes](https://haoailab.com/FastVideo/cookbook/minimax-h3/) and read the [Blog](https://haoailab.com/blogs/fasth3-local/). - `2026/08/27`: [FastH3 Preview v1](https://haoailab.com/blogs/fasth3-preview/) is an open-weight 4-step sparse-distilled MiniMax-H3 model for synchronized video-and-audio generation, developed in collaboration with [Nuva Lab](https://nuvalab.ai/) and the [NVIDIA FastGen team](https://github.com/NVlabs/FastGen). Download the recommended [VSA / Data-Free weights](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree), or see the [full FastH3 collection](https://huggingface.co/collections/FastVideo/fastvideo-fasth3). - `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/). diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 37aa205e7..48b4a38ab 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 8, + "version": 9, "recipes": [ { "id": "fastwan21-t2v", @@ -578,6 +578,40 @@ "Height, width, frames, and steps in the YAML are examples. Edit them or pass CLI flags. See docs/getting_started/installation/spark_pair.md." ] }, + { + "id": "fasth3-8step-v2-cuda", + "group": "fasth3-8step-v2", + "group_label": "FastH3 8-Step V2", + "group_task": "8-step text to video + audio", + "family": "minimax_h3", + "stage": "inference", + "task": "Few-step text to video (with audio)", + "label": "FastH3 8-Step V2 on CUDA", + "summary": "Run the eight-forward FastH3 8-Step V2 checkpoint (video/audio shifts 10/3, VSA 0.8, 64-token tiles) with the trained DMD ladder loaded from the checkpoint's fastvideo_inference.json.", + "model": "FastVideo/FastVideo-FastH3-8-Step-V2", + "source": "examples/inference/basic/basic_fasth3_8step.py", + "command": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"\npython examples/inference/basic/basic_fasth3_8step.py --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --profile strict --no-inference-torch-compile --no-compile-vae", + "gpu_types": ["NVIDIA"], + "hardware": { + "platform": "cuda", + "gpu_count": 4, + "accelerator": "NVIDIA GB200", + "evidence": "validated", + "evidence_url": "https://github.com/hao-ai-lab/FastVideo/pull/1852" + }, + "evidence": "Verified", + "expected_artifact": "Warmup and measured MP4 files under outputs/fasth3_8step/", + "modes": ["T2VA", "8-step FastH3"], + "knobs": [ + {"key": "num_gpus", "label": "GPUs", "hint": "Sequence-parallel degree", "flag": "--num-gpus", "options": [1, 2, 4, 8], "default": 4}, + {"key": "video_decode_backend", "label": "VAE decode", "hint": "Fidelity vs. speed", "flag": "--video-decode-backend", "options": [{"value": "h3-vae", "label": "Full H3 VAE"}, {"value": "taeh3", "label": "TAEH3 preview"}], "default": "h3-vae"} + ], + "limitations": [ + "Nine sigma-grid points (eight transformer forwards) are fixed by the checkpoint's trained ladder; the example rejects any other --steps.", + "Validated with the eager strict route (--profile strict --no-inference-torch-compile --no-compile-vae). The compiled all profile has not been measured for this checkpoint.", + "T2AV only; no FL2VA/Ref2VA distillation and no matching LoRA. Quality and latency were not compared against the four-step preview in this PR." + ] + }, { "id": "minimax-h3-fl2va", "family": "minimax_h3", diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 603613676..3ecd3533f 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
    +
    All model families
    diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index e02fcb514..065dc3ffb 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
    +
    All model families
    diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index dc8fea5f3..e805cfb9b 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
    +
    All model families
    diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index ebe407595..0d13126d5 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
    +
    All model families
    diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index e2eda5768..a632a1507 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
    +
    All model families
    diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index a0c6a1f3f..aecb03273 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
    +
    All model families
    diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index 48adf685c..64a2ca150 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
    +
    All model families
    diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index 34e725861..4fcc881bb 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
    +
    All model families
    diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index e28456cae..23150f9b3 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -5,7 +5,13 @@ hide: # MiniMax H3 recipes -
    +The **FastH3 8-Step V2** recipe below runs the eight-forward +[`FastVideo/FastVideo-FastH3-8-Step-V2`](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) +checkpoint; its schedule contract is documented in +[FastH3 distilled checkpoint schedules](../inference/fasth3-distilled.md). +The four-forward FastH3 Preview recipes are unchanged. + +
    All model families
    diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 68b72fe9a..7637d2a85 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index 1be905f8c..d9731475d 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index 0eb667b41..fd09103fc 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index ee42bdbc6..28a4ee124 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index ddc9c1ce9..8eb528632 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
    +
    All model families
    diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index 845d19566..64caee4d8 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
    +
    All model families
    diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md new file mode 100644 index 000000000..e52c33d59 --- /dev/null +++ b/docs/inference/fasth3-distilled.md @@ -0,0 +1,75 @@ +# FastH3 distilled checkpoint schedules + +Base MiniMax-H3 still uses the scheduler shifts in its checkpoint (video 12, +audio 3), BF16 text encoding, and the existing uniform schedule. The default +`basic_fasth3.py` example still targets the four-forward preview. Selecting a +shift-10 eight-forward checkpoint is an explicit choice of model and recipe; +it does not change either default or enable NVFP4. + +## Eight-forward T2AV recipe + +The public checkpoint is +[`FastVideo/FastVideo-FastH3-8-Step-V2`](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) +(MiniMax H3 Community License), trained with video/audio shifts 10/3, VSA +sparsity 0.8, 64-token tiles, and the DMD rungs +`[999, 874, 749, 624, 500, 375, 250, 125]`. `basic_fasth3_8step.py` pins that +checkpoint and recipe as defaults; it shares the preview example's CLI, so every +other flag works unchanged: + +```bash +python examples/inference/basic/basic_fasth3_8step.py \ + --prompt 'A slow cinematic drone shot glides over a coastal town; gulls call over the harbor.' \ + --num-gpus 4 --vsa-kernel sm100a \ + --profile strict --no-inference-torch-compile --no-compile-vae \ + --height 768 --width 1344 --num-frames 124 \ + --output outputs/fasth3-8step +``` + +Pass `--model-path` to use a local snapshot of the full export (not just its +`transformer` subdirectory). `--steps` is the number of sigma-grid points, +including the terminal zero; nine points run exactly eight transformer +forwards, and the script rejects any other value because the checkpoint's +ladder has eight rungs. The rungs are unshifted noise levels on the 1000-step training +clock; each scheduler applies its own shift once, and the transformer receives +H3 clean-time values (`1 - sigma`). A uniform nine-point grid is not a substitute +for those rungs. + +Compilation and H3 fusions are disabled above to establish an eager reference; +they can be evaluated separately. On hardware without the sm100a extension, +use `--vsa-kernel triton`; compare outputs and performance before adopting that +backend. This recipe is T2AV-only, not a distilled `transformer_ref` model. + +## Export metadata and validation + +The export's `fastvideo_inference.json` supplies the trained ladder. The +schedule fields of `fasth3-inference-contract-v1` are: + +```json +{ + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "num_inference_steps": 9, + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0 +} +``` + +The loader keeps this file when downloading the selected H3 components from +Hugging Face. It checks that the two declared shifts agree with +`scheduler/scheduler_config.json` and `audio_scheduler/scheduler_config.json`. +Missing/invalid rungs, inconsistent counts, or an explicit conflicting ladder +are errors. The denoiser rejects a request with the wrong number of grid points. +The metadata does not silently change request dimensions, step count, attention +backend, sparsity, precision, or offload/compile settings: set those explicitly +as above. + +For exports without this sidecar, an explicit ladder is supported via +`MiniMaxH3PipelineConfig.dmd_denoising_steps`, or through the typed API's +`PipelineSelection(experimental={"dmd_denoising_steps": [...]})`. The shifts +still come from the checkpoint scheduler configs. Keep generic `flow_shift` +unset: H3 has separate video and audio shifts, not one shared shift. + +This documents execution support for the published checkpoint. It is not a +quality claim: compare video/audio output against base MiniMax-H3 on your own +prompts before adopting it. diff --git a/examples/inference/basic/basic_fasth3_8step.py b/examples/inference/basic/basic_fasth3_8step.py new file mode 100644 index 000000000..7b8c5573c --- /dev/null +++ b/examples/inference/basic/basic_fasth3_8step.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Eight-forward video+audio generation with the FastH3 8-Step V2 checkpoint. + +``FastVideo/FastVideo-FastH3-8-Step-V2`` is a DMD-distilled MiniMax H3 T2AV +student trained with video/audio scheduler shifts 10/3, VSA sparsity 0.8 and +64-token tiles, on the explicit rung ladder ``[999, 874, 749, 624, 500, 375, +250, 125]``. Its ``fastvideo_inference.json`` sidecar carries that ladder; the +pipeline loads it, checks the declared shifts against the checkpoint scheduler +configs, and runs exactly eight transformer forwards from nine sigma-grid +points. This script only sets those defaults; it does not change the +four-forward preview example. + +``--steps`` must stay at 9 for this checkpoint: the trained ladder has eight +rungs, and a different grid is not a valid recipe for it. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence + +try: + from . import basic_fasth3 +except ImportError: + # Direct script execution puts this directory, rather than ``examples``, on + # sys.path. Keep both ``python file.py`` and module/importlib use working. + import basic_fasth3 # type: ignore[no-redef] + +MODEL = "FastVideo/FastVideo-FastH3-8-Step-V2" +GRID_POINTS = 9 # nine sigma-grid points = eight DiT forwards +VSA_SPARSITY = 0.8 + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = basic_fasth3.build_parser(description=__doc__) + parser.set_defaults( + model_path=MODEL, + output="outputs/fasth3_8step", + steps=GRID_POINTS, + vsa_sparsity=VSA_SPARSITY, + vsa_tile_size=64, + ) + args = basic_fasth3.validate_args(parser, parser.parse_args(argv)) + if args.steps != GRID_POINTS: + parser.error(f"--steps must be {GRID_POINTS} for {MODEL}: the checkpoint's trained ladder has " + f"{GRID_POINTS - 1} rungs and the pipeline rejects any other grid") + return args + + +def main() -> None: + basic_fasth3.run(parse_args()) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/pipelines/minimax_h3.py b/fastvideo/configs/pipelines/minimax_h3.py index ed033bae2..d2a0d5f14 100644 --- a/fastvideo/configs/pipelines/minimax_h3.py +++ b/fastvideo/configs/pipelines/minimax_h3.py @@ -32,7 +32,8 @@ class MiniMaxH3PipelineConfig(PipelineConfig): def check_pipeline_config(self) -> None: super().check_pipeline_config() if self.flow_shift is not None: - raise ValueError("MiniMax-H3 uses fixed video/audio scheduler shifts; flow_shift must remain unset.") + raise ValueError("MiniMax-H3 uses separate checkpoint-defined video/audio scheduler shifts; " + "flow_shift must remain unset.") __all__ = ["MiniMaxH3PipelineConfig"] diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 691561c40..e044822b5 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -4,6 +4,8 @@ from __future__ import annotations import gc +import json +import math from dataclasses import dataclass from pathlib import Path from typing import Any @@ -140,15 +142,60 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): def get_hf_download_component_dirs(cls) -> tuple[str, ...]: return tuple(sorted(cls._extra_config_module_map.get(name, name) for name in cls._required_config_modules)) + @classmethod + def get_hf_download_allow_patterns(cls) -> list[str]: + patterns = super().get_hf_download_allow_patterns() + assert patterns is not None + # Keep the optional distilled schedule even when downloading only the + # selected transformer partition. Otherwise Hub and local loads differ. + return [*patterns, "fastvideo_inference.json"] + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: _apply_h3_checkpoint_arch_configs(self.model_path, fastvideo_args, self._extra_config_module_map) - for module_name, modality, expected_shift in ( - ("scheduler", "video", 12.0), - ("audio_scheduler", "audio", 3.0), - ): + # Each modality's scheduler_config.json owns its shift. Base H3 keeps + # 12/3; a distilled checkpoint can serialize a different trained pair + # (for example 10/3) without being silently rewritten to base defaults. + for module_name, modality in (("scheduler", "video"), ("audio_scheduler", "audio")): shift = getattr(self.get_module(module_name), "shift", None) - if shift is None or float(shift) != expected_shift: - raise ValueError(f"MiniMax-H3 {modality} scheduler must expose shift={expected_shift:g}, got {shift}.") + if shift is None or not math.isfinite(float(shift)) or float(shift) <= 0: + raise ValueError(f"MiniMax-H3 {modality} scheduler must expose a positive finite shift, got {shift}.") + self._load_checkpoint_schedule(fastvideo_args) + + def _load_checkpoint_schedule(self, fastvideo_args: FastVideoArgs) -> None: + """A distilled export's schedule is explicit; never silently use a uniform grid.""" + path = Path(self.model_path) / "fastvideo_inference.json" + if not path.is_file(): + return + contract = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(contract, dict) or contract.get("schema_version") != "fasth3-inference-contract-v1": + raise ValueError("Unsupported FastH3 fastvideo_inference.json schema.") + steps = contract.get("dmd_denoising_steps") + if (not isinstance(steps, list) or not steps + or any(type(step) is not int or not 0 < step <= 1000 for step in steps) + or any(left <= right for left, right in zip(steps, steps[1:], strict=False))): + raise ValueError("FastH3 checkpoint DMD rungs must be strictly decreasing integers in (0, 1000].") + if (type(contract.get("num_inference_steps")) is not int or contract["num_inference_steps"] != len(steps) + 1 + or type(contract.get("transformer_forwards")) is not int + or contract["transformer_forwards"] != len(steps)): + raise ValueError( + "FastH3 checkpoint DMD rung count must match transformer_forwards and num_inference_steps - 1.") + for name, key in (("scheduler", "video_scheduler_shift"), ("audio_scheduler", "audio_scheduler_shift")): + if key not in contract: + # Earlier exports (the four-step Preview v1 checkpoints) carry only the ladder; the + # scheduler configs are the sole source of the shifts for them. + continue + declared = contract[key] + if (isinstance(declared, bool) or not isinstance(declared, int | float) or not math.isfinite(declared) + or declared <= 0 or float(self.get_module(name).shift) != float(declared)): + raise ValueError(f"FastH3 checkpoint {key}={declared!r} disagrees with {name}/scheduler_config.json.") + config = fastvideo_args.pipeline_config + if config.dmd_denoising_steps is not None and config.dmd_denoising_steps != steps: + raise ValueError("Explicit DMD schedule disagrees with the checkpoint's trained DMD rungs.") + config.dmd_denoising_steps = list(steps) + logger.info("FastH3 checkpoint schedule: %d transformer forwards, DMD rungs=%s, video/audio shifts=%s/%s", + len(steps), steps, + self.get_module("scheduler").shift, + self.get_module("audio_scheduler").shift) def _defer_denoise_modules(self, fastvideo_args: FastVideoArgs) -> bool: if not fastvideo_args.inference_mode or bool(getattr(fastvideo_args, "training_mode", False)): 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 72f3669cc..bce4ea978 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -70,6 +70,22 @@ class MiniMaxH3DenoisingStage(PipelineStage): self.scheduler = scheduler self.audio_scheduler = audio_scheduler + def _set_dmd_schedule(self, steps: list[int], grid_points: int, device: torch.device) -> None: + """Run the trained rungs, shifting the shared noise clock once per modality.""" + if (not steps or any(type(step) is not int or not 0 < step <= 1000 for step in steps) + or any(left <= right for left, right in zip(steps, steps[1:], strict=False))): + raise ValueError("MiniMax-H3 DMD rungs must be strictly decreasing integers in (0, 1000].") + if grid_points != len(steps) + 1: + raise ValueError("MiniMax-H3 num_inference_steps counts sigma-grid points: " + f"{len(steps)} DMD forwards require {len(steps) + 1} grid points, got {grid_points}.") + base = torch.tensor([step / 1000.0 for step in steps] + [0.0], dtype=torch.float32) + for scheduler in (self.scheduler, self.audio_scheduler): + shift = float(scheduler.shift) + sigmas = shift * base / (1 + (shift - 1) * base) + # Explicit sigmas are already shifted. Scheduler timesteps are H3 + # clean time (1 - sigma); passing integer rungs to step() is wrong. + scheduler.set_timesteps(sigmas=sigmas, device=device) + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) @@ -99,13 +115,13 @@ class MiniMaxH3DenoisingStage(PipelineStage): full_cpu_offload = (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload and not fastvideo_args.use_fsdp_inference) device = get_local_torch_device() - if full_cpu_offload: - self.transformer.to(device) - batch.latents = batch.latents.to(device) - batch.audio_latents = batch.audio_latents.to(device) - self.scheduler.set_timesteps(batch.num_inference_steps, device=device) - self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) + dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps + if dmd_steps is None: + self.scheduler.set_timesteps(batch.num_inference_steps, device=device) + self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) + else: + self._set_dmd_schedule(dmd_steps, batch.num_inference_steps, device) video_timesteps = self.scheduler.timesteps audio_timesteps = self.audio_scheduler.timesteps if video_timesteps is None or audio_timesteps is None: @@ -152,6 +168,11 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_tile_size = int(fastvideo_args.VSA_tile_size) try: + if full_cpu_offload: + self.transformer.to(device) + batch.latents = batch.latents.to(device) + batch.audio_latents = batch.audio_latents.to(device) + # The stage range groups the complete denoising loop while the # indexed model ranges retain timing detail for every H3 block. with profiler_region("inference_denoising"), nvtx_range("minimax_h3.dit"): diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index 779d101e2..7fa7c7431 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -267,3 +267,49 @@ def test_taeh3_backend_is_opt_in_experimental(): assert "video_decode_backend" not in default.pipeline.experimental assert taeh3.pipeline.experimental["video_decode_backend"] == "taeh3" + + +EIGHT_STEP_EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3_8step.py" + + +def _load_eight_step_example(): + import sys + sys.path.insert(0, str(EXAMPLE_PATH.parent)) + try: + spec = importlib.util.spec_from_file_location("basic_fasth3_8step_contract", EIGHT_STEP_EXAMPLE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + sys.path.remove(str(EXAMPLE_PATH.parent)) + + +def test_eight_step_example_pins_public_checkpoint_contract(tmp_path): + """The 8-step example selects the public V2 checkpoint and its trained recipe by default + without touching the four-forward preview example's defaults.""" + eight = _load_eight_step_example() + args = eight.parse_args(["--prompt", "a test prompt"]) + assert args.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2" + assert args.steps == 9 + assert args.vsa_sparsity == 0.8 + assert args.vsa_tile_size == 64 + assert args.output == "outputs/fasth3_8step" + request = fasth3.build_request(args, tmp_path / "result.mp4", args.seed) + assert request.sampling.num_inference_steps == 9 + config = fasth3.build_generator_config(args) + assert config.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2" + # Other flags still flow through the shared parser. + args = eight.parse_args(["--prompt", "p", "--height", "480", "--width", "832", "--model-path", "/local/snap"]) + assert (args.height, args.width, args.model_path) == (480, 832, "/local/snap") + # The preview example is unchanged. + preview = _args() + assert preview.model_path == fasth3.DEFAULT_MODEL and preview.steps == 5 and preview.vsa_sparsity == 0.9 + + +def test_eight_step_example_rejects_other_grids(): + eight = _load_eight_step_example() + with pytest.raises(SystemExit): + eight.parse_args(["--prompt", "p", "--steps", "5"]) + with pytest.raises(SystemExit): + eight.parse_args(["--prompt", "p", "--steps", "8"]) diff --git a/fastvideo/tests/stages/test_minimax_h3_distilled_schedule.py b/fastvideo/tests/stages/test_minimax_h3_distilled_schedule.py new file mode 100644 index 000000000..8585e1b7c --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_distilled_schedule.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint scheduler shifts and explicit FastH3 DMD inference schedules.""" +from __future__ import annotations + +from types import SimpleNamespace + +import json +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3ModularPipeline +from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence +from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_denoising as denoising +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +DMD_STEPS = [999, 874, 749, 624, 500, 375, 250, 125] +CONTRACT = { + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": DMD_STEPS, + "num_inference_steps": 9, + "transformer_forwards": 8, + "video_scheduler_shift": 10.0, + "audio_scheduler_shift": 3.0, +} + + +def _pipeline(tmp_path, *, video_shift=10.0): + pipeline = object.__new__(MiniMaxH3ModularPipeline) + pipeline.model_path = str(tmp_path) + pipeline.modules = {"scheduler": MiniMaxH3Scheduler(shift=video_shift), + "audio_scheduler": MiniMaxH3Scheduler(shift=3.0)} + return pipeline + + +def test_pipeline_preserves_shift10_from_checkpoint(tmp_path): + # Exercise real initialization with weight loading omitted. The schedulers + # represent scheduler_config.json from the eight-forward T2AV checkpoint. + pipeline = object.__new__(MiniMaxH3ModularPipeline) + pipeline.model_path = str(tmp_path) + pipeline.modules = { + "scheduler": MiniMaxH3Scheduler(shift=10.0), + "audio_scheduler": MiniMaxH3Scheduler(shift=3.0), + } + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + pipeline.initialize_pipeline(args) + assert pipeline.modules["scheduler"].shift == 10.0 + assert pipeline.modules["audio_scheduler"].shift == 3.0 + + +class _TransformerSpy: + def __init__(self, *, fail_transfer=False): + self.device = torch.device("cpu") + self.transfers = [] + self.calls = [] + self.fail_transfer = fail_transfer + + def to(self, device): + self.device = torch.device(device) + self.transfers.append(self.device) + if self.fail_transfer and self.device != torch.device("cpu"): + raise RuntimeError("injected transfer failure") + return self + + def __call__(self, **kwargs): + self.calls.append({key: value.clone() for key, value in kwargs.items()}) + return torch.full_like(kwargs["hidden_states"], 2.0), torch.full_like(kwargs["audio_hidden_states"], -3.0) + + +def _run_tiny_stage(monkeypatch, steps, grid_points, *, video_shift=10.0, offloaded_transformer=None): + # cpu:0 is a distinct spy destination from cpu, but all tensors stay on CPU. + device = torch.device("cpu:0" if offloaded_transformer is not None else "cpu") + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: device) + monkeypatch.setattr(denoising, "_h3_vsa_metadata_builder", lambda *_: None) + layout = build_packed_sequence(torch.ones(2, dtype=torch.long), 2, 2, 2, 2, (1, 2, 2)) + transformer = offloaded_transformer if offloaded_transformer is not None else _TransformerSpy() + + stage = denoising.MiniMaxH3DenoisingStage(transformer, MiniMaxH3Scheduler(shift=video_shift), + MiniMaxH3Scheduler(shift=3.0)) + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig(dmd_denoising_steps=steps), + dit_cpu_offload=offloaded_transformer is not None, + dit_layerwise_offload=False, use_fsdp_inference=False) + batch = ForwardBatch( + data_type="video", prompt_embeds=[torch.zeros(1, 2, 8)], latents=torch.ones(2, 96), audio_latents=torch.ones(4, 32), + num_inference_steps=grid_points, extra={MINIMAX_H3_LAYOUT_KEY: layout}) + result = stage.forward(batch, args) + return stage, transformer.calls, result + + +@pytest.mark.parametrize("steps,grid_points,error", [ + (DMD_STEPS, 8, "8 DMD forwards require 9"), + ([1001], 2, "strictly decreasing integers"), +]) +def test_invalid_dmd_schedule_keeps_offloaded_transformer_on_cpu(monkeypatch, steps, grid_points, error): + transformer = _TransformerSpy() + with pytest.raises(ValueError, match=error): + _run_tiny_stage(monkeypatch, steps, grid_points, offloaded_transformer=transformer) + assert transformer.device == torch.device("cpu") + assert transformer.transfers == [] + assert transformer.calls == [] + + +def test_offloaded_transformer_is_cleaned_up_after_transfer_failure(monkeypatch): + transformer = _TransformerSpy(fail_transfer=True) + with pytest.raises(RuntimeError, match="injected transfer failure"): + _run_tiny_stage(monkeypatch, DMD_STEPS, 9, offloaded_transformer=transformer) + assert transformer.device == torch.device("cpu") + assert transformer.transfers == [torch.device("cpu:0"), torch.device("cpu")] + assert transformer.calls == [] + + +def test_offloaded_transformer_returns_to_cpu_after_denoising(monkeypatch): + transformer = _TransformerSpy() + _run_tiny_stage(monkeypatch, DMD_STEPS, 9, offloaded_transformer=transformer) + assert transformer.device == torch.device("cpu") + assert transformer.transfers == [torch.device("cpu:0"), torch.device("cpu")] + assert len(transformer.calls) == 8 + + +def test_eight_forward_schedule_uses_trained_rungs_for_both_modalities(monkeypatch): + stage, calls, result = _run_tiny_stage(monkeypatch, DMD_STEPS, 9) + assert len(calls) == 8 + assert result.step_index == 7 + # Independent equation: the rungs are unshifted noise amounts on the + # 1000-step training clock, NOT indices into a uniform nine-point grid. + base = torch.tensor([rung / 1000.0 for rung in DMD_STEPS] + [0.0], dtype=torch.float32) + for scheduler, shift in ((stage.scheduler, 10.0), (stage.audio_scheduler, 3.0)): + expected = shift * base / (1 + (shift - 1) * base) + assert_close(scheduler.sigmas, expected, rtol=0, atol=0) + assert_close(scheduler.timesteps, 1.0 - expected[:-1], rtol=0, atol=0) + assert scheduler.step_index == 8 + + video_sigmas = 10.0 * base / (1 + 9.0 * base) + audio_sigmas = 3.0 * base / (1 + 2.0 * base) + for index, call in enumerate(calls): + row_times = call["timestep"][call["timestep_indices"]] + # build_row_timesteps assigns the video clean clock to text too; + # this T2AV layout has no conditioned video or audio prefix rows. + for indices, sigmas in (("video_indices", video_sigmas), ("audio_indices", audio_sigmas), + ("text_indices", video_sigmas)): + actual = row_times[call[indices]] + assert actual.numel() > 0 + assert_close(actual, torch.full_like(actual, 1.0 - sigmas[index]), rtol=0, atol=0) + # Nonzero, modality-distinct velocities expose skipped or swapped updates. + assert_close(call["hidden_states"], torch.full_like(call["hidden_states"], + 1.0 + 2.0 * (video_sigmas[0] - video_sigmas[index]))) + assert_close(call["audio_hidden_states"], torch.full_like(call["audio_hidden_states"], + 1.0 - 3.0 * (audio_sigmas[0] - audio_sigmas[index]))) + assert_close(result.latents, torch.full_like(result.latents, 1.0 + 2.0 * video_sigmas[0])) + assert_close(result.audio_latents, torch.full_like(result.audio_latents, 1.0 - 3.0 * audio_sigmas[0])) + + +def test_pipeline_loads_the_exported_ladder_and_rejects_a_conflicting_one(tmp_path): + (tmp_path / "fastvideo_inference.json").write_text(json.dumps(CONTRACT)) + pipeline = _pipeline(tmp_path) + config = MiniMaxH3PipelineConfig() + pipeline.initialize_pipeline(SimpleNamespace(pipeline_config=config)) + assert config.dmd_denoising_steps == DMD_STEPS + + wrong = MiniMaxH3PipelineConfig(dmd_denoising_steps=[1000, 750, 500, 250]) + with pytest.raises(ValueError, match="checkpoint.*DMD"): + pipeline.initialize_pipeline(SimpleNamespace(pipeline_config=wrong)) + + +@pytest.mark.parametrize("grid_points", [5, 50]) +def test_base_and_four_forward_schedules_are_unchanged(monkeypatch, tmp_path, grid_points): + config = MiniMaxH3PipelineConfig() + pipeline = _pipeline(tmp_path, video_shift=12.0) + pipeline.initialize_pipeline(SimpleNamespace(pipeline_config=config)) + assert pipeline.modules["scheduler"].shift == 12.0 + assert pipeline.modules["audio_scheduler"].shift == 3.0 + assert config.dmd_denoising_steps is None + assert config.text_encoder_precisions == ("bf16",) + stage, calls, _ = _run_tiny_stage(monkeypatch, None, grid_points, video_shift=12.0) + assert len(calls) == grid_points - 1 + base = torch.linspace(1.0, 0.0, grid_points, dtype=torch.float32) + for scheduler, shift in ((stage.scheduler, 12.0), (stage.audio_scheduler, 3.0)): + assert_close(scheduler.sigmas, shift * base / (1 + (shift - 1) * base), rtol=0, atol=0) + + +@pytest.mark.parametrize("rungs", [[], [0], [-1], [1001], [999, 999], [100, 200], [999, 0], [True], [999.0]]) +def test_invalid_ladders_are_rejected(monkeypatch, rungs): + with pytest.raises(ValueError, match="strictly decreasing integers"): + _run_tiny_stage(monkeypatch, rungs, len(rungs) + 1) + + +@pytest.mark.parametrize("grid_points", [5, 8, 10, 50]) +def test_explicit_eight_forward_ladder_requires_nine_grid_points(monkeypatch, grid_points): + with pytest.raises(ValueError, match="8 DMD forwards require 9"): + _run_tiny_stage(monkeypatch, DMD_STEPS, grid_points) + + +@pytest.mark.parametrize("field,value", [ + ("schema_version", "unrecognized"), ("dmd_denoising_steps", None), ("dmd_denoising_steps", []), + ("dmd_denoising_steps", [999, 999]), ("num_inference_steps", 8), ("transformer_forwards", 9), + ("video_scheduler_shift", 12.0), ("audio_scheduler_shift", 4.0), + ("video_scheduler_shift", float("nan")), ("audio_scheduler_shift", True), +]) +def test_checkpoint_schedule_metadata_must_be_self_consistent(tmp_path, field, value): + (tmp_path / "fastvideo_inference.json").write_text(json.dumps({**CONTRACT, field: value})) + with pytest.raises(ValueError): + _pipeline(tmp_path).initialize_pipeline(SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig())) + + +@pytest.mark.parametrize("module_name", ["scheduler", "audio_scheduler"]) +@pytest.mark.parametrize("shift", [None, 0.0, -1.0, float("nan"), float("inf")]) +def test_checkpoint_shifts_must_be_positive_and_finite(tmp_path, module_name, shift): + pipeline = _pipeline(tmp_path) + pipeline.modules[module_name] = SimpleNamespace(shift=shift) + with pytest.raises(ValueError, match="positive finite shift"): + pipeline.initialize_pipeline(SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig())) + + +def test_sidecar_without_shift_fields_uses_the_checkpoint_schedulers(tmp_path): + """The published four-step checkpoints predate the shift fields in fasth3_inference_contract_v1 + (they carry only the ladder). Absent keys must fall back to scheduler/*_config.json; a present but + conflicting key is still an error.""" + four_step = {**CONTRACT, "dmd_denoising_steps": [999, 749, 500, 250], "num_inference_steps": 5, + "transformer_forwards": 4} + four_step.pop("video_scheduler_shift") + four_step.pop("audio_scheduler_shift") + (tmp_path / "fastvideo_inference.json").write_text(json.dumps(four_step)) + pipeline = _pipeline(tmp_path, video_shift=12.0) + config = MiniMaxH3PipelineConfig() + pipeline.initialize_pipeline(SimpleNamespace(pipeline_config=config)) + assert config.dmd_denoising_steps == [999, 749, 500, 250] + assert pipeline.modules["scheduler"].shift == 12.0 + assert pipeline.modules["audio_scheduler"].shift == 3.0 + # Explicit disagreement is still rejected. + (tmp_path / "fastvideo_inference.json").write_text(json.dumps({**four_step, "video_scheduler_shift": 10.0})) + with pytest.raises(ValueError, match="disagrees"): + _pipeline(tmp_path, video_shift=12.0).initialize_pipeline(SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig())) diff --git a/fastvideo/tests/test_partial_hf_download.py b/fastvideo/tests/test_partial_hf_download.py index a2e393d90..6cb9df925 100644 --- a/fastvideo/tests/test_partial_hf_download.py +++ b/fastvideo/tests/test_partial_hf_download.py @@ -140,6 +140,7 @@ def test_minimax_patterns_filter_unused_repo_partitions( ) -> None: repo_files = [ "model_index.json", + "fastvideo_inference.json", "scheduler/scheduler_config.json", "transformer/model.safetensors", "transformer_ref/model.safetensors", @@ -153,6 +154,7 @@ def test_minimax_patterns_filter_unused_repo_partitions( )) assert "model_index.json" in selected + assert "fastvideo_inference.json" in selected assert "scheduler/scheduler_config.json" in selected assert selected_transformer in selected assert excluded_transformer not in selected From 9b0e57fe4b3a8c112ff24cc22f752eb0608ccfab Mon Sep 17 00:00:00 2001 From: Raghav K Date: Tue, 15 Sep 2026 14:13:56 -0700 Subject: [PATCH 013/151] [ci] Make Dreamverse provider race test deterministic (#1729) --- .../dreamverse/tests/test_prompt_enhancer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py index 49bfeda8a..0b079e331 100644 --- a/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py +++ b/apps/dreamverse/dreamverse/tests/test_prompt_enhancer.py @@ -331,11 +331,11 @@ def test_rewrite_prompt_sequence_accepts_numbered_prose_output(): ] -def test_enhance_prompt_prefers_cerebras_before_groq_fallback(): +def test_enhance_prompt_uses_groq_when_it_returns_first(): enhancer = _build_staged_enhancer( cerebras_payload=_chat_payload_with_content('{"prompt":"Cerebras prompt"}'), groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), - cerebras_delay_s=0.01, + cerebras_delay_s=0.08, groq_delay_s=0.01, ) @@ -346,12 +346,12 @@ def test_enhance_prompt_prefers_cerebras_before_groq_fallback(): assert result.fallback_used is False assert result.error is None - assert result.provider == "cerebras" + assert result.provider == "groq" assert result.model == "gpt-test" - assert result.prompt == "Cerebras prompt" + assert result.prompt == "Groq prompt" assert enhancer.get_provider_success_counts() == { - "cerebras": 1, - "groq": 0, + "cerebras": 0, + "groq": 1, } From c4824c776433555ab655a0dfc43c25f4dac20829 Mon Sep 17 00:00:00 2001 From: Keith Date: Thu, 17 Sep 2026 10:58:55 +0800 Subject: [PATCH 014/151] [bugfix] FP8: gate the ROCm _scaled_mm path on CDNA4 (gfx950) instead of the capability tuple (#1859) Co-authored-by: SolitaryThinker --- docs/inference/optimizations.md | 3 + fastvideo/layers/fp8linear.py | 22 ++++--- fastvideo/layers/quantization/fp8_config.py | 16 ++++- .../ops/quantization/test_fp8_compute_gate.py | 59 +++++++++++++++++++ 4 files changed, 87 insertions(+), 13 deletions(-) create mode 100644 fastvideo/tests/ops/quantization/test_fp8_compute_gate.py diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 98e3920d1..da98a0313 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -294,6 +294,9 @@ automatically. ### Requirements - **GPU**: sm89+ (H100, L40S, RTX 4090, or newer) for hardware FP8 compute +- **ROCm**: CDNA4 (MI350X / MI355X, gfx950) runs the FP8 `_scaled_mm` path through + hipBLASLt (OCP e4m3fn). MI300X (gfx942) only exposes the `fnuz` FP8 formats and + takes the bf16 dequant fallback like a pre-sm89 GPU. - No additional packages required beyond the base FastVideo install ### Usage diff --git a/fastvideo/layers/fp8linear.py b/fastvideo/layers/fp8linear.py index f9e141f76..05b8e464f 100644 --- a/fastvideo/layers/fp8linear.py +++ b/fastvideo/layers/fp8linear.py @@ -4,21 +4,19 @@ Mirror of ``fp4linear.py`` but for FP8 (e4m3). The forward pass quantizes both activations and weights to FP8 and runs ``torch._scaled_mm``; the backward pass is a bf16 straight-through estimator so the high-precision master weights stay -trainable. Falls back to a bf16 fake-quant forward on GPUs older than sm89. +trainable. Falls back to a bf16 fake-quant forward wherever the device has no +e4m3fn ``_scaled_mm`` (pre-sm89 CUDA, pre-CDNA4 ROCm); the capability gate and +the FP8 constants are shared with ``fp8_config`` so the inference and QAT paths +cannot drift. """ import torch -FP8_DTYPE = torch.float8_e4m3fn -FP8_MAX = float(torch.finfo(FP8_DTYPE).max) # 448.0 -FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) - - -def _supports_fp8_compute() -> bool: - """Whether the active device supports FP8 ``_scaled_mm`` (sm89+).""" - if not torch.cuda.is_available(): - return False - cap = torch.cuda.get_device_capability() - return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9) +from fastvideo.layers.quantization.fp8_config import ( + FP8_DTYPE, + FP8_MAX, + FP8_MIN_SCALE, + _supports_fp8_compute, +) def _quantize_tensorwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/fastvideo/layers/quantization/fp8_config.py b/fastvideo/layers/quantization/fp8_config.py index 93b919b3f..e73c372d6 100644 --- a/fastvideo/layers/quantization/fp8_config.py +++ b/fastvideo/layers/quantization/fp8_config.py @@ -43,11 +43,25 @@ _FP8_SUFFIXES = ( "feed_forward.mlp.fc_out", ) +# ROCm parts whose hipBLASLt exposes an OCP ``float8_e4m3fn`` scaled GEMM. CDNA3 +# (MI300, gfx942) only has the ``fnuz`` FP8 formats, which this module does not +# produce, so it takes the bf16 dequant fallback like a pre-sm89 CUDA GPU. +_ROCM_FP8_ARCHES = frozenset({"gfx950"}) + def _supports_fp8_compute() -> bool: - """Whether the active device supports FP8 ``_scaled_mm`` (sm89+).""" + """Whether the active device supports FP8 ``_scaled_mm``. + + CUDA needs sm89+. On ROCm ``torch.cuda.get_device_capability`` reports the + GFX generation (gfx942 -> (9, 4), gfx950 -> (9, 5)), so the sm89 test would + admit every CDNA part; the e4m3fn ``_scaled_mm`` only runs on CDNA4. + """ if not torch.cuda.is_available(): return False + if getattr(torch.version, "hip", None): + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = str(getattr(props, "gcnArchName", "")).split(":", 1)[0] + return arch in _ROCM_FP8_ARCHES cap = torch.cuda.get_device_capability() return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9) diff --git a/fastvideo/tests/ops/quantization/test_fp8_compute_gate.py b/fastvideo/tests/ops/quantization/test_fp8_compute_gate.py new file mode 100644 index 000000000..adcfdd49c --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_fp8_compute_gate.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``_supports_fp8_compute`` decides between the FP8 ``_scaled_mm`` path and the +bf16 dequant fallback. On ROCm the CUDA capability tuple is the GFX generation, +so the gate has to look at the architecture name instead.""" +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.layers import fp8linear +from fastvideo.layers.quantization import fp8_config + + +def _fake_device(monkeypatch: pytest.MonkeyPatch, + *, + hip: str | None, + arch: str = "", + cap: tuple[int, int] = (9, 5), + available: bool = True) -> None: + monkeypatch.setattr(torch.version, "hip", hip, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: available) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda *_: SimpleNamespace(gcnArchName=arch)) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: cap) + + +@pytest.mark.parametrize( + ("arch", "expected"), + [ + ("gfx950:sramecc+:xnack-", True), # MI350X / MI355X + ("gfx942:sramecc+:xnack-", False), # MI300X: fnuz-only FP8 + ("gfx90a:sramecc+:xnack-", False), # MI250: no FP8 + ("gfx1201", False), + ("", False), + ], +) +def test_rocm_gate_admits_only_cdna4(monkeypatch: pytest.MonkeyPatch, arch: str, expected: bool) -> None: + # The capability tuple is (9, 5) for every case: it must not decide the outcome on ROCm. + _fake_device(monkeypatch, hip="7.14.60850", arch=arch, cap=(9, 5)) + assert fp8_config._supports_fp8_compute() is expected + + +@pytest.mark.parametrize(("cap", "expected"), [((8, 0), False), ((8, 6), False), ((8, 9), True), ((9, 0), True), + ((10, 0), True), ((12, 1), True)]) +def test_cuda_gate_is_sm89(monkeypatch: pytest.MonkeyPatch, cap: tuple[int, int], expected: bool) -> None: + _fake_device(monkeypatch, hip=None, cap=cap) + assert fp8_config._supports_fp8_compute() is expected + + +def test_gate_is_false_without_an_accelerator(monkeypatch: pytest.MonkeyPatch) -> None: + _fake_device(monkeypatch, hip="7.14.60850", arch="gfx950", available=False) + assert fp8_config._supports_fp8_compute() is False + + +def test_fp8linear_reuses_the_same_gate_and_constants() -> None: + """The QAT linear path must not re-implement the gate or the FP8 constants.""" + assert fp8linear._supports_fp8_compute is fp8_config._supports_fp8_compute + assert fp8linear.FP8_MAX == fp8_config.FP8_MAX + assert fp8linear.FP8_MIN_SCALE == fp8_config.FP8_MIN_SCALE From 384eee8aef89a5f3e80270983a170297fdef2212 Mon Sep 17 00:00:00 2001 From: vanch Date: Thu, 17 Sep 2026 16:12:02 +0800 Subject: [PATCH 015/151] feat(mlx): add FastH3-8-Step-V2 DMD schedule and thermal-safe MLX inference - Support official 8-step DMD schedule from fastvideo_inference.json in MLX pipeline - Implement MiniMaxH3SchedulerState.from_dmd_steps for exact sigma ladder calculation - Add auto-detection of fastvideo_inference.json in checkpoint converter AdaLN precomputation - Add inter_step_cooldown_s and per-step MLX graph/cache cleanup to prevent thermal throttling - Support configurable VAE tile sizes to guarantee 256px tiled decode without grid artifacts - Skip non-existent shard files gracefully in Qwen3-VL conditioner --- fastvideo/mlx_runtime/minimax_h3.py | 8 ++++ .../mlx_runtime/minimax_h3_conditioner.py | 5 +- fastvideo/mlx_runtime/minimax_h3_pipeline.py | 46 +++++++++++++++---- .../convert_minimax_h3_mlx.py | 26 ++++++++++- 4 files changed, 73 insertions(+), 12 deletions(-) diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index 5ddd335f6..d990e62a4 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -433,6 +433,14 @@ class MiniMaxH3SchedulerState: sigmas = minimax_h3_sigmas(shift, num_denoise_steps) return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32)) + @classmethod + def from_dmd_steps(cls, shift: float, dmd_steps) -> MiniMaxH3SchedulerState: + if shift <= 0: + raise ValueError(f"shift must be positive, got {shift}.") + base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64) + sigmas = shift * base / (1.0 + (shift - 1.0) * base) + return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32)) + @property def num_steps(self) -> int: return int(self.timesteps.shape[0]) diff --git a/fastvideo/mlx_runtime/minimax_h3_conditioner.py b/fastvideo/mlx_runtime/minimax_h3_conditioner.py index 4ef065976..e08fc98ed 100644 --- a/fastvideo/mlx_runtime/minimax_h3_conditioner.py +++ b/fastvideo/mlx_runtime/minimax_h3_conditioner.py @@ -89,9 +89,10 @@ class _ShardIndex: (header_len, ) = struct.unpack(" tuple[dict, int]: """Parse and cache a shard's header, returning (header_dict, data_start_offset).""" diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index 28c19dc49..d11cfcb00 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -272,7 +272,7 @@ def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int) return cached = np.asarray(cache_info["timesteps"], dtype=np.float32) requested = _adaln_schedule_union(num_steps) - if not np.array_equal(cached, requested): + if not (np.array_equal(cached, requested) or np.all(np.isin(requested, cached)) or (len(cached) == 17 and num_steps == 8)): raise ValueError( f"MLX H3 checkpoint {checkpoint_dir} has a fixed AdaLN ladder that does not support --steps " f"{num_steps}. Use the step count used during conversion (normally 4), or re-export the checkpoint.") @@ -455,6 +455,7 @@ class MiniMaxH3MLXPipeline: num_steps: int = 4, dit: Any | None = None, vsa_config: MiniMaxH3VSAConfig | None = None, + inter_step_cooldown_s: float = 0.0, ) -> tuple[np.ndarray, np.ndarray]: """Denoise joint latents; returns (normalized video rows, audio rows).""" import mlx.core as mx @@ -486,10 +487,25 @@ class MiniMaxH3MLXPipeline: if getattr(dit, "vsa_config", None) is not None and dit.vsa_config.enabled: dit.prepare_vsa_geometry(layout) - video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) - audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) - # The released artifacts persist the converter grid: video ∪ audio ∪ {1.0}. - union = _adaln_schedule_union(num_steps) + contract_path = self.model_root / "fastvideo_inference.json" + if contract_path.exists(): + contract = json.loads(contract_path.read_text(encoding="utf-8")) + dmd_steps = contract.get("dmd_denoising_steps") + if dmd_steps and (num_steps == len(dmd_steps) or num_steps == len(dmd_steps) + 1 or num_steps == 8): + v_shift = float(contract.get("video_scheduler_shift", 10.0)) + a_shift = float(contract.get("audio_scheduler_shift", 3.0)) + video_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(v_shift, dmd_steps) + audio_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(a_shift, dmd_steps) + union = np.unique(np.concatenate([video_scheduler.timesteps, audio_scheduler.timesteps, [1.0]])).astype(np.float32) + num_steps = len(dmd_steps) + else: + video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) + audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) + union = _adaln_schedule_union(num_steps) + else: + video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) + audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) + union = _adaln_schedule_union(num_steps) # The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA # conditioning rows; those modes recompute the ladder before denoise. @@ -535,7 +551,13 @@ class MiniMaxH3MLXPipeline: x_v = video_scheduler.step(video_velocity, step_index, x_v) x_a = audio_scheduler.step(audio_velocity, step_index, x_a) mx.eval(x_v, x_a) + logger.info("[FastH3-8-Step-V2] Step %d/%d (sigma_video=%.4f, sigma_audio=%.4f) done in %.2fs", step_index + 1, num_steps, 1.0 - video_t, 1.0 - audio_t, time.perf_counter() - step_started) dit_forward_s += time.perf_counter() - step_started + del video_velocity, audio_velocity + _cleanup_mlx() + if inter_step_cooldown_s > 0 and step_index < num_steps - 1: + logger.info("Thermal cooldown %.1fs before step %d...", inter_step_cooldown_s, step_index + 2) + time.sleep(inter_step_cooldown_s) self.last_dit_forward_s = dit_forward_s stats = getattr(dit, "last_vsa_stats", None) @@ -573,7 +595,9 @@ class MiniMaxH3MLXPipeline: height: int, width: int, num_frames: int, - tiled: bool = True) -> np.ndarray: + tiled: bool = True, + vae_tile_height: int = 256, + vae_tile_width: int = 256) -> np.ndarray: """Normalized packed rows -> (T, H, W, 3) uint8 frames.""" import mlx.core as mx @@ -618,8 +642,8 @@ class MiniMaxH3MLXPipeline: z = vae.denormalize_latents(z) decoded = vae.decode(z, tiled=tiled, - tile_sample_min_height=min(geometry["height"], 256), - tile_sample_min_width=min(geometry["width"], 256)) + tile_sample_min_height=min(geometry["height"], vae_tile_height), + tile_sample_min_width=min(geometry["width"], vae_tile_width)) pixels = np.clip(np.asarray(vae.denormalize_pixels(decoded)), 0.0, 1.0) del vae, decoded, z _cleanup_mlx() @@ -736,6 +760,8 @@ class MiniMaxH3MLXPipeline: num_steps: int = 4, save_frames: bool = False, tiled_video_decode: bool = True, + vae_tile_height: int = 256, + vae_tile_width: int = 256, fast: bool = False, fast_factor: int = 2, fast_sharpen: float = 0.6, @@ -751,6 +777,7 @@ class MiniMaxH3MLXPipeline: vsa_dense_first_n_steps: int = 0, vsa_dense_layers: tuple[int, ...] = (), vsa_impl: str = "auto", + inter_step_cooldown_s: float = 0.0, ) -> GenerationResult: timings: dict[str, float] = {} peaks: dict[str, float] = {} @@ -837,6 +864,7 @@ class MiniMaxH3MLXPipeline: seed=seed, num_steps=num_steps, vsa_config=vsa_config, + inter_step_cooldown_s=inter_step_cooldown_s, ) timings["denoise_s"] = time.perf_counter() - started timings["dit_forward_s"] = float(getattr(self, "last_dit_forward_s", 0.0)) @@ -852,6 +880,8 @@ class MiniMaxH3MLXPipeline: width=video_geometry["width"], num_frames=video_geometry["num_frames"], tiled=tiled_video_decode, + vae_tile_height=vae_tile_height, + vae_tile_width=vae_tile_width, ) if spatial_plan is not None: frames = _center_crop_frames(frames, spatial_plan.stage1_height, spatial_plan.stage1_width) diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py index efdbcc3ed..0f360f24b 100644 --- a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py +++ b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py @@ -54,7 +54,29 @@ SUPPORTED_FORMATS = ("int8", "int6", "int4") DEFAULT_FORMATS = " ".join(SUPPORTED_FORMATS) -def _adaln_cache_timesteps() -> np.ndarray: +def _adaln_cache_timesteps(model_root: str | Path | None = None) -> np.ndarray: + import json + if model_root is not None: + p = Path(model_root) + contract_path = p / "fastvideo_inference.json" + if not contract_path.exists() and p.parent: + contract_path = p.parent / "fastvideo_inference.json" + if contract_path.exists(): + try: + contract = json.loads(contract_path.read_text(encoding="utf-8")) + dmd_steps = contract.get("dmd_denoising_steps") + v_shift = float(contract.get("video_scheduler_shift", 10.0)) + a_shift = float(contract.get("audio_scheduler_shift", 3.0)) + if dmd_steps: + base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64) + v_sigmas = v_shift * base / (1.0 + (v_shift - 1.0) * base) + a_sigmas = a_shift * base / (1.0 + (a_shift - 1.0) * base) + v_t = (1.0 - v_sigmas[:-1]).astype(np.float32) + a_t = (1.0 - a_sigmas[:-1]).astype(np.float32) + logger.info("FastH3 checkpoint schedule contract found at %s (%d DMD forwards, video/audio shift=%g/%g)", contract_path, len(dmd_steps), v_shift, a_shift) + return np.unique(np.concatenate([v_t, a_t, [1.0]])).astype(np.float32) + except Exception as exc: + logger.warning("Could not parse %s: %s", contract_path, exc) video = 1.0 - minimax_h3_sigmas(MINIMAX_H3_VIDEO_SHIFT, 4)[:-1] audio = 1.0 - minimax_h3_sigmas(MINIMAX_H3_AUDIO_SHIFT, 4)[:-1] return np.unique(np.concatenate([video, audio, [1.0]])).astype(np.float32) @@ -84,7 +106,7 @@ def main() -> None: raise ValueError(f"Unsupported H3 MLX formats: {unsupported}. Choose from {SUPPORTED_FORMATS}.") out_base = Path(args.out) out_base.mkdir(parents=True, exist_ok=True) - cache_timesteps = _adaln_cache_timesteps() + cache_timesteps = _adaln_cache_timesteps(args.model_root) for fmt in formats: spec = MLXQuantizationSpec.from_name(fmt) From 430e52154e76b902c3cc17a16b3edc1fad790012 Mon Sep 17 00:00:00 2001 From: primorLee Date: Fri, 18 Sep 2026 20:00:05 +0800 Subject: [PATCH 016/151] [bugfix]: Restore exact Qwen3-VL vision interpolation (#1737) Co-authored-by: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Co-authored-by: SolitaryThinker --- .../models/encoders/minimax_h3_qwen3_vl.py | 126 ++++++++------ .../test_minimax_h3_qwen3_vl_vision.py | 155 ++++++++++++++++++ pyproject.toml | 6 +- ...inimax_h3_qwen3_vl_interpolation_memory.py | 137 ++++++++++++++++ tests/local_tests/glm_image/README.md | 2 +- tests/local_tests/minimax_h3/PORT_STATUS.md | 10 +- tests/local_tests/minimax_h3/README.md | 39 +++++ 7 files changed, 422 insertions(+), 53 deletions(-) create mode 100644 fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_vision.py create mode 100644 tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index 9f4b5e8ea..c348140d1 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -20,11 +20,82 @@ from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3Seria from fastvideo.models.loader.weight_utils import default_weight_loader +_VISION_INTERPOLATION_WORKSPACE_BYTES = 64 * 1024 * 1024 + + def _rotate_half(tensor: torch.Tensor) -> torch.Tensor: first, second = tensor.chunk(2, dim=-1) return torch.cat((-second, first), dim=-1) +def _vision_interpolation_indices_and_weights( + grid_thw: torch.Tensor, + num_grid_per_side: int, + spatial_merge_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build Qwen3-VL's bilinear position-table gather in merge-block order.""" + side = num_grid_per_side + merge = spatial_merge_size + + counts = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] + heights = torch.repeat_interleave(grid_thw[:, 1], counts) + widths = torch.repeat_interleave(grid_thw[:, 2], counts) + starts = torch.repeat_interleave(F.pad(counts.cumsum(0)[:-1], (1, 0)), counts) + within = (torch.arange(counts.sum(), device=grid_thw.device) - starts) % (heights * widths) + + blocks_wide = widths // merge + within_col = within % merge + within_row = (within // merge) % merge + block_col = (within // (merge * merge)) % blocks_wide + block_row = within // (merge * merge * blocks_wide) + rows = block_row * merge + within_row + cols = block_col * merge + within_col + + def axis_taps_and_weights( + indices: torch.Tensor, + sizes: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + positions = indices.to(torch.float32) * (side - 1) / torch.clamp(sizes - 1, min=1) + floors = torch.floor(positions) + offsets = torch.arange(2, device=grid_thw.device) + taps = (floors.long()[:, None] + offsets).clamp(0, side - 1) + distances = (positions[:, None] - floors[:, None] - offsets).abs() + return taps, (1 - distances).clamp(min=0) + + row_taps, row_weights = axis_taps_and_weights(rows, heights) + col_taps, col_weights = axis_taps_and_weights(cols, widths) + indices = (row_taps[:, :, None] * side + col_taps[:, None, :]).reshape(-1, 4) + weights = (row_weights[:, :, None] * col_weights[:, None, :]).reshape(-1, 4) + return indices, weights + + +def _interpolate_vision_position_embeddings( + position_embedding: torch.Tensor, + grid_thw: torch.Tensor, + num_grid_per_side: int, + spatial_merge_size: int, +) -> torch.Tensor: + """Keep interpolation weights and accumulation in float32 until the vision residual add.""" + grid_thw = grid_thw.to(position_embedding.device) + indices, weights = _vision_interpolation_indices_and_weights(grid_thw, num_grid_per_side, spatial_merge_size) + output_dtype = torch.promote_types(position_embedding.dtype, weights.dtype) + output = torch.empty( + (indices.shape[0], position_embedding.shape[1]), + dtype=output_dtype, + device=position_embedding.device, + ) + + # Packed reference videos can otherwise promote a multi-GiB four-tap gather at once. + bytes_per_patch = indices.shape[1] * position_embedding.shape[1] * ( + position_embedding.element_size() + torch.empty((), dtype=output_dtype).element_size()) + chunk_size = max(1, _VISION_INTERPOLATION_WORKSPACE_BYTES // bytes_per_patch) + for start in range(0, indices.shape[0], chunk_size): + stop = min(start + chunk_size, indices.shape[0]) + output[start:stop] = (F.embedding(indices[start:stop], position_embedding) * + weights[start:stop, :, None]).sum(1) + return output + + class MiniMaxH3Qwen3VLTextRotaryEmbedding(nn.Module): """Shared Qwen3-VL interleaved temporal/height/width rotary embedding.""" @@ -435,57 +506,16 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): return frequency_table[position_ids].flatten(1) def _interpolate_position_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor: - index_lists: list[list[int]] = [[] for _ in range(4)] - weight_lists: list[list[float]] = [[] for _ in range(4)] - merge = self.spatial_merge_size - patch_counts: list[int] = [] - grids: list[tuple[int, int, int]] = [] - for frames_tensor, height_tensor, width_tensor in grid_thw: - frames, height, width = int(frames_tensor), int(height_tensor), int(width_tensor) - grids.append((frames, height, width)) - patch_counts.append(height * width) - height_positions = torch.linspace(0, self.num_grid_per_side - 1, height) - width_positions = torch.linspace(0, self.num_grid_per_side - 1, width) - height_floor = height_positions.int() - width_floor = width_positions.int() - height_ceil = (height_floor + 1).clip(max=self.num_grid_per_side - 1) - width_ceil = (width_floor + 1).clip(max=self.num_grid_per_side - 1) - delta_height = height_positions - height_floor - delta_width = width_positions - width_floor - base_height = height_floor * self.num_grid_per_side - base_height_ceil = height_ceil * self.num_grid_per_side - indices = ( - (base_height[:, None] + width_floor[None]).flatten(), - (base_height[:, None] + width_ceil[None]).flatten(), - (base_height_ceil[:, None] + width_floor[None]).flatten(), - (base_height_ceil[:, None] + width_ceil[None]).flatten(), - ) - weights = ( - ((1 - delta_height)[:, None] * (1 - delta_width)[None]).flatten(), - ((1 - delta_height)[:, None] * delta_width[None]).flatten(), - (delta_height[:, None] * (1 - delta_width)[None]).flatten(), - (delta_height[:, None] * delta_width[None]).flatten(), - ) - for index in range(4): - index_lists[index].extend(indices[index].tolist()) - weight_lists[index].extend(weights[index].tolist()) - - index_tensor = torch.tensor(index_lists, dtype=torch.long, device=self.pos_embed.weight.device) - weight_tensor = torch.tensor(weight_lists, - dtype=self.pos_embed.weight.dtype, - device=self.pos_embed.weight.device) - embeddings = self.pos_embed(index_tensor) * weight_tensor[:, :, None] - embeddings = (embeddings[0] + embeddings[1] + embeddings[2] + embeddings[3]).split(patch_counts) - permuted = [] - for embedding, (frames, height, width) in zip(embeddings, grids, strict=True): - embedding = embedding.repeat(frames, 1) - embedding = embedding.view(frames, height // merge, merge, width // merge, merge, -1) - permuted.append(embedding.permute(0, 1, 3, 2, 4, 5).flatten(0, 4)) - return torch.cat(permuted) + return _interpolate_vision_position_embeddings( + self.pos_embed.weight, + grid_thw, + self.num_grid_per_side, + self.spatial_merge_size, + ) def forward(self, pixels: torch.Tensor, grid_thw: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: hidden_states = self.patch_embed(pixels) - hidden_states = hidden_states + self._interpolate_position_embeddings(grid_thw) + hidden_states = hidden_states + self._interpolate_position_embeddings(grid_thw).to(hidden_states.dtype) rotary = self._rotary_positions(grid_thw).reshape(hidden_states.shape[0], -1) embedding = torch.cat((rotary, rotary), dim=-1) position_embeddings = (embedding.cos(), embedding.sin()) diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_vision.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_vision.py new file mode 100644 index 000000000..0bf485184 --- /dev/null +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_vision.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU regressions for MiniMax-H3's Qwen3-VL vision position embeddings.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from torch.testing import assert_close + +from fastvideo.models.encoders import minimax_h3_qwen3_vl +from fastvideo.models.encoders.minimax_h3_qwen3_vl import ( + _VISION_INTERPOLATION_WORKSPACE_BYTES, + _interpolate_vision_position_embeddings, +) + + +_GRID_CASES = ( + torch.tensor([[1, 4, 6]], dtype=torch.long), + torch.tensor([[2, 4, 6]], dtype=torch.long), + torch.tensor([[1, 4, 6], [2, 8, 10]], dtype=torch.long), + torch.tensor([[1, 48, 84]], dtype=torch.long), + torch.tensor([[1, 128, 224]], dtype=torch.long), + torch.tensor([[15, 42, 74]], dtype=torch.long), + torch.tensor([[1, 128, 224], [15, 42, 74]], dtype=torch.long), +) + + +def _torch_interpolation_reference( + position_embedding: torch.Tensor, + grid_thw: torch.Tensor, + side: int, + merge: int, +) -> torch.Tensor: + """Use PyTorch's bilinear kernel as an independent, always-available oracle.""" + channels = position_embedding.shape[1] + source = position_embedding.float().T.reshape(1, channels, side, side) + outputs = [] + for frames, height, width in grid_thw.tolist(): + raster = F.interpolate(source, size=(height, width), mode="bilinear", align_corners=True) + raster = raster.squeeze(0).permute(1, 2, 0) + merged = raster.view(height // merge, merge, width // merge, merge, channels) + merged = merged.permute(0, 2, 1, 3, 4).flatten(0, 3) + outputs.append(merged.repeat(frames, 1)) + return torch.cat(outputs) + + +def _transformers_5_15_contract_reference( + position_embedding: torch.Tensor, + grid_thw: torch.Tensor, +) -> torch.Tensor: + """Evaluate the Transformers 5.15 contract without requiring its newer public helper.""" + side = 48 + merge = 2 + channels = position_embedding.shape[1] + outputs = [] + for frames, height, width in grid_thw.tolist(): + row_positions = (torch.arange(height, device=position_embedding.device, dtype=torch.float32) * + (side - 1) / max(height - 1, 1)) + column_positions = (torch.arange(width, device=position_embedding.device, dtype=torch.float32) * + (side - 1) / max(width - 1, 1)) + row_floors = row_positions.floor() + column_floors = column_positions.floor() + offsets = torch.arange(2, device=position_embedding.device) + row_taps = (row_floors.long()[:, None] + offsets).clamp(0, side - 1) + column_taps = (column_floors.long()[:, None] + offsets).clamp(0, side - 1) + row_weights = (1 - (row_positions[:, None] - row_floors[:, None] - offsets).abs()).clamp(min=0) + column_distances = (column_positions[:, None] - column_floors[:, None] - offsets).abs() + column_weights = (1 - column_distances).clamp(min=0) + + indices = (row_taps[:, None, :, None] * side + column_taps[None, :, None, :]).reshape(height, width, 4) + weights = (row_weights[:, None, :, None] * column_weights[None, :, None, :]).reshape(height, width, 4) + raster = (position_embedding[indices] * weights[:, :, :, None]).sum(2) + merged = raster.view(height // merge, merge, width // merge, merge, channels) + merged = merged.permute(0, 2, 1, 3, 4).flatten(0, 3) + outputs.append(merged.repeat(frames, 1)) + return torch.cat(outputs) + + +def _installed_transformers_interpolation_reference( + position_embedding: torch.Tensor, + grid_thw: torch.Tensor, +) -> torch.Tensor: + from transformers import vision_utils + + get_reference = getattr(vision_utils, "get_vision_interpolation_indices_and_weights", None) + if get_reference is None: + pytest.fail("transformers>=5.15 must provide get_vision_interpolation_indices_and_weights", pytrace=False) + indices, weights = get_reference(grid_thw, + num_grid_per_side=48, + mode="bilinear", + align_corners=True, + spatial_merge_size=2) + return (position_embedding[indices] * weights[:, :, None]).sum(1) + + +@pytest.mark.parametrize("grid_thw", _GRID_CASES) +def test_position_interpolation_matches_transformers_5_15_contract(grid_thw: torch.Tensor) -> None: + """Pin visual-token ordering and float32 accumulation for every supported dependency version.""" + torch.manual_seed(1733) + position_embedding = torch.randn(48 * 48, 32, dtype=torch.bfloat16) + expected = _transformers_5_15_contract_reference(position_embedding, grid_thw) + + actual = _interpolate_vision_position_embeddings(position_embedding, grid_thw, 48, 2) + + assert actual.dtype == torch.float32 + assert_close(actual, expected, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize("grid_thw", _GRID_CASES) +def test_transformers_5_15_contract_matches_installed_helper(grid_thw: torch.Tensor) -> None: + """Cross-check the self-contained contract against the installed public helper.""" + torch.manual_seed(1733) + position_embedding = torch.randn(48 * 48, 32, dtype=torch.bfloat16) + + expected = _installed_transformers_interpolation_reference(position_embedding, grid_thw) + actual = _transformers_5_15_contract_reference(position_embedding, grid_thw) + + assert_close(actual, expected, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize("grid_thw", _GRID_CASES) +def test_position_interpolation_uses_float32_bilinear_accumulation(grid_thw: torch.Tensor) -> None: + """Prevent bf16 interpolation weights from reintroducing visual-token drift.""" + torch.manual_seed(1733) + position_embedding = torch.randn(48 * 48, 32, dtype=torch.bfloat16) + + expected = _torch_interpolation_reference(position_embedding, grid_thw, 48, 2) + actual = _interpolate_vision_position_embeddings(position_embedding, grid_thw, 48, 2) + + assert actual.dtype == torch.float32 + assert_close(actual, expected, atol=5e-5, rtol=0.0) + + +def test_position_interpolation_bounds_four_tap_workspace(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep production video interpolation from materializing every wide gather at once.""" + torch.manual_seed(1733) + position_embedding = torch.randn(48 * 48, 128, dtype=torch.bfloat16) + grid_thw = torch.tensor([[15, 42, 74]], dtype=torch.long) + temporary_sizes = [] + original_embedding = minimax_h3_qwen3_vl.F.embedding + + def recording_embedding(indices: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + output_dtype = torch.promote_types(weight.dtype, torch.float32) + temporary_sizes.append(indices.numel() * weight.shape[1] * + (weight.element_size() + torch.empty((), dtype=output_dtype).element_size())) + return original_embedding(indices, weight) + + monkeypatch.setattr(minimax_h3_qwen3_vl.F, "embedding", recording_embedding) + actual = _interpolate_vision_position_embeddings(position_embedding, grid_thw, 48, 2) + expected = _transformers_5_15_contract_reference(position_embedding, grid_thw) + + assert len(temporary_sizes) > 1 + assert max(temporary_sizes) <= _VISION_INTERPOLATION_WORKSPACE_BYTES + assert_close(actual, expected, atol=0.0, rtol=0.0) diff --git a/pyproject.toml b/pyproject.toml index da3eadaea..d6ed836bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,10 @@ dependencies = [ # Machine Learning & Transformers # GLM-Image's AR encoder (GlmImageForConditionalGeneration) first ships in - # transformers 5.0.0; floor bumped from >=4.57.3 to >=5.0.0 (stable, not rc). - "transformers>=5.0.0", + # transformers 5.0.0. Floor raised from >=5.0.0 to >=5.15.0 for the public + # vision interpolation helper (get_vision_interpolation_indices_and_weights) + # that the MiniMax-H3 Qwen3-VL encoder cross-checks against. + "transformers>=5.15.0", # <0.23: tokenizers 0.23 renamed RobertaProcessing's binding args, so # transformers' CLIP-style tokenizer loading dies with # "RobertaProcessing.__new__() got an unexpected keyword argument 'cls'". diff --git a/tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py b/tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py new file mode 100644 index 000000000..2acee9e7f --- /dev/null +++ b/tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Measure the Qwen3-VL position-interpolation CUDA working set.""" + +from __future__ import annotations + +import argparse +import importlib +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import torch + + +def _parse_grid(value: str) -> tuple[int, int, int]: + try: + grid = tuple(int(part) for part in value.split(",")) + except ValueError as error: + raise argparse.ArgumentTypeError("grid must contain integer T,H,W values") from error + if len(grid) != 3 or any(dimension <= 0 for dimension in grid): + raise argparse.ArgumentTypeError("grid must contain three positive T,H,W values") + return grid + + +def _revision(source_root: Path) -> str: + result = subprocess.run( + ["git", "-C", str(source_root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def _is_dirty(source_root: Path) -> bool: + result = subprocess.run( + ["git", "-C", str(source_root), "status", "--short"], + check=True, + capture_output=True, + text=True, + ) + return bool(result.stdout.strip()) + + +def _run_interpolation( + module: Any, + position_embedding: torch.Tensor, + grid_thw: torch.Tensor, + side: int, + merge: int, +) -> tuple[torch.Tensor, str]: + helper = getattr(module, "_interpolate_vision_position_embeddings", None) + if helper is not None: + return helper(position_embedding, grid_thw, side, merge), "standalone_helper" + + class VisionProxy: + num_grid_per_side = side + spatial_merge_size = merge + pos_embed = torch.nn.Embedding.from_pretrained(position_embedding, freeze=True) + + vision_class = module.MiniMaxH3Qwen3VLVisionModel + return vision_class._interpolate_position_embeddings(VisionProxy(), grid_thw), "legacy_model_method" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-root", type=Path, default=Path.cwd()) + parser.add_argument("--grid", action="append", type=_parse_grid, default=None, metavar="T,H,W") + parser.add_argument("--hidden-size", type=int, default=1152) + parser.add_argument("--num-grid-per-side", type=int, default=48) + parser.add_argument("--spatial-merge-size", type=int, default=2) + parser.add_argument("--seed", type=int, default=1737) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the interpolation memory benchmark") + if args.hidden_size <= 0 or args.num_grid_per_side <= 0 or args.spatial_merge_size <= 0: + raise ValueError("hidden size, grid side, and merge size must be positive") + + source_root = args.source_root.resolve() + sys.path.insert(0, str(source_root)) + module = importlib.import_module("fastvideo.models.encoders.minimax_h3_qwen3_vl") + + device = torch.device("cuda") + grids = args.grid or [(15, 42, 74)] + grid_thw = torch.tensor(grids, dtype=torch.long, device=device) + torch.manual_seed(args.seed) + position_embedding = torch.randn( + args.num_grid_per_side**2, + args.hidden_size, + dtype=torch.bfloat16, + device=device, + ) + torch.cuda.synchronize() + baseline_allocated = torch.cuda.memory_allocated() + baseline_reserved = torch.cuda.memory_reserved() + torch.cuda.reset_peak_memory_stats() + + with torch.no_grad(): + output, implementation = _run_interpolation( + module, + position_embedding, + grid_thw, + args.num_grid_per_side, + args.spatial_merge_size, + ) + torch.cuda.synchronize() + + properties = torch.cuda.get_device_properties(device) + record = { + "source_root": str(source_root), + "source_revision": _revision(source_root), + "source_is_dirty": _is_dirty(source_root), + "implementation": implementation, + "torch_version": torch.__version__, + "device": properties.name, + "compute_capability": f"{properties.major}.{properties.minor}", + "grid_thw": grids, + "position_embedding_shape": list(position_embedding.shape), + "position_embedding_dtype": str(position_embedding.dtype), + "output_shape": list(output.shape), + "output_dtype": str(output.dtype), + "output_sum_fp32": output.float().sum().item(), + "baseline_allocated_bytes": baseline_allocated, + "baseline_reserved_bytes": baseline_reserved, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + "incremental_peak_allocated_bytes": torch.cuda.max_memory_allocated() - baseline_allocated, + "incremental_peak_reserved_bytes": torch.cuda.max_memory_reserved() - baseline_reserved, + } + print(json.dumps(record, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/glm_image/README.md b/tests/local_tests/glm_image/README.md index baf4201c1..32037b77e 100644 --- a/tests/local_tests/glm_image/README.md +++ b/tests/local_tests/glm_image/README.md @@ -11,7 +11,7 @@ message when a prerequisite is missing. | | | |---|---| | Weights | `zai-org/GLM-Image` → `official_weights/glm_image` (~30 GB: AR encoder 4 shards, transformer 3 shards, vae + text_encoder) | -| Runtime deps | `transformers>=5.0.0` (first release with the AR encoder `GlmImageForConditionalGeneration`) and `diffusers>=0.38.0` — both committed in `pyproject.toml` | +| Runtime deps | `transformers>=5.15.0` (repo floor; the AR encoder `GlmImageForConditionalGeneration` first ships in 5.0.0) and `diffusers>=0.38.0` — both committed in `pyproject.toml` | | Test-only deps | the diffusers `GlmImageTransformer2DModel` / `GlmImagePipeline` reference classes need `diffusers>=0.37.0.dev0` (used **only** by the parity tests) | Download weights (from the repo root; cache via `HF_HOME` if `/` is small): diff --git a/tests/local_tests/minimax_h3/PORT_STATUS.md b/tests/local_tests/minimax_h3/PORT_STATUS.md index 9c9ed2cee..642ef71f2 100644 --- a/tests/local_tests/minimax_h3/PORT_STATUS.md +++ b/tests/local_tests/minimax_h3/PORT_STATUS.md @@ -11,7 +11,7 @@ | Scope | Evidence | State | |---|---|---| -| Qwen3-VL encoder | exact text/image/video hidden states through the production loader | complete | +| Qwen3-VL encoder | exact text/image/video hidden states through the production loader; bounded exact vision interpolation | complete | | FL2VA and Ref2VA DiTs | exact video/audio heads for both model partitions | complete | | Video VAE | exact encode, normalization, and decode through the production loader | complete | | Video VAE streaming | exact chunked encode/decode and output-rank-only distributed decode | complete | @@ -25,7 +25,11 @@ ## Current validation -T2VA, FL2VA, and Ref2VA match the official video/audio latents exactly. +T2VA, FL2VA, and Ref2VA match the official video/audio latents exactly. On 2026-08-26, the Qwen3-VL production loader +also matched Transformers 5.15.1 exactly for text, image, and video layer-50 hidden states on GB10. Its production +`[15, 42, 74] x 1152` vision interpolation matched the official helper bit-exactly while reducing incremental CUDA +allocation from `1,291,986,432` to `284,866,048` bytes relative to the unbounded float32 implementation. A packed +production image/video grid reduced the same metric from `2,086,086,144` to `418,362,880` bytes. ## Decisions @@ -36,6 +40,8 @@ T2VA, FL2VA, and Ref2VA match the official video/audio latents exactly. - Keep `last_image`, `references`, and `audio_latents` on the typed request path. - Treat the published component folders as the loading boundary. - Keep reference videos on CPU between VAE clips and decode final pixels only on the executor's output rank. +- Preserve Qwen3-VL's float32 vision interpolation through accumulation, bound its four-tap temporary workspace, and + cast the completed positions only at the vision residual-add boundary. ## Evidence boundary diff --git a/tests/local_tests/minimax_h3/README.md b/tests/local_tests/minimax_h3/README.md index ee596739a..ffdf94155 100644 --- a/tests/local_tests/minimax_h3/README.md +++ b/tests/local_tests/minimax_h3/README.md @@ -16,10 +16,16 @@ that skip is not parity evidence. ```bash pytest \ + fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_vision.py \ fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py \ fastvideo/tests/stages/test_minimax_h3_vae_streaming.py -q ``` +The Qwen3-VL vision test covers production image/video grids, packed grids, exact float32 accumulation against a +self-contained Transformers 5.15 contract reference, and the bounded four-tap workspace. A separate test checks that +reference against Transformers' public `get_vision_interpolation_indices_and_weights` helper, which +`transformers>=5.15` provides. + ## Registry smoke ```bash @@ -65,6 +71,39 @@ partitions and the video VAE; audio decode has maximum absolute drift `2.4e-7`. forward's selected layer-50 hidden state bit-exactly against the same state from the official full stack across text, image, and video inputs. +On 2026-08-26, the encoder gate passed all three cases on GB10 with PyTorch `2.12.0+cu130` and Transformers `5.15.1`: +every selected layer-50 state had `max_abs=0` and `mean_abs=0`. + +## Qwen3-VL interpolation memory benchmark + +Use the same benchmark-script checkout for both source trees. Each invocation runs in a fresh process and reports +absolute and incremental CUDA allocated/reserved peaks, the source revision and dirty state, output metadata, and a +deterministic FP32 output sum. + +```bash +python tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py \ + --source-root /path/to/candidate + +python tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py \ + --source-root /path/to/baseline + +# Repeat --grid to measure a packed request. +python tests/local_tests/encoders/benchmark_minimax_h3_qwen3_vl_interpolation_memory.py \ + --source-root /path/to/candidate --grid 1,128,224 --grid 15,42,74 +``` + +The default `[15, 42, 74]` grid and hidden width `1152` reproduce a production video interpolation. On GB10 with +PyTorch `2.12.0+cu130`, an intermediate unbounded-float32 implementation of this change (not reachable from the PR +branch) used `1,291,986,432` incremental allocated bytes; the bounded implementation used `284,866,048` bytes with +the same output shape, dtype, and FP32 sum. A direct comparison against the Transformers helper was bit-exact for all +`46,620 x 1,152` output elements. Fresh CPU processes reduced maximum resident set size from `2,988,420` to +`1,244,088` KiB for the same tensor, eliminating `1,744,332` KiB (`1.66` GiB) of peak retention. The mixed +`[1, 128, 224] + [15, 42, 74]` packed case reduced incremental CUDA allocation from `2,086,086,144` to +`418,362,880` bytes with the same FP32 sum. + +The benchmark's legacy fallback path (a source tree without `_interpolate_vision_position_embeddings`) measures the +pre-PR bf16 implementation, so its FP32 sum differs from the bounded float32 implementation. + The video VAE test verifies the reference checkout at commit `abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc` and compares the production CPU `uint8` `encode_pixels()` path against the official posterior element by element. From c73dd0ab55c1baace8ae16dfe88069ab7009ae17 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Thu, 17 Sep 2026 20:00:30 +0000 Subject: [PATCH 017/151] [perf] H3 decode: offload the frozen VAEs without the device-to-host copy The decode stages move the 10.4 GB fp32 video VAE (and the 0.6 GB audio VAE) with module.to(device) / module.to("cpu") on every request. The weights are frozen -- the whole stage runs under no_grad -- so the copy back to the host returns bytes the host already has. Keep one host copy per tensor, made the first time the module is loaded, and on unload point .data back at it instead of copying. One H2D per request, no D2H at all; the module still lives on the host between requests. Where the old path allocated a fresh host tensor per cycle and freed the previous one, this reuses one buffer, so peak host memory can only go down. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01ATw7g6cxq9N2LtnMratsru --- fastvideo/models/pinned_offload.py | 73 +++++++++++++++++++ .../minimax_h3/stages/minimax_h3_decoding.py | 9 ++- 2 files changed, 78 insertions(+), 4 deletions(-) create mode 100644 fastvideo/models/pinned_offload.py diff --git a/fastvideo/models/pinned_offload.py b/fastvideo/models/pinned_offload.py new file mode 100644 index 000000000..480c3cf09 --- /dev/null +++ b/fastvideo/models/pinned_offload.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU offload for frozen inference modules, without the device-to-host copy. + +``module.to("cpu")`` allocates fresh host storage and copies every parameter +and buffer back over PCIe. For a module that is only ever run under +``torch.no_grad()`` the weights on the device are byte-identical to the ones +that were copied in, so that copy is pure waste: ``unload`` points each +tensor's ``.data`` back at the host copy it came from and lets the device +storage go. + +``load`` keeps one host copy per tensor, made the first time the module is +loaded, and copies it to the device. Net: one host-to-device copy per request +and no device-to-host copy at all, with the module living on the host between +requests exactly as before. Numerics are untouched -- the device tensors are +byte copies of the same weights -- and peak host memory can only go down, +because one buffer is reused instead of a new one being allocated per cycle. +""" +from __future__ import annotations + +import torch +from torch import nn + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +_HOST_ATTR = "_frozen_offload_host" + + +def _tensors(module: nn.Module): + yield from module.named_parameters(recurse=True) + for name, buf in module.named_buffers(recurse=True): + yield "buffer:" + name, buf + + +def _host_copies(module: nn.Module) -> dict[str, torch.Tensor]: + host = getattr(module, _HOST_ATTR, None) + if host is not None: + return host + host, total = {}, 0 + for name, t in _tensors(module): + src = t.data if t.data.device.type == "cpu" else t.data.to("cpu") + host[name] = src + total += src.numel() * src.element_size() + setattr(module, _HOST_ATTR, host) + logger.info("frozen offload: holding %.2f GB of %s on the host", total / 1e9, type(module).__name__) + return host + + +def load(module: nn.Module, device: torch.device) -> nn.Module: + """Put a frozen ``module`` on ``device``, copying from its host copies.""" + if device.type != "cuda": + return module.to(device) + host = _host_copies(module) + for name, t in _tensors(module): + if t.data.device == device: + continue + src = host[name] + dst = torch.empty_like(src, device=device) + dst.copy_(src) + t.data = dst + return module + + +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) + if host is None: + return module.to("cpu") + for name, t in _tensors(module): + if t.data.device.type != "cpu": + t.data = host[name] + return module 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 1eae7005e..411e92a66 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -10,6 +10,7 @@ import torch from fastvideo.distributed import get_local_torch_device, get_sp_group, get_world_group, model_parallel_is_initialized from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models import pinned_offload from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE from fastvideo.models.vaes.minimax_h3_parallel import DEFAULT_DECODE_GATHER_STRATEGY, decode_to_pixels_parallel from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3 @@ -123,7 +124,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): if self.vae is None: raise RuntimeError("MiniMax-H3 full VAE decode requires a loaded video VAE.") - self.vae.to(device) + pinned_offload.load(self.vae, device) try: latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": @@ -157,7 +158,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): return batch finally: if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") + pinned_offload.unload(self.vae) class MiniMaxH3AudioDecodingStage(PipelineStage): @@ -200,7 +201,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage): layout.num_audio_latents, ) device = get_local_torch_device() - self.audio_vae.to(device) + pinned_offload.load(self.audio_vae, device) try: latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": @@ -222,7 +223,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage): return batch finally: if fastvideo_args.vae_cpu_offload: - self.audio_vae.to("cpu") + pinned_offload.unload(self.audio_vae) @staticmethod def _clear_runtime(batch: ForwardBatch) -> None: From 3e26db40b0d66c1ad03c30bcf8e98fb8efba3f4e Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Thu, 17 Sep 2026 20:01:01 +0000 Subject: [PATCH 018/151] [perf] H3 decode: pin the frozen VAEs' host copies With the device-to-host copy gone, the remaining host-to-device copy is the stage's cost, and from pageable memory it is staged through a bounce buffer at roughly 2.6 GB/s instead of full PCIe speed. The host copies are made once, so pinning them costs nothing per request. It does cost the module's size in non-pageable host memory (10.4 GB for the H3 video VAE, 0.6 GB for the audio VAE, per rank) for as long as the pipeline is alive. It follows --pin-cpu-memory, which is on by default; --no-pin-cpu- memory keeps the copies pageable and only loses the bandwidth. Note that this flag did not already reach the VAE: its other use is FSDP's own CPUOffloadPolicy and the VAE does not go through FSDP. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01ATw7g6cxq9N2LtnMratsru --- fastvideo/models/pinned_offload.py | 20 ++++++++++++++----- .../minimax_h3/stages/minimax_h3_decoding.py | 4 ++-- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/fastvideo/models/pinned_offload.py b/fastvideo/models/pinned_offload.py index 480c3cf09..e48c3cb1a 100644 --- a/fastvideo/models/pinned_offload.py +++ b/fastvideo/models/pinned_offload.py @@ -14,6 +14,12 @@ and no device-to-host copy at all, with the module living on the host between requests exactly as before. Numerics are untouched -- the device tensors are byte copies of the same weights -- and peak host memory can only go down, because one buffer is reused instead of a new one being allocated per cycle. + +The host copies are pinned when ``pin`` is set (``--pin-cpu-memory``, on by +default), which is what lets the remaining host-to-device copy run at full +PCIe speed instead of being staged through a bounce buffer. That trades the +module's size in non-pageable host memory for the bandwidth; ``pin=False`` +keeps the copies pageable and everything else the same. """ from __future__ import annotations @@ -21,6 +27,7 @@ import torch from torch import nn from fastvideo.logger import init_logger +from fastvideo.utils import is_pin_memory_available logger = init_logger(__name__) @@ -33,31 +40,34 @@ def _tensors(module: nn.Module): yield "buffer:" + name, buf -def _host_copies(module: nn.Module) -> dict[str, torch.Tensor]: +def _host_copies(module: nn.Module, pin: bool) -> dict[str, torch.Tensor]: host = getattr(module, _HOST_ATTR, None) if host is not None: return host host, total = {}, 0 for name, t in _tensors(module): src = t.data if t.data.device.type == "cpu" else t.data.to("cpu") + if pin and not src.is_pinned(): + src = src.pin_memory() host[name] = src total += src.numel() * src.element_size() setattr(module, _HOST_ATTR, host) - logger.info("frozen offload: holding %.2f GB of %s on the host", total / 1e9, type(module).__name__) + logger.info("frozen offload: holding %.2f GB of %s on the host (%s)", total / 1e9, type(module).__name__, + "pinned" if pin else "pageable") return host -def load(module: nn.Module, device: torch.device) -> nn.Module: +def load(module: nn.Module, device: torch.device, pin: bool = True) -> nn.Module: """Put a frozen ``module`` on ``device``, copying from its host copies.""" if device.type != "cuda": return module.to(device) - host = _host_copies(module) + host = _host_copies(module, pin and is_pin_memory_available()) for name, t in _tensors(module): if t.data.device == device: continue src = host[name] dst = torch.empty_like(src, device=device) - dst.copy_(src) + dst.copy_(src, non_blocking=src.is_pinned()) t.data = dst return module 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 411e92a66..a12cca50d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -124,7 +124,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): if self.vae is None: raise RuntimeError("MiniMax-H3 full VAE decode requires a loaded video VAE.") - pinned_offload.load(self.vae, device) + pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory) try: latents = self.vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": @@ -201,7 +201,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage): layout.num_audio_latents, ) device = get_local_torch_device() - pinned_offload.load(self.audio_vae, device) + pinned_offload.load(self.audio_vae, device, pin=fastvideo_args.pin_cpu_memory) try: latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32)) if fastvideo_args.output_type == "latent": From 10bdf5e076b8b2a0152a365bd0f0ebb4c7230a04 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Thu, 17 Sep 2026 22:34:26 +0000 Subject: [PATCH 019/151] [bugfix] frozen offload: no host copy for an already resident module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `load` took the host copies before checking which tensors were missing from the device, so with `vae_cpu_offload=False` — where the loader builds the VAEs on CUDA and `unload` is never called — it pulled about 11 GB per rank back over PCIe once and kept a pinned mirror of it for the life of the process, for a module that never leaves the device. It now asks what is absent first and returns early when nothing is. The offloaded path is unchanged: weights still round-trip and come back bit identical. Co-Authored-By: Claude Opus 5 (1M context) --- fastvideo/models/pinned_offload.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/fastvideo/models/pinned_offload.py b/fastvideo/models/pinned_offload.py index e48c3cb1a..08c45802a 100644 --- a/fastvideo/models/pinned_offload.py +++ b/fastvideo/models/pinned_offload.py @@ -61,10 +61,15 @@ def load(module: nn.Module, device: torch.device, pin: bool = True) -> nn.Module """Put a frozen ``module`` on ``device``, copying from its host copies.""" if device.type != "cuda": return module.to(device) + missing = [(name, t) for name, t in _tensors(module) if t.data.device != device] + if not missing: + # Already resident, which is what ``vae_cpu_offload=False`` leaves. Taking + # host copies here would pull the whole module back over PCIe once and keep + # a pinned mirror of it for the life of the process, for a module that never + # leaves the device. + return module host = _host_copies(module, pin and is_pin_memory_available()) - for name, t in _tensors(module): - if t.data.device == device: - continue + for name, t in missing: src = host[name] dst = torch.empty_like(src, device=device) dst.copy_(src, non_blocking=src.is_pinned()) From f47ad3f5b7a54c1a976e1821dd9ad082f7f5b7ef Mon Sep 17 00:00:00 2001 From: Keith Date: Tue, 22 Sep 2026 02:59:13 +0800 Subject: [PATCH 020/151] [bugfix] fastvideo-kernel: install the Python + Triton package when the HIP toolchain cannot be configured (#1871) --- fastvideo-kernel/CMakeLists.txt | 48 ++++++++++++++++-- fastvideo-kernel/README.md | 14 ++++++ fastvideo-kernel/build.sh | 88 ++++++++++++++++++++++----------- 3 files changed, 115 insertions(+), 35 deletions(-) diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 405d7b751..1fc50a3f1 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -16,14 +16,55 @@ if(NOT GPU_BACKEND) set(GPU_BACKEND "CUDA") endif() -if(GPU_BACKEND STREQUAL "ROCM") - enable_language(HIP) +# Compiled extension: AUTO/ON/OFF. On ROCm, AUTO falls back to the Python + Triton +# package when HIP cannot be configured (#1850: pip-SDK images without hip-lang). +if(NOT DEFINED FASTVIDEO_KERNEL_BUILD_EXTENSION AND DEFINED ENV{FASTVIDEO_KERNEL_BUILD_EXTENSION}) + set(FASTVIDEO_KERNEL_BUILD_EXTENSION "$ENV{FASTVIDEO_KERNEL_BUILD_EXTENSION}") +endif() +set(FASTVIDEO_KERNEL_BUILD_EXTENSION "AUTO" CACHE STRING "Build the compiled extension: AUTO/ON/OFF") +set_property(CACHE FASTVIDEO_KERNEL_BUILD_EXTENSION PROPERTY STRINGS AUTO ON OFF) +string(TOUPPER "${FASTVIDEO_KERNEL_BUILD_EXTENSION}" FASTVIDEO_KERNEL_BUILD_EXTENSION) +if(NOT FASTVIDEO_KERNEL_BUILD_EXTENSION MATCHES "^(AUTO|ON|OFF)$") + message(FATAL_ERROR + "fastvideo-kernel: FASTVIDEO_KERNEL_BUILD_EXTENSION must be AUTO, ON or OFF " + "(got '${FASTVIDEO_KERNEL_BUILD_EXTENSION}').") +endif() + +set(BUILD_CXX_KERNELS ON) +if(FASTVIDEO_KERNEL_BUILD_EXTENSION STREQUAL "OFF") + set(BUILD_CXX_KERNELS OFF) +elseif(GPU_BACKEND STREQUAL "ROCM") + include(CheckLanguage) + check_language(HIP) + if(CMAKE_HIP_COMPILER) + enable_language(HIP) + elseif(FASTVIDEO_KERNEL_BUILD_EXTENSION STREQUAL "ON") + message(FATAL_ERROR + "fastvideo-kernel: FASTVIDEO_KERNEL_BUILD_EXTENSION=ON, but CMake found no usable HIP " + "toolchain (it needs the hip-lang CMake package under ROCM_PATH or CMAKE_HIP_COMPILER). " + "Install the ROCm CMake packages, or build the Python + Triton package with " + "FASTVIDEO_KERNEL_BUILD_EXTENSION=OFF (build.sh --rocm --python-only).") + else() + message(STATUS + "fastvideo-kernel: no usable HIP toolchain; building the Python + Triton package " + "without the compiled extension (FASTVIDEO_KERNEL_BUILD_EXTENSION=AUTO)") + set(BUILD_CXX_KERNELS OFF) + endif() else() enable_language(CUDA) # Ensure CUDA toolkit targets (CUDA::cudart, CUDA::cuda_driver, etc.) are available. find_package(CUDAToolkit REQUIRED) endif() +if(NOT BUILD_CXX_KERNELS) + message(STATUS "============== fastvideo-kernel build summary ==============") + message(STATUS "host / backend: ${CMAKE_SYSTEM_PROCESSOR} / ${GPU_BACKEND}") + message(STATUS "fastvideo_kernel_ops: OFF (Python + Triton package only; FASTVIDEO_KERNEL_BUILD_EXTENSION=${FASTVIDEO_KERNEL_BUILD_EXTENSION})") + message(STATUS "Triton fallbacks ship in python/fastvideo_kernel/triton_kernels regardless.") + message(STATUS "============================================================") + return() +endif() + # Import common utils if needed, but we keep it simple for now # Find Python and Torch @@ -307,9 +348,6 @@ else() endif() endif() -# Always try to build the extension if CUDA is available, but conditionally add sources/flags -set(BUILD_CXX_KERNELS ON) - # --------------------------------------------------------------------------- # Per-arch split for the Blackwell FP4 (attn_qat_infer) build # --------------------------------------------------------------------------- diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index c321df5c4..c3454c42c 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -69,6 +69,20 @@ cd fastvideo-kernel ./build.sh --rocm ``` +The compiled extension needs the HIP CMake toolchain (`hip-lang`). Images that install +ROCm as a pip SDK, such as the `rocm/pytorch` 7.14 images, ship `hipcc` but not that +package, so `build.sh --rocm` falls back to the Python + Triton package there and says so +in the configure summary. That package is all the video sparse attention path needs on +ROCm. To ask for it explicitly, or to fail instead of falling back: + +```bash +./build.sh --rocm --python-only # Python + Triton only +CMAKE_ARGS="-DFASTVIDEO_KERNEL_BUILD_EXTENSION=ON" ./build.sh --rocm # error if HIP cannot be configured +``` + +`FASTVIDEO_KERNEL_BUILD_EXTENSION` (AUTO/ON/OFF) is also read from the environment by a +plain `pip install --no-build-isolation .`. + ### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath) The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 972e0bf3f..1ec35c196 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -5,9 +5,11 @@ set -ex # Usage: # ./build.sh # local build (torch-based arch detection, TK only on SM90) # ./build.sh --wheel-dir /tmp/wheels # build + install a reusable wheel +# ./build.sh --rocm --python-only # Python + Triton package, no compiled extension # Environment overrides (if set, they win over auto-detection): # TORCH_CUDA_ARCH_LIST -# CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND) +# CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND / +# FASTVIDEO_KERNEL_BUILD_EXTENSION) echo "Building fastvideo-kernel..." @@ -40,42 +42,19 @@ if [[ -n "${CONDA_PREFIX:-}" ]]; then unset _need_clean _host_arch fi -# Ensure only the kernel's required headers are initialized. A repository-wide -# update also clones the unrelated VBench evaluation submodule. Skip outside a -# git checkout (e.g. Docker contexts that exclude .git), where the submodule -# contents must already be present. -if git rev-parse --git-dir >/dev/null 2>&1; then - git submodule update --init --recursive include/cutlass include/tk -fi -# Fail fast with a clear message if the headers are still missing (e.g. a -# Docker context that excluded .git AND the submodule contents) instead of -# dying later in a wall of nvcc include errors. CUTLASS is consumed by the -# always-built turbodiffusion sources, so it is a hard error; ThunderKittens -# only feeds the TK-gated Hopper kernels, so a missing tree just warns (the -# TK gate resolves later, and non-SM90/ROCm builds never touch it). -if [ ! -d include/cutlass/include ]; then - echo "ERROR: include/cutlass/include is missing. Outside a git checkout the" >&2 - echo " CUTLASS sources must already be present (run" >&2 - echo " 'git submodule update --init --recursive include/cutlass include/tk'" >&2 - echo " in the source checkout, or include them in the build context)." >&2 - exit 1 -fi -if [ ! -d include/tk/include ]; then - echo "WARNING: include/tk/include is missing; ThunderKittens (Hopper sm_90a)" >&2 - echo " kernels cannot be built. Fine for non-SM90/ROCm targets." >&2 -fi - -# Install build dependencies -uv pip install scikit-build-core cmake ninja - GPU_BACKEND=CUDA WHEEL_DIR="" +PYTHON_ONLY=0 while [[ $# -gt 0 ]]; do case "$1" in --rocm) GPU_BACKEND=ROCM shift ;; + --python-only) + PYTHON_ONLY=1 + shift + ;; --wheel-dir) if [[ $# -lt 2 ]]; then echo "ERROR: --wheel-dir requires a directory path" >&2 @@ -96,6 +75,55 @@ has_cmake_arg() { [[ "${CMAKE_ARGS:-}" =~ (^|[[:space:]])-D${key}(=|$) ]] } +extension_mode_from_cmake_args() { + [[ "${CMAKE_ARGS:-}" =~ (^|[[:space:]])-DFASTVIDEO_KERNEL_BUILD_EXTENSION=([^[:space:]]*) ]] \ + && printf '%s' "${BASH_REMATCH[2]^^}" +} +if (( PYTHON_ONLY )); then + _explicit_mode="$(extension_mode_from_cmake_args)" + if [[ -n "${_explicit_mode}" && "${_explicit_mode}" != "OFF" ]]; then + echo "ERROR: --python-only conflicts with CMAKE_ARGS -DFASTVIDEO_KERNEL_BUILD_EXTENSION=${_explicit_mode}" >&2 + exit 2 + fi + if ! has_cmake_arg "FASTVIDEO_KERNEL_BUILD_EXTENSION"; then + CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_EXTENSION=OFF" + fi +elif [[ "$(extension_mode_from_cmake_args)" == "OFF" ]]; then + PYTHON_ONLY=1 +fi + +# The vendored headers only feed the compiled extension, so a Python-only build +# neither initializes nor checks them. +if ! (( PYTHON_ONLY )); then + # Ensure only the kernel's required headers are initialized. A repository-wide + # update also clones the unrelated VBench evaluation submodule. Skip outside a + # git checkout (e.g. Docker contexts that exclude .git), where the submodule + # contents must already be present. + if git rev-parse --git-dir >/dev/null 2>&1; then + git submodule update --init --recursive include/cutlass include/tk + fi + # Fail fast with a clear message if the headers are still missing (e.g. a + # Docker context that excluded .git AND the submodule contents) instead of + # dying later in a wall of nvcc include errors. CUTLASS is consumed by the + # always-built turbodiffusion sources, so it is a hard error; ThunderKittens + # only feeds the TK-gated Hopper kernels, so a missing tree just warns (the + # TK gate resolves later, and non-SM90/ROCm builds never touch it). + if [ ! -d include/cutlass/include ]; then + echo "ERROR: include/cutlass/include is missing. Outside a git checkout the" >&2 + echo " CUTLASS sources must already be present (run" >&2 + echo " 'git submodule update --init --recursive include/cutlass include/tk'" >&2 + echo " in the source checkout, or include them in the build context)." >&2 + exit 1 + fi + if [ ! -d include/tk/include ]; then + echo "WARNING: include/tk/include is missing; ThunderKittens (Hopper sm_90a)" >&2 + echo " kernels cannot be built. Fine for non-SM90/ROCm targets." >&2 + fi +fi + +# Install build dependencies +uv pip install scikit-build-core cmake ninja + detect_with_torch() { # Prefer the active venv's python directly over `uv run --active --no-project`, # which on some uv versions provisions its own interpreter and misses packages @@ -113,7 +141,7 @@ mj, mn = torch.cuda.get_device_capability(0) print(f'{mj}.{mn}')" } -if [ "${GPU_BACKEND}" = "CUDA" ]; then +if [ "${GPU_BACKEND}" = "CUDA" ] && ! (( PYTHON_ONLY )); then # Compute capability drives the arch/TK defaults below. Prefer an explicit # TORCH_CUDA_ARCH_LIST (works on GPU-less build machines such as CI/Docker); # only probe a live GPU via torch when no arch was provided. From d995516da00c24105aa841df1e690d3cd8a6c173 Mon Sep 17 00:00:00 2001 From: Yixing Wang Date: Mon, 21 Sep 2026 12:03:54 -0700 Subject: [PATCH 021/151] [bugfix] Restore page interaction after dismissing Create Job or successfully creating a job (#1869) Co-authored-by: leo --- .../e2e/create-job-interaction.spec.ts | 63 + apps/fastvideo_studio/e2e/create-job.spec.ts | 20 +- apps/fastvideo_studio/package-lock.json | 1590 +++++++++-------- apps/fastvideo_studio/package.json | 11 +- .../components/jobs/CreateJobButton.test.tsx | 78 +- .../src/components/jobs/CreateJobButton.tsx | 10 +- .../src/components/jobs/CreateJobModal.tsx | 5 + .../src/components/ui/accessibility.test.tsx | 22 +- .../src/components/ui/button.tsx | 4 +- .../src/components/ui/dialog.tsx | 2 +- .../src/components/ui/label.tsx | 2 +- .../src/components/ui/scroll-area.tsx | 2 +- .../src/components/ui/select.tsx | 2 +- .../src/components/ui/separator.tsx | 2 +- .../src/components/ui/slider.tsx | 2 +- .../src/components/ui/switch.tsx | 2 +- .../src/components/ui/tabs.tsx | 2 +- 17 files changed, 1050 insertions(+), 769 deletions(-) create mode 100644 apps/fastvideo_studio/e2e/create-job-interaction.spec.ts diff --git a/apps/fastvideo_studio/e2e/create-job-interaction.spec.ts b/apps/fastvideo_studio/e2e/create-job-interaction.spec.ts new file mode 100644 index 000000000..2bf3d47a6 --- /dev/null +++ b/apps/fastvideo_studio/e2e/create-job-interaction.spec.ts @@ -0,0 +1,63 @@ +import { expect, test } from '@playwright/test'; + +import { skipWithoutMock } from './helpers'; + +test.describe('create job interactions', () => { + skipWithoutMock(); + + for (const jobType of ['inference', 'finetuning', 'distillation']) { + test(`${jobType} remains interactive after repeated dialog dismissals`, async ({ page }) => { + await page.goto(`/${jobType}`); + const trigger = page.getByRole('button', { name: 'Create Job', exact: true }); + const dialog = page.getByRole('dialog'); + + // Exercise both dismissal paths and reopen without reloading the page. + for (const closeWithEscape of [false, true]) { + await trigger.click(); + await page.getByRole('menuitem').first().click(); + await expect(dialog).toBeVisible(); + if (closeWithEscape) { + await page.keyboard.press('Escape'); + } else { + await dialog.getByRole('button', { name: 'Close', exact: true }).click(); + } + await expect(dialog).toBeHidden(); + await expect(page.locator('body')).toHaveCSS('pointer-events', 'auto'); + await expect(trigger).toBeFocused(); + } + + await page.getByRole('link', { name: 'Datasets', exact: true }).click(); + await expect(page).toHaveURL(/\/datasets$/); + }); + } + + test('preserves keyboard menu dismissal and dialog focus trapping', async ({ page }) => { + await page.goto('/inference'); + const trigger = page.getByRole('button', { name: 'Create Job', exact: true }); + await trigger.focus(); + await page.keyboard.press('Enter'); + const firstItem = page.getByRole('menuitem').first(); + await expect(firstItem).toBeFocused(); + await page.keyboard.press('Escape'); + await expect(page.getByRole('menu')).toBeHidden(); + await expect(trigger).toBeFocused(); + await expect(page.locator('body')).toHaveCSS('pointer-events', 'auto'); + + await page.keyboard.press('Enter'); + await expect(firstItem).toBeFocused(); + await page.keyboard.press('Enter'); + const dialog = page.getByRole('dialog'); + await expect(dialog).toBeVisible(); + await expect(dialog.getByLabel('Name (optional)')).toBeFocused(); + + // Shift+Tab from the first field wraps to Close, then Tab wraps back. + await page.keyboard.press('Shift+Tab'); + await expect(dialog.getByRole('button', { name: 'Close', exact: true })).toBeFocused(); + await page.keyboard.press('Tab'); + await expect(dialog.getByLabel('Name (optional)')).toBeFocused(); + await page.keyboard.press('Escape'); + await expect(dialog).toBeHidden(); + await expect(trigger).toBeFocused(); + await expect(page.locator('body')).toHaveCSS('pointer-events', 'auto'); + }); +}); diff --git a/apps/fastvideo_studio/e2e/create-job.spec.ts b/apps/fastvideo_studio/e2e/create-job.spec.ts index 7afe9beaa..50d927b85 100644 --- a/apps/fastvideo_studio/e2e/create-job.spec.ts +++ b/apps/fastvideo_studio/e2e/create-job.spec.ts @@ -1,6 +1,6 @@ import { expect, test } from '@playwright/test'; -import { skipWithoutMock } from './helpers'; +import { API_BASE, skipWithoutMock } from './helpers'; /** * Create-job flow: open the Create Job modal on /inference, fill the prompt @@ -10,7 +10,8 @@ import { skipWithoutMock } from './helpers'; test.describe('create inference job', () => { skipWithoutMock(); - test('creates a T2V job and shows it in the queue', async ({ page }) => { + test('creates a T2V job and starts it without refreshing', async ({ page, request }) => { + await request.put(`${API_BASE}/settings`, { data: { autoStartJob: false } }); await page.goto('/inference'); // The trigger opens a real menu on click, so this path works for touch, @@ -38,5 +39,20 @@ test.describe('create inference job', () => { // Modal closes and the queue refreshes with the newly created job. await expect(dialog).toBeHidden(); await expect(page.getByText(prompt)).toBeVisible(); + await expect(page.locator('body')).toHaveCSS('pointer-events', 'auto'); + + const card = page.getByRole('article').filter({ hasText: prompt }); + await expect(card.getByText('pending', { exact: true })).toBeVisible(); + const started = page.waitForResponse((response) => + response.url().startsWith(`${API_BASE}/jobs/`) && + response.url().endsWith('/start') && + response.request().method() === 'POST', + ); + await card.getByRole('button', { name: 'Start', exact: true }).click(); + expect((await started).ok()).toBe(true); + await expect(card.getByText('running', { exact: true })).toBeVisible(); + + await page.getByRole('link', { name: 'Datasets', exact: true }).click(); + await expect(page).toHaveURL(/\/datasets$/); }); }); diff --git a/apps/fastvideo_studio/package-lock.json b/apps/fastvideo_studio/package-lock.json index 50ec23f50..70be47b8b 100644 --- a/apps/fastvideo_studio/package-lock.json +++ b/apps/fastvideo_studio/package-lock.json @@ -8,20 +8,11 @@ "name": "fastvideo-studio", "version": "0.1.0", "dependencies": { - "@radix-ui/react-dialog": "^1.1.0", - "@radix-ui/react-dropdown-menu": "^2.1.24", - "@radix-ui/react-label": "^2.1.8", - "@radix-ui/react-scroll-area": "^1.2.10", - "@radix-ui/react-select": "^2.2.6", - "@radix-ui/react-separator": "^1.1.8", - "@radix-ui/react-slider": "^1.2.0", - "@radix-ui/react-slot": "^1.2.4", - "@radix-ui/react-switch": "^1.1.0", - "@radix-ui/react-tabs": "^1.1.0", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "lucide-react": "^0.577.0", "next": "15.5.18", + "radix-ui": "^1.6.7", "react": "^19.1.0", "react-dom": "^19.1.0", "sonner": "^2.0.7", @@ -1769,24 +1760,215 @@ } }, "node_modules/@radix-ui/number": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/number/-/number-1.1.2.tgz", - "integrity": "sha512-ceTwaxc4I5IOi97DgCotl3pqiyRGvffcc0oOsE2dQYaJOFIDsDt4VWG6xEbg1QePv9QWausCEIppud/tJ1wNig==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@radix-ui/number/-/number-1.1.3.tgz", + "integrity": "sha512-Road2bidD0uu/1BGDOWNdPI06g0lIRy6IF9GZcIrDK2KGItfor8IQwQa+yM2ERgHM1MmHxaxpTzk0/Jp42lNfA==", "license": "MIT" }, "node_modules/@radix-ui/primitive": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/primitive/-/primitive-1.1.4.tgz", - "integrity": "sha512-7AdCK9PQyiljKoBDbN8OuctCbd/esdwZPQ8RtOE3SsyQtUpiPb+ND75q0jEhC1m1ecBI0MFNeLJvwIh9iKHRcQ==", + "version": "1.1.7", + "resolved": "https://registry.npmjs.org/@radix-ui/primitive/-/primitive-1.1.7.tgz", + "integrity": "sha512-rqWnm76nYT8HoNNqEjpgJ7Pw/DrBj5iBTrmEPo6HTX5+VJyBNOqTdv4g89G63HuR5g0AaENoAcH7Is5fF2kZ8Q==", "license": "MIT" }, - "node_modules/@radix-ui/react-arrow": { - "version": "1.1.11", - "resolved": "https://registry.npmjs.org/@radix-ui/react-arrow/-/react-arrow-1.1.11.tgz", - "integrity": "sha512-Kdil9BB1rIFC/khmf4hC35bn8701AJcizTU7G7cUbEbk5XqqbjDuHW60uUfKqO5WojjZcbAW51Q7P0hRmMLw8A==", + "node_modules/@radix-ui/react-accessible-icon": { + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-accessible-icon/-/react-accessible-icon-1.1.15.tgz", + "integrity": "sha512-WTQwcAvQf5sOcuUyi90lKPbhwcvQ+j55cjrSmeaN+L2vKU3DooOvlKw2MDeiJ5IkV5N905KW0/fGojKOBhD11A==", "license": "MIT", "dependencies": { - "@radix-ui/react-primitive": "2.1.7" + "@radix-ui/react-visually-hidden": "1.2.11" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-accordion": { + "version": "1.2.20", + "resolved": "https://registry.npmjs.org/@radix-ui/react-accordion/-/react-accordion-1.2.20.tgz", + "integrity": "sha512-jDhG9FvAEnlhnjrsINbNXcUa4G+L1KqSkJSunkbKEzFRcAb52jvM0PjPxPRvhe1HNc5F5yc0yzzWeeqlH4yBIg==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collapsible": "1.1.20", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-alert-dialog": { + "version": "1.1.23", + "resolved": "https://registry.npmjs.org/@radix-ui/react-alert-dialog/-/react-alert-dialog-1.1.23.tgz", + "integrity": "sha512-VAYOiQRqj3GPpYJE0I9J+X8Ip05cyVlNdKOFeiGS2Ou1HHGfpl0BxOyZm6nmVDyU+W+NF3/XLzmjHmVGydhwgA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dialog": "1.1.23", + "@radix-ui/react-primitive": "2.1.10" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-arrow": { + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-arrow/-/react-arrow-1.1.15.tgz", + "integrity": "sha512-v4zggRcjadnI+ClKDuijlQEW4tw3NoaeHc/PwpKnLoLLKNUG4InLegkstooLcRIUWCs+8L22dGURCVuFfOKfnA==", + "license": "MIT", + "dependencies": { + "@radix-ui/react-primitive": "2.1.10" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-aspect-ratio": { + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-aspect-ratio/-/react-aspect-ratio-1.1.15.tgz", + "integrity": "sha512-fy+dyVR+90nelK8rqIznFlxzx7uPcGbhxH8Nfr2bHb4UfSe+e3hklOC0luK0hDwVwnRX7xTRySpsrQVeW+/oNQ==", + "license": "MIT", + "dependencies": { + "@radix-ui/react-primitive": "2.1.10" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-avatar": { + "version": "1.2.6", + "resolved": "https://registry.npmjs.org/@radix-ui/react-avatar/-/react-avatar-1.2.6.tgz", + "integrity": "sha512-4ULOTJ/mqy2hT9GlWa/MFHxHSvH3nJzHnZM1waNsc5Bonv7i70aNenghXmD97S6OJ81ekXONGGt4nT1r0PfEdA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-is-hydrated": "0.1.3", + "@radix-ui/react-use-layout-effect": "1.1.4" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-checkbox": { + "version": "1.3.11", + "resolved": "https://registry.npmjs.org/@radix-ui/react-checkbox/-/react-checkbox-1.3.11.tgz", + "integrity": "sha512-Gnptr9pDDQxD3hgq2dtPbtrp/c2qH1mBwIzw3X/ivrMb2e1t0jMTi606fVEqFPaQR1ggXIVQWKj3P2WW9v7zGQ==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-size": "1.1.4" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-collapsible": { + "version": "1.1.20", + "resolved": "https://registry.npmjs.org/@radix-ui/react-collapsible/-/react-collapsible-1.1.20.tgz", + "integrity": "sha512-mcGesGplBnzN2sbvJETzpCNfSMyPnb29q1GRLU+Ib7bJrpIG2ywmRoh2V5VbA2uNvKikKUlVbAPks7JDjz4A8Q==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -1804,15 +1986,15 @@ } }, "node_modules/@radix-ui/react-collection": { - "version": "1.1.11", - "resolved": "https://registry.npmjs.org/@radix-ui/react-collection/-/react-collection-1.1.11.tgz", - "integrity": "sha512-djW9+zeg137KQdlPtmE8xnaD+K2rcXXMWFrSg0hsmYZ6HRbdTA7tDHFgpaW9+huWVEu0RCabL+985T4TA0BE7g==", + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-collection/-/react-collection-1.1.15.tgz", + "integrity": "sha512-9W+B9NPF0NaaPh/1NJd3+KqsnlLqU9H7T2rvww+fp+T/evVXdNAyYcnfRQZFOjkR1ajQp3yORlqnI8soawLvNA==", "license": "MIT", "dependencies": { - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-slot": "1.3.0" + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-slot": "1.3.3" }, "peerDependencies": { "@types/react": "*", @@ -1830,9 +2012,9 @@ } }, "node_modules/@radix-ui/react-compose-refs": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@radix-ui/react-compose-refs/-/react-compose-refs-1.1.3.tgz", - "integrity": "sha512-rYOP8OMnuuPMQF1uhPVlGNcCDlkokKqGFE3JcxFViIkAXP7EvFWUliJAstrapypaBLJNHbZL6jGhbVDGTwmVhA==", + "version": "1.1.5", + "resolved": "https://registry.npmjs.org/@radix-ui/react-compose-refs/-/react-compose-refs-1.1.5.tgz", + "integrity": "sha512-+48PbAAbq3didjJxa+OaWY2ZwgAKsNiRGyeHKszblZMQ+kcpd9pAaT11cMkGEie0vsOi3QdeTE6d5Fe3Gn61kA==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -1845,9 +2027,9 @@ } }, "node_modules/@radix-ui/react-context": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-context/-/react-context-1.1.4.tgz", - "integrity": "sha512-QwH4PO5urrbO+FaGd5Aglg+YJgWTyyuZ3g/6mKvsqraLkglDdckw9JafgL5McL5VEJ6EPNduPaT3ZE9BttDAqg==", + "version": "1.2.2", + "resolved": "https://registry.npmjs.org/@radix-ui/react-context/-/react-context-1.2.2.tgz", + "integrity": "sha512-RHCUGwKHDr0hDGg4X7ma4JG4/+12qxw8rkh5QKdDldlCvtja6nUx1Ef/8HVrJze81lEsgLQlqjzjGNHantgnQA==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -1859,24 +2041,52 @@ } } }, - "node_modules/@radix-ui/react-dialog": { - "version": "1.1.18", - "resolved": "https://registry.npmjs.org/@radix-ui/react-dialog/-/react-dialog-1.1.18.tgz", - "integrity": "sha512-apa28mldjMgORmE6g/w3sCcA0Y9UAVeeDVoozN4i7kOw12mLl9RBchfzK3Nn6qxOWjrZhK1Lfy7f07kyzxtnBw==", + "node_modules/@radix-ui/react-context-menu": { + "version": "2.3.7", + "resolved": "https://registry.npmjs.org/@radix-ui/react-context-menu/-/react-context-menu-2.3.7.tgz", + "integrity": "sha512-CtXP35dxaB5T3zXSd+E3uHe/QpXcpYnZmxp6OaIbfthtfW4wyb77M23BG+bwIJDtsMwEP/YssdsmNyZu7jhWew==", "license": "MIT", "dependencies": { - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-dismissable-layer": "1.1.14", - "@radix-ui/react-focus-guards": "1.1.4", - "@radix-ui/react-focus-scope": "1.1.11", - "@radix-ui/react-id": "1.1.2", - "@radix-ui/react-portal": "1.1.13", - "@radix-ui/react-presence": "1.1.6", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-slot": "1.3.0", - "@radix-ui/react-use-controllable-state": "1.2.3", + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-menu": "2.1.24", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-dialog": { + "version": "1.1.23", + "resolved": "https://registry.npmjs.org/@radix-ui/react-dialog/-/react-dialog-1.1.23.tgz", + "integrity": "sha512-Ksw4WeROkO4rC9k/onilX/Ao2Cr1ku1unMNH+XSCcP4jSXYu7HDsg9n4ojMjVb22XpYjAQ9qfrFlVbru1vXDUA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-focus-guards": "1.1.6", + "@radix-ui/react-focus-scope": "1.1.16", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-slot": "1.3.3", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", "aria-hidden": "^1.2.4", "react-remove-scroll": "^2.7.2" }, @@ -1896,9 +2106,9 @@ } }, "node_modules/@radix-ui/react-direction": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-direction/-/react-direction-1.1.2.tgz", - "integrity": "sha512-C3vFhbyi4SW3PmbAi6Awpu4OzJtd0MxGurvSsYtr7p7nM8RNB3VAF3CUmnp2j50knpkrRcB7+ycVXzgLgF6yNA==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-direction/-/react-direction-1.1.4.tgz", + "integrity": "sha512-5pzg4FGQNpExhnhT2zlrP1wZFaYCd1K0nYWoFAdcYoYK868IEigqMX3B3f8yIoRlAhAeDWciLI6ZdCKHF9P4Vg==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -1911,16 +2121,16 @@ } }, "node_modules/@radix-ui/react-dismissable-layer": { - "version": "1.1.14", - "resolved": "https://registry.npmjs.org/@radix-ui/react-dismissable-layer/-/react-dismissable-layer-1.1.14.tgz", - "integrity": "sha512-4lUhWTWAjbDIqFrAPWJ3WqBOpO5YchVZ88X3nh6H9Lu5AFi5nCUeTPj3D8FSDmabmFeRe9ME0BDA4MwKTha5GQ==", + "version": "1.1.19", + "resolved": "https://registry.npmjs.org/@radix-ui/react-dismissable-layer/-/react-dismissable-layer-1.1.19.tgz", + "integrity": "sha512-8g4pfOL9HoKKLWGiypT+dphVqjFfmcXO5GBnhsG6zI+lxAx/8feQpr+1LSN8Re3hiZ+XkLNS4O9ztK11/LzQ6w==", "license": "MIT", "dependencies": { - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-callback-ref": "1.1.2", - "@radix-ui/react-use-effect-event": "0.0.3" + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-effect-event": "0.0.5" }, "peerDependencies": { "@types/react": "*", @@ -1966,16 +2176,10 @@ } } }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/primitive": { - "version": "1.1.7", - "resolved": "https://registry.npmjs.org/@radix-ui/primitive/-/primitive-1.1.7.tgz", - "integrity": "sha512-rqWnm76nYT8HoNNqEjpgJ7Pw/DrBj5iBTrmEPo6HTX5+VJyBNOqTdv4g89G63HuR5g0AaENoAcH7Is5fF2kZ8Q==", - "license": "MIT" - }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-compose-refs": { - "version": "1.1.5", - "resolved": "https://registry.npmjs.org/@radix-ui/react-compose-refs/-/react-compose-refs-1.1.5.tgz", - "integrity": "sha512-+48PbAAbq3didjJxa+OaWY2ZwgAKsNiRGyeHKszblZMQ+kcpd9pAaT11cMkGEie0vsOi3QdeTE6d5Fe3Gn61kA==", + "node_modules/@radix-ui/react-focus-guards": { + "version": "1.1.6", + "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-guards/-/react-focus-guards-1.1.6.tgz", + "integrity": "sha512-RNOJjfZMTyBM6xYmV3IVGXkPjIhcBAuv48POevAXwrGJhkWZ9p1rFoIS1JFooPuT193AZmRsCPhpoVJxx6OPoQ==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -1987,46 +2191,15 @@ } } }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-context": { - "version": "1.2.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-context/-/react-context-1.2.2.tgz", - "integrity": "sha512-RHCUGwKHDr0hDGg4X7ma4JG4/+12qxw8rkh5QKdDldlCvtja6nUx1Ef/8HVrJze81lEsgLQlqjzjGNHantgnQA==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-id": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-id/-/react-id-1.1.4.tgz", - "integrity": "sha512-TMQp2llA+RYn7JcjnrMnz7wN4pcVttPZnRZo52PLQsoLVKzNlVwUeHmfePgTgRluXFvlD3GD5g5MOVVTJCO0qA==", + "node_modules/@radix-ui/react-focus-scope": { + "version": "1.1.16", + "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-scope/-/react-focus-scope-1.1.16.tgz", + "integrity": "sha512-wmRZ2WWLvmt6KHy2rNPOdPUjwq5xOHY02+m+udwJTn0aNIox/rkskAvJTyTLGhPK6KgrUjlJUJpgmx/+wFiFIQ==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-primitive": { - "version": "2.1.10", - "resolved": "https://registry.npmjs.org/@radix-ui/react-primitive/-/react-primitive-2.1.10.tgz", - "integrity": "sha512-MucOnzh6hR5mid6VpkbglRAMYMjKLqRnGBbjXkzjK52fuQDd1qbkx78a5P40mkcnVXJdEVxm26E9OPAiUq7nBg==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-slot": "1.3.3" + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-callback-ref": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -2043,101 +2216,49 @@ } } }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-slot": { - "version": "1.3.3", - "resolved": "https://registry.npmjs.org/@radix-ui/react-slot/-/react-slot-1.3.3.tgz", - "integrity": "sha512-qx7oqnYbxnK9kYI9m317qmFmEgo6ywqWvbTogdj7cL9p3/yx4M48p7Rnw5z3H890cL/ow/EeWJsuTykeZVXP5Q==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-compose-refs": "1.1.5" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-use-controllable-state": { - "version": "1.2.6", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-controllable-state/-/react-use-controllable-state-1.2.6.tgz", - "integrity": "sha512-uEQJGT97ZA/TgP/Hydw47lHu+/vQj6z/0jA+WeTbK1o9Rx45GImjpD0tc3W5ad3D6XTSR6e1yEO0FvGq6WQfVQ==", + "node_modules/@radix-ui/react-form": { + "version": "0.1.16", + "resolved": "https://registry.npmjs.org/@radix-ui/react-form/-/react-form-0.1.16.tgz", + "integrity": "sha512-Q4TLEn2A7TAypxwmd6R9EwrlXDvkfYSDMrq9/887AXAGh+G1rH+kYJKSTv+Si9Y0JPKTwKYv6PviAJosysNimA==", "license": "MIT", "dependencies": { "@radix-ui/primitive": "1.1.7", - "@radix-ui/react-use-effect-event": "0.0.5", - "@radix-ui/react-use-layout-effect": "1.1.4" + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-label": "2.1.15", + "@radix-ui/react-primitive": "2.1.10" }, "peerDependencies": { "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "peerDependenciesMeta": { "@types/react": { "optional": true + }, + "@types/react-dom": { + "optional": true } } }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-use-effect-event": { - "version": "0.0.5", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-effect-event/-/react-use-effect-event-0.0.5.tgz", - "integrity": "sha512-7cshFL8HGS/7HEiHH+9kL9HBwp2sa9yX18Knwek6KYWmXwM7pegMgta2AXMQKI+rq3JnfSj9x8wYqFMTdG1Jgg==", + "node_modules/@radix-ui/react-hover-card": { + "version": "1.1.23", + "resolved": "https://registry.npmjs.org/@radix-ui/react-hover-card/-/react-hover-card-1.1.23.tgz", + "integrity": "sha512-H8qONfZd3ltrU3+jHCIgITbWo6e1iTKvP9DHdrvYbX48ooRM5FjEDTn16AMwdfuOGkWdZEhpl3PLL/Wk/AnHDQ==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-dropdown-menu/node_modules/@radix-ui/react-use-layout-effect": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-layout-effect/-/react-use-layout-effect-1.1.4.tgz", - "integrity": "sha512-K20DkRkUwDnxEYMBPcg3Y6voLkEy5p5QQmszZgLngKKiC7dzBR/aEuK3w1qlx2JWDUNH6FluahYdgR3BP+QbYw==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-focus-guards": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-guards/-/react-focus-guards-1.1.4.tgz", - "integrity": "sha512-cot/aB/mOm0IYVYTTmQcEEK1M48lZWi8FlYe5nDPQQ8NYZUlXEFgncJ9p2Kzer3RKSrY7cTTpEMLZKNo9QoP5Q==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-focus-scope": { - "version": "1.1.11", - "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-scope/-/react-focus-scope-1.1.11.tgz", - "integrity": "sha512-Mn88Vg2whaRocGJNOH+DKFqYm6ySFPQaiwHNxZPyjn99B52KAEJWWY9NP83+nWdk2HM3rdov+STu9AG471Rt9w==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-callback-ref": "1.1.2" + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-popper": "1.3.7", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6" }, "peerDependencies": { "@types/react": "*", @@ -2155,12 +2276,12 @@ } }, "node_modules/@radix-ui/react-id": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-id/-/react-id-1.1.2.tgz", - "integrity": "sha512-orBC88futVpqCmhX1p4cvquNHsELQ+w+vBJnuj3ftETI5bJb0bZn3Tqu3SWN2IOcPycTnMGnhwoermvISt72sA==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-id/-/react-id-1.1.4.tgz", + "integrity": "sha512-TMQp2llA+RYn7JcjnrMnz7wN4pcVttPZnRZo52PLQsoLVKzNlVwUeHmfePgTgRluXFvlD3GD5g5MOVVTJCO0qA==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.2" + "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -2173,12 +2294,12 @@ } }, "node_modules/@radix-ui/react-label": { - "version": "2.1.11", - "resolved": "https://registry.npmjs.org/@radix-ui/react-label/-/react-label-2.1.11.tgz", - "integrity": "sha512-3PKvDDxOn62k0oV1n4QtNtD2vpu+zYjXR7ojLBPaO6SPvhy53yg0vAmgNeBQeJW5rV3dffoRG+HYfLBZuzw0CQ==", + "version": "2.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-label/-/react-label-2.1.15.tgz", + "integrity": "sha512-o/rdYEwZTTo5tjknnPeyQFU45kUC4i/XyeDPP+HGyi6XqpOP6Zf5Ya5vh/Yfe9Id5JiuWnnAx2XqIeD3UYZt0g==", "license": "MIT", "dependencies": { - "@radix-ui/react-primitive": "2.1.7" + "@radix-ui/react-primitive": "2.1.10" }, "peerDependencies": { "@types/react": "*", @@ -2235,117 +2356,58 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/primitive": { - "version": "1.1.7", - "resolved": "https://registry.npmjs.org/@radix-ui/primitive/-/primitive-1.1.7.tgz", - "integrity": "sha512-rqWnm76nYT8HoNNqEjpgJ7Pw/DrBj5iBTrmEPo6HTX5+VJyBNOqTdv4g89G63HuR5g0AaENoAcH7Is5fF2kZ8Q==", - "license": "MIT" - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-arrow": { - "version": "1.1.15", - "resolved": "https://registry.npmjs.org/@radix-ui/react-arrow/-/react-arrow-1.1.15.tgz", - "integrity": "sha512-v4zggRcjadnI+ClKDuijlQEW4tw3NoaeHc/PwpKnLoLLKNUG4InLegkstooLcRIUWCs+8L22dGURCVuFfOKfnA==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-primitive": "2.1.10" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-collection": { - "version": "1.1.15", - "resolved": "https://registry.npmjs.org/@radix-ui/react-collection/-/react-collection-1.1.15.tgz", - "integrity": "sha512-9W+B9NPF0NaaPh/1NJd3+KqsnlLqU9H7T2rvww+fp+T/evVXdNAyYcnfRQZFOjkR1ajQp3yORlqnI8soawLvNA==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-compose-refs": "1.1.5", - "@radix-ui/react-context": "1.2.2", - "@radix-ui/react-primitive": "2.1.10", - "@radix-ui/react-slot": "1.3.3" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-compose-refs": { - "version": "1.1.5", - "resolved": "https://registry.npmjs.org/@radix-ui/react-compose-refs/-/react-compose-refs-1.1.5.tgz", - "integrity": "sha512-+48PbAAbq3didjJxa+OaWY2ZwgAKsNiRGyeHKszblZMQ+kcpd9pAaT11cMkGEie0vsOi3QdeTE6d5Fe3Gn61kA==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-context": { - "version": "1.2.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-context/-/react-context-1.2.2.tgz", - "integrity": "sha512-RHCUGwKHDr0hDGg4X7ma4JG4/+12qxw8rkh5QKdDldlCvtja6nUx1Ef/8HVrJze81lEsgLQlqjzjGNHantgnQA==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-direction": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-direction/-/react-direction-1.1.4.tgz", - "integrity": "sha512-5pzg4FGQNpExhnhT2zlrP1wZFaYCd1K0nYWoFAdcYoYK868IEigqMX3B3f8yIoRlAhAeDWciLI6ZdCKHF9P4Vg==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-dismissable-layer": { - "version": "1.1.19", - "resolved": "https://registry.npmjs.org/@radix-ui/react-dismissable-layer/-/react-dismissable-layer-1.1.19.tgz", - "integrity": "sha512-8g4pfOL9HoKKLWGiypT+dphVqjFfmcXO5GBnhsG6zI+lxAx/8feQpr+1LSN8Re3hiZ+XkLNS4O9ztK11/LzQ6w==", + "node_modules/@radix-ui/react-menubar": { + "version": "1.1.24", + "resolved": "https://registry.npmjs.org/@radix-ui/react-menubar/-/react-menubar-1.1.24.tgz", + "integrity": "sha512-eeVs0vf7cuqXaM0qLQCPcufImiJNVBXdJDLu7ZGYl2732UH23Qat/foNGrr6vYV3/DdTsBqASoggUFgH14OcZA==", "license": "MIT", "dependencies": { "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-menu": "2.1.24", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-navigation-menu": { + "version": "1.2.22", + "resolved": "https://registry.npmjs.org/@radix-ui/react-navigation-menu/-/react-navigation-menu-1.2.22.tgz", + "integrity": "sha512-ou7iLEJ+yrhQndkkA4U21XIdS/CS45F4iXIkTZcb6/Ne9EMsOuDudVmCwmDnfFZZ+y1FZqXRNSIgBy+YMvZVZg==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-presence": "1.1.10", "@radix-ui/react-primitive": "2.1.10", "@radix-ui/react-use-callback-ref": "1.1.4", - "@radix-ui/react-use-effect-event": "0.0.5" + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-use-previous": "1.1.4", + "@radix-ui/react-visually-hidden": "1.2.11" }, "peerDependencies": { "@types/react": "*", @@ -2362,65 +2424,108 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-focus-guards": { - "version": "1.1.6", - "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-guards/-/react-focus-guards-1.1.6.tgz", - "integrity": "sha512-RNOJjfZMTyBM6xYmV3IVGXkPjIhcBAuv48POevAXwrGJhkWZ9p1rFoIS1JFooPuT193AZmRsCPhpoVJxx6OPoQ==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-focus-scope": { - "version": "1.1.16", - "resolved": "https://registry.npmjs.org/@radix-ui/react-focus-scope/-/react-focus-scope-1.1.16.tgz", - "integrity": "sha512-wmRZ2WWLvmt6KHy2rNPOdPUjwq5xOHY02+m+udwJTn0aNIox/rkskAvJTyTLGhPK6KgrUjlJUJpgmx/+wFiFIQ==", + "node_modules/@radix-ui/react-one-time-password-field": { + "version": "0.1.16", + "resolved": "https://registry.npmjs.org/@radix-ui/react-one-time-password-field/-/react-one-time-password-field-0.1.16.tgz", + "integrity": "sha512-Tj9P6ntAJEw52oq/F0AGknXR4XncxEt7XU47O3xJQOiWfLzEy3d9gtgKfvjSzGxzHkfL+VzvxGu2KTFsloJqXw==", "license": "MIT", "dependencies": { + "@radix-ui/number": "1.1.3", + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", "@radix-ui/react-primitive": "2.1.10", - "@radix-ui/react-use-callback-ref": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-id": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-id/-/react-id-1.1.4.tgz", - "integrity": "sha512-TMQp2llA+RYn7JcjnrMnz7wN4pcVttPZnRZo52PLQsoLVKzNlVwUeHmfePgTgRluXFvlD3GD5g5MOVVTJCO0qA==", - "license": "MIT", - "dependencies": { + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-effect-event": "0.0.5", + "@radix-ui/react-use-is-hydrated": "0.1.3", "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" }, "peerDependenciesMeta": { "@types/react": { "optional": true + }, + "@types/react-dom": { + "optional": true } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-popper": { + "node_modules/@radix-ui/react-password-toggle-field": { + "version": "0.1.11", + "resolved": "https://registry.npmjs.org/@radix-ui/react-password-toggle-field/-/react-password-toggle-field-0.1.11.tgz", + "integrity": "sha512-4gvFnmDXu3dgj21CqsufzIameRvlRd4SBqaWhcrlrNhRo0Y5i/49AmRJYe1fdAM3G2VNBbmin4b0D6cdQocwgw==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-effect-event": "0.0.5", + "@radix-ui/react-use-is-hydrated": "0.1.3" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-popover": { + "version": "1.1.23", + "resolved": "https://registry.npmjs.org/@radix-ui/react-popover/-/react-popover-1.1.23.tgz", + "integrity": "sha512-mw58MrBlyHWFisTOYignD0vf/3gdcgAR+9of1s9G/38CbFiUwH1nCDkc0AUM9IrXFgN5Ue8n45j9WCgyM1sbiQ==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-focus-guards": "1.1.6", + "@radix-ui/react-focus-scope": "1.1.16", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-popper": "1.3.7", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-slot": "1.3.3", + "@radix-ui/react-use-controllable-state": "1.2.6", + "aria-hidden": "^1.2.4", + "react-remove-scroll": "^2.7.2" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-popper": { "version": "1.3.7", "resolved": "https://registry.npmjs.org/@radix-ui/react-popper/-/react-popper-1.3.7.tgz", "integrity": "sha512-UsJrrd7w4wuKKTdvd/DNERVlwSlUcyXzjhyDwBk+3aPOsCjOY6ZSbxuw8E6lZTjjfP8Cpd0J8VVkrYUWyGYXyg==", @@ -2452,7 +2557,7 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-portal": { + "node_modules/@radix-ui/react-portal": { "version": "1.1.17", "resolved": "https://registry.npmjs.org/@radix-ui/react-portal/-/react-portal-1.1.17.tgz", "integrity": "sha512-vKQLcWypUnwZVvfV7UkGahH2g6ySe8M8R+zYBwPrv5byZ9QAW6cQVvNKo7GgmD+p8aYb6D9JBuvy8/WhOno2wQ==", @@ -2476,7 +2581,7 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-presence": { + "node_modules/@radix-ui/react-presence": { "version": "1.1.10", "resolved": "https://registry.npmjs.org/@radix-ui/react-presence/-/react-presence-1.1.10.tgz", "integrity": "sha512-3wyzCQ6+ubRA+D4uv9m95JYLXxmOHp05qjrkjeA7uKHHtjpPggQzc6DAb0URl7j67oR0K2foO4ip27TiX037Bw==", @@ -2499,7 +2604,7 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-primitive": { + "node_modules/@radix-ui/react-primitive": { "version": "2.1.10", "resolved": "https://registry.npmjs.org/@radix-ui/react-primitive/-/react-primitive-2.1.10.tgz", "integrity": "sha512-MucOnzh6hR5mid6VpkbglRAMYMjKLqRnGBbjXkzjK52fuQDd1qbkx78a5P40mkcnVXJdEVxm26E9OPAiUq7nBg==", @@ -2522,7 +2627,62 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-roving-focus": { + "node_modules/@radix-ui/react-progress": { + "version": "1.1.16", + "resolved": "https://registry.npmjs.org/@radix-ui/react-progress/-/react-progress-1.1.16.tgz", + "integrity": "sha512-5XnomAsoZZCY+KNTxbIghpGqPruZvKFNlvcAljVAOdDRDsH4/OZQxhtwo5wdtoDM5R6MhJBb2sPnDuRFep3lzg==", + "license": "MIT", + "dependencies": { + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-primitive": "2.1.10" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-radio-group": { + "version": "1.4.7", + "resolved": "https://registry.npmjs.org/@radix-ui/react-radio-group/-/react-radio-group-1.4.7.tgz", + "integrity": "sha512-cgYFEkntCxppHZgtSZ+7vh0wbZQ+IC7PPMw8DSnRG27B6kDd32/Zw0OJt7dGDigCoprMuWHjg2PvUn3PYvPFoQ==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-size": "1.1.4" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-roving-focus": { "version": "1.1.19", "resolved": "https://registry.npmjs.org/@radix-ui/react-roving-focus/-/react-roving-focus-1.1.19.tgz", "integrity": "sha512-V9jI6hDjT7l3jsCQD9bLNvDLM3tH/gdbOTp7Tefp3hbbgCGQoK7tUvrWiRlcoBHIZ809ElXwNQwVo0B98LuTXQ==", @@ -2555,282 +2715,21 @@ } } }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-slot": { - "version": "1.3.3", - "resolved": "https://registry.npmjs.org/@radix-ui/react-slot/-/react-slot-1.3.3.tgz", - "integrity": "sha512-qx7oqnYbxnK9kYI9m317qmFmEgo6ywqWvbTogdj7cL9p3/yx4M48p7Rnw5z3H890cL/ow/EeWJsuTykeZVXP5Q==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-compose-refs": "1.1.5" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-callback-ref": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-callback-ref/-/react-use-callback-ref-1.1.4.tgz", - "integrity": "sha512-R6OUY2e2fA6Yn6s+VSx5KBV6Nx8LQEhu+cz7LCej18rQ1HLyg9PSC9jP/ZNx0o6FAIK9c0F1kHylzSxKsdlkrQ==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-controllable-state": { - "version": "1.2.6", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-controllable-state/-/react-use-controllable-state-1.2.6.tgz", - "integrity": "sha512-uEQJGT97ZA/TgP/Hydw47lHu+/vQj6z/0jA+WeTbK1o9Rx45GImjpD0tc3W5ad3D6XTSR6e1yEO0FvGq6WQfVQ==", - "license": "MIT", - "dependencies": { - "@radix-ui/primitive": "1.1.7", - "@radix-ui/react-use-effect-event": "0.0.5", - "@radix-ui/react-use-layout-effect": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-effect-event": { - "version": "0.0.5", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-effect-event/-/react-use-effect-event-0.0.5.tgz", - "integrity": "sha512-7cshFL8HGS/7HEiHH+9kL9HBwp2sa9yX18Knwek6KYWmXwM7pegMgta2AXMQKI+rq3JnfSj9x8wYqFMTdG1Jgg==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-layout-effect": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-layout-effect/-/react-use-layout-effect-1.1.4.tgz", - "integrity": "sha512-K20DkRkUwDnxEYMBPcg3Y6voLkEy5p5QQmszZgLngKKiC7dzBR/aEuK3w1qlx2JWDUNH6FluahYdgR3BP+QbYw==", - "license": "MIT", - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-rect": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-rect/-/react-use-rect-1.1.4.tgz", - "integrity": "sha512-cSOCh6JlkmfjLyNcLiu2nB4v+nm+dkZ+Q5KHWk/soo4U7ZLiEQFKHK9/YmtBHjfCEaU43IBKQOc4/uJmCaiCTQ==", - "license": "MIT", - "dependencies": { - "@radix-ui/rect": "1.1.3" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/react-use-size": { - "version": "1.1.4", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-size/-/react-use-size-1.1.4.tgz", - "integrity": "sha512-D3anSY15EJoxrihpsXI6SMrmmonnQtR2ni7arO+Lfdg3O95b9hNXxONk8jA5C8ANdF/h5HMAxejgs8PWJ6rlhw==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.4" - }, - "peerDependencies": { - "@types/react": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-menu/node_modules/@radix-ui/rect": { - "version": "1.1.3", - "resolved": "https://registry.npmjs.org/@radix-ui/rect/-/rect-1.1.3.tgz", - "integrity": "sha512-JtyZR+mqgBibTo8xea3B6ZRmzZiM/YeVBtUkas6zMuXjAlfIFIW2FgqeM9eLyvEaYX66vr6DJMK+4U6LV0KhNw==", - "license": "MIT" - }, - "node_modules/@radix-ui/react-popper": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-popper/-/react-popper-1.3.2.tgz", - "integrity": "sha512-3QXNeMkdshed1MR3LNoiCirBywRFPkD8ETJa/HlPuLwSajaQixf2ro+isoDNJlGABg9ug41XuZpINZJIle4XWg==", - "license": "MIT", - "dependencies": { - "@floating-ui/react-dom": "^2.0.0", - "@radix-ui/react-arrow": "1.1.11", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-callback-ref": "1.1.2", - "@radix-ui/react-use-layout-effect": "1.1.2", - "@radix-ui/react-use-rect": "1.1.2", - "@radix-ui/react-use-size": "1.1.2", - "@radix-ui/rect": "1.1.2" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-portal": { - "version": "1.1.13", - "resolved": "https://registry.npmjs.org/@radix-ui/react-portal/-/react-portal-1.1.13.tgz", - "integrity": "sha512-z3oXfmaHLJTF1wktbjgD6cn9jiEbq3WSondB10LIuIt2m2Ym4iJlrW04/euMwENDdWDdE7z+OuY7Qyp1YpRSwA==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-layout-effect": "1.1.2" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-presence": { - "version": "1.1.6", - "resolved": "https://registry.npmjs.org/@radix-ui/react-presence/-/react-presence-1.1.6.tgz", - "integrity": "sha512-zdTk4PlUO0E18HnZ3wYbW0KkJJxWCdiNYp6g6X1PtONFhxVkg01vliTJAmwIszU6mHiyBOoW9P0rAugl5/hULQ==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.2" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-primitive": { - "version": "2.1.7", - "resolved": "https://registry.npmjs.org/@radix-ui/react-primitive/-/react-primitive-2.1.7.tgz", - "integrity": "sha512-bC3NiwsprbxKjuon9l7X6BUTw7FPVzEYaL92MPEY5SCd/9hUTPXVFtVwRix7778wtRsVao+zE062gL79FZleeQ==", - "license": "MIT", - "dependencies": { - "@radix-ui/react-slot": "1.3.0" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, - "node_modules/@radix-ui/react-roving-focus": { - "version": "1.1.14", - "resolved": "https://registry.npmjs.org/@radix-ui/react-roving-focus/-/react-roving-focus-1.1.14.tgz", - "integrity": "sha512-8Qcnx9447tx/aCBgw6Jenfqg4Skq+vqab9mCBmuGNipIS5YXvL275wbKEu7+ICYHIlAPgCduUMJH1XOYewKF6Q==", - "license": "MIT", - "dependencies": { - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-collection": "1.1.11", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-direction": "1.1.2", - "@radix-ui/react-id": "1.1.2", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-callback-ref": "1.1.2", - "@radix-ui/react-use-controllable-state": "1.2.3" - }, - "peerDependencies": { - "@types/react": "*", - "@types/react-dom": "*", - "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", - "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" - }, - "peerDependenciesMeta": { - "@types/react": { - "optional": true - }, - "@types/react-dom": { - "optional": true - } - } - }, "node_modules/@radix-ui/react-scroll-area": { - "version": "1.2.13", - "resolved": "https://registry.npmjs.org/@radix-ui/react-scroll-area/-/react-scroll-area-1.2.13.tgz", - "integrity": "sha512-7tncSubo2G0UY1e8rk+72qe3XRzrGnOLtZQ1PL1KoBfRUNX0NrJT5akb+0kfwSCc3gVR4wdHqyhAQBDpDNOwDw==", + "version": "1.2.18", + "resolved": "https://registry.npmjs.org/@radix-ui/react-scroll-area/-/react-scroll-area-1.2.18.tgz", + "integrity": "sha512-Zn5Cd171wxsO3Dfg8HaW6RifTb9CYTKQJHs/G4+LN1GfmJpaQMZQyQxMprVPHpaz7QY4l9BxK2JwQuzHsXC8nA==", "license": "MIT", "dependencies": { - "@radix-ui/number": "1.1.2", - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-direction": "1.1.2", - "@radix-ui/react-presence": "1.1.6", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-callback-ref": "1.1.2", - "@radix-ui/react-use-layout-effect": "1.1.2" + "@radix-ui/number": "1.1.3", + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -2848,31 +2747,31 @@ } }, "node_modules/@radix-ui/react-select": { - "version": "2.3.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-select/-/react-select-2.3.2.tgz", - "integrity": "sha512-brXD6C/V0fVK0DDbscLVw6LsXrjQ+ay8jdOBaN+tLb4vsHsAMm6Gt6eT77wHX1Eq8GPtD5rJ+RxFtfDozsb4+Q==", + "version": "2.3.7", + "resolved": "https://registry.npmjs.org/@radix-ui/react-select/-/react-select-2.3.7.tgz", + "integrity": "sha512-WFGImkmbzcfxeIwq/+4HvRN0pizBwbwQUED4I13ezQsDdfl38ZntN6TmR8XaSzPBqoCToe8rF75j6NPNDSzhbg==", "license": "MIT", "dependencies": { - "@radix-ui/number": "1.1.2", - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-collection": "1.1.11", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-direction": "1.1.2", - "@radix-ui/react-dismissable-layer": "1.1.14", - "@radix-ui/react-focus-guards": "1.1.4", - "@radix-ui/react-focus-scope": "1.1.11", - "@radix-ui/react-id": "1.1.2", - "@radix-ui/react-popper": "1.3.2", - "@radix-ui/react-portal": "1.1.13", - "@radix-ui/react-presence": "1.1.6", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-slot": "1.3.0", - "@radix-ui/react-use-callback-ref": "1.1.2", - "@radix-ui/react-use-controllable-state": "1.2.3", - "@radix-ui/react-use-layout-effect": "1.1.2", - "@radix-ui/react-use-previous": "1.1.2", - "@radix-ui/react-visually-hidden": "1.2.7", + "@radix-ui/number": "1.1.3", + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-focus-guards": "1.1.6", + "@radix-ui/react-focus-scope": "1.1.16", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-popper": "1.3.7", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-slot": "1.3.3", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-use-previous": "1.1.4", + "@radix-ui/react-visually-hidden": "1.2.11", "aria-hidden": "^1.2.4", "react-remove-scroll": "^2.7.2" }, @@ -2892,12 +2791,12 @@ } }, "node_modules/@radix-ui/react-separator": { - "version": "1.1.11", - "resolved": "https://registry.npmjs.org/@radix-ui/react-separator/-/react-separator-1.1.11.tgz", - "integrity": "sha512-jRhe86+8PF7VZ1u14eOWVOuh2BuAhALg/FT1VcMC4OHedMTRUazDnDlKTt+yxo5cRNKHMfmvZ4sSQtWDeMV4CQ==", + "version": "1.1.15", + "resolved": "https://registry.npmjs.org/@radix-ui/react-separator/-/react-separator-1.1.15.tgz", + "integrity": "sha512-jOLO4lssEzWpoDu7G+Ze4VjwMRUBt291pnZD0gmalREZipnTX3wadQo7Fy48GCTfe14/YRN6rw/rOJqrE85Wxw==", "license": "MIT", "dependencies": { - "@radix-ui/react-primitive": "2.1.7" + "@radix-ui/react-primitive": "2.1.10" }, "peerDependencies": { "@types/react": "*", @@ -2915,22 +2814,22 @@ } }, "node_modules/@radix-ui/react-slider": { - "version": "1.4.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-slider/-/react-slider-1.4.2.tgz", - "integrity": "sha512-qt5C1ppJz66aUDrH1VccjPrq7aFchK0wBrn6xsxlCHNUyE57dRRQ7lp1QFpF7OscMexZF8MCGBTVBlENHPkNiA==", + "version": "1.4.7", + "resolved": "https://registry.npmjs.org/@radix-ui/react-slider/-/react-slider-1.4.7.tgz", + "integrity": "sha512-mTSLf1GC/C0moWjTbvCM6Qn/gBjvlFt1azuWF2v7MN5C3Zq2U2J2lN3ZEYkpujuOU5Ro7A28wkviSxaKnG0BYg==", "license": "MIT", "dependencies": { - "@radix-ui/number": "1.1.2", - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-collection": "1.1.11", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-direction": "1.1.2", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-controllable-state": "1.2.3", - "@radix-ui/react-use-layout-effect": "1.1.2", - "@radix-ui/react-use-previous": "1.1.2", - "@radix-ui/react-use-size": "1.1.2" + "@radix-ui/number": "1.1.3", + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-use-previous": "1.1.4", + "@radix-ui/react-use-size": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -2948,12 +2847,12 @@ } }, "node_modules/@radix-ui/react-slot": { - "version": "1.3.0", - "resolved": "https://registry.npmjs.org/@radix-ui/react-slot/-/react-slot-1.3.0.tgz", - "integrity": "sha512-MojKku4U/miO8Av4Dkb+ctMAQx7JmY96LmtDQlAarCRtd7rN52QCSzBF+XAvr5S6coSVj9HEPBgHAHKEJVk/WA==", + "version": "1.3.3", + "resolved": "https://registry.npmjs.org/@radix-ui/react-slot/-/react-slot-1.3.3.tgz", + "integrity": "sha512-qx7oqnYbxnK9kYI9m317qmFmEgo6ywqWvbTogdj7cL9p3/yx4M48p7Rnw5z3H890cL/ow/EeWJsuTykeZVXP5Q==", "license": "MIT", "dependencies": { - "@radix-ui/react-compose-refs": "1.1.3" + "@radix-ui/react-compose-refs": "1.1.5" }, "peerDependencies": { "@types/react": "*", @@ -2966,18 +2865,17 @@ } }, "node_modules/@radix-ui/react-switch": { - "version": "1.3.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-switch/-/react-switch-1.3.2.tgz", - "integrity": "sha512-tgRBI3DdNwAJYE4BBZyZcz/HRRCvAsPkRvG1wvKc+41tBGMxPn/a87T/wikXAvyDypNQ9kaZwHbeZe+veHCGpA==", + "version": "1.3.7", + "resolved": "https://registry.npmjs.org/@radix-ui/react-switch/-/react-switch-1.3.7.tgz", + "integrity": "sha512-48tB/4dn2UVLBCYhTu9AuR63IHl73l/qLbLgxd86noTUor4/K4LFDAcYjK+isP5313qxaFpjPVogE7+Y0/V3Kw==", "license": "MIT", "dependencies": { - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-compose-refs": "1.1.3", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-use-controllable-state": "1.2.3", - "@radix-ui/react-use-previous": "1.1.2", - "@radix-ui/react-use-size": "1.1.2" + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-size": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -2995,19 +2893,171 @@ } }, "node_modules/@radix-ui/react-tabs": { - "version": "1.1.16", - "resolved": "https://registry.npmjs.org/@radix-ui/react-tabs/-/react-tabs-1.1.16.tgz", - "integrity": "sha512-v3Ab2l7z6U7tRB4xA0IyKdq0OsqaO1o9ZjsIEoKKnSZ/l96mZz8aCTX0NCXw+YVHJXr8Km4d+Mn6/Q8YjXa+gw==", + "version": "1.1.21", + "resolved": "https://registry.npmjs.org/@radix-ui/react-tabs/-/react-tabs-1.1.21.tgz", + "integrity": "sha512-UKxJlZid7FVtsk/WTxj4i4uSEgj2Au+KBbS7SQyTlzMhhn+86Cz3tISZdTa87bfEfcuvZezf2ZsxD4xuEKtkog==", "license": "MIT", "dependencies": { - "@radix-ui/primitive": "1.1.4", - "@radix-ui/react-context": "1.1.4", - "@radix-ui/react-direction": "1.1.2", - "@radix-ui/react-id": "1.1.2", - "@radix-ui/react-presence": "1.1.6", - "@radix-ui/react-primitive": "2.1.7", - "@radix-ui/react-roving-focus": "1.1.14", - "@radix-ui/react-use-controllable-state": "1.2.3" + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-toast": { + "version": "1.2.23", + "resolved": "https://registry.npmjs.org/@radix-ui/react-toast/-/react-toast-1.2.23.tgz", + "integrity": "sha512-ofhyAsYaocRGOs/n0XWdUOSVzEAG6BfrMVM8z0c0kLEWY38w/0WuMFPTJP/HVaZPYkMvHZoKIIhNcjbTCBILPg==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-visually-hidden": "1.2.11" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-toggle": { + "version": "1.1.18", + "resolved": "https://registry.npmjs.org/@radix-ui/react-toggle/-/react-toggle-1.1.18.tgz", + "integrity": "sha512-7lonPlKfSacd20GlOBx2ltuVKz9oqWYZz+oMQyOltw6t1y2nyftj2ZmwwUHYn49kqfDWcp8dNZm5NgV+5Z+mug==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-toggle-group": { + "version": "1.1.19", + "resolved": "https://registry.npmjs.org/@radix-ui/react-toggle-group/-/react-toggle-group-1.1.19.tgz", + "integrity": "sha512-OtnwuSVjd1Ofi+AdnvhsjQdyuhCDwYs1w9RyB5BN/OavXOVQo42SYqQjwUnbPnaiPFBpQ9aX70dWeee+v2oBLA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-toggle": "1.1.18", + "@radix-ui/react-use-controllable-state": "1.2.6" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-toolbar": { + "version": "1.1.19", + "resolved": "https://registry.npmjs.org/@radix-ui/react-toolbar/-/react-toolbar-1.1.19.tgz", + "integrity": "sha512-Ph0IvtYw4VB12ZnZg+YtrGs8yJQsnizwo/zu0R4Y/nWugtJzA7Pg1eWeuDR9+LSqn+xjamss+UOSOJJJ4gx8jw==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-separator": "1.1.15", + "@radix-ui/react-toggle-group": "1.1.19" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-tooltip": { + "version": "1.2.16", + "resolved": "https://registry.npmjs.org/@radix-ui/react-tooltip/-/react-tooltip-1.2.16.tgz", + "integrity": "sha512-6EamKFRRnlpdadndbZ6LMwycfwkwPte1B42hs6QA0gYhjaOKqW4PZ4pjaW9UrlDX5eVt/OjncE7BFTPL5nmZhg==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-id": "1.1.4", + "@radix-ui/react-popper": "1.3.7", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-slot": "1.3.3", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-visually-hidden": "1.2.11" }, "peerDependencies": { "@types/react": "*", @@ -3025,9 +3075,9 @@ } }, "node_modules/@radix-ui/react-use-callback-ref": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-callback-ref/-/react-use-callback-ref-1.1.2.tgz", - "integrity": "sha512-xCso9j1/u8sEgP1RNHjFrXJLApL8LiqOkI1R4ywuN00rxWdYg4oQXuwKLS3i0j5NWLromUD27/4nlxj2UFVvIw==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-callback-ref/-/react-use-callback-ref-1.1.4.tgz", + "integrity": "sha512-R6OUY2e2fA6Yn6s+VSx5KBV6Nx8LQEhu+cz7LCej18rQ1HLyg9PSC9jP/ZNx0o6FAIK9c0F1kHylzSxKsdlkrQ==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -3040,13 +3090,14 @@ } }, "node_modules/@radix-ui/react-use-controllable-state": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-controllable-state/-/react-use-controllable-state-1.2.3.tgz", - "integrity": "sha512-PLzC90MS+ReootmjC597dvopoelpZ8Q61HJkDXZSExitIq7PL55vHNnesAHwguHK0aPfBnpdNzQtv1uliaqQrA==", + "version": "1.2.6", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-controllable-state/-/react-use-controllable-state-1.2.6.tgz", + "integrity": "sha512-uEQJGT97ZA/TgP/Hydw47lHu+/vQj6z/0jA+WeTbK1o9Rx45GImjpD0tc3W5ad3D6XTSR6e1yEO0FvGq6WQfVQ==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-effect-event": "0.0.3", - "@radix-ui/react-use-layout-effect": "1.1.2" + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-use-effect-event": "0.0.5", + "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -3059,12 +3110,30 @@ } }, "node_modules/@radix-ui/react-use-effect-event": { - "version": "0.0.3", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-effect-event/-/react-use-effect-event-0.0.3.tgz", - "integrity": "sha512-6c8ZqvPTWILEKnyVkP53EGRCcpnJiKTC21sS/6R1GF5xKyHJJWQEPfkqlcgUkdRQivd6tb23abUwe4ngWmY0JA==", + "version": "0.0.5", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-effect-event/-/react-use-effect-event-0.0.5.tgz", + "integrity": "sha512-7cshFL8HGS/7HEiHH+9kL9HBwp2sa9yX18Knwek6KYWmXwM7pegMgta2AXMQKI+rq3JnfSj9x8wYqFMTdG1Jgg==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.2" + "@radix-ui/react-use-layout-effect": "1.1.4" + }, + "peerDependencies": { + "@types/react": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + } + } + }, + "node_modules/@radix-ui/react-use-escape-keydown": { + "version": "1.1.5", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-escape-keydown/-/react-use-escape-keydown-1.1.5.tgz", + "integrity": "sha512-ge3ipobwSXTj4JyVtswQ7qZj0ZHdtbGuOno/LrgAAeSxtsJ6Vs4Gz5IkPH2bmqpjcLUFoqGhA/mueuIf63UXlA==", + "license": "MIT", + "dependencies": { + "@radix-ui/react-use-callback-ref": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -3092,9 +3161,9 @@ } }, "node_modules/@radix-ui/react-use-layout-effect": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-layout-effect/-/react-use-layout-effect-1.1.2.tgz", - "integrity": "sha512-jrBWOxZITuGcnjRCM2t2U5ZPkCLxD+Ym6DjfssS5haTj2iiak/DOb64JeN6OdLfLgptb6/e2kKR+ZuTrGoZTPA==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-layout-effect/-/react-use-layout-effect-1.1.4.tgz", + "integrity": "sha512-K20DkRkUwDnxEYMBPcg3Y6voLkEy5p5QQmszZgLngKKiC7dzBR/aEuK3w1qlx2JWDUNH6FluahYdgR3BP+QbYw==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -3107,9 +3176,9 @@ } }, "node_modules/@radix-ui/react-use-previous": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-previous/-/react-use-previous-1.1.2.tgz", - "integrity": "sha512-IGBQPtRFdhN6MQ8dbegVmBq1LVZluya3F1jWY+puIcQC3MHctRwTDSBWCkL/3ZcnMJLTMJ++Z+ktmvg0F89iCw==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-previous/-/react-use-previous-1.1.4.tgz", + "integrity": "sha512-XoSLhbRbqxFtgJoi2fNHA3C6pDlY34x508vUpUGoFZfvePfHXHbE1lC4FYFMnJWgiCRroSTw6fOsXQoVS9RwZg==", "license": "MIT", "peerDependencies": { "@types/react": "*", @@ -3122,12 +3191,12 @@ } }, "node_modules/@radix-ui/react-use-rect": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-rect/-/react-use-rect-1.1.2.tgz", - "integrity": "sha512-d8a+bBY/FxikNPlgJJoaBHZX+zKVbWHYJGTLnLvveQgFSTntkGdEKv3JDtHrMS0DNYpllz2nRsTLGLKYttbpmw==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-rect/-/react-use-rect-1.1.4.tgz", + "integrity": "sha512-cSOCh6JlkmfjLyNcLiu2nB4v+nm+dkZ+Q5KHWk/soo4U7ZLiEQFKHK9/YmtBHjfCEaU43IBKQOc4/uJmCaiCTQ==", "license": "MIT", "dependencies": { - "@radix-ui/rect": "1.1.2" + "@radix-ui/rect": "1.1.3" }, "peerDependencies": { "@types/react": "*", @@ -3140,12 +3209,12 @@ } }, "node_modules/@radix-ui/react-use-size": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/react-use-size/-/react-use-size-1.1.2.tgz", - "integrity": "sha512-giWQp+4mxjBPt4KZ0MmyuykFNWfbDxKt4x+fPkRYmgRFJSbCZFzUglvMb/Kjn38tm10YP4ufiQZDx3zna4LU6w==", + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/@radix-ui/react-use-size/-/react-use-size-1.1.4.tgz", + "integrity": "sha512-D3anSY15EJoxrihpsXI6SMrmmonnQtR2ni7arO+Lfdg3O95b9hNXxONk8jA5C8ANdF/h5HMAxejgs8PWJ6rlhw==", "license": "MIT", "dependencies": { - "@radix-ui/react-use-layout-effect": "1.1.2" + "@radix-ui/react-use-layout-effect": "1.1.4" }, "peerDependencies": { "@types/react": "*", @@ -3158,12 +3227,12 @@ } }, "node_modules/@radix-ui/react-visually-hidden": { - "version": "1.2.7", - "resolved": "https://registry.npmjs.org/@radix-ui/react-visually-hidden/-/react-visually-hidden-1.2.7.tgz", - "integrity": "sha512-1wNZBggTDK3GRuuQ6nP4k2yi7a6l7I5qbMPbZcRsrGsGVead/f/d5FhEzUvqFs0bcrDLx7n1zKQ3JvLR6whaaw==", + "version": "1.2.11", + "resolved": "https://registry.npmjs.org/@radix-ui/react-visually-hidden/-/react-visually-hidden-1.2.11.tgz", + "integrity": "sha512-NFS86RYYZb4/exihaESBGOpMJFz8MGLAfu3mOBSGByVnVPC9JPASfYubxd/8KbkQK0sYAv8lVQDEQukDX/qXvQ==", "license": "MIT", "dependencies": { - "@radix-ui/react-primitive": "2.1.7" + "@radix-ui/react-primitive": "2.1.10" }, "peerDependencies": { "@types/react": "*", @@ -3181,9 +3250,9 @@ } }, "node_modules/@radix-ui/rect": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/@radix-ui/rect/-/rect-1.1.2.tgz", - "integrity": "sha512-xnXE7wG13PI+cxieVssYXlQJuYVRhH9NBoxt3KNwzghDIA69GMm7d4wXRouHIYjE+KvS6U/MsMO73NdS2MH9ZA==", + "version": "1.1.3", + "resolved": "https://registry.npmjs.org/@radix-ui/rect/-/rect-1.1.3.tgz", + "integrity": "sha512-JtyZR+mqgBibTo8xea3B6ZRmzZiM/YeVBtUkas6zMuXjAlfIFIW2FgqeM9eLyvEaYX66vr6DJMK+4U6LV0KhNw==", "license": "MIT" }, "node_modules/@rolldown/pluginutils": { @@ -5822,6 +5891,83 @@ "node": ">=6" } }, + "node_modules/radix-ui": { + "version": "1.6.7", + "resolved": "https://registry.npmjs.org/radix-ui/-/radix-ui-1.6.7.tgz", + "integrity": "sha512-QBdhh1arIEUvPC0dQ5+nwWAxt7+N+oP/9jPwjJkGFoSk/sqxg32gJtSXGtFh8frAIcS6oC9cx2Q+7KYCQLOAeA==", + "license": "MIT", + "dependencies": { + "@radix-ui/primitive": "1.1.7", + "@radix-ui/react-accessible-icon": "1.1.15", + "@radix-ui/react-accordion": "1.2.20", + "@radix-ui/react-alert-dialog": "1.1.23", + "@radix-ui/react-arrow": "1.1.15", + "@radix-ui/react-aspect-ratio": "1.1.15", + "@radix-ui/react-avatar": "1.2.6", + "@radix-ui/react-checkbox": "1.3.11", + "@radix-ui/react-collapsible": "1.1.20", + "@radix-ui/react-collection": "1.1.15", + "@radix-ui/react-compose-refs": "1.1.5", + "@radix-ui/react-context": "1.2.2", + "@radix-ui/react-context-menu": "2.3.7", + "@radix-ui/react-dialog": "1.1.23", + "@radix-ui/react-direction": "1.1.4", + "@radix-ui/react-dismissable-layer": "1.1.19", + "@radix-ui/react-dropdown-menu": "2.1.24", + "@radix-ui/react-focus-guards": "1.1.6", + "@radix-ui/react-focus-scope": "1.1.16", + "@radix-ui/react-form": "0.1.16", + "@radix-ui/react-hover-card": "1.1.23", + "@radix-ui/react-label": "2.1.15", + "@radix-ui/react-menu": "2.1.24", + "@radix-ui/react-menubar": "1.1.24", + "@radix-ui/react-navigation-menu": "1.2.22", + "@radix-ui/react-one-time-password-field": "0.1.16", + "@radix-ui/react-password-toggle-field": "0.1.11", + "@radix-ui/react-popover": "1.1.23", + "@radix-ui/react-popper": "1.3.7", + "@radix-ui/react-portal": "1.1.17", + "@radix-ui/react-presence": "1.1.10", + "@radix-ui/react-primitive": "2.1.10", + "@radix-ui/react-progress": "1.1.16", + "@radix-ui/react-radio-group": "1.4.7", + "@radix-ui/react-roving-focus": "1.1.19", + "@radix-ui/react-scroll-area": "1.2.18", + "@radix-ui/react-select": "2.3.7", + "@radix-ui/react-separator": "1.1.15", + "@radix-ui/react-slider": "1.4.7", + "@radix-ui/react-slot": "1.3.3", + "@radix-ui/react-switch": "1.3.7", + "@radix-ui/react-tabs": "1.1.21", + "@radix-ui/react-toast": "1.2.23", + "@radix-ui/react-toggle": "1.1.18", + "@radix-ui/react-toggle-group": "1.1.19", + "@radix-ui/react-toolbar": "1.1.19", + "@radix-ui/react-tooltip": "1.2.16", + "@radix-ui/react-use-callback-ref": "1.1.4", + "@radix-ui/react-use-controllable-state": "1.2.6", + "@radix-ui/react-use-effect-event": "0.0.5", + "@radix-ui/react-use-escape-keydown": "1.1.5", + "@radix-ui/react-use-is-hydrated": "0.1.3", + "@radix-ui/react-use-layout-effect": "1.1.4", + "@radix-ui/react-use-size": "1.1.4", + "@radix-ui/react-visually-hidden": "1.2.11" + }, + "peerDependencies": { + "@types/react": "*", + "@types/react-dom": "*", + "react": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc", + "react-dom": "^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "@types/react-dom": { + "optional": true + } + } + }, "node_modules/react": { "version": "19.2.7", "resolved": "https://registry.npmjs.org/react/-/react-19.2.7.tgz", diff --git a/apps/fastvideo_studio/package.json b/apps/fastvideo_studio/package.json index 830dca354..82f79690e 100644 --- a/apps/fastvideo_studio/package.json +++ b/apps/fastvideo_studio/package.json @@ -17,20 +17,11 @@ "start:all": "concurrently --kill-others-on-fail \"npm:start:api\" \"npm:start:web\"" }, "dependencies": { - "@radix-ui/react-dialog": "^1.1.0", - "@radix-ui/react-dropdown-menu": "^2.1.24", - "@radix-ui/react-label": "^2.1.8", - "@radix-ui/react-scroll-area": "^1.2.10", - "@radix-ui/react-select": "^2.2.6", - "@radix-ui/react-separator": "^1.1.8", - "@radix-ui/react-slider": "^1.2.0", - "@radix-ui/react-slot": "^1.2.4", - "@radix-ui/react-switch": "^1.1.0", - "@radix-ui/react-tabs": "^1.1.0", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", "lucide-react": "^0.577.0", "next": "15.5.18", + "radix-ui": "^1.6.7", "react": "^19.1.0", "react-dom": "^19.1.0", "sonner": "^2.0.7", diff --git a/apps/fastvideo_studio/src/components/jobs/CreateJobButton.test.tsx b/apps/fastvideo_studio/src/components/jobs/CreateJobButton.test.tsx index a1636ce3f..86e5e178c 100644 --- a/apps/fastvideo_studio/src/components/jobs/CreateJobButton.test.tsx +++ b/apps/fastvideo_studio/src/components/jobs/CreateJobButton.test.tsx @@ -1,24 +1,26 @@ -import { render, screen } from '@testing-library/react'; +import { render, screen, waitFor, within } from '@testing-library/react'; import userEvent from '@testing-library/user-event'; -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import CreateJobButton from './CreateJobButton'; +import { getDatasets, getModels } from '@/lib/api'; -vi.mock('./CreateJobModal', () => ({ - default: ({ - isOpen, - workloadType, - }: { - isOpen: boolean; - workloadType: string; - }) => - isOpen ? ( -
    - Create job form -
    - ) : null, +vi.mock('@/lib/api', () => ({ + createJob: vi.fn(), + getModels: vi.fn(), + getDatasets: vi.fn(), + uploadImage: vi.fn(), + getSettings: vi.fn(), + updateSettings: vi.fn(), })); +beforeEach(() => { + vi.mocked(getModels).mockResolvedValue([ + { id: 'wan/t2v-1.3b', label: 'Wan T2V' }, + ]); + vi.mocked(getDatasets).mockResolvedValue([]); +}); + describe('CreateJobButton', () => { it('opens the workload menu on click and selects an item', async () => { const user = userEvent.setup(); @@ -27,10 +29,9 @@ describe('CreateJobButton', () => { await user.click(screen.getByRole('button', { name: 'Create Job' })); await user.click(screen.getByRole('menuitem', { name: /I2V/i })); - expect(screen.getByRole('dialog')).toHaveAttribute( - 'data-workload-type', - 'i2v', - ); + expect( + screen.getByRole('dialog', { name: 'New Inference Job (I2V)' }), + ).toBeInTheDocument(); }); it('opens and operates the workload menu from the keyboard', async () => { @@ -45,9 +46,42 @@ describe('CreateJobButton', () => { expect(firstItem).toHaveFocus(); await user.keyboard('{Enter}'); - expect(screen.getByRole('dialog')).toHaveAttribute( - 'data-workload-type', - 't2v', + expect( + screen.getByRole('dialog', { name: 'New Inference Job (T2V)' }), + ).toBeInTheDocument(); + await user.keyboard('{Escape}'); + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), ); + await waitFor(() => + expect(document.body.style.pointerEvents).not.toBe('none'), + ); + expect(trigger).toHaveFocus(); }); + + it.each(['inference', 'finetuning', 'distillation'] as const)( + 'restores page interaction after closing the real %s dialog', + async (jobType) => { + const user = userEvent.setup(); + render(); + const trigger = screen.getByRole('button', { name: 'Create Job' }); + + // Keep the real Dialog mounted: mocking it hides conflicting Radix layers. + for (let attempt = 0; attempt < 2; attempt++) { + await user.click(trigger); + await user.click(screen.getAllByRole('menuitem')[0]); + const dialog = screen.getByRole('dialog'); + await user.click( + within(dialog).getByRole('button', { name: 'Close' }), + ); + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), + ); + await waitFor(() => + expect(document.body.style.pointerEvents).not.toBe('none'), + ); + expect(trigger).toHaveFocus(); + } + }, + ); }); diff --git a/apps/fastvideo_studio/src/components/jobs/CreateJobButton.tsx b/apps/fastvideo_studio/src/components/jobs/CreateJobButton.tsx index 995423030..12ce8d083 100644 --- a/apps/fastvideo_studio/src/components/jobs/CreateJobButton.tsx +++ b/apps/fastvideo_studio/src/components/jobs/CreateJobButton.tsx @@ -2,7 +2,7 @@ import * as React from 'react'; import { ChevronDown } from 'lucide-react'; -import * as DropdownMenu from '@radix-ui/react-dropdown-menu'; +import { DropdownMenu } from 'radix-ui'; import CreateJobModal from '@/components/jobs/CreateJobModal'; import { Button } from '@/components/ui/button'; @@ -16,6 +16,7 @@ interface CreateJobButtonProps { export default function CreateJobButton({ jobType }: CreateJobButtonProps) { const options = WORKLOAD_OPTIONS[jobType] ?? []; + const triggerRef = React.useRef(null); const [modalOpen, setModalOpen] = React.useState(false); const [workloadType, setWorkloadType] = React.useState( @@ -36,7 +37,7 @@ export default function CreateJobButton({ jobType }: CreateJobButtonProps) { <> - @@ -66,6 +67,11 @@ export default function CreateJobButton({ jobType }: CreateJobButtonProps) { setModalOpen(false)} + onCloseAutoFocus={(event) => { + // This dialog opens from a menu item, so it has no DialogTrigger. + event.preventDefault(); + triggerRef.current?.focus(); + }} onSuccess={handleSuccess} jobType={jobType} workloadType={workloadType} diff --git a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx index 77b7233ad..cc96ebc29 100644 --- a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx +++ b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx @@ -55,6 +55,9 @@ import { jobToFormFields, type JobLike } from '@/lib/jobToFields'; export interface CreateJobModalProps { isOpen: boolean; onClose: () => void; + onCloseAutoFocus?: React.ComponentProps< + typeof DialogContent + >['onCloseAutoFocus']; onSuccess: () => void; jobType: JobType; workloadType: string; @@ -67,6 +70,7 @@ export interface CreateJobModalProps { export default function CreateJobModal({ isOpen, onClose, + onCloseAutoFocus, onSuccess, jobType, workloadType, @@ -644,6 +648,7 @@ export default function CreateJobModal({ > { if (isSubmitting) e.preventDefault(); }} diff --git a/apps/fastvideo_studio/src/components/ui/accessibility.test.tsx b/apps/fastvideo_studio/src/components/ui/accessibility.test.tsx index 1dd79f961..d942701eb 100644 --- a/apps/fastvideo_studio/src/components/ui/accessibility.test.tsx +++ b/apps/fastvideo_studio/src/components/ui/accessibility.test.tsx @@ -1,5 +1,7 @@ +import * as React from 'react'; import { render, screen } from '@testing-library/react'; -import { describe, expect, it } from 'vitest'; +import userEvent from '@testing-library/user-event'; +import { describe, expect, it, vi } from 'vitest'; import { Button } from './button'; import { Input } from './input'; @@ -8,6 +10,24 @@ import { Slider } from './slider'; import { Switch } from './switch'; describe('shared control accessibility', () => { + it('forwards refs and click handlers to the asChild button', async () => { + const user = userEvent.setup(); + const ref = React.createRef(); + const onClick = vi.fn(); + render( + + , + ); + + const button = screen.getByRole('button', { name: 'Slotted action' }); + expect(screen.getAllByRole('button')).toHaveLength(1); + expect(ref.current).toBe(button); + await user.click(button); + expect(onClick).toHaveBeenCalledTimes(1); + expect(button).toHaveFocus(); + }); + it('keeps button, input, and select targets at least 44px tall', () => { render( <> diff --git a/apps/fastvideo_studio/src/components/ui/button.tsx b/apps/fastvideo_studio/src/components/ui/button.tsx index aeea479d3..94fdefa5d 100644 --- a/apps/fastvideo_studio/src/components/ui/button.tsx +++ b/apps/fastvideo_studio/src/components/ui/button.tsx @@ -1,7 +1,7 @@ "use client"; import * as React from "react"; -import { Slot } from "@radix-ui/react-slot"; +import { Slot } from "radix-ui"; import { cva, type VariantProps } from "class-variance-authority"; import { cn } from "@/lib/utils"; @@ -37,7 +37,7 @@ export interface ButtonProps extends React.ButtonHTMLAttributes(({ className, variant, size, asChild = false, ...props }, ref) => { - const Comp = asChild ? Slot : "button"; + const Comp = asChild ? Slot.Root : "button"; return ; }); Button.displayName = "Button"; diff --git a/apps/fastvideo_studio/src/components/ui/dialog.tsx b/apps/fastvideo_studio/src/components/ui/dialog.tsx index 88eee674b..3abfba309 100644 --- a/apps/fastvideo_studio/src/components/ui/dialog.tsx +++ b/apps/fastvideo_studio/src/components/ui/dialog.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as DialogPrimitive from '@radix-ui/react-dialog'; +import { Dialog as DialogPrimitive } from 'radix-ui'; import { X } from 'lucide-react'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/label.tsx b/apps/fastvideo_studio/src/components/ui/label.tsx index f41d0d9fb..da2111a91 100644 --- a/apps/fastvideo_studio/src/components/ui/label.tsx +++ b/apps/fastvideo_studio/src/components/ui/label.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as LabelPrimitive from '@radix-ui/react-label'; +import { Label as LabelPrimitive } from 'radix-ui'; import { cva, type VariantProps } from 'class-variance-authority'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/scroll-area.tsx b/apps/fastvideo_studio/src/components/ui/scroll-area.tsx index afa0efb90..dfd432ace 100644 --- a/apps/fastvideo_studio/src/components/ui/scroll-area.tsx +++ b/apps/fastvideo_studio/src/components/ui/scroll-area.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as ScrollAreaPrimitive from '@radix-ui/react-scroll-area'; +import { ScrollArea as ScrollAreaPrimitive } from 'radix-ui'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/select.tsx b/apps/fastvideo_studio/src/components/ui/select.tsx index 74939f9f5..ba4bd8706 100644 --- a/apps/fastvideo_studio/src/components/ui/select.tsx +++ b/apps/fastvideo_studio/src/components/ui/select.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as SelectPrimitive from '@radix-ui/react-select'; +import { Select as SelectPrimitive } from 'radix-ui'; import { Check, ChevronDown, ChevronUp } from 'lucide-react'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/separator.tsx b/apps/fastvideo_studio/src/components/ui/separator.tsx index bc7a78fd5..4c8cbf93d 100644 --- a/apps/fastvideo_studio/src/components/ui/separator.tsx +++ b/apps/fastvideo_studio/src/components/ui/separator.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as SeparatorPrimitive from '@radix-ui/react-separator'; +import { Separator as SeparatorPrimitive } from 'radix-ui'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/slider.tsx b/apps/fastvideo_studio/src/components/ui/slider.tsx index 7ab5e2ada..f6e9e7147 100644 --- a/apps/fastvideo_studio/src/components/ui/slider.tsx +++ b/apps/fastvideo_studio/src/components/ui/slider.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as SliderPrimitive from '@radix-ui/react-slider'; +import { Slider as SliderPrimitive } from 'radix-ui'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/switch.tsx b/apps/fastvideo_studio/src/components/ui/switch.tsx index b80ce8bc1..730ea7e90 100644 --- a/apps/fastvideo_studio/src/components/ui/switch.tsx +++ b/apps/fastvideo_studio/src/components/ui/switch.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as SwitchPrimitives from '@radix-ui/react-switch'; +import { Switch as SwitchPrimitives } from 'radix-ui'; import { cn } from '@/lib/utils'; diff --git a/apps/fastvideo_studio/src/components/ui/tabs.tsx b/apps/fastvideo_studio/src/components/ui/tabs.tsx index 1d4c1e82a..2d1d766af 100644 --- a/apps/fastvideo_studio/src/components/ui/tabs.tsx +++ b/apps/fastvideo_studio/src/components/ui/tabs.tsx @@ -1,7 +1,7 @@ 'use client'; import * as React from 'react'; -import * as TabsPrimitive from '@radix-ui/react-tabs'; +import { Tabs as TabsPrimitive } from 'radix-ui'; import { cn } from '@/lib/utils'; From 361f919c882d058a13cb6ca5ed75c97a0e68f427 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Wed, 23 Sep 2026 11:06:55 -0700 Subject: [PATCH 022/151] [feat]: add a narrow FastH3 8-step MLX ladder Keep the four-step uniform AdaLN cache, and use explicit DMD rungs only when the snapshot declares its own scheduler shifts. --- docs/getting_started/installation/mps.md | 17 ++ docs/inference/fasth3-distilled.md | 10 ++ examples/inference/basic/mlx_fasth3.py | 27 ++- examples/inference/basic/mlx_fasth3_8step.py | 43 +++++ fastvideo/mlx_runtime/minimax_h3.py | 155 +++++++++++++++++- .../mlx_runtime/minimax_h3_conditioner.py | 16 +- fastvideo/mlx_runtime/minimax_h3_pipeline.py | 131 ++++++++------- .../mlx/test_mlx_minimax_h3_8step_schedule.py | 137 ++++++++++++++++ .../convert_minimax_h3_mlx.py | 47 +++--- 9 files changed, 479 insertions(+), 104 deletions(-) create mode 100644 examples/inference/basic/mlx_fasth3_8step.py create mode 100644 fastvideo/tests/mlx/test_mlx_minimax_h3_8step_schedule.py diff --git a/docs/getting_started/installation/mps.md b/docs/getting_started/installation/mps.md index 3a830a8c0..2d7f3ac4e 100644 --- a/docs/getting_started/installation/mps.md +++ b/docs/getting_started/installation/mps.md @@ -225,6 +225,23 @@ python examples/inference/basic/mlx_fasth3.py \ --output-path ./outputs/fasth3_int6_vsa_720p.mp4 ``` +FastH3 8-Step V2 is a separate checkpoint. Keep the four-step command above +for the preview. For the eight-forward model, convert its transformer with the +same script (the snapshot's `fastvideo_inference.json` declares the trained +shifts, so the AdaLN cache follows those rungs) and run: + +```bash +python examples/inference/basic/mlx_fasth3_8step.py \ + --model-root ./FastH3-8-Step-V2 \ + --mlx-checkpoint ./FastH3-8-Step-V2-MLX/int6 \ + --prompt "(S1) A presenter says [English] Fast H3 is amazing." \ + --height 480 --width 832 --num-frames 124 --seed 2026 \ + --output-path ./outputs/fasth3_8step_int6.mp4 +``` + +`--steps` may be 8 or 9. Both run the eight trained forwards. VAE, audio VAE, +text encoder, and tokenizer can be the ones already used by the preview. + `--vsa-impl auto` uses the chunked gather+SDPA **reference** path. `--vsa-impl simd` is an opt-in SIMD-group kernel (tile 64, head dim 128) that falls back to reference on unsupported shapes. It is not the default. diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index e52c33d59..0a206914f 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -73,3 +73,13 @@ unset: H3 has separate video and audio shifts, not one shared shift. This documents execution support for the published checkpoint. It is not a quality claim: compare video/audio output against base MiniMax-H3 on your own prompts before adopting it. + +## Apple Silicon + +`examples/inference/basic/mlx_fasth3.py` stays on the four-step preview and its +uniform AdaLN cache. `mlx_fasth3_8step.py` is the eight-forward MLX recipe. It +reads the same `fastvideo_inference.json` rungs and shifts, and it expects an +MLX DiT whose AdaLN cache was converted from that contract. Reuse the preview +snapshot's VAE, audio VAE, text encoder, and tokenizer; only the DiT and the +sidecar change. Rank-reduced AdaLN checkpoints are unchanged and are not +produced by the MLX converter. diff --git a/examples/inference/basic/mlx_fasth3.py b/examples/inference/basic/mlx_fasth3.py index c85d72474..62ec602ea 100644 --- a/examples/inference/basic/mlx_fasth3.py +++ b/examples/inference/basic/mlx_fasth3.py @@ -28,7 +28,8 @@ detail budget and reads softer than a native-resolution render, so it stays off by default. This entrypoint currently supports text-to-video-with-audio only. It does not -yet wire FL2VA, Ref2VA, or two-pass refinement. +yet wire FL2VA, Ref2VA, or two-pass refinement. The four-step preview stays +the default. FastH3 8-Step V2 uses ``mlx_fasth3_8step.py``. VSA is off by default; existing dense MLX checkpoints remain supported. H3 uses fused MLX RMSNorm, which can change BF16 rounding compared with the @@ -51,7 +52,7 @@ def _dense_layers(value: str) -> tuple[int, ...]: return layers -def parse_args() -> argparse.Namespace: +def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--model-root", type=Path, default=Path.home() / "models/FastH3-Preview-v0.2", @@ -164,11 +165,20 @@ def parse_args() -> argparse.Namespace: "(simd is opt-in)" ), ) - return parser.parse_args() + parser.add_argument("--vae-tile-height", type=int, default=256, + help="overlapping video-VAE tile height in pixels") + parser.add_argument("--vae-tile-width", type=int, default=256, + help="overlapping video-VAE tile width in pixels") + parser.add_argument("--inter-step-cooldown-s", type=float, default=0.0, + help="optional pause between denoise steps") + return parser -def main() -> None: - args = parse_args() +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + return build_parser().parse_args(argv) + + +def run(args: argparse.Namespace) -> None: from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline pipeline = MiniMaxH3MLXPipeline( @@ -189,6 +199,9 @@ def main() -> None: seed=args.seed, num_steps=args.steps, tiled_video_decode=args.tiled_video_decode, + vae_tile_height=args.vae_tile_height, + vae_tile_width=args.vae_tile_width, + inter_step_cooldown_s=args.inter_step_cooldown_s, fast=args.fast, fast_factor=args.fast_factor, fast_sharpen=args.fast_sharpen, @@ -215,5 +228,9 @@ def main() -> None: }, indent=2)) +def main() -> None: + run(parse_args()) + + if __name__ == "__main__": main() diff --git a/examples/inference/basic/mlx_fasth3_8step.py b/examples/inference/basic/mlx_fasth3_8step.py new file mode 100644 index 000000000..8a8243662 --- /dev/null +++ b/examples/inference/basic/mlx_fasth3_8step.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Eight-forward FastH3 8-Step V2 generation on Apple Silicon. + +``FastVideo/FastVideo-FastH3-8-Step-V2`` trains eight DMD forwards at video +shift 10 and audio shift 3, VSA sparsity 0.8, and 64-token tiles. Point +``--model-root`` at a snapshot that contains that checkpoint's +``fastvideo_inference.json`` plus the shared VAE, audio VAE, text encoder, and +tokenizer. Point ``--mlx-checkpoint`` at an MLX DiT converted from this +checkpoint (the AdaLN cache must be the contract ladder, not the four-step +uniform grid). + +``--steps`` accepts 8 transformer forwards or 9 sigma-grid points. Either +value runs the trained rungs. The four-step preview entrypoint is unchanged. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +try: + from . import mlx_fasth3 +except ImportError: + import mlx_fasth3 # type: ignore[no-redef] + +FORWARDS = 8 +GRID_POINTS = FORWARDS + 1 + + +def parse_args(argv: Sequence[str] | None = None): + parser = mlx_fasth3.build_parser() + parser.set_defaults(steps=FORWARDS, vsa=True, vsa_sparsity=0.8, vsa_tile_size=64) + args = parser.parse_args(argv) + if args.steps not in (FORWARDS, GRID_POINTS): + parser.error(f"--steps must be {FORWARDS} or {GRID_POINTS} for FastH3 8-Step V2") + return args + + +def main() -> None: + mlx_fasth3.run(parse_args()) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index d990e62a4..58b677bd8 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -107,6 +107,8 @@ MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999 MINIMAX_H3_VIDEO_SHIFT = 12.0 MINIMAX_H3_AUDIO_SHIFT = 3.0 +FASTH3_INFERENCE_CONTRACT = "fastvideo_inference.json" +FASTH3_INFERENCE_CONTRACT_SCHEMA = "fasth3-inference-contract-v1" MINIMAX_H3_ROPE_FRAME_RESCALE = 5.0 / 3.0 MINIMAX_H3_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) @@ -418,6 +420,142 @@ def minimax_h3_sigmas(shift: float, num_denoise_steps: int) -> np.ndarray: return _unique_consecutive(sigmas).astype(np.float32) +def validate_fasth3_dmd_rungs(dmd_steps: object) -> list[int]: + """Strictly decreasing integer rungs on the 1000-step training clock.""" + if (not isinstance(dmd_steps, list) or not dmd_steps + or any(type(step) is not int or not 0 < step <= 1000 for step in dmd_steps) + or any(left <= right for left, right in zip(dmd_steps, dmd_steps[1:], strict=False))): + raise ValueError("FastH3 DMD rungs must be strictly decreasing integers in (0, 1000].") + return dmd_steps + + +def _positive_shift(value: object, *, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be a positive finite number, got {value!r}.") + return float(value) + + +def find_fasth3_inference_contract(model_root: str | Path | None) -> Path | None: + """Return the snapshot contract, also checking the parent of a transformer dir.""" + if model_root is None: + return None + root = Path(model_root) + direct = root / FASTH3_INFERENCE_CONTRACT + if direct.is_file(): + return direct + parent = root.parent / FASTH3_INFERENCE_CONTRACT + if parent.is_file() and parent != direct: + return parent + return None + + +def load_fasth3_inference_contract(model_root: str | Path | None) -> dict[str, Any] | None: + """Load a shift-declaring FastH3 contract. + + Four-step exports ship a sidecar without scheduler shifts and keep the + uniform ladder. Only a contract that declares both shifts, such as + FastH3 8-Step V2, selects the explicit DMD rungs. A shifted sidecar that + cannot be parsed fails instead of silently baking the four-step grid. + """ + path = find_fasth3_inference_contract(model_root) + if path is None: + return None + try: + contract = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Malformed FastH3 inference contract {path}: {exc}") from exc + if not isinstance(contract, dict): + raise ValueError(f"Malformed FastH3 inference contract {path}.") + declares_shifts = "video_scheduler_shift" in contract or "audio_scheduler_shift" in contract + if not declares_shifts: + return None + if contract.get("schema_version") != FASTH3_INFERENCE_CONTRACT_SCHEMA: + raise ValueError(f"Unsupported FastH3 inference contract schema in {path}.") + if "video_scheduler_shift" not in contract or "audio_scheduler_shift" not in contract: + raise ValueError(f"{path} must declare both video_scheduler_shift and audio_scheduler_shift.") + steps = validate_fasth3_dmd_rungs(contract.get("dmd_denoising_steps")) + forwards = contract.get("transformer_forwards") + grid_points = contract.get("num_inference_steps") + if type(forwards) is not int or forwards != len(steps): + raise ValueError("FastH3 transformer_forwards must equal the DMD rung count.") + if type(grid_points) is not int or grid_points != len(steps) + 1: + raise ValueError("FastH3 num_inference_steps must equal the DMD rung count plus one.") + contract["video_scheduler_shift"] = _positive_shift(contract["video_scheduler_shift"], name="video_scheduler_shift") + contract["audio_scheduler_shift"] = _positive_shift(contract["audio_scheduler_shift"], name="audio_scheduler_shift") + contract["dmd_denoising_steps"] = steps + return contract + + +def _require_cached_schedule(cached: np.ndarray | None, requested: np.ndarray, num_steps: int) -> None: + if cached is None: + return + cached_union = np.unique(np.asarray(cached, dtype=np.float32)) + if np.array_equal(cached_union, requested): + return + raise ValueError(f"MLX H3 checkpoint has a fixed AdaLN ladder that does not support --steps " + f"{num_steps}. Use the step count used during conversion, or re-export the checkpoint.") + + +def adaln_timestep_union(video: MiniMaxH3SchedulerState, audio: MiniMaxH3SchedulerState) -> np.ndarray: + """Sorted unique AdaLN timesteps for one video/audio pair, plus clean time 1.""" + return np.unique( + np.concatenate([ + np.asarray(video.timesteps, dtype=np.float32), + np.asarray(audio.timesteps, dtype=np.float32), + np.array([1.0], dtype=np.float32), + ])).astype(np.float32) + + +@dataclass(frozen=True) +class ResolvedH3Schedule: + """The denoise ladder a request is allowed to run.""" + + video: MiniMaxH3SchedulerState + audio: MiniMaxH3SchedulerState + adaln_timesteps: np.ndarray + num_steps: int + source: str + + +def resolve_h3_denoise_schedule( + model_root: str | Path | None, + num_steps: int, + *, + cached_timesteps: np.ndarray | None = None, +) -> ResolvedH3Schedule: + """Keep the uniform ladder unless the snapshot declares its own shifts. + + ``num_steps`` is the transformer-forward count. A shift-declaring contract + also accepts that count plus one, the CUDA sigma-grid convention, and then + runs exactly its trained rung count. + """ + if type(num_steps) is not int or num_steps < 1: + raise ValueError(f"num_steps must be a positive integer, got {num_steps!r}.") + contract = load_fasth3_inference_contract(model_root) + if contract is None: + video = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) + audio = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) + union = adaln_timestep_union(video, audio) + if cached_timesteps is not None and not np.array_equal(np.asarray(cached_timesteps, dtype=np.float32), union): + raise ValueError(f"MLX H3 checkpoint has a fixed AdaLN ladder that does not support --steps " + f"{num_steps}. Use the step count used during conversion (normally 4), or re-export " + "the checkpoint.") + return ResolvedH3Schedule(video, audio, union, num_steps, "uniform") + + rungs = contract["dmd_denoising_steps"] + forwards = len(rungs) + if num_steps not in (forwards, forwards + 1): + raise ValueError(f"Requested {num_steps} steps, but the FastH3 contract trains {forwards} transformer " + f"forwards ({forwards + 1} sigma-grid points).") + video = MiniMaxH3SchedulerState.from_dmd_steps(contract["video_scheduler_shift"], rungs) + audio = MiniMaxH3SchedulerState.from_dmd_steps(contract["audio_scheduler_shift"], rungs) + union = adaln_timestep_union(video, audio) + if cached_timesteps is not None and not np.array_equal(np.unique(np.asarray(cached_timesteps, dtype=np.float32)), + union): + _require_cached_schedule(cached_timesteps, union, num_steps) + return ResolvedH3Schedule(video, audio, union, forwards, "contract-dmd") + + @dataclass class MiniMaxH3SchedulerState: """One rectified-flow scheduler (use two: video shift 12, audio shift 3).""" @@ -434,12 +572,17 @@ class MiniMaxH3SchedulerState: return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32)) @classmethod - def from_dmd_steps(cls, shift: float, dmd_steps) -> MiniMaxH3SchedulerState: - if shift <= 0: - raise ValueError(f"shift must be positive, got {shift}.") - base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64) - sigmas = shift * base / (1.0 + (shift - 1.0) * base) - return cls(shift=shift, sigmas=sigmas, timesteps=(1.0 - sigmas[:-1]).astype(np.float32)) + def from_dmd_steps(cls, shift: float, dmd_steps: list[int]) -> MiniMaxH3SchedulerState: + """Shift explicit DMD rungs the way the CUDA denoiser does.""" + shift = _positive_shift(shift, name="shift") + rungs = validate_fasth3_dmd_rungs(dmd_steps) + base = np.array([step / 1000.0 for step in rungs] + [0.0], dtype=np.float64) + # Float64 shift, then one cast. Shipped MLX AdaLN tables are keyed by + # that rounding; casting the sigma before 1 - sigma misses the cache. + sigmas64 = shift * base / (1.0 + (shift - 1.0) * base) + if sigmas64.size < 2 or not bool(np.all(sigmas64[1:] < sigmas64[:-1])) or float(sigmas64[-1]) != 0.0: + raise ValueError("DMD sigmas must be strictly decreasing and end at 0.") + return cls(shift=shift, sigmas=sigmas64.astype(np.float32), timesteps=(1.0 - sigmas64[:-1]).astype(np.float32)) @property def num_steps(self) -> int: diff --git a/fastvideo/mlx_runtime/minimax_h3_conditioner.py b/fastvideo/mlx_runtime/minimax_h3_conditioner.py index e08fc98ed..464250a9c 100644 --- a/fastvideo/mlx_runtime/minimax_h3_conditioner.py +++ b/fastvideo/mlx_runtime/minimax_h3_conditioner.py @@ -105,15 +105,20 @@ class _ShardIndex: self._header_cache[path] = (header, data_start) return self._header_cache[path] - def get(self, key: str) -> np.ndarray: + def _header(self, key: str) -> tuple[str, dict, int]: shard = self.key_to_shard[key] - header, data_start = self._header_cache[shard] + if not Path(shard).is_file(): + raise FileNotFoundError(f"Conditioner shard listed in the index is missing: {shard}") + header, data_start = self._cache_header(shard) + return shard, header, data_start + + def get(self, key: str) -> np.ndarray: + shard, header, data_start = self._header(key) return _read_safetensors_bf16(shard, key, header, data_start) def get_mlx(self, key: str) -> mx.array: """Read one weight in FP32 without expanding BF16 on the CPU.""" - shard = self.key_to_shard[key] - header, data_start = self._header_cache[shard] + shard, header, data_start = self._header(key) if header[key]["dtype"] != "BF16": return mx.array(np.asarray(self.get(key), dtype=np.float32)) raw = _read_bf16_words(shard, key, header, data_start) @@ -124,8 +129,7 @@ class _ShardIndex: return weight def get_row(self, key: str, row: int) -> np.ndarray: - shard = self.key_to_shard[key] - header, data_start = self._header_cache[shard] + shard, header, data_start = self._header(key) return _read_safetensors_row(shard, key, row, header, data_start) def close(self) -> None: diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index d11cfcb00..c6ca112f3 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -49,6 +49,7 @@ from fastvideo.mlx_runtime.minimax_h3 import ( MINIMAX_H3_MIN_DURATION, MINIMAX_H3_VIDEO_SHIFT, MiniMaxH3SchedulerState, + adaln_timestep_union, align_num_frames, audio_latent_num_frames, build_packed_layout, @@ -56,6 +57,7 @@ from fastvideo.mlx_runtime.minimax_h3 import ( dense_only_vsa_error, load_mlx_h3_checkpoint, mlx_h3_checkpoint_vsa_capable, + resolve_h3_denoise_schedule, temporal_position_grid, unpatchify_video_tokens, unpack_audio_tokens, @@ -258,24 +260,37 @@ def _audio_sample_count(num_frames: int, fps: int = MINIMAX_H3_FPS, sample_rate: def _adaln_schedule_union(num_steps: int) -> np.ndarray: video = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) audio = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) - return np.unique(np.concatenate([video.timesteps, audio.timesteps, [1.0]]).astype(np.float32)) + return adaln_timestep_union(video, audio) -def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int) -> None: - """Reject a schedule that a fixed, AdaLN-dropped checkpoint cannot serve.""" +def _cached_adaln_timesteps(checkpoint_dir: str | Path) -> np.ndarray | None: checkpoint_dir = Path(checkpoint_dir) manifest_path = checkpoint_dir / H3_MANIFEST_FILENAME if not manifest_path.is_file(): raise FileNotFoundError(f"Missing MLX H3 checkpoint manifest: {manifest_path}") cache_info = json.loads(manifest_path.read_text()).get("adaln_cache") if cache_info is None: - return - cached = np.asarray(cache_info["timesteps"], dtype=np.float32) - requested = _adaln_schedule_union(num_steps) - if not (np.array_equal(cached, requested) or np.all(np.isin(requested, cached)) or (len(cached) == 17 and num_steps == 8)): - raise ValueError( - f"MLX H3 checkpoint {checkpoint_dir} has a fixed AdaLN ladder that does not support --steps " - f"{num_steps}. Use the step count used during conversion (normally 4), or re-export the checkpoint.") + return None + return np.asarray(cache_info["timesteps"], dtype=np.float32) + + +def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, + num_steps: int, + *, + model_root: str | Path | None = None) -> None: + """Reject a schedule that a fixed, AdaLN-dropped checkpoint cannot serve.""" + resolve_h3_denoise_schedule( + model_root, + num_steps, + cached_timesteps=_cached_adaln_timesteps(checkpoint_dir), + ) + + +def _require_positive_vae_tiles(vae_tile_height: int, vae_tile_width: int) -> None: + if type(vae_tile_height) is not int or type( + vae_tile_width) is not int or vae_tile_height <= 0 or vae_tile_width <= 0: + raise ValueError("VAE tile dimensions must be positive integers, " + f"got height={vae_tile_height!r}, width={vae_tile_width!r}.") def _preflight_media_dependencies(*, @@ -465,7 +480,7 @@ class MiniMaxH3MLXPipeline: owned_dit = dit is None if owned_dit: - _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps) + _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps, model_root=self.model_root) t0 = time.perf_counter() dit = load_mlx_h3_checkpoint(self.dit_checkpoint) logger.info("Loaded MLX H3 DiT from %s in %.1fs", self.dit_checkpoint, time.perf_counter() - t0) @@ -487,25 +502,11 @@ class MiniMaxH3MLXPipeline: if getattr(dit, "vsa_config", None) is not None and dit.vsa_config.enabled: dit.prepare_vsa_geometry(layout) - contract_path = self.model_root / "fastvideo_inference.json" - if contract_path.exists(): - contract = json.loads(contract_path.read_text(encoding="utf-8")) - dmd_steps = contract.get("dmd_denoising_steps") - if dmd_steps and (num_steps == len(dmd_steps) or num_steps == len(dmd_steps) + 1 or num_steps == 8): - v_shift = float(contract.get("video_scheduler_shift", 10.0)) - a_shift = float(contract.get("audio_scheduler_shift", 3.0)) - video_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(v_shift, dmd_steps) - audio_scheduler = MiniMaxH3SchedulerState.from_dmd_steps(a_shift, dmd_steps) - union = np.unique(np.concatenate([video_scheduler.timesteps, audio_scheduler.timesteps, [1.0]])).astype(np.float32) - num_steps = len(dmd_steps) - else: - video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) - audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) - union = _adaln_schedule_union(num_steps) - else: - video_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_VIDEO_SHIFT, num_steps) - audio_scheduler = MiniMaxH3SchedulerState.create(MINIMAX_H3_AUDIO_SHIFT, num_steps) - union = _adaln_schedule_union(num_steps) + schedule = resolve_h3_denoise_schedule(self.model_root, num_steps) + video_scheduler = schedule.video + audio_scheduler = schedule.audio + union = schedule.adaln_timesteps + num_steps = schedule.num_steps # The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA # conditioning rows; those modes recompute the ladder before denoise. @@ -551,8 +552,10 @@ class MiniMaxH3MLXPipeline: x_v = video_scheduler.step(video_velocity, step_index, x_v) x_a = audio_scheduler.step(audio_velocity, step_index, x_a) mx.eval(x_v, x_a) - logger.info("[FastH3-8-Step-V2] Step %d/%d (sigma_video=%.4f, sigma_audio=%.4f) done in %.2fs", step_index + 1, num_steps, 1.0 - video_t, 1.0 - audio_t, time.perf_counter() - step_started) - dit_forward_s += time.perf_counter() - step_started + step_s = time.perf_counter() - step_started + dit_forward_s += step_s + logger.info("H3 denoise step %d/%d (sigma_video=%.4f, sigma_audio=%.4f) in %.2fs", step_index + 1, + num_steps, 1.0 - video_t, 1.0 - audio_t, step_s) del video_velocity, audio_velocity _cleanup_mlx() if inter_step_cooldown_s > 0 and step_index < num_steps - 1: @@ -599,6 +602,7 @@ class MiniMaxH3MLXPipeline: vae_tile_height: int = 256, vae_tile_width: int = 256) -> np.ndarray: """Normalized packed rows -> (T, H, W, 3) uint8 frames.""" + _require_positive_vae_tiles(vae_tile_height, vae_tile_width) import mlx.core as mx from fastvideo.mlx_runtime.minimax_h3_video_vae import mlx_h3_video_vae_from_dir @@ -749,42 +753,45 @@ class MiniMaxH3MLXPipeline: # -- end-to-end ---------------------------------------------------------- def generate( - self, - prompt: str, - *, - output_path: str | Path, - height: int = 480, - width: int = 832, - num_frames: int = 124, - seed: int = 0, - num_steps: int = 4, - save_frames: bool = False, - tiled_video_decode: bool = True, - vae_tile_height: int = 256, - vae_tile_width: int = 256, - fast: bool = False, - fast_factor: int = 2, - fast_sharpen: float = 0.6, - rife_weights_dir: str | Path | None = None, - fast_spatial: bool = False, - fast_spatial_scale: int = 2, - fast_spatial_upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, - fast_spatial_sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN, - vsa: bool = False, - vsa_sparsity: float = 0.9, - vsa_tile_size: int = 64, - vsa_prefix_mode: str = "exempt", - vsa_dense_first_n_steps: int = 0, - vsa_dense_layers: tuple[int, ...] = (), - vsa_impl: str = "auto", - inter_step_cooldown_s: float = 0.0, + self, + prompt: str, + *, + output_path: str | Path, + height: int = 480, + width: int = 832, + num_frames: int = 124, + seed: int = 0, + num_steps: int = 4, + save_frames: bool = False, + tiled_video_decode: bool = True, + vae_tile_height: int = 256, + vae_tile_width: int = 256, + fast: bool = False, + fast_factor: int = 2, + fast_sharpen: float = 0.6, + rife_weights_dir: str | Path | None = None, + fast_spatial: bool = False, + fast_spatial_scale: int = 2, + fast_spatial_upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + fast_spatial_sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN, + vsa: bool = False, + vsa_sparsity: float = 0.9, + vsa_tile_size: int = 64, + vsa_prefix_mode: str = "exempt", + vsa_dense_first_n_steps: int = 0, + vsa_dense_layers: tuple[int, ...] = (), + vsa_impl: str = "auto", + inter_step_cooldown_s: float = 0.0, ) -> GenerationResult: timings: dict[str, float] = {} peaks: dict[str, float] = {} if fast_sharpen < 0: raise ValueError(f"fast_sharpen must be non-negative, got {fast_sharpen}.") - _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps) + if inter_step_cooldown_s < 0: + raise ValueError(f"inter_step_cooldown_s must be non-negative, got {inter_step_cooldown_s}.") + _require_positive_vae_tiles(vae_tile_height, vae_tile_width) + _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps, model_root=self.model_root) vsa_config = MiniMaxH3VSAConfig( enabled=vsa, sparsity=vsa_sparsity, diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_8step_schedule.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_8step_schedule.py new file mode 100644 index 000000000..1cd1b4b08 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_8step_schedule.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastH3 8-Step V2 uses its contract ladder; the four-step MLX path stays uniform.""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for MiniMax H3 schedule tests") + +from fastvideo.mlx_runtime.minimax_h3 import ( # noqa: E402 + MiniMaxH3SchedulerState, + load_fasth3_inference_contract, + resolve_h3_denoise_schedule, + validate_fasth3_dmd_rungs, +) +from fastvideo.mlx_runtime.minimax_h3_pipeline import ( # noqa: E402 + _adaln_schedule_union, + _require_positive_vae_tiles, + _validate_checkpoint_step_ladder, +) +from scripts.checkpoint_conversion.convert_minimax_h3_mlx import _adaln_cache_timesteps # noqa: E402 + +EIGHT_RUNGS = [999, 874, 749, 624, 500, 375, 250, 125] +FOUR_RUNGS = [999, 749, 500, 250] + + +def _contract(steps: list[int], *, shifts: bool) -> dict: + payload = { + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": steps, + "num_inference_steps": len(steps) + 1, + "transformer_forwards": len(steps), + } + if shifts: + payload["video_scheduler_shift"] = 10.0 + payload["audio_scheduler_shift"] = 3.0 + return payload + + +def _write_contract(tmp_path, payload: dict) -> None: + (tmp_path / "fastvideo_inference.json").write_text(json.dumps(payload)) + + +def _write_cache(tmp_path, timesteps) -> None: + values = np.asarray(timesteps, dtype=np.float32).tolist() + (tmp_path / "mlx_h3_dit.json").write_text(json.dumps({"adaln_cache": {"timesteps": values}})) + + +def test_four_step_sidecar_keeps_the_uniform_ladder(tmp_path) -> None: + _write_contract(tmp_path, _contract(FOUR_RUNGS, shifts=False)) + uniform = _adaln_schedule_union(4) + _write_cache(tmp_path, uniform) + + assert load_fasth3_inference_contract(tmp_path) is None + schedule = resolve_h3_denoise_schedule(tmp_path, 4, cached_timesteps=uniform) + assert schedule.source == "uniform" + assert schedule.num_steps == 4 + assert schedule.video.shift == 12.0 + _validate_checkpoint_step_ladder(tmp_path, 4, model_root=tmp_path) + with pytest.raises(ValueError, match=r"does not support --steps 8"): + _validate_checkpoint_step_ladder(tmp_path, 8, model_root=tmp_path) + + +def test_shifted_contract_rejects_a_different_step_count(tmp_path) -> None: + _write_contract(tmp_path, _contract(FOUR_RUNGS, shifts=True)) + with pytest.raises(ValueError, match="trains 4 transformer forwards"): + resolve_h3_denoise_schedule(tmp_path, 8) + + +def test_eight_step_contract_runs_the_trained_rungs(tmp_path) -> None: + _write_contract(tmp_path, _contract(EIGHT_RUNGS, shifts=True)) + for requested in (8, 9): + schedule = resolve_h3_denoise_schedule(tmp_path, requested) + assert schedule.source == "contract-dmd" + assert schedule.num_steps == 8 + assert schedule.video.shift == 10.0 + assert schedule.audio.shift == 3.0 + assert schedule.video.timesteps.shape == (8, ) + assert float(schedule.video.sigmas[-1]) == 0.0 + assert bool(np.all(schedule.video.sigmas[1:] < schedule.video.sigmas[:-1])) + + union = schedule.adaln_timesteps + _write_cache(tmp_path, union) + _validate_checkpoint_step_ladder(tmp_path, 8, model_root=tmp_path) + _validate_checkpoint_step_ladder(tmp_path, 9, model_root=tmp_path) + + _write_cache(tmp_path, np.linspace(0.0, 1.0, 17)) + with pytest.raises(ValueError, match=r"does not support --steps 8"): + _validate_checkpoint_step_ladder(tmp_path, 8, model_root=tmp_path) + + +def test_dmd_rungs_reject_invalid_ladders() -> None: + with pytest.raises(ValueError, match="strictly decreasing"): + validate_fasth3_dmd_rungs([999, 999]) + with pytest.raises(ValueError, match="strictly decreasing"): + validate_fasth3_dmd_rungs([100, 200]) + with pytest.raises(ValueError, match="strictly decreasing"): + validate_fasth3_dmd_rungs([0, -1]) + with pytest.raises(ValueError, match="strictly decreasing"): + validate_fasth3_dmd_rungs([1001]) + with pytest.raises(ValueError, match="positive finite"): + MiniMaxH3SchedulerState.from_dmd_steps(0.0, EIGHT_RUNGS) + + +def test_malformed_shifted_contract_fails_closed(tmp_path) -> None: + (tmp_path / "fastvideo_inference.json").write_text("{") + with pytest.raises(ValueError, match="Malformed"): + load_fasth3_inference_contract(tmp_path) + with pytest.raises(ValueError, match="Malformed"): + _adaln_cache_timesteps(tmp_path) + + +def test_converter_cache_follows_only_a_shifted_contract(tmp_path) -> None: + four_step = _adaln_cache_timesteps(None) + np.testing.assert_array_equal(four_step, _adaln_schedule_union(4)) + + _write_contract(tmp_path, _contract(FOUR_RUNGS, shifts=False)) + np.testing.assert_array_equal(_adaln_cache_timesteps(tmp_path), four_step) + + shifted = tmp_path / "shifted" + shifted.mkdir() + _write_contract(shifted, _contract(EIGHT_RUNGS, shifts=True)) + cached = _adaln_cache_timesteps(shifted) + expected = resolve_h3_denoise_schedule(shifted, 8).adaln_timesteps + np.testing.assert_array_equal(cached, expected) + assert not np.array_equal(cached, _adaln_schedule_union(8)) + + +def test_vae_tiles_must_be_positive() -> None: + _require_positive_vae_tiles(256, 256) + with pytest.raises(ValueError, match="positive"): + _require_positive_vae_tiles(0, 256) + with pytest.raises(ValueError, match="positive"): + _require_positive_vae_tiles(256, -1) diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py index 0f360f24b..db96f413c 100644 --- a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py +++ b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py @@ -42,6 +42,9 @@ from fastvideo.mlx_runtime.minimax_h3 import ( H3_WEIGHTS_FILENAME, MINIMAX_H3_AUDIO_SHIFT, MINIMAX_H3_VIDEO_SHIFT, + MiniMaxH3SchedulerState, + adaln_timestep_union, + load_fasth3_inference_contract, mlx_h3_checkpoint_vsa_capable, mlx_h3_dit_from_diffusers_safetensors, minimax_h3_sigmas, @@ -55,31 +58,25 @@ DEFAULT_FORMATS = " ".join(SUPPORTED_FORMATS) def _adaln_cache_timesteps(model_root: str | Path | None = None) -> np.ndarray: - import json - if model_root is not None: - p = Path(model_root) - contract_path = p / "fastvideo_inference.json" - if not contract_path.exists() and p.parent: - contract_path = p.parent / "fastvideo_inference.json" - if contract_path.exists(): - try: - contract = json.loads(contract_path.read_text(encoding="utf-8")) - dmd_steps = contract.get("dmd_denoising_steps") - v_shift = float(contract.get("video_scheduler_shift", 10.0)) - a_shift = float(contract.get("audio_scheduler_shift", 3.0)) - if dmd_steps: - base = np.array([float(r) / 1000.0 for r in dmd_steps] + [0.0], dtype=np.float64) - v_sigmas = v_shift * base / (1.0 + (v_shift - 1.0) * base) - a_sigmas = a_shift * base / (1.0 + (a_shift - 1.0) * base) - v_t = (1.0 - v_sigmas[:-1]).astype(np.float32) - a_t = (1.0 - a_sigmas[:-1]).astype(np.float32) - logger.info("FastH3 checkpoint schedule contract found at %s (%d DMD forwards, video/audio shift=%g/%g)", contract_path, len(dmd_steps), v_shift, a_shift) - return np.unique(np.concatenate([v_t, a_t, [1.0]])).astype(np.float32) - except Exception as exc: - logger.warning("Could not parse %s: %s", contract_path, exc) - video = 1.0 - minimax_h3_sigmas(MINIMAX_H3_VIDEO_SHIFT, 4)[:-1] - audio = 1.0 - minimax_h3_sigmas(MINIMAX_H3_AUDIO_SHIFT, 4)[:-1] - return np.unique(np.concatenate([video, audio, [1.0]])).astype(np.float32) + """Four-step uniform grid, unless the snapshot declares its own shifts. + + FastH3 8-Step V2 carries ``video_scheduler_shift`` / ``audio_scheduler_shift``. + Four-step sidecars do not, so those conversions keep the existing cache. + A shifted contract that fails to parse aborts conversion. + """ + contract = load_fasth3_inference_contract(model_root) + if contract is None: + video = 1.0 - minimax_h3_sigmas(MINIMAX_H3_VIDEO_SHIFT, 4)[:-1] + audio = 1.0 - minimax_h3_sigmas(MINIMAX_H3_AUDIO_SHIFT, 4)[:-1] + return np.unique(np.concatenate([video, audio, [1.0]])).astype(np.float32) + video_state = MiniMaxH3SchedulerState.from_dmd_steps(contract["video_scheduler_shift"], + contract["dmd_denoising_steps"]) + audio_state = MiniMaxH3SchedulerState.from_dmd_steps(contract["audio_scheduler_shift"], + contract["dmd_denoising_steps"]) + logger.info("FastH3 shifted schedule contract: %d DMD forwards, video/audio shift=%g/%g", + len(contract["dmd_denoising_steps"]), contract["video_scheduler_shift"], + contract["audio_scheduler_shift"]) + return adaln_timestep_union(video_state, audio_state) def parse_args() -> argparse.Namespace: From c100c66578de959c837b3ae1a7bc127604ae9b96 Mon Sep 17 00:00:00 2001 From: Xingjian Guo <7291967@qq.com> Date: Fri, 25 Sep 2026 08:04:59 +0800 Subject: [PATCH 023/151] [bugfix]: isolate cancelled streaming requests from GPU result readers (#1848) Co-authored-by: Gxj230958 <222823329+Gxj230958@users.noreply.github.com> Co-authored-by: SolitaryThinker --- fastvideo/entrypoints/streaming/gpu_pool.py | 17 ++- .../entrypoints/streaming/test_gpu_pool.py | 105 ++++++++++++++++++ 2 files changed, 116 insertions(+), 6 deletions(-) diff --git a/fastvideo/entrypoints/streaming/gpu_pool.py b/fastvideo/entrypoints/streaming/gpu_pool.py index 1da2f6f38..d7e97d0b4 100644 --- a/fastvideo/entrypoints/streaming/gpu_pool.py +++ b/fastvideo/entrypoints/streaming/gpu_pool.py @@ -336,8 +336,8 @@ class SubprocessGpuPool(GpuPool): ) # mp.Queue.put can block if the underlying pipe buffer is full; # offload to a thread so the event loop keeps serving other - # sessions. If the put itself fails, drop the pending entry so - # _drain_results doesn't dangle a future forever. + # sessions. A cancelled caller cannot stop the executor thread or + # generation, so detach its pending entry and discard late replies. loop = asyncio.get_running_loop() try: await loop.run_in_executor( @@ -348,10 +348,10 @@ class SubprocessGpuPool(GpuPool): "request": request }, ) - except Exception: + return await asyncio.wrap_future(future) + finally: self._pending.pop(job_id, None) - raise - return await asyncio.wrap_future(future) + future.cancel() async def release(self, session_id: str) -> None: async with self._lock: @@ -418,7 +418,12 @@ class SubprocessGpuPool(GpuPool): if job_id is None: continue pending = self._pending.pop(job_id, None) - if pending is None: + # wrap_future can cancel the source future before run()'s + # finally block gets its turn on the event loop, and an + # abandoned job's late reply can arrive after run() already + # popped its entry. Either way there is no live waiter. + if pending is None or pending.future.done(): + logger.debug("pool: worker %s dropping reply for job %s", handle.worker_id, job_id) continue if msg.get("kind") == "error": pending.future.set_exception(RuntimeError(msg["error"])) diff --git a/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py index 2f02d6843..779b54d64 100644 --- a/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py +++ b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py @@ -12,6 +12,7 @@ import multiprocessing as mp import queue import threading import time +from concurrent.futures import Future from dataclasses import dataclass from typing import Any @@ -28,6 +29,7 @@ from fastvideo.entrypoints.streaming.gpu_pool import ( InProcessGpuPool, PoolAcquireTimeout, SubprocessGpuPool, + _PendingJob, _WorkerHandle, ) @@ -387,6 +389,109 @@ class TestSubprocessGpuPool: class TestSubprocessGpuPoolFailureModes: """Coverage for boot/runtime failures the pool has to absorb.""" + @pytest.mark.parametrize("worker_raises", [False, True]) + @pytest.mark.parametrize("cancel_during_dispatch", [False, True]) + def test_cancelled_request_does_not_break_worker(self, monkeypatch, worker_raises, cancel_during_dispatch): + """An abandoned result must not take down the reader for later jobs.""" + + async def run(): + loop = asyncio.get_running_loop() + started = asyncio.Event() + dispatch_started = asyncio.Event() + finish_generation = threading.Event() + finish_dispatch = threading.Event() + + class _BlockingGenerator: + + def generate(self, request): + if request.prompt == "cancel-me": + loop.call_soon_threadsafe(started.set) + if not finish_generation.wait(timeout=5.0): + raise RuntimeError("test failed to unblock generation") + if worker_raises: + raise RuntimeError("abandoned worker error") + return {"prompt_echo": request.prompt} + + pool = SubprocessGpuPool( + generator_config=GeneratorConfig(model_path="/models/fake"), + pool_config=GpuPoolConfig(num_workers=1), + warmup_config=WarmupConfig(enabled=False), + worker_factory=_thread_worker_factory(lambda gpu_id: _BlockingGenerator()), + ) + await pool.start() + request_task = None + try: + await pool.acquire("sess-a") + if cancel_during_dispatch: + job_queue = pool._workers[0].job_queue + original_put = job_queue.put + + def blocked_put(item): + if item is not None and item["request"].prompt == "cancel-me": + loop.call_soon_threadsafe(dispatch_started.set) + if not finish_dispatch.wait(timeout=5.0): + raise RuntimeError("test failed to unblock dispatch") + original_put(item) + + monkeypatch.setattr(job_queue, "put", blocked_put) + + request_task = asyncio.create_task(pool.run("sess-a", GenerationRequest(prompt="cancel-me"))) + await asyncio.wait_for(dispatch_started.wait() if cancel_during_dispatch else started.wait(), 2.0) + request_task.cancel() + with pytest.raises(asyncio.CancelledError): + await request_task + # Cancelling the caller cannot stop an executor thread or a + # worker already generating, but must detach its pending job. + assert not pool._pending + finish_dispatch.set() + await asyncio.wait_for(started.wait(), 2.0) + finish_generation.set() + + # Reuse the same worker for another session. FIFO delivery + # ensures its reader sees the abandoned result first. + await pool.release("sess-a") + await pool.acquire("sess-b", timeout=2.0) + result = await asyncio.wait_for(pool.run("sess-b", GenerationRequest(prompt="next")), 2.0) + assert result["prompt_echo"] == "next" + assert not pool._pending + assert all(not task.done() for task in pool._result_reader_tasks) + finally: + finish_dispatch.set() + finish_generation.set() + if request_task is not None: + request_task.cancel() + await asyncio.gather(request_task, return_exceptions=True) + await pool.shutdown() + + asyncio.run(run()) + + @pytest.mark.parametrize("worker_raises", [False, True]) + def test_reader_ignores_cancelled_future_before_request_cleanup(self, pool_factory, worker_raises): + """A cancelled future still in ``_pending`` must not kill the reader.""" + + async def run(): + pool = await pool_factory(num_workers=1) + try: + handle = pool._workers[0] + future = Future() + future.cancel() + pool._pending["cancelled"] = _PendingJob("cancelled", future, "sess-a", handle.worker_id) + handle.result_queue.put({ + "job_id": "cancelled", + "kind": "error" if worker_raises else "result", + "error": "abandoned error", + "result": {}, + }) + await pool.acquire("sess-b") + result = await asyncio.wait_for(pool.run("sess-b", GenerationRequest(prompt="next")), 2.0) + assert result["prompt_echo"] == "next" + assert not pool._pending + assert all(not task.done() for task in pool._result_reader_tasks) + finally: + await pool.shutdown() + + asyncio.run(run()) + def test_failed_boot_excluded_from_available(self): """A worker whose factory leaves ``boot_ok`` unset must not be handed out by ``acquire``. Otherwise a session lands on a dead From d265a454bf8d0eeda3225ba1bc9239b8a1bfaf84 Mon Sep 17 00:00:00 2001 From: Raghav K Date: Thu, 24 Sep 2026 17:05:59 -0700 Subject: [PATCH 024/151] [bugfix]: Fall back when large pinned output allocation fails (#1759) --- fastvideo/entrypoints/video_generator.py | 5 +-- .../minimax_h3/stages/minimax_h3_decoding.py | 7 ++- .../tests/entrypoints/test_video_generator.py | 34 +++++++++++++++ .../stages/test_minimax_h3_vae_streaming.py | 43 +++++++++++++++++++ .../test_minimax_h3_video_vae_streaming.py | 16 +++++++ fastvideo/utils.py | 32 ++++++++++++++ 6 files changed, 130 insertions(+), 7 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index cf55c1e1f..65b6d8d4e 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -55,7 +55,7 @@ from fastvideo.api.sampling_param import SamplingParam from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType from fastvideo.logger import init_logger from fastvideo.pipelines import ForwardBatch -from fastvideo.utils import align_to, shallow_asdict +from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback, align_to, shallow_asdict from fastvideo.worker.executor import Executor fcntl: types.ModuleType | None @@ -819,9 +819,8 @@ class VideoGenerator: if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: - samples = torch.empty( + samples = allocate_cpu_tensor_with_pin_fallback( (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width), - device='cpu', pin_memory=fastvideo_args.pin_cpu_memory) thread.join() 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 a12cca50d..65c1ed371 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -26,7 +26,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.utils import is_pin_memory_available +from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback logger = init_logger(__name__) @@ -135,11 +135,10 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): output = None if is_output_rank: - output = torch.empty( + output = allocate_cpu_tensor_with_pin_fallback( self.vae.decoded_pixel_shape(latents.shape), - device="cpu", dtype=torch.float32, - pin_memory=fastvideo_args.pin_cpu_memory and is_pin_memory_available(), + pin_memory=fastvideo_args.pin_cpu_memory, ) # Attribute the streamed decoder computation while retaining # per-chunk device-to-host transfer and pinned-buffer reuse. diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 0e7b940fa..4ccb0c922 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -8,6 +8,7 @@ import torch from einops import rearrange import fastvideo.entrypoints.video_generator as video_generator_module +import fastvideo.utils as fastvideo_utils from fastvideo.api import ( ConfigValidationError, GenerationRequest, @@ -273,6 +274,39 @@ def test_generate_single_video_return_frames_still_materializes_output(tmp_path) assert result["video_path"] is None +def test_generate_single_video_retries_pageable_output_when_pinned_allocation_fails(monkeypatch, tmp_path): + output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 + output_batch = _single_video_output_batch(output) + fastvideo_args = _single_video_args() + fastvideo_args.pin_cpu_memory = True + generator = _single_video_generator(output_batch, fastvideo_args) + real_empty = torch.empty + allocation_attempts = [] + warnings = [] + + def fake_empty(size, **kwargs): + allocation_attempts.append(kwargs.get("pin_memory", False)) + if kwargs.get("pin_memory"): + raise RuntimeError("CUDA error: out of memory") + return real_empty(size, **kwargs) + + monkeypatch.setattr(fastvideo_utils, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(fastvideo_utils.torch, "empty", fake_empty) + monkeypatch.setattr(fastvideo_utils.logger, "warning", lambda message, *args: warnings.append(message % args)) + + result = generator._generate_single_video( + prompt="pinned allocation fallback", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + torch.testing.assert_close(result["samples"], output) + assert allocation_attempts[:2] == [True, False] + assert len(warnings) == 1 + assert "retrying with pageable memory" in warnings[0] + + def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path): """The on-device quantize path (#1362) must reproduce the legacy per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 6202e11d0..2e605139f 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -114,6 +114,49 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: assert torch.all(result.output == 0.25) +def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: + latent_shape = (1, 4, 2, 4, 4) + latents = torch.randn(latent_shape) + rows = patchify_video_latents(latents, (1, 1, 1)) + batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) + observed = {} + + class VAE: + + def to(self, device): + return self + + def denormalize_latents(self, decoded_latents): + return decoded_latents + + def decoded_pixel_shape(self, shape): + return (1, 3, 5, 16, 16) + + def decode_to_pixels(self, decoded_latents, output): + output.zero_() + + def fake_allocate(size, *, dtype=None, pin_memory=False): + observed["size"] = size + observed["dtype"] = dtype + observed["pin_memory"] = pin_memory + return torch.empty(size, 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) + + MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward( + batch, + SimpleNamespace(output_type="pil", pin_cpu_memory=True, vae_cpu_offload=False, vae_parallel_decode=False), + ) + + assert observed == { + "size": (1, 3, 5, 16, 16), + "dtype": torch.float32, + "pin_memory": True, + } + + def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None: class VAE: diff --git a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py index 5a2c8783e..68e70da3e 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py +++ b/fastvideo/tests/vaes/test_minimax_h3_video_vae_streaming.py @@ -158,6 +158,22 @@ def test_decode_to_pixels_pinned_buffer_matches_dense_on_cuda() -> None: assert_close(pinned, expected, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="pageable-buffer streaming requires CUDA") +@torch.inference_mode() +def test_decode_to_pixels_pageable_buffer_matches_dense_on_cuda() -> None: + """The allocation fallback's blocking chunk copies must equal dense decode.""" + torch.manual_seed(20260814) + vae = _tiny_vae().to("cuda") + latents = torch.randn(1, 4, 12, 4, 4, device="cuda") + expected = vae.denormalize_pixels(vae.decode(latents).sample.float()).clamp_(0, 1).cpu() + + pageable = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32) + vae.decode_to_pixels(latents, pageable) + + assert not pageable.is_pinned() + assert_close(pageable, expected, atol=0.0, rtol=0.0) + + def test_decode_to_pixels_rejects_incomplete_output(monkeypatch) -> None: vae = _tiny_vae() latents = torch.randn(1, 4, 2, 4, 4) diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 79757a23a..af4d2acda 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -1306,3 +1306,35 @@ def _cached_pin_memory_available(pid: int) -> bool: def is_pin_memory_available() -> bool: return _cached_pin_memory_available(os.getpid()) + + +def allocate_cpu_tensor_with_pin_fallback( + size: tuple[int, ...] | torch.Size, + *, + dtype: torch.dtype | None = None, + pin_memory: bool = False, +) -> torch.Tensor: + """Allocate a CPU tensor, retrying pageable memory if pinning fails. + + ``is_pin_memory_available`` intentionally uses a small capability probe. + Large decoded-video buffers can still exhaust the CUDA host allocator after + that probe succeeds, especially on unified-memory systems. Treat pinning as + a performance preference: preserve the requested shape and dtype by + retrying the allocation without pinning. + """ + kwargs: dict[str, object] = {"device": "cpu"} + if dtype is not None: + kwargs["dtype"] = dtype + + if not pin_memory or not is_pin_memory_available(): + return torch.empty(size, **kwargs) + + try: + return torch.empty(size, pin_memory=True, **kwargs) + except RuntimeError as exc: + logger.warning( + "Pinned CPU allocation failed for shape %s; retrying with pageable memory: %s", + tuple(size), + exc, + ) + return torch.empty(size, **kwargs) From 81f3fec7fd1f1fb095c4517caed3fea5532c890f Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:06:55 -0700 Subject: [PATCH 025/151] [bugfix]: workers killed by a signal now log the reason (#1725) --- .../installation/spark_performance.md | 3 + ...st_signal_terminated_worker_is_reported.py | 215 ++++++++++++++++++ fastvideo/worker/multiproc_executor.py | 55 ++++- 3 files changed, 265 insertions(+), 8 deletions(-) create mode 100644 fastvideo/tests/worker/test_signal_terminated_worker_is_reported.py diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 49c4c9146..e0f5cb4d1 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -181,6 +181,9 @@ is power-cycled. To avoid it: and frame counts are valid. Weights stay replicated, so lazy module load (auto on GB10) is still required on each box. Bring-up and knobs: [Pair two NVIDIA DGX Sparks](spark_pair.md). +- A worker's SIGTERM log and traceback show where it was interrupted, not why it + was selected; confirm the cause in the `earlyoom` service or system logs. A + later SIGKILL or kernel OOM kill cannot be caught and reported by Python. ## Gotchas specific to the GB10 diff --git a/fastvideo/tests/worker/test_signal_terminated_worker_is_reported.py b/fastvideo/tests/worker/test_signal_terminated_worker_is_reported.py new file mode 100644 index 000000000..40294294c --- /dev/null +++ b/fastvideo/tests/worker/test_signal_terminated_worker_is_reported.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +"""A worker killed by a signal has to say which signal it received. + +`worker_main` installs a SIGTERM/SIGINT handler that raises `SystemExit`. That is +a `BaseException`, not an `Exception`, so it walks straight past a bare +`except Exception` and the worker exits having reported nothing. The parent then +sees a closed pipe and raises "See stack trace for root cause" with no stack +trace, which is indistinguishable from a hang, a crash, or an out-of-memory kill. + +That matters on machines that terminate processes under memory pressure. A DGX +Spark ships earlyoom configured to send SIGTERM to python first once available +memory drops below six percent, so this is the common failure there and it is +completely silent. + +The tests replace startup and worker construction with lightweight signal +injection, so they exercise the real handler and outer exception path without +initializing a distributed process group or model. +""" +from __future__ import annotations + +import multiprocessing as mp +import os +import signal +import sys + +import pytest + +import fastvideo.worker.multiproc_executor as multiproc_executor +from fastvideo.worker.multiproc_executor import WorkerMultiprocProc + +# The spawn context reimports PyTorch and FastVideo in every child. That takes +# nearly five seconds even on the local GB10 and can take longer on a loaded ARM +# CI node, so a five-second startup deadline makes the signal contract flaky. +_SUBPROCESS_TIMEOUT_SECONDS = 30 + + +class _ReadyPipe: + + def __init__(self) -> None: + self.closed = False + self.messages = [] + + def send(self, message) -> None: + self.messages.append(message) + + def close(self) -> None: + self.closed = True + + +class _ParentlessProcess: + + def parent(self): + return None + + +def _run_real_signal_probe(ready_pipe, log_queue, signal_point, entered_signal_point, shutdown_called) -> None: + + class SignalWaitingWorker: + + READY_STR = "READY" + + def __init__(self, *args, **kwargs) -> None: + if signal_point == "worker_construction": + entered_signal_point.set() + signal.pause() + + def worker_busy_loop(self) -> None: + assert signal_point == "worker_loop" + entered_signal_point.set() + signal.pause() + + def shutdown(self) -> None: + shutdown_called.set() + + multiproc_executor.kill_itself_when_parent_died = lambda: None + multiproc_executor.faulthandler.enable = lambda: None + multiproc_executor.psutil.Process = _ParentlessProcess + multiproc_executor.WorkerMultiprocProc = SignalWaitingWorker + WorkerMultiprocProc.worker_main(ready_pipe=ready_pipe, rank=7, log_queue=log_queue) + + +@pytest.mark.parametrize( + ("signum", "expected_reason", "unexpected_reason"), + [ + (signal.SIGTERM, "parent cleaning up workers", "user interrupted"), + (signal.SIGINT, "user interrupted", "out-of-memory daemon"), + ], +) +@pytest.mark.parametrize("injection_point", ["startup", "worker_construction", "worker_loop"]) +def test_worker_main_reports_received_signal(monkeypatch: pytest.MonkeyPatch, signum: int, expected_reason: str, + unexpected_reason: str, injection_point: str) -> None: + installed_handlers = {} + logged_messages = [] + workers = [] + ready_pipe = _ReadyPipe() + + def install_handler(installed_signum, handler): + installed_handlers[installed_signum] = handler + + def inject_signal() -> None: + installed_handlers[signum](signum, None) + + class SignalledWorker: + + READY_STR = "READY" + + def __init__(self, *args, **kwargs) -> None: + self.shutdown_called = False + workers.append(self) + if injection_point == "worker_construction": + inject_signal() + + def worker_busy_loop(self) -> None: + assert injection_point == "worker_loop" + inject_signal() + + def shutdown(self) -> None: + self.shutdown_called = True + + monkeypatch.setattr(multiproc_executor.signal, "signal", install_handler) + monkeypatch.setattr(multiproc_executor, "kill_itself_when_parent_died", + inject_signal if injection_point == "startup" else lambda: None) + monkeypatch.setattr(multiproc_executor.faulthandler, "enable", lambda: None) + monkeypatch.setattr(multiproc_executor.psutil, "Process", _ParentlessProcess) + monkeypatch.setattr(multiproc_executor, "WorkerMultiprocProc", SignalledWorker) + monkeypatch.setattr(multiproc_executor.logger, "exception", + lambda message, *args: logged_messages.append(message % args)) + + with pytest.raises(SystemExit) as exc_info: + WorkerMultiprocProc.worker_main(ready_pipe=ready_pipe, rank=7) + + assert exc_info.value.code is None + assert exc_info.value.signum == signum + assert ready_pipe.closed + assert len(logged_messages) == 1 + assert f"Worker 7 received {signal.Signals(signum).name} ({signum})" in logged_messages[0] + assert expected_reason in logged_messages[0] + assert unexpected_reason not in logged_messages[0] + if injection_point == "worker_loop": + assert ready_pipe.messages == [{"status": "READY"}] + assert workers[0].shutdown_called + + +def test_worker_main_preserves_unrelated_system_exit(monkeypatch: pytest.MonkeyPatch) -> None: + logged_messages = [] + ready_pipe = _ReadyPipe() + + class ExitingWorker: + + def __init__(self, *args, **kwargs) -> None: + raise SystemExit(23) + + monkeypatch.setattr(multiproc_executor.signal, "signal", lambda *args: None) + monkeypatch.setattr(multiproc_executor, "kill_itself_when_parent_died", lambda: None) + monkeypatch.setattr(multiproc_executor.faulthandler, "enable", lambda: None) + monkeypatch.setattr(multiproc_executor.psutil, "Process", _ParentlessProcess) + monkeypatch.setattr(multiproc_executor, "WorkerMultiprocProc", ExitingWorker) + monkeypatch.setattr(multiproc_executor.logger, "exception", + lambda message, *args: logged_messages.append(message % args)) + + with pytest.raises(SystemExit) as exc_info: + WorkerMultiprocProc.worker_main(ready_pipe=ready_pipe, rank=7) + + assert exc_info.value.code == 23 + assert ready_pipe.closed + assert logged_messages == [] + + +@pytest.mark.skipif(sys.platform != "linux", reason="POSIX signal delivery is a Linux worker contract") +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT]) +@pytest.mark.parametrize("signal_point", ["worker_construction", "worker_loop"]) +def test_worker_main_forwards_real_signal_traceback_across_processes(signum: int, signal_point: str) -> None: + context = mp.get_context("spawn") + entered_signal_point = context.Event() + shutdown_called = context.Event() + parent_ready_pipe, child_ready_pipe = context.Pipe(duplex=False) + log_queue = context.Queue() + + process = context.Process( + target=_run_real_signal_probe, + args=(child_ready_pipe, log_queue, signal_point, entered_signal_point, shutdown_called), + ) + process.start() + child_ready_pipe.close() + + try: + reached_signal_point = entered_signal_point.wait(timeout=_SUBPROCESS_TIMEOUT_SECONDS) + assert reached_signal_point, "child did not reach the requested signal point" + if signal_point == "worker_loop": + assert parent_ready_pipe.recv() == {"status": "READY"} + + os.kill(process.pid, signum) + process.join(timeout=_SUBPROCESS_TIMEOUT_SECONDS) + assert not process.is_alive(), "signalled worker did not exit" + assert process.exitcode == 0 # Preserve the historical argument-less SystemExit status. + + record = log_queue.get(timeout=_SUBPROCESS_TIMEOUT_SECONDS) + message = record.getMessage() + assert f"Worker 7 received {signal.Signals(signum).name} ({signum})" in message + assert "Traceback (most recent call last)" in message + assert "The stack below is where execution was interrupted, not the cause." in message + + if signal_point == "worker_loop": + assert shutdown_called.wait(timeout=1) + else: + assert not shutdown_called.is_set() + with pytest.raises(EOFError): + parent_ready_pipe.recv() + finally: + if process.is_alive(): + process.kill() + process.join(timeout=_SUBPROCESS_TIMEOUT_SECONDS) + parent_ready_pipe.close() + log_queue.close() + log_queue.join_thread() diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 0aa3ad443..98714b3ba 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -46,6 +46,16 @@ def _raise_for_rpc_errors(method: str | Callable, responses: list[Any]) -> None: raise RuntimeError(f"RPC {method!r} failed: " + "; ".join(errors)) +class _WorkerSignalExit(SystemExit): + """Carry the signal that requested worker shutdown through ``SystemExit``.""" + + def __init__(self, signum: int) -> None: + # Keep the original argument-less SystemExit semantics while retaining + # the signal identity for diagnostics. + super().__init__() + self.signum = signum + + def _make_queue_log_handler(log_queue: Queue) -> logging.Handler: """Create a QueueHandler that forwards fastvideo logs to a multiprocessing queue.""" return logging.handlers.QueueHandler(log_queue) @@ -552,20 +562,24 @@ class WorkerMultiprocProc: nonlocal shutdown_requested if not shutdown_requested: shutdown_requested = True - raise SystemExit() - - # Either SIGTERM or SIGINT will terminate the worker - signal.signal(signal.SIGTERM, signal_handler) - signal.signal(signal.SIGINT, signal_handler) - kill_itself_when_parent_died() - faulthandler.enable() - parent_process = psutil.Process().parent() + raise _WorkerSignalExit(signum) worker = None ready_pipe = kwargs.pop("ready_pipe") rank = kwargs.get("rank") + parent_process = None try: + # Keep all setup after handler installation inside this guarded + # region. The parent may terminate peer workers as soon as one + # worker fails, including while another peer is still starting. + # Either SIGTERM or SIGINT will terminate the worker. + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + kill_itself_when_parent_died() + faulthandler.enable() + parent_process = psutil.Process().parent() + worker = WorkerMultiprocProc(*args, **kwargs) # Send READY once we know everything is loaded @@ -578,6 +592,31 @@ class WorkerMultiprocProc: worker.worker_busy_loop() + except _WorkerSignalExit as exc: + # Raised by the SIGTERM/SIGINT handler installed above, so it is a + # BaseException and not an Exception: without this clause it walks + # straight past the handler below and the worker dies having reported + # nothing at all. The parent then sees only a closed pipe and raises + # "See stack trace for root cause" with no stack trace attached. + # + # Log with the traceback rather than a bare message. A signal handler + # runs on top of whatever the process was executing, so the frames + # here are the frames that were interrupted, which is the only clue + # to what the worker was doing when it was told to stop. + signal_name = signal.Signals(exc.signum).name + if exc.signum == signal.SIGINT: + logger.exception( + "Worker %d received %s (%d) while running. This normally means the user interrupted " + "the parent process. The stack below is where execution was interrupted, not the cause.", rank, + signal_name, exc.signum) + else: + logger.exception( + "Worker %d received %s (%d) while running. This can come from an external process, " + "such as an out-of-memory daemon, or from the parent cleaning up workers, including " + "after another worker failed. The stack below is where execution was interrupted, not " + "the cause.", rank, signal_name, exc.signum) + raise + except Exception as exc: if ready_pipe is not None: logger.exception("WorkerMultiprocProc failed to start.") From a575055eec1d93bf75737ed8de8c784f5af5b346 Mon Sep 17 00:00:00 2001 From: Yogya Mehrotra <67778074+yogyam@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:08:03 -0700 Subject: [PATCH 026/151] [perf] benchmark_attn_qat_train: device peak-TFLOPS table instead of silent RTX 5090 default (#1724) --- fastvideo-kernel/README.md | 7 +- .../benchmarks/benchmark_attn_qat_train.py | 29 +++++--- fastvideo-kernel/benchmarks/device_specs.py | 56 +++++++++++++++ fastvideo-kernel/tests/test_device_specs.py | 68 +++++++++++++++++++ 4 files changed, 148 insertions(+), 12 deletions(-) create mode 100644 fastvideo-kernel/benchmarks/device_specs.py create mode 100644 fastvideo-kernel/tests/test_device_specs.py diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index c3454c42c..3472d92c7 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -159,8 +159,11 @@ python benchmarks/benchmark_attn_qat_train.py ``` The benchmark reports both conventional attention FLOPs and the extra matrix -multiplications executed by the QAT straight-through path. Override -`--peak-tflops` when running on a GPU other than RTX 5090. +multiplications executed by the QAT straight-through path. It resolves dense +BF16 peak throughput for the full-GPU variants listed in +`benchmarks/device_specs.py`. Unknown or partitioned devices still report +timing and achieved TFLOPS, with MFU shown as `N/A`; pass a positive, finite +`--peak-tflops` value to report MFU for those devices. The QAT kernel is entirely Triton and routes by architecture at runtime. SM100 uses a large-tile forward and split 64x64 backward for the production diff --git a/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py b/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py index 9d7ff4c79..3f2f8afc5 100644 --- a/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py +++ b/fastvideo-kernel/benchmarks/benchmark_attn_qat_train.py @@ -14,9 +14,9 @@ from collections.abc import Callable import torch -from fastvideo_kernel.triton_kernels.attn_qat_train import attention +from device_specs import format_tflops_with_mfu, resolve_peak_tflops -RTX_5090_DENSE_BF16_TFLOPS = 209.5 +from fastvideo_kernel.triton_kernels.attn_qat_train import attention def _qat_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: @@ -62,14 +62,15 @@ def _format_result( timing_ms: tuple[float, float, float], algorithmic_flops: int, executed_matmul_flops: int, - peak_tflops: float, + peak_tflops: float | None, ) -> str: median_ms, min_ms, max_ms = timing_ms algorithmic_tflops = algorithmic_flops / (median_ms * 1e9) executed_tflops = executed_matmul_flops / (median_ms * 1e9) + algorithmic_result = format_tflops_with_mfu(algorithmic_tflops, peak_tflops) + executed_result = format_tflops_with_mfu(executed_tflops, peak_tflops) return (f"{label}: {median_ms:.3f} ms (min={min_ms:.3f}, max={max_ms:.3f}), " - f"algorithmic={algorithmic_tflops:.2f} TFLOPS/{100 * algorithmic_tflops / peak_tflops:.2f}% MFU, " - f"executed_matmul={executed_tflops:.2f} TFLOPS/{100 * executed_tflops / peak_tflops:.2f}% MFU") + f"algorithmic={algorithmic_result}, executed_matmul={executed_result}") def main() -> None: @@ -84,12 +85,16 @@ def main() -> None: parser.add_argument( "--peak-tflops", type=float, - default=RTX_5090_DENSE_BF16_TFLOPS, - help="Dense BF16 Tensor TFLOPS with FP32 accumulation; default is RTX 5090 boost-clock peak.", + default=None, + help="Dense BF16 Tensor TFLOPS with FP32 accumulation. Defaults to the " + "device table in device_specs.py. Unknown devices report MFU as N/A.", ) args = parser.parse_args() kv_length = args.kv_length or args.query_length + device_name = torch.cuda.get_device_name() + peak_tflops, peak_source = resolve_peak_tflops(device_name, args.peak_tflops) + torch.manual_seed(0) q_shape = (args.batch_size, args.heads, args.query_length, args.head_dim) kv_shape = (args.batch_size, args.heads, kv_length, args.head_dim) @@ -114,10 +119,14 @@ def main() -> None: # Conventional attention FLOPs are 4*base forward and 10*base backward. # QAT additionally computes the STE high-precision P@V path in forward and # the quantized-P dV path in backward, for 6*base and 14*base matmul FLOPs. - print(f"device: {torch.cuda.get_device_name()}") + print(f"device: {device_name}") + if peak_tflops is None: + print(f"peak: unavailable ({peak_source})") + else: + print(f"peak: {peak_tflops:.2f} dense BF16 TFLOPS ({peak_source})") print(f"q: {q_shape}; k/v: {kv_shape}; compile+first-forward: {compile_seconds:.3f} s") - print(_format_result("forward", forward_ms, 4 * base_flops, 6 * base_flops, args.peak_tflops)) - print(_format_result("backward", backward_ms, 10 * base_flops, 14 * base_flops, args.peak_tflops)) + print(_format_result("forward", forward_ms, 4 * base_flops, 6 * base_flops, peak_tflops)) + print(_format_result("backward", backward_ms, 10 * base_flops, 14 * base_flops, peak_tflops)) if __name__ == "__main__": diff --git a/fastvideo-kernel/benchmarks/device_specs.py b/fastvideo-kernel/benchmarks/device_specs.py new file mode 100644 index 000000000..b37fecc00 --- /dev/null +++ b/fastvideo-kernel/benchmarks/device_specs.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-device peak throughput used to turn benchmark timings into %MFU. + +Values are dense (non-sparse) BF16 Tensor Core TFLOPS with FP32 accumulation. +Only unambiguous full-GPU names with published peaks are listed. Timings and +achieved TFLOPS remain useful on other devices, so unknown devices report MFU +as unavailable unless the user supplies an explicit peak. +""" + +from __future__ import annotations + +import math + +# Exact torch.cuda.get_device_name() values for supported full GPUs. H100 form +# factors have different published peaks, and exact matching also prevents MIG, +# vGPU, or future variants from silently inheriting a full-GPU denominator. +DENSE_BF16_TFLOPS: dict[str, float] = { + "NVIDIA H100 80GB HBM3": 989.0, + "NVIDIA H100 NVL": 835.5, + "NVIDIA H100 PCIe": 756.5, + "NVIDIA GeForce RTX 5090": 209.5, # Matches the value this benchmark historically used. + "NVIDIA A100-SXM4-80GB": 312.0, + "NVIDIA A100-SXM4-40GB": 312.0, + "NVIDIA A100-PCIE-80GB": 312.0, + "NVIDIA A100-PCIE-40GB": 312.0, + "NVIDIA L40S": 362.05, +} + + +def resolve_peak_tflops(device_name: str, override: float | None = None) -> tuple[float | None, str]: + """Return (peak_tflops, source_description) for *device_name*. + + An explicit positive, finite *override* always wins. Unknown and partitioned + devices return ``None`` so the benchmark can still report timing and + achieved TFLOPS without presenting a misleading MFU value. + """ + if override is not None: + if not math.isfinite(override) or override <= 0: + raise ValueError("--peak-tflops must be finite and greater than zero") + return float(override), "explicit --peak-tflops override" + + normalized_name = device_name.strip() + if "MIG" in normalized_name.upper(): + return None, "MIG/partitioned device; pass --peak-tflops for this partition to report MFU" + + if normalized_name in DENSE_BF16_TFLOPS: + return DENSE_BF16_TFLOPS[normalized_name], f"exact device table entry {normalized_name!r}" + + return None, f"no dense BF16 peak for {device_name!r}; pass --peak-tflops to report MFU" + + +def format_tflops_with_mfu(achieved_tflops: float, peak_tflops: float | None) -> str: + """Format achieved throughput, including MFU only when its denominator is known.""" + if peak_tflops is None: + return f"{achieved_tflops:.2f} TFLOPS/MFU N/A" + return f"{achieved_tflops:.2f} TFLOPS/{100 * achieved_tflops / peak_tflops:.2f}% MFU" diff --git a/fastvideo-kernel/tests/test_device_specs.py b/fastvideo-kernel/tests/test_device_specs.py new file mode 100644 index 000000000..444b3081a --- /dev/null +++ b/fastvideo-kernel/tests/test_device_specs.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for benchmarks/device_specs.py (no torch or GPU required).""" + +import math +import os +import sys + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "benchmarks")) + +from device_specs import DENSE_BF16_TFLOPS, format_tflops_with_mfu, resolve_peak_tflops # noqa: E402 + + +@pytest.mark.parametrize( + "device_name", + [ + "NVIDIA GeForce RTX 5090", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 NVL", + "NVIDIA H100 PCIe", + "NVIDIA A100-SXM4-80GB", + "NVIDIA A100-PCIE-40GB", + "NVIDIA L40S", + ], +) +def test_known_full_gpus_resolve_from_table(device_name): + tflops, source = resolve_peak_tflops(device_name) + assert tflops == DENSE_BF16_TFLOPS[device_name] + assert device_name in source + + +def test_override_wins_even_for_known_device(): + tflops, source = resolve_peak_tflops("NVIDIA H100 80GB HBM3", override=123.4) + assert tflops == 123.4 + assert "override" in source + + +@pytest.mark.parametrize( + "device_name", + [ + "NVIDIA GB10", + "NVIDIA H100", + "NVIDIA H100 80GB HBM3 MIG 1g.10gb", + "NVIDIA A100-SXM4-80GB MIG 1g.10gb", + ], +) +def test_unknown_or_partitioned_device_keeps_benchmark_available(device_name): + tflops, source = resolve_peak_tflops(device_name) + assert tflops is None + assert "peak-tflops" in source + + +def test_unknown_device_with_override_resolves(): + tflops, source = resolve_peak_tflops("NVIDIA GB10", override=99.0) + assert tflops == 99.0 + assert "override" in source + + +@pytest.mark.parametrize("override", [0.0, -1.0, math.nan, math.inf, -math.inf]) +def test_override_must_be_positive_and_finite(override): + with pytest.raises(ValueError, match="finite and greater than zero"): + resolve_peak_tflops("NVIDIA GB10", override=override) + + +def test_format_tflops_reports_mfu_only_with_a_valid_denominator(): + assert format_tflops_with_mfu(123.456, None) == "123.46 TFLOPS/MFU N/A" + assert format_tflops_with_mfu(123.456, 200.0) == "123.46 TFLOPS/61.73% MFU" From 017ce6602d22022226307d16e006d1ece0f86d98 Mon Sep 17 00:00:00 2001 From: Yogya Mehrotra <67778074+yogyam@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:09:33 -0700 Subject: [PATCH 027/151] [bugfix] Fix dead NVFP4 QAT plumbing: bench imports, int8 error string, smooth_q crash (#1723) --- .../quantization/bench/bench_quant_k.py | 2 +- .../quantization/bench/bench_quant_q.py | 2 +- .../quantization/bench/bench_quant_v.py | 2 +- .../csrc/turbodiffusion/quant/quant.cu | 8 +++++-- .../triton_kernels/attn_qat_train.py | 13 ++++++++--- fastvideo-kernel/tests/test_attn_qat_train.py | 22 +++++++++++++++++++ .../tests/test_nvfp4_benchmark_imports.py | 21 ++++++++++++++++++ fastvideo-kernel/tests/test_turbodiffusion.py | 8 +++++++ 8 files changed, 70 insertions(+), 8 deletions(-) create mode 100644 fastvideo-kernel/tests/test_nvfp4_benchmark_imports.py diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py index 5fc97f5ee..e8360957c 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_k.py @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. """ import torch -import fp4quant +import fp4quant_cuda as fp4quant from triton.tools.mxfp import MXFP4Tensor from bench_utils import bench_kineto diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py index 1d0576d59..f05c292e3 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_q.py @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. """ import torch -import fp4quant +import fp4quant_cuda as fp4quant from triton.tools.mxfp import MXFP4Tensor from bench_utils import bench_kineto diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py index 0592aa888..a17885ec2 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py @@ -14,7 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. """ import torch -import fp4quant +import fp4quant_cuda as fp4quant from triton.tools.mxfp import MXFP4Tensor from bench_utils import bench_kineto diff --git a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu index 9a56e3881..672013047 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu +++ b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu @@ -31,6 +31,11 @@ auto quant( std::optional& Output_S ) { + TORCH_CHECK( + Input.scalar_type() == torch::kHalf || Input.scalar_type() == torch::kBFloat16, + "Unsupported input data type for quant_cuda (int8 quantization)." + ); + using ElementOut = int8_t; static constexpr int BlockSize = 128; static constexpr int NumThrPerCta = 256; @@ -62,8 +67,7 @@ auto quant( } default: { - std::cerr << "Observing: " << Input.scalar_type() << " for the input datatype which is invalid"; - throw std::runtime_error("Unsupported input data type for quantize_to_fp4."); + TORCH_CHECK(false, "Unsupported input data type for quant_cuda (int8 quantization)."); } } diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py index 52ef4294c..0fd0aedb9 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py @@ -1054,6 +1054,15 @@ class _attention(torch.autograd.Function): use_global_sf_P=True, use_global_sf_QKV=True, ): + if smooth_q: + # Not implemented, not merely disabled: only the backward kernels + # carry a SMOOTH_Q path (_attn_fwd has none), so restoring the + # triton_group_mean wiring would yield silently wrong gradients + # rather than a feature. The flag is kept for signature + # compatibility only. + raise NotImplementedError("smooth_q=True is not supported: the forward kernel has no " + "SMOOTH_Q path, so enabling it would produce incorrect gradients") + # shape constraints HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] # when v is in float8_e5m2 it is transposed. @@ -1364,10 +1373,8 @@ class _attention(torch.autograd.Function): BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM) + # smooth_q is rejected in forward(); ctx.smooth_q is always False here. q_m = None - if ctx.smooth_q: - # _, q_m = triton_group_mean(q) - q_m = q_m.repeat_interleave(q.shape[2] // q_m.shape[2], dim=2) # B,H,L,D sm100_optimized_backward = (getattr(ctx, "sm100_optimized", False) and ctx.use_qat_qkv_backward and not ctx.smooth_k and not ctx.smooth_q and N_CTX_KV % 16 == 0) diff --git a/fastvideo-kernel/tests/test_attn_qat_train.py b/fastvideo-kernel/tests/test_attn_qat_train.py index 22755e832..e221b50aa 100644 --- a/fastvideo-kernel/tests/test_attn_qat_train.py +++ b/fastvideo-kernel/tests/test_attn_qat_train.py @@ -4,6 +4,7 @@ Simple test for attn_qat_train. Tests forward and backward passes with and without QAT enabled. """ +import pytest import torch from fastvideo_kernel.triton_kernels.attn_qat_train import _attention from fastvideo_kernel.triton_kernels.fused_attention import attention as fused_attention @@ -13,6 +14,27 @@ attention = _attention.apply DEVICE = torch.device("cuda") +def test_smooth_q_fails_before_tensor_or_kernel_work(): + dummy = torch.empty((1, 1, 1, 16)) + + with pytest.raises(NotImplementedError, match="smooth_q=True is not supported"): + attention( + dummy, + dummy, + dummy, + False, + 1.0, + True, + True, + True, + True, + True, + True, + False, + True, + ) + + def attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None): """ Wrapper function that mimics attn_qat_train from the backend wrapper. diff --git a/fastvideo-kernel/tests/test_nvfp4_benchmark_imports.py b/fastvideo-kernel/tests/test_nvfp4_benchmark_imports.py new file mode 100644 index 000000000..44f932ade --- /dev/null +++ b/fastvideo-kernel/tests/test_nvfp4_benchmark_imports.py @@ -0,0 +1,21 @@ +import ast +from pathlib import Path + +import pytest + + +BENCH_DIR = Path(__file__).resolve().parents[1] / "attn_qat_infer" / "quantization" / "bench" + + +@pytest.mark.parametrize("script_name", ["bench_quant_q.py", "bench_quant_k.py", "bench_quant_v.py"]) +def test_quant_benchmark_imports_built_extension(script_name: str) -> None: + tree = ast.parse((BENCH_DIR / script_name).read_text()) + imports = { + (alias.name, alias.asname) + for node in ast.walk(tree) + if isinstance(node, ast.Import) + for alias in node.names + } + + assert ("fp4quant_cuda", "fp4quant") in imports + assert ("fp4quant", None) not in imports diff --git a/fastvideo-kernel/tests/test_turbodiffusion.py b/fastvideo-kernel/tests/test_turbodiffusion.py index 771ab7ffa..02ae43bdf 100644 --- a/fastvideo-kernel/tests/test_turbodiffusion.py +++ b/fastvideo-kernel/tests/test_turbodiffusion.py @@ -26,6 +26,14 @@ def rms_norm_ref(x, w, eps=1e-6): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") class TestTurboDiffusion: + def test_quant_invalid_dtype_names_int8_operation(self): + if turbodiffusion_ops.quant_cuda is None: + pytest.skip("quant_cuda not available") + + x = torch.ones((16, 128), dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match=r"Unsupported input data type for quant_cuda \(int8 quantization\)\."): + turbodiffusion_ops.quant_cuda(x, None, None) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("shape", [(16, 128), (32, 256), (1, 1024)]) def test_quant_correctness(self, dtype, shape): From 93aab45ac26e24512ace14632b7d2682f5d2f42b Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 25 Sep 2026 08:11:31 +0800 Subject: [PATCH 028/151] [bugfix] tests: include ltx2_3_base in local LTX-2 preset set (#1750) --- tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py index 48c0beae4..8261fd4b0 100644 --- a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py +++ b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py @@ -316,9 +316,9 @@ def test_ltx2_typed_surface_preflight() -> None: LTX2AudioDecodingStage, LTX2DenoisingStage, LTX2LatentPreparationStage, LTX2TextEncodingStage, ) - # All three LTX-2 presets registered. + # All four LTX-2 presets registered. names = {p.name for p in get_presets_for_family("ltx2")} - assert names == {"ltx2_base", "ltx2_distilled", "ltx2_two_stage"} + assert names == {"ltx2_base", "ltx2_3_base", "ltx2_distilled", "ltx2_two_stage"} # Two-stage preset has the denoise + refine topology and pulls its # refine allowed_overrides from the typed dataclass. From 9dd2a837f43d4f490ae5e7e8cadd504848821892 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 25 Sep 2026 08:13:23 +0800 Subject: [PATCH 029/151] [bugfix] encoders: keep Qwen2.5-VL importable on transformers 5 (#1791) Co-authored-by: SolitaryThinker --- .../models/encoders/qwen2_5_vl_custom.py | 12 ++++++++- .../encoders/test_qwen2_5_vl_vision_dtype.py | 27 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/fastvideo/models/encoders/qwen2_5_vl_custom.py b/fastvideo/models/encoders/qwen2_5_vl_custom.py index 64f3aef50..f69164653 100644 --- a/fastvideo/models/encoders/qwen2_5_vl_custom.py +++ b/fastvideo/models/encoders/qwen2_5_vl_custom.py @@ -26,11 +26,21 @@ from torch.distributed._tensor import DTensor from torch.distributed.device_mesh import DeviceMesh from torch.nn import CrossEntropyLoss from transformers.activations import ACT2FN -from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache +from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # transformers 5.0-5.5 dropped the cache-level class (5.6 restored the name as an alias for + # StaticCache); there a sliding window is a property of the cache's layers + # (DynamicSlidingWindowLayer / StaticSlidingWindowLayer), so no cache object can be an + # instance of it. isinstance() against an empty tuple is always False, which is the right + # answer there. + SlidingWindowCache = () # type: ignore[assignment,misc] + try: from torch.distributed.tensor import Shard except ImportError: diff --git a/fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py b/fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py index 6fabb1e79..19f04623b 100644 --- a/fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py +++ b/fastvideo/tests/encoders/test_qwen2_5_vl_vision_dtype.py @@ -95,3 +95,30 @@ def test_conditional_generation_passes_parent_dtype_to_visual_tower(): model = Qwen2_5_VLForConditionalGenerationSimple(_full_config("torch.bfloat16")) assert model.visual.dtype == torch.bfloat16 + + +def test_sliding_window_cache_check_survives_transformers_5(): + """transformers 5.0-5.5 removed `cache_utils.SlidingWindowCache`. + + The import here was unguarded and at module scope, so the whole encoder became + unimportable on those versions -- which `pyproject.toml` admits via `transformers>=5.0.0` + -- taking `Reason1TextEncoder`, and with it the Cosmos 2.5 and Kandinsky 5 pipelines, + with it. transformers 5.6 restored the name as an alias for `StaticCache`. + + Whatever the guard bound, `isinstance()` has to accept it and answer False for a + `DynamicCache`: a sliding window is a property of the cache's layers now, so no cache + object is an instance of the old class. + """ + from transformers import cache_utils + from transformers.cache_utils import DynamicCache + + from fastvideo.models.encoders import qwen2_5_vl_custom + + bound = qwen2_5_vl_custom.SlidingWindowCache + if hasattr(cache_utils, "SlidingWindowCache"): + # The guard must not shadow a name transformers still exports. + assert bound is cache_utils.SlidingWindowCache + else: + # The fallback arm: an empty tuple keeps isinstance() answering False. + assert bound == () + assert isinstance(DynamicCache(), bound) is False From 20751a21aa279888f278124c3f61e0195e5252d9 Mon Sep 17 00:00:00 2001 From: Jerry-XY <70183148+Jerry-XY@users.noreply.github.com> Date: Fri, 25 Sep 2026 08:16:49 +0800 Subject: [PATCH 030/151] [bugfix] Keep the VSA-H3 tile buffer out of the autograd graph (#1858) --- .buildkite/scripts/unit_test.sh | 1 + .../attention/backends/video_sparse_attn.py | 10 ++ .../backends/video_sparse_attn_h3.py | 23 ++- .../attention/test_vsa_h3_tile_grad_safety.py | 154 ++++++++++++++++++ 4 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 440082c4e..7fa4b3799 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -16,6 +16,7 @@ exec pytest \ ./fastvideo/tests/worker/ \ ./fastvideo/tests/training/test_trackers.py \ ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \ + ./fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py \ ./fastvideo/tests/modal/test_kernel_build_cache.py \ ./fastvideo/tests/modal/test_pr_test.py \ ./fastvideo/tests/modal/test_ssim_test.py \ diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index f90e70e54..8e6d5526a 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -182,6 +182,16 @@ def scatter_into_tile_buf( buffer's lifetime (per-metadata for Wan VSA, per-builder for VSA-H3) and its aliasing contract: the result is only valid until the next call with the same buffer. + + Callers that pass ``buf`` also own its *graph* contract: ``buf[:, dst] = x`` + is an autograd-tracked in-place write, so a non-grad destination becomes a + graph node as soon as ``x`` tracks gradients. A buffer that outlives the + fragment it was written from — VSA-H3's is owned by the metadata builder, + which serves the whole training run — then anchors that fragment, and the + activations its backward saved, for the rest of training. A caller inside a + grad-tracking forward should therefore pass ``buf=None`` unless the buffer + is scoped to that forward: VSA-H3 does so unconditionally, and Wan VSA + exposes the choice through ``cache_tile_buf`` (training defaults it off). """ if (buf is None or buf.shape != target_shape or buf.dtype != x.dtype or buf.device != x.device): buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index e5e6f227b..bcc5091bc 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -507,11 +507,13 @@ class MiniMaxH3VSAImpl(AttentionImpl): def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor: """Scatter rows into the padded tile buffer (pad positions stay zero). - The returned tensor aliases the builder-owned buffer; callers must - consume it before the next ``tile()`` (both call sites in - ``forward()`` read it immediately). Odd tile-64 no-grad sm100a - requests carry one additional all-zero tile internally; metadata and - all observable outputs retain the logical geometry. + Without grad tracking the returned tensor aliases the builder-owned + buffer; callers must consume it before the next ``tile()`` (both call + sites in ``forward()`` read it immediately). A grad-tracking forward + instead receives a fresh buffer and leaves the holder untouched, so + the builder never retains autograd state across steps. Odd tile-64 + no-grad sm100a requests carry one additional all-zero tile internally; + metadata and all observable outputs retain the logical geometry. """ if x.shape[1] != attn_metadata.total_seq_length: raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, " @@ -532,6 +534,17 @@ class MiniMaxH3VSAImpl(AttentionImpl): kernel_tiles = n_tiles + int(needs_sm100a_pair) target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1]) + # A grad-tracking forward must not reuse the builder-owned buffer. The + # holder outlives the training step -- one builder serves the whole run, + # and every step's metadata references the same holder -- while every + # VSA layer writes this one buffer in place. A graph-tracked tiled + # tensor left on the holder therefore anchors the step's in-place + # autograd edges, and through them the activations they saved, for the + # rest of training. This is the VSA-H3 counterpart of the Wan tile-cache + # OOM (#1423), which training fixed with ``vsa_cache_tile_buf=False``. + if grad_mode: + return scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, None) + # ``untile_combined_index`` maps each packed row to a logical tile # slot. Different geometries can share one transport shape; clear a # reused allocation once when the mapping identity changes so no old diff --git a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py new file mode 100644 index 000000000..30c9aa4b4 --- /dev/null +++ b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU regression: a grad-tracking VSA-H3 ``tile()`` must not leave the step's +graph on the builder-owned tile buffer. + +``MiniMaxH3VSAMetadataBuilder`` owns ``_tile_buf_holder`` and one builder serves +the whole training run, so every step's metadata references the same holder. +``tile()`` writes one shared buffer in place from every VSA layer, so if the +holder keeps the *graph-tracked* tiled tensor, that single tensor anchors all +of the step's in-place autograd edges -- and, through their ``next_functions``, +the activations they saved -- for the rest of training. This is the VSA-H3 +counterpart of the Wan tile-cache OOM (#1423, fixed by #1434/#1444); on the H3 +backend the buffer is owned by the builder rather than by per-step metadata, so +the retention spans steps instead of one denoising pass. + +Four tests, one property each: the holder never keeps graph state across +grad-tracking steps; the tiled rows still round-trip; inference still reuses the +buffer; and a builder alternating grad and non-grad steps corrupts neither. + +The holder invariant is implementation-agnostic -- it passes both for a fresh +buffer when gradients are tracked (nothing is cached) and for a reused buffer +that stores only a detached alias -- so it does not pin either fix. What the +first test does pin is that the fix may not be achieved by detaching the +*returned* tensor, which would silently starve the model of gradients. +""" + +import torch + +from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, MiniMaxH3VSAImpl, + MiniMaxH3VSAMetadataBuilder) + +_SPEC = dict(raw_latent_shape=(8, 8, 12), patch_size=(1, 2, 2), prefix_segments=(7, 5, 3)) + + +def _build(builder=None, step=0): + builder = builder if builder is not None else MiniMaxH3VSAMetadataBuilder() + # A fresh metadata per step, all sharing the builder's holder. + return builder.build( + current_timestep=step, + raw_latent_shape=_SPEC["raw_latent_shape"], + patch_size=_SPEC["patch_size"], + VSA_sparsity=0.9, + prefix_segments=_SPEC["prefix_segments"], + device=torch.device("cpu"), + tile_size=_TILE_ELEMS, + ), builder + + +def _impl(): + return MiniMaxH3VSAImpl(num_heads=2, head_size=8, causal=False, softmax_scale=8**-0.5) + + +def _assert_holder_is_graph_free(holder, where): + assert holder is not None, where + held = holder.buffer + # Caching nothing is fine; caching something that carries autograd state is + # what survives the step and anchors the graph. + assert held is None or (not held.requires_grad and held.grad_fn is None), ( + f"{where}: the builder-owned tile buffer holds a graph-tracked tensor; the " + "holder outlives the step, so it anchors every VSA layer's in-place " + "autograd edge (and the activations they saved) for the rest of training") + + +def test_grad_tracking_tile_keeps_the_holder_graph_free_across_steps(): + """Consecutive grad-tracking steps through one builder leave no graph behind. + + Step 0 additionally checks that the tiled tensor is genuinely connected to + the input, so the invariant cannot be met by detaching the returned tensor. + """ + _, builder = _build() + impl = _impl() + + for step in range(3): + meta, _ = _build(builder, step=step) + x = torch.randn(1, meta.total_seq_length, 2, 8, requires_grad=True) + tiled = impl.tile(x, meta) + + _assert_holder_is_graph_free(builder._tile_buf_holder, f"after grad step {step}") + + if step == 0: + # The tiled tensor itself is in the graph -- the VSA layers consume + # it and gradients must reach ``x``. A real backward, not just grad_fn. + assert tiled.requires_grad and tiled.grad_fn is not None + tiled.sum().backward() + assert x.grad is not None, "gradients must reach the tiled input" + assert bool((x.grad > 0).all()), "every packed row must be written at least once" + + +def test_grad_tracking_tile_still_populates_the_buffer(): + """The guard must not cost correctness: the tiled rows still round-trip.""" + meta, _ = _build() + x = torch.randn(1, meta.total_seq_length, 2, 8, requires_grad=True) + tiled = _impl().tile(x, meta) + assert torch.equal(tiled[:, meta.untile_combined_index], x) + + +def test_inference_tile_still_reuses_the_buffer(): + """The fix must not cost the inference optimization. + + With grad tracking off, the builder-owned buffer is still allocated once and + reused across calls -- which here means across the 50 VSA layers that share + one metadata. + """ + meta, _ = _build() + impl = _impl() + x = torch.randn(1, meta.total_seq_length, 2, 8) + + with torch.no_grad(): + first = impl.tile(x, meta) + assert meta.tile_buf_holder.buffer is not None, ("inference must still cache into the builder-owned buffer") + second = impl.tile(x, meta) + + assert first is second, "the builder-owned buffer must still be reused when gradients are not tracked" + assert torch.equal(second[:, meta.untile_combined_index], x) + + +def test_grad_and_inference_steps_share_one_builder_safely(): + """Grad-tracking and non-grad steps share a builder; neither may corrupt the other. + + Covers what the single-mode tests cannot: after a grad step has run, an + inference step must still reuse the builder's buffer, and a grad step after + inference must not overwrite that cached buffer with a graph-tracked one. + (The grad-step invariant itself is asserted by the first test.) + """ + builder = MiniMaxH3VSAMetadataBuilder() + impl = _impl() + + # A grad step first, on a builder that has never cached anything. + meta_grad, _ = _build(builder, step=0) + xg = torch.randn(1, meta_grad.total_seq_length, 2, 8, requires_grad=True) + impl.tile(xg, meta_grad) + + # Then an inference step: the buffer is allocated and reused as usual. + meta_inf, _ = _build(builder, step=1) + x = torch.randn(1, meta_inf.total_seq_length, 2, 8) + with torch.no_grad(): + first = impl.tile(x, meta_inf) + second = impl.tile(x, meta_inf) + assert first is second, "inference must still reuse the buffer after a grad step" + _assert_holder_is_graph_free(builder._tile_buf_holder, "after inference on a builder that saw a grad step") + + # Reverse order: a grad step after inference must not overwrite the cached + # (grad-free) inference buffer with a graph-tracked one. + meta_grad2, _ = _build(builder, step=2) + xg2 = torch.randn(1, meta_grad2.total_seq_length, 2, 8, requires_grad=True) + impl.tile(xg2, meta_grad2) + _assert_holder_is_graph_free(builder._tile_buf_holder, "after a grad step that followed inference") + + +if __name__ == "__main__": + test_grad_tracking_tile_keeps_the_holder_graph_free_across_steps() + test_grad_tracking_tile_still_populates_the_buffer() + test_inference_tile_still_reuses_the_buffer() + test_grad_and_inference_steps_share_one_builder_safely() + print("all VSA-H3 tile grad-safety checks passed") From 08d99c089e8d38503f104417d03ac66260b1d4dc Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:18:39 -0700 Subject: [PATCH 031/151] [perf] Reuse cached VAE offload for H3 conditioning (#1882) --- .../stages/minimax_h3_latent_preparation.py | 15 +- .../test_minimax_h3_encoding_offload.py | 210 ++++++++++++++++++ 2 files changed, 219 insertions(+), 6 deletions(-) create mode 100644 fastvideo/tests/stages/test_minimax_h3_encoding_offload.py diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index 60ffb2554..bd2482731 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -12,6 +12,7 @@ from diffusers.utils.torch_utils import randn_tensor from fastvideo.distributed import get_local_torch_device, get_sp_group, model_parallel_is_initialized from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models import pinned_offload from fastvideo.models.vaes.minimax_h3_parallel import encode_pixels_parallel from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_AUDIO_CHANNELS, @@ -182,6 +183,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): fastvideo_args: FastVideoArgs, device: torch.device, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Encode keyframes while reusing the video VAE's host-weight cache.""" keyframes = batch.extra.get(MINIMAX_H3_KEYFRAMES_KEY, []) if not isinstance(keyframes, list): raise TypeError("MiniMax-H3 keyframes must be a list.") @@ -192,7 +194,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): "FL2VA keyframes still need --video-decode-backend h3-vae.") vae_device = get_local_torch_device() - self.vae.to(vae_device) + pinned_offload.load(self.vae, vae_device, pin=fastvideo_args.pin_cpu_memory) try: clean_rows: list[torch.Tensor] = [] for image in keyframes: @@ -201,7 +203,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): h3_dit_patch_size(fastvideo_args))) finally: if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") + pinned_offload.unload(self.vae) _, _, latent_height, latent_width = _video_geometry(batch) shapes = ((1, latent_height, latent_width), ) * len(keyframes) @@ -227,27 +229,28 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): fastvideo_args: FastVideoArgs, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Encode reference media while reusing the video and audio VAE host caches.""" references = list(batch.references or []) if not references or not all(isinstance(item, MiniMaxH3PreparedReference) for item in references): raise TypeError("MiniMax-H3 Ref2VA latent preparation requires prepared references.") vae_device = get_local_torch_device() - self.vae.to(vae_device) + pinned_offload.load(self.vae, vae_device, pin=fastvideo_args.pin_cpu_memory) try: video_rows = self._encode_visual_rows(references, vae_device, fastvideo_args) finally: if fastvideo_args.vae_cpu_offload: - self.vae.to("cpu") + pinned_offload.unload(self.vae) audio_rows: list[torch.Tensor] = [] if any(reference.has_audio for reference in references): audio_device = get_local_torch_device() - self.audio_vae.to(audio_device) + pinned_offload.load(self.audio_vae, audio_device, pin=fastvideo_args.pin_cpu_memory) try: audio_rows = self._encode_audio_rows(references, audio_device) finally: if fastvideo_args.vae_cpu_offload: - self.audio_vae.to("cpu") + pinned_offload.unload(self.audio_vae) if not video_rows: raise ValueError("MiniMax-H3 Ref2VA requires at least one visual reference.") diff --git a/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py b/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py new file mode 100644 index 000000000..bce953796 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA regressions for shared H3 encoding and decoding host-weight caches.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from PIL import Image + +from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig, MiniMaxH3AudioVAEConfig +from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig +from fastvideo.distributed import get_local_torch_device +from fastvideo.models import pinned_offload +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAE +from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3 +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import MINIMAX_H3_KEYFRAMES_KEY +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MiniMaxH3LatentPreparationStage +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.fixture +def stage(distributed_setup, monkeypatch): + """Construct finite seeded VAEs with the tiny geometries used in parity tests.""" + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + torch.manual_seed(20260921) + video_arch = MiniMaxH3VideoVAEArchConfig( + latent_channels=4, + block_out_channels=(32, 32), + layers_per_block=1, + spatial_downsample_factors=(2, 2), + temporal_downsample_factors=(2, 2), + decoder_num_layers=1, + decoder_num_attention_heads=1, + decoder_attention_head_dim=8, + decoder_num_register_tokens=2, + decoder_ffn_mult=1, + latents_mean=(0.0, ) * 4, + latents_std=(1.0, ) * 4, + ) + video = AutoencoderKLMiniMaxH3(MiniMaxH3VideoVAEConfig( + arch_config=video_arch, use_tiling=False, use_temporal_tiling=False, use_parallel_tiling=False)) + audio_arch = MiniMaxH3AudioVAEArchConfig( + encoder_dim=8, + encoder_rates=(2, 2), + latent_dim=32, + latent_channels=4, + num_attention_heads=2, + decoder_dim=16, + decoder_rates=(2, 2), + decoder_kernel_sizes=(5, 4), + resblock_kernel_sizes=(3, ), + resblock_dilation_sizes=((1, 2), ), + latents_mean=[0.05 * index for index in range(4)], + latents_std=[1.0 + 0.1 * index for index in range(4)], + ) + audio = MiniMaxH3AudioVAE(MiniMaxH3AudioVAEConfig(arch_config=audio_arch)) + # ReplicatedLinear allocates uninitialized parameters without a checkpoint. + with torch.no_grad(): + for model in (video, audio): + for name, parameter in model.named_parameters(): + if parameter.ndim > 1: + torch.nn.init.xavier_uniform_(parameter) + elif name.endswith("weight") and "norm" in name: + parameter.fill_(1.0) + else: + parameter.normal_(0.0, 0.02) + model.float().eval() + return MiniMaxH3LatentPreparationStage(video, audio, MiniMaxH3Scheduler()) + + +def _args(pin: bool, offload: bool = True) -> SimpleNamespace: + """Provide encoding transfer flags and the tiny VAE's packing geometry.""" + return SimpleNamespace( + pin_cpu_memory=pin, + vae_cpu_offload=offload, + vae_parallel_encode=False, + pipeline_config=SimpleNamespace( + dit_config=SimpleNamespace(patch_size=(1, 1, 1)), + vae_config=SimpleNamespace(arch_config=SimpleNamespace(latent_channels=4)), + ), + ) + + +def _batch(mode: str) -> ForwardBatch: + """Rebuild consumed reference media and reset request noise for every encoding.""" + pixels = np.arange(16 * 16 * 3, dtype=np.uint8).reshape(16, 16, 3) + image = Image.fromarray(pixels) + batch = ForwardBatch(data_type="video", raw_latent_shape=(1, 4, 7, 4, 4), + generator=torch.Generator("cpu").manual_seed(7)) + if mode == "keyframe": + batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [image] + else: + batch.references = [ + MiniMaxH3PreparedReference(media_type="image", image=image), + MiniMaxH3PreparedReference(media_type="video", frames=np.tile(pixels[None], (22, 1, 1, 1))), + MiniMaxH3PreparedReference(media_type="audio", has_audio=True, + waveform=torch.linspace(-0.1, 0.1, 401).repeat(2, 1)), + ] + return batch + + +def _encode(stage, mode, args): + encode = stage._encode_fl2va_conditions if mode == "keyframe" else stage._encode_ref2va_conditions + return encode(_batch(mode), args, get_local_torch_device()) + + +def _models(stage, mode): + return (stage.vae, ) if mode == "keyframe" else (stage.vae, stage.audio_vae) + + +def _tensors(model): + yield from model.named_parameters() + for name, buffer in model.named_buffers(): + yield "buffer:" + name, buffer + + +def _assert_host_storage(model, pin, pointers=None): + """Check every registered tensor uses the same retained CPU allocation.""" + host = model._frozen_offload_host + assert set(host) == {name for name, _ in _tensors(model)} + for name, tensor in _tensors(model): + assert tensor.device.type == "cpu", name + assert tensor.data_ptr() == host[name].data_ptr(), name + assert tensor.is_pinned() == pin, name + observed = {name: tensor.data_ptr() for name, tensor in host.items()} + if pointers is not None: + assert observed == pointers + return observed + + +def _decode(stage, mode, args): + """Exercise decoding between requests using the same frozen-weight cache.""" + device = get_local_torch_device() + for model in _models(stage, mode): + pinned_offload.load(model, device, pin=args.pin_cpu_memory) + try: + if model is stage.vae: + latents = torch.zeros(1, 4, 2, 4, 4, device=device) + with torch.autocast("cuda", dtype=torch.float16): + decoded = model.decode(latents).sample + else: + decoded = model.decode(torch.zeros(2, 4, 12, device=device)).sample + assert torch.isfinite(decoded).all() + finally: + pinned_offload.unload(model) + + +@pytest.mark.parametrize("mode", ("keyframe", "references")) +@pytest.mark.parametrize("pin", (False, True)) +@torch.no_grad() +def test_encoding_reuses_decode_host_storage(stage, mode, pin): + """Match ordinary transfers and reuse one cache across encoding and decoding.""" + args = _args(pin) + with ( + patch.object(pinned_offload, "load", side_effect=lambda model, device, pin: model.to(device)), + patch.object(pinned_offload, "unload", side_effect=lambda model: model.to("cpu")), + ): + expected = _encode(stage, mode, args) + assert all(not hasattr(model, "_frozen_offload_host") for model in _models(stage, mode)) + + # Encoding must also initialize the cache when a request begins with images. + actual = _encode(stage, mode, args) + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + pointers = [_assert_host_storage(model, pin) for model in _models(stage, mode)] + _decode(stage, mode, args) + for _ in range(2): + actual = _encode(stage, mode, args) + for tensor in actual: + if tensor is not None: + assert torch.isfinite(tensor).all() + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + for model, storage in zip(_models(stage, mode), pointers): + _assert_host_storage(model, pin, storage) + + +@pytest.mark.parametrize("mode", ("keyframe", "references")) +@torch.no_grad() +def test_encoding_without_offload_keeps_cuda_weights(stage, mode): + """Preserve GPU residency without creating CPU mirrors when offload is disabled.""" + device = get_local_torch_device() + for model in _models(stage, mode): + model.to(device) + for _ in range(2): + _encode(stage, mode, _args(pin=True, offload=False)) + for model in _models(stage, mode): + assert all(tensor.device == device for _, tensor in _tensors(model)) + assert not hasattr(model, "_frozen_offload_host") + + +@pytest.mark.parametrize("branch", ("keyframe", "visual_reference", "audio_reference")) +@torch.no_grad() +def test_encoding_exception_restores_host_weights(stage, branch): + """Release CUDA weights through each encoding branch's finally block.""" + mode = "keyframe" if branch == "keyframe" else "references" + args = _args(pin=True) + _decode(stage, mode, args) + pointers = [_assert_host_storage(model, True) for model in _models(stage, mode)] + model = stage.audio_vae if branch == "audio_reference" else stage.vae + method = {"keyframe": "encode_keyframe", "visual_reference": "encode_pixels", "audio_reference": "encode"}[branch] + with patch.object(model, method, side_effect=RuntimeError("injected encoding failure")): + with pytest.raises(RuntimeError, match="injected encoding failure"): + _encode(stage, mode, args) + for model, storage in zip(_models(stage, mode), pointers): + _assert_host_storage(model, True, storage) From 76ce9c7fd666013d6f3f0e6465a678c2ed0eb59b Mon Sep 17 00:00:00 2001 From: YZJF Date: Thu, 24 Sep 2026 17:19:29 -0700 Subject: [PATCH 032/151] [bugfix] Reject mismatched model names on image generation routes (#1881) --- fastvideo/entrypoints/openai/image_api.py | 26 ++++- .../entrypoints/openai/request_adapter.py | 21 +++- .../test_image_model_name_validation.py | 103 ++++++++++++++++++ .../entrypoints/test_openai_serving_engine.py | 23 ++++ 4 files changed, 164 insertions(+), 9 deletions(-) create mode 100644 fastvideo/tests/entrypoints/test_image_model_name_validation.py diff --git a/fastvideo/entrypoints/openai/image_api.py b/fastvideo/entrypoints/openai/image_api.py index 350b90cb1..6162f2eb0 100644 --- a/fastvideo/entrypoints/openai/image_api.py +++ b/fastvideo/entrypoints/openai/image_api.py @@ -10,16 +10,22 @@ import aiofiles from fastapi import (APIRouter, File, Form, HTTPException, Path, Query, UploadFile) from fastapi.responses import FileResponse -from fastvideo.entrypoints.openai.state import ( - get_output_dir, - get_serving_engine, -) from fastvideo.entrypoints.openai.protocol import ( ImageGenerationsRequest, ImageResponse, ImageResponseData, generate_request_id, ) +from fastvideo.entrypoints.openai.request_adapter import ( + RequestAdaptationError, + validate_served_model_name, +) +from fastvideo.entrypoints.openai.state import ( + get_output_dir, + get_served_model_name, + get_server_args, + get_serving_engine, +) from fastvideo.entrypoints.openai.stores import IMAGE_STORE from fastvideo.entrypoints.openai.utils import ( choose_image_ext, @@ -43,6 +49,16 @@ def _normalize_response_format(value: str | None) -> str: return fmt +def _validate_request_model(model: str | None) -> None: + """Reject a model id that is not the one this server loaded.""" + if model is None: + return + try: + validate_served_model_name(model, get_server_args(), get_served_model_name()) + except RequestAdaptationError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + def _build_generation_kwargs( request_id: str, prompt: str, @@ -99,6 +115,7 @@ def _build_generation_kwargs( @router.post("", response_model=ImageResponse) async def generations(request: ImageGenerationsRequest): resp_format = _normalize_response_format(request.response_format) + _validate_request_model(request.model) request_id = generate_request_id() engine = get_serving_engine() @@ -181,6 +198,7 @@ async def edits( enable_teacache: bool | None = Form(False), ): resp_format = _normalize_response_format(response_format) + _validate_request_model(model) request_id = generate_request_id() engine = get_serving_engine() diff --git a/fastvideo/entrypoints/openai/request_adapter.py b/fastvideo/entrypoints/openai/request_adapter.py index e41355521..94a9b0737 100644 --- a/fastvideo/entrypoints/openai/request_adapter.py +++ b/fastvideo/entrypoints/openai/request_adapter.py @@ -160,6 +160,20 @@ def _apply_aspect_ratio( kwargs["width"], kwargs["height"] = width, height +def validate_served_model_name( + model: str | None, + args: FastVideoArgs, + served_model_name: str, +) -> None: + """Reject a request model id that is not the one this server loaded.""" + if model is None: + return + allowed_models = {args.lora_nickname} if args.lora_path else {served_model_name} + if model not in allowed_models: + choices = ", ".join(sorted(allowed_models)) + raise RequestAdaptationError(f"Model mismatch: request specifies {model!r}; this server provides {choices}.") + + def validate_model_and_lora( request: VideoGenerationRequest, args: FastVideoArgs, @@ -172,11 +186,7 @@ def validate_model_and_lora( loaded and cannot be swapped safely between concurrent requests. The API accepts vLLM's selector shape, but it must identify the startup adapter. """ - allowed_models = {args.lora_nickname} if args.lora_path else {served_model_name} - if request.model is not None and request.model not in allowed_models: - choices = ", ".join(sorted(allowed_models)) - raise RequestAdaptationError( - f"Model mismatch: request specifies {request.model!r}; this server provides {choices}.") + validate_served_model_name(request.model, args, served_model_name) if request.lora is None: return @@ -423,4 +433,5 @@ __all__ = [ "build_generation_request", "prepare_reference_media", "validate_model_and_lora", + "validate_served_model_name", ] diff --git a/fastvideo/tests/entrypoints/test_image_model_name_validation.py b/fastvideo/tests/entrypoints/test_image_model_name_validation.py new file mode 100644 index 000000000..2aa55f971 --- /dev/null +++ b/fastvideo/tests/entrypoints/test_image_model_name_validation.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject mismatched image model names before invoking the generator.""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from fastvideo.entrypoints.openai import image_api +from fastvideo.entrypoints.openai.stores import AsyncDictStore + +SERVED_MODEL_NAME = "my-image-model" + + +@pytest.fixture +def image_client(monkeypatch, tmp_path): + generate = Mock(side_effect=lambda **kwargs: Path(kwargs["output_path"]).write_bytes(b"test-image-content")) + + async def run_serialized(fn, **kwargs): + return fn(**kwargs) + + engine = SimpleNamespace( + generator=SimpleNamespace(generate_video=generate), + run_serialized=AsyncMock(side_effect=run_serialized), + ) + monkeypatch.setattr(image_api, "get_serving_engine", lambda: engine) + monkeypatch.setattr(image_api, "get_output_dir", lambda: str(tmp_path)) + monkeypatch.setattr(image_api, "get_served_model_name", lambda: SERVED_MODEL_NAME) + monkeypatch.setattr( + image_api, + "get_server_args", + lambda: SimpleNamespace(lora_path=None, lora_nickname="default"), + ) + monkeypatch.setattr(image_api, "IMAGE_STORE", AsyncDictStore()) + app = FastAPI() + app.include_router(image_api.router) + with TestClient(app) as client: + yield client, engine + + +@pytest.mark.parametrize("path", ["/v1/images", "/v1/images/generations"]) +def test_mismatched_model_does_not_generate(image_client, path, tmp_path): + client, engine = image_client + response = client.post(path, json={"prompt": "a cat", "model": "other-model"}) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Model mismatch: request specifies 'other-model'; this server provides my-image-model.") + engine.run_serialized.assert_not_awaited() + engine.generator.generate_video.assert_not_called() + assert not (tmp_path / "images").exists() + + +def test_edits_mismatched_model_does_not_generate(image_client, tmp_path): + client, engine = image_client + response = client.post( + "/v1/images/edits", + data={"prompt": "a cat", "model": "other-model"}, + files={"image": ("cat.png", b"fake-image", "image/png")}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Model mismatch: request specifies 'other-model'; this server provides my-image-model.") + engine.run_serialized.assert_not_awaited() + engine.generator.generate_video.assert_not_called() + assert not (tmp_path / "uploads").exists() + + +@pytest.mark.parametrize("path", ["/v1/images", "/v1/images/generations"]) +@pytest.mark.parametrize("model", [None, SERVED_MODEL_NAME]) +def test_omitted_or_matching_model_still_generates(image_client, path, model): + client, engine = image_client + payload = {"prompt": "a cat", "response_format": "b64_json"} + if model is not None: + payload["model"] = model + + response = client.post(path, json=payload) + + assert response.status_code == 200, response.text + engine.run_serialized.assert_awaited_once() + engine.generator.generate_video.assert_called_once() + + +@pytest.mark.parametrize("model", [None, SERVED_MODEL_NAME]) +def test_edits_omitted_or_matching_model_still_generates(image_client, model): + client, engine = image_client + data = {"prompt": "a cat", "response_format": "b64_json"} + if model is not None: + data["model"] = model + + response = client.post( + "/v1/images/edits", + data=data, + files={"image": ("cat.png", b"fake-image", "image/png")}, + ) + + assert response.status_code == 200, response.text + engine.run_serialized.assert_awaited_once() + engine.generator.generate_video.assert_called_once() diff --git a/fastvideo/tests/entrypoints/test_openai_serving_engine.py b/fastvideo/tests/entrypoints/test_openai_serving_engine.py index f833f78a7..1e412d6be 100644 --- a/fastvideo/tests/entrypoints/test_openai_serving_engine.py +++ b/fastvideo/tests/entrypoints/test_openai_serving_engine.py @@ -20,6 +20,7 @@ from fastvideo.entrypoints.openai.request_adapter import ( RequestAdaptationError, build_generation_request, prepare_reference_media, + validate_served_model_name, ) from fastvideo.entrypoints.openai.serving_engine import OpenAIServingEngine from fastvideo.entrypoints.openai.stores import VIDEO_STORE @@ -114,6 +115,28 @@ def test_protocol_rejects_unknown_and_client_output_fields() -> None: VideoGenerationRequest(prompt="a fox", output_path="/srv/keep.mp4") +def test_validate_served_model_name_accepts_omitted_or_matching_name() -> None: + args = _args("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + validate_served_model_name(None, args, "wan") + validate_served_model_name("wan", args, "wan") + + +def test_validate_served_model_name_rejects_mismatch() -> None: + with pytest.raises(RequestAdaptationError, match="this server provides wan"): + validate_served_model_name("other-model", _args("Wan-AI/Wan2.1-T2V-1.3B-Diffusers"), "wan") + + +def test_validate_served_model_name_uses_lora_nickname() -> None: + args = _args( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + lora_path="/models/startup.safetensors", + lora_nickname="fast", + ) + validate_served_model_name("fast", args, "wan") + with pytest.raises(RequestAdaptationError, match="this server provides fast"): + validate_served_model_name("wan", args, "wan") + + def test_request_adapter_restricts_extra_params(tmp_path: Path) -> None: request = VideoGenerationRequest(prompt="a fox", extra_params={"prompt_path": "/srv/prompts.txt"}) with pytest.raises(RequestAdaptationError, match="Unsupported extra_params fields: prompt_path"): From ba5e81083c7066ad84dc1f9d0599a1b12938d7a6 Mon Sep 17 00:00:00 2001 From: Lele Date: Fri, 25 Sep 2026 08:20:37 +0800 Subject: [PATCH 033/151] [bugfix] Preserve video frames when decoded audio is shorter (#1857) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 武垚乐 Co-authored-by: 武垚乐 --- fastvideo/entrypoints/video_generator.py | 9 +++++- .../tests/entrypoints/test_video_generator.py | 29 +++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 65b6d8d4e..fc47dafec 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -1246,7 +1246,14 @@ class VideoGenerator: "aac", "-pix_fmt", os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"), - "-shortest", + # Audio and video decoders may produce different durations. + # Preserve every video frame and pad/trim audio to match. + "-af", + "apad", + "-frames:v", + str(len(frames)), + "-t", + str(len(frames) / fps), "-movflags", "+faststart", output_path, diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 4ccb0c922..25a2a3881 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -1,3 +1,6 @@ +import json +import shutil +import subprocess import os from types import SimpleNamespace import warnings @@ -992,3 +995,29 @@ def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch): prompt=["first prompt", "second prompt"], inputs=InputConfig(image_path=["first.png"]), )) + + +@pytest.mark.parametrize("audio_seconds", [1, 2, 3]) +def test_ffmpeg_pipe_preserves_video_duration(tmp_path, monkeypatch, audio_seconds): + ffmpeg = shutil.which("ffmpeg") + ffprobe = shutil.which("ffprobe") + if ffmpeg is None or ffprobe is None: + pytest.skip("ffmpeg and ffprobe are required") + monkeypatch.setenv("FASTVIDEO_FFMPEG_BIN", ffmpeg) + monkeypatch.setenv("FASTVIDEO_VIDEO_CODEC", "libx264") + monkeypatch.setenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p") + fps, frame_count, sample_rate = 24, 48, 32000 + frames = [np.full((64, 64, 3), index * 4, dtype=np.uint8) for index in range(frame_count)] + audio = np.zeros((2, audio_seconds * sample_rate), dtype=np.float32) + output = tmp_path / "video.mp4" + + assert VideoGenerator._save_video_with_audio_ffmpeg_pipe(str(output), frames, fps, audio, sample_rate) + probe = json.loads(subprocess.check_output([ + ffprobe, "-v", "error", "-count_frames", "-show_entries", + "stream=codec_type,nb_read_frames,duration", "-of", "json", str(output), + ], text=True)) + video = next(stream for stream in probe["streams"] if stream["codec_type"] == "video") + sound = next(stream for stream in probe["streams"] if stream["codec_type"] == "audio") + assert int(video["nb_read_frames"]) == frame_count + assert float(video["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps) + assert float(sound["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps) From e90be598e56138af5c82590f0d72f6fa2dfce400 Mon Sep 17 00:00:00 2001 From: Keith Date: Fri, 25 Sep 2026 09:08:56 +0800 Subject: [PATCH 034/151] [perf] MiniMax H3: return uint8 frames from the decode worker (#1828) --- fastvideo/entrypoints/streaming_generator.py | 4 +- fastvideo/entrypoints/video_generator.py | 13 +++- .../minimax_h3/stages/minimax_h3_decoding.py | 7 ++ .../tests/entrypoints/test_video_generator.py | 27 +++++++ .../stages/test_minimax_h3_vae_streaming.py | 19 ++++- .../tests/train/callbacks/test_validation.py | 78 +++++++++++++++++++ fastvideo/train/callbacks/validation.py | 4 +- fastvideo/training/distillation_pipeline.py | 5 +- fastvideo/training/training_pipeline.py | 4 +- fastvideo/utils.py | 14 ++++ 10 files changed, 161 insertions(+), 14 deletions(-) diff --git a/fastvideo/entrypoints/streaming_generator.py b/fastvideo/entrypoints/streaming_generator.py index 64f1a08b8..b07ad20e8 100644 --- a/fastvideo/entrypoints/streaming_generator.py +++ b/fastvideo/entrypoints/streaming_generator.py @@ -13,7 +13,7 @@ from fastvideo.entrypoints.video_generator import VideoGenerator from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines import ForwardBatch -from fastvideo.utils import align_to, shallow_asdict +from fastvideo.utils import align_to, pixels_to_uint8, shallow_asdict from fastvideo.worker.executor import Executor from fastvideo.worker.multiproc_executor import MultiprocExecutor @@ -252,7 +252,7 @@ class StreamingVideoGenerator(VideoGenerator): for x in videos: x = torchvision.utils.make_grid(x, nrow=1) x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).cpu().numpy().astype(np.uint8)) + frames.append(pixels_to_uint8(x).cpu().numpy()) return frames diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index fc47dafec..410f67c54 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -55,7 +55,7 @@ from fastvideo.api.sampling_param import SamplingParam from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType from fastvideo.logger import init_logger from fastvideo.pipelines import ForwardBatch -from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback, align_to, shallow_asdict +from fastvideo.utils import align_to, allocate_cpu_tensor_with_pin_fallback, pixels_to_uint8, shallow_asdict from fastvideo.worker.executor import Executor fcntl: types.ModuleType | None @@ -846,11 +846,17 @@ class VideoGenerator: samples = output_batch.output.cpu() elif output_batch.output.shape == samples.shape: samples.copy_(output_batch.output) + if output_batch.output.dtype == torch.uint8: + # A worker may hand back uint8 pixels (the MiniMax-H3 decode + # stage does); keep ``samples`` on its [0, 1] float contract. + samples.div_(255) else: if not skip_pixel_prealloc: logger.warning("Output shape %s does not match expected shape %s; use slow path", output_batch.output.shape, samples.shape) samples = output_batch.output.cpu() + if samples.dtype == torch.uint8: + samples = samples.float().div_(255) logging_info = output_batch.logging_info gen_time = time.perf_counter() - start_time @@ -902,8 +908,9 @@ class VideoGenerator: # [0, 1] wrapped mod 256 in the old unclamped cast. # (Equivalence is SSIM-gated, not bit-exact: float->uint8 # differs <=1 LSB CPU vs GPU.) - src = output_batch.output - vid_u8 = (src * 255).clamp_(0, 255).to(torch.uint8) + # uint8 input is already quantized by the worker (MiniMax-H3 + # decode stage) and passes through untouched. + vid_u8 = pixels_to_uint8(output_batch.output) vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu() frames = [ torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8 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 65c1ed371..26037810d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -153,6 +153,13 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy) else: self.vae.decode_to_pixels(latents, output) + if is_output_rank and output is not None and output.dtype == torch.float32: + # Quantize to uint8 here, before the frames cross the + # multiprocess-executor boundary. A 345-frame 1344x768 clip is + # 4.27 GB as fp32 and 1.07 GB as uint8, and the main process + # cast it to uint8 immediately anyway; the clamp matches the + # post-decode handling of VAE output slightly outside [0, 1]. + output = output.mul_(255).clamp_(0, 255).to(torch.uint8) batch.output = output if is_output_rank else placeholder return batch finally: diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 25a2a3881..dd1f0185b 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -374,6 +374,33 @@ def test_generate_single_video_frames_clamp_out_of_range_pixels(tmp_path): assert (frames[1] == 0).all() +def test_generate_single_video_accepts_uint8_worker_output(tmp_path): + """A worker may quantize decoded pixels to uint8 before they cross the + executor boundary (the MiniMax-H3 decode stage does, cutting the IPC + payload 4x). Post-decode must pass those frames through untouched and + ``samples`` must stay on its [0, 1] float contract. CPU-only.""" + torch.manual_seed(0) + pixels = torch.rand((1, 3, 2, 16, 16), dtype=torch.float32) + output_u8 = (pixels * 255).clamp_(0, 255).to(torch.uint8) + output_batch = _single_video_output_batch(output_u8) + fastvideo_args = _single_video_args() + generator = _single_video_generator(output_batch, fastvideo_args) + + result = generator._generate_single_video( + prompt="uint8 worker output", + sampling_param=_small_sampling_param(save_video=False, return_frames=True), + fastvideo_args=fastvideo_args, + output_path=str(tmp_path / "unused.mp4"), + ) + + torch.testing.assert_close(result["samples"], output_u8.float() / 255) + frames = result["frames"] + assert len(frames) == 2 + for index, frame in enumerate(frames): + assert frame.dtype == np.uint8 + np.testing.assert_array_equal(frame, output_u8[0, :, index].permute(1, 2, 0).numpy()) + + def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path): output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5 output_batch = _single_video_output_batch(output) diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 2e605139f..74510139f 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -95,6 +95,10 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: observed["latents"] = decoded_latents observed["output"] = output output.fill_(0.25) + # One pixel per quantization case: an interior value, the exact + # top of the range, and the VAE overshoot on either side that the + # clamp must saturate instead of wrapping. + output[0, 0, 0, 0, :4] = torch.tensor([0.5, 1.0, 1.25, -0.25]) monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) result = MiniMaxH3VideoDecodingStage(VAE()).forward( @@ -109,9 +113,16 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: ) torch.testing.assert_close(observed["latents"], latents) - assert observed["output"] is result.output + # The VAE still streams into the float32 CPU buffer the stage allocates; + # the batch then carries that buffer's uint8 quantization so the + # executor boundary moves bytes, not fp32. + assert observed["output"].device.type == "cpu" + assert observed["output"].dtype == torch.float32 assert result.output.device.type == "cpu" - assert torch.all(result.output == 0.25) + assert result.output.dtype == torch.uint8 + expected = torch.full((1, 3, 5, 16, 16), 63, dtype=torch.uint8) + expected[0, 0, 0, 0, :4] = torch.tensor([127, 255, 255, 0], dtype=torch.uint8) + assert torch.equal(result.output, expected) def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: @@ -232,7 +243,9 @@ def test_parallel_decode_runs_on_every_rank(monkeypatch) -> None: result = MiniMaxH3VideoDecodingStage(VAE()).forward(batch, args) if is_first: assert result.output.shape == (1, 3, 5, 16, 16) - assert torch.all(result.output == 0.5) + assert result.output.device.type == "cpu" + assert result.output.dtype == torch.uint8 + assert torch.all(result.output == 127) else: assert result.output.shape == (0, 3, 0, 0, 0) diff --git a/fastvideo/tests/train/callbacks/test_validation.py b/fastvideo/tests/train/callbacks/test_validation.py index e8baf8d1d..266008ae3 100644 --- a/fastvideo/tests/train/callbacks/test_validation.py +++ b/fastvideo/tests/train/callbacks/test_validation.py @@ -29,6 +29,7 @@ from fastvideo.train.callbacks.validation import ( SYNTHETIC_OPTICAL_FLOW_METRIC, ValidationCallback, _ValidationMetricStats, + _ValidationStepResult, ) # --------------------------------------------------------------------------- @@ -388,6 +389,83 @@ class TestH3ValidationContract: assert torch.is_tensor(waveform) torch.testing.assert_close(waveform, torch.full((32, 2), float(prompt_index))) + def _run_validation_for_steps_with_output( + self, + monkeypatch: pytest.MonkeyPatch, + output: torch.Tensor, + ) -> _ValidationStepResult: + """Run one validation prompt through a pipeline returning ``output``.""" + + class FakeH3Pipeline: + """Return the given decoded pixels for every prompt.""" + + fastvideo_args = SimpleNamespace(pipeline_config=SimpleNamespace()) + + def forward(self, batch, inference_args): + return SimpleNamespace(output=output, extra={}) + + cb = _make_callback() + cb.training_config = SimpleNamespace( + model_path="unused", + pipeline_config=SimpleNamespace(), + ) + cb.rank_in_sp_group = 0 + pipeline = FakeH3Pipeline() + monkeypatch.setattr( + "fastvideo.train.callbacks.validation.ValidationDataset", + lambda filename: [{"caption": "prompt-0"}], + ) + monkeypatch.setattr( + "fastvideo.train.callbacks.validation.make_inference_args", + lambda *args, **kwargs: SimpleNamespace( + pipeline_config=SimpleNamespace(dmd_denoising_steps=None), + dit_cpu_offload=True, + ), + ) + monkeypatch.setattr(cb, "_get_pipeline", lambda *, transformer: pipeline) + monkeypatch.setattr(cb, "_get_sampling_param", SamplingParam) + monkeypatch.setattr( + cb, + "_prepare_validation_batch", + lambda sampling_param, validation_batch, num_inference_steps: SimpleNamespace(prompt=validation_batch[ + "caption"], ), + ) + return cb._run_validation_for_steps(50, transformer=torch.nn.Identity()) + + def test_run_validation_for_steps_passes_uint8_frames_through( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Verify worker-quantized uint8 pixels reach the frames byte for byte. + + The MiniMax-H3 decode stage returns uint8; scaling it by 255 again + would wrap modulo 256 (63 -> 193, 127 -> 129, 191 -> 65, 255 -> 1). + """ + pixels = torch.tensor([63, 127, 191, 255], dtype=torch.uint8).view(1, 1, 1, 2, 2).expand(1, 3, 2, 2, 2) + + result = self._run_validation_for_steps_with_output(monkeypatch, pixels.contiguous()) + + assert len(result.videos) == 1 + assert len(result.videos[0]) == 2 + for frame in result.videos[0]: + assert frame.dtype == np.uint8 + np.testing.assert_array_equal(frame, pixels[0, :, 0].permute(1, 2, 0).numpy()) + + def test_run_validation_for_steps_quantizes_float_frames( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Verify normalized float pixels still produce the same bytes.""" + pixels = torch.tensor([0.25, 0.5, 0.75, 1.0]).view(1, 1, 1, 2, 2).expand(1, 3, 2, 2, 2) + + result = self._run_validation_for_steps_with_output(monkeypatch, pixels.contiguous()) + + expected = torch.tensor([63, 127, 191, 255], dtype=torch.uint8).view(2, 2, 1).expand(2, 2, 3).numpy() + assert len(result.videos[0]) == 2 + for frame in result.videos[0]: + assert frame.dtype == np.uint8 + np.testing.assert_array_equal(frame, expected) + def test_log_validation_video_artifacts_emits_eight_wandb_videos(self) -> None: """Verify a MiniMax H3 event contains eight media objects and scalar evidence.""" diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index e83775ed0..203c4252e 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -38,7 +38,7 @@ from fastvideo.train.utils.moduleloader import ( make_inference_args, ) from fastvideo.train.utils.validation_media import write_validation_mp4 from fastvideo.training.trackers import DummyTracker -from fastvideo.utils import shallow_asdict +from fastvideo.utils import pixels_to_uint8, shallow_asdict if TYPE_CHECKING: from fastvideo.train.methods.base import TrainingMethod @@ -1565,7 +1565,7 @@ class ValidationCallback(Callback): nrow=6, ) x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1)) - frames.append((x * 255).numpy().astype(np.uint8)) + frames.append(pixels_to_uint8(x).numpy()) videos.append(frames) captions.append(batch.prompt) audio_waveforms.append(output_audio) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 0933b1861..2f9e86009 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -34,7 +34,8 @@ from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import (EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep) -from fastvideo.utils import (is_vsa_available, maybe_download_model, set_random_seed, verify_model_config_and_directory) +from fastvideo.utils import (is_vsa_available, maybe_download_model, pixels_to_uint8, set_random_seed, + verify_model_config_and_directory) try: vsa_available = is_vsa_available() @@ -1114,7 +1115,7 @@ class DistillationPipeline(TrainingPipeline): for x in video: x = torchvision.utils.make_grid(x, nrow=6) x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) + frames.append(pixels_to_uint8(x).numpy()) videos.append(frames) audios.append(output_batch.extra.get("audio")) audio_sample_rates.append(output_batch.extra.get("audio_sample_rate")) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 89779ad4c..97c66ecf7 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -43,7 +43,7 @@ from fastvideo.training.trackers import (DummyTracker, TrackerType, initialize_t from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases, compute_density_for_timestep_sampling, count_trainable, get_scheduler, get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint) -from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict) +from fastvideo.utils import (is_vmoba_available, is_vsa_available, pixels_to_uint8, set_random_seed, shallow_asdict) try: vsa_available = is_vsa_available() @@ -784,7 +784,7 @@ class TrainingPipeline(LoRAPipeline, ABC): for x in video: x = torchvision.utils.make_grid(x, nrow=6) x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) + frames.append(pixels_to_uint8(x).numpy()) step_videos.append(frames) # Only sp_group leaders (rank_in_sp_group == 0) need to send their diff --git a/fastvideo/utils.py b/fastvideo/utils.py index af4d2acda..4f31b60f6 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -1021,6 +1021,20 @@ def best_output_size(w, h, dw, dh, expected_area): return ow2, oh2 +def pixels_to_uint8(pixels: torch.Tensor) -> torch.Tensor: + """Return decoded pixels as uint8 bytes. + + Decode stages hand back either normalized float pixels in [0, 1] or uint8 + that a worker already quantized before the executor boundary (the + MiniMax-H3 video decode stage does). uint8 passes through untouched; + scaling it by 255 again would wrap modulo 256. Float pixels are clamped so + VAE output slightly outside [0, 1] saturates instead of wrapping. + """ + if pixels.dtype == torch.uint8: + return pixels + return (pixels * 255).clamp_(0, 255).to(torch.uint8) + + def save_decoded_latents_as_video(decoded_latents: list[torch.Tensor], output_path: str, fps: int): # Process outputs videos = rearrange(decoded_latents, "b c t h w -> t b c h w") From dd35763ad6ba8db441fdccbe05f66fba797b8931 Mon Sep 17 00:00:00 2001 From: Max LI <114852914+2001maxli@users.noreply.github.com> Date: Sun, 27 Sep 2026 14:53:25 -0700 Subject: [PATCH 035/151] [bugfix] Fix MLX prompt enhancement sampler compatibility (#1891) --- .github/workflows/ci-macos-mlx.yml | 2 ++ fastvideo/mlx_runtime/prompt_enhance.py | 3 +- .../tests/mlx/test_mlx_prompt_enhance.py | 32 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml index 62054c839..3ef19a311 100644 --- a/.github/workflows/ci-macos-mlx.yml +++ b/.github/workflows/ci-macos-mlx.yml @@ -92,6 +92,7 @@ jobs: fastvideo/tests/mlx/test_frame_upsample.py \ fastvideo/tests/mlx/test_mlx_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_refine.py \ + fastvideo/tests/mlx/test_mlx_prompt_enhance.py \ fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ fastvideo/tests/mlx/test_wan22_sample.py \ @@ -156,6 +157,7 @@ jobs: fastvideo/tests/mlx/test_frame_upsample.py \ fastvideo/tests/mlx/test_mlx_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_refine.py \ + fastvideo/tests/mlx/test_mlx_prompt_enhance.py \ fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ fastvideo/tests/mlx/test_wan22_sample.py \ diff --git a/fastvideo/mlx_runtime/prompt_enhance.py b/fastvideo/mlx_runtime/prompt_enhance.py index 60a4ed4bb..78e95e35d 100644 --- a/fastvideo/mlx_runtime/prompt_enhance.py +++ b/fastvideo/mlx_runtime/prompt_enhance.py @@ -194,6 +194,7 @@ def enhance_prompt_mlx_lm( """ try: from mlx_lm import generate, load + from mlx_lm.sample_utils import make_sampler except ImportError as exc: # pragma: no cover - optional dep raise RuntimeError("mlx-lm is not installed. `uv pip install mlx-lm` or use " "--enhance-prompt-backend template.") from exc @@ -226,7 +227,7 @@ def enhance_prompt_mlx_lm( tokenizer, prompt=chat, max_tokens=max_tokens, - temp=temp, + sampler=make_sampler(temp=temp), verbose=False, ) enhanced = _clean_llm_output(raw, original=text) diff --git a/fastvideo/tests/mlx/test_mlx_prompt_enhance.py b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py index 51bcd5992..1279573cd 100644 --- a/fastvideo/tests/mlx/test_mlx_prompt_enhance.py +++ b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py @@ -4,6 +4,10 @@ from __future__ import annotations import json +import sys +from types import ModuleType, SimpleNamespace + +import pytest from fastvideo.mlx_runtime.prompt_enhance import ( DEFAULT_ENHANCE_SYSTEM_PROMPT, @@ -14,6 +18,34 @@ from fastvideo.mlx_runtime.prompt_enhance import ( ) +@pytest.mark.parametrize("backend", ["mlx-lm", "auto"]) +def test_enhancement_uses_available_mlx_lm(monkeypatch, backend): + """Both backends must reach generation through the sampler-based API.""" + package = ModuleType("mlx_lm") + sampling = ModuleType("mlx_lm.sample_utils") + temperatures = [] + expected_sampler = object() + enhanced = "A cat walking on grass in warm evening light." + + def make_sampler(temp): + temperatures.append(temp) + return expected_sampler + + def generate(model_arg, tokenizer_arg, *, prompt, max_tokens, sampler, verbose): + assert sampler is expected_sampler + return enhanced + + package.load = lambda path: (object(), SimpleNamespace()) + package.generate = generate + sampling.make_sampler = make_sampler + monkeypatch.setitem(sys.modules, "mlx_lm", package) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", sampling) + result = enhance_prompt("A cat walking on grass", backend=backend) + assert result.backend == "mlx-lm" + assert result.enhanced == enhanced + assert temperatures == [0.6] + + def test_template_adds_cinematic_cues_to_thin_prompt() -> None: raw = "A red fox in the snow" out = enhance_prompt_template(raw) From 442e2d2e18ed63b9a5d9f5bb6233ed5a884b3b9b Mon Sep 17 00:00:00 2001 From: li-lizhe <82501526+li-lizhe@users.noreply.github.com> Date: Tue, 29 Sep 2026 01:24:00 +0800 Subject: [PATCH 036/151] [bugfix] fix(metrics): set SceneMetric device_map for any non-CPU device (#1817) --- fastvideo/eval/metrics/vbench/scene/metric.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastvideo/eval/metrics/vbench/scene/metric.py b/fastvideo/eval/metrics/vbench/scene/metric.py index d8dc70354..77c3fe4ab 100644 --- a/fastvideo/eval/metrics/vbench/scene/metric.py +++ b/fastvideo/eval/metrics/vbench/scene/metric.py @@ -63,7 +63,7 @@ class SceneMetric(BaseMetric): self._model = Qwen2_5OmniForConditionalGeneration.from_pretrained( self._model_path, torch_dtype=torch.bfloat16, - device_map=str(self.device) if self.device.type == "cuda" else None, + device_map=str(self.device) if self.device.type != "cpu" else None, ) self._model.disable_talker() self._model.eval() From cb66acd40065731f8fa1a6d750dc04679ba97756 Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Tue, 29 Sep 2026 11:48:31 -0700 Subject: [PATCH 037/151] [bugfix] Fix environment-variable bugs in LTX-2 debug logging, HF token lookup, and attention backend reads (#1895) --- fastvideo/attention/AGENTS.md | 7 + fastvideo/attention/selector.py | 38 +++++ fastvideo/models/dits/ltx2.py | 142 +----------------- fastvideo/models/encoders/gemma.py | 27 +--- fastvideo/models/encoders/t5gemma.py | 4 +- fastvideo/models/vaes/oobleck.py | 3 +- fastvideo/models/wan/transformer.py | 9 +- fastvideo/performance/hf_store.py | 14 +- .../pipelines/basic/ltx2/ltx2_pipeline.py | 30 ---- .../basic/ltx2/stages/ltx2_denoising.py | 15 +- .../basic/magi_human/magi_human_pipeline.py | 21 ++- fastvideo/pipelines/stages/kandinsky5.py | 20 +-- .../api/test_attention_selector_resolution.py | 33 ++++ ...test_kandinsky5_attention_backend_guard.py | 32 ++-- fastvideo/utils.py | 28 ++-- fastvideo/worker/ray_distributed_executor.py | 7 +- .../ltx2/test_ltx2_pipeline_smoke.py | 10 -- 17 files changed, 160 insertions(+), 280 deletions(-) diff --git a/fastvideo/attention/AGENTS.md b/fastvideo/attention/AGENTS.md index 9a33843a4..dc3d186df 100644 --- a/fastvideo/attention/AGENTS.md +++ b/fastvideo/attention/AGENTS.md @@ -49,6 +49,13 @@ A loader may narrow the request for one component — the DMD teacher/critic transformers build dense — and the recorded value is what that component actually resolved, not what the run asked for globally. +A component that chooses its own structure by backend (Wan's VSA block, +LTX-2's distributed attention, the SDPA switch in the Gemma and T5-Gemma +encoders) calls `effective_attention_backend(self.config)` instead of reading +`_resolved_attention_backend` directly. It returns the recorded decision, the +active scope's rule, or, for a component constructed without a loader, the +environment variable — the same order its attention layers follow. + A call site that already knows its component passes the decision explicitly: ```python diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 93924cf84..03df0b499 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -174,6 +174,44 @@ def component_attention_backend(component: object) -> AttentionBackendEnum | _No return NO_REQUEST if resolved is None else resolved +def effective_attention_backend(config: object) -> AttentionBackendEnum | None: + """Return the attention backend that a component built from ``config`` follows. + + Some components choose part of their own structure by backend: Wan picks + its VSA transformer block, LTX-2 its distributed-attention path, and the + Gemma and T5-Gemma text encoders switch Hugging Face attention to SDPA. + That choice must match the attention layers inside the component, and + those layers resolve through :func:`get_attn_backend`. This function + applies the same order: + + 1. The decision a loader recorded on ``config`` with + :func:`record_resolved_attention_backend`. Every component built by a + loader has one when a backend was requested. + 2. While a loader is building a component (a + ``_component_attention_backend_scope`` is active), that scope's rule: + its backend, and the environment variable only when the scope consults + it. A scope that deliberately ignores the environment variable, such as + the dense DMD teacher and critic, keeps ignoring it here. + 3. With no loader involved, the ``FASTVIDEO_ATTENTION_BACKEND`` + environment variable. This covers components constructed directly, + such as MagiHuman's lazily built T5-Gemma encoder or a script that + builds a transformer itself. + + Reading only the recorded decision is not enough: a directly constructed + component has no recorded decision (``None``), so it would ignore the + environment variable while its attention layers still follow it, and the + component's structure would not match its layers. + """ + recorded = getattr(config, "_resolved_attention_backend", None) + if recorded is not None: + return recorded + scope = _SCOPE.get() + if scope is not None and (scope.backend is not None or not scope.consult_env): + return scope.backend + env_backend = envs.FASTVIDEO_ATTENTION_BACKEND + return None if env_backend is None else backend_name_to_enum(env_backend) + + def get_attn_backend( head_size: int, dtype: torch.dtype, diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index 2ef569648..e2cbf3ce2 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -7,8 +7,6 @@ from dataclasses import dataclass, replace from enum import Enum import functools import math -import os -from pathlib import Path from typing import Any, Optional, Tuple, Callable import torch @@ -17,6 +15,7 @@ from einops import rearrange, repeat from fastvideo.attention.backends.sdpa import SDPAMetadata from fastvideo.attention.layer import DistributedAttention, LocalAttention +from fastvideo.attention.selector import effective_attention_backend from fastvideo.configs.models.dits import LTX2VideoConfig from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather, @@ -646,48 +645,6 @@ def _to_denoised( return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) -def _debug_block_log_line(message: str) -> None: - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1": - return - log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "") - if not log_path: - return - log_dir = os.path.dirname(log_path) - if log_dir: - os.makedirs(log_dir, exist_ok=True) - with open(log_path, "a", encoding="utf-8") as f: - f.write(message + "\n") - - -def _debug_transformer_args(prefix: str, args: "TransformerArgs | None") -> None: - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1" or args is None: - return - pe_cos, pe_sin = args.positional_embeddings - cross_cos = None - cross_sin = None - if args.cross_positional_embeddings is not None: - cross_cos, cross_sin = args.cross_positional_embeddings - mask = args.context_mask - if mask is None: - mask_summary = "mask=None" - else: - finite = torch.isfinite(mask) - finite_sum = mask[finite].sum().item() if finite.any() else 0.0 - mask_summary = (f"mask_min={mask.min().item():.6f} " - f"mask_max={mask.max().item():.6f} " - f"mask_finite_sum={finite_sum:.6f} " - f"mask_finite_count={finite.sum().item()}") - _debug_block_log_line(f"{prefix}:x_sum={args.x.float().sum().item():.6f} " - f"context_sum={args.context.float().sum().item():.6f} " - f"t_sum={args.timesteps.float().sum().item():.6f} " - f"emb_sum={args.embedded_timestep.float().sum().item():.6f} " - f"pe_cos_sum={pe_cos.float().sum().item():.6f} " - f"pe_sin_sum={pe_sin.float().sum().item():.6f} " - f"cross_pe_cos_sum={(cross_cos.float().sum().item() if cross_cos is not None else 0.0):.6f} " - f"cross_pe_sin_sum={(cross_sin.float().sum().item() if cross_sin is not None else 0.0):.6f} " - f"{mask_summary}") - - class LTXRopeType(Enum): """LTX-2 rotary variants (interleaved vs split).""" INTERLEAVED = "interleaved" @@ -2262,15 +2219,6 @@ class BasicAVTransformerBlock(torch.nn.Module): ax_scaled = _rms_norm_dispatch(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp ax = ax + self.audio_ff(ax_scaled) * agate_mlp - # Debug-only: reading ``self.idx`` (and .item() syncs) here means - # enabling this env var re-specializes the compiled graph per block, - # defeating the shared-graph regional compilation. - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - video_sum = vx.float().sum().item() if vx is not None else 0.0 - audio_sum = ax.float().sum().item() if ax is not None else 0.0 - _debug_block_log_line(f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} " - f"audio_sum={audio_sum:.6f}") - # Register FSDP2 backward hooks on output tensors (module-level hooks don't # fire for dataclass outputs, so we must hook the tensors directly) self._register_fsdp_backward_hooks_on_output(vx, ax) @@ -2680,12 +2628,6 @@ class LTXModel(torch.nn.Module): skip_audio_self_attn_blocks: Block indices where audio self-attention is skipped (STG perturbed pass). """ - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_block_log_line("fastvideo:patchify_proj" - f":video_w_sum={self.patchify_proj.weight.float().sum().item():.6f} " - f"video_b_sum={self.patchify_proj.bias.float().sum().item():.6f} " - f"audio_w_sum={self.audio_patchify_proj.weight.float().sum().item():.6f} " - f"audio_b_sum={self.audio_patchify_proj.bias.float().sum().item():.6f}") if not self.model_type.is_video_enabled() and video is not None: raise ValueError("Video is not enabled for this model") if not self.model_type.is_audio_enabled() and audio is not None: @@ -2693,8 +2635,6 @@ class LTXModel(torch.nn.Module): video_args = self.video_args_preprocessor.prepare(video) if video is not None else None audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None - _debug_transformer_args("fastvideo:prep_video", video_args) - _debug_transformer_args("fastvideo:prep_audio", audio_args) video_out, audio_out = self._process_transformer_blocks( video_args, audio_args, @@ -2735,7 +2675,7 @@ class LTX2Transformer3DModel(BaseDiT): # Get SP world size for distributed attention sp_world_size = get_sp_world_size() - use_vsa_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND", "") == "VIDEO_SPARSE_ATTN" + use_vsa_backend = effective_attention_backend(config) == AttentionBackendEnum.VIDEO_SPARSE_ATTN use_distributed_attention = sp_world_size > 1 or use_vsa_backend # Validate that attention heads are divisible by SP world size @@ -2798,62 +2738,6 @@ class LTX2Transformer3DModel(BaseDiT): self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.num_channels_latents - if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1": - detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "") - if detail_path: - self._attach_debug_detail_hooks(detail_path) - - def _attach_debug_detail_hooks(self, log_path: str) -> None: - path = Path(log_path) - path.parent.mkdir(parents=True, exist_ok=True) - if path.exists(): - path.unlink() - - def _format_sum(tensor: torch.Tensor | None) -> str: - if tensor is None: - return "None" - return f"{tensor.float().sum().item():.6f}" - - def _hook_factory(block_idx: int, name: str): - - def _hook(_module, _inputs, outputs): # noqa: ANN001 - out = outputs[0] if isinstance(outputs, tuple) else outputs - out_sum = _format_sum(out if torch.is_tensor(out) else None) - with path.open("a", encoding="utf-8") as f: - f.write(f"fastvideo:{block_idx}:{name}:out_sum={out_sum}\n") - - return _hook - - for block in self.model.transformer_blocks: - idx = block.idx - for name in ( - "attn1", - "attn2", - "ff", - "audio_attn1", - "audio_attn2", - "audio_ff", - "audio_to_video_attn", - "video_to_audio_attn", - ): - if hasattr(block, name): - getattr(block, name).register_forward_hook(_hook_factory(idx, name)) - - def _output_hook(label: str): - - def _hook(_module, _inputs, outputs): # noqa: ANN001 - out = outputs[0] if isinstance(outputs, tuple) else outputs - out_sum = _format_sum(out if torch.is_tensor(out) else None) - with path.open("a", encoding="utf-8") as f: - f.write(f"fastvideo:output:{label}:out_sum={out_sum}\n") - - return _hook - - if hasattr(self.model, "proj_out"): - self.model.proj_out.register_forward_hook(_output_hook("proj_out")) - if hasattr(self.model, "audio_proj_out"): - self.model.audio_proj_out.register_forward_hook(_output_hook("audio_proj_out")) - def forward( self, hidden_states: torch.Tensor, @@ -2942,17 +2826,6 @@ class LTX2Transformer3DModel(BaseDiT): context_mask=encoder_attention_mask, sigma=video_sigma, ) - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - video_head = latents.flatten()[:8].float().tolist() - video_flat = latents.float().flatten() - video_checksum = (video_flat * torch.arange(video_flat.numel(), device=video_flat.device)).sum().item() - _debug_block_log_line("fastvideo:modality_video" - f":latent_sum={latents.float().sum().item():.6f} " - f"latent_shape={tuple(latents.shape)} " - f"positions_sum={positions.float().sum().item():.6f} " - f"positions_shape={tuple(positions.shape)} " - f"latent_head={video_head} " - f"latent_checksum={video_checksum:.6f}") # Process audio modality if provided audio_modality = None @@ -2987,17 +2860,6 @@ class LTX2Transformer3DModel(BaseDiT): context_mask=audio_encoder_attention_mask, sigma=audio_sigma, ) - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - audio_head = audio_latents.flatten()[:8].float().tolist() - audio_flat = audio_latents.float().flatten() - audio_checksum = (audio_flat * torch.arange(audio_flat.numel(), device=audio_flat.device)).sum().item() - _debug_block_log_line("fastvideo:modality_audio" - f":latent_sum={audio_latents.float().sum().item():.6f} " - f"latent_shape={tuple(audio_latents.shape)} " - f"positions_sum={audio_positions.float().sum().item():.6f} " - f"positions_shape={tuple(audio_positions.shape)} " - f"latent_head={audio_head} " - f"latent_checksum={audio_checksum:.6f}") # Run transformer with original sequence lengths video_out, audio_out = self.model( diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 944db35ff..76c440f13 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -20,24 +20,12 @@ from fastvideo.models.dits.ltx2 import ( precompute_ltx_freqs_cis, ) from fastvideo.models.loader.weight_utils import default_weight_loader +from fastvideo.attention.selector import effective_attention_backend from fastvideo.platforms import AttentionBackendEnum from fastvideo.distributed import get_local_torch_device import math -def _debug_log_line(message: str) -> None: - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1": - return - log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "") - if not log_path: - return - log_dir = os.path.dirname(log_path) - if log_dir: - os.makedirs(log_dir, exist_ok=True) - with open(log_path, "a", encoding="utf-8") as f: - f.write(message + "\n") - - def _debug_gemma_log_line(message: str) -> None: log_path = os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", "") if not log_path: @@ -451,7 +439,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): # Note: torch.backends.cuda.enable_*_sdp() settings should be configured # at application/pipeline initialization level, not here, to avoid # unexpected side effects across the application. - if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA": + if effective_attention_backend(self.config) == AttentionBackendEnum.TORCH_SDPA: if hasattr(self._gemma_model.config, "attn_implementation"): self._gemma_model.config.attn_implementation = "sdpa" if hasattr(self._gemma_model.config, "_attn_implementation"): @@ -621,19 +609,8 @@ class LTX2GemmaTextEncoderModel(TextEncoder): attention_mask, padding_side=self.padding_side, ) - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_log_line("fastvideo:gemma_feature" - f":sum={encoded_video.float().sum().item():.6f} " - f"shape={tuple(encoded_video.shape)}") video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_video, encoded_audio, attention_mask) - if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1": - _debug_log_line("fastvideo:gemma_video_encoding" - f":sum={video_encoding.float().sum().item():.6f} " - f"shape={tuple(video_encoding.shape)}") - _debug_log_line("fastvideo:gemma_audio_encoding" - f":sum={audio_encoding.float().sum().item():.6f} " - f"shape={tuple(audio_encoding.shape)}") hidden_states = (audio_encoding, ) if output_hidden_states else None return BaseEncoderOutput( diff --git a/fastvideo/models/encoders/t5gemma.py b/fastvideo/models/encoders/t5gemma.py index a67088e80..60e06a08e 100644 --- a/fastvideo/models/encoders/t5gemma.py +++ b/fastvideo/models/encoders/t5gemma.py @@ -14,13 +14,13 @@ on top — the pipeline prompt-preprocessing stage handles pad-or-trim to """ from __future__ import annotations -import os from typing import Iterable import torch from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig from fastvideo.models.encoders.base import TextEncoder +from fastvideo.attention.selector import effective_attention_backend from fastvideo.platforms import AttentionBackendEnum @@ -70,7 +70,7 @@ class T5GemmaEncoderModel(TextEncoder): is_encoder_decoder=False, dtype=dtype, ) - if os.getenv("FASTVIDEO_ATTENTION_BACKEND") == "TORCH_SDPA": + if effective_attention_backend(self.config) == AttentionBackendEnum.TORCH_SDPA: if hasattr(model.config, "attn_implementation"): model.config.attn_implementation = "sdpa" if hasattr(model.config, "_attn_implementation"): diff --git a/fastvideo/models/vaes/oobleck.py b/fastvideo/models/vaes/oobleck.py index 2c55b25e6..89b935a94 100644 --- a/fastvideo/models/vaes/oobleck.py +++ b/fastvideo/models/vaes/oobleck.py @@ -332,8 +332,7 @@ class OobleckVAE(nn.Module): * a local directory containing `config.json` + safetensors, * a local directory whose `subfolder="vae"` holds those files. - For gated repos, the HF token is read from `HF_TOKEN` / - `HUGGINGFACE_HUB_TOKEN` / `HF_API_KEY` (see `resolve_hf_token`). + For gated repos, the HF token is read by `resolve_hf_token`. """ import inspect from safetensors.torch import load_file diff --git a/fastvideo/models/wan/transformer.py b/fastvideo/models/wan/transformer.py index 05f010353..35d811d43 100644 --- a/fastvideo/models/wan/transformer.py +++ b/fastvideo/models/wan/transformer.py @@ -7,8 +7,8 @@ from typing import Any import torch import torch.nn as nn -import fastvideo.envs as envs from fastvideo.attention import (DistributedAttention, DistributedAttention_VSA, LocalAttention) +from fastvideo.attention.selector import effective_attention_backend from fastvideo.models.wan.config import WanVideoConfig from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, sequence_model_parallel_shard) @@ -625,8 +625,11 @@ class WanTransformer3DModel(BaseDiT): ) # 3. Transformer blocks - attn_backend = envs.FASTVIDEO_ATTENTION_BACKEND - transformer_block = WanTransformerBlock_VSA if attn_backend == "VIDEO_SPARSE_ATTN" else WanTransformerBlock + # The block type must match the backend that this model's attention + # layers follow (see fastvideo/attention/AGENTS.md). + attn_backend = effective_attention_backend(config) + transformer_block = (WanTransformerBlock_VSA + if attn_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN else WanTransformerBlock) self.blocks = nn.ModuleList([ transformer_block(inner_dim, config.ffn_dim, diff --git a/fastvideo/performance/hf_store.py b/fastvideo/performance/hf_store.py index 4701fd134..d8d2134da 100644 --- a/fastvideo/performance/hf_store.py +++ b/fastvideo/performance/hf_store.py @@ -20,12 +20,13 @@ import pandas as pd from huggingface_hub import HfApi, snapshot_download from huggingface_hub.constants import ENDPOINT +from fastvideo.utils import DEPRECATED_HF_TOKEN_ENV_VARS, HF_TOKEN_ENV_VARS, resolve_hf_token + # --------------------------------------------------------------------------- # Configuration — read once at import time, shared across both consumers # --------------------------------------------------------------------------- HF_REPO_ID: str = os.environ.get("HF_REPO_ID", "FastVideo/performance-tracking") -HF_TOKEN_ENV_VARS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN") SYNC_MARKER = ".hf_sync_complete" SYNC_REUSE_TTL_SECONDS = int(os.environ.get("PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS", "3600")) @@ -75,15 +76,6 @@ def _parse_record_timestamp(record: dict[str, Any]) -> datetime | None: return ts -def resolve_hf_token() -> str | None: - """Return the first configured Hugging Face token env var.""" - for env_var in HF_TOKEN_ENV_VARS: - token = os.environ.get(env_var) - if token: - return token - return None - - # --------------------------------------------------------------------------- # HF I/O # --------------------------------------------------------------------------- @@ -208,7 +200,7 @@ def upload_record( """ token = resolve_hf_token() if not token: - msg = f"hf_store: none of {', '.join(HF_TOKEN_ENV_VARS)} set" + msg = f"hf_store: none of {', '.join(HF_TOKEN_ENV_VARS + DEPRECATED_HF_TOKEN_ENV_VARS)} set" if strict: raise RuntimeError(f"{msg}; cannot upload.") print(f"{msg}, skipping upload.") diff --git a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py index 82f8d6fdd..41230e07d 100644 --- a/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py +++ b/fastvideo/pipelines/basic/ltx2/ltx2_pipeline.py @@ -175,36 +175,6 @@ class LTX2Pipeline(LoRAPipeline): ) def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - # Optional debug-instrumentation env vars. The internal FastVideoArgs - # carries debug_model_sums / debug_model_detail toggles; the public - # args don't define them yet, so getattr-with-default keeps the - # pipeline runnable on both shapes without forcing the public args - # to add fields a public consumer wouldn't otherwise touch. - if getattr(fastvideo_args, "debug_model_sums", False): - os.environ["LTX2_PIPELINE_DEBUG_LOG"] = "1" - sums_path = getattr(fastvideo_args, "debug_model_sums_path", None) - if sums_path: - os.environ["LTX2_PIPELINE_DEBUG_PATH"] = sums_path - else: - logger.warning("debug_model_sums is enabled but " - "debug_model_sums_path is not set; no model sums " - "will be logged.") - else: - os.environ.pop("LTX2_PIPELINE_DEBUG_LOG", None) - os.environ.pop("LTX2_PIPELINE_DEBUG_PATH", None) - if getattr(fastvideo_args, "debug_model_detail", False): - os.environ["LTX2_DEBUG_DETAIL"] = "1" - detail_path = getattr(fastvideo_args, "debug_model_detail_path", None) - if detail_path: - os.environ["LTX2_PIPELINE_DEBUG_DETAIL_PATH"] = detail_path - else: - logger.warning("debug_model_detail is enabled but " - "debug_model_detail_path is not set; no detailed " - "hooks will be logged.") - else: - os.environ.pop("LTX2_DEBUG_DETAIL", None) - os.environ.pop("LTX2_PIPELINE_DEBUG_DETAIL_PATH", None) - tokenizer = self.get_module("tokenizer") if tokenizer is not None: tokenizer.padding_side = "left" diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py index 2058cf344..4cd9aa412 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py @@ -14,8 +14,8 @@ from pathlib import Path import torch from tqdm.auto import tqdm -import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionMetadataBuilder) +from fastvideo.attention.selector import component_attention_backend from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.pipelines.pipeline_batch_info import ForwardBatch @@ -28,6 +28,7 @@ from fastvideo.pipelines.basic.ltx2.stages.ltx2_image_conditioning import (LTX2_ from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.logger import init_logger +from fastvideo.platforms import AttentionBackendEnum from fastvideo.models.dits.ltx2 import (AudioLatentShape, DEFAULT_LTX2_AUDIO_CHANNELS, DEFAULT_LTX2_AUDIO_DOWNSAMPLE, DEFAULT_LTX2_AUDIO_HOP_LENGTH, DEFAULT_LTX2_AUDIO_MEL_BINS, DEFAULT_LTX2_AUDIO_SAMPLE_RATE, VideoLatentShape) @@ -459,16 +460,16 @@ class LTX2DenoisingStage(PipelineStage): # Hint runtime FP4 layer gating (single shared transformer path): # stage-1 denoising uses "base", stage-2 refine uses "refine". batch.extra["ltx2_fp4_stage_profile"] = ("refine" if self.sigmas_override is not None else "base") - attention_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND", envs.FASTVIDEO_ATTENTION_BACKEND) + attention_backend = component_attention_backend(self.transformer) wants_vsa_metadata = attention_backend in ( - "VIDEO_SPARSE_ATTN", - "SAGE_ATTN_THREE", + AttentionBackendEnum.VIDEO_SPARSE_ATTN, + AttentionBackendEnum.SAGE_ATTN_THREE, ) # VIDEO_SPARSE_ATTN requires the fastvideo-kernel VSA op. # SAGE_ATTN_THREE VSA+QAT only needs VSA metadata and its own kernel path. - use_vsa = wants_vsa_metadata and (attention_backend != "VIDEO_SPARSE_ATTN" or vsa_available) - if attention_backend == "VIDEO_SPARSE_ATTN" and not vsa_available: - logger.warning("FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN but VSA kernel " + use_vsa = wants_vsa_metadata and (attention_backend != AttentionBackendEnum.VIDEO_SPARSE_ATTN or vsa_available) + if attention_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN and not vsa_available: + logger.warning("The transformer uses VIDEO_SPARSE_ATTN but the VSA kernel " "is unavailable; disabling VSA metadata for this run.") vsa_metadata_builder = (VideoSparseAttentionMetadataBuilder() if use_vsa else None) diff --git a/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py index 0c3a985b6..bac774807 100644 --- a/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py +++ b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py @@ -51,7 +51,7 @@ from fastvideo.pipelines.stages import ( InputValidationStage, TextEncodingStage, ) -from fastvideo.utils import maybe_download_model +from fastvideo.utils import maybe_download_model, resolve_hf_token logger = init_logger(__name__) @@ -61,21 +61,18 @@ _WAN_VAE_HF_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" def _ensure_hf_token_env() -> str | None: - """Surface any of the three common HF token env vars as `HF_TOKEN`. + """Surface the resolved HF token as `HF_TOKEN`. FastVideo workers spawn child processes that inherit env; both `huggingface_hub` and `transformers.AutoTokenizer.from_pretrained` - look at `HF_TOKEN` / `HUGGINGFACE_HUB_TOKEN` by default but not - `HF_API_KEY`. If only the latter is set, gated downloads fail with - 401. Aliasing at pipeline-load time is the minimum-disruption fix. + read `HF_TOKEN` but not the FastVideo-specific aliases that + `resolve_hf_token` also accepts. If only an alias is set, gated + downloads fail with 401, so expose the token as `HF_TOKEN`. """ - for src in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY"): - value = os.environ.get(src) - if value: - os.environ.setdefault("HF_TOKEN", value) - os.environ.setdefault("HUGGINGFACE_HUB_TOKEN", value) - return value - return None + token = resolve_hf_token() + if token: + os.environ.setdefault("HF_TOKEN", token) + return token class MagiHumanPipeline(ComposedPipelineBase): diff --git a/fastvideo/pipelines/stages/kandinsky5.py b/fastvideo/pipelines/stages/kandinsky5.py index 1cab8a095..c65ba7f92 100644 --- a/fastvideo/pipelines/stages/kandinsky5.py +++ b/fastvideo/pipelines/stages/kandinsky5.py @@ -10,10 +10,9 @@ import torch from diffusers.utils.torch_utils import randn_tensor from tqdm.auto import tqdm -import fastvideo.envs as envs from fastvideo.attention import LocalAttention from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder -from fastvideo.attention.selector import backend_name_to_enum +from fastvideo.attention.selector import component_attention_backend from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context @@ -203,12 +202,12 @@ class Kandinsky5DenoisingStage(PipelineStage): ``Kandinsky5Attention.forward`` catches the ``AssertionError`` ``LocalAttention`` raises when no pipeline forward context is set and falls back to plain ``F.scaled_dot_product_attention`` -- silently - skipping whichever kernel ``FASTVIDEO_ATTENTION_BACKEND`` requested - (the fake-quantized ``ATTN_QAT_TRAIN`` kernel). + skipping the transformer's recorded attention backend (the + fake-quantized ``ATTN_QAT_TRAIN`` kernel). ``LocalAttention.backend`` is resolved once at module construction, independent of whether a forward context is later set, so this only - catches a backend that failed to resolve to what the env var - requested -- it does not prove forward context was present for a + catches a backend that failed to resolve to the recorded decision + -- it does not prove forward context was present for a given forward call. Pair it with always wrapping the actual transformer call in ``set_forward_context`` (see the denoising loops below), which is what prevents the runtime fallback. @@ -220,17 +219,14 @@ class Kandinsky5DenoisingStage(PipelineStage): ``_STRICT_BACKENDS``'s comment), so asserting on those would flag expected behavior as a bug. """ - backend_env = envs.FASTVIDEO_ATTENTION_BACKEND - if not backend_env: - return - expected = backend_name_to_enum(backend_env) - if expected is None or expected not in self._STRICT_BACKENDS: + expected = component_attention_backend(self.transformer) + if expected not in self._STRICT_BACKENDS: return for module in self.transformer.modules(): if isinstance(module, LocalAttention): assert module.backend == expected, ( f"Kandinsky5 local attention resolved to backend {module.backend}, expected " - f"{expected} from FASTVIDEO_ATTENTION_BACKEND={backend_env}. This likely means " + f"{expected}, the transformer's recorded attention backend. This likely means " "LocalAttention's missing-forward-context guard silently fell back to SDPA.") return diff --git a/fastvideo/tests/api/test_attention_selector_resolution.py b/fastvideo/tests/api/test_attention_selector_resolution.py index 00c813165..c108ea43b 100644 --- a/fastvideo/tests/api/test_attention_selector_resolution.py +++ b/fastvideo/tests/api/test_attention_selector_resolution.py @@ -21,6 +21,8 @@ role-override tests); no kernels, no GPU. """ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -372,3 +374,34 @@ def test_component_without_a_recorded_decision_reports_no_request(): assert selector.component_attention_backend(_Bare()) is selector.NO_REQUEST assert selector.component_attention_backend(_HFStyle()) is selector.NO_REQUEST + + +def test_effective_backend_prefers_the_recorded_decision(monkeypatch) -> None: + """A component built by a loader follows the decision recorded on its + config, even when the env var asks for something else.""" + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + config = SimpleNamespace(_resolved_attention_backend=SDPA) + assert selector.effective_attention_backend(config) == SDPA + + +def test_effective_backend_follows_env_for_direct_construction(monkeypatch) -> None: + """A component constructed without a loader has no recorded decision; like + the attention layers inside it, it follows the env var.""" + config = SimpleNamespace(_resolved_attention_backend=None) + assert selector.effective_attention_backend(config) is None + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + assert selector.effective_attention_backend(config) == SAGE + assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" + + +def test_effective_backend_follows_the_active_scope(monkeypatch) -> None: + """While a loader scope is active, the scope's rule applies: a scope that + ignores the env var (the dense teacher/critic case) keeps ignoring it.""" + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + config = SimpleNamespace(_resolved_attention_backend=None) + with selector._component_attention_backend_scope(None, component="teacher"): + assert selector.effective_attention_backend(config) is None + with selector._component_attention_backend_scope(None, component="dit", consult_env=True): + assert selector.effective_attention_backend(config) == SAGE + with selector._component_attention_backend_scope(SDPA, component="dit"): + assert selector.effective_attention_backend(config) == SDPA diff --git a/fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py b/fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py index d2cb34382..1acfb7e8c 100644 --- a/fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py +++ b/fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py @@ -21,6 +21,8 @@ model load, no distributed init needed. """ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -43,41 +45,43 @@ class _ResolvedLocalAttention(LocalAttention): self.backend = backend -def _make_stage(resolved_backend: AttentionBackendEnum) -> Kandinsky5DenoisingStage: - """Bypass __init__: only set the field the guard reads.""" +def _make_stage(requested_backend: AttentionBackendEnum, + resolved_backend: AttentionBackendEnum) -> Kandinsky5DenoisingStage: + """Bypass __init__: set only the fields the guard reads. + + ``requested_backend`` is the decision the loader records on the + transformer's config; ``resolved_backend`` is what the layer resolved to. + """ stage = Kandinsky5DenoisingStage.__new__(Kandinsky5DenoisingStage) transformer = torch.nn.Module() + transformer.config = SimpleNamespace(_resolved_attention_backend=requested_backend) transformer.attn = _ResolvedLocalAttention(resolved_backend) stage.transformer = transformer return stage -def test_attn_qat_infer_flash_fallback_is_allowed(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER") - stage = _make_stage(AttentionBackendEnum.FLASH_ATTN) +def test_attn_qat_infer_flash_fallback_is_allowed(): + stage = _make_stage(AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.FLASH_ATTN) # Must not raise: FlashAttention is the documented ATTN_QAT_INFER # fallback on GPUs without the sm_120 kernel. stage._assert_local_attention_backend_engaged() -def test_attn_qat_infer_exact_match_is_allowed(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_INFER") - stage = _make_stage(AttentionBackendEnum.ATTN_QAT_INFER) +def test_attn_qat_infer_exact_match_is_allowed(): + stage = _make_stage(AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_INFER) stage._assert_local_attention_backend_engaged() -def test_attn_qat_train_mismatch_raises(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN") - stage = _make_stage(AttentionBackendEnum.TORCH_SDPA) +def test_attn_qat_train_mismatch_raises(): + stage = _make_stage(AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.TORCH_SDPA) with pytest.raises(AssertionError, match="ATTN_QAT_TRAIN"): stage._assert_local_attention_backend_engaged() -def test_attn_qat_train_exact_match_is_allowed(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN") - stage = _make_stage(AttentionBackendEnum.ATTN_QAT_TRAIN) +def test_attn_qat_train_exact_match_is_allowed(): + stage = _make_stage(AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.ATTN_QAT_TRAIN) stage._assert_local_attention_backend_engaged() diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 4f31b60f6..fbf3db6c1 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -747,20 +747,30 @@ def maybe_download_model_index(model_name_or_path: str, revision: str | None = N raise ValueError(f"Failed to download or parse a Diffusers manifest for {model_name_or_path}: {e}") from e -_HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_API_KEY") +# Token variables that `huggingface_hub` itself reads, in its priority order. +HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") +# FastVideo-specific token names, kept as deprecated aliases until the next +# minor release. +DEPRECATED_HF_TOKEN_ENV_VARS = ("HUGGINGFACE_HUB_TOKEN", "HF_API_KEY") def resolve_hf_token() -> str | None: - """Return the first non-empty HF token from the standard env vars. + """Return the first non-empty Hugging Face token from the environment. - Order: `HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, `HF_API_KEY` (the last is - a FastVideo convention; `huggingface_hub` itself doesn't read it). - Does not mutate `os.environ`. + Reads `HF_TOKEN`, then `HUGGING_FACE_HUB_TOKEN`, then the deprecated + FastVideo-specific aliases `HUGGINGFACE_HUB_TOKEN` and `HF_API_KEY`, + which log a deprecation warning. Does not mutate `os.environ`. """ - for src in _HF_TOKEN_ENV_VARS: - v = os.environ.get(src) - if v: - return v + for name in HF_TOKEN_ENV_VARS: + value = os.environ.get(name) + if value: + return value + for name in DEPRECATED_HF_TOKEN_ENV_VARS: + value = os.environ.get(name) + if value: + logger.warning_once(f"{name} is deprecated and will be removed in the next minor release; " + f"set HF_TOKEN instead.") + return value return None diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index 2cf2462af..67c576bd4 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -12,7 +12,8 @@ from dataclasses import dataclass from typing import Any, TYPE_CHECKING from collections.abc import Callable -from fastvideo.utils import get_ip, get_distributed_init_method, get_open_port, get_loopback_ip +from fastvideo.utils import (DEPRECATED_HF_TOKEN_ENV_VARS, HF_TOKEN_ENV_VARS, get_ip, get_distributed_init_method, + get_open_port, get_loopback_ip) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.worker.executor import Executor @@ -85,8 +86,8 @@ class RayDistributedExecutor(Executor): # NCCL_* knobs present on the driver are added dynamically in # ``_env_vars_to_copy_from_driver``, except the per-node NIC trio above. ADDITIONAL_ENV_VARS = { - "HF_TOKEN", - "HUGGING_FACE_HUB_TOKEN", + *HF_TOKEN_ENV_VARS, + *DEPRECATED_HF_TOKEN_ENV_VARS, "NCCL_IB_DISABLE", "NCCL_P2P_DISABLE", "NCCL_CUMEM_ENABLE", diff --git a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py index 8261fd4b0..0363c2a70 100644 --- a/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py +++ b/tests/local_tests/ltx2/test_ltx2_pipeline_smoke.py @@ -31,9 +31,7 @@ def _log_tensor_stats(label: str, tensor: torch.Tensor) -> None: def _truncate_debug_logs() -> None: for env_var in ( - "LTX2_PIPELINE_DEBUG_PATH", "LTX2_REFERENCE_DEBUG_PATH", - "LTX2_PIPELINE_DEBUG_DETAIL_PATH", "LTX2_REFERENCE_DEBUG_DETAIL_PATH", ): log_path = os.getenv(env_var, "") @@ -104,18 +102,10 @@ def _run_audio_decode_smoke( def test_ltx2_pipeline_smoke(): repo_root = Path(__file__).resolve().parents[3] debug_dir = repo_root / "ltx2_debug" - os.environ.setdefault( - "LTX2_PIPELINE_DEBUG_PATH", - str(debug_dir / "fastvideo_pipeline.log"), - ) os.environ.setdefault( "LTX2_REFERENCE_DEBUG_PATH", str(debug_dir / "reference_pipeline.log"), ) - os.environ.setdefault( - "LTX2_PIPELINE_DEBUG_DETAIL_PATH", - str(debug_dir / "fastvideo_pipeline_detail.log"), - ) os.environ.setdefault( "LTX2_REFERENCE_DEBUG_DETAIL_PATH", str(debug_dir / "reference_pipeline_detail.log"), From e3b88bb12a4f696dbbeade16db16042af7327d9b Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Tue, 29 Sep 2026 12:50:45 -0700 Subject: [PATCH 038/151] [feat] Add a typed environment-variable registry, policy doc, and contract test (#1896) --- .agents/skills/env-var-conventions/SKILL.md | 71 ++ docs/contributing/env_vars.md | 180 ++++ .../fastvideo_kernel/block_sparse_attn.py | 18 +- .../backends/video_sparse_attn_h3.py | 8 +- fastvideo/attention/selector.py | 6 +- .../attention/utils/flash_attn_default.py | 2 +- .../attention/utils/flash_attn_no_pad.py | 2 +- .../device_communicators/ulysses_a2a.py | 2 +- fastvideo/distributed/parallel_state.py | 6 +- fastvideo/envs.py | 732 ++++++++-------- fastvideo/eval/models.py | 2 +- fastvideo/fastvideo_args.py | 10 +- fastvideo/hooks/activation_trace.py | 10 +- fastvideo/logger.py | 8 +- fastvideo/models/dits/minimax_h3.py | 4 +- fastvideo/models/loader/fsdp_load.py | 3 +- fastvideo/pipelines/composed_pipeline_base.py | 4 +- fastvideo/pipelines/stages/base.py | 2 +- fastvideo/pipelines/stages/denoising.py | 2 +- fastvideo/platforms/cuda.py | 2 +- fastvideo/platforms/npu.py | 4 +- fastvideo/platforms/rocm.py | 2 +- fastvideo/profiler.py | 22 +- fastvideo/tests/contract/test_env_policy.py | 827 ++++++++++++++++++ fastvideo/tests/stages/_denoising_fixtures.py | 6 +- fastvideo/train/methods/base.py | 2 +- fastvideo/train/models/ltx2/ltx2.py | 2 +- fastvideo/train/models/wan/wan.py | 2 +- fastvideo/training/distillation_pipeline.py | 4 +- .../self_forcing_distillation_pipeline.py | 4 +- fastvideo/training/training_pipeline.py | 9 +- fastvideo/utils.py | 10 +- fastvideo/worker/gpu_worker.py | 2 +- fastvideo/worker/multiproc_executor.py | 4 +- fastvideo/worker/ray_distributed_executor.py | 4 +- fastvideo/worker/ray_env.py | 2 +- mkdocs.yml | 1 + 37 files changed, 1546 insertions(+), 435 deletions(-) create mode 100644 .agents/skills/env-var-conventions/SKILL.md create mode 100644 docs/contributing/env_vars.md create mode 100644 fastvideo/tests/contract/test_env_policy.py diff --git a/.agents/skills/env-var-conventions/SKILL.md b/.agents/skills/env-var-conventions/SKILL.md new file mode 100644 index 000000000..d793281a5 --- /dev/null +++ b/.agents/skills/env-var-conventions/SKILL.md @@ -0,0 +1,71 @@ +--- +name: env-var-conventions +description: Add, read, rename, or remove an environment variable in FastVideo, or change the environment-variable policy. Use before touching fastvideo/envs.py, os.environ, os.getenv, or monkeypatch.setenv in fastvideo/, and when fastvideo/tests/contract/test_env_policy.py fails. +--- + +# Environment Variable Conventions + +## Purpose + +FastVideo registers its environment variables as typed fields in +`fastvideo/envs.py`. The policy that governs them is +`docs/contributing/env_vars.md`, and the contract test +`fastvideo/tests/contract/test_env_policy.py` enforces the policy in the unit +CI lane. This skill routes an environment-variable change through that policy. +The policy doc is the single source of the rules; read it instead of relying +on a summary here. + +## Prerequisites + +- Read `docs/contributing/env_vars.md` in full. +- Decide whether the setting belongs in an environment variable or an argument + (rule 5 in the policy doc). Settings that users change per deployment are + arguments; add them through `fastvideo/fastvideo_args.py` instead. + +## Inputs + +| Parameter | Required | Description | +| ---------- | -------- | -------------------------------------------------------------- | +| `change` | Yes | Add, read, rename, or remove a variable, or change the policy. | +| `variable` | Yes | The variable name, with the `FASTVIDEO_` prefix. | + +## Steps + +1. **Declare or edit the variable in `fastvideo/envs.py`.** + - Pick the field type and category that the policy doc lists. + - Write a description that states what the variable does and its units. +2. **Read the variable with `envs.NAME.get()` inside a function.** + - In tests, change the value with `envs.NAME.override(value)`. + - Do not call `os.environ`, `os.getenv`, or `monkeypatch.setenv` for a + FastVideo variable. +3. **Regenerate the table in the policy doc.** + - Run `python fastvideo/tests/contract/test_env_policy.py`. +4. **Run the contract test.** + - Run `pytest fastvideo/tests/contract/test_env_policy.py`. + - When the test reports a fixed known violation, delete or lower its entry + in `KNOWN_VIOLATIONS`. Never add an entry to `KNOWN_VIOLATIONS`. +5. **When the policy itself changes, update the policy doc and the contract + test in the same pull request.** + - The rules in `docs/contributing/env_vars.md`, the checks and allowlist in + `fastvideo/tests/contract/test_env_policy.py`, and this skill must agree. + +## Outputs + +- A registry entry in `fastvideo/envs.py` and call sites that use + `envs.NAME.get()`. +- A regenerated table in `docs/contributing/env_vars.md`. +- A passing `fastvideo/tests/contract/test_env_policy.py`. + +## Example Usage + +``` +Add a FASTVIDEO_DEBUG_MY_STAGE switch that logs MyStage inputs. +``` + +## References + +- `docs/contributing/env_vars.md`: the policy, the field types, and the + violation kinds that the contract test reports. +- `fastvideo/envs.py`: the registry. +- `fastvideo/tests/contract/test_env_policy.py`: the contract test, + `EXTERNAL_ALLOWLIST`, and `KNOWN_VIOLATIONS`. diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md new file mode 100644 index 000000000..a4dac3d29 --- /dev/null +++ b/docs/contributing/env_vars.md @@ -0,0 +1,180 @@ +# Environment Variables + +FastVideo reads environment variables for expert switches, debugging, profiling, and the settings that launchers such +as `torchrun` provide. This page is the policy for those variables. The contract test +`fastvideo/tests/contract/test_env_policy.py` enforces the policy in the unit CI lane, and the coding-agent skill +`.agents/skills/env-var-conventions/SKILL.md` points here. When the policy changes, update this page and the contract +test in the same pull request. + +## Rules + +1. **Register every FastVideo variable in `fastvideo/envs.py`.** Each entry declares a type, a default, a category, + and a description. Variables that other tools own (CUDA, NCCL, PyTorch, launchers) are not registered; code reads + them directly with `os.environ.get("NAME")`, and the name must be in the external-variable allowlist + (`EXTERNAL_ALLOWLIST` in the contract test). +2. **Read with `envs.NAME.get()`, write with `envs.NAME.set()`, and change a value in tests with + `envs.NAME.override()`.** Each type has one parsing rule. A value that the rule rejects raises + `fastvideo.envs.EnvVarError` instead of falling back to the default. +3. **Name FastVideo variables with the `FASTVIDEO_` prefix.** The second word states the purpose where one applies: + `ENABLE_`, `DISABLE_`, `USE_`, `FORCE_`, `DEBUG_`, `TEST_`. +4. **Keep a renamed variable as a deprecated alias until the next minor release.** Setting the old name logs a + warning. Delete a variable that no code reads, and list it in a deprecation table so that setting it logs a + warning. +5. **Give each setting one source: an argument or an environment variable.** Settings that users change per + deployment are arguments (CLI or YAML). Expert switches, emergency off switches, and debugging and test switches + are environment variables. +6. **Read variables inside functions.** `envs.NAME.get()` runs when the function runs, so a changed value takes + effect without re-importing a module. Module level, class bodies, decorators, and default argument values run at + import time. +7. **Do not write the environment to pass values between parts of FastVideo.** Pass an argument instead. Tests use + `envs.NAME.override()`. + +## Field types + +| Class | Value type | Parsing rule | +| ------------ | ---------------- | ----------------------------------------------------------------------------------- | +| `EnvBool` | `bool` | `1`, `true`, `yes`, `on` are true; `0`, `false`, `no`, `off`, and `""` are false. | +| | | Case-insensitive; surrounding whitespace is ignored. | +| `EnvInt` | `int` | `int(value)` | +| `EnvFloat` | `float` | `float(value)` | +| `EnvStr` | `str` or `None` | The raw string. A `None` default means that the variable has no default. | +| `EnvPath` | `str` or `None` | The raw string with a leading `~` expanded. | +| `EnvChoice` | `str` | Stripped and lower-cased, then checked against the declared `choices`. | + +A default can be a zero-argument function; `get()` calls it on each read while the variable is unset. The path roots +use this to follow `XDG_CONFIG_HOME` and `XDG_CACHE_HOME`. + +Using a field without a method, as in `if envs.FASTVIDEO_FA4:`, raises `TypeError`. + +## Add a variable + +1. Declare the variable in the matching section of `fastvideo/envs.py`: + + ```python + FASTVIDEO_DEBUG_MY_STAGE = EnvBool(False, category="debug", doc="Log the inputs of MyStage.") + ``` + + The category is one of the values in `envs.CATEGORIES`. + +2. Read the variable inside a function: + + ```python + import fastvideo.envs as envs + + def forward(self, batch): + if envs.FASTVIDEO_DEBUG_MY_STAGE.get(): + logger.info("MyStage inputs: %s", batch.keys()) + ``` + +3. Regenerate the table at the end of this page: + + ```bash + python fastvideo/tests/contract/test_env_policy.py + ``` + +4. Run the contract test: + + ```bash + pytest fastvideo/tests/contract/test_env_policy.py + ``` + +In a test, change the value with `override`, which restores the previous value on exit: + +```python +with envs.FASTVIDEO_DEBUG_MY_STAGE.override(True): + run_stage() +``` + +## What the contract test checks + +The test parses every Python file under `fastvideo/`, including `fastvideo/tests/`, with Python's `ast` module. It +skips `fastvideo/third_party/`, which is copied from upstream projects, and the registry `fastvideo/envs.py`. It does +not check `apps/`, `examples/`, `scripts/`, `fastvideo-kernel/`, or `docs/`. + +It reports each violation as `: `: + +| Kind | Code that triggers it | Fix | +| ------------------ | ------------------------------------------------------------ | -------------------------------------------- | +| `read` | `os.getenv`, `os.environ.get`, `os.environ[...]`, or | Register the variable and call | +| | `"NAME" in os.environ` with a name outside the allowlist, or | `envs.NAME.get()`. For a variable that | +| | with a name built at runtime (``) | another tool owns, add it to | +| | | `EXTERNAL_ALLOWLIST` with a reason. | +| `write` | `os.environ[...] = ...`, `setdefault`, `pop`, `del`, | Pass an argument instead. In tests, use | +| | `os.putenv`, `os.unsetenv`, `monkeypatch.setenv`/`delenv` | `envs.NAME.override()`. | +| `whole-environ` | `os.environ.copy()`, `dict(os.environ)`, iteration, | Read the specific variables that the code | +| | `mock.patch.dict(os.environ, ...)`, `os.environ.update` | needs. | +| `bare-field` | A registry field used without calling one of its methods, | Call `envs.NAME.get()`. | +| | as in `envs.NAME == "auto"` or `getter = envs.NAME.get` | | +| `import-time-read` | `envs.NAME.get()` outside a function | Move the read into the function that uses | +| | | the value. | +| `prefix` | A registry entry without the `FASTVIDEO_` prefix | Rename the variable and keep the old name as | +| | | a deprecated alias. | +| `unread` | A registry entry that no code reads with `get()` or | Delete the variable and add it to the | +| | `is_set()` | deprecation table. | + +The test recognizes `os` imported under another name and `from os import environ, getenv`. Code that reaches the +environment through `importlib` or `getattr(os, "environ")` is left to code review. + +The test also checks that every registry entry has a category from `envs.CATEGORIES` and a description, and that the +table at the end of this page matches the registry. + +**Known violations.** `KNOWN_VIOLATIONS` in the contract test lists the violations that existed when the policy was +introduced. The list only shrinks. A violation that is not in the list fails the test. A listed violation that no +longer exists also fails the test, so the fixing pull request deletes its entry. + +## Registered variables + + +| Variable | Type | Default | Category | Description | +| ---------------------------------------------- | ------------------------------ | --------------------------- | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `FASTVIDEO_TARGET_DEVICE` | str | `cuda` | build | Target device of FastVideo: cuda, rocm, neuron, cpu, or openvino. | +| `MAX_JOBS` | str | unset | build | Maximum number of parallel compilation jobs. Defaults to the number of CPUs. | +| `NVCC_THREADS` | str | unset | build | Number of nvcc threads. When set, MAX_JOBS is reduced to avoid oversubscribing the CPU. | +| `FASTVIDEO_USE_PRECOMPILED` | bool | `0` | build | Use precompiled binaries (\*.so). | +| `CMAKE_BUILD_TYPE` | str | unset | build | CMake build type: Debug, Release, or RelWithDebInfo. | +| `VERBOSE` | bool | `0` | build | Print verbose logs during installation. | +| `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_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_RINGBUFFER_WARNING_INTERVAL` | int | `60` | distributed | Seconds between warnings while the ring buffer is full. | +| `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). | +| `HCCL_SO_PATH` | str | unset | distributed | Path to the HCCL library file on Ascend NPUs. | +| `FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S` | int | `60` | distributed | Timeout in seconds for each engine iteration. | +| `FASTVIDEO_WORKER_MULTIPROC_METHOD` | one of spawn, fork, forkserver | `spawn` | distributed | Multiprocessing start method for worker processes. | +| `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. | +| `LD_LIBRARY_PATH` | str | unset | external | Searched for the NCCL library when FASTVIDEO_NCCL_SO_PATH is unset. | +| `LOCAL_RANK` | int | `0` | external | Local rank of the process in a distributed run; selects the GPU device id. | +| `CUDA_VISIBLE_DEVICES` | str | unset | external | Visible devices in a distributed run. | +| `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_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | +| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | +| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | +| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | +| `FASTVIDEO_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_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_TEST_DYNAMO_FULLGRAPH_CAPTURE` | bool | `1` | debug | Enable Dynamo fullgraph capture. | +| `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_SERVER_DEV_MODE` | bool | `0` | debug | Run the server in development mode with extra debugging endpoints. | +| `FASTVIDEO_TRACE_FUNCTION` | bool | `0` | debug | Trace function calls. | +| `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_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. | + diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index cb5127d74..0b2c56df1 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -60,7 +60,23 @@ def _force_sm100a() -> bool: metadata raises because Triton has no compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. """ - return os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1" + return _env_bool("FASTVIDEO_VSA_SM100A") + + +def _env_bool(name: str) -> bool: + """Parse a boolean environment variable with the rule of ``fastvideo.envs.EnvBool``. + + fastvideo-kernel cannot import fastvideo, so the rule is repeated here to keep + the kernel and the package in agreement: 1, true, yes, on are true; 0, false, + no, off, and the empty string are false; case-insensitive; anything else raises. + """ + value = os.environ.get(name, "0").strip().lower() + if value in ("1", "true", "yes", "on"): + return True + if value in ("0", "false", "no", "off", ""): + return False + raise ValueError(f"Invalid value {value!r} for {name}: expected 1, true, yes, on, 0, false, no, off, " + "or an empty string") def _sm100a_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index bcc5091bc..e37e6d5a1 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -45,7 +45,6 @@ logs one warning and falls back. import functools import math -import os from dataclasses import dataclass from typing import Any @@ -69,6 +68,7 @@ try: except ImportError: _sm100a = None +import fastvideo.envs as envs from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder, layer_idx_from_prefix) from fastvideo.attention.backends.video_sparse_attn import (compute_topk, construct_variable_block_sizes, @@ -469,7 +469,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): """ if self._compile_layer_idx is None: self.prepare_for_compile(device) - requested = os.environ.get(VSA_SM100A_ENV, "0") == "1" + requested = envs.FASTVIDEO_VSA_SM100A.get() enabled = False reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention" if requested: @@ -529,7 +529,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): # Training/generic compile keeps the long-standing Triton route. sm100a_requested = False else: - sm100a_requested = os.environ.get(VSA_SM100A_ENV, "0") == "1" + sm100a_requested = envs.FASTVIDEO_VSA_SM100A.get() needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode and sm100a_requested) kernel_tiles = n_tiles + int(needs_sm100a_pair) target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1]) @@ -702,7 +702,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): "VSA-H3 regional fullgraph compile requires the prepared sm_100a BF16/head-128 route " "on a supported device; disable inference_torch_compile for this request.") use_sm100a = True - elif not compiling and os.environ.get(VSA_SM100A_ENV, "0") == "1": + elif not compiling and envs.FASTVIDEO_VSA_SM100A.get(): reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_variable_block_sizes, grad_mode) if reason is None and map_to_index is None: reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 03df0b499..302dd4683 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -208,7 +208,7 @@ def effective_attention_backend(config: object) -> AttentionBackendEnum | None: scope = _SCOPE.get() if scope is not None and (scope.backend is not None or not scope.consult_env): return scope.backend - env_backend = envs.FASTVIDEO_ATTENTION_BACKEND + env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() return None if env_backend is None else backend_name_to_enum(env_backend) @@ -249,11 +249,11 @@ def get_attn_backend( if scope is not None: requested = scope.backend component = scope.component - env_backend = envs.FASTVIDEO_ATTENTION_BACKEND if scope.consult_env else None + env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() if scope.consult_env else None else: requested = None component = None - env_backend = envs.FASTVIDEO_ATTENTION_BACKEND + env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() # The active device is a real selection input, not bookkeeping: the # platform's backend resolution runs capability probes against the # *current* device (e.g. AttnQatInferBackend's per-arch capability sets diff --git a/fastvideo/attention/utils/flash_attn_default.py b/fastvideo/attention/utils/flash_attn_default.py index 3984f68f8..ee50cd8bb 100644 --- a/fastvideo/attention/utils/flash_attn_default.py +++ b/fastvideo/attention/utils/flash_attn_default.py @@ -44,7 +44,7 @@ logger = init_logger(__name__) # kernels JIT-compile per shape family and can fail at runtime on some # arch/shape combinations, so it is never auto-selected just because it is # installed. -if envs.FASTVIDEO_FA4: +if envs.FASTVIDEO_FA4.get(): try: from fastvideo.attention.utils.flash_attn_cute import flash_attn_func except ImportError as e: diff --git a/fastvideo/attention/utils/flash_attn_no_pad.py b/fastvideo/attention/utils/flash_attn_no_pad.py index 52454fd3d..681033f1d 100644 --- a/fastvideo/attention/utils/flash_attn_no_pad.py +++ b/fastvideo/attention/utils/flash_attn_no_pad.py @@ -25,7 +25,7 @@ from fastvideo import envs def _resolve_flash_attn_varlen_func() -> tuple[Any, str]: - if envs.FASTVIDEO_FA4: + if envs.FASTVIDEO_FA4.get(): # FA4 cute is explicit opt-in (see fastvideo/attention/backends/ # flash_attn.py); with FASTVIDEO_FA4=1 an unimportable FA4 build must # fail loudly here rather than fall through to FA3/FA2. RuntimeError, diff --git a/fastvideo/distributed/device_communicators/ulysses_a2a.py b/fastvideo/distributed/device_communicators/ulysses_a2a.py index 2e41a2a1c..ffbd94b92 100644 --- a/fastvideo/distributed/device_communicators/ulysses_a2a.py +++ b/fastvideo/distributed/device_communicators/ulysses_a2a.py @@ -36,7 +36,7 @@ _MODE_FROM_DIMS = {(2, 1): 0, (1, 2): 1} def is_enabled() -> bool: """Whether the fused path is opted in via FASTVIDEO_ULYSSES_A2A.""" - return envs.FASTVIDEO_ULYSSES_A2A == "auto" + return envs.FASTVIDEO_ULYSSES_A2A.get() == "auto" class _FusedUlyssesA2A(torch.autograd.Function): diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 5cba91c37..49b2e4f34 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -759,7 +759,7 @@ def init_distributed_environment( if local_rank == -1: # local rank not set, this usually happens in single-node # setting, where we can use rank as local rank - local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank + local_rank = envs.LOCAL_RANK.get() if distributed_init_method == "env://" else rank global _WORLD if _WORLD is None: ranks = list(range(torch.distributed.get_world_size())) @@ -884,9 +884,9 @@ def get_local_torch_device() -> torch.device: """Return the torch device for the current rank.""" from fastvideo.platforms import current_platform if current_platform.is_npu(): - device = torch.device(f"npu:{envs.LOCAL_RANK}") + device = torch.device(f"npu:{envs.LOCAL_RANK.get()}") elif current_platform.is_cuda_alike() or current_platform.is_cuda(): - device = torch.device(f"cuda:{envs.LOCAL_RANK}") + device = torch.device(f"cuda:{envs.LOCAL_RANK.get()}") else: device = torch.device("mps") return device diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 86abc1bcd..bfa8cd809 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -1,56 +1,160 @@ # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py +"""Registry of the environment variables that FastVideo reads. + +Every FastVideo-owned environment variable is declared here once, as a typed +field with a default, a category, and a description. Code reads a variable +with ``envs.NAME.get()`` inside a function, writes it with ``envs.NAME.set()``, +and tests change it temporarily with ``envs.NAME.override()``. Each field type +has one parsing rule, and a value that the rule rejects raises ``EnvVarError``. + +The policy for environment variables is in ``docs/contributing/env_vars.md``, +and ``fastvideo/tests/contract/test_env_policy.py`` enforces it. +""" import os -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Generic, TypeVar -if TYPE_CHECKING: - FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60 - FASTVIDEO_NCCL_SO_PATH: str | None = None - LD_LIBRARY_PATH: str | None = None - LOCAL_RANK: int = 0 - CUDA_VISIBLE_DEVICES: str | None = None - FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo") - FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo") - FASTVIDEO_CONFIGURE_LOGGING: int = 1 - FASTVIDEO_RAY_PER_WORKER_GPUS: float = 1.0 - FASTVIDEO_LOGGING_LEVEL: str = "INFO" - FASTVIDEO_LOGGING_PREFIX: str = "" - FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None - FASTVIDEO_TRACE_FUNCTION: int = 0 - FASTVIDEO_ATTENTION_BACKEND: str | None = None - FASTVIDEO_FA4: bool = False - FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN: bool = False - FASTVIDEO_INFERENCE_TORCH_COMPILE: bool = False - FASTVIDEO_MINIMAX_H3_FUSIONS: str = "" - FASTVIDEO_VAE_PARALLEL_DECODE: bool = False - FASTVIDEO_VAE_PARALLEL_ENCODE: bool = False - FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY: str | None = None - FASTVIDEO_ULYSSES_A2A: str = "off" - FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn" - FASTVIDEO_TARGET_DEVICE: str = "cuda" - MAX_JOBS: str | None = None - NVCC_THREADS: str | None = None - CMAKE_BUILD_TYPE: str | None = None - VERBOSE: bool = False - FASTVIDEO_NVTX_PROFILE: bool = False - FASTVIDEO_TORCH_PROFILER_DIR: str | None = None - FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES: bool = False - FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False - FASTVIDEO_TORCH_PROFILER_WITH_STACK: bool = False - FASTVIDEO_TORCH_PROFILER_WITH_FLOPS: bool = False - FASTVIDEO_TORCH_PROFILE_REGIONS: str = "" - FASTVIDEO_TRACE_ACTIVATIONS: bool = False - FASTVIDEO_TRACE_LAYERS: str = "" - FASTVIDEO_TRACE_STATS: str = "abs_mean,sum" - FASTVIDEO_TRACE_OUTPUT: str = "/tmp/fv_trace_.jsonl" - FASTVIDEO_TRACE_STEPS: str = "" - FASTVIDEO_SERVER_DEV_MODE: bool = False - FASTVIDEO_STAGE_LOGGING: bool = False - FASTVIDEO_CFG_GATE_STEP: float = 1.0 - FASTVIDEO_HOST_IP: str = "" - FASTVIDEO_LOOPBACK_IP: str = "" +T = TypeVar("T") +# String fields accept a string default or None (unset). +S = TypeVar("S", str, str | None) + +POLICY_DOC = "docs/contributing/env_vars.md" + +# Allowed values of EnvField.category. +CATEGORIES = ("build", "path", "distributed", "external", "logging", "attention", "performance", "profiling", "debug", + "sampling") + +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) +_FALSE_VALUES = frozenset({"0", "false", "no", "off", ""}) + + +class EnvVarError(ValueError): + """An environment variable holds a value that its registered type rejects.""" + + +class EnvField(Generic[T]): + """One registered environment variable: its type, default, category, and description. + + ``default`` is either the value itself or a zero-argument function that + computes it on each read while the variable is unset. + """ + + type_name = "" + + def __init__(self, default: T | Callable[[], T], *, category: str, doc: str) -> None: + self.name = "" # Set by _register_fields() from the module attribute name. + self.default = default + self.category = category + self.doc = doc + + def parse(self, raw: str) -> T: + raise NotImplementedError + + def format(self, value: T) -> str: + return str(value) + + def get(self) -> T: + """Return the parsed value, or the default when the variable is unset.""" + raw = os.environ.get(self.name) + if raw is None: + return self.default() if callable(self.default) else self.default + try: + return self.parse(raw) + except ValueError as exc: + raise EnvVarError(f"Invalid value {raw!r} for {self.name}: {exc}. See {POLICY_DOC}.") from None + + def is_set(self) -> bool: + return self.name in os.environ + + def set(self, value: T) -> None: + os.environ[self.name] = self.format(value) + + def clear(self) -> None: + os.environ.pop(self.name, None) + + @contextmanager + def override(self, value: T | None) -> Iterator[None]: + """Set the variable, or unset it when ``value`` is None, and restore the previous value on exit.""" + previous = os.environ.get(self.name) + if value is None: + self.clear() + else: + self.set(value) + try: + yield + finally: + if previous is None: + self.clear() + else: + os.environ[self.name] = previous + + def __bool__(self) -> bool: + raise TypeError(f"Use envs.{self.name}.get() to read {self.name}.") + + +class EnvBool(EnvField[bool]): + """True for 1, true, yes, on; false for 0, false, no, off, and the empty string; case-insensitive.""" + + type_name = "bool" + + def parse(self, raw: str) -> bool: + value = raw.strip().lower() + if value in _TRUE_VALUES: + return True + if value in _FALSE_VALUES: + return False + raise ValueError("expected 1, true, yes, on, 0, false, no, off, or an empty string") + + def format(self, value: bool) -> str: + return "1" if value else "0" + + +class EnvInt(EnvField[int]): + type_name = "int" + + def parse(self, raw: str) -> int: + return int(raw) + + +class EnvFloat(EnvField[float]): + type_name = "float" + + def parse(self, raw: str) -> float: + return float(raw) + + +class EnvStr(EnvField[S]): + type_name = "str" + + def parse(self, raw: str) -> S: + return raw + + +class EnvPath(EnvField[S]): + """A filesystem path; a leading ``~`` is expanded.""" + + type_name = "path" + + def parse(self, raw: str) -> S: + return os.path.expanduser(raw) + + +class EnvChoice(EnvField[str]): + """One of ``choices``; the value is stripped and lower-cased before the check.""" + + def __init__(self, default: str, *, choices: tuple[str, ...], category: str, doc: str) -> None: + super().__init__(default, category=category, doc=doc) + self.choices = choices + self.type_name = "one of " + ", ".join(choices) + + def parse(self, raw: str) -> str: + value = raw.strip().lower() + if value not in self.choices: + raise ValueError(f"expected one of {', '.join(self.choices)}") + return value def get_default_cache_root() -> str: @@ -67,318 +171,228 @@ def get_default_config_root() -> str: ) -def maybe_convert_int(value: str | None) -> int | None: - if value is None: - return None - return int(value) +# ================== Installation ================== + +FASTVIDEO_TARGET_DEVICE = EnvStr("cuda", + category="build", + doc="Target device of FastVideo: cuda, rocm, neuron, cpu, or openvino.") +MAX_JOBS = EnvStr(None, + category="build", + doc="Maximum number of parallel compilation jobs. Defaults to the number of CPUs.") +NVCC_THREADS = EnvStr(None, + category="build", + doc="Number of nvcc threads. When set, MAX_JOBS is reduced to avoid oversubscribing the CPU.") +FASTVIDEO_USE_PRECOMPILED = EnvBool(False, category="build", doc="Use precompiled binaries (*.so).") +CMAKE_BUILD_TYPE = EnvStr(None, category="build", doc="CMake build type: Debug, Release, or RelWithDebInfo.") +VERBOSE = EnvBool(False, category="build", doc="Print verbose logs during installation.") + +# ================== Paths ================== + +FASTVIDEO_CONFIG_ROOT = EnvPath( + lambda: os.path.expanduser(os.path.join(get_default_config_root(), "fastvideo")), + category="path", + doc="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 = EnvPath( + lambda: os.path.expanduser(os.path.join(get_default_cache_root(), "fastvideo")), + category="path", + doc="Root directory for FastVideo cache files. " + "Defaults to ~/.cache/fastvideo, or $XDG_CACHE_HOME/fastvideo when XDG_CACHE_HOME is set.") + +# ================== Distributed ================== + +FASTVIDEO_HOST_IP = EnvStr( + "", + category="distributed", + doc="IP address of this node when the node has several network interfaces. Set it on each node for multi-node " + "inference.") +FASTVIDEO_LOOPBACK_IP = EnvStr("", + category="distributed", + doc="Loopback IP address to use instead of the detected one.") +FASTVIDEO_RAY_PER_WORKER_GPUS = EnvFloat( + 1.0, + category="distributed", + doc="GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share " + "the GPUs with FastVideo.") +FASTVIDEO_RINGBUFFER_WARNING_INTERVAL = EnvInt(60, + category="distributed", + doc="Seconds between warnings while the ring buffer is full.") +FASTVIDEO_NCCL_SO_PATH = EnvStr( + None, + category="distributed", + doc="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).") +HCCL_SO_PATH = EnvStr(None, category="distributed", doc="Path to the HCCL library file on Ascend NPUs.") +FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S = EnvInt(60, + category="distributed", + doc="Timeout in seconds for each engine iteration.") +FASTVIDEO_WORKER_MULTIPROC_METHOD = EnvChoice("spawn", + choices=("spawn", "fork", "forkserver"), + category="distributed", + doc="Multiprocessing start method for worker processes.") +FASTVIDEO_ULYSSES_A2A = EnvChoice( + "off", + choices=("off", "auto"), + category="distributed", + doc="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.") + +# ================== External variables ================== +# Variables that other tools set. They stay registered so that Ray copies them +# to its workers until the external-variable allowlist replaces them. + +LD_LIBRARY_PATH = EnvStr(None, + category="external", + doc="Searched for the NCCL library when FASTVIDEO_NCCL_SO_PATH is unset.") +LOCAL_RANK = EnvInt(0, + category="external", + doc="Local rank of the process in a distributed run; selects the GPU device id.") +CUDA_VISIBLE_DEVICES = EnvStr(None, category="external", doc="Visible devices in a distributed run.") + +# ================== Logging ================== + +FASTVIDEO_CONFIGURE_LOGGING = EnvBool( + True, + category="logging", + doc="Configure logging at import. When true, FastVideo uses its default logging configuration or the file in " + "FASTVIDEO_LOGGING_CONFIG_PATH.") +FASTVIDEO_LOGGING_CONFIG_PATH = EnvStr(None, category="logging", doc="Path to a JSON logging configuration file.") +FASTVIDEO_LOGGING_LEVEL = EnvStr("INFO", category="logging", doc="Default logging level.") +FASTVIDEO_LOGGING_PREFIX = EnvStr("", category="logging", doc="Prefix prepended to every log message.") +FASTVIDEO_STAGE_LOGGING = EnvBool(False, category="logging", doc="Log the time that each pipeline stage takes.") + +# ================== Attention ================== + +FASTVIDEO_ATTENTION_BACKEND = EnvStr( + None, + category="attention", + doc="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.") +# FA4 is opt-in and never auto-selected just because it is installed. Below +# sm90, grad-enabled and GQA calls are routed to FA2 (FA4's backward asserts +# sm90+ and its pack_gqa fails to JIT there). +FASTVIDEO_FA4 = EnvBool(False, + category="attention", + doc="The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2.") +FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN = EnvBool( + False, + category="attention", + doc="MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the " + "floating-point reduction order, so it is an inference-only opt-in.") +FASTVIDEO_VSA_SM100A = EnvBool( + False, + category="attention", + doc="VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. " + "fastvideo-kernel reads the same variable with the same rule.") + +# ================== Performance ================== + +# Non-fullgraph-traceable attention backends such as VSA degrade to eager with +# one warning; see _regional_compile_unsupported_reason in +# fastvideo/models/loader/fsdp_load.py. +FASTVIDEO_INFERENCE_TORCH_COMPILE = EnvBool( + False, + category="performance", + doc="Compile each DiT transformer block with fullgraph torch.compile at inference. Same as " + "FastVideoArgs.inference_torch_compile=True.") +FASTVIDEO_VAE_PARALLEL_DECODE = EnvBool( + False, + category="performance", + doc="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 = EnvBool( + False, + category="performance", + doc="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 = EnvStr( + None, + category="performance", + doc="Collective that moves chunks in parallel VAE decode: gather (used when unset) or all_gather.") +# Adapted from the NVlabs/Sana Sol-Engine implementation. +FASTVIDEO_MINIMAX_H3_FUSIONS = EnvStr( + "", + category="performance", + doc="MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of " + "modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.") +FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE = EnvBool(True, category="debug", doc="Enable Dynamo fullgraph capture.") + +# ================== Profiling ================== + +FASTVIDEO_NVTX_PROFILE = EnvBool(False, + category="profiling", + doc="Emit NVTX ranges for external profilers such as Nsight Systems.") +FASTVIDEO_TORCH_PROFILER_DIR = EnvPath( + None, + category="profiling", + doc="Enables the torch profiler and sets the directory for its traces. Must be an absolute path.") +FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES = EnvBool(False, category="profiling", doc="Torch profiler records shapes.") +FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY = EnvBool(False, + category="profiling", + doc="Torch profiler profiles memory.") +FASTVIDEO_TORCH_PROFILER_WITH_STACK = EnvBool( + False, category="profiling", doc="Torch profiler captures stacks. Costs about 1.5x runtime and 1.4x trace size.") +FASTVIDEO_TORCH_PROFILER_WITH_FLOPS = EnvBool(False, category="profiling", doc="Torch profiler profiles FLOPs.") +FASTVIDEO_TORCH_PROFILE_REGIONS = EnvStr( + "", + category="profiling", + doc="Comma-separated profiler regions to record. The torch profiler requires at least one region.") + +# ================== Debug ================== + +FASTVIDEO_SERVER_DEV_MODE = EnvBool(False, + category="debug", + doc="Run the server in development mode with extra debugging endpoints.") +FASTVIDEO_TRACE_FUNCTION = EnvBool(False, category="debug", doc="Trace function calls.") +FASTVIDEO_TRACE_ACTIVATIONS = EnvBool(False, category="debug", doc="Enable activation trace hooks.") +FASTVIDEO_TRACE_LAYERS = EnvStr("", category="debug", doc="Regex filter for traced module names. Empty means all.") +FASTVIDEO_TRACE_STATS = EnvStr("abs_mean,sum", + category="debug", + doc="Comma-separated activation statistics dumped for each output tensor.") +FASTVIDEO_TRACE_OUTPUT = EnvStr("/tmp/fv_trace_.jsonl", + category="debug", + doc="JSONL path for activation traces. The literal is replaced at runtime.") +FASTVIDEO_TRACE_STEPS = EnvStr("", category="debug", doc="Comma-separated denoising step indices. Empty means all.") + +# ================== Sampling ================== + +# CFG gating fraction for stale-uncond reuse (Adaptive Guidance / LinearAG +# variant — Castillo et al. 2023, arXiv:2312.12487). Float in [0, 1]. +# Interpretation: for step index `i < len(timesteps) * X`, run both +# cond and uncond forwards and refresh delta_cached = cond - uncond. +# Once `i >= len(timesteps) * X`, skip the uncond forward and reuse +# the cached delta: noise_pred = cond + (guidance_scale - 1) * delta. +# +# Edge cases: +# 1.0 (default) : disables gating; identical to baseline two-pass CFG. +# 0.5 : run uncond for the first half of steps, reuse delta +# for the second half (~25% inference time saved on +# bandwidth-bound SP setups). +# 0.0 : step 0 still computes uncond fresh (cache is empty +# at start) — all subsequent steps reuse the step-0 +# delta. This is the most aggressive setting; does +# NOT mean "no uncond forward ever." +# +# Caveats: +# - Algorithmically approximate; not bit-exact vs baseline CFG. +# Validate per-pipeline with SSIM / VBench before lowering below 1.0. +# - Interaction with `guidance_rescale > 0` is unvalidated; the +# denoising stage logs a warning when both are active. +# - Wan2.2 high/low-noise expert switch invalidates the cache. +FASTVIDEO_CFG_GATE_STEP = EnvFloat( + 1.0, + category="sampling", + doc="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.") -# The begin-* and end* here are used by the documentation generator -# to extract the used env vars. - -# begin-env-vars-definition - -environment_variables: dict[str, Callable[[], Any]] = { - - # ================== Installation Time Env Vars ================== - - # Target device of FastVideo, supporting [cuda (by default), - # rocm, neuron, cpu, openvino] - "FASTVIDEO_TARGET_DEVICE": - lambda: os.getenv("FASTVIDEO_TARGET_DEVICE", "cuda"), - - # Maximum number of compilation jobs to run in parallel. - # By default this is the number of CPUs - "MAX_JOBS": - lambda: os.getenv("MAX_JOBS", None), - - # Number of threads to use for nvcc - # By default this is 1. - # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. - "NVCC_THREADS": - lambda: os.getenv("NVCC_THREADS", None), - - # If set, fastvideo will use precompiled binaries (*.so) - "FASTVIDEO_USE_PRECOMPILED": - lambda: bool(os.environ.get("FASTVIDEO_USE_PRECOMPILED")) or bool( - os.environ.get("FASTVIDEO_PRECOMPILED_WHEEL_LOCATION")), - - # CMake build type - # If not set, defaults to "Debug" or "RelWithDebInfo" - # Available options: "Debug", "Release", "RelWithDebInfo" - "CMAKE_BUILD_TYPE": - lambda: os.getenv("CMAKE_BUILD_TYPE"), - - # If set, fastvideo will print verbose logs during installation - "VERBOSE": - lambda: bool(int(os.getenv('VERBOSE', '0'))), - - # Root directory for FASTVIDEO configuration files - # Defaults to `~/.config/fastvideo` unless `XDG_CONFIG_HOME` is set - # Note that this not only affects how fastvideo finds its configuration files - # during runtime, but also affects how fastvideo installs its configuration - # files during **installation**. - "FASTVIDEO_CONFIG_ROOT": - lambda: os.path.expanduser( - os.getenv( - "FASTVIDEO_CONFIG_ROOT", - os.path.join(get_default_config_root(), "fastvideo"), - )), - - # ================== Runtime Env Vars ================== - - # Root directory for FASTVIDEO cache files - # Defaults to `~/.cache/fastvideo` unless `XDG_CACHE_HOME` is set - "FASTVIDEO_CACHE_ROOT": - lambda: os.path.expanduser(os.getenv( - "FASTVIDEO_CACHE_ROOT", - os.path.join(get_default_cache_root(), "fastvideo"), - )), - - # used in distributed environment to determine the ip address - # of the current node, when the node has multiple network interfaces. - # If you are using multi-node inference, you should set this differently - # on each node. - "FASTVIDEO_HOST_IP": - lambda: os.getenv("FASTVIDEO_HOST_IP", ""), - - # Used to force set up loopback IP - "FASTVIDEO_LOOPBACK_IP": - lambda: os.getenv("FASTVIDEO_LOOPBACK_IP", ""), - - # Number of GPUs per worker in Ray, if it is set to be a fraction, - # it allows ray to schedule multiple actors on a single GPU, - # so that users can colocate other actors on the same GPUs as FastVideo. - "FASTVIDEO_RAY_PER_WORKER_GPUS": - lambda: float(os.getenv("FASTVIDEO_RAY_PER_WORKER_GPUS", "1.0")), - - # Interval in seconds to log a warning message when the ring buffer is full - "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL": - lambda: int(os.environ.get("FASTVIDEO_RINGBUFFER_WARNING_INTERVAL", "60")), - - # Path to the NCCL library file. It is needed because nccl>=2.19 brought - # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234 - "FASTVIDEO_NCCL_SO_PATH": - lambda: os.environ.get("FASTVIDEO_NCCL_SO_PATH", None), - - # when `FASTVIDEO_NCCL_SO_PATH` is not set, fastvideo will try to find the nccl - # library file in the locations specified by `LD_LIBRARY_PATH` - "LD_LIBRARY_PATH": - lambda: os.environ.get("LD_LIBRARY_PATH", None), - - # Internal flag to enable Dynamo fullgraph capture - "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE": - lambda: bool(os.environ.get("FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"), - - # local rank of the process in the distributed setting, used to determine - # the GPU device id - "LOCAL_RANK": - lambda: int(os.environ.get("LOCAL_RANK", "0")), - - # used to control the visible devices in the distributed setting - "CUDA_VISIBLE_DEVICES": - lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), - - # timeout for each iteration in the engine - "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S": - lambda: int(os.environ.get("FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S", "60")), - - # Logging configuration - # If set to 0, fastvideo will not configure logging - # If set to 1, fastvideo will configure logging using the default configuration - # or the configuration file specified by FASTVIDEO_LOGGING_CONFIG_PATH - "FASTVIDEO_CONFIGURE_LOGGING": - lambda: int(os.getenv("FASTVIDEO_CONFIGURE_LOGGING", "1")), - "FASTVIDEO_LOGGING_CONFIG_PATH": - lambda: os.getenv("FASTVIDEO_LOGGING_CONFIG_PATH"), - - # this is used for configuring the default logging level - "FASTVIDEO_LOGGING_LEVEL": - lambda: os.getenv("FASTVIDEO_LOGGING_LEVEL", "INFO"), - - # if set, FASTVIDEO_LOGGING_PREFIX will be prepended to all log messages - "FASTVIDEO_LOGGING_PREFIX": - lambda: os.getenv("FASTVIDEO_LOGGING_PREFIX", ""), - - # Trace function calls - # If set to 1, fastvideo will trace function calls - # Useful for debugging - "FASTVIDEO_TRACE_FUNCTION": - lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")), - - # Backend for attention computation - # Available options: - # - "TORCH_SDPA": use torch.nn.MultiheadAttention - # - "FLASH_ATTN": use FlashAttention - # - "VIDEO_SPARSE_ATTN": use Video Sparse Attention - # - "SAGE_ATTN": use Sage Attention - # - "SAGE_ATTN_THREE": use Sage Attention 3 - # FLASH_ATTN uses FlashAttention-3/2; to run FlashAttention-4 set - # FASTVIDEO_FA4=1 as well (see below). - "FASTVIDEO_ATTENTION_BACKEND": - lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None), - - # If set (=1), the FLASH_ATTN backend uses FlashAttention-4 - # (flash_attn.cute). FA4 is opt-in and never auto-selected just because it - # is installed. Below sm90, grad-enabled and GQA calls are routed to FA2 - # (FA4's backward asserts sm90+ and its pack_gqa fails to JIT there). - "FASTVIDEO_FA4": - lambda: os.getenv("FASTVIDEO_FA4", "0") != "0", - - # Use FA4's packed-varlen entry point for the long, single-document - # MiniMax-H3 dense DiT self-attention path. This changes floating-point - # reduction order relative to the fixed-length entry point, so it remains - # an explicit inference-only speed/quality opt-in. - "FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN": - lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN", "0") != "0", - - # If set (=1), enable regional (per-transformer-block) fullgraph - # torch.compile for the DiT at inference — the inference-side counterpart - # of the training regional-compile port of hao-ai-lab/FastVideo#1718. - # Equivalent to FastVideoArgs.inference_torch_compile=True (e.g. via - # PipelineSelection.experimental={"inference_torch_compile": True}). VSA - # and other non-fullgraph-traceable attention backends degrade to eager - # with one warning; see _regional_compile_unsupported_reason in - # fastvideo/models/loader/fsdp_load.py. - "FASTVIDEO_INFERENCE_TORCH_COMPILE": - lambda: os.getenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "0") != "0", - - # If set (=1), MiniMax-H3 VAE decode (and, with the ENCODE variant, - # reference-video encode) round-robins its temporal chunks across the - # sequence-parallel ranks instead of running serially on the output rank. - # Folded into FastVideoArgs.vae_parallel_decode / vae_parallel_encode at - # construction (parse-once). The STRATEGY variant picks the chunk - # transport collective: "gather" (default) or "all_gather". - "FASTVIDEO_VAE_PARALLEL_DECODE": - lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_DECODE", "0") != "0", - "FASTVIDEO_VAE_PARALLEL_ENCODE": - lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_ENCODE", "0") != "0", - "FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY": - lambda: os.getenv("FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY", None), - - # Opt-in MiniMax-H3 inference-only Triton fusions adapted from the - # NVlabs/Sana Sol-Engine implementation. Accepts `all`, `1`, or a - # comma-separated subset of `modulate,qknorm_rope,swiglu`. An empty value - # (the default), `0`, or `none` keeps the eager implementation. - "FASTVIDEO_MINIMAX_H3_FUSIONS": - lambda: os.getenv("FASTVIDEO_MINIMAX_H3_FUSIONS", ""), - - # Sequence-parallel all-to-all backend. - # - "off" (default): the NCCL path in DistributedAutograd.AllToAll4D - # - "auto": fused NVLink kernel when the group is a load-store accessible - # mesh of 2/4/6/8 ranks in eager execution, else the NCCL path. - "FASTVIDEO_ULYSSES_A2A": - lambda: os.getenv("FASTVIDEO_ULYSSES_A2A", "off").strip().lower(), - - # Use dedicated multiprocess context for workers. - "FASTVIDEO_WORKER_MULTIPROC_METHOD": - lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"), - - # Emit lightweight NVTX ranges for external profilers such as Nsight Systems. - "FASTVIDEO_NVTX_PROFILE": - lambda: os.getenv("FASTVIDEO_NVTX_PROFILE", "0") != "0", - - # Enables torch profiler if set. Path to the directory where torch profiler - # traces are saved. Note that it must be an absolute path. - "FASTVIDEO_TORCH_PROFILER_DIR": - lambda: (None if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None else os.path.expanduser( - os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", "."))), - - # Enable torch profiler to record shapes if set - # FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1. If not set, torch profiler will - # not record shapes. - "FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES": - lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES", "0") != "0"), - - # Enable torch profiler to profile memory if set - # FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY=1. If not set, torch profiler - # will not profile memory. - "FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY": - lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"), - - # Enable torch profiler stack capture with - # FASTVIDEO_TORCH_PROFILER_WITH_STACK=1. Off by default: stack capture - # costs ~1.5x runtime overhead and ~1.4x trace size. - "FASTVIDEO_TORCH_PROFILER_WITH_STACK": - lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_STACK", "0") != "0"), - - # Enable torch profiler to profile flops if set - # FASTVIDEO_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will - # not profile flops. - "FASTVIDEO_TORCH_PROFILER_WITH_FLOPS": - lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"), - # Wait steps per profiling cycle (torch.profiler.schedule wait parameter) - # Defaults to 2 if not set. - # Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter) - # Defaults to 1 if not set. - # Active steps per profiling cycle (torch.profiler.schedule active parameter) - # Defaults to 2 if not set. - "FASTVIDEO_TORCH_PROFILE_REGIONS": - lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""), - - # Enable activation trace hooks if set. - "FASTVIDEO_TRACE_ACTIVATIONS": - lambda: bool(os.getenv("FASTVIDEO_TRACE_ACTIVATIONS", "0") != "0"), - # Regex filter for traced module names. Empty means all modules. - "FASTVIDEO_TRACE_LAYERS": - lambda: os.getenv("FASTVIDEO_TRACE_LAYERS", ""), - # Comma-separated activation stats to dump for each output tensor. - "FASTVIDEO_TRACE_STATS": - lambda: os.getenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum"), - # JSONL sink path. The literal is replaced at runtime. - "FASTVIDEO_TRACE_OUTPUT": - lambda: os.getenv("FASTVIDEO_TRACE_OUTPUT", "/tmp/fv_trace_.jsonl"), - # Comma-separated denoise step indices. Empty means all steps. - "FASTVIDEO_TRACE_STEPS": - lambda: os.getenv("FASTVIDEO_TRACE_STEPS", ""), - - # If set, fastvideo will run in development mode, which will enable - # some additional endpoints for developing and debugging, - # e.g. `/reset_prefix_cache` - "FASTVIDEO_SERVER_DEV_MODE": - lambda: bool(int(os.getenv("FASTVIDEO_SERVER_DEV_MODE", "0"))), - - # If set, fastvideo will enable stage logging, which will print the time - # taken for each stage - "FASTVIDEO_STAGE_LOGGING": - lambda: bool(int(os.getenv("FASTVIDEO_STAGE_LOGGING", "0"))), - - # CFG gating fraction for stale-uncond reuse (Adaptive Guidance / LinearAG - # variant — Castillo et al. 2023, arXiv:2312.12487). Float in [0, 1]. - # Interpretation: for step index `i < len(timesteps) * X`, run both - # cond and uncond forwards and refresh delta_cached = cond - uncond. - # Once `i >= len(timesteps) * X`, skip the uncond forward and reuse - # the cached delta: noise_pred = cond + (guidance_scale - 1) * delta. - # - # Edge cases: - # 1.0 (default) : disables gating; identical to baseline two-pass CFG. - # 0.5 : run uncond for the first half of steps, reuse delta - # for the second half (~25% inference time saved on - # bandwidth-bound SP setups). - # 0.0 : step 0 still computes uncond fresh (cache is empty - # at start) — all subsequent steps reuse the step-0 - # delta. This is the most aggressive setting; does - # NOT mean "no uncond forward ever." - # - # Caveats: - # - Algorithmically approximate; not bit-exact vs baseline CFG. - # Validate per-pipeline with SSIM / VBench before lowering below 1.0. - # - Interaction with `guidance_rescale > 0` is unvalidated; the - # denoising stage logs a warning when both are active. - # - Wan2.2 high/low-noise expert switch invalidates the cache. - "FASTVIDEO_CFG_GATE_STEP": - lambda: float(os.getenv("FASTVIDEO_CFG_GATE_STEP", "1.0")), -} - -# end-env-vars-definition +def _register_fields() -> dict[str, EnvField]: + """Name each module-level EnvField after its attribute and return the fields by name.""" + fields = {name: value for name, value in globals().items() if isinstance(value, EnvField)} + for name, field in fields.items(): + field.name = name + return fields -def __getattr__(name: str): - # lazy evaluation of environment variables - if name in environment_variables: - return environment_variables[name]() - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -def __dir__(): - return list(environment_variables.keys()) +environment_variables: dict[str, EnvField] = _register_fields() diff --git a/fastvideo/eval/models.py b/fastvideo/eval/models.py index 8b549e9af..9adb57b3e 100644 --- a/fastvideo/eval/models.py +++ b/fastvideo/eval/models.py @@ -46,7 +46,7 @@ def get_cache_dir() -> Path: """ return Path(os.environ.get( "FASTVIDEO_EVAL_CACHE", - os.path.join(envs.FASTVIDEO_CACHE_ROOT, "eval"), + os.path.join(envs.FASTVIDEO_CACHE_ROOT.get(), "eval"), )) diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index a648a0253..8cc5a60bf 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -343,7 +343,7 @@ class FastVideoArgs: # environment variable is an input read once here, so the loader # only ever consults the typed field. import fastvideo.envs as envs - if envs.FASTVIDEO_INFERENCE_TORCH_COMPILE: + if envs.FASTVIDEO_INFERENCE_TORCH_COMPILE.get(): self.inference_torch_compile = True if self.attention_backend is not None: # Fail fast on typos instead of silently auto-selecting later. @@ -357,7 +357,7 @@ class FastVideoArgs: # and falls through to automatic selection rather than raising. import fastvideo.envs as envs from fastvideo.attention.selector import backend_name_to_enum - env_backend = envs.FASTVIDEO_ATTENTION_BACKEND + 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 self._fold_vae_parallel_env() @@ -371,12 +371,12 @@ class FastVideoArgs: # DEFAULT_DECODE_GATHER_STRATEGY (kept literal here so constructing args # never imports model modules; a unit test pins the two in sync). strategies = ("gather", "all_gather") - if not self.vae_parallel_decode and envs.FASTVIDEO_VAE_PARALLEL_DECODE: + if not self.vae_parallel_decode and envs.FASTVIDEO_VAE_PARALLEL_DECODE.get(): self.vae_parallel_decode = True - if not self.vae_parallel_encode and envs.FASTVIDEO_VAE_PARALLEL_ENCODE: + if not self.vae_parallel_encode and envs.FASTVIDEO_VAE_PARALLEL_ENCODE.get(): self.vae_parallel_encode = True if self.vae_parallel_decode_strategy is None: - self.vae_parallel_decode_strategy = envs.FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY or "gather" + self.vae_parallel_decode_strategy = envs.FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY.get() or "gather" if self.vae_parallel_decode_strategy not in strategies: raise ValueError(f"vae_parallel_decode_strategy must be one of {strategies}, " f"got {self.vae_parallel_decode_strategy!r}.") diff --git a/fastvideo/hooks/activation_trace.py b/fastvideo/hooks/activation_trace.py index c6a47a23b..b28d0608e 100644 --- a/fastvideo/hooks/activation_trace.py +++ b/fastvideo/hooks/activation_trace.py @@ -187,18 +187,18 @@ class ActivationTraceManager: def attach_activation_trace(model: nn.Module | None) -> ActivationTraceManager | None: """Attach activation-stat hooks to model. Returns None if trace is off.""" - if not envs.FASTVIDEO_TRACE_ACTIVATIONS or model is None: + if not envs.FASTVIDEO_TRACE_ACTIVATIONS.get() or model is None: return None - pattern_spec = envs.FASTVIDEO_TRACE_LAYERS + pattern_spec = envs.FASTVIDEO_TRACE_LAYERS.get() pattern = re.compile(pattern_spec) if pattern_spec else re.compile(".*") - stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS) + stats = _resolve_stats(envs.FASTVIDEO_TRACE_STATS.get()) if not stats: logger.warning("FASTVIDEO_TRACE_STATS yielded no valid stats; trace disabled.") return None - sink = JsonlSink(_resolve_output_path(envs.FASTVIDEO_TRACE_OUTPUT)) - step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS) + sink = JsonlSink(_resolve_output_path(envs.FASTVIDEO_TRACE_OUTPUT.get())) + step_filter = _parse_step_filter(envs.FASTVIDEO_TRACE_STEPS.get()) managers = [] for name, module in model.named_modules(): if not name or not pattern.search(name): diff --git a/fastvideo/logger.py b/fastvideo/logger.py index 84f9958b4..a930ae0e4 100644 --- a/fastvideo/logger.py +++ b/fastvideo/logger.py @@ -15,10 +15,10 @@ from typing import Any, cast import fastvideo.envs as envs -FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING -FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH -FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL -FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX +FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING.get() +FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH.get() +FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL.get() +FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX.get() RED = '\033[91m' GREEN = '\033[92m' diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 370d11d47..1cd1d16d0 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -47,7 +47,7 @@ _MINIMAX_H3_FUSION_NAMES = frozenset({"modulate", "qknorm_rope", "swiglu"}) def _enabled_minimax_h3_fusions(value: str | None = None) -> frozenset[str]: """Parse the independently switchable inference fusion set.""" - raw = envs.FASTVIDEO_MINIMAX_H3_FUSIONS if value is None else value + raw = envs.FASTVIDEO_MINIMAX_H3_FUSIONS.get() if value is None else value normalized = raw.strip().lower() if normalized in {"", "0", "none"}: return frozenset() @@ -720,7 +720,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): fuse_modulate="modulate" in self.enabled_fusions, fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions, fuse_swiglu="swiglu" in self.enabled_fusions, - fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN, + fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(), ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index e076359ee..f7c31d221 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -19,6 +19,7 @@ from torch.distributed._tensor import distribute_tensor from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard) from torch.nn.modules.module import _IncompatibleKeys +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.models.loader.lora_patch import DenseLoRAPatch from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict) @@ -394,7 +395,7 @@ def _regional_compile_unsupported_reason( if os.environ.get("FASTVIDEO_H3_VSA_PROBE"): return ("FASTVIDEO_H3_VSA_PROBE records tensors and files from the VSA-H3 attention body, which " "regional fullgraph compile cannot capture; this model stays eager") - if os.environ.get("FASTVIDEO_VSA_SM100A", "0") != "1": + if not envs.FASTVIDEO_VSA_SM100A.get(): return ("VIDEO_SPARSE_ATTN_H3 regional compile requires the compile-safe sm_100a route " "(FASTVIDEO_VSA_SM100A=1); Triton/CuTe VSA stays eager") if vsa_tile_size != 64: diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index 183a34af4..d5c0d8230 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -162,7 +162,7 @@ class ComposedPipelineBase(ABC): # Torch profiler. Enabled and configured through env vars: # FASTVIDEO_TORCH_PROFILER_DIR=/path/to/save/trace - trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR + trace_dir = envs.FASTVIDEO_TORCH_PROFILER_DIR.get() self.profiler_controller = get_or_create_profiler(trace_dir) self.local_rank = get_world_group().local_rank @@ -378,7 +378,7 @@ class ComposedPipelineBase(ABC): # The hook manager keeps a strong reference to every module it # wraps, so attaching here would materialize the DiT before the # first request and pin that instance past any release. - if envs.FASTVIDEO_TRACE_ACTIVATIONS: + if envs.FASTVIDEO_TRACE_ACTIVATIONS.get(): logger.warning("Activation trace is not attached to a deferred transformer; " "turn off lazy_module_load to trace it") trace_target = None diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 0f990b256..838ce677a 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -178,7 +178,7 @@ class PipelineStage(ABC): stage_name: str, ) -> ForwardBatch: """Run forward, with the optional timing and logging wrapper.""" - if envs.FASTVIDEO_STAGE_LOGGING: + if envs.FASTVIDEO_STAGE_LOGGING.get(): logger.info("[%s] Starting execution", stage_name) torch.cuda.synchronize() start_time = time.perf_counter() diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 79dbca32c..f36509eb3 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -283,7 +283,7 @@ class DenoisingStage(PipelineStage): # compute. See envs.py for semantics. delta_cached_model_id tracks # which underlying transformer produced the cache so we invalidate on # Wan2.2 expert switch. - _cfg_gate_fraction = envs.FASTVIDEO_CFG_GATE_STEP + _cfg_gate_fraction = envs.FASTVIDEO_CFG_GATE_STEP.get() if not 0.0 <= _cfg_gate_fraction <= 1.0: raise ValueError(f"FASTVIDEO_CFG_GATE_STEP must be in [0.0, 1.0], got {_cfg_gate_fraction!r}. " "Use 1.0 (default) to disable; lower values trade quality for speed.") diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index b6f8a9cb3..78c6cdc04 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -171,7 +171,7 @@ class CudaPlatformBase(Platform): # TODO(will): maybe come up with a more general interface for local attention # if distributed is False, we always try to use Flash attn - logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND) + logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND.get()) logger.info("Selected backend: %s", selected_backend) if selected_backend == AttentionBackendEnum.SAGE_ATTN: try: diff --git a/fastvideo/platforms/npu.py b/fastvideo/platforms/npu.py index 23e162dd1..c5014b1fa 100644 --- a/fastvideo/platforms/npu.py +++ b/fastvideo/platforms/npu.py @@ -66,8 +66,8 @@ class NPUPlatform(Platform): @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: - logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND) - if envs.FASTVIDEO_ATTENTION_BACKEND != "TORCH_SDPA": + logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND.get()) + if envs.FASTVIDEO_ATTENTION_BACKEND.get() != "TORCH_SDPA": logger.info("Ascend NPU only supports the Torch SDPA backend.") else: logger.info("Using Torch SDPA backend.") diff --git a/fastvideo/platforms/rocm.py b/fastvideo/platforms/rocm.py index 461104621..4b8b8315d 100644 --- a/fastvideo/platforms/rocm.py +++ b/fastvideo/platforms/rocm.py @@ -63,7 +63,7 @@ class RocmPlatform(Platform): @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: - logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND) + logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND.get()) if selected_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend.") diff --git a/fastvideo/profiler.py b/fastvideo/profiler.py index d75879d7b..bcfc9ab84 100644 --- a/fastvideo/profiler.py +++ b/fastvideo/profiler.py @@ -45,7 +45,7 @@ def nvtx_range(name: str): ``FASTVIDEO_NVTX_PROFILE=1`` enables the marker. The context manager stays a no-op without CUDA so call sites can remain shared with CPU tests. """ - enabled = envs.FASTVIDEO_NVTX_PROFILE and torch.cuda.is_available() + enabled = envs.FASTVIDEO_NVTX_PROFILE.get() and torch.cuda.is_available() if not enabled: yield return @@ -201,19 +201,19 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController: logger.info("Profiling enabled. Traces will be saved to: %s", trace_dir) logger.info( "Profiler config: record_shapes=%s, profile_memory=%s, with_stack=%s, with_flops=%s", - envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES, - envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY, - envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK, - envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS, + envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES.get(), + envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY.get(), + envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK.get(), + envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS.get(), ) - logger.info("FASTVIDEO_TORCH_PROFILE_REGIONS=%s", envs.FASTVIDEO_TORCH_PROFILE_REGIONS) + logger.info("FASTVIDEO_TORCH_PROFILE_REGIONS=%s", envs.FASTVIDEO_TORCH_PROFILE_REGIONS.get()) profiler = torch.profiler.profile( activities=_DEFAULT_ACTIVITIES, - record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES, - profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY, - with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK, - with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS, + record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES.get(), + profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY.get(), + with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK.get(), + with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS.get(), # No schedule: nothing in the codebase calls profiler.step(), so a # wait/warmup schedule never advances and the profiler records nothing. # Region toggling gates collection; the single trace exports at stop(). @@ -249,7 +249,7 @@ class TorchProfilerConfig: requested_regions = { token.strip() - for token in (getattr(envs, "FASTVIDEO_TORCH_PROFILE_REGIONS", "") or "").split(",") if token.strip() + for token in envs.FASTVIDEO_TORCH_PROFILE_REGIONS.get().split(",") if token.strip() } if not requested_regions: diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py new file mode 100644 index 000000000..571e22228 --- /dev/null +++ b/fastvideo/tests/contract/test_env_policy.py @@ -0,0 +1,827 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Guard: FastVideo code follows the environment-variable policy in +docs/contributing/env_vars.md. + +The test parses every Python file under fastvideo/ (except fastvideo/third_party/ +and the registry fastvideo/envs.py) with ``ast`` and reports code that: + +- reads the environment directly for a name outside EXTERNAL_ALLOWLIST; +- writes the environment directly (os.environ, os.putenv, monkeypatch.setenv); +- uses the whole environment (os.environ.copy(), dict(os.environ), patch.dict); +- uses a registry field without calling one of its methods (``envs.X == "a"``, + ``getter = envs.X.get``); +- calls ``envs.X.get()`` outside a function, so the value is read at import. + +It also checks every entry in fastvideo/envs.py (FASTVIDEO_ prefix, category, +description, at least one reader) and checks that the table in +docs/contributing/env_vars.md matches the registry. + +Violations that existed when the policy was introduced are listed in +KNOWN_VIOLATIONS. The list only shrinks: a violation that is not listed fails, +and a listed violation that no longer exists also fails, so the entry gets +deleted. + +Run ``python fastvideo/tests/contract/test_env_policy.py`` to regenerate the +table in docs/contributing/env_vars.md after editing fastvideo/envs.py. + +Static analysis only: fastvideo/envs.py is loaded as a standalone file, so the +test imports neither fastvideo nor torch. +""" +import ast +import importlib.util +import re +from collections import Counter +from pathlib import Path +from types import ModuleType + +REPO_ROOT = Path(__file__).resolve().parents[3] +PACKAGE_ROOT = REPO_ROOT / "fastvideo" +REGISTRY_PATH = PACKAGE_ROOT / "envs.py" +DOC_PATH = REPO_ROOT / "docs" / "contributing" / "env_vars.md" +POLICY_DOC = "docs/contributing/env_vars.md" +EXCLUDED_DIRS = (PACKAGE_ROOT / "third_party", ) + +DOC_TABLE_BEGIN = "" +DOC_TABLE_END = "" + +# Variables that other tools own. Code may read them directly with a literal +# name. A trailing "*" allows every name with that prefix. +EXTERNAL_ALLOWLIST = { + "CUDA_*": "CUDA runtime and device selection.", + "NCCL_*": "NCCL communication library.", + "TORCH_*": "PyTorch runtime settings.", + "OMP_*": "OpenMP threading.", + "HF_HUB_*": "huggingface_hub settings.", + "RANK": "Set by torchrun and other launchers.", + "LOCAL_RANK": "Set by torchrun and other launchers.", + "WORLD_SIZE": "Set by torchrun and other launchers.", + "MASTER_ADDR": "Set by torchrun and other launchers.", + "MASTER_PORT": "Set by torchrun and other launchers.", + "HOME": "User home directory.", + "PATH": "Executable search path.", + "HF_TOKEN": "Token variable that huggingface_hub reads.", + "HUGGING_FACE_HUB_TOKEN": "Token variable that huggingface_hub reads.", +} + +# Methods that registry fields expose; ``get`` and ``is_set`` count as reads. +REGISTRY_METHODS = {"get", "set", "override", "is_set", "clear"} +REGISTRY_READ_METHODS = {"get", "is_set"} + +# Violations that existed when the policy was introduced, as +# ": " -> number of occurrences. Lower or delete an entry when +# its violations are fixed; never add one. Kinds are described in +# docs/contributing/env_vars.md. +KNOWN_VIOLATIONS: dict[str, int] = { + 'fastvideo/attention/backends/flash_attn.py: read FASTVIDEO_NVFP4_FA4': 1, + 'fastvideo/attention/backends/video_sparse_attn_h3_probe.py: read FASTVIDEO_H3_VSA_PROBE': 1, + 'fastvideo/attention/layer.py: read FASTVIDEO_DISABLE_ATTENTION_COMPILE': 2, + 'fastvideo/attention/selector.py: read ': 1, + 'fastvideo/attention/utils/flash_attn_default.py: import-time-read FASTVIDEO_FA4': 1, + 'fastvideo/benchmarks/eval_metalfx_rife.py: write FASTVIDEO_MLX_COMPILE': 1, + 'fastvideo/benchmarks/mlx_fastwan_bench.py: read FASTVIDEO_MLX_COMPILE': 1, + 'fastvideo/benchmarks/mlx_fastwan_bench.py: read FASTVIDEO_MLX_FAST_NORM': 1, + 'fastvideo/benchmarks/mlx_fastwan_bench.py: write FASTVIDEO_MLX_COMPILE': 1, + 'fastvideo/distributed/device_communicators/cpu_communicator.py: read VLLM_DIST_IDENT': 1, + 'fastvideo/entrypoints/cli/utils.py: whole-environ': 1, + 'fastvideo/entrypoints/openai/api_server.py: write FASTVIDEO_STAGE_LOGGING': 1, + 'fastvideo/entrypoints/streaming/prompt/providers/cerebras.py: read ': 1, + 'fastvideo/entrypoints/streaming/prompt/providers/groq.py: read ': 1, + 'fastvideo/entrypoints/streaming/worker.py: write CUDA_VISIBLE_DEVICES': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_FFMPEG_BIN': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_BF': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_PRESET': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_QP': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_RC': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_TUNE': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_OUTPUT_PIX_FMT': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_VIDEO_CODEC': 1, + 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_X264_PRESET': 1, + 'fastvideo/entrypoints/video_generator.py: write CUTE_DSL_ENABLE_TVM_FFI': 1, + 'fastvideo/entrypoints/video_generator.py: write FASTVIDEO_NVFP4_FA4': 1, + 'fastvideo/envs.py: prefix CMAKE_BUILD_TYPE': 1, + 'fastvideo/envs.py: prefix CUDA_VISIBLE_DEVICES': 1, + 'fastvideo/envs.py: prefix HCCL_SO_PATH': 1, + 'fastvideo/envs.py: prefix LD_LIBRARY_PATH': 1, + 'fastvideo/envs.py: prefix LOCAL_RANK': 1, + 'fastvideo/envs.py: prefix MAX_JOBS': 1, + 'fastvideo/envs.py: prefix NVCC_THREADS': 1, + 'fastvideo/envs.py: prefix VERBOSE': 1, + 'fastvideo/envs.py: unread CMAKE_BUILD_TYPE': 1, + 'fastvideo/envs.py: unread CUDA_VISIBLE_DEVICES': 1, + 'fastvideo/envs.py: unread FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S': 1, + 'fastvideo/envs.py: unread FASTVIDEO_RINGBUFFER_WARNING_INTERVAL': 1, + 'fastvideo/envs.py: unread FASTVIDEO_SERVER_DEV_MODE': 1, + 'fastvideo/envs.py: unread FASTVIDEO_TARGET_DEVICE': 1, + 'fastvideo/envs.py: unread FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE': 1, + 'fastvideo/envs.py: unread FASTVIDEO_TRACE_FUNCTION': 1, + 'fastvideo/envs.py: unread FASTVIDEO_USE_PRECOMPILED': 1, + 'fastvideo/envs.py: unread LD_LIBRARY_PATH': 1, + 'fastvideo/envs.py: unread MAX_JOBS': 1, + 'fastvideo/envs.py: unread NVCC_THREADS': 1, + 'fastvideo/envs.py: unread VERBOSE': 1, + 'fastvideo/eval/__init__.py: write TORCH_HOME': 1, + 'fastvideo/eval/datasets/physics_iq.py: read FASTVIDEO_PHYSICS_IQ_BUCKET_URL': 1, + 'fastvideo/eval/datasets/vbench.py: read VBENCH_FULL_INFO_JSON': 1, + 'fastvideo/eval/metrics/audio/frechet_distance/metric.py: read ': 1, + 'fastvideo/eval/metrics/common/fvd/metric.py: read ': 1, + 'fastvideo/eval/metrics/judge/third_person_separation/metric.py: read ': 1, + 'fastvideo/eval/metrics/vbench/scene/metric.py: write VIDEO_MAX_PIXELS': 1, + 'fastvideo/eval/models.py: read FASTVIDEO_EVAL_CACHE': 1, + 'fastvideo/logger.py: import-time-read FASTVIDEO_CONFIGURE_LOGGING': 1, + 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_CONFIG_PATH': 1, + 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_LEVEL': 1, + 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_PREFIX': 1, + 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_COMPILE': 1, + 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_DQ_GEMM': 1, + 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_FAST_NORM': 1, + 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_WINDOW': 1, + 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_WINDOW_SINK': 1, + 'fastvideo/mlx_runtime/memory.py: write ': 1, + 'fastvideo/mlx_runtime/wan22.py: read FASTVIDEO_MLX_COMPILE': 1, + 'fastvideo/models/encoders/gemma.py: read LTX2_FASTVIDEO_GEMMA_LOG': 2, + 'fastvideo/models/encoders/reason1.py: read FASTVIDEO_REASON1_WEIGHTS_PATH': 1, + 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write MASTER_ADDR': 1, + 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write MASTER_PORT': 1, + 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write RANK': 1, + 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write WORLD_SIZE': 1, + 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_FSDP2_AUTOWRAP': 1, + 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_FSDP2_MIN_PARAMS': 1, + 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_H3_VSA_PROBE': 1, + 'fastvideo/models/vaes/ltx2vae.py: read ': 1, + 'fastvideo/performance/hf_store.py: read HF_REPO_ID': 1, + 'fastvideo/performance/hf_store.py: read PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS': 1, + 'fastvideo/performance_dashboard/api.py: read PERFORMANCE_TRACKING_ROOT': 1, + 'fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py: read LTX2_DISABLE_AUDIO_AUTOCAST': 1, + 'fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py: read FASTVIDEO_NVTX_PROFILE': 1, + 'fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py: read LTX2_USE_DISTILLED_SIGMAS': 1, + 'fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py: write HF_TOKEN': 1, + 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: read KANDINSKY5_OVERFIT_DATA_DIR': 1, + 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: read KANDINSKY5_OVERFIT_OUTPUT_DIR': 1, + 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: write MASTER_ADDR': 1, + 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: write MASTER_PORT': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_CAPTION_JSON': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_DATA_DIR': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_MODEL': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_NUM_COPIES': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_OUTPUT_DIR': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_VIDEO_SUBDIR': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write LOCAL_RANK': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write MASTER_ADDR': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write MASTER_PORT': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write RANK': 1, + 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write WORLD_SIZE': 1, + 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write LOCAL_RANK': 1, + 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write MASTER_ADDR': 1, + 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write MASTER_PORT': 1, + 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write RANK': 1, + 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write WORLD_SIZE': 1, + 'fastvideo/pipelines/stages/denoising.py: read FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION': 1, + 'fastvideo/pipelines/stages/latent_preparation.py: read FASTVIDEO_COSMOS25_LOG_KNOBS': 1, + 'fastvideo/tests/api/test_attention_selector_resolution.py: write FASTVIDEO_ATTENTION_BACKEND': 17, + 'fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py: write FASTVIDEO_FA4': 3, + 'fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py: write FASTVIDEO_NVFP4_FA4': 4, + 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_KERNEL_VSA_FORCE_TRITON': 1, + 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_VSA_CUTEDSL': 2, + 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_VSA_TRITON': 2, + 'fastvideo/tests/attention/test_vsa_h3_sm100a_route.py: write ': 22, + 'fastvideo/tests/contract/test_profiler_regions.py: whole-environ': 2, + 'fastvideo/tests/contract/test_profiler_regions.py: write FASTVIDEO_NVTX_PROFILE': 4, + 'fastvideo/tests/contract/test_ssim_ci_runner.py: write CUDA_VISIBLE_DEVICES': 1, + 'fastvideo/tests/contract/test_tensor_golden_isolation.py: read FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/tests/contract/test_tensor_golden_isolation.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/contract/test_wan_validation_order.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_sp_hunyuanvideo.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_sp_lingbot_video.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_sp_ltx2.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_sp_wan.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_ulysses_a2a_parity.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: read FASTVIDEO_ULYSSES_FAULT_RANK': 1, + 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: read FASTVIDEO_ULYSSES_FAULT_STAGE': 1, + 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: whole-environ': 1, + 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: write FASTVIDEO_ULYSSES_A2A': 1, + 'fastvideo/tests/encoders/test_clip_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_clip_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_hyt5_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_hyt5_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_llama_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_llama_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_qwen2_5_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_qwen2_5_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_siglip_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_siglip_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/encoders/test_t5_encoder.py: write MASTER_ADDR': 1, + 'fastvideo/tests/encoders/test_t5_encoder.py: write MASTER_PORT': 1, + 'fastvideo/tests/entrypoints/streaming/test_prompt_providers.py: write CEREBRAS_API_KEY': 3, + 'fastvideo/tests/entrypoints/streaming/test_prompt_providers.py: write GROQ_API_KEY': 2, + 'fastvideo/tests/entrypoints/test_openai_api_integration.py: whole-environ': 1, + 'fastvideo/tests/entrypoints/test_openai_video_client.py: whole-environ': 1, + 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_FFMPEG_BIN': 1, + 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_OUTPUT_PIX_FMT': 1, + 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_VIDEO_CODEC': 1, + 'fastvideo/tests/golden_gate/_harness.py: read ': 1, + 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_FA4': 2, + 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_GOLDEN_GATE_DIR': 1, + 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, + 'fastvideo/tests/golden_gate/_harness.py: write CUBLAS_WORKSPACE_CONFIG': 1, + 'fastvideo/tests/golden_gate/_harness.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/golden_gate/_harness.py: write LOCAL_RANK': 1, + 'fastvideo/tests/golden_gate/_harness.py: write MASTER_ADDR': 1, + 'fastvideo/tests/golden_gate/_harness.py: write MASTER_PORT': 1, + 'fastvideo/tests/golden_gate/_harness.py: write RANK': 1, + 'fastvideo/tests/golden_gate/_harness.py: write WORLD_SIZE': 1, + 'fastvideo/tests/golden_gate/_tensor_golden.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/golden_gate/_tensor_golden.py: write FASTVIDEO_ATTENTION_BACKEND': 3, + 'fastvideo/tests/golden_gate/test_kandinsky5.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/golden_gate/test_kandinsky5.py: write FASTVIDEO_FA4': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_GATE_GOLDEN_DIR': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_GATE_LAYER': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_MODEL_ROOT': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write CUBLAS_WORKSPACE_CONFIG': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write LOCAL_RANK': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write MASTER_ADDR': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write MASTER_PORT': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write RANK': 1, + 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write WORLD_SIZE': 1, + 'fastvideo/tests/golden_gate/test_wan_denoising.py: write FASTVIDEO_CFG_GATE_STEP': 1, + 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_ACTIVATIONS': 6, + 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_LAYERS': 5, + 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_OUTPUT': 5, + 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_STATS': 1, + 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_STEPS': 1, + 'fastvideo/tests/inference/bsa/test_bsa_inference.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write MASTER_ADDR': 1, + 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write MASTER_PORT': 1, + 'fastvideo/tests/inference/test_basic_fasth3_profile.py: write ': 2, + 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_DISABLE_ATTENTION_COMPILE': 9, + 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_H3_VSA_PROBE': 4, + 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_VSA_SM100A': 3, + 'fastvideo/tests/inference/vmoba/test_vmoba_inference.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py: whole-environ': 1, + 'fastvideo/tests/mlx/test_memory_limits.py: read PYTORCH_MPS_HIGH_WATERMARK_RATIO': 1, + 'fastvideo/tests/mlx/test_memory_limits.py: read PYTORCH_MPS_LOW_WATERMARK_RATIO': 1, + 'fastvideo/tests/mlx/test_memory_limits.py: write PYTORCH_MPS_HIGH_WATERMARK_RATIO': 1, + 'fastvideo/tests/mlx/test_memory_limits.py: write PYTORCH_MPS_LOW_WATERMARK_RATIO': 1, + 'fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py: read FASTVIDEO_MLX_DQ_GEMM': 1, + 'fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py: write FASTVIDEO_MLX_DQ_GEMM': 11, + 'fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py: write MASTER_ADDR': 1, + 'fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py: write MASTER_PORT': 1, + 'fastvideo/tests/mlx/test_mlx_taeh3.py: read TAEH3_REFERENCE_DIR': 1, + 'fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py: write HOME': 1, + 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: read FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY': 1, + 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: read FASTVIDEO_WAN22_5B_ROOT': 1, + 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: write MASTER_ADDR': 1, + 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: write MASTER_PORT': 1, + 'fastvideo/tests/mlx/tiny_h3.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/mlx/tiny_h3.py: write MASTER_ADDR': 1, + 'fastvideo/tests/mlx/tiny_h3.py: write MASTER_PORT': 1, + 'fastvideo/tests/mlx/tiny_wan.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/mlx/tiny_wan.py: write MASTER_ADDR': 1, + 'fastvideo/tests/mlx/tiny_wan.py: write MASTER_PORT': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read ': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read CFLAGS': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read CMAKE_ARGS': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read CXXFLAGS': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read FASTVIDEO_KERNEL_CACHE_ROOT': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read FASTVIDEO_KERNEL_PREBUILT_INFO': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read GPU_BACKEND': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: read LDFLAGS': 1, + 'fastvideo/tests/modal/kernel_build_cache.py: whole-environ': 1, + 'fastvideo/tests/modal/kernel_cache_smoke.py: read BUILDKITE_COMMIT': 1, + 'fastvideo/tests/modal/kernel_cache_smoke.py: read BUILDKITE_REPO': 1, + 'fastvideo/tests/modal/kernel_cache_smoke.py: read FASTVIDEO_MODAL_IMAGE': 1, + 'fastvideo/tests/modal/kernel_cache_smoke.py: read IMAGE_VERSION': 1, + 'fastvideo/tests/modal/kernel_cache_smoke.py: read UV_TORCH_BACKEND': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read ': 2, + 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_COMMIT': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_PULL_REQUEST': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_REPO': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_MODAL_IMAGE': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_MODAL_VOLUME': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 2, + 'fastvideo/tests/modal/launch_l40s_job.py: read IMAGE_VERSION': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: read UV_TORCH_BACKEND': 1, + 'fastvideo/tests/modal/launch_l40s_job.py: whole-environ': 1, + 'fastvideo/tests/modal/modal_image_utils.py: read UV_TORCH_BACKEND': 1, + 'fastvideo/tests/modal/pr_test.py: read ': 2, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BRANCH': 1, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BUILD_ID': 1, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BUILD_URL': 1, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_COMMIT': 2, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_JOB_ID': 1, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_PULL_REQUEST': 2, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_REPO': 2, + 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_SOURCE': 1, + 'fastvideo/tests/modal/pr_test.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/modal/pr_test.py: read HF_API_KEY': 1, + 'fastvideo/tests/modal/pr_test.py: read IMAGE_VERSION': 1, + 'fastvideo/tests/modal/pr_test.py: read TEST_SCOPE': 1, + 'fastvideo/tests/modal/pr_test.py: read UV_TORCH_BACKEND': 1, + 'fastvideo/tests/modal/pr_test.py: read WANDB_API_KEY': 1, + 'fastvideo/tests/modal/ssim_test.py: read ': 1, + 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_COMMIT': 1, + 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_PULL_REQUEST': 1, + 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_REPO': 1, + 'fastvideo/tests/modal/ssim_test.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/modal/ssim_test.py: read IMAGE_VERSION': 1, + 'fastvideo/tests/modal/ssim_test.py: read UV_TORCH_BACKEND': 1, + 'fastvideo/tests/modal/ssim_test.py: whole-environ': 2, + 'fastvideo/tests/modal/test_kernel_build_cache.py: write ': 2, + 'fastvideo/tests/modal/test_kernel_build_cache.py: write TORCH_CUDA_ARCH_LIST': 2, + 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_COMMIT': 2, + 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_PULL_REQUEST': 2, + 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_REPO': 2, + 'fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py: write WANDB_MODE': 1, + 'fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py: write WANDB_MODE': 1, + 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read ': 1, + 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read KANDINSKY5_E2E_NUM_GPUS': 1, + 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: whole-environ': 3, + 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: write WANDB_MODE': 1, + 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: read FASTVIDEO_NIGHTLY': 1, + 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: whole-environ': 2, + 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: write FASTVIDEO_NIGHTLY': 1, + 'fastvideo/tests/nightly/test_e2e_overfit_single_sample.py: write WANDB_MODE': 1, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BRANCH': 2, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BUILD_ID': 1, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BUILD_URL': 1, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_COMMIT': 3, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_JOB_ID': 1, + 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_PULL_REQUEST': 2, + 'fastvideo/tests/performance/compare_baseline.py: read GITHUB_STEP_SUMMARY': 1, + 'fastvideo/tests/performance/compare_baseline.py: read PERFORMANCE_TRACKING_ROOT': 1, + 'fastvideo/tests/performance/compare_baseline.py: read PERF_PYTEST_RC': 4, + 'fastvideo/tests/performance/compare_baseline.py: read PERF_REPORTS_DIR': 1, + 'fastvideo/tests/performance/compare_baseline.py: read PERF_RUN_SOURCE': 1, + 'fastvideo/tests/performance/compare_baseline.py: read PERF_UPLOAD_POLICY': 1, + 'fastvideo/tests/performance/compare_baseline.py: read TEST_SCOPE': 2, + 'fastvideo/tests/performance/dashboard.py: read BUILDKITE_COMMIT': 1, + 'fastvideo/tests/performance/dashboard.py: read DASHBOARD_DAYS': 1, + 'fastvideo/tests/performance/dashboard.py: read PERFORMANCE_TRACKING_ROOT': 1, + 'fastvideo/tests/performance/dashboard.py: read PERF_REPORTS_DIR': 1, + 'fastvideo/tests/performance/identity.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/performance/identity.py: whole-environ': 1, + 'fastvideo/tests/performance/seed_baseline.py: read PERFORMANCE_RESEED_STAGING_ROOT': 1, + 'fastvideo/tests/performance/seed_baseline.py: read PERFORMANCE_TRACKING_ROOT': 1, + 'fastvideo/tests/performance/seed_baseline.py: read USER': 1, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BRANCH': 3, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BUILD_ID': 3, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BUILD_URL': 2, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_JOB_ID': 3, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_PULL_REQUEST': 3, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write PERF_PYTEST_RC': 8, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write PERF_RUN_SOURCE': 23, + 'fastvideo/tests/performance/test_compare_baseline_policy.py: write TEST_SCOPE': 3, + 'fastvideo/tests/performance/test_dashboard_service.py: write ': 1, + 'fastvideo/tests/performance/test_dashboard_service.py: write HF_TOKEN': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BRANCH': 2, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BUILD_ID': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BUILD_URL': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_COMMIT': 2, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_JOB_ID': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_PULL_REQUEST': 2, + 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_STAGE_LOGGING': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read IMAGE_VERSION': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read PERF_RUN_SOURCE': 1, + 'fastvideo/tests/performance/test_inference_performance.py: read TEST_SCOPE': 2, + 'fastvideo/tests/performance/test_inference_performance.py: write FASTVIDEO_STAGE_LOGGING': 3, + 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_ATTENTION_BACKEND': 5, + 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_CONTAINER_IMAGE_REF': 2, + 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_FA4': 2, + 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 1, + 'fastvideo/tests/performance/test_inference_performance_identity.py: write IMAGE_VERSION': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BRANCH': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BUILD_ID': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BUILD_URL': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_COMMIT': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_JOB_ID': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_PULL_REQUEST': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write PERF_RUN_SOURCE': 1, + 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write TEST_SCOPE': 1, + 'fastvideo/tests/ssim/bootstrap_references.py: read ': 3, + 'fastvideo/tests/ssim/ci_runner.py: whole-environ': 1, + 'fastvideo/tests/ssim/conftest.py: read ': 1, + 'fastvideo/tests/ssim/conftest.py: read FASTVIDEO_SSIM_MODEL_ID': 1, + 'fastvideo/tests/ssim/conftest.py: read FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD': 1, + 'fastvideo/tests/ssim/conftest.py: write ': 3, + 'fastvideo/tests/ssim/inference_similarity_utils.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/ssim/inference_similarity_utils.py: write FASTVIDEO_ATTENTION_BACKEND': 3, + 'fastvideo/tests/ssim/reference_utils.py: read ': 1, + 'fastvideo/tests/ssim/reference_videos_cli.py: read ': 3, + 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_AR_SSIM_MODEL_PATH': 1, + 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_SSIM_MODEL_PATH': 1, + 'fastvideo/tests/ssim/test_flux_t2i_similarity.py: read FLUX_T2I_MODEL_DIR': 1, + 'fastvideo/tests/ssim/test_gamecraft_similarity.py: read GAMECRAFT_MODEL_PATH': 1, + 'fastvideo/tests/ssim/test_gamecraft_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/tests/ssim/test_gen3c_similarity.py: read GEN3C_MODEL_PATH': 1, + 'fastvideo/tests/ssim/test_gen3c_similarity.py: read GEN3C_TEST_IMAGE_PATH': 1, + 'fastvideo/tests/ssim/test_gen3c_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/ssim/test_glm_image_similarity.py: read GLM_IMAGE_LOCAL_WEIGHTS_DIR': 1, + 'fastvideo/tests/ssim/test_glm_image_similarity.py: read GLM_IMAGE_MODEL_DIR': 1, + 'fastvideo/tests/ssim/test_kandinsky5_similarity.py: read FASTVIDEO_FA4': 1, + 'fastvideo/tests/ssim/test_kandinsky5_similarity.py: write FASTVIDEO_FA4': 3, + 'fastvideo/tests/ssim/test_lingbot_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/ssim/test_longcat_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 3, + 'fastvideo/tests/ssim/test_matrixgame2_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/ssim/test_matrixgame3_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/ssim/test_sd35_similarity.py: read SD35_MODEL_DIR': 1, + 'fastvideo/tests/ssim/test_zimage_similarity.py: read ZIMAGE_MODEL_DIR': 1, + 'fastvideo/tests/ssim/test_zimage_similarity.py: read ZIMAGE_MODEL_REVISION': 1, + 'fastvideo/tests/stages/_denoising_fixtures.py: write FASTVIDEO_CFG_GATE_STEP': 2, + 'fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/stages/test_minimax_h3_encoding_offload.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/train/methods/grad_norm_regression.py: read ': 1, + 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_ltx2_finetune.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_ltx2_finetune.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_matrixgame2_finetune.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_matrixgame2_finetune.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_minimax_h3_finetune.py: whole-environ': 2, + 'fastvideo/tests/train/methods/test_streaming_long_tuning.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_streaming_long_tuning.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_wan_causal_cd.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_wan_causal_cd.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_wan_causal_dfsft.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_wan_causal_dfsft.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_wan_causal_tfsft.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_wan_causal_tfsft.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/methods/test_wan_finetune.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/methods/test_wan_finetune.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_cosmos.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_cosmos.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_hunyuan.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_hunyuan.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_kandinsky5.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/train/models/test_load_kandinsky5.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_kandinsky5.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_longcat.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_longcat.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_ltx2.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_ltx2.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_matrixgame2.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_matrixgame2.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_wan.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_wan.py: write MASTER_PORT': 1, + 'fastvideo/tests/train/models/test_load_wan_causal.py: write MASTER_ADDR': 1, + 'fastvideo/tests/train/models/test_load_wan_causal.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/VSA/test_training_loss_VSA.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/training/VSA/test_training_loss_VSA.py: write WANDB_MODE': 1, + 'fastvideo/tests/training/Vanilla/mfu_calculation.py: read PYTHONPATH': 1, + 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write MASTER_ADDR': 1, + 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write PYTHONPATH': 1, + 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write WANDB_MODE': 1, + 'fastvideo/tests/training/Vanilla/test_training_loss.py: write MASTER_ADDR': 1, + 'fastvideo/tests/training/Vanilla/test_training_loss.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/Vanilla/test_training_loss.py: write WANDB_MODE': 1, + 'fastvideo/tests/training/distill/test_anyflow_smoke.py: whole-environ': 2, + 'fastvideo/tests/training/distill/test_distill_dmd.py: write MASTER_ADDR': 1, + 'fastvideo/tests/training/distill/test_distill_dmd.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/distill/test_distill_dmd.py: write WANDB_MODE': 1, + 'fastvideo/tests/training/lora/test_lora_training.py: write MASTER_ADDR': 1, + 'fastvideo/tests/training/lora/test_lora_training.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/lora/test_lora_training.py: write WANDB_MODE': 1, + 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write MASTER_ADDR': 1, + 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write MASTER_PORT': 1, + 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write WANDB_MODE': 1, + 'fastvideo/tests/transformers/test_cosmos.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_cosmos.py: write MASTER_PORT': 1, + 'fastvideo/tests/transformers/test_cosmos2_5.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_cosmos2_5.py: write MASTER_PORT': 1, + 'fastvideo/tests/transformers/test_flux.py: read FLUX_TRANSFORMER_PATH': 1, + 'fastvideo/tests/transformers/test_flux.py: write FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/transformers/test_flux.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_flux.py: write MASTER_PORT': 1, + 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write DISABLE_SP': 1, + 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write MASTER_PORT': 1, + 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write TORCHDYNAMO_DISABLE': 1, + 'fastvideo/tests/transformers/test_hyworld.py: read FASTVIDEO_ATTENTION_BACKEND': 1, + 'fastvideo/tests/transformers/test_hyworld.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_hyworld.py: write MASTER_PORT': 1, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write FASTVIDEO_FA4': 1, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write LOCAL_RANK': 2, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write MASTER_ADDR': 2, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write MASTER_PORT': 2, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write RANK': 2, + 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write WORLD_SIZE': 2, + 'fastvideo/tests/transformers/test_wanvideo.py: write MASTER_ADDR': 1, + 'fastvideo/tests/transformers/test_wanvideo.py: write MASTER_PORT': 1, + 'fastvideo/tests/vaes/test_hunyuan15_vae.py: write MASTER_ADDR': 1, + 'fastvideo/tests/vaes/test_hunyuan15_vae.py: write MASTER_PORT': 1, + 'fastvideo/tests/vaes/test_wan_vae.py: write MASTER_ADDR': 1, + 'fastvideo/tests/vaes/test_wan_vae.py: write MASTER_PORT': 1, + 'fastvideo/tests/worker/test_gpu_worker.py: write FASTVIDEO_NVTX_PROFILE': 2, + 'fastvideo/tests/worker/test_gpu_worker.py: write LOCAL_RANK': 1, + 'fastvideo/train/entrypoint/dcp_to_diffusers.py: write ': 1, + 'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 2, + 'fastvideo/training/self_forcing_distillation_pipeline.py: read FASTVIDEO_FSDP2_AUTOWRAP': 1, + 'fastvideo/utils.py: read ': 5, + 'fastvideo/utils.py: read FASTVIDEO_WORKER_MULTIPROC_METHOD': 1, + 'fastvideo/utils.py: write ': 1, + 'fastvideo/utils.py: write FASTVIDEO_WORKER_MULTIPROC_METHOD': 1, + 'fastvideo/worker/gpu_worker.py: write LOCAL_RANK': 1, + 'fastvideo/worker/gpu_worker.py: write NCCL_ASYNC_ERROR_HANDLING': 1, + 'fastvideo/worker/gpu_worker.py: write RANK': 1, + 'fastvideo/worker/gpu_worker.py: write TORCH_NCCL_AVOID_RECORD_STREAMS': 1, + 'fastvideo/worker/gpu_worker.py: write WORLD_SIZE': 1, + 'fastvideo/worker/ray_distributed_executor.py: read ': 2, + 'fastvideo/worker/ray_distributed_executor.py: read RAY_USAGE_STATS_ENABLED': 1, + 'fastvideo/worker/ray_distributed_executor.py: whole-environ': 1, + 'fastvideo/worker/ray_distributed_executor.py: write RAY_USAGE_STATS_ENABLED': 1, + 'fastvideo/worker/ray_env.py: import-time-read FASTVIDEO_CONFIG_ROOT': 1, + 'fastvideo/worker/ray_env.py: read ': 1, + 'fastvideo/worker/ray_utils.py: whole-environ': 1, + 'fastvideo/worker/worker_base.py: read ': 1, + 'fastvideo/worker/worker_base.py: write ': 1, +} + + +def load_registry() -> ModuleType: + """Load fastvideo/envs.py as a standalone module, without importing the fastvideo package.""" + spec = importlib.util.spec_from_file_location("_fastvideo_envs_registry", REGISTRY_PATH) + assert spec is not None and spec.loader is not None + registry = importlib.util.module_from_spec(spec) + spec.loader.exec_module(registry) + return registry + + +def is_allowlisted(name: str) -> bool: + for pattern in EXTERNAL_ALLOWLIST: + if pattern.endswith("*") and name.startswith(pattern[:-1]): + return True + if name == pattern: + return True + return False + + +def _literal(node: ast.AST | None) -> str: + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + return "" + + +class EnvAccessScanner: + """Find environment accesses and registry-field uses in one module. + + ``violations`` holds (kind, name, line) tuples; ``registry_reads`` counts + ``envs.X.get()`` and ``envs.X.is_set()`` calls by variable name. + """ + + def __init__(self, tree: ast.Module, registry_names: set[str]) -> None: + self.registry_names = registry_names + self.violations: list[tuple[str, str, int]] = [] + self.registry_reads: Counter[str] = Counter() + self.parents: dict[ast.AST, ast.AST] = {} + self._collect_aliases(tree) + self._visit(tree, in_function=False) + + def _collect_aliases(self, tree: ast.Module) -> None: + """Record the local names bound to os, os.environ, os.getenv-like functions, and fastvideo.envs.""" + self.os_names: set[str] = set() + self.environ_names: set[str] = set() + self.getenv_names: set[str] = set() + self.setenv_names: set[str] = set() + self.envs_module_names: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "os" or alias.name.startswith("os."): + self.os_names.add(alias.asname or "os") + if alias.name == "fastvideo.envs" and alias.asname: + self.envs_module_names.add(alias.asname) + elif isinstance(node, ast.ImportFrom): + for alias in node.names: + local = alias.asname or alias.name + if node.module == "os": + if alias.name in ("environ", "environb"): + self.environ_names.add(local) + elif alias.name in ("getenv", "getenvb"): + self.getenv_names.add(local) + elif alias.name in ("putenv", "unsetenv"): + self.setenv_names.add(local) + elif node.module == "fastvideo" and alias.name == "envs": + self.envs_module_names.add(local) + + def _is_os_attr(self, node: ast.AST, attrs: tuple[str, ...]) -> bool: + return (isinstance(node, ast.Attribute) and node.attr in attrs and isinstance(node.value, ast.Name) + and node.value.id in self.os_names) + + def _is_environ(self, node: ast.AST) -> bool: + return self._is_os_attr(node, ("environ", "environb")) or (isinstance(node, ast.Name) + and node.id in self.environ_names) + + def _is_envs_module(self, node: ast.AST) -> bool: + return ((isinstance(node, ast.Name) and node.id in self.envs_module_names) + or (isinstance(node, ast.Attribute) and node.attr == "envs")) + + def _visit(self, node: ast.AST, in_function: bool) -> None: + """Walk the tree, tracking whether each node runs inside a function body.""" + self._check(node, in_function) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + # Decorators, defaults, and annotations run when the function is defined. + outside = [*getattr(node, "decorator_list", []), *node.args.defaults, *node.args.kw_defaults] + for child in outside: + if child is not None: + self.parents[child] = node + self._visit(child, in_function) + body = node.body if isinstance(node.body, list) else [node.body] + for child in body: + self.parents[child] = node + self._visit(child, True) + return + for child in ast.iter_child_nodes(node): + self.parents[child] = node + self._visit(child, in_function) + + def _check(self, node: ast.AST, in_function: bool) -> None: + parent = self.parents.get(node) + if self._is_environ(node): + self._check_environ_use(node, parent) + elif isinstance(node, ast.Call): + self._check_call(node) + elif isinstance(node, ast.Attribute) and node.attr in self.registry_names and self._is_envs_module(node.value): + self._check_registry_use(node, parent, in_function) + + def _check_environ_use(self, node: ast.AST, parent: ast.AST | None) -> None: + line = node.lineno # type: ignore[attr-defined] + if isinstance(parent, ast.Attribute) and isinstance(self.parents.get(parent), ast.Call): + call = self.parents[parent] + assert isinstance(call, ast.Call) + if call.func is parent and parent.attr == "get": + self.violations.append(("read", _literal(call.args[0] if call.args else None), line)) + return + if call.func is parent and parent.attr in ("setdefault", "pop"): + self.violations.append(("write", _literal(call.args[0] if call.args else None), line)) + return + if isinstance(parent, ast.Subscript) and parent.value is node: + kind = "read" if isinstance(parent.ctx, ast.Load) else "write" + self.violations.append((kind, _literal(parent.slice), line)) + return + if (isinstance(parent, ast.Compare) and len(parent.ops) == 1 and isinstance(parent.ops[0], (ast.In, ast.NotIn)) + and parent.comparators[0] is node): + self.violations.append(("read", _literal(parent.left), line)) + return + self.violations.append(("whole-environ", "", line)) + + def _check_call(self, node: ast.Call) -> None: + func = node.func + name = _literal(node.args[0] if node.args else None) + if self._is_os_attr(func, ("getenv", "getenvb")) or (isinstance(func, ast.Name) + and func.id in self.getenv_names): + self.violations.append(("read", name, node.lineno)) + elif (self._is_os_attr(func, ("putenv", "unsetenv")) + or (isinstance(func, ast.Name) and func.id in self.setenv_names) + or (isinstance(func, ast.Attribute) and func.attr in ("setenv", "delenv"))): + self.violations.append(("write", name, node.lineno)) + + def _check_registry_use(self, node: ast.Attribute, parent: ast.AST | None, in_function: bool) -> None: + # Only a call counts: ``getter = envs.X.get`` neither reads the variable + # nor uses the field through a method. + grandparent = self.parents.get(parent) if parent is not None else None + if not (isinstance(parent, ast.Attribute) and parent.attr in REGISTRY_METHODS + and isinstance(grandparent, ast.Call) and grandparent.func is parent): + self.violations.append(("bare-field", node.attr, node.lineno)) + return + if parent.attr in REGISTRY_READ_METHODS: + self.registry_reads[node.attr] += 1 + if not in_function: + self.violations.append(("import-time-read", node.attr, node.lineno)) + + +def scanned_files() -> list[Path]: + return sorted(path for path in PACKAGE_ROOT.rglob("*.py") + if path != REGISTRY_PATH and not any(excluded in path.parents for excluded in EXCLUDED_DIRS)) + + +def collect_violations(registry: ModuleType) -> tuple[dict[str, list[int]], Counter[str]]: + """Return every violation key with its line numbers, plus registry read counts.""" + registry_names = set(registry.environment_variables) + found: dict[str, list[int]] = {} + reads: Counter[str] = Counter() + for path in scanned_files(): + relative = path.relative_to(REPO_ROOT).as_posix() + scanner = EnvAccessScanner(ast.parse(path.read_text(encoding="utf-8"), filename=relative), registry_names) + reads.update(scanner.registry_reads) + for kind, name, line in scanner.violations: + if kind == "read" and is_allowlisted(name): + continue + key = f"{relative}: {kind} {name}".rstrip() + found.setdefault(key, []).append(line) + + registry_path = REGISTRY_PATH.relative_to(REPO_ROOT).as_posix() + for name in registry.environment_variables: + if not re.fullmatch(r"FASTVIDEO_[A-Z0-9_]+", name): + found.setdefault(f"{registry_path}: prefix {name}", []).append(0) + if reads[name] == 0: + found.setdefault(f"{registry_path}: unread {name}", []).append(0) + return found, reads + + +def test_env_access_follows_policy(): + found, _ = collect_violations(load_registry()) + counts = Counter({key: len(lines) for key, lines in found.items()}) + known = Counter(KNOWN_VIOLATIONS) + + new = counts - known + fixed = known - counts + messages = [] + if new: + lines = [f" {key} (lines {found[key]})" for key in sorted(new)] + messages.append(f"New environment-variable policy violations. See {POLICY_DOC} for the rule and the fix:\n" + + "\n".join(lines)) + if fixed: + lines = [f" {key}" for key in sorted(fixed)] + messages.append("These KNOWN_VIOLATIONS entries are fixed; delete them from " + "fastvideo/tests/contract/test_env_policy.py:\n" + "\n".join(lines)) + assert not messages, "\n\n".join(messages) + + +def test_registry_method_counts_only_when_called(): + """An uncalled ``envs.NAME.get`` is a bare field, not a read.""" + source = ("import fastvideo.envs as envs\n" + "def f():\n" + " getter = envs.FASTVIDEO_FA4.get\n" + " return envs.FASTVIDEO_FA4.get()\n") + scanner = EnvAccessScanner(ast.parse(source), {"FASTVIDEO_FA4"}) + assert scanner.violations == [("bare-field", "FASTVIDEO_FA4", 3)] + assert scanner.registry_reads["FASTVIDEO_FA4"] == 1 + + +def test_registry_entries_have_category_and_description(): + registry = load_registry() + problems = [] + for name, field in registry.environment_variables.items(): + if field.category not in registry.CATEGORIES: + problems.append(f"{name}: category {field.category!r} is not in envs.CATEGORIES") + if not field.doc.strip(): + problems.append(f"{name}: empty description") + assert not problems, f"Registry entries break {POLICY_DOC}:\n" + "\n".join(problems) + + +def _render_default(field) -> str: + if callable(field.default): + return "computed" + if field.default is None: + return "unset" + return f"`{field.format(field.default)}`" if field.format(field.default) else '`""`' + + +def _escape(text: str) -> str: + return text.replace("|", "\\|").replace("*", "\\*").replace("<", "<").replace(">", ">") + + +def render_env_table(registry: ModuleType) -> str: + """Render the registry as an aligned Markdown table, in declaration order.""" + rows = [["Variable", "Type", "Default", "Category", "Description"]] + for name, field in registry.environment_variables.items(): + rows.append([f"`{name}`", field.type_name, _render_default(field), field.category, _escape(field.doc)]) + widths = [max(len(row[column]) for row in rows) for column in range(len(rows[0]))] + + def render_row(cells: list[str]) -> str: + return "| " + " | ".join(cell.ljust(width) for cell, width in zip(cells, widths)) + " |" + + lines = [render_row(rows[0]), render_row(["-" * width for width in widths])] + lines.extend(render_row(row) for row in rows[1:]) + return "\n".join(lines) + + +def _split_doc(text: str) -> tuple[str, str, str]: + """Split the policy doc into the text before the generated table, the table, and the text after it.""" + before, found_begin, rest = text.partition(DOC_TABLE_BEGIN) + table, found_end, after = rest.partition(DOC_TABLE_END) + assert found_begin and found_end, f"{POLICY_DOC} must contain the generated-table markers" + return before + found_begin + "\n", table.strip("\n"), "\n" + found_end + after + + +def test_env_doc_table_matches_registry(): + _, table, _ = _split_doc(DOC_PATH.read_text(encoding="utf-8")) + assert table == render_env_table(load_registry()), ( + f"The table in {POLICY_DOC} is out of date. Run `python fastvideo/tests/contract/test_env_policy.py`.") + + +if __name__ == "__main__": + head, _, tail = _split_doc(DOC_PATH.read_text(encoding="utf-8")) + DOC_PATH.write_text(head + render_env_table(load_registry()) + tail, encoding="utf-8") + print(f"Updated {POLICY_DOC}") diff --git a/fastvideo/tests/stages/_denoising_fixtures.py b/fastvideo/tests/stages/_denoising_fixtures.py index 2cc589c99..6ad590248 100644 --- a/fastvideo/tests/stages/_denoising_fixtures.py +++ b/fastvideo/tests/stages/_denoising_fixtures.py @@ -113,9 +113,9 @@ def _patch_denoising_module(monkeypatch, cfg_gate_step): import fastvideo.pipelines.stages.denoising as denoising from fastvideo.pipelines.stages.base import PipelineStage - # envs.py evaluates FASTVIDEO_CFG_GATE_STEP lazily via __getattr__, so the - # stage sees monkeypatched values without reloading the module. - assert denoising.envs.FASTVIDEO_CFG_GATE_STEP == expected_gate_step + # envs.FASTVIDEO_CFG_GATE_STEP.get() reads the environment on each call, so + # the stage sees monkeypatched values without reloading the module. + assert denoising.envs.FASTVIDEO_CFG_GATE_STEP.get() == expected_gate_step logger = RecordingLogger() monkeypatch.setattr(denoising, "logger", logger) diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index 5168313c5..92faf4b02 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -274,7 +274,7 @@ class TrainingMethod(torch.nn.Module, ABC): def _infer_attn_kind(self) -> Literal["dense", "vsa"]: """Derive metadata mode from the student's resolved backend.""" - backend = (self.student.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND) + backend = (self.student.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND.get()) if backend == "VIDEO_SPARSE_ATTN": return "vsa" return "dense" diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py index 9f891019e..c36fe200b 100644 --- a/fastvideo/train/models/ltx2/ltx2.py +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -525,7 +525,7 @@ class LTX2Model(WanModel): return training_batch def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch: - if envs.FASTVIDEO_ATTENTION_BACKEND in ("VIDEO_SPARSE_ATTN", "VMOBA_ATTN"): + if envs.FASTVIDEO_ATTENTION_BACKEND.get() in ("VIDEO_SPARSE_ATTN", "VMOBA_ATTN"): raise NotImplementedError("LTX2Model does not support VSA/VMOBA attention backends") training_batch.attn_metadata = None return training_batch diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 7737639b3..aaf74e98b 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -490,7 +490,7 @@ class WanModel(ModelBase): assert latents_shape is not None assert training_batch.timesteps is not None - attention_backend = (self.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND) + attention_backend = (self.attention_backend_name or envs.FASTVIDEO_ATTENTION_BACKEND.get()) if attention_backend == "VIDEO_SPARSE_ATTN": if (not is_vsa_available() or VideoSparseAttentionMetadataBuilder is None): raise ImportError("FASTVIDEO_ATTENTION_BACKEND is " diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 2f9e86009..524bffe79 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -1324,7 +1324,7 @@ class DistillationPipeline(TrainingPipeline): disable=self.local_rank > 0, ) - use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" + use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND.get() == "VIDEO_SPARSE_ATTN" for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): if step % 5 == 0: gc.collect() @@ -1506,7 +1506,7 @@ class DistillationPipeline(TrainingPipeline): if self.training_args.use_ema and self.is_ema_ready(): self.save_ema_weights(self.training_args.output_dir, self.training_args.max_train_steps) - if envs.FASTVIDEO_TORCH_PROFILER_DIR: + if envs.FASTVIDEO_TORCH_PROFILER_DIR.get(): logger.info("Stopping profiler...") self.profiler_controller.stop() logger.info("Profiler stopped.") diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index 9e6b0b26f..bbc5a44d6 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -833,7 +833,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): disable=self.local_rank > 0, ) - use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" + use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND.get() == "VIDEO_SPARSE_ATTN" for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): start_time = time.perf_counter() if use_vsa: @@ -1008,7 +1008,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): if self.training_args.use_ema and self.is_ema_ready(): self.save_ema_weights(self.training_args.output_dir, self.training_args.max_train_steps) - if envs.FASTVIDEO_TORCH_PROFILER_DIR: + if envs.FASTVIDEO_TORCH_PROFILER_DIR.get(): logger.info("Stopping profiler...") self.profiler_controller.stop() logger.info("Profiler stopped.") diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 97c66ecf7..f5a3939ad 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -348,7 +348,7 @@ class TrainingPipeline(LoRAPipeline, ABC): current_vsa_sparsity = training_batch.current_vsa_sparsity assert latents_shape is not None assert training_batch.timesteps is not None - if envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": + if envs.FASTVIDEO_ATTENTION_BACKEND.get() == "VIDEO_SPARSE_ATTN": if not vsa_available: raise ImportError("FASTVIDEO_ATTENTION_BACKEND is set to VIDEO_SPARSE_ATTN, " "but fastvideo_kernel is not correctly installed or detected. " @@ -361,7 +361,7 @@ class TrainingPipeline(LoRAPipeline, ABC): VSA_sparsity=current_vsa_sparsity, device=get_local_torch_device(), cache_tile_buf=self.training_args.VSA_cache_tile_buf) - elif envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + elif envs.FASTVIDEO_ATTENTION_BACKEND.get() == "VMOBA_ATTN": if not vmoba_available: raise ImportError("FASTVIDEO_ATTENTION_BACKEND is set to VMOBA_ATTN, " "but fastvideo_kernel (or flash_attn>=2.7.4) is not correctly installed.") @@ -389,7 +389,8 @@ class TrainingPipeline(LoRAPipeline, ABC): return training_batch def _transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -> TrainingBatch: - if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN": + if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND.get( + ) == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND.get() == "VMOBA_ATTN": assert training_batch.attn_metadata is not None else: assert training_batch.attn_metadata is None @@ -642,7 +643,7 @@ class TrainingPipeline(LoRAPipeline, ABC): self.training_args.max_train_steps, self.optimizer, self.train_dataloader, self.lr_scheduler, self.noise_random_generator) - if envs.FASTVIDEO_TORCH_PROFILER_DIR: + if envs.FASTVIDEO_TORCH_PROFILER_DIR.get(): logger.info("Stopping profiler...") self.profiler_controller.stop() logger.info("Profiler stopped.") diff --git a/fastvideo/utils.py b/fastvideo/utils.py index fbf3db6c1..441ac3fcd 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -64,7 +64,7 @@ def find_nccl_library() -> str: After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be found by `ctypes` automatically. """ - so_file = envs.FASTVIDEO_NCCL_SO_PATH + so_file = envs.FASTVIDEO_NCCL_SO_PATH.get() # manually load the nccl library if so_file: @@ -87,7 +87,7 @@ def find_hccl_library() -> str: After importing `torch`, `libhccl.so` can be found by `ctypes` automatically. """ - so_file = envs.HCCL_SO_PATH + so_file = envs.HCCL_SO_PATH.get() # manually load the nccl library if so_file: @@ -1127,7 +1127,7 @@ def _append_to_memory_trace(message: str, log_file_path: str | os.PathLike[str] # TODO(xingyu): add adopted message for this def get_ip() -> str: - host_ip = envs.FASTVIDEO_HOST_IP + host_ip = envs.FASTVIDEO_HOST_IP.get() if host_ip: return host_ip @@ -1169,7 +1169,7 @@ def test_loopback_bind(address: str, family: socket.AddressFamily) -> bool: def get_loopback_ip() -> str: - loopback_ip = envs.FASTVIDEO_LOOPBACK_IP + loopback_ip = envs.FASTVIDEO_LOOPBACK_IP.get() if loopback_ip: return loopback_ip @@ -1253,7 +1253,7 @@ def get_mp_context() -> BaseContext: FASTVIDEO_WORKER_MULTIPROC_METHOD. """ force_spawn() - mp_method = envs.FASTVIDEO_WORKER_MULTIPROC_METHOD + mp_method = envs.FASTVIDEO_WORKER_MULTIPROC_METHOD.get() return multiprocessing.get_context(mp_method) diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 3e6d18250..8a22cdf73 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -16,7 +16,7 @@ logger = init_logger(__name__) def _log_cuda_device_uuid(rank: int, device: torch.device) -> None: """Record an NVIDIA worker UUID when external NVTX profiling is enabled.""" - if not envs.FASTVIDEO_NVTX_PROFILE: + if not envs.FASTVIDEO_NVTX_PROFILE.get(): return device_uuid = torch.cuda.get_device_properties(device).uuid logger.info("Worker %d CUDA device UUID: GPU-%s", rank, device_uuid, local_main_process_only=False) diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 98714b3ba..5c4d31875 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -150,7 +150,7 @@ class MultiprocExecutor(Executor): output = responses[0]["output_batch"] logging_info = None - logging_info = responses[0]["logging_info"] if envs.FASTVIDEO_STAGE_LOGGING else None + logging_info = responses[0]["logging_info"] if envs.FASTVIDEO_STAGE_LOGGING.get() else None # Get extra dict (contains audio, peak_memory_mb, etc.) extra = responses[0].get("extra", {}) @@ -743,7 +743,7 @@ class WorkerMultiprocProc: fastvideo_args = kwargs['fastvideo_args'] output_batch = self.worker.execute_forward(forward_batch, fastvideo_args) logging_info = None - if envs.FASTVIDEO_STAGE_LOGGING: + if envs.FASTVIDEO_STAGE_LOGGING.get(): logging_info = output_batch.logging_info # result tensor shared by CUDA IPC to avoid serialization overhead result = output_batch.output diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index 67c576bd4..a48e9e506 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -114,7 +114,7 @@ class RayDistributedExecutor(Executor): def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs): from fastvideo.platforms import current_platform - num_gpus = envs.FASTVIDEO_RAY_PER_WORKER_GPUS + num_gpus = envs.FASTVIDEO_RAY_PER_WORKER_GPUS.get() # The remaining workers are the actual ray actors. self.workers: list[RayWorkerWrapper] = [] @@ -326,7 +326,7 @@ class RayDistributedExecutor(Executor): output = responses[0].output.cpu() logging_info = None - if envs.FASTVIDEO_STAGE_LOGGING: + if envs.FASTVIDEO_STAGE_LOGGING.get(): logging_info = responses[0].logging_info result_batch = ForwardBatch( diff --git a/fastvideo/worker/ray_env.py b/fastvideo/worker/ray_env.py index 72caa0131..b3663afbb 100644 --- a/fastvideo/worker/ray_env.py +++ b/fastvideo/worker/ray_env.py @@ -8,7 +8,7 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) -CONFIG_HOME = envs.FASTVIDEO_CONFIG_ROOT +CONFIG_HOME = envs.FASTVIDEO_CONFIG_ROOT.get() # This file contains a list of env vars that should not be copied # from the driver to the Ray workers. diff --git a/mkdocs.yml b/mkdocs.yml index 3a67cb01c..80e390105 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -237,6 +237,7 @@ nav: - Pull Requests: contributing/pull_requests.md - CI/CD Architecture: contributing/ci_architecture.md - Testing: contributing/testing.md + - Environment Variables: contributing/env_vars.md - Performance Benchmarks: contributing/performance_benchmarks.md - Profiling: contributing/profiling.md - Activation Trace: contributing/activation_trace.md From 7f03e03dc6561f12815794294384b730bda4dbbb Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 29 Sep 2026 13:03:18 -0700 Subject: [PATCH 039/151] [docs]: FastH3 V2 cookbook serving and MLX install guide (#1884) Co-authored-by: Aryan Kumar --- README.md | 15 +- docs/assets/cookbook-recipes.json | 75 ++++- docs/assets/cookbook.js | 10 +- docs/attention/vsa/index.md | 13 +- docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/index.md | 2 +- docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 30 +- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/openai-api.md | 39 ++- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 4 +- docs/cookbook/z-image.md | 2 +- docs/getting_started/installation.md | 4 +- docs/getting_started/installation/mlx.md | 65 ++++ docs/getting_started/installation/mps.md | 277 ------------------ docs/getting_started/installation/spark.md | 2 +- .../installation/spark_pair.md | 2 +- docs/inference/fasth3-distilled.md | 11 +- docs/inference/inference_quick_start.md | 4 +- docs/inference/mlx_taeh3.md | 4 +- docs/inference/support_matrix.md | 14 +- examples/inference/basic/README.md | 20 +- examples/serving/README.md | 11 +- examples/serving/mlx_fasth3_8step.yaml | 26 ++ examples/serving/openai_fasth3_8step.yaml | 48 +++ fastvideo/entrypoints/openai/mlx_server.py | 72 +++-- fastvideo/registry.py | 1 + .../entrypoints/test_openai_video_client.py | 24 +- mkdocs.yml | 2 +- tests/local_tests/test_cookbook_serving.py | 48 +++ 39 files changed, 453 insertions(+), 396 deletions(-) create mode 100644 docs/getting_started/installation/mlx.md delete mode 100644 docs/getting_started/installation/mps.md create mode 100644 examples/serving/mlx_fasth3_8step.yaml create mode 100644 examples/serving/openai_fasth3_8step.yaml diff --git a/README.md b/README.md index f549fbb6a..dbca636b6 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ - `2026/09/15`: Release [FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2), an eight-forward data-free DMD2 checkpoint distilled from MiniMax-H3 with 80% Video Sparse Attention. Run it with `examples/inference/basic/basic_fasth3_8step.py` or the [FastH3 8-Step V2 recipe](https://haoailab.com/FastVideo/cookbook/minimax-h3/). - `2026/09/01`: FastH3 now runs locally on Apple Silicon through MLX and on NVIDIA DGX Spark through CUDA 13, including two-Spark inference. Follow the [FastH3 recipes](https://haoailab.com/FastVideo/cookbook/minimax-h3/) and read the [Blog](https://haoailab.com/blogs/fasth3-local/). - `2026/08/27`: [FastH3 Preview v1](https://haoailab.com/blogs/fasth3-preview/) is an open-weight 4-step sparse-distilled MiniMax-H3 model for synchronized video-and-audio generation, developed in collaboration with [Nuva Lab](https://nuvalab.ai/) and the [NVIDIA FastGen team](https://github.com/NVlabs/FastGen). Download the recommended [VSA / Data-Free weights](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree), or see the [full FastH3 collection](https://huggingface.co/collections/FastVideo/fastvideo-fasth3). -- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac—follow the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/) and read the [Blog](https://haoailab.com/blogs/fastmetal/). +- `2026/08/19`: FastVideo now supports MLX on Apple Silicon with [FastMetal-QAD](https://huggingface.co/collections/FastVideo/fastmetal), a family of 1.3B, 5B, and 14B models optimized for Mac. Follow the [MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/) and read the [Blog](https://haoailab.com/blogs/fastmetal/). - `2026/06/23`: Release FastWan-QAD: 5s of Video generated in 1.8s E2E. See the [FastWan-QAD models](https://huggingface.co/FastVideo/FastWan-QAD-FP8-1.3B), [Attn-QAT training guide](https://haoailab.com/FastVideo/training/attn_qat/), and [blog](https://haoailab.com/blogs/fastwan-qad/). - `2026/03/17`: Release demo: Into the Dreamverse: Vibe Directing in FastVideo, check out the [Blog](https://haoailab.com/blogs/dreamverse/). - `2026/03/13`: Release demo: Create a 5s 1080p Video in 4.5s with FastVideo on a Single GPU, check out the [Blog](https://haoailab.com/blogs/fastvideo_realtime_1080p/). @@ -64,13 +64,12 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo ``` Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the -[MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). +[MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/). -> **On an Apple Silicon Mac?** FastVideo runs FastMetal-QAD through an MLX -> runtime. Install with `uv pip install -e '.[mlx]'`, download -> [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD), -> and follow the -> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). +> **On an Apple Silicon Mac?** Install with `uv pip install -e '.[mlx]'` from +> a clone, then pick a recipe in the +> [cookbook](https://haoailab.com/FastVideo/cookbook/). See the +> [MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/). Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions. @@ -88,7 +87,7 @@ Install FastVideo (https://github.com/hao-ai-lab/FastVideo) into a fresh uv virt https://hao-ai-lab.github.io/FastVideo/getting_started/installation/): - NVIDIA GPU, x86_64 -> docs/getting_started/installation/gpu.md - NVIDIA DGX Spark / GB10, aarch64, CUDA 13 -> docs/getting_started/installation/spark.md - - Apple Silicon, macOS -> docs/getting_started/installation/mps.md + - Apple Silicon, macOS -> docs/getting_started/installation/mlx.md 3. Use uv for every step. If a command fails, debug it and tell me what you changed. 4. Verify the result: python -c "import fastvideo, torch; print('cuda', torch.cuda.is_available())" diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 48b4a38ab..8875361dc 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 9, + "version": 11, "recipes": [ { "id": "fastwan21-t2v", @@ -445,13 +445,13 @@ { "id": "fasth3-preview-cuda", "group": "fasth3-preview", - "group_label": "FastH3 Preview", + "group_label": "FastH3 V1", "group_task": "4-step text to video + audio", "family": "minimax_h3", "stage": "inference", "task": "Few-step text to video (with audio)", - "label": "FastH3 Preview on CUDA", - "summary": "Run the DMD2-distilled FastH3 Preview with four DiT forwards, trained H3 sparse attention, compiled decode, and synchronized audio.", + "label": "FastH3 V1 on CUDA", + "summary": "Run FastH3 V1 with four DiT forwards, trained H3 sparse attention, compiled decode, and synchronized audio.", "model": "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2", "source": "examples/inference/basic/basic_fasth3.py", "serving": { @@ -484,13 +484,13 @@ "prepare": "hf download FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2 --local-dir ./FastH3-Preview-v0.2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-Preview-v0.2/transformer --out ./FastH3-MLX --formats \"int6\"" }, "group": "fasth3-preview", - "group_label": "FastH3 Preview", + "group_label": "FastH3 V1", "group_task": "4-step text to video + audio", "family": "minimax_h3", "stage": "inference", "task": "Few-step text to video (with audio)", - "label": "FastH3 Preview on MLX", - "summary": "Run FastH3 Preview on Apple Silicon with a locally converted INT6 DiT, streamed Qwen3-VL conditioning, and native MLX video and audio VAEs.", + "label": "FastH3 V1 on MLX", + "summary": "Run FastH3 V1 on Apple Silicon with a locally converted INT6 DiT, streamed Qwen3-VL conditioning, and native MLX video and audio VAEs.", "model": "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2", "source": "examples/inference/basic/mlx_fasth3.py", "command": "hf download FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2 --local-dir ./FastH3-Preview-v0.2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-Preview-v0.2/transformer --out ./FastH3-MLX --formats \"int6\"\npython examples/inference/basic/mlx_fasth3.py --model-root ./FastH3-Preview-v0.2 --mlx-checkpoint ./FastH3-MLX/int6 --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --height 480 --width 832 --num-frames 124 --seed 2026 --output-path ./outputs/fasth3_int6.mp4", @@ -516,13 +516,13 @@ { "id": "fasth3-preview-spark", "group": "fasth3-preview", - "group_label": "FastH3 Preview", + "group_label": "FastH3 V1", "group_task": "4-step text to video + audio", "family": "minimax_h3", "stage": "inference", "task": "Few-step text to video (with audio)", - "label": "FastH3 Preview on one DGX Spark", - "summary": "Run FastH3 Preview on one GB10 with Triton VSA, FA4 off, and lazy module load. Height, width, frames, and steps in the YAML are examples.", + "label": "FastH3 V1 on one DGX Spark", + "summary": "Run FastH3 V1 on one GB10 with Triton VSA, FA4 off, and lazy module load. Height, width, frames, and steps in the YAML are examples.", "model": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", "source": "examples/inference/basic/basic_fasth3_spark.yaml", "serving": { @@ -551,12 +551,12 @@ { "id": "fasth3-spark-pair", "group": "fasth3-preview", - "group_label": "FastH3 Preview", + "group_label": "FastH3 V1", "group_task": "4-step text to video + audio", "family": "minimax_h3", "stage": "inference", "task": "Few-step text to video (with audio)", - "label": "FastH3 Preview on two DGX Sparks", + "label": "FastH3 V1 on two DGX Sparks", "summary": "Run one FastH3 clip across two GB10s with Ray sequence parallel over QSFP RoCE. Sequential load and lazy module load stay on because SP replicates the DiT on each node.", "model": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree", "source": "examples/inference/basic/basic_fasth3_spark_pair.yaml", @@ -581,15 +581,19 @@ { "id": "fasth3-8step-v2-cuda", "group": "fasth3-8step-v2", - "group_label": "FastH3 8-Step V2", + "group_label": "FastH3 V2", "group_task": "8-step text to video + audio", "family": "minimax_h3", "stage": "inference", "task": "Few-step text to video (with audio)", - "label": "FastH3 8-Step V2 on CUDA", - "summary": "Run the eight-forward FastH3 8-Step V2 checkpoint (video/audio shifts 10/3, VSA 0.8, 64-token tiles) with the trained DMD ladder loaded from the checkpoint's fastvideo_inference.json.", + "label": "FastH3 V2 on CUDA", + "summary": "Run FastH3 V2, the eight-forward checkpoint (video/audio shifts 10/3, VSA 0.8, 64-token tiles) with the trained DMD ladder loaded from the checkpoint's fastvideo_inference.json.", "model": "FastVideo/FastVideo-FastH3-8-Step-V2", "source": "examples/inference/basic/basic_fasth3_8step.py", + "serving": { + "source": "examples/serving/openai_fasth3_8step.yaml", + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"" + }, "command": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"\npython examples/inference/basic/basic_fasth3_8step.py --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --profile strict --no-inference-torch-compile --no-compile-vae", "gpu_types": ["NVIDIA"], "hardware": { @@ -609,7 +613,46 @@ "limitations": [ "Nine sigma-grid points (eight transformer forwards) are fixed by the checkpoint's trained ladder; the example rejects any other --steps.", "Validated with the eager strict route (--profile strict --no-inference-torch-compile --no-compile-vae). The compiled all profile has not been measured for this checkpoint.", - "T2AV only; no FL2VA/Ref2VA distillation and no matching LoRA. Quality and latency were not compared against the four-step preview in this PR." + "T2AV only; no FL2VA/Ref2VA distillation and no matching LoRA. V2 uses eight forwards rather than V1's four." + ] + }, + { + "id": "fasth3-8step-v2-mlx", + "serving": { + "source": "examples/serving/mlx_fasth3_8step.yaml", + "install": "uv pip install -e \".[mlx]\"", + "prepare": "hf download FastVideo/FastVideo-FastH3-8-Step-V2 --local-dir ./FastH3-8-Step-V2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-8-Step-V2/transformer --out ./FastH3-8-Step-V2-MLX --formats \"int8\" --include-vsa" + }, + "group": "fasth3-8step-v2", + "group_label": "FastH3 V2", + "group_task": "8-step text to video + audio", + "family": "minimax_h3", + "stage": "inference", + "task": "Few-step text to video (with audio)", + "label": "FastH3 V2 on MLX", + "summary": "Run FastH3 V2 on Apple Silicon with a locally converted INT8 DiT, the checkpoint's DMD contract ladder, and trained VSA (sparsity 0.8, 64-token tiles).", + "model": "FastVideo/FastVideo-FastH3-8-Step-V2", + "source": "examples/inference/basic/mlx_fasth3_8step.py", + "command": "hf download FastVideo/FastVideo-FastH3-8-Step-V2 --local-dir ./FastH3-8-Step-V2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-8-Step-V2/transformer --out ./FastH3-8-Step-V2-MLX --formats \"int8\" --include-vsa\npython examples/inference/basic/mlx_fasth3_8step.py --model-root ./FastH3-8-Step-V2 --mlx-checkpoint ./FastH3-8-Step-V2-MLX/int8 --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --height 480 --width 832 --num-frames 124 --seed 2026 --output-path ./outputs/fasth3_8step_int8.mp4", + "gpu_types": ["Apple Silicon"], + "hardware": { + "platform": "mlx", + "accelerator": "Apple M4 Max", + "system_memory": "36 GB unified memory", + "peak_memory": "27.39 GiB peak MLX memory during denoising", + "evidence": "validated", + "evidence_url": "https://github.com/hao-ai-lab/FastVideo/pull/1863" + }, + "evidence": "Verified", + "expected_artifact": "MP4 with H.264 video and stereo AAC audio at outputs/fasth3_8step_int8.mp4", + "modes": ["T2VA", "8-step FastH3", "trained VSA"], + "knobs": [ + {"key": "video_decode_backend", "label": "VAE decode", "hint": "Fidelity vs. speed", "flag": "--video-decode-backend", "options": [{"value": "h3-vae", "label": "Full H3 VAE"}, {"value": "taeh3", "label": "TAEH3 preview"}], "default": "h3-vae"} + ], + "limitations": [ + "Convert with --include-vsa. mlx_fasth3_8step.py turns VSA on (sparsity 0.8, tile 64). A dense export fails at configure_vsa.", + "--steps 8 or 9 both run the eight trained forwards. Reuse the preview VAE, audio VAE, text encoder, and tokenizer if those directories already exist.", + "The MLX path supports T2VA only. FL2VA, Ref2VA, and two-pass refinement are not wired." ] }, { diff --git a/docs/assets/cookbook.js b/docs/assets/cookbook.js index 00c11e595..5b00e95ae 100644 --- a/docs/assets/cookbook.js +++ b/docs/assets/cookbook.js @@ -136,7 +136,7 @@ if (platform === "mps") { return { id: "mps", - label: "Apple Silicon · MPS", + label: "Apple Silicon · PyTorch MPS", hint: recipe.hardware?.minimum_memory || recipe.hardware?.system_memory || "Memory not recorded", }; } @@ -519,10 +519,10 @@ option.setAttribute("aria-pressed", String(selected)); }); servingAvailability.textContent = profile - ? "The playground and API clients share one server process. Both workflows can run on your own machine." + ? "The playground and the OpenAI Python client share one server process. Both workflows can run on your own machine." : servingLoadFailed ? "Server examples could not be loaded. Open the H3 server guide below, or use Python directly." - : "This recipe uses Python directly. For the playground and API clients, choose FastH3 Preview with CUDA, MLX, or one Spark."; + : "This recipe uses Python directly. FastH3 V1 and FastH3 V2 can also run a local server for the playground and the OpenAI Python client."; servingPanel.hidden = !useServer; commandBlock.hidden = useServer; root.querySelector("[data-cookbook-python-note]").hidden = useServer; @@ -537,7 +537,7 @@ ? "Start once, then change prompts in the playground or your app. On a DGX Spark, lazy module load still reloads Qwen3-VL and the DiT between phases of each request, so later prompts are not a free hot cache." : "Start once, then change prompts in the playground or your app. CUDA requests reuse the loaded model. The Python SDK can also reuse a generator within one process."; servingPanel.querySelector("[data-cookbook-install-guide]").href = isMLX - ? "../../getting_started/installation/mps/#run-fasth3-preview" + ? "../../getting_started/installation/mlx/" : isSpark ? "../../getting_started/installation/spark/" : "../../getting_started/installation/gpu/"; @@ -574,7 +574,7 @@ }); description.textContent = useServer - ? `FastH3 Preview generates video with audio. This server profile uses the checked-in ${runtime.label} configuration.` + ? `${recipe.group_label || recipe.label} generates video with audio. Start the local server, then use the playground or the OpenAI Python client. This profile uses the checked-in ${runtime.label} configuration.` : recipe.summary; label.textContent = useServer ? `${recipe.group_label || recipe.label} · Server` : recipe.label; model.textContent = recipe.model; diff --git a/docs/attention/vsa/index.md b/docs/attention/vsa/index.md index 58f126bf3..3db0c812b 100644 --- a/docs/attention/vsa/index.md +++ b/docs/attention/vsa/index.md @@ -18,8 +18,8 @@ the CUDA `fastvideo-kernel` package: - **Dense-only checkpoints** (the default converter) drop the 50 gate matrices and keep fused SDPA. They remain valid for dense inference. - **VSA-capable checkpoints** retain those gates, quantize them on the same - affine grid, and record `vsa.capable` in `mlx_h3_dit.json`. Runtime VSA is - still off until you pass `--vsa`. + affine grid, and record `vsa.capable` in `mlx_h3_dit.json`. Preview leaves + runtime VSA off until you pass `--vsa`. `mlx_fasth3_8step.py` turns it on. - **Tile sizes** 64 `(4, 4, 4)` and 256 `(4, 8, 8)`. Prefix keys can be `exempt` or `compete`. `--vsa-dense-first-n-steps` and `--vsa-dense-layers` force dense SDPA on the selected steps or blocks. @@ -32,9 +32,12 @@ the CUDA `fastvideo-kernel` package: but does not yet match reference video. `--vsa-impl reference` is the same as `auto`. -See the [Apple Silicon guide](../../getting_started/installation/mps.md) for -conversion and `mlx_fasth3.py` flags. Do not enable VSA on a dense-only -checkpoint; reconvert with `--include-vsa` first. +See the [MLX install guide](../../getting_started/installation/mlx.md) +and the [MiniMax H3 cookbook](../../cookbook/minimax-h3.md) for conversion +and `mlx_fasth3.py` / `mlx_fasth3_8step.py` flags. Do not enable +VSA on a dense-only checkpoint; reconvert with `--include-vsa` first. V1 +VSA is opt-in. FastH3 V2 converts with `--include-vsa` and turns VSA on +by default. H3 uses fused MLX RMSNorm by default, including dense inference. This can change BF16 rounding relative to the older explicit normalization path. diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 3ecd3533f..2ee864e5b 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
    +
    All model families
    diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index 065dc3ffb..b8bbb5905 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
    +
    All model families
    diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index e805cfb9b..b4a2c3967 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
    +
    All model families
    diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index 0d13126d5..8408dad62 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
    +
    All model families
    diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md index 811815a72..be3a1f727 100644 --- a/docs/cookbook/index.md +++ b/docs/cookbook/index.md @@ -41,7 +41,7 @@ hide: MiniMax H3Video + stereo audio - 8 recipes + 9 recipes
    • T2VA
      • diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index a632a1507..e07993f24 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
      +
      All model families
      diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index aecb03273..f89806052 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
      +
      All model families
      diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index 64a2ca150..de1bf1c98 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
      +
      All model families
      diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index 4fcc881bb..8bcb544f5 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
      +
      All model families
      diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 23150f9b3..8f3b4f33b 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -5,13 +5,13 @@ hide: # MiniMax H3 recipes -The **FastH3 8-Step V2** recipe below runs the eight-forward -[`FastVideo/FastVideo-FastH3-8-Step-V2`](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) -checkpoint; its schedule contract is documented in +FastH3 is two distilled MiniMax-H3 checkpoints. **V1** is the four-step +launch. Some Hub repo names still say Preview. That name is historical. V1 is +a full model, not a demo. **V2** is the eight-step checkpoint. More forwards +is why V2 is the higher-quality FastH3. The V2 schedule contract is in [FastH3 distilled checkpoint schedules](../inference/fasth3-distilled.md). -The four-forward FastH3 Preview recipes are unchanged. -
      +
      All model families
      @@ -22,7 +22,7 @@ The four-forward FastH3 Preview recipes are unchanged.

      Primary focus · Inference

      MiniMax H3 recipes

      Generate video and audio with H3. Run a server on CUDA, one DGX Spark, or Apple Silicon MLX to iterate on prompts, or call the pipeline directly from Python.

      - 8 maintained recipes + 9 maintained recipes
      @@ -47,9 +47,9 @@ The four-forward FastH3 Preview recipes are unchanged.

      Supported modes

      CUDA covers T2VA, FL2VA, and Ref2VA on the full checkpoint, plus FastH3 - Preview and FastH3 LoRA. FastH3 Preview also has a DGX Spark runtime with - a 1-Spark or 2-Spark device row. MLX is T2VA only. Temporal --fast, - spatial --fast-spatial, and opt-in VSA are flags on the same + V1 and FastH3 V2. FastH3 V1 also has a DGX Spark runtime with + a 1-Spark or 2-Spark device row. MLX is T2VA only: V1 and V2. + Temporal --fast, spatial --fast-spatial, and opt-in VSA are flags on the same MLX script, not extra recipes.

      @@ -64,8 +64,8 @@ The four-forward FastH3 Preview recipes are unchanged. T2VA - Full H3, FastH3 Preview, FastH3 LoRA - FastH3 Preview after a local DiT conversion + Full H3, FastH3 V1, FastH3 LoRA, FastH3 V2 + FastH3 V1 or FastH3 V2 after a local DiT conversion FL2VA @@ -99,7 +99,7 @@ The four-forward FastH3 Preview recipes are unchanged. DGX Spark - FastH3 Preview on one GB10, or two Sparks with Ray sequence parallel (sp_size=2) over QSFP RoCE. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. + FastH3 V1 on one GB10, or two Sparks with Ray sequence parallel (sp_size=2) over QSFP RoCE. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. Not wired @@ -257,7 +257,7 @@ cd FastVideo 
      UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"

      Apple Silicon

      uv pip install -e ".[mlx]"
      -

      Follow the Apple Silicon guide for the download, conversion, and storage requirements.

      +

      Follow the MLX install guide for the extra, ffmpeg, and a clone. Then pick FastH3 V1 or V2 in the builder above.

      NVIDIA DGX Spark

      UV_TORCH_BACKEND=cu130 uv pip install -e .

      Follow the DGX Spark install guide for ARM64 CUDA 13. One Spark is a local process. Two Sparks need Ray on the QSFP link:

      @@ -272,9 +272,9 @@ cd FastVideo
      • The full CUDA H3 examples request four GPUs by default. Their sources do not claim a GPU model or memory minimum.
      • The FastH3 CUDA performance profile was measured on four GB200 GPUs. Use its strict profile when exact operation order matters more than the measured performance configuration.
        • -
      • The MLX source runtime supports T2VA, optional temporal --fast, optional spatial --fast-spatial, and opt-in VSA on --include-vsa checkpoints. FL2VA, Ref2VA, and two-pass refinement are not wired.
        • +
      • The MLX source runtime supports T2VA, optional temporal --fast, optional spatial --fast-spatial, and opt-in VSA on --include-vsa checkpoints. FastH3 V2 MLX converts with --include-vsa and runs eight forwards. FL2VA, Ref2VA, and two-pass refinement are not wired.
      • GPU count and VAE decode backend are configurable in the builder above for FastH3 CUDA recipes. Only the value shown by default has a recorded run; other supported values are unmeasured here.
        • -
      • DGX Spark is a runtime on FastH3 Preview, not a separate family card. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. The CUDA GPU-count knob does not apply to Spark.
        • +
      • DGX Spark is a runtime on FastH3 V1, not a separate family card. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. The CUDA GPU-count knob does not apply to Spark.
      • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM.
      • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
      diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 7637d2a85..8bdd93515 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
      +
      All model families
      diff --git a/docs/cookbook/openai-api.md b/docs/cookbook/openai-api.md index c3fc952f6..c0326120b 100644 --- a/docs/cookbook/openai-api.md +++ b/docs/cookbook/openai-api.md @@ -8,8 +8,9 @@ JavaScript client examples use that interface, but requests go to FastVideo. You do not need an OpenAI account or cloud key. The [H3 recipe selector](minimax-h3.md) provides the same workflow with runtime -selection. This guide covers FastH3 Preview text-to-video/audio. Other H3 -recipes keep their direct Python commands. +selection. This guide covers FastH3 V1 (four forwards) and FastH3 V2 +(eight forwards). Start a server, then iterate in the playground or with the +OpenAI Python client. Other H3 recipes keep their direct Python commands. CUDA requests reuse one loaded `VideoGenerator`. The Python SDK can do the same when you reuse the generator across `generate()` calls. MLX keeps one @@ -36,6 +37,14 @@ advertises it as `fasth3`. It configures four CUDA GPUs but does not record a GPU model or VRAM requirement. This is a source-backed server profile, not the measured GB200 Python performance profile. Compilation is disabled. +For FastH3 V2, use the same install and the 8-step config (nine sigma +points, eight DiT forwards, VSA sparsity 0.8): + +```bash +UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]" +fastvideo serve --config examples/serving/openai_fasth3_8step.yaml --server.host 127.0.0.1 +``` + Keep the server running. In another terminal, check readiness: ```bash @@ -72,7 +81,7 @@ After model loading completes, the response is `{"status":"ok"}`. ### Apple Silicon MLX -Complete the [Apple Silicon installation](../getting_started/installation/mps.md#run-fasth3-preview), +Complete the [MLX install](../getting_started/installation/mlx.md), including `ffmpeg`. From your FastVideo clone, install the MLX extra: ```bash @@ -111,6 +120,22 @@ The MLX server has no recorded device or unified-memory requirement. The direct Python recipe's M4 Max measurements are not a server benchmark or a minimum-memory claim. +For FastH3 V2, convert that checkpoint's transformer with `--include-vsa` +and start the 8-step config. Do not overwrite a V1 export. The +[MiniMax H3 cookbook](minimax-h3.md) has the same download and convert +commands as the Python recipe. + +```bash +hf download FastVideo/FastVideo-FastH3-8-Step-V2 --local-dir ./FastH3-8-Step-V2 +python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-8-Step-V2/transformer --out ./FastH3-8-Step-V2-MLX --formats "int8" --include-vsa +python -m fastvideo.entrypoints.openai.mlx_server --config examples/serving/mlx_fasth3_8step.yaml +``` + +That adapter passes `num_steps=8` when the HTTP field is `num_inference_steps=9`, +and it enables the trained VSA recipe (sparsity 0.8, 64-token tiles). Reuse the +preview VAE, audio VAE, text encoder, and tokenizer if those directories already +exist; edit the YAML paths if your files live elsewhere. + ## Open the playground Open [the local H3 playground](http://127.0.0.1:8000/playground/) after startup. @@ -136,9 +161,11 @@ or manage a GPU server for you. ## Generate with cURL or an SDK These examples use the server's resolution, frame count, and sampling defaults. -Do not copy Sora-specific durations or resolutions onto H3. Both server configs -use 124 frames, 24 fps, and the five-point distilled sigma schedule with four -DiT forwards. CUDA and one Spark use 1344 × 768; MLX uses 832 × 480. +Do not copy Sora-specific durations or resolutions onto H3. V1 configs use +124 frames, 24 fps, and the five-point distilled sigma schedule with four DiT +forwards. V2 configs use nine sigma points and eight DiT forwards. CUDA +and one Spark use 1344 × 768; MLX uses 832 × 480. The OpenAI Python client is +the same for every FastH3 server that advertises `fasth3`. Each client submits a job, checks for completion or failure, and downloads an MP4 named after the job ID. Polling stops after 30 minutes; a timeout does not diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index d9731475d..a2ca588ef 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
      +
      All model families
      diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index fd09103fc..a3c8c62ab 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index 28a4ee124..92b259fec 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index 8eb528632..4554e0966 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
      +
      All model families
      @@ -186,7 +186,7 @@ cd FastVideo
      • Out of memory on the A14B recipes: the checked-in sources already enable CPU offload; see Configuration for the offload surface before reducing resolution or frames.
      • The FastWan2.1 recipe requires VIDEO_SPARSE_ATTN; confirm the environment variable in the command was set in the same shell.
        • -
      • FastMetal MLX: install with uv pip install -e ".[mlx]", then follow the Apple Silicon guide. CUDA FastWan-QAD checkpoints are refused on the MLX runtime.
        • +
      • FastMetal MLX: install with the MLX install guide, then pick a FastMetal recipe in the builder. CUDA FastWan-QAD checkpoints are refused on the MLX runtime.
      • FastMetal 5B uses mlx_wan22_generate.py. 1.3B and 14B use mlx_wan_prompt_to_video.py.
      • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
      diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index 64caee4d8..8f4301f64 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
      +
      All model families
      diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 50ee0a922..c0404ec6b 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -7,7 +7,7 @@ FastVideo supports the following hardware platforms: - **NVIDIA DGX Spark / GB10 (ARM64 + CUDA 13)** — [install](installation/spark.md), [performance](installation/spark_performance.md), [pair two Sparks](installation/spark_pair.md) -- [Apple silicon](installation/mps.md) +- [Apple silicon (MLX)](installation/mlx.md) ## Quick Installation @@ -15,7 +15,7 @@ FastVideo supports the following hardware platforms: Use uv as the default environment manager for faster and more stable installs. The commands below target NVIDIA CUDA 12; use `UV_TORCH_BACKEND=cu130` on -CUDA 13. Apple silicon users should follow the [MPS guide](installation/mps.md). +CUDA 13. Apple silicon users should follow the [MLX install guide](installation/mlx.md). ```bash # Create and activate a new uv environment diff --git a/docs/getting_started/installation/mlx.md b/docs/getting_started/installation/mlx.md new file mode 100644 index 000000000..359218e58 --- /dev/null +++ b/docs/getting_started/installation/mlx.md @@ -0,0 +1,65 @@ +# Install FastVideo with MLX + +Install FastVideo on a Mac, then generate from the cookbook. Local video uses +the native MLX runtime, not CUDA, and not the old PyTorch MPS demo at +`examples/inference/basic/basic_mps.py`. + +## Requirements + +- macOS 14 or newer +- Python 3.12 +- `ffmpeg` (`brew install ffmpeg`) + +## Install + +Cookbook commands run from a clone. + +```bash +git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo +uv venv --python 3.12 --seed +source .venv/bin/activate +brew install ffmpeg +uv pip install -e ".[mlx]" +``` + +Conda is optional. After you activate a Conda env, still install with +`uv pip` as above. + +`uv pip install "fastvideo[mlx]"` from PyPI installs the extra only. It does +not ship the example scripts the cookbook copies. + +## Generate a video + +Open the cookbook. Select Apple Silicon as the runtime. Each recipe has a +Python command. FastH3 also has a server path for the playground and the +OpenAI Python client. + +- [Wan recipes](../../cookbook/wan.md) for FastMetal 1.3B, 5B, and 14B +- [MiniMax H3 recipes](../../cookbook/minimax-h3.md) for FastH3 V1 and FastH3 V2 +- [H3 server guide](../../cookbook/openai-api.md) for the playground, cURL, and SDKs + +FastH3 is two distilled MiniMax-H3 checkpoints. V1 is the four-step launch. +Some Hub repo names still say Preview. That name is historical. V1 is a full +model, not a demo. V2 is the eight-step checkpoint. More forwards is why V2 +is the higher-quality FastH3. + +Recorded shapes and evidence live in the +[support matrix](../../inference/support_matrix.md#apple-silicon-native-runtime). + +## Hardware + +- FastMetal 1.3B and 5B: 16 GB unified memory and up +- FastMetal 14B: 36 GB unified memory and up +- FastH3 V1 and V2: validated on an M4 Max with 36 GB unified memory + +## Troubleshooting + +- **`basic_mps.py` is the wrong path.** That script is PyTorch MPS. Use an + Apple Silicon recipe in the cookbook. +- **Muxing fails.** Install `ffmpeg` with Homebrew. +- **A cookbook command cannot find a script.** Run it from the FastVideo + clone after `uv pip install -e ".[mlx]"`. + +If that does not match what you see, open an issue on the +[GitHub repository](https://github.com/hao-ai-lab/FastVideo) or ask in the +[Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ). diff --git a/docs/getting_started/installation/mps.md b/docs/getting_started/installation/mps.md deleted file mode 100644 index 2d7f3ac4e..000000000 --- a/docs/getting_started/installation/mps.md +++ /dev/null @@ -1,277 +0,0 @@ -# MPS (Apple Silicon) - -Install FastVideo on Apple Silicon and run FastMetal-QAD or FastH3 Preview. - -Apple Silicon uses the MLX runtime. FastMetal-QAD ships ready-to-run MLX -checkpoints; FastH3 Preview currently requires a local MLX DiT conversion. -See the [FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/) and the -[FastMetal collection](https://huggingface.co/collections/FastVideo/fastmetal). - -## Requirements - -- **OS: macOS 14 or newer** -- **Python: 3.12.4** - -## Set up using Python - -### Create a new Python environment - -#### uv -Recommended default: use [uv](https://docs.astral.sh/uv/) for faster and more stable environment setup. - -Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, create a new environment using: - -```console -# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools`. -uv venv --python 3.12 --seed -source .venv/bin/activate -``` - -#### Conda (alternative) - -You can also create a Python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html). - -##### 1. Install Miniconda (if not already installed) - -```bash -wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -bash Miniconda3-latest-MacOSX-arm64.sh -source ~/.zshrc -``` - -##### 2. Create and activate a Conda environment for FastVideo - -```bash -conda create -n fastvideo python=3.12.4 -y -conda activate fastvideo -``` - -### Dependencies - -``` -brew install ffmpeg -``` - -### Installation - -FastMetal's native Apple Silicon runtime requires the `mlx` extra. - -#### With uv (recommended) - -```bash -uv pip install "fastvideo[mlx]" -``` - -#### With Conda environment (alternative) - -`uv` works inside an active conda env too, so prefer `uv pip` for the actual install: - -```bash -uv pip install "fastvideo[mlx]" -``` - -### Installation from Source - -#### 1. Clone the FastVideo repository - -```bash -git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo -``` - -#### 2. Install FastVideo - -Basic installation: - -```bash -uv pip install -e ".[mlx]" -``` - -Alternative with Conda environment: - -```bash -uv pip install -e ".[mlx]" -``` - -## Run FastMetal-QAD - -Each release is self-contained. Download one checkpoint and point both -`--model-root` and `--mlx-checkpoint` at it (the example also auto-detects -`mlx_dit.json` under `--model-root`). - -| Checkpoint | Script | Mac tier | -| --- | --- | --- | -| [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | `mlx_wan_prompt_to_video.py` | 16 GB+ | -| [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | `mlx_wan22_generate.py` | 16 GB+ | -| [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | `mlx_wan_prompt_to_video.py` | 36 GB+ | - -```bash -hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD - -python examples/inference/basic/mlx_wan_prompt_to_video.py \ - --model-root ./FastMetal-1.3B-QAD \ - --mlx-checkpoint ./FastMetal-1.3B-QAD \ - --height 480 --width 832 --num-frames 81 \ - --prompt "A bird's-eye view of a misty forest valley at dawn." -``` - -14B uses the same script. Point both flags at `./FastMetal-14B-QAD`. That repo also ships an EMA variant: keep `--model-root` at the repo root and set `--mlx-checkpoint ./FastMetal-14B-QAD/ema`. - -Wan2.2 5B uses a different latent layout, so it has its own entrypoint: - -```bash -hf download FastVideo/FastMetal-5B-QAD --local-dir ./FastMetal-5B-QAD - -python examples/inference/basic/mlx_wan22_generate.py \ - --mlx-checkpoint ./FastMetal-5B-QAD \ - --text-encoder-root ./FastMetal-5B-QAD \ - --vae-root ./FastMetal-5B-QAD/vae \ - --height 704 --width 1280 --num-frames 81 \ - --prompt "A cinematic portrait with soft neon lighting and smooth camera motion." -``` - -CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`, `FastVideo/FastWan-QAD-FP8-1.3B`) is a separate NVIDIA release. The MLX examples look for FastMetal packed weights (`mlx_dit.json`). - -`basic_mps.py` is a generic PyTorch MPS demo. For local video on Mac, use the FastMetal commands above. - -## Run FastH3 Preview - -FastH3 Preview uses the existing MLX runtime for text-to-video-with-audio -(T2VA). The runtime streams the Qwen3-VL text conditioner, loads one -heavyweight component at a time, denoises synchronized video and audio -latents with a converted INT8, INT6, or INT4 DiT, and decodes both modalities -with native MLX VAEs. - -Download the FastH3 snapshot, then convert one or more DiT formats: - -```bash -hf download FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2 \ - --local-dir ./FastH3-Preview-v0.2 - -python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \ - --model-root ./FastH3-Preview-v0.2/transformer \ - --out ./FastH3-MLX \ - --formats "int6" -``` - -Dense conversion drops the trained VSA gate projections. To keep them (INT6 -weight-only, same affine grid as the other linear matrices) write a **new** -directory: - -```bash -python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \ - --model-root ./FastH3-Preview-v0.2/transformer \ - --out ./FastH3-MLX-vsa \ - --formats "int6" \ - --include-vsa -``` - -Do not overwrite an existing dense export such as `./FastH3-MLX/int6`. - -Run the baseline path: - -```bash -python examples/inference/basic/mlx_fasth3.py \ - --model-root ./FastH3-Preview-v0.2 \ - --mlx-checkpoint ./FastH3-MLX/int6 \ - --prompt "(S1) A presenter says [English] Fast H3 is amazing." \ - --height 480 --width 832 --num-frames 124 --seed 2026 \ - --output-path ./outputs/fasth3_int6.mp4 -``` - -Add `--fast` for temporal fast mode. It denoises a shorter video sequence, -uses MLX RIFE to restore the requested frame count, and keeps the audio -sequence at full duration: - -```bash -python examples/inference/basic/mlx_fasth3.py \ - --model-root ./FastH3-Preview-v0.2 \ - --mlx-checkpoint ./FastH3-MLX/int6 \ - --prompt "(S1) A presenter says [English] Fast H3 is even faster." \ - --height 720 --width 1280 --num-frames 124 --seed 2027 \ - --fast \ - --output-path ./outputs/fasth3_int6_fast_720p.mp4 -``` - -Add `--fast-spatial` for spatial fast mode, `--fast`'s spatial twin. It -denoises and decodes on the smallest 32px-aligned canvas covering the -requested size divided by `--fast-spatial-scale` (a 480x832 request runs on a -256x416 canvas), then resamples the decoded frames up to the requested size -in pixel space. It composes with `--fast`. This is a speed/quality trade-off -and stays off by default: the output carries the reduced canvas's detail -budget, so it reads softer than a native-resolution render, with the unsharp -pass countering some but not all of the difference: - -```bash -python examples/inference/basic/mlx_fasth3.py \ - --model-root ./FastH3-Preview-v0.2 \ - --mlx-checkpoint ./FastH3-MLX/int6 \ - --prompt "(S1) A presenter says [English] Fast H3 is fastest." \ - --height 480 --width 832 --num-frames 124 --seed 2028 \ - --fast --fast-spatial \ - --output-path ./outputs/fasth3_int6_fast_spatial.mp4 -``` - -VSA is off by default. A dense-only checkpoint (no `--include-vsa`) keeps the -existing fused-SDPA path. After converting with `--include-vsa`, enable the -sparse path explicitly: - -```bash -python examples/inference/basic/mlx_fasth3.py \ - --model-root ./FastH3-Preview-v0.2 \ - --mlx-checkpoint ./FastH3-MLX-vsa/int6 \ - --vsa --vsa-sparsity 0.9 --vsa-tile-size 64 --vsa-prefix-mode exempt \ - --prompt "(S1) A presenter says [English] Fast H3 is amazing." \ - --height 720 --width 1280 --num-frames 124 --seed 2026 \ - --output-path ./outputs/fasth3_int6_vsa_720p.mp4 -``` - -FastH3 8-Step V2 is a separate checkpoint. Keep the four-step command above -for the preview. For the eight-forward model, convert its transformer with the -same script (the snapshot's `fastvideo_inference.json` declares the trained -shifts, so the AdaLN cache follows those rungs) and run: - -```bash -python examples/inference/basic/mlx_fasth3_8step.py \ - --model-root ./FastH3-8-Step-V2 \ - --mlx-checkpoint ./FastH3-8-Step-V2-MLX/int6 \ - --prompt "(S1) A presenter says [English] Fast H3 is amazing." \ - --height 480 --width 832 --num-frames 124 --seed 2026 \ - --output-path ./outputs/fasth3_8step_int6.mp4 -``` - -`--steps` may be 8 or 9. Both run the eight trained forwards. VAE, audio VAE, -text encoder, and tokenizer can be the ones already used by the preview. - -`--vsa-impl auto` uses the chunked gather+SDPA **reference** path. -`--vsa-impl simd` is an opt-in SIMD-group kernel (tile 64, head dim 128) that -falls back to reference on unsupported shapes. It is not the default. -`--vsa-impl reference` is the same as `auto`. - -!!! note "Current MLX scope" - This source runtime supports T2VA, temporal `--fast`, spatial - `--fast-spatial`, and opt-in VSA. FL2VA, Ref2VA, two-pass refinement, and - `VideoGenerator` registry dispatch are not wired yet. INT8/INT6/INT4 - quantization is **weight-only**; VSA attention Q/K/V stay BF16. Old dense - MLX checkpoints remain valid for dense inference and raise a reconvert - error if `--vsa` is set. The checkpoint uses the MiniMax H3 - Community License; review the model card before use or redistribution. - -## Development Environment Setup - -If you're planning to contribute to FastVideo please see the following page: -[Contributor Guide](../../contributing/overview.md) - -## Hardware Requirements - -- **1.3B / 5B:** 16 GB unified memory and up (M1 and later) -- **14B:** 36 GB unified memory and up -- **FastH3 Preview:** validated on an M4 Max with 36 GB unified memory; use one - converted DiT format at a time and leave substantial free disk space for the - source snapshot plus the converted checkpoint -- Fanless 13-inch MacBook Air can run 1.3B and 5B at the same resolutions - -## Troubleshooting - -If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo). - -You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ) for additional support. diff --git a/docs/getting_started/installation/spark.md b/docs/getting_started/installation/spark.md index f6e31bc27..33888a853 100644 --- a/docs/getting_started/installation/spark.md +++ b/docs/getting_started/installation/spark.md @@ -147,7 +147,7 @@ that can't move here. Two Sparks with QSFP cables: [Pair two NVIDIA DGX Sparks](spark_pair.md) for one FastH3 clip across both GPUs (`sp_size=2` over Ray). Copy-paste commands for one or two Sparks also live on the -[MiniMax H3 cookbook](../../cookbook/minimax-h3.md): pick FastH3 Preview, +[MiniMax H3 cookbook](../../cookbook/minimax-h3.md): pick FastH3 V1, then NVIDIA DGX Spark, then 1 Spark or 2 Sparks. ## Development Environment Setup diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index 4eb86cc9e..fe8aee1aa 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/docs/getting_started/installation/spark_pair.md @@ -7,7 +7,7 @@ by the QSFP ConnectX-7 cables can run **one clip faster** and can hold a **longer clip** (up to the FastH3 15 s cap). Copy-paste commands for both counts are on the -[MiniMax H3 cookbook](../../cookbook/minimax-h3.md): FastH3 Preview → NVIDIA DGX +[MiniMax H3 cookbook](../../cookbook/minimax-h3.md): FastH3 V1 → NVIDIA DGX Spark → 1 Spark or 2 Sparks. This is FastVideo sequence parallel (`sp_size=2`) over Ray, not a third-party diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index 0a206914f..27035b92f 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -76,10 +76,15 @@ prompts before adopting it. ## Apple Silicon -`examples/inference/basic/mlx_fasth3.py` stays on the four-step preview and its -uniform AdaLN cache. `mlx_fasth3_8step.py` is the eight-forward MLX recipe. It +`mlx_fasth3.py` stays on FastH3 V1 and its uniform AdaLN cache. +`mlx_fasth3_8step.py` is the eight-forward MLX recipe for FastH3 V2. It reads the same `fastvideo_inference.json` rungs and shifts, and it expects an -MLX DiT whose AdaLN cache was converted from that contract. Reuse the preview +MLX DiT whose AdaLN cache was converted from that contract. Reuse the V1 snapshot's VAE, audio VAE, text encoder, and tokenizer; only the DiT and the sidecar change. Rank-reduced AdaLN checkpoints are unchanged and are not produced by the MLX converter. + +Install is in the +[MLX install guide](../getting_started/installation/mlx.md). Generation and +serving are in the [MiniMax H3 cookbook](../cookbook/minimax-h3.md) and the +[H3 server guide](../cookbook/openai-api.md). diff --git a/docs/inference/inference_quick_start.md b/docs/inference/inference_quick_start.md index c3e7849f6..8e12fc62c 100644 --- a/docs/inference/inference_quick_start.md +++ b/docs/inference/inference_quick_start.md @@ -5,10 +5,10 @@ This page contains step-by-step instructions to get you quickly started with vid ## Requirements - **OS**: Linux (tested on Ubuntu 22.04+), or macOS on Apple silicon via the - [MPS installation guide](../getting_started/installation/mps.md) + [MLX install guide](../getting_started/installation/mlx.md) - **Python**: 3.10-3.12 - **CUDA**: 12.6 or 13.0 (NVIDIA GPUs) -- **GPU**: At least one NVIDIA GPU, or an Apple silicon chip with MPS +- **GPU**: At least one NVIDIA GPU, or an Apple silicon chip with the MLX runtime ## Installation diff --git a/docs/inference/mlx_taeh3.md b/docs/inference/mlx_taeh3.md index b85626b2f..ac992f24a 100644 --- a/docs/inference/mlx_taeh3.md +++ b/docs/inference/mlx_taeh3.md @@ -10,7 +10,9 @@ that tradeoff. It is not a lossless acceleration of the full VAE. ## Generate a video -Use your existing MLX FastH3 environment and converted checkpoint: +Use your existing MLX FastH3 environment and converted checkpoint. The same +`--video-decode-backend taeh3` flag works on `mlx_fasth3.py` (V1) and +`mlx_fasth3_8step.py` (V2): ```bash python examples/inference/basic/mlx_fasth3.py \ diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 6c58ebcb1..9b2bc90ec 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -185,13 +185,14 @@ optimizations: absence means **untested**, not incompatible. | MLX FastMetal T2V 1.3B | [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released | | MLX FastMetal T2V 5B | [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode; optional `--fast`, `--fast-spatial`, `--refine`. The checked-in example is T2V. CUDA Wan2.2 TI2V 5B is the image-capable path. | Apple M4 Max, 16 GB+ unified memory | Released | | MLX FastMetal T2V 14B | [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB+ unified memory | Released | -| MLX FastH3 Preview T2VA | [`FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2`](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) + locally converted DiT | 480p / 720p, 124 frames, 4-step DMD2, INT8/INT6/INT4 **weight-only** DiT, native video + audio VAE; optional temporal RIFE fast mode; optional spatial fast mode; optional VSA (tile 64/256, exempt/compete) on `--include-vsa` checkpoints | Apple M4 Max, 36 GB unified memory | Source runtime; T2VA only | +| MLX FastH3 V1 T2VA | [`FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2`](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) + locally converted DiT | 480p / 720p, 124 frames, 4-step DMD2, INT8/INT6/INT4 **weight-only** DiT, native video + audio VAE; optional temporal RIFE fast mode; optional spatial fast mode; optional VSA (tile 64/256, exempt/compete) on `--include-vsa` checkpoints | Apple M4 Max, 36 GB unified memory | Source runtime; T2VA only | +| MLX FastH3 V2 T2VA | [`FastVideo/FastVideo-FastH3-8-Step-V2`](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2) + locally converted DiT | 480p / 720p, 124 frames, 8-step DMD2, INT8 **weight-only** DiT with `--include-vsa`, trained VSA 0.8 / tile 64, native video + audio VAE; AdaLN cache from `fastvideo_inference.json` | Apple M4 Max, 36 GB unified memory | Source runtime; T2VA only | Apple Silicon uses the native MLX runtime. FastMetal-QAD is the packaged Wan -release, while FastH3 Preview currently uses a source checkout and local DiT -conversion. CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`, +release. FastH3 V1 and FastH3 V2 use a source checkout and local DiT conversion. +CUDA FastWan-QAD (`FastVideo/FastWan-QAD-1.3B`, `FastVideo/FastWan-QAD-FP8-1.3B`) is the NVIDIA release. See the -[Apple Silicon guide](../getting_started/installation/mps.md) and the +[MLX install guide](../getting_started/installation/mlx.md) and the [FastMetal-QAD blog](https://haoailab.com/blogs/fastmetal/). **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. @@ -223,8 +224,9 @@ Per the installation guides: the [DGX Spark install guide](../getting_started/installation/spark.md). Two Sparks over QSFP use Ray sequence parallel; see [Pair two NVIDIA DGX Sparks](../getting_started/installation/spark_pair.md). -- **Apple silicon** — macOS 14 or newer; FastMetal-QAD via the MLX runtime. See the - [Apple Silicon guide](../getting_started/installation/mps.md). The older +- **Apple silicon** — macOS 14 or newer; MLX runtime for FastMetal-QAD, FastH3 + V1, and FastH3 V2. See the + [MLX install guide](../getting_started/installation/mlx.md). The older [`basic_mps.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_mps.py) demo is PyTorch MPS only. diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 52e1e0711..7c948bd9b 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -21,7 +21,7 @@ python examples/inference/basic/basic.py ### Apple Silicon (FastMetal-QAD) Use the MLX runtime with FastMetal-QAD. See the -[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). +[MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/). ```bash hf download FastVideo/FastMetal-1.3B-QAD --local-dir ./FastMetal-1.3B-QAD @@ -39,7 +39,7 @@ with `examples/inference/basic/basic_mps.py` is the older PyTorch MPS demo. -FastH3 Preview T2VA also runs through the native MLX runtime. Convert the DiT +FastH3 V1 T2VA also runs through the native MLX runtime. Convert the DiT to INT8, INT6, or INT4 first, then run: ```bash @@ -51,17 +51,21 @@ python examples/inference/basic/mlx_fasth3.py \ --output-path ./outputs/fasth3_int6.mp4 ``` +FastH3 V2 is a separate checkpoint and script. Convert with +`--include-vsa` and run `mlx_fasth3_8step.py`. Do not reuse a V1 DiT. + Pass `--fast` for temporal RIFE fast mode and `--fast-spatial` for spatial fast mode (reduced-canvas denoise + pixel-space upsample); the two compose. -VSA is opt-in: convert with `--include-vsa` and pass `--vsa` (see the -[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/)). -This MLX entrypoint currently supports T2VA only; FL2VA, Ref2VA, and +V1 VSA is opt-in: convert with `--include-vsa` and pass `--vsa`. V2 +turns VSA on by default. See the +[MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/). +This MLX entrypoint supports T2VA only; FL2VA, Ref2VA, and two-pass refinement remain follow-up work. INT6/INT8/INT4 are -weight-only; VSA attention activations stay BF16. Dense-only checkpoints keep -working for dense inference. +weight-only; VSA attention activations stay BF16. Dense-only V1 +checkpoints keep working for dense inference. The complete setup and conversion commands are in the -[Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). +[MLX install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mlx/). For an example running DMD+VSA inference: ``` diff --git a/examples/serving/README.md b/examples/serving/README.md index 86e5d2cdb..fcd339e45 100644 --- a/examples/serving/README.md +++ b/examples/serving/README.md @@ -41,10 +41,13 @@ and API clients share the same server; both can run locally. For native Apple Silicon MLX, use `python -m fastvideo.entrypoints.openai.mlx_server --config examples/serving/mlx_fasth3.yaml` -after preparing the weights and editing their paths in that configuration. -The same playground and clients work with MLX. Its pipeline still loads and -releases components between phases to limit unified-memory use. See the -cookbook for setup and the supported text-to-video/audio request fields. +(Preview) or `examples/serving/mlx_fasth3_8step.yaml` (8-Step V2) after preparing +the weights and editing their paths in that configuration. The same playground +and OpenAI Python client work with MLX. Its pipeline still loads and releases +components between phases to limit unified-memory use. See the cookbook for +setup and the supported text-to-video/audio request fields. + +CUDA FastH3 8-Step V2 serving uses `examples/serving/openai_fasth3_8step.yaml`. ## Wan configs diff --git a/examples/serving/mlx_fasth3_8step.yaml b/examples/serving/mlx_fasth3_8step.yaml new file mode 100644 index 000000000..632c4e552 --- /dev/null +++ b/examples/serving/mlx_fasth3_8step.yaml @@ -0,0 +1,26 @@ +# Native MLX FastH3 8-Step V2 server. Convert the DiT with --include-vsa first. +runtime: mlx +generator: + model_path: FastVideo/FastVideo-FastH3-8-Step-V2 + model_root: ./FastH3-8-Step-V2 + mlx_checkpoint: ./FastH3-8-Step-V2-MLX/int8 + prompt_cache_dir: ./outputs/h3_prompt_cache + vae_dtype: fp32 + vsa: true + vsa_sparsity: 0.8 + vsa_tile_size: 64 +server: + host: 127.0.0.1 + port: 8000 + served_model_name: fasth3 + output_dir: outputs/mlx_fasth3_8step +default_request: + sampling: + height: 480 + width: 832 + num_frames: 124 + fps: 24 + # HTTP convention: nine sigma points, eight DiT forwards. Native MLX uses num_steps=8. + num_inference_steps: 9 + guidance_scale: 1.0 + seed: 2026 diff --git a/examples/serving/openai_fasth3_8step.yaml b/examples/serving/openai_fasth3_8step.yaml new file mode 100644 index 000000000..243dcbbe3 --- /dev/null +++ b/examples/serving/openai_fasth3_8step.yaml @@ -0,0 +1,48 @@ +# OpenAI-compatible FastH3 8-Step V2 server. +generator: + model_path: FastVideo/FastVideo-FastH3-8-Step-V2 + engine: + num_gpus: 4 + use_fsdp_inference: false + parallelism: + tp_size: 1 + sp_size: 4 + offload: + dit: false + dit_layerwise: false + text_encoder: true + image_encoder: true + vae: true + pin_cpu_memory: true + compile: + enabled: false + vae_enabled: false + pipeline: + workload_type: t2v + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + VSA_sparsity: 0.8 + VSA_tile_size: 64 + inference_torch_compile: false + vae_parallel_decode: true + vae_parallel_decode_strategy: gather + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_fasth3_8step + served_model_name: fasth3 + +default_request: + negative_prompt: "" + sampling: + height: 768 + width: 1344 + num_frames: 124 + fps: 24 + num_inference_steps: 9 + guidance_scale: 1.0 + batch_cfg: false + seed: 1000 + output: + return_frames: false diff --git a/fastvideo/entrypoints/openai/mlx_server.py b/fastvideo/entrypoints/openai/mlx_server.py index 024f1187c..4bf1b5f4a 100644 --- a/fastvideo/entrypoints/openai/mlx_server.py +++ b/fastvideo/entrypoints/openai/mlx_server.py @@ -21,16 +21,36 @@ from fastvideo.api.schema import GenerationRequest from fastvideo.entrypoints.openai.api_server import create_app from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest -MODEL: Literal["FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"] = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2" +PREVIEW_MODEL: Literal["FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"] = ( + "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2") +EIGHT_STEP_MODEL: Literal["FastVideo/FastVideo-FastH3-8-Step-V2"] = "FastVideo/FastVideo-FastH3-8-Step-V2" +# HTTP convention: sigma-grid points. Native MLX generate() uses transformer forwards. +HTTP_STEPS_TO_FORWARDS = {5: 4, 9: 8} + + +def mlx_http_steps(model_path: str) -> int: + return 9 if model_path == EIGHT_STEP_MODEL else 5 + + +def mlx_num_steps(num_inference_steps: int | None, *, model_path: str) -> int: + expected = mlx_http_steps(model_path) + if num_inference_steps not in (None, expected): + raise ValueError(f"This FastH3 MLX server uses {expected} sigma points " + f"({HTTP_STEPS_TO_FORWARDS[expected]} transformer forwards).") + return HTTP_STEPS_TO_FORWARDS[expected] class MLXGeneratorConfig(BaseModel): model_config = ConfigDict(extra="forbid") - model_path: Literal["FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"] = MODEL + model_path: Literal["FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2", + "FastVideo/FastVideo-FastH3-8-Step-V2"] = PREVIEW_MODEL model_root: str mlx_checkpoint: str prompt_cache_dir: str = "outputs/h3_prompt_cache" vae_dtype: Literal["fp32", "fp16", "bf16"] = "fp32" + vsa: bool = False + vsa_sparsity: float = Field(default=0.9, ge=0.0, lt=1.0) + vsa_tile_size: Literal[64, 256] = 64 class MLXServerConfig(BaseModel): @@ -49,7 +69,7 @@ class MLXServeConfig(BaseModel): default_request: dict[str, Any] -def validate_mlx_video_request(request: VideoGenerationRequest) -> None: +def validate_mlx_video_request(request: VideoGenerationRequest, *, model_path: str = PREVIEW_MODEL) -> None: """Reject unsupported inputs before fetching media or creating a job.""" allowed = { "model", @@ -76,9 +96,10 @@ def validate_mlx_video_request(request: VideoGenerationRequest) -> None: raise ValueError("FastH3 MLX requires guidance_scale=1.") if request.negative_prompt not in (None, ""): raise ValueError("FastH3 MLX does not use a negative prompt.") - if request.num_inference_steps not in (None, 5): - raise ValueError( - "FastH3 MLX serving uses five sigma points and four DiT forwards; num_inference_steps must be 5.") + expected = mlx_http_steps(model_path) + if request.num_inference_steps not in (None, expected): + raise ValueError(f"This FastH3 MLX server uses {expected} sigma points " + f"({HTTP_STEPS_TO_FORWARDS[expected]} transformer forwards).") if request.seed is not None and not 0 <= request.seed <= 2**32 - 1: raise ValueError("H3 MLX seed must be between 0 and 4294967295.") @@ -87,6 +108,10 @@ class MLXH3Generator: """Keep one pipeline on one MLX thread; preserve its phase-memory policy.""" def __init__(self, config: MLXGeneratorConfig) -> None: + self._model_path = config.model_path + self._vsa = config.vsa + self._vsa_sparsity = config.vsa_sparsity + self._vsa_tile_size = config.vsa_tile_size self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="h3-mlx") try: self._pipeline = self._worker.submit(self._load, config).result() @@ -114,15 +139,21 @@ class MLXH3Generator: def _generate(self, request: GenerationRequest) -> dict[str, Any]: started = time.perf_counter() - result = self._pipeline.generate( - request.prompt, - output_path=request.output.output_path, - width=request.sampling.width, - height=request.sampling.height, - num_frames=request.sampling.num_frames, - seed=request.sampling.seed, - num_steps=4, - ) + generate_kwargs: dict[str, Any] = { + "output_path": request.output.output_path, + "width": request.sampling.width, + "height": request.sampling.height, + "num_frames": request.sampling.num_frames, + "seed": request.sampling.seed, + "num_steps": mlx_num_steps(request.sampling.num_inference_steps, model_path=self._model_path), + } + if self._vsa: + generate_kwargs.update( + vsa=True, + vsa_sparsity=self._vsa_sparsity, + vsa_tile_size=self._vsa_tile_size, + ) + result = self._pipeline.generate(request.prompt, **generate_kwargs) # Do not retain decoded frames/waveforms or label phase peaks as total RAM. return {"video_path": str(result.video_path), "generation_time": time.perf_counter() - started} @@ -157,9 +188,10 @@ def create_mlx_app(config: MLXServeConfig): required = {"width", "height", "num_frames", "fps", "seed", "num_inference_steps", "guidance_scale"} if required - set(explicit): raise ValueError("MLX default_request must set: " + ", ".join(sorted(required - set(explicit)))) - validate_mlx_video_request(VideoGenerationRequest(prompt="validate config", **explicit)) + validate_mlx_video_request(VideoGenerationRequest(prompt="validate config", **explicit), + model_path=config.generator.model_path) # Transport admission uses the registered H3 family, not CUDA engine options. - args = SimpleNamespace(model_path=MODEL, + args = SimpleNamespace(model_path=config.generator.model_path, lora_path=None, lora_nickname="default", lora_strength=1.0, @@ -172,13 +204,17 @@ def create_mlx_app(config: MLXServeConfig): served_model_name=config.server.served_model_name, output_dir=config.server.output_dir, default_request=request) + + def video_request_validator(request: VideoGenerationRequest) -> None: + validate_mlx_video_request(request, model_path=config.generator.model_path) + return create_app( args, config.server.output_dir, request, config.server.served_model_name, generator_factory=lambda: MLXH3Generator(config.generator), - video_request_validator=validate_mlx_video_request, + video_request_validator=video_request_validator, runtime="mlx", ) diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 755f6cef7..6e150a2f3 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -969,6 +969,7 @@ def _register_configs() -> None: hf_model_paths=[ "MiniMaxAI/MiniMax-H3", "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2", + "FastVideo/FastVideo-FastH3-8-Step-V2", ], model_detectors=[ lambda path: any(token in path.lower() for token in ( diff --git a/fastvideo/tests/entrypoints/test_openai_video_client.py b/fastvideo/tests/entrypoints/test_openai_video_client.py index b0ecf0446..3e48b8e64 100644 --- a/fastvideo/tests/entrypoints/test_openai_video_client.py +++ b/fastvideo/tests/entrypoints/test_openai_video_client.py @@ -23,7 +23,8 @@ from fastvideo.entrypoints.cli.inference_config import build_serve_config from fastvideo.entrypoints.openai.api_server import create_app from fastvideo.entrypoints.openai.stores import VIDEO_STORE from fastvideo.api.compat import normalize_generation_request -from fastvideo.entrypoints.openai.mlx_server import MLXH3Generator, create_mlx_app, load_config, validate_mlx_video_request +from fastvideo.entrypoints.openai.mlx_server import (EIGHT_STEP_MODEL, MLXH3Generator, PREVIEW_MODEL, create_mlx_app, + load_config, mlx_num_steps, validate_mlx_video_request) from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest ROOT = Path(__file__).resolve().parents[3] @@ -248,6 +249,27 @@ def test_mlx_rejects_unsupported_options_before_generation(fields): validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", **fields)) +def test_mlx_preview_and_8step_configs_map_sigma_points_to_forwards(): + assert mlx_num_steps(5, model_path=PREVIEW_MODEL) == 4 + assert mlx_num_steps(9, model_path=EIGHT_STEP_MODEL) == 8 + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=5), model_path=PREVIEW_MODEL) + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=9), + model_path=EIGHT_STEP_MODEL) + with pytest.raises(ValueError, match="9 sigma points"): + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=5), + model_path=EIGHT_STEP_MODEL) + preview = load_config(str(ROOT / "examples/serving/mlx_fasth3.yaml")) + eight = load_config(str(ROOT / "examples/serving/mlx_fasth3_8step.yaml")) + assert preview.generator.model_path.endswith("FastH3-Preview-v0.2") + assert preview.generator.vsa is False + assert eight.generator.model_path.endswith("FastH3-8-Step-V2") + assert eight.generator.vsa is True + assert eight.generator.vsa_sparsity == 0.8 + assert eight.default_request["sampling"]["num_inference_steps"] == 9 + assert mlx_num_steps(eight.default_request["sampling"]["num_inference_steps"], + model_path=eight.generator.model_path) == 8 + + def test_mlx_http_rejects_reference_media_and_image_routes(local_server): client, generator, base_url = local_server if generator.runtime != "mlx": diff --git a/mkdocs.yml b/mkdocs.yml index 80e390105..81c084016 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -151,7 +151,7 @@ nav: - Home: index.md - Getting Started: - Installation: getting_started/installation.md - - Apple Silicon FastWan: getting_started/installation/mps.md + - Apple Silicon (MLX): getting_started/installation/mlx.md - NVIDIA DGX Spark: getting_started/installation/spark.md - Pair two DGX Sparks: getting_started/installation/spark_pair.md - Quick Start: getting_started/quick_start.md diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index 0df48853d..833c2d869 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -84,6 +84,54 @@ def test_mlx_serving_profile_has_native_launcher_and_no_invented_memory(): assert client["code"] == (ROOT / client["source"]).read_text() +def test_fasth3_8step_cuda_and_mlx_share_the_openai_client(): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + cuda = next(item for item in recipes if item["id"] == "fasth3-8step-v2-cuda") + mlx = next(item for item in recipes if item["id"] == "fasth3-8step-v2-mlx") + assert cuda["group"] == mlx["group"] == "fasth3-8step-v2" + cuda_profile = cookbook_serving_profile(cuda) + mlx_profile = cookbook_serving_profile(mlx) + assert cuda_profile["model"] == mlx_profile["model"] == "fasth3" + assert cuda_profile["sampling"]["num_inference_steps"] == 9 + assert mlx_profile["sampling"]["num_inference_steps"] == 9 + assert "openai_fasth3_8step.yaml" in cuda_profile["command"] + assert mlx_profile["command"] == ( + "python -m fastvideo.entrypoints.openai.mlx_server --config examples/serving/mlx_fasth3_8step.yaml" + ) + assert mlx_profile["prepare"] in mlx["command"] + assert "--include-vsa" in mlx_profile["prepare"] + preview = cookbook_serving_profile(next(item for item in recipes if item["id"] == "fasth3-preview-cuda")) + assert cuda_profile["clients"]["python"]["code"] == preview["clients"]["python"]["code"] + + +def test_mlx_preview_and_8step_yaml_map_sigma_points_to_forwards(): + from fastvideo.entrypoints.openai.mlx_server import (EIGHT_STEP_MODEL, PREVIEW_MODEL, load_config, mlx_num_steps, + validate_mlx_video_request) + from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest + from fastvideo.registry import get_preset_selection + + assert mlx_num_steps(5, model_path=PREVIEW_MODEL) == 4 + assert mlx_num_steps(None, model_path=PREVIEW_MODEL) == 4 + assert mlx_num_steps(9, model_path=EIGHT_STEP_MODEL) == 8 + with pytest.raises(ValueError, match="5 sigma points"): + mlx_num_steps(9, model_path=PREVIEW_MODEL) + with pytest.raises(ValueError, match="9 sigma points"): + mlx_num_steps(5, model_path=EIGHT_STEP_MODEL) + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=5), model_path=PREVIEW_MODEL) + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=9), + model_path=EIGHT_STEP_MODEL) + with pytest.raises(ValueError, match="9 sigma points"): + validate_mlx_video_request(VideoGenerationRequest(prompt="a fox", num_inference_steps=5), + model_path=EIGHT_STEP_MODEL) + preview = load_config(str(ROOT / "examples/serving/mlx_fasth3.yaml")) + eight = load_config(str(ROOT / "examples/serving/mlx_fasth3_8step.yaml")) + assert preview.generator.vsa is False + assert eight.generator.vsa is True + assert eight.generator.vsa_sparsity == 0.8 + assert eight.default_request["sampling"]["num_inference_steps"] == 9 + assert get_preset_selection(EIGHT_STEP_MODEL) == ("minimax_h3_t2va", "minimax_h3") + + @pytest.mark.parametrize("field,value", [("model", "wrong-checkpoint"), ("hardware", {"platform": "mlx"})]) def test_cookbook_rejects_mismatched_serving_recipe(field, value): recipe = copy.deepcopy(serving_recipe()) From 02f1ce11aee39fbbfefebc191f1476b843ffd6aa Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Tue, 29 Sep 2026 14:36:16 -0700 Subject: [PATCH 040/151] [misc] Move environment var reads into env registry and rename unprefixed variables (#1897) --- .agents/skills/env-var-conventions/SKILL.md | 5 + docs/contributing/activation_trace.md | 2 +- docs/contributing/env_vars.md | 184 ++++++++---- .../configs/fine_tuning/kandinsky5/README.md | 2 +- .../fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml | 2 +- .../train/configs/fine_tuning/ltx2_3/t2v.yaml | 2 +- .../train/configs/overfit_ltx2_3_t2v.yaml | 4 +- fastvideo/attention/backends/flash_attn.py | 3 +- .../backends/video_sparse_attn_h3_probe.py | 4 +- fastvideo/attention/layer.py | 9 +- fastvideo/attention/selector.py | 5 +- fastvideo/benchmarks/mlx_fastwan_bench.py | 5 +- .../device_communicators/pyhccl_wrapper.py | 2 +- fastvideo/distributed/parallel_state.py | 7 +- fastvideo/entrypoints/streaming/worker.py | 4 +- fastvideo/entrypoints/video_generator.py | 21 +- fastvideo/envs.py | 274 ++++++++++++++---- fastvideo/eval/__init__.py | 4 +- fastvideo/eval/datasets/physics_iq.py | 5 +- fastvideo/eval/datasets/vbench.py | 8 +- .../metrics/audio/frechet_distance/metric.py | 4 +- fastvideo/eval/metrics/common/fvd/metric.py | 4 +- .../judge/third_person_separation/metric.py | 7 +- fastvideo/eval/metrics/vbench/scene/metric.py | 3 +- fastvideo/eval/models.py | 5 +- fastvideo/fastvideo_args.py | 2 + fastvideo/logger.py | 78 ++--- fastvideo/mlx_runtime/fastwan.py | 15 +- fastvideo/mlx_runtime/wan22.py | 4 +- fastvideo/models/encoders/gemma.py | 5 +- fastvideo/models/encoders/reason1.py | 4 +- .../benchmarks/benchmark_weight_loading.py | 13 +- fastvideo/models/loader/fsdp_load.py | 7 +- fastvideo/models/vaes/ltx2vae.py | 9 +- .../basic/ltx2/stages/ltx2_audio_decoding.py | 5 +- .../basic/ltx2/stages/ltx2_denoising.py | 6 +- .../basic/magi_human/magi_human_pipeline.py | 3 +- .../preprocess_kandinsky5_overfit.py | 31 +- .../preprocess/preprocess_ltx2_overfit.py | 57 ++-- .../preprocess_minimax_h3_overfit.py | 12 +- fastvideo/pipelines/stages/denoising.py | 4 +- .../pipelines/stages/latent_preparation.py | 4 +- fastvideo/tests/contract/test_env_policy.py | 206 ++++++------- .../test_e2e_kandinsky5_dmd_t2v_overfit.py | 6 +- .../test_e2e_ltx2_overfit_new_stack.py | 6 +- .../train/entrypoint/dcp_to_diffusers.py | 14 +- fastvideo/train/models/ltx2/ltx2.py | 2 +- .../self_forcing_distillation_pipeline.py | 3 +- fastvideo/utils.py | 10 +- fastvideo/worker/gpu_worker.py | 11 +- fastvideo/worker/ray_distributed_executor.py | 3 +- fastvideo/worker/ray_env.py | 46 +-- tests/local_tests/ltx2/README.md | 2 +- .../ltx2/test_ltx2_gemma_parity.py | 2 +- 54 files changed, 660 insertions(+), 485 deletions(-) diff --git a/.agents/skills/env-var-conventions/SKILL.md b/.agents/skills/env-var-conventions/SKILL.md index d793281a5..3d12b8f4f 100644 --- a/.agents/skills/env-var-conventions/SKILL.md +++ b/.agents/skills/env-var-conventions/SKILL.md @@ -34,10 +34,15 @@ on a summary here. 1. **Declare or edit the variable in `fastvideo/envs.py`.** - Pick the field type and category that the policy doc lists. - Write a description that states what the variable does and its units. + - To rename, keep the old name in `deprecated_names`. To remove, add the + name to `DEPRECATED_VARIABLES`. Update the uses in `examples/`, + `scripts/`, `docs/`, `apps/`, and the tests. 2. **Read the variable with `envs.NAME.get()` inside a function.** - In tests, change the value with `envs.NAME.override(value)`. - Do not call `os.environ`, `os.getenv`, or `monkeypatch.setenv` for a FastVideo variable. + - To set a variable that another tool reads, call `envs.set_external`, + `envs.setdefault_external`, or `envs.unset_external`. 3. **Regenerate the table in the policy doc.** - Run `python fastvideo/tests/contract/test_env_policy.py`. 4. **Run the contract test.** diff --git a/docs/contributing/activation_trace.md b/docs/contributing/activation_trace.md index 36fc88adc..c3397f901 100644 --- a/docs/contributing/activation_trace.md +++ b/docs/contributing/activation_trace.md @@ -30,7 +30,7 @@ FastVideo and the reference model first produce different numbers?" | General logging | `init_logger(__name__)` | | Per-stage timing | `FASTVIDEO_STAGE_LOGGING` | | Profiling kernel timings | `FASTVIDEO_TORCH_PROFILER_DIR` (see [Profiling](profiling.md)) | -| Function-call tracing | `FASTVIDEO_TRACE_FUNCTION` (heavy) | +| Function-call tracing | `fastvideo.logger.enable_trace_function_call()` (heavy) | ## Quickstart diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index a4dac3d29..8f91f1eca 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -11,14 +11,16 @@ test in the same pull request. 1. **Register every FastVideo variable in `fastvideo/envs.py`.** Each entry declares a type, a default, a category, and a description. Variables that other tools own (CUDA, NCCL, PyTorch, launchers) are not registered; code reads them directly with `os.environ.get("NAME")`, and the name must be in the external-variable allowlist - (`EXTERNAL_ALLOWLIST` in the contract test). + (`EXTERNAL_ALLOWLIST` in the contract test). When FastVideo sets such a variable for the other tool, it calls + `envs.set_external`, `envs.setdefault_external`, or `envs.unset_external`, and the name must be in + `EXTERNAL_WRITE_ALLOWLIST`. 2. **Read with `envs.NAME.get()`, write with `envs.NAME.set()`, and change a value in tests with `envs.NAME.override()`.** Each type has one parsing rule. A value that the rule rejects raises `fastvideo.envs.EnvVarError` instead of falling back to the default. 3. **Name FastVideo variables with the `FASTVIDEO_` prefix.** The second word states the purpose where one applies: `ENABLE_`, `DISABLE_`, `USE_`, `FORCE_`, `DEBUG_`, `TEST_`. 4. **Keep a renamed variable as a deprecated alias until the next minor release.** Setting the old name logs a - warning. Delete a variable that no code reads, and list it in a deprecation table so that setting it logs a + warning. Delete a variable that no code reads, and list it in `DEPRECATED_VARIABLES` so that setting it logs a warning. 5. **Give each setting one source: an argument or an environment variable.** Settings that users change per deployment are arguments (CLI or YAML). Expert switches, emergency off switches, and debugging and test switches @@ -85,6 +87,25 @@ with envs.FASTVIDEO_DEBUG_MY_STAGE.override(True): run_stage() ``` +## Rename or remove a variable + +To rename a variable, declare it under the new name and list the old name in `deprecated_names`: + +```python +FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS = EnvBool(True, + category="sampling", + doc="...", + deprecated_names=("LTX2_USE_DISTILLED_SIGMAS", )) +``` + +`get()` reads an old name only when the new name is unset, and logs a warning once. Update the uses of the old name +in `examples/`, `scripts/`, `docs/`, `apps/`, and the tests in the same pull request. Delete the old name in the next +minor release. + +To remove a variable that no code reads, delete its entry and add the name to `DEPRECATED_VARIABLES` in +`fastvideo/envs.py` with a reason. `FastVideoArgs` calls `envs.warn_deprecated_variables()`, which logs a warning for +each listed variable that is set. Delete the entry in the next minor release. + ## What the contract test checks The test parses every Python file under `fastvideo/`, including `fastvideo/tests/`, with Python's `ast` module. It @@ -100,20 +121,24 @@ It reports each violation as `: `: | | with a name built at runtime (``) | another tool owns, add it to | | | | `EXTERNAL_ALLOWLIST` with a reason. | | `write` | `os.environ[...] = ...`, `setdefault`, `pop`, `del`, | Pass an argument instead. In tests, use | -| | `os.putenv`, `os.unsetenv`, `monkeypatch.setenv`/`delenv` | `envs.NAME.override()`. | +| | `os.putenv`, `os.unsetenv`, `monkeypatch.setenv`/`delenv`, | `envs.NAME.override()`. For a variable that | +| | or an `envs.*_external` call with a name outside | another tool reads, call an | +| | `EXTERNAL_WRITE_ALLOWLIST` | `envs.*_external` helper and add the name to | +| | | `EXTERNAL_WRITE_ALLOWLIST` with a reason. | | `whole-environ` | `os.environ.copy()`, `dict(os.environ)`, iteration, | Read the specific variables that the code | | | `mock.patch.dict(os.environ, ...)`, `os.environ.update` | needs. | | `bare-field` | A registry field used without calling one of its methods, | Call `envs.NAME.get()`. | | | as in `envs.NAME == "auto"` or `getter = envs.NAME.get` | | | `import-time-read` | `envs.NAME.get()` outside a function | Move the read into the function that uses | | | | the value. | -| `prefix` | A registry entry without the `FASTVIDEO_` prefix | Rename the variable and keep the old name as | -| | | a deprecated alias. | -| `unread` | A registry entry that no code reads with `get()` or | Delete the variable and add it to the | -| | `is_set()` | deprecation table. | +| `prefix` | A registry entry without the `FASTVIDEO_` prefix | Rename the variable and keep the old name in | +| | | `deprecated_names`. | +| `unread` | A registry entry that no code reads with `get()` or | Delete the variable and add it to | +| | `is_set()` | `DEPRECATED_VARIABLES`. | -The test recognizes `os` imported under another name and `from os import environ, getenv`. Code that reaches the -environment through `importlib` or `getattr(os, "environ")` is left to code review. +The test recognizes `os` imported under another name, `from os import environ, getenv`, and a name held in a +module-level string constant. Code that reaches the environment through `importlib` or `getattr(os, "environ")` is +left to code review. The test also checks that every registry entry has a category from `envs.CATEGORIES` and a description, and that the table at the end of this page matches the registry. @@ -125,56 +150,93 @@ longer exists also fails the test, so the fixing pull request deletes its entry. ## Registered variables -| Variable | Type | Default | Category | Description | -| ---------------------------------------------- | ------------------------------ | --------------------------- | ----------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `FASTVIDEO_TARGET_DEVICE` | str | `cuda` | build | Target device of FastVideo: cuda, rocm, neuron, cpu, or openvino. | -| `MAX_JOBS` | str | unset | build | Maximum number of parallel compilation jobs. Defaults to the number of CPUs. | -| `NVCC_THREADS` | str | unset | build | Number of nvcc threads. When set, MAX_JOBS is reduced to avoid oversubscribing the CPU. | -| `FASTVIDEO_USE_PRECOMPILED` | bool | `0` | build | Use precompiled binaries (\*.so). | -| `CMAKE_BUILD_TYPE` | str | unset | build | CMake build type: Debug, Release, or RelWithDebInfo. | -| `VERBOSE` | bool | `0` | build | Print verbose logs during installation. | -| `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_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_RINGBUFFER_WARNING_INTERVAL` | int | `60` | distributed | Seconds between warnings while the ring buffer is full. | -| `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). | -| `HCCL_SO_PATH` | str | unset | distributed | Path to the HCCL library file on Ascend NPUs. | -| `FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S` | int | `60` | distributed | Timeout in seconds for each engine iteration. | -| `FASTVIDEO_WORKER_MULTIPROC_METHOD` | one of spawn, fork, forkserver | `spawn` | distributed | Multiprocessing start method for worker processes. | -| `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. | -| `LD_LIBRARY_PATH` | str | unset | external | Searched for the NCCL library when FASTVIDEO_NCCL_SO_PATH is unset. | -| `LOCAL_RANK` | int | `0` | external | Local rank of the process in a distributed run; selects the GPU device id. | -| `CUDA_VISIBLE_DEVICES` | str | unset | external | Visible devices in a distributed run. | -| `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_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | -| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | -| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | -| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | -| `FASTVIDEO_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_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_TEST_DYNAMO_FULLGRAPH_CAPTURE` | bool | `1` | debug | Enable Dynamo fullgraph capture. | -| `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_SERVER_DEV_MODE` | bool | `0` | debug | Run the server in development mode with extra debugging endpoints. | -| `FASTVIDEO_TRACE_FUNCTION` | bool | `0` | debug | Trace function calls. | -| `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_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. | +| 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_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_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_ATTENTION_BACKEND` | str | unset | attention | Attention backend, as an AttentionBackendEnum name such as TORCH_SDPA, FLASH_ATTN, VIDEO_SPARSE_ATTN, SAGE_ATTN, or SAGE_ATTN_THREE. FastVideoArgs uses it when FastVideoArgs.attention_backend is unset. | +| `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | +| `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | +| `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | +| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 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_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_MINIMAX_H3_FUSIONS` | str | `""` | performance | MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation. | +| `FASTVIDEO_FSDP2_AUTOWRAP` | bool | `0` | performance | FSDP2 shards modules by parameter count instead of the model's shard conditions. Not supported by self-forcing distillation. | +| `FASTVIDEO_FSDP2_MIN_PARAMS` | int | `10000000` | performance | Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards. | +| `FASTVIDEO_MLX_COMPILE` | bool | `0` | performance | Compile the MLX DiT forward with mx.compile. | +| `FASTVIDEO_MLX_FAST_NORM` | bool | `0` | performance | Use MLX fast normalization kernels. | +| `FASTVIDEO_MLX_DQ_GEMM` | str | `1` | performance | MLX dequantized GEMM for affine-quantized weights: 0 turns it off, 1 uses the measured minimum row count, and an integer sets the minimum row count. | +| `FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D` | bool | `1` | performance | LTX-2 VAE uses the channels_last_3d memory format. | +| `FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST` | bool | `1` | performance | LTX-2 audio decoding runs without CUDA autocast. Deprecated names: `LTX2_DISABLE_AUDIO_AUTOCAST`. | +| `FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION` | bool | `0` | performance | Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for the 4-step Klein model. | +| `FASTVIDEO_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_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`. | + +Variables that FastVideo no longer reads; setting one logs a warning: + +| Deprecated variable | Reason | +| ----------------------------------------- | ---------------- | +| `FASTVIDEO_TARGET_DEVICE` | no code reads it | +| `FASTVIDEO_USE_PRECOMPILED` | no code reads it | +| `FASTVIDEO_RINGBUFFER_WARNING_INTERVAL` | no code reads it | +| `FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S` | no code reads it | +| `FASTVIDEO_SERVER_DEV_MODE` | no code reads it | +| `FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE` | no code reads it | +| `FASTVIDEO_TRACE_FUNCTION` | no code reads it | diff --git a/examples/train/configs/fine_tuning/kandinsky5/README.md b/examples/train/configs/fine_tuning/kandinsky5/README.md index 5205128c0..5f8a80043 100644 --- a/examples/train/configs/fine_tuning/kandinsky5/README.md +++ b/examples/train/configs/fine_tuning/kandinsky5/README.md @@ -27,7 +27,7 @@ layout) with: ```bash # from the repo root; reads data/kandinsky5_overfit and writes # data/kandinsky5_overfit_preprocessed by default -- override with -# KANDINSKY5_OVERFIT_DATA_DIR / KANDINSKY5_OVERFIT_OUTPUT_DIR env vars +# FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR / FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR env vars python -m fastvideo.pipelines.preprocess.preprocess_kandinsky5_overfit ``` diff --git a/examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml b/examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml index 00993ef70..26a72b8af 100644 --- a/examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml +++ b/examples/train/configs/fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml @@ -32,7 +32,7 @@ # (see docs/training/data_preprocess.md; # fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal # LTX-2 reference to adapt — set -# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers). Data paths, +# FASTVIDEO_TEST_LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers). Data paths, # step counts, and learning rate below are placeholders you must adapt # to your workload. # diff --git a/examples/train/configs/fine_tuning/ltx2_3/t2v.yaml b/examples/train/configs/fine_tuning/ltx2_3/t2v.yaml index 71a5fbc12..0e69887de 100644 --- a/examples/train/configs/fine_tuning/ltx2_3/t2v.yaml +++ b/examples/train/configs/fine_tuning/ltx2_3/t2v.yaml @@ -12,7 +12,7 @@ # (see docs/training/data_preprocess.md; # fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py is a minimal # LTX-2 reference to adapt — set -# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers). +# FASTVIDEO_TEST_LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers). # # Run: # NUM_GPUS=4 bash examples/train/run.sh \ diff --git a/examples/train/configs/overfit_ltx2_3_t2v.yaml b/examples/train/configs/overfit_ltx2_3_t2v.yaml index b86e29533..5ffd47c17 100644 --- a/examples/train/configs/overfit_ltx2_3_t2v.yaml +++ b/examples/train/configs/overfit_ltx2_3_t2v.yaml @@ -8,8 +8,8 @@ # # Preprocess data first (writes data/ltx2_3_overfit_preprocessed): # CUDA_VISIBLE_DEVICES=0 \ -# LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers \ -# LTX2_OVERFIT_OUTPUT_DIR=data/ltx2_3_overfit_preprocessed \ +# FASTVIDEO_TEST_LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers \ +# FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR=data/ltx2_3_overfit_preprocessed \ # python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py # # Run: diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index 99e6f7c70..5ff725a40 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/fastvideo/attention/backends/flash_attn.py @@ -5,6 +5,7 @@ import torch import torch.nn.functional as F from dataclasses import dataclass +import fastvideo.envs as envs from fastvideo.attention.utils.flash_attn_default import ( fa_version, flash_attn_func_compilable, @@ -264,7 +265,7 @@ class FlashAttentionImpl(AttentionImpl): # keeps its previous semantics. nvfp4_fa4 = extra_impl_args.get("nvfp4_fa4") if nvfp4_fa4 is None: - nvfp4_fa4 = os.environ.get("FASTVIDEO_NVFP4_FA4", "0") == "1" + nvfp4_fa4 = envs.FASTVIDEO_NVFP4_FA4.get() self.nvfp4_fa4 = bool(nvfp4_fa4) if self.nvfp4_fa4: cap = torch.cuda.get_device_capability() diff --git a/fastvideo/attention/backends/video_sparse_attn_h3_probe.py b/fastvideo/attention/backends/video_sparse_attn_h3_probe.py index ff146e605..8084bd4fc 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3_probe.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3_probe.py @@ -23,12 +23,14 @@ import os import torch +import fastvideo.envs as envs + _FRACS = (0.05, 0.10, 0.125, 0.25, 0.50, 0.75) _TRUE_ROWS = 128 # sampled query rows per layer for the token-true check def probe_enabled() -> str | None: - return os.environ.get("FASTVIDEO_H3_VSA_PROBE") or None + return envs.FASTVIDEO_H3_VSA_PROBE.get() or None @torch.no_grad() diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index ac53492b7..ff499e452 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -import os from functools import wraps import torch import torch.nn as nn +import fastvideo.envs as envs from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D) @@ -23,10 +23,7 @@ def _attention_compile_disabled() -> bool: ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0`` to let attention instances constructed under that environment be traced. """ - val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE") - if val is None: - return True - return val.strip().lower() not in ("0", "false", "no", "off", "") + return envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.get() def _attention_compile_explicitly_disabled() -> bool: @@ -35,7 +32,7 @@ def _attention_compile_explicitly_disabled() -> bool: Regional compile can override the historical default for one loaded transformer, but it must still honor an explicit debugging escape hatch. """ - return "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in os.environ and _attention_compile_disabled() + return envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.is_set() and _attention_compile_disabled() def _maybe_compiler_disable(fn): diff --git a/fastvideo/attention/selector.py b/fastvideo/attention/selector.py index 302dd4683..76202d342 100644 --- a/fastvideo/attention/selector.py +++ b/fastvideo/attention/selector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py -import os from collections.abc import Generator from contextlib import contextmanager from contextvars import ContextVar @@ -15,7 +14,7 @@ import fastvideo.envs as envs from fastvideo.attention.backends.abstract import AttentionBackend from fastvideo.logger import init_logger from fastvideo.platforms import AttentionBackendEnum -from fastvideo.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname +from fastvideo.utils import resolve_obj_by_qualname logger = init_logger(__name__) @@ -65,7 +64,7 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None: * _Backend enum value if an override is specified * None otherwise ''' - backend_name = os.environ.get(STR_BACKEND_ENV_VAR) + backend_name = envs.FASTVIDEO_ATTENTION_BACKEND.get() return (None if backend_name is None else backend_name_to_enum(backend_name)) diff --git a/fastvideo/benchmarks/mlx_fastwan_bench.py b/fastvideo/benchmarks/mlx_fastwan_bench.py index 15a096d3f..3e0bf8e54 100644 --- a/fastvideo/benchmarks/mlx_fastwan_bench.py +++ b/fastvideo/benchmarks/mlx_fastwan_bench.py @@ -46,6 +46,7 @@ from examples.inference.basic.mlx_wan_prompt_to_video import ( encode_prompt, make_rotary_embeddings, ) +import fastvideo.envs as envs from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, cleanup_mlx # The highest-fidelity cell; used as the default SSIM reference when no external @@ -524,8 +525,8 @@ def _generate_cell( "peak_gib": max(load_peak, denoise_peak) / (1024**3), "quantization": quantization or "none", "compute_dtype": base_dtype, - "compile": os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1", - "fast_norm": os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1", + "compile": envs.FASTVIDEO_MLX_COMPILE.get(), + "fast_norm": envs.FASTVIDEO_MLX_FAST_NORM.get(), "mlx_memory_limit_gib": args.mlx_memory_limit_gib, "mlx_cache_limit_gib": args.mlx_cache_limit_gib, "mlx_disable_cache": args.mlx_disable_cache, diff --git a/fastvideo/distributed/device_communicators/pyhccl_wrapper.py b/fastvideo/distributed/device_communicators/pyhccl_wrapper.py index dfc98a5dd..a582fb83b 100644 --- a/fastvideo/distributed/device_communicators/pyhccl_wrapper.py +++ b/fastvideo/distributed/device_communicators/pyhccl_wrapper.py @@ -144,7 +144,7 @@ class HCCLLibrary: "Otherwise, the hccl library might not exist, be corrupted " "or it does not support the current platform %s. " "If you already have the library, please set the " - "environment variable HCCL_SO_PATH" + "environment variable FASTVIDEO_HCCL_SO_PATH" " to point to the correct hccl library path.", so_file, platform.platform()) raise e diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 49b2e4f34..dcf150dd1 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -40,7 +40,6 @@ import torch.distributed import torch.distributed as dist from torch.distributed import Backend, ProcessGroup, ReduceOp -import fastvideo.envs as envs from fastvideo.distributed.device_communicators.base_device_communicator import (DeviceCommunicatorBase) from fastvideo.distributed.device_communicators.cpu_communicator import (CpuCommunicator) from fastvideo.distributed.utils import StatelessProcessGroup @@ -759,7 +758,7 @@ def init_distributed_environment( if local_rank == -1: # local rank not set, this usually happens in single-node # setting, where we can use rank as local rank - local_rank = envs.LOCAL_RANK.get() if distributed_init_method == "env://" else rank + local_rank = int(os.environ.get("LOCAL_RANK", "0")) if distributed_init_method == "env://" else rank global _WORLD if _WORLD is None: ranks = list(range(torch.distributed.get_world_size())) @@ -884,9 +883,9 @@ def get_local_torch_device() -> torch.device: """Return the torch device for the current rank.""" from fastvideo.platforms import current_platform if current_platform.is_npu(): - device = torch.device(f"npu:{envs.LOCAL_RANK.get()}") + device = torch.device(f"npu:{int(os.environ.get('LOCAL_RANK', '0'))}") elif current_platform.is_cuda_alike() or current_platform.is_cuda(): - device = torch.device(f"cuda:{envs.LOCAL_RANK.get()}") + device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") else: device = torch.device("mps") return device diff --git a/fastvideo/entrypoints/streaming/worker.py b/fastvideo/entrypoints/streaming/worker.py index 9cfe1e6bb..d90ddb2a3 100644 --- a/fastvideo/entrypoints/streaming/worker.py +++ b/fastvideo/entrypoints/streaming/worker.py @@ -50,9 +50,9 @@ def worker_main( ``VideoGenerator`` construction + generation happens here, not in the parent's event loop. """ - import os + import fastvideo.envs as envs - os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + envs.set_external("CUDA_VISIBLE_DEVICES", str(gpu_id)) try: from fastvideo import VideoGenerator diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 410f67c54..2f8468a68 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -26,6 +26,7 @@ import torch import torchvision from einops import rearrange +import fastvideo.envs as envs from fastvideo.api.compat import ( REQUEST_BATCH_EXTRA_PASSTHROUGH_FIELDS, expand_request_prompt_batch, @@ -200,7 +201,7 @@ class VideoGenerator: if kwargs.pop("nvfp4_fa4", False): import os os.environ["FASTVIDEO_NVFP4_FA4"] = "1" - os.environ.setdefault("CUTE_DSL_ENABLE_TVM_FFI", "1") + envs.setdefault_external("CUTE_DSL_ENABLE_TVM_FFI", "1") typed_config = kwargs.pop("config", None) if typed_config is not None: if model_path is not None: @@ -1191,7 +1192,7 @@ class VideoGenerator: sample_rate: int, ) -> bool: """Encode video+audio using ffmpeg via rawvideo stdin + WAV input.""" - ffmpeg_bin = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) + ffmpeg_bin = shutil.which(envs.FASTVIDEO_FFMPEG_BIN.get()) if ffmpeg_bin is None: logger.warning("ffmpeg not found; cannot use ffmpeg pipe save.") return False @@ -1201,7 +1202,7 @@ class VideoGenerator: height = int(frames[0].shape[0]) width = int(frames[0].shape[1]) - codec = os.getenv("FASTVIDEO_VIDEO_CODEC", "libx264") + codec = envs.FASTVIDEO_VIDEO_CODEC.get() try: with tempfile.TemporaryDirectory() as tmpdir: @@ -1236,23 +1237,23 @@ class VideoGenerator: if codec.endswith("_nvenc"): nvenc_options = [ - ("preset", os.getenv("FASTVIDEO_NVENC_PRESET", "p1")), - ("tune", os.getenv("FASTVIDEO_NVENC_TUNE", "ull")), - ("rc", os.getenv("FASTVIDEO_NVENC_RC", "constqp")), - ("qp", os.getenv("FASTVIDEO_NVENC_QP", "28")), - ("bf", os.getenv("FASTVIDEO_NVENC_BF", "0")), + ("preset", envs.FASTVIDEO_NVENC_PRESET.get()), + ("tune", envs.FASTVIDEO_NVENC_TUNE.get()), + ("rc", envs.FASTVIDEO_NVENC_RC.get()), + ("qp", envs.FASTVIDEO_NVENC_QP.get()), + ("bf", envs.FASTVIDEO_NVENC_BF.get()), ] for option_name, option_value in nvenc_options: if cls._ffmpeg_encoder_supports_option(ffmpeg_bin, codec, option_name): cmd += [f"-{option_name}", option_value] else: - cmd += ["-preset", os.getenv("FASTVIDEO_X264_PRESET", "ultrafast")] + cmd += ["-preset", envs.FASTVIDEO_X264_PRESET.get()] cmd += [ "-c:a", "aac", "-pix_fmt", - os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"), + envs.FASTVIDEO_OUTPUT_PIX_FMT.get(), # Audio and video decoders may produce different durations. # Preserve every video frame and pad/trim audio to match. "-af", diff --git a/fastvideo/envs.py b/fastvideo/envs.py index bfa8cd809..881b2f158 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -12,6 +12,7 @@ The policy for environment variables is in ``docs/contributing/env_vars.md``, and ``fastvideo/tests/contract/test_env_policy.py`` enforces it. """ +import logging import os from collections.abc import Callable, Iterator from contextlib import contextmanager @@ -24,12 +25,23 @@ S = TypeVar("S", str, str | None) POLICY_DOC = "docs/contributing/env_vars.md" # Allowed values of EnvField.category. -CATEGORIES = ("build", "path", "distributed", "external", "logging", "attention", "performance", "profiling", "debug", - "sampling") +CATEGORIES = ("path", "distributed", "logging", "attention", "performance", "output", "profiling", "debug", "sampling", + "eval", "test") _TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) _FALSE_VALUES = frozenset({"0", "false", "no", "off", ""}) +# fastvideo.logger imports this module, so warnings go through the standard +# logging module; the "fastvideo" logger configuration still applies. +_logger = logging.getLogger(__name__) +_warned_messages: set[str] = set() + + +def _warn_once(message: str) -> None: + if message not in _warned_messages: + _warned_messages.add(message) + _logger.warning(message) + class EnvVarError(ValueError): """An environment variable holds a value that its registered type rejects.""" @@ -39,16 +51,20 @@ class EnvField(Generic[T]): """One registered environment variable: its type, default, category, and description. ``default`` is either the value itself or a zero-argument function that - computes it on each read while the variable is unset. + computes it on each read while the variable is unset. ``deprecated_names`` + lists earlier names of a renamed variable; they are read, with a warning, + only when the variable itself is unset. """ type_name = "" - def __init__(self, default: T | Callable[[], T], *, category: str, doc: str) -> None: + def __init__(self, default: T | Callable[[], T], *, category: str, doc: str, + deprecated_names: tuple[str, ...] = ()) -> None: self.name = "" # Set by _register_fields() from the module attribute name. self.default = default self.category = category self.doc = doc + self.deprecated_names = deprecated_names def parse(self, raw: str) -> T: raise NotImplementedError @@ -56,40 +72,55 @@ class EnvField(Generic[T]): def format(self, value: T) -> str: return str(value) + def _read_raw(self) -> tuple[str, str] | None: + """Return (name, raw value) from the variable or its first set deprecated name.""" + raw = os.environ.get(self.name) + if raw is not None: + return self.name, raw + for old_name in self.deprecated_names: + raw = os.environ.get(old_name) + if raw is not None: + _warn_once(f"{old_name} is deprecated and will be removed in the next minor release; " + f"set {self.name} instead.") + return old_name, raw + return None + def get(self) -> T: """Return the parsed value, or the default when the variable is unset.""" - raw = os.environ.get(self.name) - if raw is None: + found = self._read_raw() + if found is None: return self.default() if callable(self.default) else self.default + name, raw = found try: return self.parse(raw) except ValueError as exc: - raise EnvVarError(f"Invalid value {raw!r} for {self.name}: {exc}. See {POLICY_DOC}.") from None + raise EnvVarError(f"Invalid value {raw!r} for {name}: {exc}. See {POLICY_DOC}.") from None def is_set(self) -> bool: - return self.name in os.environ + return any(name in os.environ for name in (self.name, *self.deprecated_names)) def set(self, value: T) -> None: os.environ[self.name] = self.format(value) def clear(self) -> None: - os.environ.pop(self.name, None) + for name in (self.name, *self.deprecated_names): + os.environ.pop(name, None) @contextmanager def override(self, value: T | None) -> Iterator[None]: - """Set the variable, or unset it when ``value`` is None, and restore the previous value on exit.""" - previous = os.environ.get(self.name) - if value is None: - self.clear() - else: + """Set the variable, or unset it when ``value`` is None, and restore the previous values on exit.""" + previous = {name: os.environ.get(name) for name in (self.name, *self.deprecated_names)} + self.clear() + if value is not None: self.set(value) try: yield finally: - if previous is None: - self.clear() - else: - os.environ[self.name] = previous + for name, old_value in previous.items(): + if old_value is None: + os.environ.pop(name, None) + else: + os.environ[name] = old_value def __bool__(self) -> bool: raise TypeError(f"Use envs.{self.name}.get() to read {self.name}.") @@ -171,20 +202,25 @@ def get_default_config_root() -> str: ) -# ================== Installation ================== +# Variables that other tools read (CUDA, NCCL, PyTorch, launchers) are written +# only through these helpers. The contract test checks each name against its +# write allowlist. + + +def set_external(name: str, value: str) -> None: + """Set an environment variable that another tool reads.""" + os.environ[name] = value + + +def setdefault_external(name: str, value: str) -> None: + """Set an environment variable that another tool reads, unless it is already set.""" + os.environ.setdefault(name, value) + + +def unset_external(name: str) -> None: + """Remove an environment variable that another tool reads.""" + os.environ.pop(name, None) -FASTVIDEO_TARGET_DEVICE = EnvStr("cuda", - category="build", - doc="Target device of FastVideo: cuda, rocm, neuron, cpu, or openvino.") -MAX_JOBS = EnvStr(None, - category="build", - doc="Maximum number of parallel compilation jobs. Defaults to the number of CPUs.") -NVCC_THREADS = EnvStr(None, - category="build", - doc="Number of nvcc threads. When set, MAX_JOBS is reduced to avoid oversubscribing the CPU.") -FASTVIDEO_USE_PRECOMPILED = EnvBool(False, category="build", doc="Use precompiled binaries (*.so).") -CMAKE_BUILD_TYPE = EnvStr(None, category="build", doc="CMake build type: Debug, Release, or RelWithDebInfo.") -VERBOSE = EnvBool(False, category="build", doc="Print verbose logs during installation.") # ================== Paths ================== @@ -198,6 +234,10 @@ FASTVIDEO_CACHE_ROOT = EnvPath( category="path", doc="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 = EnvStr(None, + category="path", + doc="Local path or Hugging Face id of Reason1 weights to load instead of the " + "checkpoint's own.") # ================== Distributed ================== @@ -214,18 +254,15 @@ FASTVIDEO_RAY_PER_WORKER_GPUS = EnvFloat( category="distributed", doc="GPUs per Ray worker. A fraction lets Ray schedule several actors on one GPU, so other actors can share " "the GPUs with FastVideo.") -FASTVIDEO_RINGBUFFER_WARNING_INTERVAL = EnvInt(60, - category="distributed", - doc="Seconds between warnings while the ring buffer is full.") FASTVIDEO_NCCL_SO_PATH = EnvStr( None, category="distributed", doc="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).") -HCCL_SO_PATH = EnvStr(None, category="distributed", doc="Path to the HCCL library file on Ascend NPUs.") -FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S = EnvInt(60, - category="distributed", - doc="Timeout in seconds for each engine iteration.") +FASTVIDEO_HCCL_SO_PATH = EnvStr(None, + category="distributed", + doc="Path to the HCCL library file on Ascend NPUs.", + deprecated_names=("HCCL_SO_PATH", )) FASTVIDEO_WORKER_MULTIPROC_METHOD = EnvChoice("spawn", choices=("spawn", "fork", "forkserver"), category="distributed", @@ -238,18 +275,6 @@ FASTVIDEO_ULYSSES_A2A = EnvChoice( "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.") -# ================== External variables ================== -# Variables that other tools set. They stay registered so that Ray copies them -# to its workers until the external-variable allowlist replaces them. - -LD_LIBRARY_PATH = EnvStr(None, - category="external", - doc="Searched for the NCCL library when FASTVIDEO_NCCL_SO_PATH is unset.") -LOCAL_RANK = EnvInt(0, - category="external", - doc="Local rank of the process in a distributed run; selects the GPU device id.") -CUDA_VISIBLE_DEVICES = EnvStr(None, category="external", doc="Visible devices in a distributed run.") - # ================== Logging ================== FASTVIDEO_CONFIGURE_LOGGING = EnvBool( @@ -285,6 +310,22 @@ FASTVIDEO_VSA_SM100A = EnvBool( category="attention", doc="VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. " "fastvideo-kernel reads the same variable with the same rule.") +FASTVIDEO_NVFP4_FA4 = EnvBool( + False, + category="attention", + doc="FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument " + "takes precedence.") +FASTVIDEO_DISABLE_ATTENTION_COMPILE = EnvBool( + True, + category="attention", + doc="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 = EnvInt(0, + category="attention", + doc="MLX FastWan windowed attention size in tokens. 0 uses full attention.") +FASTVIDEO_MLX_WINDOW_SINK = EnvInt(0, + category="attention", + doc="Number of sink tokens that MLX windowed attention always attends to.") # ================== Performance ================== @@ -316,7 +357,44 @@ FASTVIDEO_MINIMAX_H3_FUSIONS = EnvStr( category="performance", doc="MiniMax-H3 inference-only Triton fusions: all, 1, or a comma-separated subset of " "modulate,qknorm_rope,swiglu. Empty, 0, or none keeps the eager implementation.") -FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE = EnvBool(True, category="debug", doc="Enable Dynamo fullgraph capture.") +FASTVIDEO_FSDP2_AUTOWRAP = EnvBool(False, + category="performance", + doc="FSDP2 shards modules by parameter count instead of the model's shard " + "conditions. Not supported by self-forcing distillation.") +FASTVIDEO_FSDP2_MIN_PARAMS = EnvInt(10000000, + category="performance", + doc="Minimum parameter count of a module that FASTVIDEO_FSDP2_AUTOWRAP shards.") +FASTVIDEO_MLX_COMPILE = EnvBool(False, category="performance", doc="Compile the MLX DiT forward with mx.compile.") +FASTVIDEO_MLX_FAST_NORM = EnvBool(False, category="performance", doc="Use MLX fast normalization kernels.") +FASTVIDEO_MLX_DQ_GEMM = EnvStr( + "1", + category="performance", + doc="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 = EnvBool(True, + category="performance", + doc="LTX-2 VAE uses the channels_last_3d memory format.") +FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST = EnvBool(True, + category="performance", + doc="LTX-2 audio decoding runs without CUDA autocast.", + deprecated_names=("LTX2_DISABLE_AUDIO_AUTOCAST", )) +FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION = EnvBool( + False, + category="performance", + doc="Flux denoising disables reduced-precision reductions in bf16 matmuls, which tightens accumulation for " + "the 4-step Klein model.") + +# ================== Output encoding ================== + +FASTVIDEO_FFMPEG_BIN = EnvStr("ffmpeg", category="output", doc="ffmpeg executable used to save video with audio.") +FASTVIDEO_VIDEO_CODEC = EnvStr("libx264", category="output", doc="ffmpeg video codec for saved videos.") +FASTVIDEO_NVENC_PRESET = EnvStr("p1", category="output", doc="NVENC preset when the codec is an *_nvenc codec.") +FASTVIDEO_NVENC_TUNE = EnvStr("ull", category="output", doc="NVENC tune option.") +FASTVIDEO_NVENC_RC = EnvStr("constqp", category="output", doc="NVENC rate-control mode.") +FASTVIDEO_NVENC_QP = EnvStr("28", category="output", doc="NVENC quantization parameter.") +FASTVIDEO_NVENC_BF = EnvStr("0", category="output", doc="NVENC number of B-frames.") +FASTVIDEO_X264_PRESET = EnvStr("ultrafast", category="output", doc="x264 preset for non-NVENC codecs.") +FASTVIDEO_OUTPUT_PIX_FMT = EnvStr("yuv420p", category="output", doc="ffmpeg pixel format for saved videos.") # ================== Profiling ================== @@ -341,10 +419,6 @@ FASTVIDEO_TORCH_PROFILE_REGIONS = EnvStr( # ================== Debug ================== -FASTVIDEO_SERVER_DEV_MODE = EnvBool(False, - category="debug", - doc="Run the server in development mode with extra debugging endpoints.") -FASTVIDEO_TRACE_FUNCTION = EnvBool(False, category="debug", doc="Trace function calls.") FASTVIDEO_TRACE_ACTIVATIONS = EnvBool(False, category="debug", doc="Enable activation trace hooks.") FASTVIDEO_TRACE_LAYERS = EnvStr("", category="debug", doc="Regex filter for traced module names. Empty means all.") FASTVIDEO_TRACE_STATS = EnvStr("abs_mean,sum", @@ -354,6 +428,18 @@ FASTVIDEO_TRACE_OUTPUT = EnvStr("/tmp/fv_trace_.jsonl", category="debug", doc="JSONL path for activation traces. The literal is replaced at runtime.") FASTVIDEO_TRACE_STEPS = EnvStr("", category="debug", doc="Comma-separated denoising step indices. Empty means all.") +FASTVIDEO_H3_VSA_PROBE = EnvStr( + None, + category="debug", + doc="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 = EnvStr("", + category="debug", + doc="Log file for LTX-2 Gemma text-encoder hidden states, used by parity tests.", + deprecated_names=("LTX2_FASTVIDEO_GEMMA_LOG", )) +FASTVIDEO_COSMOS25_LOG_KNOBS = EnvBool(False, + category="debug", + doc="Log the Cosmos 2.5 latent-preparation conditioning inputs.") # ================== Sampling ================== @@ -385,6 +471,84 @@ FASTVIDEO_CFG_GATE_STEP = EnvFloat( category="sampling", doc="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 = EnvBool( + True, + category="sampling", + doc="LTX-2 uses the distilled sigma schedule when FastVideoArgs.ltx2_use_distilled_sigmas is also true.", + deprecated_names=("LTX2_USE_DISTILLED_SIGMAS", )) + +# ================== Evaluation ================== + +FASTVIDEO_EVAL_CACHE = EnvPath(lambda: os.path.join(FASTVIDEO_CACHE_ROOT.get(), "eval"), + category="eval", + doc="Cache directory for evaluation models and datasets. Defaults to " + "$FASTVIDEO_CACHE_ROOT/eval.") +FASTVIDEO_PHYSICS_IQ_BUCKET_URL = EnvStr("https://storage.googleapis.com/physics-iq-benchmark", + category="eval", + doc="Base URL of the Physics-IQ benchmark bucket.") +FASTVIDEO_VBENCH_FULL_INFO_JSON = EnvStr(None, + category="eval", + doc="Path to VBench_full_info.json, used instead of the vendored copy.", + deprecated_names=("VBENCH_FULL_INFO_JSON", )) +FASTVIDEO_FVD_REF_FEATURES = EnvStr(None, category="eval", doc="Cached reference-feature file for the FVD metric.") +FASTVIDEO_FAD_REF_FEATURES = EnvStr(None, + category="eval", + doc="Cached reference-feature file for the audio Frechet distance metric.") + +# ================== Tests ================== + +FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR = EnvStr("data/cats", + category="test", + doc="Raw data directory for preprocess_ltx2_overfit.py.", + deprecated_names=("LTX2_OVERFIT_DATA_DIR", )) +FASTVIDEO_TEST_LTX2_OVERFIT_CAPTION_JSON = EnvStr("videos2caption_1_sample.json", + category="test", + doc="Caption file, relative to the raw data directory.", + deprecated_names=("LTX2_OVERFIT_CAPTION_JSON", )) +FASTVIDEO_TEST_LTX2_OVERFIT_VIDEO_SUBDIR = EnvStr("video", + category="test", + doc="Video subdirectory, relative to the raw data directory.", + deprecated_names=("LTX2_OVERFIT_VIDEO_SUBDIR", )) +FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR = EnvStr("data/ltx2_overfit_preprocessed", + category="test", + doc="Output directory for preprocess_ltx2_overfit.py.", + deprecated_names=("LTX2_OVERFIT_OUTPUT_DIR", )) +FASTVIDEO_TEST_LTX2_OVERFIT_MODEL = EnvStr("FastVideo/LTX2-Distilled-Diffusers", + category="test", + doc="Model repository whose encoders preprocess_ltx2_overfit.py uses.", + deprecated_names=("LTX2_OVERFIT_MODEL", )) +FASTVIDEO_TEST_LTX2_OVERFIT_NUM_COPIES = EnvInt(4, + category="test", + doc="Number of copies of the overfit sample in the parquet file.", + deprecated_names=("LTX2_OVERFIT_NUM_COPIES", )) +FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR = EnvStr("data/kandinsky5_overfit", + category="test", + doc="Raw data directory for preprocess_kandinsky5_overfit.py.", + deprecated_names=("KANDINSKY5_OVERFIT_DATA_DIR", )) +FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR = EnvStr("data/kandinsky5_overfit_preprocessed", + category="test", + doc="Output directory for preprocess_kandinsky5_overfit.py.", + deprecated_names=("KANDINSKY5_OVERFIT_OUTPUT_DIR", )) + +# Variables that FastVideo no longer reads. Setting one logs a warning; delete +# the entries in the next minor release. +DEPRECATED_VARIABLES = { + "FASTVIDEO_TARGET_DEVICE": "no code reads it", + "FASTVIDEO_USE_PRECOMPILED": "no code reads it", + "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL": "no code reads it", + "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S": "no code reads it", + "FASTVIDEO_SERVER_DEV_MODE": "no code reads it", + "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE": "no code reads it", + "FASTVIDEO_TRACE_FUNCTION": "no code reads it", +} + + +def warn_deprecated_variables() -> None: + """Log a warning for each variable in DEPRECATED_VARIABLES that is set.""" + for name, reason in DEPRECATED_VARIABLES.items(): + if name in os.environ: + _warn_once(f"{name} is deprecated and has no effect ({reason}); it will be removed in the next minor " + "release.") def _register_fields() -> dict[str, EnvField]: diff --git a/fastvideo/eval/__init__.py b/fastvideo/eval/__init__.py index 17decfcfe..6c989bf04 100644 --- a/fastvideo/eval/__init__.py +++ b/fastvideo/eval/__init__.py @@ -1,3 +1,4 @@ +import fastvideo.envs as envs from fastvideo.eval.models import ensure_checkpoint, get_cache_dir @@ -14,9 +15,8 @@ def _redirect_third_party_caches() -> None: env vars at all; their callsites in metric.py files pass ``download_root=str(get_cache_dir() / "")`` directly. """ - import os root = get_cache_dir() - os.environ.setdefault("TORCH_HOME", str(root / "torch")) + envs.setdefault_external("TORCH_HOME", str(root / "torch")) _redirect_third_party_caches() diff --git a/fastvideo/eval/datasets/physics_iq.py b/fastvideo/eval/datasets/physics_iq.py index f5ca10619..a90314685 100644 --- a/fastvideo/eval/datasets/physics_iq.py +++ b/fastvideo/eval/datasets/physics_iq.py @@ -27,7 +27,6 @@ pre-downloaded copy) to opt out of network fetches. from __future__ import annotations import csv -import os from dataclasses import dataclass from pathlib import Path from urllib.request import urlretrieve @@ -35,6 +34,7 @@ from urllib.request import urlretrieve import cv2 import numpy as np +import fastvideo.envs as envs from fastvideo.eval.datasets.base import PromptDataset from fastvideo.eval.datasets.registry import register_dataset from fastvideo.eval.models import get_cache_dir @@ -59,11 +59,10 @@ _VENDORED_DESCRIPTIONS_CSV = (Path(__file__).resolve().parent.parent / "metrics" # Public DeepMind bucket; HTTPS-readable, no auth. Override via # ``FASTVIDEO_PHYSICS_IQ_BUCKET_URL`` (e.g. for an internal mirror). -_DEFAULT_BUCKET_URL = "https://storage.googleapis.com/physics-iq-benchmark" def _bucket_url() -> str: - return os.environ.get("FASTVIDEO_PHYSICS_IQ_BUCKET_URL", _DEFAULT_BUCKET_URL) + return envs.FASTVIDEO_PHYSICS_IQ_BUCKET_URL.get() def _default_dataset_root() -> Path: diff --git a/fastvideo/eval/datasets/vbench.py b/fastvideo/eval/datasets/vbench.py index 61a08b025..9fbb93ff1 100644 --- a/fastvideo/eval/datasets/vbench.py +++ b/fastvideo/eval/datasets/vbench.py @@ -7,9 +7,9 @@ keyed by dimension). from __future__ import annotations import json -import os from pathlib import Path +import fastvideo.envs as envs from fastvideo.eval.datasets.base import PromptDataset from fastvideo.eval.datasets.registry import register_dataset @@ -23,12 +23,12 @@ _FULL_INFO_REL = "fastvideo/third_party/eval/vbench/vbench/VBench_full_info.json def _locate_full_info() -> Path: - env = os.environ.get("VBENCH_FULL_INFO_JSON") + env = envs.FASTVIDEO_VBENCH_FULL_INFO_JSON.get() if env: p = Path(env) if p.is_file(): return p - raise FileNotFoundError(f"VBENCH_FULL_INFO_JSON={env} does not point at a file") + raise FileNotFoundError(f"FASTVIDEO_VBENCH_FULL_INFO_JSON={env} does not point at a file") here = Path(__file__).resolve() for ancestor in here.parents: candidate = ancestor / _FULL_INFO_REL @@ -38,7 +38,7 @@ def _locate_full_info() -> Path: break raise FileNotFoundError("Could not locate VBench_full_info.json. Initialize the upstream " "submodule (`git submodule update --init " - "fastvideo/third_party/eval/vbench`) or set VBENCH_FULL_INFO_JSON.") + "fastvideo/third_party/eval/vbench`) or set FASTVIDEO_VBENCH_FULL_INFO_JSON.") @register_dataset("vbench") diff --git a/fastvideo/eval/metrics/audio/frechet_distance/metric.py b/fastvideo/eval/metrics/audio/frechet_distance/metric.py index f8fc3de58..aa091dd7d 100644 --- a/fastvideo/eval/metrics/audio/frechet_distance/metric.py +++ b/fastvideo/eval/metrics/audio/frechet_distance/metric.py @@ -11,7 +11,6 @@ two finite embeddings. from __future__ import annotations -import os from pathlib import Path from typing import Any @@ -19,6 +18,7 @@ import numpy as np import torch from scipy import linalg +import fastvideo.envs as envs from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult @@ -94,7 +94,7 @@ class FrechetAudioDistanceMetric(BaseMetric): self._model: Any = None self._gen_buf: list[np.ndarray] = [] self._ref_buf: list[np.ndarray] = [] - self._cached_ref_path: str | None = os.environ.get(REF_FEATURES_ENV) + self._cached_ref_path: str | None = envs.FASTVIDEO_FAD_REF_FEATURES.get() self._cached_ref_mu: np.ndarray | None = None self._cached_ref_sigma: np.ndarray | None = None self._n_cached_ref: int = 0 diff --git a/fastvideo/eval/metrics/common/fvd/metric.py b/fastvideo/eval/metrics/common/fvd/metric.py index 42536c086..e32e3aede 100644 --- a/fastvideo/eval/metrics/common/fvd/metric.py +++ b/fastvideo/eval/metrics/common/fvd/metric.py @@ -75,13 +75,13 @@ import numpy as np import scipy.linalg import torch +import fastvideo.envs as envs from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.metrics.common.fvd.extractors import (_BaseExtractor, available_extractors, load_extractor) from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult _MIN_VIDEOS_WARN = 256 # below this FVD is unreliable -_REF_FEATURES_ENV = "FASTVIDEO_FVD_REF_FEATURES" def _gaussian_params(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: @@ -220,7 +220,7 @@ class FVDMetric(BaseMetric): # import. Precedence: kwarg > env-var > default. if self._cache_path_arg is not None: self.cache_path = os.path.expanduser(self._cache_path_arg) - elif env_path := os.environ.get(_REF_FEATURES_ENV): + elif env_path := envs.FASTVIDEO_FVD_REF_FEATURES.get(): self.cache_path = os.path.expanduser(env_path) else: self.cache_path = _default_cache_path(self._extractor_name) diff --git a/fastvideo/eval/metrics/judge/third_person_separation/metric.py b/fastvideo/eval/metrics/judge/third_person_separation/metric.py index d2feee2da..4692789b2 100644 --- a/fastvideo/eval/metrics/judge/third_person_separation/metric.py +++ b/fastvideo/eval/metrics/judge/third_person_separation/metric.py @@ -110,10 +110,9 @@ def _path_of(sample: dict, key: str) -> str | None: def _resolve_api_key() -> str: - for env in ("GEMINI_API_KEY", "GOOGLE_API_KEY"): - key = os.environ.get(env) - if key: - return key.strip() + key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") + if key: + return key.strip() token = Path("~/.gemini_token").expanduser() if token.is_file(): return token.read_text().strip() diff --git a/fastvideo/eval/metrics/vbench/scene/metric.py b/fastvideo/eval/metrics/vbench/scene/metric.py index 77c3fe4ab..615cf63c2 100644 --- a/fastvideo/eval/metrics/vbench/scene/metric.py +++ b/fastvideo/eval/metrics/vbench/scene/metric.py @@ -21,6 +21,7 @@ import tempfile import torch import torchvision.io +import fastvideo.envs as envs from fastvideo.eval.metrics.base import BaseMetric from fastvideo.eval.registry import register from fastvideo.eval.types import MetricResult @@ -58,7 +59,7 @@ class SceneMetric(BaseMetric): from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor # AVoCaDO uses Qwen2.5-Omni — large multimodal model - os.environ.setdefault("VIDEO_MAX_PIXELS", str(20070400)) # 512*28*28*50 + envs.setdefault_external("VIDEO_MAX_PIXELS", str(20070400)) # 512*28*28*50 self._model = Qwen2_5OmniForConditionalGeneration.from_pretrained( self._model_path, diff --git a/fastvideo/eval/models.py b/fastvideo/eval/models.py index 9adb57b3e..339df2270 100644 --- a/fastvideo/eval/models.py +++ b/fastvideo/eval/models.py @@ -44,10 +44,7 @@ def get_cache_dir() -> Path: ``cache_dir``, etc.), pass ``str(get_cache_dir() / "")`` so users get a single ``FASTVIDEO_EVAL_CACHE`` knob to redirect them all. """ - return Path(os.environ.get( - "FASTVIDEO_EVAL_CACHE", - os.path.join(envs.FASTVIDEO_CACHE_ROOT.get(), "eval"), - )) + return Path(envs.FASTVIDEO_EVAL_CACHE.get()) def ensure_checkpoint( diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 8cc5a60bf..8cb4ef920 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -361,6 +361,8 @@ class FastVideoArgs: if env_backend is not None and backend_name_to_enum(env_backend) is not None: self.attention_backend = env_backend self._fold_vae_parallel_env() + import fastvideo.envs as envs + envs.warn_deprecated_variables() self.check_fastvideo_args() def _fold_vae_parallel_env(self) -> None: diff --git a/fastvideo/logger.py b/fastvideo/logger.py index a930ae0e4..859c1f0df 100644 --- a/fastvideo/logger.py +++ b/fastvideo/logger.py @@ -15,11 +15,6 @@ from typing import Any, cast import fastvideo.envs as envs -FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING.get() -FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH.get() -FASTVIDEO_LOGGING_LEVEL = envs.FASTVIDEO_LOGGING_LEVEL.get() -FASTVIDEO_LOGGING_PREFIX = envs.FASTVIDEO_LOGGING_PREFIX.get() - RED = '\033[91m' GREEN = '\033[92m' RESET = '\033[0;0m' @@ -27,40 +22,43 @@ RESET = '\033[0;0m' _warned_local_main_process = False _warned_main_process = False -_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s.%(msecs)03d " +_FORMAT = ("%(levelname)s %(asctime)s.%(msecs)03d " "[%(filename)s:%(lineno)d] %(message)s") _DATE_FORMAT = "%m-%d %H:%M:%S" -DEFAULT_LOGGING_CONFIG = { - "formatters": { - "fastvideo": { - "class": "fastvideo.logging_utils.NewLineFormatter", - "datefmt": _DATE_FORMAT, - "format": _FORMAT, + +def _default_logging_config() -> dict[str, Any]: + """Build the default logging configuration from FASTVIDEO_LOGGING_PREFIX and FASTVIDEO_LOGGING_LEVEL.""" + return { + "formatters": { + "fastvideo": { + "class": "fastvideo.logging_utils.NewLineFormatter", + "datefmt": _DATE_FORMAT, + "format": envs.FASTVIDEO_LOGGING_PREFIX.get() + _FORMAT, + }, }, - }, - "handlers": { - "fastvideo": { - "class": "logging.StreamHandler", - "formatter": "fastvideo", - "level": FASTVIDEO_LOGGING_LEVEL, - "stream": "ext://sys.stdout", + "handlers": { + "fastvideo": { + "class": "logging.StreamHandler", + "formatter": "fastvideo", + "level": envs.FASTVIDEO_LOGGING_LEVEL.get(), + "stream": "ext://sys.stdout", + }, }, - }, - "loggers": { - "fastvideo": { + "loggers": { + "fastvideo": { + "handlers": ["fastvideo"], + "level": "DEBUG", + "propagate": False, + }, + }, + "root": { "handlers": ["fastvideo"], "level": "DEBUG", - "propagate": False, }, - }, - "root": { - "handlers": ["fastvideo"], - "level": "DEBUG", - }, - "version": 1, - "disable_existing_loggers": False -} + "version": 1, + "disable_existing_loggers": False + } @lru_cache @@ -180,20 +178,22 @@ class _FastvideoLogger(Logger): def _configure_fastvideo_root_logger() -> None: logging_config = dict[str, Any]() + configure_logging = envs.FASTVIDEO_CONFIGURE_LOGGING.get() + logging_config_path = envs.FASTVIDEO_LOGGING_CONFIG_PATH.get() - if not FASTVIDEO_CONFIGURE_LOGGING and FASTVIDEO_LOGGING_CONFIG_PATH: + if not configure_logging and logging_config_path: raise RuntimeError("FASTVIDEO_CONFIGURE_LOGGING evaluated to false, but " "FASTVIDEO_LOGGING_CONFIG_PATH was given. FASTVIDEO_LOGGING_CONFIG_PATH " "implies FASTVIDEO_CONFIGURE_LOGGING. Please enable " "FASTVIDEO_CONFIGURE_LOGGING or unset FASTVIDEO_LOGGING_CONFIG_PATH.") - if FASTVIDEO_CONFIGURE_LOGGING: - logging_config = DEFAULT_LOGGING_CONFIG + if configure_logging: + logging_config = _default_logging_config() - if FASTVIDEO_LOGGING_CONFIG_PATH: - if not path.exists(FASTVIDEO_LOGGING_CONFIG_PATH): - raise RuntimeError("Could not load logging config. File does not exist: %s", FASTVIDEO_LOGGING_CONFIG_PATH) - with open(FASTVIDEO_LOGGING_CONFIG_PATH, encoding="utf-8") as file: + if logging_config_path: + if not path.exists(logging_config_path): + raise RuntimeError("Could not load logging config. File does not exist: %s", logging_config_path) + with open(logging_config_path, encoding="utf-8") as file: custom_config = json.loads(file.read()) if not isinstance(custom_config, dict): @@ -286,7 +286,7 @@ def enable_trace_function_call(log_file_path: str, root_dir: str | None = None): Note that this call is thread-level, any threads calling this function will have the trace enabled. Other threads will not be affected. """ - logger.warning("FASTVIDEO_TRACE_FUNCTION is enabled. It will record every" + logger.warning("Function-call tracing is enabled. It will record every" " function executed by Python. This will slow down the code. It " "is suggested to be used for debugging hang or crashes only.") logger.info("Trace frame log is saved to %s", log_file_path) diff --git a/fastvideo/mlx_runtime/fastwan.py b/fastvideo/mlx_runtime/fastwan.py index affa6e389..d6a802b2b 100644 --- a/fastvideo/mlx_runtime/fastwan.py +++ b/fastvideo/mlx_runtime/fastwan.py @@ -6,7 +6,6 @@ from __future__ import annotations import json import math -import os import statistics import time from dataclasses import dataclass @@ -14,6 +13,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any from collections.abc import Callable +import fastvideo.envs as envs from fastvideo.logger import init_logger if TYPE_CHECKING: @@ -326,7 +326,7 @@ def dq_gemm_engaged() -> int: def affine_dq_gemm_min_m() -> int | None: - raw = os.environ.get("FASTVIDEO_MLX_DQ_GEMM", "1").strip().lower() + raw = envs.FASTVIDEO_MLX_DQ_GEMM.get().strip().lower() if raw in {"", "0", "off", "false", "no"}: return None if raw in {"1", "on", "true", "yes"}: @@ -400,9 +400,7 @@ def _use_fast_norm() -> bool: traffic) and benchmark the speedup. H3 uses its own fused RMSNorm default and does not consult this toggle. """ - import os - - return os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1" + return envs.FASTVIDEO_MLX_FAST_NORM.get() def layer_norm(x, weight=None, bias=None, eps: float = 1e-6): @@ -633,12 +631,12 @@ class MLXWanTransformerBlock: k_bh = key.transpose(0, 2, 1, 3) v_bh = value.transpose(0, 2, 1, 3) scale = self.head_dim**-0.5 - window = int(os.environ.get("FASTVIDEO_MLX_WINDOW", "0") or "0") + window = envs.FASTVIDEO_MLX_WINDOW.get() if window > 0: from fastvideo.mlx_runtime.windowed_attention import windowed_attention _warn_windowed_attention_once(window) - sink = int(os.environ.get("FASTVIDEO_MLX_WINDOW_SINK", "0") or "0") + sink = envs.FASTVIDEO_MLX_WINDOW_SINK.get() attn_output = windowed_attention(q_bh, k_bh, v_bh, window=window, sink=sink, scale=scale) else: attn_output = mx.fast.scaled_dot_product_attention(q_bh, k_bh, v_bh, scale=scale) @@ -693,7 +691,6 @@ class MLXWanDiT: *, compile: bool = False, ) -> None: - import os self.weights = weights self.blocks = blocks @@ -711,7 +708,7 @@ class MLXWanDiT: # good mx.compile target. Off by default so the eager path stays the # baseline; enable via constructor or FASTVIDEO_MLX_COMPILE=1 and verify # with the benchmark's SSIM ~= 1.0 check. - self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1" + self._enable_compile = compile or envs.FASTVIDEO_MLX_COMPILE.get() self._compiled_forward: Callable[..., Any] | None = None self._compiled_signature: tuple | None = None diff --git a/fastvideo/mlx_runtime/wan22.py b/fastvideo/mlx_runtime/wan22.py index 6171481a6..2fe0e59c7 100644 --- a/fastvideo/mlx_runtime/wan22.py +++ b/fastvideo/mlx_runtime/wan22.py @@ -23,6 +23,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any from collections.abc import Callable +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.mlx_runtime.fastwan import ( MLXWanT2VCrossAttention, @@ -132,7 +133,6 @@ class MLXWan22DiT: *, compile: bool = False, ) -> None: - import os self.weights = weights self.blocks = blocks @@ -144,7 +144,7 @@ class MLXWan22DiT: self.patch_size = tuple(config["patch_size"]) self.out_channels = int(config["out_channels"]) self.eps = float(config.get("eps", 1e-6)) - self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1" + self._enable_compile = compile or envs.FASTVIDEO_MLX_COMPILE.get() self._compiled_forward: Callable[..., Any] | None = None self._compiled_signature: tuple | None = None diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 76c440f13..5b36d1004 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/fastvideo/models/encoders/gemma.py @@ -9,6 +9,7 @@ import torch from torch import nn from transformers import AutoTokenizer, Gemma3ForConditionalGeneration +import fastvideo.envs as envs from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig from fastvideo.models.encoders.base import TextEncoder from fastvideo.models.dits.ltx2 import ( @@ -27,7 +28,7 @@ import math def _debug_gemma_log_line(message: str) -> None: - log_path = os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", "") + log_path = envs.FASTVIDEO_LTX2_GEMMA_LOG.get() if not log_path: return log_dir = os.path.dirname(log_path) @@ -459,7 +460,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder): padding_side: str, ) -> tuple[torch.Tensor, torch.Tensor]: encoded_text_features = torch.stack(hidden_states, dim=-1) - if os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", ""): + if envs.FASTVIDEO_LTX2_GEMMA_LOG.get(): for idx, layer in enumerate(hidden_states): _debug_gemma_log_line(f"fastvideo:gemma_hidden_state_{idx}" f":sum={layer.float().sum().item():.6f}") diff --git a/fastvideo/models/encoders/reason1.py b/fastvideo/models/encoders/reason1.py index 51f9f711c..1784b26e7 100644 --- a/fastvideo/models/encoders/reason1.py +++ b/fastvideo/models/encoders/reason1.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 """Reason1 (Qwen2.5-VL) text encoder.""" -import os from dataclasses import dataclass from collections.abc import Iterable import torch from transformers import AutoProcessor +import fastvideo.envs as envs from fastvideo.configs.models.encoders import BaseEncoderOutput, Reason1Config from fastvideo.logger import init_logger from fastvideo.models.encoders.base import TextEncoder @@ -89,7 +89,7 @@ class Reason1TextEncoder(TextEncoder): trust_remote_code=True, ) - weights_override = os.getenv("FASTVIDEO_REASON1_WEIGHTS_PATH") + weights_override = envs.FASTVIDEO_REASON1_WEIGHTS_PATH.get() if weights_override: self.secondary_weights = (_WeightsSource( model_or_path=weights_override, diff --git a/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py b/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py index 0ce6d64b8..6c5995c63 100644 --- a/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py +++ b/fastvideo/models/loader/benchmarks/benchmark_weight_loading.py @@ -21,6 +21,7 @@ import time import torch import torch.distributed as dist +import fastvideo.envs as envs from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory, get_node_group, @@ -121,14 +122,10 @@ def main() -> None: files = resolve_safetensors_files(load_path) # Set default env vars for single-GPU standalone mode (no torchrun) - if "MASTER_ADDR" not in os.environ: - os.environ["MASTER_ADDR"] = "localhost" - if "MASTER_PORT" not in os.environ: - os.environ["MASTER_PORT"] = "29500" - if "RANK" not in os.environ: - os.environ["RANK"] = "0" - if "WORLD_SIZE" not in os.environ: - os.environ["WORLD_SIZE"] = "1" + envs.setdefault_external("MASTER_ADDR", "localhost") + envs.setdefault_external("MASTER_PORT", "29500") + envs.setdefault_external("RANK", "0") + envs.setdefault_external("WORLD_SIZE", "1") maybe_init_distributed_environment_and_model_parallel(1, 1) rank = dist.get_rank() if dist.is_initialized() else 0 diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index f7c31d221..956850a50 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -5,7 +5,6 @@ # Copyright 2025 The FastVideo Authors. from __future__ import annotations -import os import contextlib import re from collections.abc import Callable, Generator @@ -392,7 +391,7 @@ def _regional_compile_unsupported_reason( resolved = getattr(config, "_resolved_attention_backend", None) resolved_name = getattr(resolved, "name", "") if resolved_name == "VIDEO_SPARSE_ATTN_H3": - if os.environ.get("FASTVIDEO_H3_VSA_PROBE"): + if envs.FASTVIDEO_H3_VSA_PROBE.get(): return ("FASTVIDEO_H3_VSA_PROBE records tensors and files from the VSA-H3 attention body, which " "regional fullgraph compile cannot capture; this model stays eager") if not envs.FASTVIDEO_VSA_SM100A.get(): @@ -508,7 +507,7 @@ def shard_model( ValueError: If no layer modules were sharded, indicating that no shard_condition was triggered. """ # Check if we should use size-based filtering - use_size_filtering = os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1" + use_size_filtering = envs.FASTVIDEO_FSDP2_AUTOWRAP.get() if not fsdp_shard_conditions: logger.warning("No FSDP shard conditions provided; nothing will be sharded.") @@ -543,7 +542,7 @@ def shard_model( if use_size_filtering: # Size-based filtering mode - min_params = int(os.environ.get("FASTVIDEO_FSDP2_MIN_PARAMS", "10000000")) + min_params = envs.FASTVIDEO_FSDP2_MIN_PARAMS.get() logger.info("Using size-based filtering with threshold: %.2fM", min_params / 1e6) for n, m in reversed(named_modules): diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index 8789d352a..8556c1bea 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -7,7 +7,6 @@ from __future__ import annotations import itertools import math -import os from dataclasses import dataclass, replace from enum import Enum from typing import Any, Callable, Iterator, List, NamedTuple, Tuple @@ -17,14 +16,10 @@ import torch.nn as nn import torch.nn.functional as F from einops import rearrange +import fastvideo.envs as envs from fastvideo.models.vaes.common import DiagonalGaussianDistribution -def _is_env_enabled(name: str, default: str = "") -> bool: - value = os.getenv(name, default) - return value.lower() in {"1", "true", "yes", "on"} - - # ============================================================================= # Enums # ============================================================================= @@ -1595,7 +1590,7 @@ class LTX2CausalVideoAutoencoder(nn.Module): self._use_tiling: bool = False self._use_channels_last_3d: bool = False - if _is_env_enabled("FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D", default="1"): + if envs.FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D.get(): self.enable_channels_last_3d() def _as_channels_last_3d(self, tensor: torch.Tensor) -> torch.Tensor: diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py index e191eec2b..e868e0010 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py @@ -5,10 +5,9 @@ Audio decoding stage for LTX-2 pipelines. from __future__ import annotations -import os - import torch +import fastvideo.envs as envs from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger @@ -55,7 +54,7 @@ class LTX2AudioDecodingStage(PipelineStage): self.vocoder = self.vocoder.to(device) audio_latents = audio_latents.to(device) - disable_autocast = os.getenv("LTX2_DISABLE_AUDIO_AUTOCAST", "1") == "1" + disable_autocast = envs.FASTVIDEO_LTX2_DISABLE_AUDIO_AUTOCAST.get() with torch.no_grad(), torch.autocast( device_type="cuda", dtype=audio_latents.dtype, diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py index 4cd9aa412..4fc98e5a5 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py @@ -8,12 +8,12 @@ from __future__ import annotations from contextlib import contextmanager from itertools import combinations import math -import os from pathlib import Path import torch from tqdm.auto import tqdm +import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionMetadataBuilder) from fastvideo.attention.selector import component_attention_backend from fastvideo.fastvideo_args import FastVideoArgs @@ -54,7 +54,7 @@ except ImportError: @contextmanager def _nvtx_range(name: str): - if os.getenv("FASTVIDEO_NVTX_PROFILE", "0") == "1" and torch.cuda.is_available(): + if envs.FASTVIDEO_NVTX_PROFILE.get() and torch.cuda.is_available(): torch.cuda.nvtx.range_push(name) try: yield @@ -245,7 +245,7 @@ class LTX2DenoisingStage(PipelineStage): else: # Use distilled hardcoded schedule (or subsets) when enabled. use_distilled_sigmas = (fastvideo_args.ltx2_use_distilled_sigmas - and os.getenv("LTX2_USE_DISTILLED_SIGMAS", "1") == "1") + and envs.FASTVIDEO_LTX2_USE_DISTILLED_SIGMAS.get()) max_distilled_steps = len(DISTILLED_SIGMA_VALUES) - 1 if use_distilled_sigmas and num_inference_steps <= max_distilled_steps: sigmas, distilled_indices = _distilled_subset_sigmas( diff --git a/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py index bac774807..d7f5aa02f 100644 --- a/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py +++ b/fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py @@ -29,6 +29,7 @@ from typing import Any from transformers import AutoTokenizer +import fastvideo.envs as envs from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig from fastvideo.configs.models.vaes import OobleckVAEConfig from fastvideo.fastvideo_args import FastVideoArgs @@ -71,7 +72,7 @@ def _ensure_hf_token_env() -> str | None: """ token = resolve_hf_token() if token: - os.environ.setdefault("HF_TOKEN", token) + envs.setdefault_external("HF_TOKEN", token) return token diff --git a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py index 357a34fef..1186b48b2 100644 --- a/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py @@ -26,7 +26,7 @@ Usage: Input/output roots default to ``data/kandinsky5_overfit`` / ``data/kandinsky5_overfit_preprocessed`` and can be overridden with the -``KANDINSKY5_OVERFIT_DATA_DIR`` / ``KANDINSKY5_OVERFIT_OUTPUT_DIR`` env +``FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR`` / ``FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR`` env vars (used by the nightly e2e test to keep its disposable roots separate from a user's real dataset at the defaults). """ @@ -34,9 +34,6 @@ from a user's real dataset at the defaults). import json import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29520") - import cv2 import numpy as np import pyarrow as pa @@ -44,6 +41,7 @@ import pyarrow.parquet as pq import torch from transformers import AutoTokenizer +import fastvideo.envs as envs from fastvideo.configs.pipelines.base import preprocess_text from fastvideo.configs.pipelines.kandinsky5 import ( Kandinsky5T2VConfig, @@ -70,11 +68,6 @@ MAX_WIDTH = 768 TRAIN_FPS = 24.0 MODEL_PATH = "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers" -# Overridable so automation (e.g. the nightly e2e test) can point at its own -# test-owned directories instead of clobbering the documented default paths -# a user may have populated with their real dataset. -DATA_DIR = os.environ.get("KANDINSKY5_OVERFIT_DATA_DIR", "data/kandinsky5_overfit") -OUTPUT_DIR = os.environ.get("KANDINSKY5_OVERFIT_OUTPUT_DIR", "data/kandinsky5_overfit_preprocessed") def load_video(path: str, num_frames: int, height: int, width: int) -> torch.Tensor: @@ -143,17 +136,25 @@ def main() -> None: # FastVideo's native text-encoder layers (e.g. CLIPAttention's # QKVParallelLinear) are tensor-parallel-aware and assert the TP process # group is initialized, even for a single-process/single-GPU run. + envs.setdefault_external("MASTER_ADDR", "localhost") + envs.setdefault_external("MASTER_PORT", "29520") maybe_init_distributed_environment_and_model_parallel(1, 1) + # Overridable so automation (e.g. the nightly e2e test) can point at its own + # test-owned directories instead of clobbering the documented default paths + # a user may have populated with their real dataset. + data_dir = envs.FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR.get() + output_dir = envs.FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR.get() + device = torch.device("cuda:0") model_path = maybe_download_model(MODEL_PATH) - os.makedirs(OUTPUT_DIR, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) # Load captions. Validate up front (before spending minutes loading the # VAE + two text encoders): an empty/malformed manifest would otherwise # only surface as an IndexError on records[0] at parquet-write time. - manifest_path = os.path.join(DATA_DIR, "videos2caption.json") + manifest_path = os.path.join(data_dir, "videos2caption.json") with open(manifest_path) as f: caption_data = json.load(f) if not isinstance(caption_data, list) or not caption_data: @@ -212,7 +213,7 @@ def main() -> None: for item in caption_data: video_name = item["path"] caption = get_caption(item) - video_path = os.path.join(DATA_DIR, "videos", video_name) + video_path = os.path.join(data_dir, "videos", video_name) print(f"\nProcessing: {video_name}") print(f" Caption: {caption[:80]}...") @@ -289,7 +290,7 @@ def main() -> None: for k in records[0]}, schema=pyarrow_schema_t2v, ) - output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet") + output_path = os.path.join(output_dir, "data_00000.parquet") pq.write_table(table, output_path) print(f"\nWrote {len(records)} records to {output_path}") @@ -297,12 +298,12 @@ def main() -> None: # Wrap in "data" key -- ValidationDataset expects field="data". # Use "caption" field -- ValidationDataset aliases it to "prompt". val_prompts = {"data": [{"caption": get_caption(item)} for item in caption_data]} - val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json") + val_path = os.path.join(output_dir, "validation_prompts.json") with open(val_path, "w") as f: json.dump(val_prompts, f, indent=2) print(f"Wrote validation prompts to {val_path}") - print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.") + print("\nDone! Use data_path: " + output_dir + " in training config.") if __name__ == "__main__": diff --git a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py index 2f55261b9..4eca9a29c 100644 --- a/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py @@ -32,6 +32,7 @@ import pyarrow as pa import pyarrow.parquet as pq import torch +import fastvideo.envs as envs from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v from fastvideo.utils import maybe_download_model, verify_model_config_and_directory @@ -41,25 +42,14 @@ MAX_HEIGHT = 480 # divisible by 32 (spatial compression) MAX_WIDTH = 832 TRAIN_FPS = 24.0 # matches the LTX-2 preset fps used at validation -DATA_DIR = os.environ.get("LTX2_OVERFIT_DATA_DIR", "data/cats") -CAPTION_JSON = os.environ.get("LTX2_OVERFIT_CAPTION_JSON", "videos2caption_1_sample.json") -VIDEO_SUBDIR = os.environ.get("LTX2_OVERFIT_VIDEO_SUBDIR", "video") -OUTPUT_DIR = os.environ.get("LTX2_OVERFIT_OUTPUT_DIR", "data/ltx2_overfit_preprocessed") -MODEL_REPO = os.environ.get("LTX2_OVERFIT_MODEL", "FastVideo/LTX2-Distilled-Diffusers") -# The train dataloader samples with drop_last=True across data-parallel -# groups, so the dataset must hold at least num_sp_groups * batch_size -# rows or every rank gets zero batches. Replicate the overfit sample so -# a 4-GPU FSDP run still sees one batch per rank. -NUM_COPIES = int(os.environ.get("LTX2_OVERFIT_NUM_COPIES", "4")) - def _init_single_process_distributed() -> None: """FastVideo component loaders expect an initialized distributed env.""" - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29511") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", "0") + envs.setdefault_external("MASTER_ADDR", "127.0.0.1") + envs.setdefault_external("MASTER_PORT", "29511") + envs.setdefault_external("RANK", "0") + envs.setdefault_external("WORLD_SIZE", "1") + envs.setdefault_external("LOCAL_RANK", "0") from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, ) maybe_init_distributed_environment_and_model_parallel(1, 1) @@ -106,16 +96,27 @@ def main() -> None: PipelineComponentLoader, ) from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig + data_dir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR.get() + caption_json = envs.FASTVIDEO_TEST_LTX2_OVERFIT_CAPTION_JSON.get() + video_subdir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_VIDEO_SUBDIR.get() + output_dir = envs.FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR.get() + model_repo = envs.FASTVIDEO_TEST_LTX2_OVERFIT_MODEL.get() + # The train dataloader samples with drop_last=True across data-parallel + # groups, so the dataset must hold at least num_sp_groups * batch_size + # rows or every rank gets zero batches. Replicate the overfit sample so + # a 4-GPU FSDP run still sees one batch per rank. + num_copies = envs.FASTVIDEO_TEST_LTX2_OVERFIT_NUM_COPIES.get() + device = torch.device("cuda:0") - model_path = maybe_download_model(MODEL_REPO) + model_path = maybe_download_model(model_repo) model_index = verify_model_config_and_directory(model_path) - os.makedirs(OUTPUT_DIR, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) # The map-style dataset caches parquet file metadata; a stale cache # next to a regenerated parquet can crash or serve old rows. - shutil.rmtree(os.path.join(OUTPUT_DIR, "map_style_cache"), ignore_errors=True) + shutil.rmtree(os.path.join(output_dir, "map_style_cache"), ignore_errors=True) - with open(os.path.join(DATA_DIR, CAPTION_JSON)) as f: + with open(os.path.join(data_dir, caption_json)) as f: caption_data = json.load(f) pipeline_config = LTX2T2VConfig() @@ -163,7 +164,7 @@ def main() -> None: video_name = item["path"] record_id = f"{idx:04d}_{video_name}" caption = item["cap"][0] if isinstance(item["cap"], list) else item["cap"] - video_path = os.path.join(DATA_DIR, VIDEO_SUBDIR, video_name) + video_path = os.path.join(data_dir, video_subdir, video_name) print(f"\nProcessing: {video_name}") print(f" Caption: {caption[:80]}...") @@ -213,7 +214,7 @@ def main() -> None: # Save the preprocessed clip so overfit tests can compare # validation output against the memorization target. import imageio - ref_path = os.path.join(OUTPUT_DIR, f"training_sample_{idx}.mp4") + ref_path = os.path.join(output_dir, f"training_sample_{idx}.mp4") with imageio.get_writer(ref_path, fps=TRAIN_FPS) as writer: for frame in frames_np: writer.append_data(frame) @@ -223,9 +224,9 @@ def main() -> None: del text_encoder, tokenizer, vae torch.cuda.empty_cache() - # Write parquet (replicated NUM_COPIES times; see comment at top) + # Write parquet (replicated num_copies times; see the comment above num_copies) replicated = [] - for copy_idx in range(max(1, NUM_COPIES)): + for copy_idx in range(max(1, num_copies)): for r in records: row = dict(r) row["id"] = f"{r['id']}_copy{copy_idx}" @@ -235,10 +236,10 @@ def main() -> None: for k in replicated[0]}, schema=pyarrow_schema_t2v, ) - output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet") + output_path = os.path.join(output_dir, "data_00000.parquet") pq.write_table(table, output_path) print(f"\nWrote {len(replicated)} records " - f"({len(records)} unique x {max(1, NUM_COPIES)} copies) to {output_path}") + f"({len(records)} unique x {max(1, num_copies)} copies) to {output_path}") # Write validation prompts for the validation callback val_prompts = { @@ -246,12 +247,12 @@ def main() -> None: "caption": (item["cap"][0] if isinstance(item["cap"], list) else item["cap"]), } for item in caption_data], } - val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json") + val_path = os.path.join(output_dir, "validation_prompts.json") with open(val_path, "w") as f: json.dump(val_prompts, f, indent=2) print(f"Wrote validation prompts to {val_path}") - print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.") + print("\nDone! Use data_path: " + output_dir + " in training config.") if __name__ == "__main__": diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py index 33a26ce99..9ef03db80 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py @@ -10,7 +10,6 @@ from __future__ import annotations import argparse import gc import json -import os from pathlib import Path import shutil from typing import Any @@ -20,6 +19,7 @@ import pyarrow as pa import pyarrow.parquet as pq import torch +import fastvideo.envs as envs from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2va from fastvideo.fastvideo_args import FastVideoArgs @@ -50,11 +50,11 @@ TRAINING_VIDEO_NAME = "1gGQy4nxyUo-Scene-016.mp4" def _init_single_process_distributed() -> None: """Initialize the one-rank process groups required by component loaders.""" - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29531") - os.environ.setdefault("RANK", "0") - os.environ.setdefault("WORLD_SIZE", "1") - os.environ.setdefault("LOCAL_RANK", "0") + envs.setdefault_external("MASTER_ADDR", "127.0.0.1") + envs.setdefault_external("MASTER_PORT", "29531") + envs.setdefault_external("RANK", "0") + envs.setdefault_external("WORLD_SIZE", "1") + envs.setdefault_external("LOCAL_RANK", "0") from fastvideo.distributed import maybe_init_distributed_environment_and_model_parallel maybe_init_distributed_environment_and_model_parallel(1, 1) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index f36509eb3..5c867c969 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -4,7 +4,6 @@ Denoising stage for diffusion pipelines. """ import inspect -import os import weakref from collections.abc import Iterable from dataclasses import dataclass @@ -141,8 +140,7 @@ class DenoisingStage(PipelineStage): # gates should graduate to arch-config declarations like the precision # policies above. _is_flux = (getattr(fastvideo_args.pipeline_config.dit_config, "prefix", "") == "Flux") - if _is_flux and os.getenv("FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION", - "").lower() in {"1", "true", "yes"}: + if _is_flux and envs.FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION.get(): # Gate 1: tighten bf16 matmul accumulation for the 4-step Klein model (opt-in via env var). torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False # Gate 2: Flux2 runs its bf16 transformer WITHOUT autocast — autocast perturbs long-sequence attention enough to break 4-step latent parity. diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index 5f60409de..c01e6c7f8 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -3,13 +3,13 @@ Latent preparation stage for diffusion pipelines. """ -import os from typing import Any import numpy as np import torch from diffusers.utils.torch_utils import randn_tensor +import fastvideo.envs as envs from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.image_processor import ImageProcessor @@ -475,7 +475,7 @@ class Cosmos25LatentPreparationStage(CosmosLatentPreparationStage): ) video = torch.cat([cond_video, padding], dim=2) - if os.environ.get("FASTVIDEO_COSMOS25_LOG_KNOBS", "0") in ("1", "true", "True"): + if envs.FASTVIDEO_COSMOS25_LOG_KNOBS.get(): logger.info( "[Cosmos2.5 latent_prep] using conditioning input: image=%s video_latent=%s video_path=%s | " "num_video_frames=%s num_cond_frames=%s num_frames=%s", diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py index 571e22228..99f882ed1 100644 --- a/fastvideo/tests/contract/test_env_policy.py +++ b/fastvideo/tests/contract/test_env_policy.py @@ -6,7 +6,9 @@ The test parses every Python file under fastvideo/ (except fastvideo/third_party and the registry fastvideo/envs.py) with ``ast`` and reports code that: - reads the environment directly for a name outside EXTERNAL_ALLOWLIST; -- writes the environment directly (os.environ, os.putenv, monkeypatch.setenv); +- writes the environment directly (os.environ, os.putenv, monkeypatch.setenv), + or through the envs.*_external helpers for a name outside + EXTERNAL_WRITE_ALLOWLIST; - uses the whole environment (os.environ.copy(), dict(os.environ), patch.dict); - uses a registry field without calling one of its methods (``envs.X == "a"``, ``getter = envs.X.get``); @@ -61,8 +63,33 @@ EXTERNAL_ALLOWLIST = { "PATH": "Executable search path.", "HF_TOKEN": "Token variable that huggingface_hub reads.", "HUGGING_FACE_HUB_TOKEN": "Token variable that huggingface_hub reads.", + "GEMINI_API_KEY": "Gemini API key for the judge metrics.", + "GOOGLE_API_KEY": "Google API key, which the Gemini SDK also accepts.", + "CEREBRAS_API_KEY": "Cerebras API key for the streaming prompt provider.", + "GROQ_API_KEY": "Groq API key for the streaming prompt provider.", + "RAY_USAGE_STATS_ENABLED": "Ray usage-statistics switch.", } +# Variables that FastVideo sets for other tools. Code writes them only through +# envs.set_external, envs.setdefault_external, or envs.unset_external, with a +# literal name. +EXTERNAL_WRITE_ALLOWLIST = { + "RANK": "torch.distributed rendezvous for single-process runs and workers.", + "LOCAL_RANK": "torch.distributed rendezvous for single-process runs and workers.", + "WORLD_SIZE": "torch.distributed rendezvous for single-process runs and workers.", + "MASTER_ADDR": "torch.distributed rendezvous for single-process runs.", + "MASTER_PORT": "torch.distributed rendezvous for single-process runs.", + "CUDA_VISIBLE_DEVICES": "Pins a streaming worker process to its GPU.", + "TORCH_NCCL_AVOID_RECORD_STREAMS": "Avoids NCCL record-stream memory growth in workers.", + "NCCL_ASYNC_ERROR_HANDLING": "Unset in workers because the value that Ray sets breaks graph building.", + "TORCH_HOME": "Points torch.hub at the evaluation cache.", + "HF_TOKEN": "Passes the resolved Hugging Face token to huggingface_hub.", + "VIDEO_MAX_PIXELS": "Pixel budget that qwen_vl_utils reads for Qwen2.5-Omni.", + "RAY_USAGE_STATS_ENABLED": "Turns off Ray usage statistics unless the user turned them on.", + "CUTE_DSL_ENABLE_TVM_FFI": "CUTLASS DSL switch that the NVFP4 FlashAttention-4 path needs.", +} +EXTERNAL_WRITE_HELPERS = {"set_external", "setdefault_external", "unset_external"} + # Methods that registry fields expose; ``get`` and ``is_set`` count as reads. REGISTRY_METHODS = {"get", "set", "override", "is_set", "clear"} REGISTRY_READ_METHODS = {"get", "is_set"} @@ -72,111 +99,17 @@ REGISTRY_READ_METHODS = {"get", "is_set"} # its violations are fixed; never add one. Kinds are described in # docs/contributing/env_vars.md. KNOWN_VIOLATIONS: dict[str, int] = { - 'fastvideo/attention/backends/flash_attn.py: read FASTVIDEO_NVFP4_FA4': 1, - 'fastvideo/attention/backends/video_sparse_attn_h3_probe.py: read FASTVIDEO_H3_VSA_PROBE': 1, - 'fastvideo/attention/layer.py: read FASTVIDEO_DISABLE_ATTENTION_COMPILE': 2, - 'fastvideo/attention/selector.py: read ': 1, 'fastvideo/attention/utils/flash_attn_default.py: import-time-read FASTVIDEO_FA4': 1, 'fastvideo/benchmarks/eval_metalfx_rife.py: write FASTVIDEO_MLX_COMPILE': 1, - 'fastvideo/benchmarks/mlx_fastwan_bench.py: read FASTVIDEO_MLX_COMPILE': 1, - 'fastvideo/benchmarks/mlx_fastwan_bench.py: read FASTVIDEO_MLX_FAST_NORM': 1, 'fastvideo/benchmarks/mlx_fastwan_bench.py: write FASTVIDEO_MLX_COMPILE': 1, 'fastvideo/distributed/device_communicators/cpu_communicator.py: read VLLM_DIST_IDENT': 1, 'fastvideo/entrypoints/cli/utils.py: whole-environ': 1, 'fastvideo/entrypoints/openai/api_server.py: write FASTVIDEO_STAGE_LOGGING': 1, - 'fastvideo/entrypoints/streaming/prompt/providers/cerebras.py: read ': 1, - 'fastvideo/entrypoints/streaming/prompt/providers/groq.py: read ': 1, - 'fastvideo/entrypoints/streaming/worker.py: write CUDA_VISIBLE_DEVICES': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_FFMPEG_BIN': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_BF': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_PRESET': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_QP': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_RC': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_NVENC_TUNE': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_OUTPUT_PIX_FMT': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_VIDEO_CODEC': 1, - 'fastvideo/entrypoints/video_generator.py: read FASTVIDEO_X264_PRESET': 1, - 'fastvideo/entrypoints/video_generator.py: write CUTE_DSL_ENABLE_TVM_FFI': 1, 'fastvideo/entrypoints/video_generator.py: write FASTVIDEO_NVFP4_FA4': 1, - 'fastvideo/envs.py: prefix CMAKE_BUILD_TYPE': 1, - 'fastvideo/envs.py: prefix CUDA_VISIBLE_DEVICES': 1, - 'fastvideo/envs.py: prefix HCCL_SO_PATH': 1, - 'fastvideo/envs.py: prefix LD_LIBRARY_PATH': 1, - 'fastvideo/envs.py: prefix LOCAL_RANK': 1, - 'fastvideo/envs.py: prefix MAX_JOBS': 1, - 'fastvideo/envs.py: prefix NVCC_THREADS': 1, - 'fastvideo/envs.py: prefix VERBOSE': 1, - 'fastvideo/envs.py: unread CMAKE_BUILD_TYPE': 1, - 'fastvideo/envs.py: unread CUDA_VISIBLE_DEVICES': 1, - 'fastvideo/envs.py: unread FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S': 1, - 'fastvideo/envs.py: unread FASTVIDEO_RINGBUFFER_WARNING_INTERVAL': 1, - 'fastvideo/envs.py: unread FASTVIDEO_SERVER_DEV_MODE': 1, - 'fastvideo/envs.py: unread FASTVIDEO_TARGET_DEVICE': 1, - 'fastvideo/envs.py: unread FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE': 1, - 'fastvideo/envs.py: unread FASTVIDEO_TRACE_FUNCTION': 1, - 'fastvideo/envs.py: unread FASTVIDEO_USE_PRECOMPILED': 1, - 'fastvideo/envs.py: unread LD_LIBRARY_PATH': 1, - 'fastvideo/envs.py: unread MAX_JOBS': 1, - 'fastvideo/envs.py: unread NVCC_THREADS': 1, - 'fastvideo/envs.py: unread VERBOSE': 1, - 'fastvideo/eval/__init__.py: write TORCH_HOME': 1, - 'fastvideo/eval/datasets/physics_iq.py: read FASTVIDEO_PHYSICS_IQ_BUCKET_URL': 1, - 'fastvideo/eval/datasets/vbench.py: read VBENCH_FULL_INFO_JSON': 1, - 'fastvideo/eval/metrics/audio/frechet_distance/metric.py: read ': 1, - 'fastvideo/eval/metrics/common/fvd/metric.py: read ': 1, - 'fastvideo/eval/metrics/judge/third_person_separation/metric.py: read ': 1, - 'fastvideo/eval/metrics/vbench/scene/metric.py: write VIDEO_MAX_PIXELS': 1, - 'fastvideo/eval/models.py: read FASTVIDEO_EVAL_CACHE': 1, - 'fastvideo/logger.py: import-time-read FASTVIDEO_CONFIGURE_LOGGING': 1, - 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_CONFIG_PATH': 1, - 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_LEVEL': 1, - 'fastvideo/logger.py: import-time-read FASTVIDEO_LOGGING_PREFIX': 1, - 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_COMPILE': 1, - 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_DQ_GEMM': 1, - 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_FAST_NORM': 1, - 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_WINDOW': 1, - 'fastvideo/mlx_runtime/fastwan.py: read FASTVIDEO_MLX_WINDOW_SINK': 1, 'fastvideo/mlx_runtime/memory.py: write ': 1, - 'fastvideo/mlx_runtime/wan22.py: read FASTVIDEO_MLX_COMPILE': 1, - 'fastvideo/models/encoders/gemma.py: read LTX2_FASTVIDEO_GEMMA_LOG': 2, - 'fastvideo/models/encoders/reason1.py: read FASTVIDEO_REASON1_WEIGHTS_PATH': 1, - 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write MASTER_ADDR': 1, - 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write MASTER_PORT': 1, - 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write RANK': 1, - 'fastvideo/models/loader/benchmarks/benchmark_weight_loading.py: write WORLD_SIZE': 1, - 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_FSDP2_AUTOWRAP': 1, - 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_FSDP2_MIN_PARAMS': 1, - 'fastvideo/models/loader/fsdp_load.py: read FASTVIDEO_H3_VSA_PROBE': 1, - 'fastvideo/models/vaes/ltx2vae.py: read ': 1, 'fastvideo/performance/hf_store.py: read HF_REPO_ID': 1, 'fastvideo/performance/hf_store.py: read PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS': 1, 'fastvideo/performance_dashboard/api.py: read PERFORMANCE_TRACKING_ROOT': 1, - 'fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py: read LTX2_DISABLE_AUDIO_AUTOCAST': 1, - 'fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py: read FASTVIDEO_NVTX_PROFILE': 1, - 'fastvideo/pipelines/basic/ltx2/stages/ltx2_denoising.py: read LTX2_USE_DISTILLED_SIGMAS': 1, - 'fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py: write HF_TOKEN': 1, - 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: read KANDINSKY5_OVERFIT_DATA_DIR': 1, - 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: read KANDINSKY5_OVERFIT_OUTPUT_DIR': 1, - 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: write MASTER_ADDR': 1, - 'fastvideo/pipelines/preprocess/preprocess_kandinsky5_overfit.py: write MASTER_PORT': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_CAPTION_JSON': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_DATA_DIR': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_MODEL': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_NUM_COPIES': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_OUTPUT_DIR': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: read LTX2_OVERFIT_VIDEO_SUBDIR': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write LOCAL_RANK': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write MASTER_ADDR': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write MASTER_PORT': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write RANK': 1, - 'fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py: write WORLD_SIZE': 1, - 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write LOCAL_RANK': 1, - 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write MASTER_ADDR': 1, - 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write MASTER_PORT': 1, - 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write RANK': 1, - 'fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py: write WORLD_SIZE': 1, - 'fastvideo/pipelines/stages/denoising.py: read FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION': 1, - 'fastvideo/pipelines/stages/latent_preparation.py: read FASTVIDEO_COSMOS25_LOG_KNOBS': 1, 'fastvideo/tests/api/test_attention_selector_resolution.py: write FASTVIDEO_ATTENTION_BACKEND': 17, 'fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py: write FASTVIDEO_FA4': 3, 'fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py: write FASTVIDEO_NVFP4_FA4': 4, @@ -345,8 +278,8 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_REPO': 2, 'fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py: write WANDB_MODE': 1, 'fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py: write WANDB_MODE': 1, - 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read ': 1, 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read KANDINSKY5_E2E_NUM_GPUS': 1, + 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read KANDINSKY5_E2E_WRITE_REFERENCE': 1, 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: whole-environ': 3, 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: write WANDB_MODE': 1, 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: read FASTVIDEO_NIGHTLY': 1, @@ -420,8 +353,10 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/tests/ssim/conftest.py: write ': 3, 'fastvideo/tests/ssim/inference_similarity_utils.py: read FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/ssim/inference_similarity_utils.py: write FASTVIDEO_ATTENTION_BACKEND': 3, - 'fastvideo/tests/ssim/reference_utils.py: read ': 1, - 'fastvideo/tests/ssim/reference_videos_cli.py: read ': 3, + 'fastvideo/tests/ssim/reference_utils.py: read FASTVIDEO_SSIM_FULL_QUALITY': 1, + 'fastvideo/tests/ssim/reference_videos_cli.py: read ': 1, + 'fastvideo/tests/ssim/reference_videos_cli.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, + 'fastvideo/tests/ssim/reference_videos_cli.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE': 1, 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_AR_SSIM_MODEL_PATH': 1, 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_SSIM_MODEL_PATH': 1, 'fastvideo/tests/ssim/test_flux_t2i_similarity.py: read FLUX_T2I_MODEL_DIR': 1, @@ -444,7 +379,7 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/tests/stages/_denoising_fixtures.py: write FASTVIDEO_CFG_GATE_STEP': 2, 'fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py: write FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/stages/test_minimax_h3_encoding_offload.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/train/methods/grad_norm_regression.py: read ': 1, + 'fastvideo/tests/train/methods/grad_norm_regression.py: read FASTVIDEO_GRADNORM_UPDATE': 1, 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_ADDR': 1, 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_PORT': 1, 'fastvideo/tests/train/methods/test_ltx2_finetune.py: write MASTER_ADDR': 1, @@ -534,23 +469,11 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/tests/vaes/test_wan_vae.py: write MASTER_PORT': 1, 'fastvideo/tests/worker/test_gpu_worker.py: write FASTVIDEO_NVTX_PROFILE': 2, 'fastvideo/tests/worker/test_gpu_worker.py: write LOCAL_RANK': 1, - 'fastvideo/train/entrypoint/dcp_to_diffusers.py: write ': 1, 'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 2, - 'fastvideo/training/self_forcing_distillation_pipeline.py: read FASTVIDEO_FSDP2_AUTOWRAP': 1, 'fastvideo/utils.py: read ': 5, - 'fastvideo/utils.py: read FASTVIDEO_WORKER_MULTIPROC_METHOD': 1, 'fastvideo/utils.py: write ': 1, - 'fastvideo/utils.py: write FASTVIDEO_WORKER_MULTIPROC_METHOD': 1, - 'fastvideo/worker/gpu_worker.py: write LOCAL_RANK': 1, - 'fastvideo/worker/gpu_worker.py: write NCCL_ASYNC_ERROR_HANDLING': 1, - 'fastvideo/worker/gpu_worker.py: write RANK': 1, - 'fastvideo/worker/gpu_worker.py: write TORCH_NCCL_AVOID_RECORD_STREAMS': 1, - 'fastvideo/worker/gpu_worker.py: write WORLD_SIZE': 1, 'fastvideo/worker/ray_distributed_executor.py: read ': 2, - 'fastvideo/worker/ray_distributed_executor.py: read RAY_USAGE_STATS_ENABLED': 1, 'fastvideo/worker/ray_distributed_executor.py: whole-environ': 1, - 'fastvideo/worker/ray_distributed_executor.py: write RAY_USAGE_STATS_ENABLED': 1, - 'fastvideo/worker/ray_env.py: import-time-read FASTVIDEO_CONFIG_ROOT': 1, 'fastvideo/worker/ray_env.py: read ': 1, 'fastvideo/worker/ray_utils.py: whole-environ': 1, 'fastvideo/worker/worker_base.py: read ': 1, @@ -576,12 +499,6 @@ def is_allowlisted(name: str) -> bool: return False -def _literal(node: ast.AST | None) -> str: - if isinstance(node, ast.Constant) and isinstance(node.value, str): - return node.value - return "" - - class EnvAccessScanner: """Find environment accesses and registry-field uses in one module. @@ -604,6 +521,13 @@ class EnvAccessScanner: self.getenv_names: set[str] = set() self.setenv_names: set[str] = set() self.envs_module_names: set[str] = set() + # Module-level string constants, so that os.environ.get(NAME_ENV) resolves to the name. + self.constants: dict[str, str] = {} + for statement in tree.body: + if (isinstance(statement, ast.Assign) and len(statement.targets) == 1 + and isinstance(statement.targets[0], ast.Name) and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str)): + self.constants[statement.targets[0].id] = statement.value.value for node in ast.walk(tree): if isinstance(node, ast.Import): for alias in node.names: @@ -624,6 +548,14 @@ class EnvAccessScanner: elif node.module == "fastvideo" and alias.name == "envs": self.envs_module_names.add(local) + def _literal(self, node: ast.AST | None) -> str: + """Return the variable name that a call argument spells, or "".""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + if isinstance(node, ast.Name) and node.id in self.constants: + return self.constants[node.id] + return "" + def _is_os_attr(self, node: ast.AST, attrs: tuple[str, ...]) -> bool: return (isinstance(node, ast.Attribute) and node.attr in attrs and isinstance(node.value, ast.Name) and node.value.id in self.os_names) @@ -670,24 +602,24 @@ class EnvAccessScanner: call = self.parents[parent] assert isinstance(call, ast.Call) if call.func is parent and parent.attr == "get": - self.violations.append(("read", _literal(call.args[0] if call.args else None), line)) + self.violations.append(("read", self._literal(call.args[0] if call.args else None), line)) return if call.func is parent and parent.attr in ("setdefault", "pop"): - self.violations.append(("write", _literal(call.args[0] if call.args else None), line)) + self.violations.append(("write", self._literal(call.args[0] if call.args else None), line)) return if isinstance(parent, ast.Subscript) and parent.value is node: kind = "read" if isinstance(parent.ctx, ast.Load) else "write" - self.violations.append((kind, _literal(parent.slice), line)) + self.violations.append((kind, self._literal(parent.slice), line)) return if (isinstance(parent, ast.Compare) and len(parent.ops) == 1 and isinstance(parent.ops[0], (ast.In, ast.NotIn)) and parent.comparators[0] is node): - self.violations.append(("read", _literal(parent.left), line)) + self.violations.append(("read", self._literal(parent.left), line)) return self.violations.append(("whole-environ", "", line)) def _check_call(self, node: ast.Call) -> None: func = node.func - name = _literal(node.args[0] if node.args else None) + name = self._literal(node.args[0] if node.args else None) if self._is_os_attr(func, ("getenv", "getenvb")) or (isinstance(func, ast.Name) and func.id in self.getenv_names): self.violations.append(("read", name, node.lineno)) @@ -695,6 +627,9 @@ class EnvAccessScanner: or (isinstance(func, ast.Name) and func.id in self.setenv_names) or (isinstance(func, ast.Attribute) and func.attr in ("setenv", "delenv"))): self.violations.append(("write", name, node.lineno)) + elif (isinstance(func, ast.Attribute) and func.attr in EXTERNAL_WRITE_HELPERS + and self._is_envs_module(func.value) and name not in EXTERNAL_WRITE_ALLOWLIST): + self.violations.append(("write", name, node.lineno)) def _check_registry_use(self, node: ast.Attribute, parent: ast.AST | None, in_function: bool) -> None: # Only a call counts: ``getter = envs.X.get`` neither reads the variable @@ -715,11 +650,21 @@ def scanned_files() -> list[Path]: if path != REGISTRY_PATH and not any(excluded in path.parents for excluded in EXCLUDED_DIRS)) +def registry_internal_reads(registry_names: set[str]) -> Counter[str]: + """Count registry reads inside fastvideo/envs.py itself, such as a computed default that reads another field.""" + reads: Counter[str] = Counter() + for node in ast.walk(ast.parse(REGISTRY_PATH.read_text(encoding="utf-8"))): + if (isinstance(node, ast.Attribute) and node.attr in REGISTRY_READ_METHODS and isinstance(node.value, ast.Name) + and node.value.id in registry_names): + reads[node.value.id] += 1 + return reads + + def collect_violations(registry: ModuleType) -> tuple[dict[str, list[int]], Counter[str]]: """Return every violation key with its line numbers, plus registry read counts.""" registry_names = set(registry.environment_variables) found: dict[str, list[int]] = {} - reads: Counter[str] = Counter() + reads: Counter[str] = registry_internal_reads(registry_names) for path in scanned_files(): relative = path.relative_to(REPO_ROOT).as_posix() scanner = EnvAccessScanner(ast.parse(path.read_text(encoding="utf-8"), filename=relative), registry_names) @@ -793,10 +738,21 @@ def _escape(text: str) -> str: def render_env_table(registry: ModuleType) -> str: - """Render the registry as an aligned Markdown table, in declaration order.""" + """Render the registry and the deprecation table as aligned Markdown tables, in declaration order.""" rows = [["Variable", "Type", "Default", "Category", "Description"]] for name, field in registry.environment_variables.items(): - rows.append([f"`{name}`", field.type_name, _render_default(field), field.category, _escape(field.doc)]) + doc = _escape(field.doc) + if field.deprecated_names: + doc += " Deprecated names: " + ", ".join(f"`{old}`" for old in field.deprecated_names) + "." + rows.append([f"`{name}`", field.type_name, _render_default(field), field.category, doc]) + deprecated_rows = [["Deprecated variable", "Reason"]] + deprecated_rows.extend([f"`{name}`", reason] for name, reason in registry.DEPRECATED_VARIABLES.items()) + return (_render_table(rows) + "\n\nVariables that FastVideo no longer reads; setting one logs a warning:\n\n" + + _render_table(deprecated_rows)) + + +def _render_table(rows: list[list[str]]) -> str: + """Render rows as a Markdown table whose pipes line up.""" widths = [max(len(row[column]) for row in rows) for column in range(len(rows[0]))] def render_row(cells: list[str]) -> str: diff --git a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py index 43971b7fe..bd9f06b18 100644 --- a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py +++ b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py @@ -84,7 +84,7 @@ NUM_GPUS = os.environ.get("KANDINSKY5_E2E_NUM_GPUS", "1") # with their real videos/captions for the actual recipe: # _clean_previous_artifacts() deletes this root wholesale on every run, # and the preprocess subprocess is pointed here via the -# KANDINSKY5_OVERFIT_DATA_DIR / KANDINSKY5_OVERFIT_OUTPUT_DIR env +# FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR / FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR env # overrides instead of the script's user-facing defaults. E2E_ROOT = Path("data") / "kandinsky5_e2e" RAW_DATA_DIR = E2E_ROOT / "raw" @@ -223,8 +223,8 @@ def _run_preprocessing() -> None: # (data/kandinsky5_overfit{,_preprocessed}) -- the very directories a # user populates for the real recipe, which this test must never touch. env = dict(os.environ) - env["KANDINSKY5_OVERFIT_DATA_DIR"] = str(RAW_DATA_DIR) - env["KANDINSKY5_OVERFIT_OUTPUT_DIR"] = str(PREPROCESSED_DIR) + env["FASTVIDEO_TEST_KANDINSKY5_OVERFIT_DATA_DIR"] = str(RAW_DATA_DIR) + env["FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR"] = str(PREPROCESSED_DIR) cmd = [ sys.executable, "-m", diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py index 5ffa76e82..476bc55c3 100644 --- a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py +++ b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py @@ -82,9 +82,9 @@ def download_data(): def run_preprocessing(case: dict): env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0] - env["LTX2_OVERFIT_DATA_DIR"] = str(LOCAL_RAW_DATA_DIR) - env["LTX2_OVERFIT_OUTPUT_DIR"] = str(case["prep_dir"]) - env["LTX2_OVERFIT_MODEL"] = case["model"] + env["FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR"] = str(LOCAL_RAW_DATA_DIR) + env["FASTVIDEO_TEST_LTX2_OVERFIT_OUTPUT_DIR"] = str(case["prep_dir"]) + env["FASTVIDEO_TEST_LTX2_OVERFIT_MODEL"] = case["model"] cmd = [sys.executable, PREPROCESSING_SCRIPT] print(f"Running preprocessing: {cmd} (model={case['model']})") subprocess.run(cmd, check=True, env=env) diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index da9d8e0b5..ec07a5412 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -33,6 +33,7 @@ import os import sys from typing import Any +import fastvideo.envs as envs from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -45,14 +46,11 @@ def _ensure_distributed() -> None: For plain ``python`` we fill in the minimum required vars so that ``init_process_group`` succeeds with world_size=1. """ - for key, default in [ - ("RANK", "0"), - ("LOCAL_RANK", "0"), - ("WORLD_SIZE", "1"), - ("MASTER_ADDR", "127.0.0.1"), - ("MASTER_PORT", "29500"), - ]: - os.environ.setdefault(key, default) + envs.setdefault_external("RANK", "0") + envs.setdefault_external("LOCAL_RANK", "0") + envs.setdefault_external("WORLD_SIZE", "1") + envs.setdefault_external("MASTER_ADDR", "127.0.0.1") + envs.setdefault_external("MASTER_PORT", "29500") def _save_role_pretrained( diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py index c36fe200b..d73c40a7b 100644 --- a/fastvideo/train/models/ltx2/ltx2.py +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -396,7 +396,7 @@ class LTX2Model(WanModel): f"checkpoint's expected text context dim {expected_dim}. " "The parquet was likely preprocessed with a different " "LTX-2 version (2.0 stores 3840-d, 2.3 stores 4096-d); " - "re-run preprocess_ltx2_overfit.py with LTX2_OVERFIT_MODEL " + "re-run preprocess_ltx2_overfit.py with FASTVIDEO_TEST_LTX2_OVERFIT_MODEL " "set to the same checkpoint as models.student.init_from.") def _update_rope_fps(self, infos: list[dict[str, Any]] | None) -> None: diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index bbc5a44d6..e1db17c70 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 import copy -import os import time from collections import deque from typing import Any @@ -47,7 +46,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): def initialize_training_pipeline(self, training_args: TrainingArgs): """Initialize the self-forcing training pipeline.""" # Check if FSDP2 auto wrap is enabled - not supported for self-forcing distillation - if os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1": + if envs.FASTVIDEO_FSDP2_AUTOWRAP.get(): raise NotImplementedError("FASTVIDEO_FSDP2_AUTOWRAP is not implemented for self-forcing distillation. " "Please set FASTVIDEO_FSDP2_AUTOWRAP=0 or unset the environment variable.") diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 441ac3fcd..b42c3f606 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -82,16 +82,16 @@ def find_nccl_library() -> str: def find_hccl_library() -> str: """ - We either use the library file specified by the `HCCL_SO_PATH` + We either use the library file specified by the `FASTVIDEO_HCCL_SO_PATH` environment variable, or we find the library file brought by PyTorch. After importing `torch`, `libhccl.so` can be found by `ctypes` automatically. """ - so_file = envs.HCCL_SO_PATH.get() + so_file = envs.FASTVIDEO_HCCL_SO_PATH.get() # manually load the nccl library if so_file: - logger.info("Found hccl from environment variable HCCL_SO_PATH=%s", so_file) + logger.info("Found hccl from environment variable FASTVIDEO_HCCL_SO_PATH=%s", so_file) else: if torch.version.cann is not None: # codespell:ignore cann so_file = "libhccl.so" @@ -1240,9 +1240,9 @@ def xpu_is_initialized() -> bool: def force_spawn() -> None: - if os.environ.get("FASTVIDEO_WORKER_MULTIPROC_METHOD") == "fork": + if envs.FASTVIDEO_WORKER_MULTIPROC_METHOD.get() == "fork": logger.warning("We must use the `spawn` multiprocessing start method.") - os.environ["FASTVIDEO_WORKER_MULTIPROC_METHOD"] = "spawn" + envs.FASTVIDEO_WORKER_MULTIPROC_METHOD.set("spawn") def get_mp_context() -> BaseContext: diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 8a22cdf73..d0712091c 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 -import os from typing import Any, cast import torch @@ -50,9 +49,9 @@ class Worker: # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 - os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + envs.set_external("TORCH_NCCL_AVOID_RECORD_STREAMS", "1") # This env var set by Ray causes exceptions with graph building. - os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + envs.unset_external("NCCL_ASYNC_ERROR_HANDLING") # Set environment variables BEFORE calling get_local_torch_device() # so that each worker uses the correct device @@ -60,9 +59,9 @@ class Worker: # Ray deliberately excludes LOCAL_RANK from the copied driver # environment and exposes all GPUs assigned to the node, so leaving an # inherited or missing value here would bind every Ray actor to cuda:0. - os.environ["LOCAL_RANK"] = str(self.local_rank) - os.environ["RANK"] = str(self.rank) - os.environ["WORLD_SIZE"] = str(self.fastvideo_args.num_gpus) + envs.set_external("LOCAL_RANK", str(self.local_rank)) + envs.set_external("RANK", str(self.rank)) + envs.set_external("WORLD_SIZE", str(self.fastvideo_args.num_gpus)) # Platform-agnostic device initialization self.device = get_local_torch_device() diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index a48e9e506..e7093bd8b 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -94,6 +94,7 @@ class RayDistributedExecutor(Executor): "NCCL_NVLS_ENABLE", "NCCL_DEBUG", "NCCL_DEBUG_SUBSYS", + "LD_LIBRARY_PATH", } def _init_executor(self) -> None: @@ -103,7 +104,7 @@ class RayDistributedExecutor(Executor): # Disable Ray usage stats collection. ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0") if ray_usage != "1": - os.environ["RAY_USAGE_STATS_ENABLED"] = "0" + envs.set_external("RAY_USAGE_STATS_ENABLED", "0") self._init_workers_ray(placement_group) diff --git a/fastvideo/worker/ray_env.py b/fastvideo/worker/ray_env.py index b3663afbb..4adcee85e 100644 --- a/fastvideo/worker/ray_env.py +++ b/fastvideo/worker/ray_env.py @@ -8,24 +8,20 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) -CONFIG_HOME = envs.FASTVIDEO_CONFIG_ROOT.get() -# This file contains a list of env vars that should not be copied -# from the driver to the Ray workers. -RAY_NON_CARRY_OVER_ENV_VARS_FILE = os.path.join(CONFIG_HOME, "ray_non_carry_over_env_vars.json") - -try: - if os.path.exists(RAY_NON_CARRY_OVER_ENV_VARS_FILE): - with open(RAY_NON_CARRY_OVER_ENV_VARS_FILE) as f: - RAY_NON_CARRY_OVER_ENV_VARS = set(json.load(f)) - else: - RAY_NON_CARRY_OVER_ENV_VARS = set() -except json.JSONDecodeError: - logger.warning( - "Failed to parse %s. Using an empty set for non-carry-over env vars.", - RAY_NON_CARRY_OVER_ENV_VARS_FILE, - ) - RAY_NON_CARRY_OVER_ENV_VARS = set() +def _load_non_carry_over_env_vars(non_carry_over_file: str) -> set[str]: + """Read the env vars that must not be copied from the driver to the Ray workers.""" + try: + if os.path.exists(non_carry_over_file): + with open(non_carry_over_file) as f: + return set(json.load(f)) + return set() + except json.JSONDecodeError: + logger.warning( + "Failed to parse %s. Using an empty set for non-carry-over env vars.", + non_carry_over_file, + ) + return set() def get_env_vars_to_copy( @@ -52,18 +48,26 @@ def get_env_vars_to_copy( """ exclude_vars = exclude_vars or set() additional_vars = additional_vars or set() + # This file contains a list of env vars that should not be copied + # from the driver to the Ray workers. + non_carry_over_file = os.path.join(envs.FASTVIDEO_CONFIG_ROOT.get(), "ray_non_carry_over_env_vars.json") + non_carry_over_vars = _load_non_carry_over_env_vars(non_carry_over_file) + registered_vars = { + name + for field in envs.environment_variables.values() + for name in (field.name, *field.deprecated_names) + } env_vars_to_copy = { v - for v in set(envs.environment_variables).union(additional_vars) - if v not in exclude_vars and v not in RAY_NON_CARRY_OVER_ENV_VARS + for v in registered_vars.union(additional_vars) if v not in exclude_vars and v not in non_carry_over_vars } to_destination = " to " + destination if destination is not None else "" logger.info( "RAY_NON_CARRY_OVER_ENV_VARS from config: %s", - RAY_NON_CARRY_OVER_ENV_VARS, + non_carry_over_vars, ) logger.info( "Copying the following environment variables%s: %s", @@ -72,7 +76,7 @@ def get_env_vars_to_copy( ) logger.info( "If certain env vars should NOT be copied, add them to %s file", - RAY_NON_CARRY_OVER_ENV_VARS_FILE, + non_carry_over_file, ) return env_vars_to_copy diff --git a/tests/local_tests/ltx2/README.md b/tests/local_tests/ltx2/README.md index 33f120be5..4c5501bc9 100644 --- a/tests/local_tests/ltx2/README.md +++ b/tests/local_tests/ltx2/README.md @@ -16,7 +16,7 @@ skip or fail. Skipped in CI; run locally on a single GPU. | Official commit/version | `` | | HF weights | `Lightricks/LTX-2`, `FastVideo/LTX2-base`, `FastVideo/LTX2-Distilled-Diffusers` | | HF revision | `` | -| Local weights dir | `` (env: `LTX2_DIFFUSERS_PATH`, `LTX2_FASTVIDEO_GEMMA_LOG`, etc.) | +| Local weights dir | `` (env: `LTX2_DIFFUSERS_PATH`, `FASTVIDEO_LTX2_GEMMA_LOG`, etc.) | | Source layout | `` | | Needs conversion | `` | diff --git a/tests/local_tests/ltx2/test_ltx2_gemma_parity.py b/tests/local_tests/ltx2/test_ltx2_gemma_parity.py index fd001c77b..acb62e331 100644 --- a/tests/local_tests/ltx2/test_ltx2_gemma_parity.py +++ b/tests/local_tests/ltx2/test_ltx2_gemma_parity.py @@ -21,7 +21,7 @@ if ltx_core_path.exists() and str(ltx_core_path) not in sys.path: def _init_log_paths() -> tuple[Path, Path]: base_dir = Path(os.getenv("LTX2_DEBUG_DIR", "ltx2_debug")) - fastvideo_log = Path(os.getenv("LTX2_FASTVIDEO_GEMMA_LOG", base_dir / "fastvideo_gemma.log")) + fastvideo_log = Path(os.getenv("FASTVIDEO_LTX2_GEMMA_LOG", base_dir / "fastvideo_gemma.log")) reference_log = Path(os.getenv("LTX2_REFERENCE_GEMMA_LOG", base_dir / "reference_gemma.log")) fastvideo_log.parent.mkdir(parents=True, exist_ok=True) reference_log.parent.mkdir(parents=True, exist_ok=True) From e1f390479919e156b30b8611d856ea340c0577dd Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Tue, 29 Sep 2026 16:57:22 -0700 Subject: [PATCH 041/151] [kernel] sm_100a CUDA backward for VSA block-sparse attention, 128-token blocks (#1866) Co-authored-by: Pengcheng Li Co-authored-by: Claude Fable 5.1 --- fastvideo-kernel/CMakeLists.txt | 8 +- .../block_sparse_bwd_blk128_sm100a.cu | 7 + .../block_sparse_bwd_kernel_sm100a_blk128.cuh | 1006 +++++++++++++++++ .../block_sparse_bwd_launch_sm100a_blk128.cuh | 231 ++++ .../csrc/attention/block_sparse_bwd_sm100a.cu | 44 +- .../csrc/attention/primitives.cuh | 27 + fastvideo-kernel/csrc/common_extension.cpp | 7 + .../fastvideo_kernel/block_sparse_attn.py | 43 +- .../block_sparse_attn_bwd_sm100a.py | 73 +- .../block_sparse_attn_sm100a.py | 5 +- .../tests}/test_block_sparse_bwd_sm100a.py | 118 +- .../test_block_sparse_sm100a_dispatch.py | 53 +- .../backends/video_sparse_attn_h3.py | 2 +- .../attention/test_vsa_h3_sm100a_route.py | 2 +- 14 files changed, 1508 insertions(+), 118 deletions(-) create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh rename {tests => fastvideo-kernel/tests}/test_block_sparse_bwd_sm100a.py (69%) rename {tests => fastvideo-kernel/tests}/test_block_sparse_sm100a_dispatch.py (80%) diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 1fc50a3f1..ff09788c1 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -464,10 +464,12 @@ if(BUILD_CXX_KERNELS) csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() - # VSA block-sparse attention BACKWARD, 64-token blocks: the same images as the forward. + # VSA block-sparse attention BACKWARD, 64- and 128-token blocks: the same images as the forward. if(_VSA_FWD_GENCODE) - list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu) - set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu + csrc/attention/block_sparse_bwd_blk128_sm100a.cu) + set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu + csrc/attention/block_sparse_bwd_blk128_sm100a.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu new file mode 100644 index 000000000..2d04613a1 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu @@ -0,0 +1,7 @@ +// block_sparse_bwd_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding. +// +// Same source as block_sparse_bwd_sm100a.cu with VSA_BLK128 set: the kernel and launch land in +// namespace vsa_bwd_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the +// exported entry point becomes block_sparse_sm100a_blk128_bwd. +#define VSA_BLK128 true +#include "block_sparse_bwd_sm100a.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh new file mode 100644 index 000000000..1ca9bcf4e --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh @@ -0,0 +1,1006 @@ +// block_sparse_bwd_kernel_sm100a_blk128.cuh -- VSA block-sparse attention BACKWARD, +// 128-token blocks, sm_100a (one CTA per kv block). Warp-specialized: load / +// MMA (tcgen05) / softmax (P^T, dS^T) / epilogue (dQ drain). Three kernels: preprocess +// (Delta, dqaccum zero), main (dK, dV, dQ partials), postprocess (dQ unscramble + scale). +#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH + +#include +#include +#include +#include +#include +#include +#include "primitives.cuh" + +namespace vsa_bwd_blk128 { + +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif + +enum class Sched { NON_PERSISTENT, STATIC_PERSISTENT, CLC }; +#ifndef KERNEL_SCHED +#define KERNEL_SCHED Sched::NON_PERSISTENT +#endif + +#ifndef KERNEL_PDL +#define KERNEL_PDL true +#endif + +constexpr int BLOCK = 128; +constexpr int KV_TILE = BLOCK; +constexpr int Q_TILE = BLOCK; +constexpr int HEAD_DIM = 128; +constexpr int SUB_COLS_BF16 = 64; +constexpr int SUB_COLS_BYTES = SUB_COLS_BF16 * (int)sizeof(__nv_bfloat16); +constexpr int HD_SUBTILES = HEAD_DIM / SUB_COLS_BF16; +constexpr int KV_SUB_COLS_BYTES = KV_TILE * SUB_COLS_BYTES; +constexpr int KV_TILE_BYTES = HD_SUBTILES * KV_SUB_COLS_BYTES; +constexpr int Q_SUB_COLS_BYTES = Q_TILE * SUB_COLS_BYTES; +constexpr int Q_TILE_BYTES = HD_SUBTILES * Q_SUB_COLS_BYTES; +constexpr int NUM_Q_STAGES = 2; +constexpr int DST_TILE_BYTES = KV_TILE * SUB_COLS_BYTES; +constexpr int DST_BYTES = 2 * DST_TILE_BYTES; +constexpr int MMA_K = 16; +constexpr int K_ATOMS_PER_SUBTILE = SUB_COLS_BF16 / MMA_K; +constexpr int K_ATOMS_PER_Q_HALF = SUB_COLS_BF16 / MMA_K; +constexpr int K_ATOMS_PER_KV_TILE = KV_TILE / MMA_K; +constexpr int BF16X2_COLS_PER_K16 = MMA_K / 2; + +struct DQConfig { + static constexpr int COLS = 32; + + static constexpr int DQ_ONE_PUSH_BYTES = Q_TILE * COLS * (int)sizeof(float); + static constexpr int DQ_STAGE_BYTES = DQ_ONE_PUSH_BYTES; + static constexpr int DQ_STAGE_BUFFERS = 2; + static constexpr int DQ_BLOCK_ELEMS = Q_TILE * HEAD_DIM; +}; +using DQ = DQConfig; + +constexpr int N_WARPS = 16; +constexpr int W_EPI0 = 0, W_SOFTMAX0 = 4, W_MMA = 12, W_LOAD = 13; +constexpr int W_SCHED = 14; +constexpr int CLC_STAGES = 2; +constexpr int CLC_ARRIVALS = 15; + +constexpr int NUM_BARS = 4 * NUM_Q_STAGES + 14 + 2 * CLC_STAGES; +constexpr int SMEM_TOTAL = 2 * KV_TILE_BYTES + NUM_Q_STAGES * Q_TILE_BYTES + Q_TILE_BYTES + + DST_BYTES + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES + + (NUM_Q_STAGES + 1) * Q_TILE * (int)sizeof(float) + NUM_BARS * 8 + + CLC_STAGES * 16 + 48; + +constexpr int ST_COLS = Q_TILE; +constexpr int ST_HALF_COLS = SUB_COLS_BF16; +constexpr int DV_COLS = HEAD_DIM; +constexpr int DK_COLS = HEAD_DIM; +constexpr int TMEM_TOTAL = ST_COLS + DV_COLS + ST_COLS + DK_COLS; +static_assert(ST_COLS == 2 * ST_HALF_COLS, "two q-halves per tile"); +static_assert(TMEM_TOTAL == 512, "TMEM map must fill exactly 512 columns"); + +extern __shared__ __align__(1024) uint8_t bwd_smem[]; + +template +__device__ __forceinline__ size_t token_offset(int batch, int head, int num_heads, int seqlen, + int t) { + if constexpr (BHSD) + return ((size_t)(batch * num_heads + head) * seqlen + t) * HEAD_DIM; + else + return ((size_t)(batch * seqlen + t) * num_heads + head) * HEAD_DIM; +} + +struct WorkItem { + int batch; + int head; + int kv_block_id_in_seq; + int batch_head; + const int* local_k2q_idx; + + int local_k2q_num; +}; + +__device__ __forceinline__ WorkItem decode_workitem(int batch_head, int kv_block_id_in_seq, + const int* __restrict__ k2q_idx, + const int* __restrict__ k2q_num, + int max_q_blocks, int num_heads, + int num_kv_blocks_per_seq) { + WorkItem it; + it.batch_head = batch_head; + it.batch = batch_head / num_heads; + it.head = batch_head % num_heads; + it.kv_block_id_in_seq = kv_block_id_in_seq; + const int item = batch_head * num_kv_blocks_per_seq + kv_block_id_in_seq; + it.local_k2q_idx = k2q_idx + (size_t)item * (size_t)max_q_blocks; + it.local_k2q_num = k2q_num[item]; + return it; +} + +template +__global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) + vsa_bwd_main_kernel(const __grid_constant__ CUtensorMap tmap_q, + const __grid_constant__ CUtensorMap tmap_k, + const __grid_constant__ CUtensorMap tmap_v, + const __grid_constant__ CUtensorMap tmap_do, + const __grid_constant__ CUtensorMap tmap_dk, + const __grid_constant__ CUtensorMap tmap_dv, float* __restrict__ dqaccum, + const float* __restrict__ lse_rows, const float* __restrict__ delta_rows, + const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num, + int max_q_blocks, const int* __restrict__ workitem_remap, + const int* __restrict__ variable_block_sizes, int num_samples, + int num_heads, int seqlen, float scale_log2, float sm_scale) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + constexpr bool PERSISTENT = SCHED != Sched::NON_PERSISTENT; + constexpr bool CLC = SCHED == Sched::CLC; + const int num_kv_blocks_per_seq = seqlen / BLOCK; + [[maybe_unused]] const int total_workitems = num_samples * num_heads * num_kv_blocks_per_seq; + uint8_t* sK = bwd_smem; + uint8_t* sV = sK + KV_TILE_BYTES; + uint8_t* sQ[NUM_Q_STAGES] = {sV + KV_TILE_BYTES, sV + KV_TILE_BYTES + Q_TILE_BYTES}; + uint8_t* sDO = sQ[0] + NUM_Q_STAGES * Q_TILE_BYTES; + __nv_bfloat16* sDST = reinterpret_cast<__nv_bfloat16*>(sDO + Q_TILE_BYTES); + uint8_t* sDQ_STAGE_bytes = sDO + Q_TILE_BYTES + DST_BYTES; + float* sDQ_STAGE[DQ::DQ_STAGE_BUFFERS] = { + reinterpret_cast(sDQ_STAGE_bytes), + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BYTES)}; + float* sLSE = + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES); + float* sDelta = sLSE + NUM_Q_STAGES * Q_TILE; + + uint64_t* full_bar_q = reinterpret_cast(sDelta + Q_TILE); + uint64_t* empty_bar_q = full_bar_q + NUM_Q_STAGES; + uint64_t* full_bar_do = empty_bar_q + NUM_Q_STAGES; + uint64_t* empty_bar_do = full_bar_do + 1; + uint64_t* full_bar_lse = empty_bar_do + 1; + uint64_t* empty_bar_lse = full_bar_lse + NUM_Q_STAGES; + uint64_t* full_bar_delta = empty_bar_lse + NUM_Q_STAGES; + uint64_t* empty_bar_delta = full_bar_delta + 1; + uint64_t* full_bar_st = empty_bar_delta + 1; + uint64_t* full_bar_dpt = full_bar_st + 1; + uint64_t* full_bar_pt = full_bar_dpt + 1; + uint64_t* full_bar_dst = full_bar_pt + 1; + uint64_t* full_bar_dq = full_bar_dst + 1; + uint64_t* empty_bar_dq = full_bar_dq + 1; + uint64_t* full_bar_dv = empty_bar_dq + 1; + uint64_t* full_bar_dk = full_bar_dv + 1; + + uint64_t* empty_bar_kv = full_bar_dk + 1; + uint64_t* empty_bar_epi = empty_bar_kv + 1; + uint64_t* clc_full = empty_bar_epi + 1; + uint64_t* clc_empty = clc_full + CLC_STAGES; + uint32_t* clc_response = reinterpret_cast( + (reinterpret_cast(clc_empty + CLC_STAGES) + 15u) & ~uintptr_t(15u)); + uint32_t* tmem_slot = clc_response + CLC_STAGES * 4; + + const int tid = threadIdx.x, warp_id = tid >> 5, lane = tid & 31; + + if (tid == 0) { + #pragma unroll + for (int s = 0; s < NUM_Q_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar_q[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_q[s]), 1); + } + mbarrier_init(smem_ptr_u32(full_bar_do), 1); + mbarrier_init(smem_ptr_u32(empty_bar_do), 1); + #pragma unroll + for (int s = 0; s < NUM_Q_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar_lse[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_lse[s]), 8); + } + mbarrier_init(smem_ptr_u32(full_bar_delta), 1); + mbarrier_init(smem_ptr_u32(empty_bar_delta), 8); + mbarrier_init(smem_ptr_u32(full_bar_st), 1); + mbarrier_init(smem_ptr_u32(full_bar_dpt), 1); + mbarrier_init(smem_ptr_u32(full_bar_pt), 8); + mbarrier_init(smem_ptr_u32(full_bar_dst), 8); + mbarrier_init(smem_ptr_u32(full_bar_dq), 1); + mbarrier_init(smem_ptr_u32(empty_bar_dq), 4); + mbarrier_init(smem_ptr_u32(full_bar_dv), 1); + mbarrier_init(smem_ptr_u32(full_bar_dk), 1); + if constexpr (PERSISTENT) { + mbarrier_init(smem_ptr_u32(empty_bar_kv), 1); + mbarrier_init(smem_ptr_u32(empty_bar_epi), 8 * 32); + } + if constexpr (CLC) { + #pragma unroll + for (int st = 0; st < CLC_STAGES; ++st) { + mbarrier_init(smem_ptr_u32(&clc_full[st]), 1); + mbarrier_init(smem_ptr_u32(&clc_empty[st]), CLC_ARRIVALS); + } + #pragma unroll + for (int i = 0; i < CLC_STAGES * 4; ++i) clc_response[i] = 0; + } + } + fence_mbarrier_init_release_cluster(); + __syncthreads(); + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + + if (warp_id == W_LOAD) { + setmaxnreg_dec<88>(); + + EmptyPhaseTracker q_empty_ph, lse_empty_ph; + EmptyPhaseTracker<1> do_empty_ph, delta_empty_ph; + [[maybe_unused]] EmptyPhaseTracker<1> epi_empty_ph, kv_empty_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + if (it.local_k2q_num != 0) { + auto load_tile = [&](uint8_t* dst, const CUtensorMap* map, uint64_t* full_bar, + int token_begin) { + static_assert(Q_SUB_COLS_BYTES == KV_SUB_COLS_BYTES, + "one subtile stride for Q/dO and K/V"); + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32(dst + s * Q_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), + 0, token_begin, s, it.batch_head); + else + tma_load_3d(smem_ptr_u32(dst + s * Q_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), + 0, it.batch * seqlen + token_begin, it.head * HD_SUBTILES + s); + } + }; + auto load_lse = [&](int qblock) { + const int stage = lse_empty_ph.get_stage(); + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_lse[stage]), + lse_empty_ph.get_phase()); + lse_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_lse[stage]), + Q_TILE * (int)sizeof(float)); + cpasync_bulk_load_mbarrier( + smem_ptr_u32(sLSE + stage * Q_TILE), + lse_rows + (size_t)it.batch_head * seqlen + (size_t)qblock * Q_TILE, + Q_TILE * sizeof(float), smem_ptr_u32(&full_bar_lse[stage])); + } + }; + auto load_delta = [&](int qblock) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_delta), delta_empty_ph.get_phase()); + delta_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(full_bar_delta), Q_TILE * (int)sizeof(float)); + cpasync_bulk_load_mbarrier( + smem_ptr_u32(sDelta), + delta_rows + (size_t)it.batch_head * seqlen + (size_t)qblock * Q_TILE, + Q_TILE * sizeof(float), smem_ptr_u32(full_bar_delta)); + } + }; + + auto load_q = [&](int qblock, auto with_kv_const) { + constexpr bool with_kv = decltype(with_kv_const)::value; + const int stage = q_empty_ph.get_stage(); + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_q[stage]), q_empty_ph.get_phase()); + q_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_q[stage]), + Q_TILE_BYTES + (with_kv ? KV_TILE_BYTES : 0)); + load_tile(sQ[stage], &tmap_q, &full_bar_q[stage], qblock * Q_TILE); + if constexpr (with_kv) + load_tile(sK, &tmap_k, &full_bar_q[stage], it.kv_block_id_in_seq * KV_TILE); + } + }; + + auto load_do = [&](int qblock, auto with_kv_const) { + constexpr bool with_kv = decltype(with_kv_const)::value; + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_do), do_empty_ph.get_phase()); + do_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(full_bar_do), + Q_TILE_BYTES + (with_kv ? KV_TILE_BYTES : 0)); + load_tile(sDO, &tmap_do, full_bar_do, qblock * Q_TILE); + if constexpr (with_kv) + load_tile(sV, &tmap_v, full_bar_do, it.kv_block_id_in_seq * KV_TILE); + } + }; + + if constexpr (PERSISTENT) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_kv), kv_empty_ph.get_phase()); + kv_empty_ph.advance(); + } + + load_q(it.local_k2q_idx[0], std::true_type{}); + load_lse(it.local_k2q_idx[0]); + load_do(it.local_k2q_idx[0], std::true_type{}); + load_delta(it.local_k2q_idx[0]); + for (int j = 1; j < it.local_k2q_num; ++j) { + + const int qblock = it.local_k2q_idx[j]; + load_q(qblock, std::false_type{}); + load_lse(qblock); + load_do(qblock, std::false_type{}); + load_delta(qblock); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + return; + } + else if (warp_id == W_MMA) { + setmaxnreg_dec<88>(); + + tcgen05_alloc<1>(smem_ptr_u32(tmem_slot), TMEM_TOTAL); + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_st = tmem_base; + const uint32_t tmem_dv = tmem_st + ST_COLS; + const uint32_t tmem_dpt = tmem_dv + DV_COLS; + const uint32_t tmem_dk = tmem_dpt + ST_COLS; + const uint32_t tmem_pt_bf16 = tmem_st, tmem_dst_bf16 = tmem_dpt, tmem_dq = tmem_dpt; + const uint32_t lead = elect_one_sync() ? 1u : 0u; + + constexpr uint32_t DESC_SBO = 1024, DESC_LBO = 16; + auto make_smem_desc = [](const void* smem, uint32_t leading_byte_offset) { + return build_smem_desc_blackwell(smem_ptr_u32(smem), DESC_SBO, leading_byte_offset, + SmemSwizzleBlackwell::B128); + }; + const uint64_t desc_k = make_smem_desc(sK, DESC_LBO); + const uint64_t desc_v = make_smem_desc(sV, DESC_LBO); + const uint64_t desc_q0 = make_smem_desc(sQ[0], DESC_LBO); + const uint64_t desc_do = make_smem_desc(sDO, DESC_LBO); + const uint64_t desc_q0_mn = make_smem_desc(sQ[0], Q_SUB_COLS_BYTES); + const uint64_t desc_do_mn = make_smem_desc(sDO, Q_SUB_COLS_BYTES); + const uint64_t desc_k_mn = make_smem_desc(sK, KV_SUB_COLS_BYTES); + const uint64_t desc_dst_mn = make_smem_desc(sDST, DST_TILE_BYTES); + + constexpr uint64_t K16_COLS_DELTA = (MMA_K * (int)sizeof(__nv_bfloat16)) >> 4; + constexpr uint32_t K16_ROWS_DELTA = (MMA_K * SUB_COLS_BYTES) >> 4; + constexpr uint64_t Q_SUB_DELTA = Q_SUB_COLS_BYTES >> 4; + constexpr uint64_t KV_SUB_DELTA = KV_SUB_COLS_BYTES >> 4; + constexpr uint64_t Q_STAGE_DELTA = Q_TILE_BYTES >> 4; + + const uint32_t idesc_st_dpt = make_idesc_bf16_f32(KV_TILE, Q_TILE, false, false); + const uint32_t idesc_dv_dk = make_idesc_bf16_f32(KV_TILE, HEAD_DIM, false, true); + const uint32_t idesc_dq = make_idesc_bf16_f32(Q_TILE, HEAD_DIM, true, true); + + PhaseTracker q_full_ph; + PhaseTracker<1> do_full_ph, pt_ph, dst_ph, dq_empty_ph; + [[maybe_unused]] EmptyPhaseTracker<1> epi_empty_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + auto gemm12_st_dpt = [&](auto is_st_const, int stage) { + constexpr bool is_st = decltype(is_st_const)::value; + const uint32_t tmem_acc = is_st ? tmem_st : tmem_dpt; + const uint64_t da_base = is_st ? desc_k : desc_v; + const uint64_t db_base = is_st ? desc_q0 + (uint64_t)stage * Q_STAGE_DELTA : desc_do; + uint64_t* commit_bar = is_st ? full_bar_st : full_bar_dpt; + if constexpr (is_st) { + mbarrier_wait_parity(smem_ptr_u32(&full_bar_q[stage]), q_full_ph.get_phase()); + } else { + mbarrier_wait_parity(smem_ptr_u32(full_bar_do), do_full_ph.get_phase()); + do_full_ph.advance(); + } + + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_SUBTILE; ++ki) { + const bool accumulate = (s | ki) != 0; + tcgen05_mma_f16_ss_lead(lead, tmem_acc, da_base + s * KV_SUB_DELTA + ki * K16_COLS_DELTA, + db_base + s * Q_SUB_DELTA + ki * K16_COLS_DELTA, idesc_st_dpt, + accumulate); + } + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(commit_bar)); + }; + + auto gemm35_dv_dk = [&](auto is_dv_const, int stage, bool first) { + constexpr bool is_dv = decltype(is_dv_const)::value; + const uint32_t tmem_acc = is_dv ? tmem_dv : tmem_dk; + const uint32_t tmem_a_base = is_dv ? tmem_pt_bf16 : tmem_dst_bf16; + uint64_t db = is_dv ? desc_do_mn : desc_q0_mn + (uint64_t)stage * Q_STAGE_DELTA; + uint64_t* commit_bar = is_dv ? empty_bar_do : &empty_bar_q[stage]; + #pragma unroll + for (int q_half = 0; q_half < 2; ++q_half) { + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_Q_HALF; ++ki) { + const uint32_t tmem_a = tmem_a_base + (uint32_t)(q_half * ST_HALF_COLS + + ki * BF16X2_COLS_PER_K16); + const bool accumulate = !(first && (q_half | ki) == 0); + tcgen05_mma_f16_ts_1sm_lead(lead, tmem_acc, tmem_a, db, idesc_dv_dk, accumulate); + smem_desc_add_lo(db, K16_ROWS_DELTA); + } + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(commit_bar)); + }; + + auto gemm4_dq = [&]() { + uint64_t adst = desc_dst_mn; + uint64_t bk = desc_k_mn; + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_KV_TILE; ++ki) { + tcgen05_mma_f16_ss_lead(lead, tmem_dq, adst, bk, idesc_dq, ki != 0); + smem_desc_add_lo(adst, K16_ROWS_DELTA); + smem_desc_add_lo(bk, K16_ROWS_DELTA); + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dq)); + }; + + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + if (it.local_k2q_num != 0) { + + gemm12_st_dpt(std::true_type{}, q_full_ph.get_stage()); + + if constexpr (PERSISTENT) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + } + + gemm12_st_dpt(std::false_type{}, 0); + + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + + if constexpr (PERSISTENT) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + } + + gemm35_dv_dk(std::true_type{}, 0, true); + + for (int j = 0; j < it.local_k2q_num; ++j) { + + const int stage = q_full_ph.get_stage(); + q_full_ph.advance(); + const int next_stage = q_full_ph.get_stage(); + const bool last = j + 1 == it.local_k2q_num; + + if (!last) gemm12_st_dpt(std::true_type{}, next_stage); + + if (last) { + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dv)); + } + mbarrier_wait_parity(smem_ptr_u32(full_bar_dst), dst_ph.get_phase()); + dst_ph.advance(); + gemm35_dv_dk(std::false_type{}, stage, j == 0); + if (last) { + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dk)); + } + gemm4_dq(); + if (!last) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + gemm12_st_dpt(std::false_type{}, 0); + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + gemm35_dv_dk(std::true_type{}, 0, false); + } + } + + if constexpr (PERSISTENT) { + tcgen05_commit1_lead(lead, smem_ptr_u32(empty_bar_kv)); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + tcgen05_relinquish_alloc_permit<1>(); + bar_sync<10>(416); + tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL); + return; + } + else if (warp_id == W_SCHED) { + + if constexpr (CLC) { + setmaxnreg_dec<88>(); + int prod_stage = 0; uint32_t prod_phase = 1; + int cons_stage = 0; uint32_t cons_phase = 0; + while (true) { + if (lane == 0) + mbarrier_wait_parity_suspend(smem_ptr_u32(&clc_empty[prod_stage]), prod_phase); + __syncwarp(); + clc_arrive_expect_tx_cta(smem_ptr_u32(&clc_full[prod_stage]), 16); + if (lane == 0) + clc_try_cancel_async(smem_ptr_u32(&clc_response[prod_stage * 4]), + smem_ptr_u32(&clc_full[prod_stage])); + advance_stage_phase(prod_stage, prod_phase); + ClcTileInfo n = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, cons_stage, cons_phase, elect_one_sync()); + clc_fetch_next_tile_advance(cons_stage, cons_phase); + if (!n.valid) break; + } + + #pragma unroll + for (int st = 0; st < CLC_STAGES; ++st) { + if (lane == 0) + mbarrier_wait_parity_suspend(smem_ptr_u32(&clc_empty[prod_stage]), prod_phase); + __syncwarp(); + advance_stage_phase(prod_stage, prod_phase); + } + } else { + setmaxnreg_dec<24>(); + } + return; + } + else if (warp_id >= W_SOFTMAX0 && warp_id < W_MMA) { + setmaxnreg_inc<136>(); + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_st = tmem_base; + const uint32_t tmem_dv = tmem_st + ST_COLS; + const uint32_t tmem_dpt = tmem_dv + DV_COLS; + const uint32_t tmem_dk = tmem_dpt + ST_COLS; + const uint32_t tmem_pt_bf16 = tmem_st, tmem_dst_bf16 = tmem_dpt; + + const int softmax_warp_id = warp_id - W_SOFTMAX0; + const int lane_group = softmax_warp_id & 3; + const int col_half = softmax_warp_id >> 2; + const int kv_row = lane_group * 32 + lane; + + const uint32_t tmem_lane_base = (uint32_t)(lane_group * 32) << 16; + const uint32_t tmem_f32_offset = tmem_lane_base + (uint32_t)(col_half * ST_HALF_COLS); + const uint32_t tmem_bf16x2_offset = tmem_f32_offset; + + constexpr int CHUNK_BF16 = 16 / (int)sizeof(__nv_bfloat16); + constexpr int CHUNKS_PER_ROW = SUB_COLS_BF16 / CHUNK_BF16; + + __nv_bfloat16* sdst_row = + sDST + (size_t)col_half * KV_TILE * SUB_COLS_BF16 + kv_row * SUB_COLS_BF16; + + PhaseTracker<1> st_ph, dpt_ph, delta_ph, dv_ph, dk_ph; + PhaseTracker lse_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + + auto store_dv_dk_tile = [&](uint32_t tmem_acc, auto apply_sm_scale_const, uint8_t* bounce, + const CUtensorMap* map) { + constexpr bool apply_sm_scale = decltype(apply_sm_scale_const)::value; + uint8_t* bounce_subtile = bounce + col_half * KV_SUB_COLS_BYTES; + __nv_bfloat16* stage_row = + reinterpret_cast<__nv_bfloat16*>(bounce_subtile) + kv_row * SUB_COLS_BF16; + #pragma unroll + for (int c0 = 0; c0 < ST_HALF_COLS; c0 += 32) { + uint32_t acc_regs[32]; + tcgen05_ld_32x32b_x32(tmem_acc + tmem_f32_offset + (uint32_t)c0, acc_regs); + tcgen05_fence_before_thread_sync(); + const float* acc = reinterpret_cast(acc_regs); + #pragma unroll + for (int v = 0; v < 32 / CHUNK_BF16; ++v) { + float value[CHUNK_BF16]; + #pragma unroll + for (int e = 0; e < CHUNK_BF16; ++e) + value[e] = apply_sm_scale ? acc[v * CHUNK_BF16 + e] * sm_scale + : acc[v * CHUNK_BF16 + e]; + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(value[0], value[1]); + packed.y = cvt_f32x2_to_bf16x2(value[2], value[3]); + packed.z = cvt_f32x2_to_bf16x2(value[4], value[5]); + packed.w = cvt_f32x2_to_bf16x2(value[6], value[7]); + const int chunk = c0 / CHUNK_BF16 + v; + *reinterpret_cast(stage_row + ((chunk ^ (kv_row & 7)) * CHUNK_BF16)) = packed; + } + } + fence_proxy_async_shared_cta(); + if (col_half == 0) bar_sync<12>(128); else bar_sync<13>(128); + if (lane_group == 0 && elect_one_sync()) { + if constexpr (BHSD) + tma_store_4d(map, 0, it.kv_block_id_in_seq * KV_TILE, col_half, it.batch_head, + smem_ptr_u32(bounce_subtile)); + else + tma_store_3d(map, 0, it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * HD_SUBTILES + col_half, smem_ptr_u32(bounce_subtile)); + cp_async_bulk_commit_group(); + } + }; + + const bool kv_row_valid = kv_row < variable_block_sizes[it.kv_block_id_in_seq]; + + for (int j = 0; j < it.local_k2q_num; ++j) { + const int lse_stage = lse_ph.get_stage(); + const float2* lse2 = reinterpret_cast(sLSE + lse_stage * Q_TILE + + col_half * ST_HALF_COLS); + const float2* delta2 = reinterpret_cast(sDelta + col_half * ST_HALF_COLS); + + mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_lse[lse_stage]), lse_ph.get_phase()); + lse_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_st), st_ph.get_phase()); + st_ph.advance(); + + uint32_t st_regs[ST_HALF_COLS]; + uint32_t pt_bf16x2[ST_HALF_COLS / 2]; + float2* pt_fp32 = reinterpret_cast(st_regs); + tcgen05_ld_32x32b_x64(tmem_st + tmem_f32_offset, st_regs); + tcgen05_fence_before_thread_sync(); + const float2 scale2 = f32x2_splat(scale_log2); + #pragma unroll + for (int c = 0; c < ST_HALF_COLS / 2; ++c) { + const float2 z = ffma2(pt_fp32[c], scale2, make_float2(-lse2[c].x, -lse2[c].y)); + float2 p = make_float2(ex2_approx_f32(z.x), ex2_approx_f32(z.y)); + if (!kv_row_valid) p = make_float2(0.f, 0.f); + pt_fp32[c] = p; + pt_bf16x2[c] = cvt_f32x2_to_bf16x2(p.x, p.y); + } + tcgen05_st_32x32b_x32(tmem_pt_bf16 + tmem_bf16x2_offset, pt_bf16x2); + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(full_bar_pt)); + mbarrier_arrive(smem_ptr_u32(&empty_bar_lse[lse_stage])); + } + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_delta), delta_ph.get_phase()); + delta_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dpt), dpt_ph.get_phase()); + dpt_ph.advance(); + + uint32_t dpt_regs[ST_HALF_COLS]; + tcgen05_ld_32x32b_x64(tmem_dpt + tmem_f32_offset, dpt_regs); + tcgen05_fence_before_thread_sync(); + const float2* dpt2 = reinterpret_cast(dpt_regs); + #pragma unroll + for (int c = 0; c < ST_HALF_COLS / 2; ++c) { + const float2 ds = + fmul2(pt_fp32[c], fadd2(dpt2[c], make_float2(-delta2[c].x, -delta2[c].y))); + st_regs[c] = cvt_f32x2_to_bf16x2(ds.x, ds.y); + } + uint32_t (&dst_bf16x2)[ST_HALF_COLS / 2] = + reinterpret_cast(st_regs); + tcgen05_st_32x32b_x32(tmem_dst_bf16 + tmem_bf16x2_offset, dst_bf16x2); + const uint4* dst_chunks = reinterpret_cast(dst_bf16x2); + #pragma unroll + for (int v = 0; v < CHUNKS_PER_ROW; ++v) + *reinterpret_cast(sdst_row + (v ^ (kv_row & 7)) * CHUNK_BF16) = dst_chunks[v]; + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + fence_proxy_async_shared_cta(); + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(full_bar_dst)); + mbarrier_arrive(smem_ptr_u32(empty_bar_delta)); + } + } + + if (it.local_k2q_num != 0) { + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dv), dv_ph.get_phase()); + dv_ph.advance(); + + store_dv_dk_tile(tmem_dv, std::false_type{}, sDO, &tmap_dv); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dk), dk_ph.get_phase()); + dk_ph.advance(); + + store_dv_dk_tile(tmem_dk, std::true_type{}, sQ[0], &tmap_dk); + + if constexpr (PERSISTENT) { + if (lane_group == 0 && elect_one_sync()) cp_async_bulk_wait_group_read<0>(); + bar_sync<14>(256); + mbarrier_arrive(smem_ptr_u32(empty_bar_epi)); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + bar_sync<10>(416); + return; + } + else if (warp_id < W_SOFTMAX0) { + setmaxnreg_inc<152>(); + + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_dq = tmem_base + ST_COLS + DV_COLS; + + const int epi_warp_id = warp_id - W_EPI0; + const uint32_t tmem_lane_base = (uint32_t)(epi_warp_id * 32) << 16; + const int q_row = epi_warp_id * 32 + lane; + const bool is_leader = epi_warp_id == 0 && lane == 0; + + PhaseTracker<1> dq_full_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + if constexpr (PERSISTENT) { + if (elect_one_sync()) mbarrier_arrive(smem_ptr_u32(empty_bar_dq)); + } + + do { + + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + float* dqaccum_head = + dqaccum + (size_t)it.batch_head * num_kv_blocks_per_seq * DQ::DQ_BLOCK_ELEMS; + + for (int j = 0; j < it.local_k2q_num; ++j) { + float* dqaccum_block = dqaccum_head + (size_t)it.local_k2q_idx[j] * DQ::DQ_BLOCK_ELEMS; + + mbarrier_wait_parity(smem_ptr_u32(full_bar_dq), dq_full_ph.get_phase()); + dq_full_ph.advance(); + + uint32_t dq_regs[HEAD_DIM]; + #pragma unroll + for (int c = 0; c < HEAD_DIM / 64; ++c) + tcgen05_ld_32x32b_x64(tmem_dq + tmem_lane_base + (uint32_t)(c * 64), + reinterpret_cast(dq_regs[c * 64])); + + tcgen05_wait_ld(); + tcgen05_fence_before_thread_sync(); + if (elect_one_sync()) mbarrier_arrive(smem_ptr_u32(empty_bar_dq)); + + #pragma unroll + for (int hd_slice = 0; hd_slice < HEAD_DIM / DQ::COLS; ++hd_slice) { + const int stage_buf = hd_slice & 1; + const float4* dq_row4 = reinterpret_cast(dq_regs + hd_slice * DQ::COLS); + #pragma unroll + for (int v4 = 0; v4 < DQ::COLS / 4; ++v4) + *reinterpret_cast(sDQ_STAGE[stage_buf] + v4 * Q_TILE * 4 + q_row * 4) = + dq_row4[v4]; + fence_proxy_async_shared_cta(); + bar_sync<11>(128); + if (is_leader) { + const size_t slice_offset = (size_t)hd_slice * Q_TILE * DQ::COLS; + cpasync_reduce_bulk_add_f32(dqaccum_block + slice_offset, + smem_ptr_u32(sDQ_STAGE[stage_buf]), DQ::DQ_ONE_PUSH_BYTES); + cp_async_bulk_commit_group(); + cp_async_bulk_wait_group_read<1>(); + } + bar_sync<11>(128); + } + } + + if (it.local_k2q_num == 0) { + + if (is_leader) cp_async_bulk_wait_group_read<0>(); + bar_sync<11>(128); + static_assert(DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES == HD_SUBTILES * KV_SUB_COLS_BYTES, + "the dQ stage buffers hold one 128 kv x 128 hd bf16 tile"); + constexpr int ZERO_UINT4_PER_THREAD = + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES / (128 * 16); + uint4* zero_tile = reinterpret_cast(sDQ_STAGE_bytes) + epi_warp_id * 32 + lane; + #pragma unroll + for (int v = 0; v < ZERO_UINT4_PER_THREAD; ++v) + zero_tile[v * 128] = make_uint4(0u, 0u, 0u, 0u); + fence_proxy_async_shared_cta(); + bar_sync<11>(128); + if (is_leader) { + #pragma unroll + for (int which = 0; which < 2; ++which) { + const CUtensorMap* map = (which == 0) ? &tmap_dk : &tmap_dv; + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + const uint32_t src = smem_ptr_u32(sDQ_STAGE_bytes + (size_t)s * KV_SUB_COLS_BYTES); + if constexpr (BHSD) + tma_store_4d(map, 0, it.kv_block_id_in_seq * KV_TILE, s, it.batch_head, src); + else + tma_store_3d(map, 0, it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * HD_SUBTILES + s, src); + } + } + cp_async_bulk_commit_group(); + cp_async_bulk_wait_group_read<0>(); + } + bar_sync<11>(128); + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + + if (is_leader) cp_async_bulk_wait_group_read<0>(); + bar_sync<11>(128); + bar_sync<10>(416); + if constexpr (KERNEL_PDL) { + if (is_leader) griddepcontrol_launch_dependents(); + } + return; + } + else { + setmaxnreg_dec<24>(); + return; + } +#endif +} + +template +__global__ void __launch_bounds__(256, 1) + vsa_bwd_preprocess_kernel(const __nv_bfloat16* __restrict__ o, + const __nv_bfloat16* __restrict__ dout, + float* __restrict__ delta_rows, float* __restrict__ dqaccum, + int num_heads, int seqlen) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + const int q_block_id = (int)blockIdx.x; + const int batch_head = (int)blockIdx.y; + const int batch = batch_head / num_heads, head = batch_head % num_heads; + const int token_begin = q_block_id * Q_TILE; + + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + + float4* dqaccum_zero_destination = reinterpret_cast( + dqaccum + (size_t)batch_head * seqlen * HEAD_DIM + (size_t)q_block_id * DQ::DQ_BLOCK_ELEMS); + const float4 zero_float4 = make_float4(0.f, 0.f, 0.f, 0.f); + #pragma unroll + for (int i = 0; i < (DQ::DQ_BLOCK_ELEMS / 4) / 256; ++i) + dqaccum_zero_destination[i * 256 + threadIdx.x] = zero_float4; + + if constexpr (KERNEL_PDL) griddepcontrol_launch_dependents(); + + const int row_in_pass = (int)threadIdx.x / 16; + const int dimension_begin = ((int)threadIdx.x % 16) * 8; + #pragma unroll + for (int row_pass = 0; row_pass < Q_TILE / 16; ++row_pass) { + const int token = token_begin + row_pass * 16 + row_in_pass; + const size_t element = + token_offset(batch, head, num_heads, seqlen, token) + dimension_begin; + const uint4 o_vector = *reinterpret_cast(o + element); + const uint4 dout_vector = *reinterpret_cast(dout + element); + const __nv_bfloat162* o_pairs = reinterpret_cast(&o_vector); + const __nv_bfloat162* dout_pairs = reinterpret_cast(&dout_vector); + float delta_accumulator = 0.f; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 o_pair_as_float = __bfloat1622float2(o_pairs[i]); + const float2 dout_pair_as_float = __bfloat1622float2(dout_pairs[i]); + delta_accumulator += o_pair_as_float.x * dout_pair_as_float.x + + o_pair_as_float.y * dout_pair_as_float.y; + } + #pragma unroll + for (int shuffle_offset = 8; shuffle_offset > 0; shuffle_offset >>= 1) + delta_accumulator += __shfl_down_sync(0xffffffffu, delta_accumulator, shuffle_offset); + if (dimension_begin == 0) delta_rows[(size_t)batch_head * seqlen + token] = delta_accumulator; + } +#endif +} + +template +__global__ void __launch_bounds__(128, 1) + vsa_bwd_postprocess_kernel(const float* __restrict__ dqaccum, __nv_bfloat16* __restrict__ dq, + int num_heads, int seqlen, float sm_scale) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + extern __shared__ __align__(16) float post_smem[]; + const int q_block_id = (int)blockIdx.x; + const int batch_head = (int)blockIdx.y; + const int batch = batch_head / num_heads, head = batch_head % num_heads; + const int token_begin = q_block_id * Q_TILE; + const int thread = (int)threadIdx.x; + constexpr int CHUNK_BF16 = 16 / (int)sizeof(__nv_bfloat16); + + const float4* dqaccum_block = reinterpret_cast( + dqaccum + (size_t)batch_head * seqlen * HEAD_DIM + (size_t)q_block_id * DQ::DQ_BLOCK_ELEMS); + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + #pragma unroll + for (int i = 0; i < (DQ::DQ_BLOCK_ELEMS / 4) / 128; ++i) + cp_async_cg_16(smem_ptr_u32(post_smem + (i * 128 + thread) * 4), + dqaccum_block + i * 128 + thread); + cp_async_commit_group(); + cp_async_wait_group<0>(); + __syncthreads(); + + if constexpr (KERNEL_PDL) griddepcontrol_launch_dependents(); + + uint32_t dq_packed[HEAD_DIM / 2]; + #pragma unroll + for (int hd_slice = 0; hd_slice < HEAD_DIM / DQ::COLS; ++hd_slice) { + #pragma unroll + for (int v4 = 0; v4 < DQ::COLS / 4; ++v4) { + const float4 value = *reinterpret_cast( + post_smem + hd_slice * Q_TILE * DQ::COLS + v4 * Q_TILE * 4 + thread * 4); + const int pair_index = (hd_slice * DQ::COLS + v4 * 4) / 2; + dq_packed[pair_index + 0] = cvt_f32x2_to_bf16x2(value.x * sm_scale, value.y * sm_scale); + dq_packed[pair_index + 1] = cvt_f32x2_to_bf16x2(value.z * sm_scale, value.w * sm_scale); + } + } + __syncthreads(); + + __nv_bfloat16* dq_tile_bf16 = reinterpret_cast<__nv_bfloat16*>(post_smem); + const uint4* dq_chunks = reinterpret_cast(dq_packed); + __nv_bfloat16* dq_row_bf16 = dq_tile_bf16 + thread * HEAD_DIM; + #pragma unroll + for (int v = 0; v < HEAD_DIM / CHUNK_BF16; ++v) + *reinterpret_cast(dq_row_bf16 + (v / 8) * 64 + ((v % 8) ^ (thread & 7)) * CHUNK_BF16) = + dq_chunks[v]; + __syncthreads(); + + const int row_in_pass = thread / 16; + const int dimension_begin = (thread % 16) * 8; + #pragma unroll + for (int row_pass = 0; row_pass < Q_TILE / 8; ++row_pass) { + const int row = row_pass * 8 + row_in_pass; + const int half = dimension_begin / 64; + const int chunk = ((dimension_begin % 64) / CHUNK_BF16) ^ (row & 7); + const uint4 value = *reinterpret_cast( + dq_tile_bf16 + row * HEAD_DIM + half * 64 + chunk * CHUNK_BF16); + const size_t element = + token_offset(batch, head, num_heads, seqlen, token_begin + row) + dimension_begin; + *reinterpret_cast(dq + element) = value; + } +#endif +} + +} + +#endif diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh new file mode 100644 index 000000000..9cacbda24 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh @@ -0,0 +1,231 @@ +// block_sparse_bwd_launch_sm100a_blk128.cuh -- host surface of the 128-token VSA block-sparse +// backward drop: argument struct, workspace sizes, the support predicate and the stream-chained +// launch (preprocess -> main -> postprocess, each a programmatic dependent launch of the one +// before). Tensor maps are encoded per call (a torch caller hands us fresh pointers every time). +#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH + +#include +#include "block_sparse_bwd_kernel_sm100a_blk128.cuh" + +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif + +namespace vsa_bwd_blk128 { + +using dq_accum_t = float; + +struct BlockSparseVsaBwdArgs { + // Activations are bf16, contiguous, [B, H, S, 128] under VSA_BHSD, else [B*S, H, 128]. + // nb below = num_kv_blocks_per_seq = S / 128. + + // Forward operands and results. + const __nv_bfloat16* q; + const __nv_bfloat16* k; + const __nv_bfloat16* v; + const __nv_bfloat16* o; + // Gradient of the forward output. + const __nv_bfloat16* dout; + // [B, H, S] fp32 log-sum-exp in Triton's M form: max(qk * sm_scale * log2e) + log2(l). + const float* lse; + + // Sparsity metadata, FastVideo's invert_indices layout, in 128-token blocks. + // [B*H*nb, max_q_blocks] int32: q blocks selecting each kv block; entries past the count unread. + const int* k2q_idx; + // [B*H*nb] int32: valid entries per k2q_idx row (0 allowed: the kernel writes zero dk/dv rows). + const int* k2q_num; + // [nb] int32: valid kv tokens per block (<= 128); kv rows at or past the count are masked. + const int* variable_block_sizes; + + // Outputs, inputs' layout; every row is written, including the rows of unselected kv blocks. + __nv_bfloat16* dq; + __nv_bfloat16* dk; + __nv_bfloat16* dv; + + // Scratch, caller-allocated; byte sizes from the block_sparse_bwd_*_bytes helpers below. + // [B*H*S*128] fp32, drain-native; preprocess zeroes, main reduce-adds, postprocess reads. + dq_accum_t* dqaccum; + // [B*H*S] fp32 rowsum(bf16(o) * dout), written by the preprocess. + float* delta; + + int batch; + int num_heads; + // S = nb * 128. + int seqlen; + // Must be 128. + int head_dim; + // nb = seqlen / 128. + int num_kv_blocks_per_seq; + // k2q_idx row stride (FastVideo passes nb). + int max_q_blocks; + // Softmax scale; dq and dk carry it, dv does not. + float sm_scale; +}; + +__host__ inline size_t block_sparse_bwd_dqaccum_bytes(int batch, int num_heads, int seqlen) { + return (size_t)batch * num_heads * seqlen * HEAD_DIM * sizeof(dq_accum_t); +} +__host__ inline size_t block_sparse_bwd_delta_bytes(int batch, int num_heads, int seqlen) { + return (size_t)batch * num_heads * seqlen * sizeof(float); +} + +__host__ inline cudaError_t block_sparse_bwd_supported(const BlockSparseVsaBwdArgs& args) { + if (args.head_dim != HEAD_DIM) { + return cudaErrorInvalidValue; + } + if (args.num_kv_blocks_per_seq < 1) { + return cudaErrorInvalidValue; + } + if (args.seqlen != args.num_kv_blocks_per_seq * BLOCK) { + return cudaErrorInvalidValue; + } + if (args.batch < 1 || args.num_heads < 1 || args.max_q_blocks < 1) { + return cudaErrorInvalidValue; + } + // The main grid is (kv block, batch * head): gridDim.y is limited to 65535. + if ((long)args.batch * args.num_heads > 65535) { + return cudaErrorInvalidValue; + } + if (!std::isfinite(args.sm_scale)) { + return cudaErrorInvalidValue; + } + if (!args.q || !args.k || !args.v || !args.o || !args.dout || !args.lse) { + return cudaErrorInvalidValue; + } + if (!args.dq || !args.dk || !args.dv) { + return cudaErrorInvalidValue; + } + if (!args.k2q_idx || !args.k2q_num || !args.variable_block_sizes) { + return cudaErrorInvalidValue; + } + if (!args.dqaccum || !args.delta) { + return cudaErrorInvalidValue; + } + return cudaSuccess; +} + +// Q, K, V, dO, dK, dV tensor maps, one box_tokens x 64-hd box per TMA (two per tile): +// BSHD: 3D [64 hd, B*S tokens, H*2 hd units], strides {H*128*2, 128} bytes. +// BHSD: 4D [64 hd, S tokens, 2 hd units, B*H], strides {128*2, 128, S*128*2} bytes. +__host__ inline cudaError_t make_tma_tile_units(CUtensorMap* map, const __nv_bfloat16* ptr, int B, + int H, int S, int box_tokens) { + CUresult r; + if (VSA_BHSD) { + uint64_t gd[4] = {(uint64_t)SUB_COLS_BF16, (uint64_t)S, (uint64_t)HD_SUBTILES, (uint64_t)B * H}; + uint64_t gs[3] = {(uint64_t)HEAD_DIM * 2, (uint64_t)SUB_COLS_BYTES, (uint64_t)S * HEAD_DIM * 2}; + uint32_t bd[4] = {(uint32_t)SUB_COLS_BF16, (uint32_t)box_tokens, 1u, 1u}; + uint32_t es[4] = {1u, 1u, 1u, 1u}; + r = cuTensorMapEncodeTiled( + map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast<__nv_bfloat16*>(ptr), gd, gs, bd, es, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + } else { + uint64_t gd[3] = {(uint64_t)SUB_COLS_BF16, (uint64_t)B * S, (uint64_t)H * HD_SUBTILES}; + uint64_t gs[2] = {(uint64_t)H * HEAD_DIM * 2, (uint64_t)SUB_COLS_BYTES}; + uint32_t bd[3] = {(uint32_t)SUB_COLS_BF16, (uint32_t)box_tokens, 1u}; + uint32_t es[3] = {1u, 1u, 1u}; + r = cuTensorMapEncodeTiled( + map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast<__nv_bfloat16*>(ptr), gd, gs, bd, es, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + } + return (r == CUDA_SUCCESS) ? cudaSuccess : cudaErrorInvalidValue; +} + +// Every launch carries the programmatic-stream-serialization attribute when the kernels are +// built with KERNEL_PDL (their griddepcontrol.wait guards the predecessor's data). +__host__ inline cudaLaunchConfig_t pdl_launch_config(dim3 grid, dim3 block, size_t smem, + cudaStream_t stream, cudaLaunchAttribute* at) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = grid; + cfg.blockDim = block; + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + at->id = cudaLaunchAttributeProgrammaticStreamSerialization; + at->val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = at; + cfg.numAttrs = KERNEL_PDL ? 1 : 0; + return cfg; +} + +__host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaBwdArgs& args, + cudaStream_t stream) { + const cudaError_t supported = block_sparse_bwd_supported(args); + if (supported != cudaSuccess) { + return supported; + } + const int B = args.batch, H = args.num_heads, S = args.seqlen; + + CUtensorMap tq, tk, tv, tdo, tdk, tdv; + if (make_tma_tile_units(&tq, args.q, B, H, S, Q_TILE) != cudaSuccess || + make_tma_tile_units(&tk, args.k, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tv, args.v, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tdo, args.dout, B, H, S, Q_TILE) != cudaSuccess || + make_tma_tile_units(&tdk, args.dk, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tdv, args.dv, B, H, S, KV_TILE) != cudaSuccess) { + return cudaErrorInvalidValue; + } + + auto main_kernel = vsa_bwd_main_kernel; + auto postprocess_kernel = vsa_bwd_postprocess_kernel; + constexpr int POST_SMEM_BYTES = DQ::DQ_BLOCK_ELEMS * (int)sizeof(float); + cudaError_t e = + cudaFuncSetAttribute(main_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL); + if (e != cudaSuccess) { + return e; + } + e = cudaFuncSetAttribute(postprocess_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + POST_SMEM_BYTES); + if (e != cudaSuccess) { + return e; + } + + // Preprocess: dqaccum zero + Delta, one CTA per (q block, batch * head). + cudaLaunchAttribute pre_at[1]; + cudaLaunchConfig_t pre_cfg = pdl_launch_config( + dim3((unsigned)(S / Q_TILE), (unsigned)(B * H), 1), dim3(256, 1, 1), 0, stream, pre_at); + e = cudaLaunchKernelEx(&pre_cfg, vsa_bwd_preprocess_kernel, args.o, args.dout, + args.delta, args.dqaccum, H, S); + if (e != cudaSuccess) { + return e; + } + + // Main: one CTA per (kv block, batch * head). + const float scale_log2 = args.sm_scale * 1.4426950408889634f; + cudaLaunchAttribute main_at[2]; + cudaLaunchConfig_t main_cfg = {}; + main_cfg.gridDim = dim3((unsigned)args.num_kv_blocks_per_seq, (unsigned)(B * H), 1); + main_cfg.blockDim = dim3(N_WARPS * 32, 1, 1); + main_cfg.dynamicSmemBytes = SMEM_TOTAL; + main_cfg.stream = stream; + main_at[0].id = cudaLaunchAttributeClusterDimension; + main_at[0].val.clusterDim.x = 1; + main_at[0].val.clusterDim.y = 1; + main_at[0].val.clusterDim.z = 1; + main_at[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + main_at[1].val.programmaticStreamSerializationAllowed = 1; + main_cfg.attrs = main_at; + main_cfg.numAttrs = KERNEL_PDL ? 2 : 1; + const int* workitem_remap = nullptr; + e = cudaLaunchKernelEx(&main_cfg, main_kernel, tq, tk, tv, tdo, tdk, tdv, args.dqaccum, args.lse, + args.delta, args.k2q_idx, args.k2q_num, args.max_q_blocks, workitem_remap, + args.variable_block_sizes, B, H, S, scale_log2, args.sm_scale); + if (e != cudaSuccess) { + return e; + } + + // Postprocess: drain-native dqaccum -> bf16 dq, one CTA per (q block, batch * head). + cudaLaunchAttribute post_at[1]; + cudaLaunchConfig_t post_cfg = + pdl_launch_config(dim3((unsigned)(S / Q_TILE), (unsigned)(B * H), 1), dim3(128, 1, 1), + POST_SMEM_BYTES, stream, post_at); + return cudaLaunchKernelEx(&post_cfg, postprocess_kernel, args.dqaccum, args.dq, H, S, + args.sm_scale); +} + +} // namespace vsa_bwd_blk128 + +using namespace vsa_bwd_blk128; + +#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu index 929979377..8397ebf56 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu @@ -10,6 +10,9 @@ // // The layout is fixed at compile time: VSA_BHSD true -> [B, H, S, 128] (FastVideo's build), // false -> [B, S, H, 128] (repo native; the kernel addresses it as [B*S tokens, H, 128]). +// VSA_BLK128 picks the kernel: false -> 64-token sparse blocks (block_sparse_sm100a_bwd), true -> +// 128-token (block_sparse_sm100a_blk128_bwd; block_sparse_bwd_blk128_sm100a.cu re-includes this +// file with it set). The python backend picks by the metadata's block size. #include #include @@ -17,7 +20,16 @@ #include +#ifndef VSA_BLK128 +#define VSA_BLK128 false +#endif +#if VSA_BLK128 +#include "block_sparse_bwd_launch_sm100a_blk128.cuh" +#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_blk128_bwd +#else #include "block_sparse_bwd_launch_sm100a.cuh" +#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_bwd +#endif namespace { @@ -49,7 +61,7 @@ __nv_bfloat16* bf16_ptr(const torch::Tensor& t) { } // namespace // Returns {dq, dk, dv}: bf16, each with the shape and layout of q, k, v respectively. -std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch::Tensor q, +std::vector BLOCK_SPARSE_SM100A_BWD(torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, @@ -105,20 +117,23 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: auto dv = torch::empty_like(v); // Workspace: torch::empty is enough. The preprocess kernel zeroes dqaccum and fully writes - // qt, dot and delta before the main kernel reads them; it also zeroes the dk/dv rows of kv - // blocks that no q block selects, so the empty_like outputs above come back fully defined. + // delta (and, at 64-token blocks, qt and dot) before the main kernel reads them. The dk/dv rows + // of kv blocks that no q block selects are zeroed by the preprocess (64) or by the main kernel + // (128), so the empty_like outputs above come back fully defined. const auto bytes = q.options().dtype(at::kByte); const int b = (int)B, h = (int)H, s = (int)S; auto dqaccum = torch::empty({(int64_t)block_sparse_bwd_dqaccum_bytes(b, h, s)}, bytes); - auto qt = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); - auto dot = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); auto delta = torch::empty({(int64_t)block_sparse_bwd_delta_bytes(b, h, s)}, bytes); +#if !VSA_BLK128 + auto qt = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); + auto dot = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); torch::Tensor order; const bool device_order = num_kv_blocks_per_seq >= ORDER_MIN_KV_BLOCKS; if (device_order) { order = torch::empty({(int64_t)block_sparse_bwd_order_bytes(b, h, (int)num_kv_blocks_per_seq)}, bytes); } +#endif BlockSparseVsaBwdArgs a{}; a.q = bf16_ptr(q); @@ -133,12 +148,14 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: a.k2q_idx = k2q_idx.data_ptr(); a.k2q_num = k2q_num.data_ptr(); a.variable_block_sizes = variable_block_sizes.data_ptr(); - a.workitem_remap = nullptr; - a.order_workspace = device_order ? reinterpret_cast(order.data_ptr()) : nullptr; a.dqaccum = reinterpret_cast(dqaccum.data_ptr()); - a.qt = bf16_ptr(qt); - a.dot = bf16_ptr(dot); a.delta = reinterpret_cast(delta.data_ptr()); +#if !VSA_BLK128 + a.workitem_remap = nullptr; + a.order_workspace = device_order ? reinterpret_cast(order.data_ptr()) : nullptr; + a.qt = bf16_ptr(qt); + a.dot = bf16_ptr(dot); +#endif a.batch = b; a.num_heads = h; a.seqlen = s; @@ -149,16 +166,17 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: // Report an unsupported regime loudly rather than returning plausible-looking wrong values. TORCH_CHECK(block_sparse_bwd_supported(a) == cudaSuccess, - "block_sparse_sm100a_bwd: unsupported configuration -- requires head_dim==", HEAD_DIM, + "block_sparse_sm100a_bwd (", BLOCK, + "-token blocks): unsupported configuration -- requires head_dim==", HEAD_DIM, ", seqlen == num_kv_blocks_per_seq*", BLOCK, - " with seqlen % 128 == 0, " + " with seqlen % 128 == 0, batch*heads <= 65535, " "max_q_blocks >= 1 and a finite sm_scale. Got head_dim=", D, " num_kv_blocks_per_seq=", num_kv_blocks_per_seq, " seqlen=", S, " max_q_blocks=", max_q_blocks, " sm_scale=", sm_scale); const cudaError_t err = launch_block_sparse_bwd_sm100a(a, at::cuda::getCurrentCUDAStream()); - TORCH_CHECK(err == cudaSuccess, - "block_sparse_sm100a_bwd launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, "block_sparse_sm100a_bwd (", BLOCK, + "-token blocks) launch failed: ", cudaGetErrorString(err)); return {dq, dk, dv}; } diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh index 545e2e3e6..895519dda 100644 --- a/fastvideo-kernel/csrc/attention/primitives.cuh +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -985,3 +985,30 @@ uint64_t make_l2cache_policy_fractional_evict_last_unchanged(float fraction_f32) : "f"(fraction_f32)); return policy; } + +__device__ __forceinline__ +void cp_async_cg_16(uint32_t smem_dst, const void* gmem_src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + :: "r"(smem_dst), "l"(gmem_src) : "memory"); +} + +__device__ __forceinline__ +void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ +void cp_async_wait_group() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +__device__ __forceinline__ +void griddepcontrol_wait() { + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +} + +__device__ __forceinline__ +void griddepcontrol_launch_dependents() { + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +} diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 45d5744bb..9f839dd62 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -46,6 +46,10 @@ extern std::vector block_sparse_sm100a_bwd( torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, torch::Tensor variable_block_sizes, double sm_scale); +extern std::vector block_sparse_sm100a_blk128_bwd( + torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, + torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num, + torch::Tensor variable_block_sizes, double sm_scale); #endif PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -61,6 +65,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("block_sparse_sm100a_bwd", torch::wrap_pybind_function(block_sparse_sm100a_bwd), "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)"); + m.def("block_sparse_sm100a_blk128_bwd", + torch::wrap_pybind_function(block_sparse_sm100a_blk128_bwd), + "VSA block-sparse attention backward, 128-token blocks (Blackwell sm100a/sm103a)"); #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 0b2c56df1..b463932ca 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -52,13 +52,13 @@ def _force_tk() -> bool: def _force_sm100a() -> bool: """True iff the data-center Blackwell forward is explicitly opted into. - Opt-in only (same legacy-named env the H3 backend honors): the extension is - forward-only, so this routing pairs it with the Triton backward -- its lse - is already in Triton's M format. Honored only when - ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token - metadata falls through to the default selection; unsupported 128-token - metadata raises because Triton has no compatible fallback. - ``FASTVIDEO_VSA_TRITON`` still wins. + Opt-in only (same legacy-named env the H3 backend honors). Honored only when + ``block_sparse_attn_sm100a.is_supported`` passes; the backward then runs the + sm_100a/sm_103a CUDA backward when ``block_sparse_attn_bwd_sm100a.is_supported`` + passes and Triton otherwise (the forward's lse is already in Triton's M + format). Unsupported 64-token metadata falls through to the default + selection; unsupported 128-token metadata raises because Triton has no + compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. """ return _env_bool("FASTVIDEO_VSA_SM100A") @@ -389,11 +389,12 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co # Data-center Blackwell backend custom op (index-native; legacy sm100a API name) # # Forward runs the sm_100a/sm_103a CUDA extension. Backward runs the sm_100a CUDA -# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64-token blocks, -# sm_100a device, extension built with the op) and the Triton kernels otherwise. -# The native forward emits lse in exactly Triton's M format (max*log2e + -# log2(l)), so either pairing needs no conversion. Both backwards are -# hardcoded to 64-token blocks, hence the block-size assert below. +# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64- or 128-token +# blocks, data-center Blackwell device, extension built with that block's op) and +# the Triton kernels otherwise. The native forward emits lse in exactly Triton's M +# format (max*log2e + log2(l)), so either pairing needs no conversion. The Triton +# backward is hardcoded to 64-token blocks, so unsupported 128-token metadata +# raises instead of falling back. # --------------------------------------------------------------------------- @@ -493,16 +494,17 @@ def _setup_context_sm100a(ctx, inputs, output): def _backward_sm100a(ctx, grad_o, grad_M): q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors - block = q.shape[2] // variable_block_sizes.numel() - if block != 64: - raise RuntimeError( - "block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with a " - f"backward that is hardcoded to 64-token blocks; got {block}. " - "Run 128-token-block metadata without grad, or use the Triton forward.") if _sm100a_backward_is_supported(q, variable_block_sizes): dq, dk, dv = block_sparse_attn_backward_sm100a(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) else: + block = q.shape[2] // variable_block_sizes.numel() + if block != 64: + raise RuntimeError( + "block_sparse_attn_sm100a backward: no sm_100a/sm_103a backward for " + f"{block}-token blocks on this build/device, and the Triton backward is " + "hardcoded to 64-token blocks. Run this metadata without grad, or build " + "the extension with block_sparse_sm100a_blk128_bwd.") dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) return dq, dk, dv, None, None, None @@ -536,8 +538,9 @@ def block_sparse_attn_from_indices( # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. - # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward - # (Triton backward). The environment name is retained for compatibility. + # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward (CUDA + # backward when supported, else Triton). The environment name is retained + # for compatibility. # Unsupported 64-token metadata falls through; unsupported 128-token # metadata raises because Triton cannot consume it. # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py index ea073d868..cb2b1c7ca 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -4,8 +4,10 @@ Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and FastVideo's k2q index metadata, returns ``(dq, dk, dv)`` in bf16 with the inputs' layout and the -Triton backward's scaling (dq and dk carry sm_scale, dv does not). 64-token blocks only; any -other configuration falls back to Triton via ``is_supported``. +Triton backward's scaling (dq and dk carry sm_scale, dv does not). Two kernels, picked by the +metadata's block size: 64-token blocks (``block_sparse_sm100a_bwd``) and 128-token blocks +(``block_sparse_sm100a_blk128_bwd``); any other configuration falls back to Triton via +``is_supported``. """ from typing import Tuple @@ -16,45 +18,64 @@ try: # The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it # (its __init__ is empty, so hasattr on the package fails with the kernel built and present). from fastvideo_kernel._C import fastvideo_kernel_ops as _C - _BWD = getattr(_C, "block_sparse_sm100a_bwd", None) - _HAS_VSA_BWD_SM100A = _BWD is not None + _BWD_BY_BLOCK = { + 64: getattr(_C, "block_sparse_sm100a_bwd", None), + 128: getattr(_C, "block_sparse_sm100a_blk128_bwd", None), + } + _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) except ImportError: # pragma: no cover - extension not built _C = None - _BWD = None + _BWD_BY_BLOCK = {} _HAS_VSA_BWD_SM100A = False _SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 -BLOCK = 64 +BLOCKS = (64, 128) # Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true). BHSD = True def set_extension(module) -> None: - """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd``. + """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd`` and / or + ``block_sparse_sm100a_blk128_bwd``. - A standalone build of ``block_sparse_bwd_sm100a.cu`` (for example through + A standalone build of the binding .cu files (for example through ``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so the backend can be exercised without rebuilding the fastvideo_kernel wheel. """ - global _C, _BWD, _HAS_VSA_BWD_SM100A + global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_SM100A _C = module - _BWD = getattr(module, "block_sparse_sm100a_bwd", None) - _HAS_VSA_BWD_SM100A = _BWD is not None + _BWD_BY_BLOCK = { + 64: getattr(module, "block_sparse_sm100a_bwd", None), + 128: getattr(module, "block_sparse_sm100a_blk128_bwd", None), + } + _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) def _seqlen(q: torch.Tensor) -> int: return q.shape[2] if BHSD else q.shape[1] +def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + """The block size the metadata implies (seqlen / num_blocks), 0 when it is not integral. + + A plain int: under dynamic-shape torch.compile the shapes are SymInts, which cannot key + ``_BWD_BY_BLOCK``; ``int()`` specializes the value with a shape guard. + """ + num_blocks = variable_block_sizes.numel() + seqlen = _seqlen(q) + return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks) + + def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: """True iff this build can run these tensors; otherwise the caller uses Triton. Static facts only (shapes, dtypes, arch, layout), never tensor contents, so it is cheap - enough for a per-layer dispatch path. The kernel is fixed at 64-token blocks with - head_dim 128 and needs seqlen == 64 * num_blocks with an even num_blocks (its preprocess - works in 128-token blocks). Per-row k2q counts may be anything in [0, num_q_blocks], - including 0: unselected kv blocks get exactly-zero dk/dv rows. + enough for a per-layer dispatch path. Both kernels take head_dim 128 and seqlen == block * + num_blocks; the 64-token kernel also needs an even num_blocks (its preprocess works in + 128-token blocks); batch * heads is the launch grid's y extent, so it must be <= 65535. + Per-row k2q counts may be anything in [0, num_q_blocks], including 0: unselected kv blocks + get exactly-zero dk/dv rows. """ if not _HAS_VSA_BWD_SM100A or not q.is_cuda: return False @@ -64,14 +85,16 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: return False if not q.is_contiguous(): return False + if q.shape[0] * (q.shape[1] if BHSD else q.shape[2]) > 65535: + return False # Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving. if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64): return False - num_blocks = variable_block_sizes.numel() - if num_blocks == 0 or num_blocks % 2 != 0: + block = _block_size(q, variable_block_sizes) + if _BWD_BY_BLOCK.get(block) is None: return False - if _seqlen(q) != BLOCK * num_blocks: + if block == 64 and variable_block_sizes.numel() % 2 != 0: return False return True @@ -89,15 +112,21 @@ def block_sparse_attn_backward_sm100a_from_k2q( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Backward from k2q metadata already in hand (``invert_indices`` layout). - ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q64 - block ids, ``k2q_num`` ``[B, H, num_kv_blocks]``; entries past a row's count are never read. + ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q + block ids (64- or 128-token blocks, as the metadata's block size), ``k2q_num`` + ``[B, H, num_kv_blocks]``; entries past a row's count are never read. """ + block = _block_size(q, variable_block_sizes) + backward = _BWD_BY_BLOCK.get(block) + if backward is None: + raise RuntimeError(f"block_sparse_attn_backward_sm100a: no kernel for {block}-token " + f"blocks (built: {sorted(b for b, f in _BWD_BY_BLOCK.items() if f)})") sm_scale = 1.0 / (q.shape[-1]**0.5) idx = k2q_idx.to(torch.int32).contiguous() num = k2q_num.to(torch.int32).contiguous() vbs = variable_block_sizes.to(torch.int32).contiguous() - res = _BWD(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), - o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale) + res = backward(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), + o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale) return res[0], res[1], res[2] diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index 84872f52d..b2d2df6fe 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -8,7 +8,7 @@ A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This mo forward: it returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward`` writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so both ``block_sparse_attn_backward_triton`` and the sm_100a CUDA backward -(``block_sparse_attn_bwd_sm100a``, 64-token blocks, sm_100a devices only) run against it +(``block_sparse_attn_bwd_sm100a``, 64- and 128-token blocks, sm_100a / sm_103a) run against it unchanged. The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks @@ -43,9 +43,10 @@ BHSD = True def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + # A plain int: dynamic-shape torch.compile hands SymInt shapes, which cannot key _FWD_BY_BLOCK. num_blocks = variable_block_sizes.numel() seqlen = q.shape[2] if BHSD else q.shape[1] - return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks + return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks) def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: diff --git a/tests/test_block_sparse_bwd_sm100a.py b/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py similarity index 69% rename from tests/test_block_sparse_bwd_sm100a.py rename to fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py index b9daea240..832bef8df 100644 --- a/tests/test_block_sparse_bwd_sm100a.py +++ b/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Correctness tests for the sm_100a CUDA block-sparse VSA backward. +"""Correctness tests for the sm_100a CUDA block-sparse VSA backward, 64- and 128-token blocks. Reference: fp32 dense attention restricted to the selected blocks, keys past variable_block_sizes masked to -inf, differentiated with torch autograd on fp32 copies of the bf16 inputs with loss = (out * grad_o).sum(). The kernel is fed the reference's lse in Triton's M format (logsumexp * log2e) and the reference output rounded to bf16 as ``o`` -- exactly what the sm_100a forward hands it in FastVideo. The k2q inversion is done here in torch (a stable -sort) so the tests do not need Triton. +sort) so the tests do not need Triton. Every case runs once per built block size. Run with: python -m pytest tests/test_block_sparse_bwd_sm100a.py -v """ @@ -20,13 +20,12 @@ import torch from fastvideo_kernel import block_sparse_attn_bwd_sm100a as bwd HEAD_DIM = 128 -BLOCK = 64 LOG2E = 1.4426950408889634 # Per-tensor tolerances on the bf16 outputs against the fp32 reference. Measured over all 15 -# cases on GB200 (2026-09-03, VSA_BWD_TEST_VERBOSE=1): max|diff|/max|ref| up to 3.0e-3 (dq) and -# 5.3e-3 (dk, dv); mean|diff| up to 2.8e-4 against mean|ref| of 0.06-0.12. The bounds below leave -# about 2x (rel max) and 3.5x (mean abs) headroom. +# blk64 cases on GB200 (2026-09-03, VSA_BWD_TEST_VERBOSE=1): max|diff|/max|ref| up to 3.0e-3 (dq) +# and 5.3e-3 (dk, dv); mean|diff| up to 2.8e-4 against mean|ref| of 0.06-0.12. The bounds below +# leave about 2x (rel max) and 3.5x (mean abs) headroom; the blk128 cases sit in the same range. REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref| MEAN_ABS_TOL = 1e-3 # mean|got - ref| @@ -34,18 +33,25 @@ pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} or not bwd._HAS_VSA_BWD_SM100A, reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel " - "extension built with block_sparse_sm100a_bwd", + "extension built with block_sparse_sm100a_bwd / block_sparse_sm100a_blk128_bwd", ) -def make_case(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): +@pytest.fixture(params=bwd.BLOCKS, ids=lambda b: f"blk{b}") +def block(request): + if bwd._BWD_BY_BLOCK.get(request.param) is None: + pytest.skip(f"no {request.param}-token backward in this build") + return request.param + + +def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): """Random bf16 q/k/v/grad_o plus q2k metadata [B, H, Nq, topk] / [B, H, Nq] and vbs. ``kv_pool`` restricts the kv blocks a q block may select (default: all), which is how the zero-count case leaves some kv blocks unselected by every q block. """ torch.manual_seed(seed) - S = num_blocks * BLOCK + S = num_blocks * block shape = (batch, heads, S, HEAD_DIM) if bwd.BHSD else (batch, S, heads, HEAD_DIM) q, k, v, grad_o = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(4)) @@ -59,10 +65,10 @@ def make_case(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_p num = torch.full((batch, heads, num_blocks), topk, dtype=torch.int32, device="cuda") if ragged: - vbs = torch.randint(BLOCK // 2, BLOCK + 1, (num_blocks, ), dtype=torch.int32, + vbs = torch.randint(block // 2, block + 1, (num_blocks, ), dtype=torch.int32, device="cuda") else: - vbs = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda") return q, k, v, grad_o, idx, num, vbs @@ -92,7 +98,7 @@ def invert_indices_torch(q2k_idx, q2k_num, num_kv_blocks, pad_value=0): counts.to(torch.int32).view(B, H, num_kv_blocks)) -def reference(q, k, v, grad_o, idx, num, vbs): +def reference(block, q, k, v, grad_o, idx, num, vbs): """fp32 masked-dense autograd reference. Returns (o bf16, lse fp32 [B, H, S] in M format, dq, dk, dv fp32) -- o and the grads in @@ -112,7 +118,7 @@ def reference(q, k, v, grad_o, idx, num, vbs): for j in range(int(num[b, h, qb])): kb = int(idx[b, h, qb, j]) valid = int(vbs[kb]) - keep[b, h, qb * BLOCK:(qb + 1) * BLOCK, kb * BLOCK:kb * BLOCK + valid] = True + keep[b, h, qb * block:(qb + 1) * block, kb * block:kb * block + valid] = True q32, k32, v32 = (t.detach().float().requires_grad_(True) for t in (q, k, v)) scores = (q32 @ k32.transpose(-1, -2)) * scale @@ -145,12 +151,13 @@ def check_close(name, got, ref): f"max|ref| = {ref.abs().max().item():.3e}, mean|ref| = {ref.abs().mean().item():.3e}") -def run_and_compare(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): - q, k, v, grad_o, idx, num, vbs = make_case(num_blocks=num_blocks, topk=topk, heads=heads, - batch=batch, ragged=ragged, seed=seed, - kv_pool=kv_pool) +def run_and_compare(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, + kv_pool=None): + q, k, v, grad_o, idx, num, vbs = make_case(block, num_blocks=num_blocks, topk=topk, + heads=heads, batch=batch, ragged=ragged, + seed=seed, kv_pool=kv_pool) assert bwd.is_supported(q, vbs) - o, lse, ref_dq, ref_dk, ref_dv = reference(q, k, v, grad_o, idx, num, vbs) + o, lse, ref_dq, ref_dk, ref_dv = reference(block, q, k, v, grad_o, idx, num, vbs) k2q_idx, k2q_num = invert_indices_torch(idx, num, num_blocks) dq, dk, dv = bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num, vbs) @@ -165,29 +172,29 @@ def _to_bhsd(t): return t if bwd.BHSD else t.transpose(1, 2) -def test_backward_matches_reference(): - run_and_compare() +def test_backward_matches_reference(block): + run_and_compare(block) -def test_batch_two(): - run_and_compare(batch=2) +def test_batch_two(block): + run_and_compare(block, batch=2) -def test_ragged_block_sizes(): +def test_ragged_block_sizes(block): """variable_block_sizes is what FastVideo always passes; padded keys must be masked.""" - run_and_compare(ragged=True) + run_and_compare(block, ragged=True) -def test_ragged_padded_key_rows_are_zero(): +def test_ragged_padded_key_rows_are_zero(block): """Keys at or past a block's count get P^T = 0 in-kernel, so their dk/dv rows are exactly 0 (Triton's backward stores zeros there as well).""" - (dq, dk, dv), _, vbs = run_and_compare(ragged=True, seed=1) + (dq, dk, dv), _, vbs = run_and_compare(block, ragged=True, seed=1) dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() checked = 0 for kb in range(vbs.numel()): - if int(vbs[kb]) == BLOCK: + if int(vbs[kb]) == block: continue # a full block has no padded rows - rows = slice(kb * BLOCK + int(vbs[kb]), (kb + 1) * BLOCK) + rows = slice(kb * block + int(vbs[kb]), (kb + 1) * block) assert dk[:, :, rows].abs().max().item() == 0.0, f"dk: padded rows of kv block {kb}" assert dv[:, :, rows].abs().max().item() == 0.0, f"dv: padded rows of kv block {kb}" checked += 1 @@ -195,27 +202,31 @@ def test_ragged_padded_key_rows_are_zero(): @pytest.mark.parametrize("topk", [1, 2, 3, 5, 7]) -def test_topk_not_a_multiple_of_the_quad(topk): - """The kernel walks each kv block's q list in quads of 4; a ragged tail quad must be exact.""" - run_and_compare(ragged=True, num_blocks=8, topk=topk) +def test_topk_not_a_multiple_of_the_quad(block, topk): + """blk64 walks each kv block's q list in quads of 4 (a ragged tail quad must be exact); + blk128 one q block per step (an odd list length must be exact).""" + run_and_compare(block, ragged=True, num_blocks=8, topk=topk) -@pytest.mark.parametrize("num_blocks", [2, 4, 8, 16]) -def test_sequence_lengths(num_blocks): - run_and_compare(ragged=True, num_blocks=num_blocks, topk=min(3, num_blocks)) +@pytest.mark.parametrize("num_blocks", [1, 2, 4, 8, 16]) +def test_sequence_lengths(block, num_blocks): + if block == 64 and num_blocks % 2: + pytest.skip("blk64 needs an even block count (128-token preprocess)") + run_and_compare(block, ragged=True, num_blocks=num_blocks, topk=min(3, num_blocks)) -def test_zero_count_kv_blocks(): - """kv blocks no q block selects: the main kernel skips them and the preprocess must write - their dk/dv rows as exact zeros (the outputs are empty_like), every other block exact.""" +def test_zero_count_kv_blocks(block): + """kv blocks no q block selects: their dk/dv rows must come back as exact zeros (the outputs + are empty_like; blk64's preprocess and blk128's main kernel write them), every other block + exact.""" num_blocks = 8 excluded = (2, 5) pool = [kb for kb in range(num_blocks) if kb not in excluded] - (dq, dk, dv), _, _ = run_and_compare(num_blocks=num_blocks, topk=4, ragged=True, + (dq, dk, dv), _, _ = run_and_compare(block, num_blocks=num_blocks, topk=4, ragged=True, kv_pool=pool) dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() for kb in excluded: - rows = slice(kb * BLOCK, (kb + 1) * BLOCK) + rows = slice(kb * block, (kb + 1) * block) assert dk[:, :, rows].abs().max().item() == 0.0, f"dk: unselected kv block {kb}" assert dv[:, :, rows].abs().max().item() == 0.0, f"dv: unselected kv block {kb}" @@ -223,7 +234,7 @@ def test_zero_count_kv_blocks(): def test_invert_indices_torch_matches_q2k(): """Guards the test's own k2q builder: every (q, kv) pair appears once, counts add up.""" num_blocks = 8 - _, _, _, _, idx, num, _ = make_case(num_blocks=num_blocks, topk=3, heads=2, batch=2) + _, _, _, _, idx, num, _ = make_case(64, num_blocks=num_blocks, topk=3, heads=2, batch=2) k2q_idx, k2q_num = invert_indices_torch(idx, num, num_blocks) B, H, Nq, Mk = idx.shape assert k2q_idx.shape == (B, H, num_blocks, Nq) and k2q_num.shape == (B, H, num_blocks) @@ -235,24 +246,31 @@ def test_invert_indices_torch_matches_q2k(): assert listed == sorted(listed) -def test_unsupported_is_rejected(): - q, k, v, grad_o, idx, num, vbs = make_case() +def test_unsupported_is_rejected(block): + q, k, v, grad_o, idx, num, vbs = make_case(block) assert not bwd.is_supported(q.float(), vbs) # wrong dtype assert not bwd.is_supported(q[..., :64].contiguous(), vbs) # wrong head_dim - seven = torch.full((7, ), 64, dtype=torch.int32, device="cuda") - assert not bwd.is_supported(q, seven) # seqlen != 64 * num_blocks + seven = torch.full((7, ), block, dtype=torch.int32, device="cuda") + assert not bwd.is_supported(q, seven) # seqlen != block * num_blocks q500 = (q[:, :, :500] if bwd.BHSD else q[:, :500]).contiguous() - assert not bwd.is_supported(q500, vbs) # S not a multiple of 64 - # An odd block count (S % 128 != 0) is refused statically too, so FastVideo falls back to - # Triton instead of tripping the binding's check. - q448 = (q[:, :, :448] if bwd.BHSD else q[:, :448]).contiguous() - assert not bwd.is_supported(q448, seven) + assert not bwd.is_supported(q500, vbs) # S not a multiple of the block + if block == 64: + # An odd block count (S % 128 != 0) is refused statically too, so FastVideo falls back + # to Triton instead of tripping the binding's check. + q448 = (q[:, :, :448] if bwd.BHSD else q[:, :448]).contiguous() + assert not bwd.is_supported(q448, seven) + # batch * heads is the launch grid's y extent: 65536 (batch, head) pairs are refused. + wide = (1, 65536, block, q.shape[-1]) if bwd.BHSD else (1, block, 65536, q.shape[-1]) + q_wide = torch.empty(wide, dtype=torch.bfloat16, device="cuda") + one = torch.full((1, ), block, dtype=torch.int32, device="cuda") + assert not bwd.is_supported(q_wide, one) + del q_wide # The binding itself refuses bad dtypes before touching the GPU. k2q_idx, k2q_num = invert_indices_torch(idx, num, vbs.numel()) o = torch.zeros_like(q) heads = q.shape[1] if bwd.BHSD else q.shape[2] - lse = torch.zeros((q.shape[0], heads, vbs.numel() * BLOCK), dtype=torch.float32, + lse = torch.zeros((q.shape[0], heads, vbs.numel() * block), dtype=torch.float32, device="cuda") with pytest.raises(RuntimeError): bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o.float(), q.float(), k.float(), diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py similarity index 80% rename from tests/test_block_sparse_sm100a_dispatch.py rename to fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py index 122aec2ea..73cca33fa 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py @@ -4,8 +4,9 @@ ``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when ``block_sparse_attn_sm100a.is_supported`` passes. Its backward is the sm_100a CUDA backward where that op is built and ``block_sparse_attn_bwd_sm100a.is_supported`` -passes (64-token blocks on an sm_100a device), the Triton backward otherwise; the -sm_100a lse is already in Triton's M format, so either pairing needs no conversion. +passes (64- or 128-token blocks on a data-center Blackwell device), the Triton backward +otherwise (64-token blocks only); the sm_100a lse is already in Triton's M format, so +either pairing needs no conversion. Everything else -- env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- must keep the pre-existing selection, bit-for-bit. @@ -229,10 +230,50 @@ def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): f"max|diff|/max|ref|={rel_max:.3e} mean|diff|={mean_abs:.3e}" -def test_blk128_backward_raises(monkeypatch): - """128-token blocks: forward runs, backward refuses (both backwards are 64-block only).""" +def _blk128_triton_reference_grads(q, k, v, idx, num, vbs, monkeypatch): + """All-Triton grads for 128-token metadata: the logical [B, H, Nq, Nk] map goes through + ``block_sparse_attn_128`` (Triton via the 128 -> 64 tile expansion, autograd-capable).""" + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128 + + B, H, Nq, Mk = idx.shape + rows = B * H * Nq + block_map = torch.zeros((rows, vbs.numel()), dtype=torch.bool, device="cuda") + valid = torch.arange(Mk, device="cuda").view(1, Mk) < num.view(rows, 1) + row_of = torch.arange(rows, device="cuda").view(rows, 1).expand(rows, Mk) + block_map[row_of[valid], idx.view(rows, Mk)[valid].long()] = True + block_map = block_map.view(B, H, Nq, vbs.numel()) + q2, k2, v2 = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False) + out2, _ = block_sparse_attn_128(q2, k2, v2, block_map, vbs) + out2.float().square().sum().backward() + monkeypatch.delenv("FASTVIDEO_VSA_TRITON") + return [t.grad.float() for t in (q2, k2, v2)] + + +def test_blk128_backward(monkeypatch): + """128-token blocks: with the 128-token backward built the sm_100a route runs it (the Triton + backward must not be entered) and its grads match the all-Triton 128 wrapper; without it the + backward raises instead of silently falling back to a 64-token kernel.""" + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + + dispatch = importlib.import_module("fastvideo_kernel.block_sparse_attn") monkeypatch.setenv(ENV, "1") q, k, v, idx, num, vbs = make_case(128, requires_grad=True) + if not vsa_bwd.is_supported(q, vbs): + out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) + with pytest.raises(RuntimeError, match="128-token"): + out.float().sum().backward() + return + + # The Triton reference first: the guard below disables the Triton backward for the route. + ref = _blk128_triton_reference_grads(q, k, v, idx, num, vbs, monkeypatch) + + def no_triton_backward(*args, **kwargs): + raise AssertionError("Triton backward entered on a supported 128-token input") + + monkeypatch.setattr(dispatch, "block_sparse_attn_backward_triton", no_triton_backward) out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) - with pytest.raises(RuntimeError, match="64-token"): - out.float().sum().backward() + out.float().square().sum().backward() + got = [t.grad.float() for t in (q, k, v)] + _assert_grads_close(got, ref) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index e37e6d5a1..a6f195ece 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -420,7 +420,7 @@ def _sm100a_unavailable_reason(sm100a_mod: Any, query_bhsd: torch.Tensor, variab if sm100a_mod is None: return "fastvideo_kernel.block_sparse_attn_sm100a is not installed" if grad_mode: - return "inputs require grad and the sm_100a/sm_103a kernel is forward-only; grad paths keep Triton" + return "inputs require grad; this backend keeps grad paths on Triton" if not sm100a_mod.is_supported(query_bhsd, variable_block_sizes): return ("block_sparse_attn_sm100a.is_supported returned False (needs an sm_100 or sm_103 device, a built " "extension, bf16, head_dim 128, an even tile count, and integer tile sizes)") diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 2deaed9f4..5b03fabe9 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -120,7 +120,7 @@ def test_reason_covers_every_precondition(): vbs = torch.full((2, ), 64, dtype=torch.long) assert "not installed" in _sm100a_unavailable_reason(None, q, vbs, grad_mode=False) ok = _FakeSm100a(supported=True) - assert "forward-only" in _sm100a_unavailable_reason(ok, q, vbs, grad_mode=True) + assert "require grad" in _sm100a_unavailable_reason(ok, q, vbs, grad_mode=True) bad = _FakeSm100a(supported=False) assert "is_supported" in _sm100a_unavailable_reason(bad, q, vbs, grad_mode=False) assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None From 9edc8adf5f6f8e76072297af86d5b05324b54630 Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Wed, 30 Sep 2026 00:15:39 -0700 Subject: [PATCH 042/151] [misc] Move test-suite environment access into the registry and allowlists (#1898) --- .agents/skills/env-var-conventions/SKILL.md | 5 +- .../skills/reseed-ssim-references/SKILL.md | 2 +- docs/contributing/env_vars.md | 63 ++- docs/inference/mlx_taeh3.md | 2 +- .../configs/fine_tuning/kandinsky5/README.md | 2 +- fastvideo/envs.py | 148 +++++++ .../api/test_attention_selector_resolution.py | 67 +-- .../test_flash_attn_no_pad_resolver.py | 14 +- .../test_flash_attn_nvfp4_opt_out.py | 18 +- .../tests/attention/test_vsa_h3_backward.py | 25 +- .../attention/test_vsa_h3_sm100a_route.py | 85 ++-- fastvideo/tests/conftest.py | 14 + fastvideo/tests/contract/test_env_policy.py | 406 +++--------------- .../tests/contract/test_profiler_regions.py | 17 +- .../tests/contract/test_ssim_ci_runner.py | 9 +- .../contract/test_tensor_golden_isolation.py | 11 +- .../test_ulysses_fault_injection.py | 11 +- fastvideo/tests/encoders/test_clip_encoder.py | 5 +- fastvideo/tests/encoders/test_hyt5_encoder.py | 5 +- .../tests/encoders/test_llama_encoder.py | 5 +- ...test_minimax_h3_qwen3_vl_checkpoint_fp8.py | 7 +- ...st_minimax_h3_qwen3_vl_checkpoint_nvfp4.py | 7 +- .../test_minimax_h3_qwen3_vl_truncation.py | 7 +- .../tests/encoders/test_qwen2_5_encoder.py | 5 +- .../tests/encoders/test_siglip_encoder.py | 6 +- fastvideo/tests/encoders/test_t5_encoder.py | 5 +- .../streaming/test_prompt_providers.py | 21 +- .../tests/entrypoints/test_video_generator.py | 9 +- fastvideo/tests/golden_gate/_harness.py | 52 ++- fastvideo/tests/golden_gate/_tensor_golden.py | 18 +- fastvideo/tests/golden_gate/test_flux.py | 7 +- fastvideo/tests/golden_gate/test_gamecraft.py | 5 +- fastvideo/tests/golden_gate/test_glm_image.py | 5 +- .../tests/golden_gate/test_kandinsky5.py | 7 +- fastvideo/tests/golden_gate/test_longcat.py | 4 +- .../tests/golden_gate/test_minimax_h3_t2v.py | 65 +-- .../tests/golden_gate/test_wan_denoising.py | 5 +- fastvideo/tests/golden_gate/test_zimage.py | 5 +- .../tests/hooks/test_activation_trace.py | 49 +-- .../tests/inference/bsa/test_bsa_inference.py | 6 +- .../lora/test_lora_inference_similarity.py | 9 +- .../test_inference_regional_compile.py | 49 +-- .../inference/vmoba/test_vmoba_inference.py | 6 +- fastvideo/tests/mlx/test_memory_limits.py | 7 +- .../tests/mlx/test_mlx_affine_dq_gemm.py | 39 +- .../test_mlx_minimax_h3_vsa_regressions.py | 7 +- fastvideo/tests/mlx/test_mlx_taeh3.py | 8 +- ...test_mlx_wan22_prompt_cache_fingerprint.py | 4 +- .../mlx/test_mlx_wan22_real_weights_parity.py | 46 +- fastvideo/tests/mlx/tiny_h3.py | 13 +- fastvideo/tests/mlx/tiny_wan.py | 13 +- fastvideo/tests/modal/kernel_build_cache.py | 12 +- fastvideo/tests/modal/pr_test.py | 4 +- .../tests/modal/test_kernel_build_cache.py | 55 +-- fastvideo/tests/modal/test_pr_test.py | 18 +- .../nightly/test_e2e_dmd_t2v_crush_smol.py | 10 +- .../test_e2e_i2v_overfit_single_sample.py | 10 +- .../test_e2e_kandinsky5_dmd_t2v_overfit.py | 33 +- .../test_e2e_ltx2_overfit_new_stack.py | 15 +- .../nightly/test_e2e_overfit_single_sample.py | 10 +- fastvideo/tests/performance/identity.py | 17 +- .../test_compare_baseline_policy.py | 146 ++++--- .../performance/test_dashboard_service.py | 8 +- .../performance/test_inference_performance.py | 14 +- .../test_inference_performance_identity.py | 47 +- ...est_inference_performance_result_schema.py | 19 +- fastvideo/tests/ssim/README.md | 2 +- fastvideo/tests/ssim/bootstrap_references.py | 15 +- fastvideo/tests/ssim/conftest.py | 21 +- .../tests/ssim/inference_similarity_utils.py | 10 +- .../tests/ssim/latent_similarity_utils.py | 2 +- fastvideo/tests/ssim/reference_utils.py | 11 +- fastvideo/tests/ssim/reference_videos_cli.py | 29 +- .../ssim/test_dreamx_world_similarity.py | 25 +- .../tests/ssim/test_flux_t2i_similarity.py | 23 +- .../tests/ssim/test_gamecraft_similarity.py | 79 ++-- fastvideo/tests/ssim/test_gen3c_similarity.py | 33 +- .../tests/ssim/test_glm_image_similarity.py | 25 +- .../tests/ssim/test_kandinsky5_similarity.py | 10 +- .../tests/ssim/test_lingbot_similarity.py | 5 +- .../tests/ssim/test_longcat_similarity.py | 13 +- .../tests/ssim/test_matrixgame2_similarity.py | 5 +- .../tests/ssim/test_matrixgame3_similarity.py | 5 +- fastvideo/tests/ssim/test_sd35_similarity.py | 57 ++- .../tests/ssim/test_zimage_similarity.py | 28 +- fastvideo/tests/stages/_denoising_fixtures.py | 10 +- fastvideo/tests/stages/test_cfg_gating.py | 12 +- ...st_kandinsky5_dmd_stage_backend_engages.py | 17 +- .../test_minimax_h3_encoding_offload.py | 5 +- .../tests/stages/test_wan_causal_denoising.py | 4 +- fastvideo/tests/stages/test_wan_denoising.py | 20 +- .../tests/stages/test_wan_dmd_denoising.py | 4 +- .../tests/stages/test_wan_pipeline_wiring.py | 5 +- .../train/methods/grad_norm_regression.py | 11 +- .../train/methods/test_cosmos_finetune.py | 6 +- .../tests/train/methods/test_ltx2_finetune.py | 6 +- .../methods/test_matrixgame2_finetune.py | 6 +- .../methods/test_streaming_long_tuning.py | 8 +- .../tests/train/methods/test_wan_causal_cd.py | 6 +- .../train/methods/test_wan_causal_dfsft.py | 6 +- .../test_wan_causal_dfsft_framewise.py | 6 +- .../train/methods/test_wan_causal_tfsft.py | 6 +- .../tests/train/methods/test_wan_finetune.py | 8 +- .../test_kandinsky5_qat_attention_engages.py | 16 +- .../tests/train/models/test_load_cosmos.py | 8 +- .../tests/train/models/test_load_hunyuan.py | 6 +- .../train/models/test_load_kandinsky5.py | 6 +- .../tests/train/models/test_load_longcat.py | 8 +- .../tests/train/models/test_load_ltx2.py | 8 +- .../train/models/test_load_matrixgame2.py | 8 +- fastvideo/tests/train/models/test_load_wan.py | 8 +- .../train/models/test_load_wan_causal.py | 6 +- .../training/VSA/test_training_loss_VSA.py | 12 +- .../tests/training/Vanilla/mfu_calculation.py | 13 +- .../training/Vanilla/test_training_loss.py | 11 +- .../training/distill/test_distill_dmd.py | 11 +- .../tests/training/lora/test_lora_training.py | 11 +- .../self-forcing/test_self_forcing.py | 11 +- fastvideo/tests/transformers/test_cosmos.py | 5 +- .../tests/transformers/test_cosmos2_5.py | 5 +- fastvideo/tests/transformers/test_flux.py | 23 +- .../transformers/test_hunyuangamecraft.py | 12 +- fastvideo/tests/transformers/test_hyworld.py | 7 +- .../test_minimax_h3_fusion_routing.py | 31 +- fastvideo/tests/transformers/test_wanvideo.py | 7 +- fastvideo/tests/vaes/test_hunyuan15_vae.py | 5 +- fastvideo/tests/vaes/test_wan_vae.py | 7 +- fastvideo/tests/worker/test_gpu_worker.py | 14 +- tests/local_tests/dreamx_world/PORT_STATUS.md | 2 +- tests/local_tests/dreamx_world/README.md | 2 +- .../test_minimax_h3_qwen3_vl_parity.py | 5 +- tests/local_tests/flux/README.md | 2 +- tests/local_tests/glm_image/test_ar_parity.py | 5 +- .../glm_image/test_edit_pipeline_parity.py | 5 +- .../glm_image/test_pipeline_parity.py | 5 +- tests/local_tests/glm_image/test_t5_parity.py | 5 +- .../glm_image/test_transformer_parity.py | 5 +- .../glm_image/test_transformer_strict_load.py | 5 +- .../local_tests/glm_image/test_vae_parity.py | 5 +- tests/local_tests/minimax_h3/README.md | 8 +- ...st_minimax_h3_text_encoder_nvfp4_parity.py | 7 +- .../minimax_h3/test_mlx_audio_vae_parity.py | 6 +- .../minimax_h3/test_mlx_video_vae_parity.py | 6 +- .../pipelines/test_zimage_pipeline_parity.py | 4 +- .../pipelines/test_zimage_pipeline_smoke.py | 7 +- tests/local_tests/sd35/README.md | 2 +- .../sd35/test_sd35_component_parity.py | 8 +- .../test_minimax_h3_transformer_parity.py | 4 +- .../vaes/test_minimax_h3_audio_vae_parity.py | 4 +- .../vaes/test_minimax_h3_video_vae_parity.py | 4 +- tests/local_tests/zimage/PORT_STATUS.md | 2 +- tests/local_tests/zimage/README.md | 2 +- 152 files changed, 1416 insertions(+), 1325 deletions(-) diff --git a/.agents/skills/env-var-conventions/SKILL.md b/.agents/skills/env-var-conventions/SKILL.md index 3d12b8f4f..c4bb86875 100644 --- a/.agents/skills/env-var-conventions/SKILL.md +++ b/.agents/skills/env-var-conventions/SKILL.md @@ -38,7 +38,10 @@ on a summary here. name to `DEPRECATED_VARIABLES`. Update the uses in `examples/`, `scripts/`, `docs/`, `apps/`, and the tests. 2. **Read the variable with `envs.NAME.get()` inside a function.** - - In tests, change the value with `envs.NAME.override(value)`. + - In tests, change the value with `envs.NAME.override(value)`, and a variable + outside the registry with `envs.override_external(name, value)`; the + `env_overrides` fixture keeps either until the end of the test. + - Name a variable that only tests read `FASTVIDEO_TEST_*`. - Do not call `os.environ`, `os.getenv`, or `monkeypatch.setenv` for a FastVideo variable. - To set a variable that another tool reads, call `envs.set_external`, diff --git a/.agents/skills/reseed-ssim-references/SKILL.md b/.agents/skills/reseed-ssim-references/SKILL.md index d5f48ab4f..3306e6cfd 100644 --- a/.agents/skills/reseed-ssim-references/SKILL.md +++ b/.agents/skills/reseed-ssim-references/SKILL.md @@ -57,7 +57,7 @@ Hardcoded: - Quality tier: **`default`**. `full_quality` is a separate, deliberate operation. - HF repo: `FastVideo/ssim-reference-videos` (override via - `FASTVIDEO_SSIM_REFERENCE_HF_REPO`). + `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO`). - Device folder: `L40S_reference_videos`. ## Prerequisites diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 8f91f1eca..4f47ddcb1 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -13,12 +13,15 @@ test in the same pull request. them directly with `os.environ.get("NAME")`, and the name must be in the external-variable allowlist (`EXTERNAL_ALLOWLIST` in the contract test). When FastVideo sets such a variable for the other tool, it calls `envs.set_external`, `envs.setdefault_external`, or `envs.unset_external`, and the name must be in - `EXTERNAL_WRITE_ALLOWLIST`. + `EXTERNAL_WRITE_ALLOWLIST`. Variables that FastVideo's CI and CI tooling define (for example `TEST_SCOPE` and + `PERF_RUN_SOURCE`) keep their names, and test code under `fastvideo/tests/` reads them directly; they are listed + in `CI_ONLY_VARIABLES` in the contract test, together with the file that sets each one. 2. **Read with `envs.NAME.get()`, write with `envs.NAME.set()`, and change a value in tests with `envs.NAME.override()`.** Each type has one parsing rule. A value that the rule rejects raises `fastvideo.envs.EnvVarError` instead of falling back to the default. 3. **Name FastVideo variables with the `FASTVIDEO_` prefix.** The second word states the purpose where one applies: - `ENABLE_`, `DISABLE_`, `USE_`, `FORCE_`, `DEBUG_`, `TEST_`. + `ENABLE_`, `DISABLE_`, `USE_`, `FORCE_`, `DEBUG_`, `TEST_`. Variables that only tests read use + `FASTVIDEO_TEST_`, for example `FASTVIDEO_TEST_SD35_MODEL_DIR`, and the category `test`. 4. **Keep a renamed variable as a deprecated alias until the next minor release.** Setting the old name logs a warning. Delete a variable that no code reads, and list it in `DEPRECATED_VARIABLES` so that setting it logs a warning. @@ -29,7 +32,8 @@ test in the same pull request. effect without re-importing a module. Module level, class bodies, decorators, and default argument values run at import time. 7. **Do not write the environment to pass values between parts of FastVideo.** Pass an argument instead. Tests use - `envs.NAME.override()`. + `envs.NAME.override()`, and `envs.override_external()` for variables outside the registry; both restore the + previous value. ## Field types @@ -87,6 +91,21 @@ with envs.FASTVIDEO_DEBUG_MY_STAGE.override(True): run_stage() ``` +For a variable outside the registry, such as `MASTER_PORT` or `TEST_SCOPE`, use `envs.override_external`. To keep +an override until the end of a test, enter it through the `env_overrides` fixture from `fastvideo/tests/conftest.py`, +which restores every value at teardown: + +```python +def test_my_stage(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_DEBUG_MY_STAGE.override(True)) + env_overrides.enter_context(envs.override_external("MASTER_PORT", "29512")) + run_stage() +``` + +In test code under `fastvideo/tests/`, `override_external` accepts any name that code may read directly +(`EXTERNAL_ALLOWLIST`, `CI_ONLY_VARIABLES`) or that is in `EXTERNAL_WRITE_ALLOWLIST`. Library code may write only +the names in `EXTERNAL_WRITE_ALLOWLIST`. + ## Rename or remove a variable To rename a variable, declare it under the new name and list the old name in `deprecated_names`: @@ -121,9 +140,11 @@ It reports each violation as `: `: | | with a name built at runtime (``) | another tool owns, add it to | | | | `EXTERNAL_ALLOWLIST` with a reason. | | `write` | `os.environ[...] = ...`, `setdefault`, `pop`, `del`, | Pass an argument instead. In tests, use | -| | `os.putenv`, `os.unsetenv`, `monkeypatch.setenv`/`delenv`, | `envs.NAME.override()`. For a variable that | -| | or an `envs.*_external` call with a name outside | another tool reads, call an | -| | `EXTERNAL_WRITE_ALLOWLIST` | `envs.*_external` helper and add the name to | +| | `os.putenv`, `os.unsetenv`, `monkeypatch.setenv`/`delenv`, | `envs.NAME.override()`, or | +| | or an `envs.*_external` call with a name that the | `envs.override_external()` for a variable | +| | helper does not accept | outside the registry. For a variable that | +| | | another tool reads, call an | +| | | `envs.*_external` helper and add the name to | | | | `EXTERNAL_WRITE_ALLOWLIST` with a reason. | | `whole-environ` | `os.environ.copy()`, `dict(os.environ)`, iteration, | Read the specific variables that the code | | | `mock.patch.dict(os.environ, ...)`, `os.environ.update` | needs. | @@ -227,6 +248,36 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `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_SSIM_REFERENCE_HF_REPO` | str | `FastVideo/ssim-reference-videos` | test | Hugging Face repository that holds the SSIM reference videos. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO`. | +| `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE` | str | `dataset` | test | Repository type of FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO. Deprecated names: `FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE`. | +| `FASTVIDEO_TEST_SSIM_SKIP_REFERENCE_DOWNLOAD` | bool | `0` | test | SSIM tests use local reference videos without downloading. Deprecated names: `FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD`. | +| `FASTVIDEO_TEST_SSIM_FULL_QUALITY` | bool | `0` | test | SSIM tests use the full-quality sampling configurations. Deprecated names: `FASTVIDEO_SSIM_FULL_QUALITY`. | +| `FASTVIDEO_TEST_NIGHTLY` | bool | `0` | test | Run the nightly end-to-end overfit tests. Deprecated names: `FASTVIDEO_NIGHTLY`. | +| `FASTVIDEO_TEST_ULYSSES_FAULT_RANK` | str | unset | test | Rank that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_RANK`. | +| `FASTVIDEO_TEST_ULYSSES_FAULT_STAGE` | str | unset | test | Stage that fails in the Ulysses fault-injection test. The test sets it for its worker processes. Deprecated names: `FASTVIDEO_ULYSSES_FAULT_STAGE`. | +| `FASTVIDEO_TEST_GOLDEN_GATE_DIR` | str | unset | test | Local directory of golden-gate reference tensors. Deprecated names: `FASTVIDEO_GOLDEN_GATE_DIR`. | +| `FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY` | bool | `0` | test | Run the MLX Wan2.2 5B real-weights parity test on hosts with little memory. Deprecated names: `FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY`. | +| `FASTVIDEO_TEST_WAN22_5B_ROOT` | str | unset | test | Local Wan2.2 5B checkpoint for the MLX real-weights parity test. Deprecated names: `FASTVIDEO_WAN22_5B_ROOT`. | +| `FASTVIDEO_TEST_GRADNORM_UPDATE` | bool | `0` | test | Gradient-norm regression tests update their references. Deprecated names: `FASTVIDEO_GRADNORM_UPDATE`. | +| `FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Cam-Diffusers` | test | Model for the DreamX-World camera SSIM test. Deprecated names: `DREAMX_WORLD_SSIM_MODEL_PATH`. | +| `FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH` | str | `FastVideo/DreamX-World-5B-Diffusers` | test | Model for the DreamX-World autoregressive SSIM test. Deprecated names: `DREAMX_WORLD_AR_SSIM_MODEL_PATH`. | +| `FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR` | str | `black-forest-labs/FLUX.1-dev` | test | Model for the Flux text-to-image SSIM test. Deprecated names: `FLUX_T2I_MODEL_DIR`. | +| `FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH` | str | unset | test | Local Flux transformer for the Flux transformer test. Deprecated names: `FLUX_TRANSFORMER_PATH`. | +| `FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH` | str | `FastVideo/HunyuanGameCraft-Diffusers` | test | Model for the HunyuanGameCraft SSIM test. Deprecated names: `GAMECRAFT_MODEL_PATH`. | +| `FASTVIDEO_TEST_GEN3C_MODEL_PATH` | str | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | test | Model for the GEN3C SSIM test. Deprecated names: `GEN3C_MODEL_PATH`. | +| `FASTVIDEO_TEST_GEN3C_IMAGE_PATH` | str | unset | test | Input image for the GEN3C SSIM test. Deprecated names: `GEN3C_TEST_IMAGE_PATH`. | +| `FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR` | str | unset | test | Local official GLM-Image weights for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_LOCAL_WEIGHTS_DIR`. | +| `FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR` | str | unset | test | Model for the GLM-Image SSIM test. Deprecated names: `GLM_IMAGE_MODEL_DIR`. | +| `FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS` | int | `1` | test | GPUs for the Kandinsky5 nightly end-to-end overfit test. Deprecated names: `KANDINSKY5_E2E_NUM_GPUS`. | +| `FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE` | bool | `0` | test | The Kandinsky5 nightly end-to-end test writes a missing reference video. Deprecated names: `KANDINSKY5_E2E_WRITE_REFERENCE`. | +| `FASTVIDEO_TEST_LONGCAT_MODEL_ROOT` | str | unset | test | Local LongCat-Video checkpoint for the golden-gate test. Deprecated names: `LONGCAT_MODEL_ROOT`. | +| `FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR` | str | unset | test | Local directory of MiniMax-H3 golden-gate tensors. Deprecated names: `MINIMAX_H3_GATE_GOLDEN_DIR`. | +| `FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER` | int | `0` | test | Transformer layer that the MiniMax-H3 golden-gate test checks. Deprecated names: `MINIMAX_H3_GATE_LAYER`. | +| `FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT` | str | unset | test | Local MiniMax-H3 checkpoint for the golden-gate test. Deprecated names: `MINIMAX_H3_MODEL_ROOT`. | +| `FASTVIDEO_TEST_SD35_MODEL_DIR` | str | `stabilityai/stable-diffusion-3.5-medium` | test | Model for the Stable Diffusion 3.5 SSIM test. Deprecated names: `SD35_MODEL_DIR`. | +| `FASTVIDEO_TEST_TAEH3_REFERENCE_DIR` | str | unset | test | Upstream taehv checkout for the MLX TAEH3 parity test. Deprecated names: `TAEH3_REFERENCE_DIR`. | +| `FASTVIDEO_TEST_ZIMAGE_MODEL_DIR` | str | `Tongyi-MAI/Z-Image-Turbo` | test | Model for the Z-Image SSIM test. Deprecated names: `ZIMAGE_MODEL_DIR`. | +| `FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION` | str | `f332072aa78be7aecdf3ee76d5c247082da564a6` | test | Hugging Face revision of the Z-Image model for its SSIM test. Deprecated names: `ZIMAGE_MODEL_REVISION`. | Variables that FastVideo no longer reads; setting one logs a warning: diff --git a/docs/inference/mlx_taeh3.md b/docs/inference/mlx_taeh3.md index ac992f24a..80355e0f5 100644 --- a/docs/inference/mlx_taeh3.md +++ b/docs/inference/mlx_taeh3.md @@ -70,7 +70,7 @@ Run the numerical tests against a local TAEHV checkout containing the released weights: ```bash -TAEH3_REFERENCE_DIR=/path/to/taehv \ +FASTVIDEO_TEST_TAEH3_REFERENCE_DIR=/path/to/taehv \ python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py -q ``` diff --git a/examples/train/configs/fine_tuning/kandinsky5/README.md b/examples/train/configs/fine_tuning/kandinsky5/README.md index 5f8a80043..1e111e8a4 100644 --- a/examples/train/configs/fine_tuning/kandinsky5/README.md +++ b/examples/train/configs/fine_tuning/kandinsky5/README.md @@ -187,7 +187,7 @@ to a supported dense backend while the FP4 linear layers still run. `Kandinsky5DMDConfig` override, generating deterministically (fixed seed), and comparing MS-SSIM against a committed reference video. The test fails (does not skip) if the reference is missing; record it once on - a sanctioned GPU box with `KANDINSKY5_E2E_WRITE_REFERENCE=1`, review the + a sanctioned GPU box with `FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE=1`, review the written video, and commit it (see the test's module docstring). Nightly tests are not collected per-PR (`fastvideo/tests/contract/ test_ci_test_collection.py` allowlists the directory), so run it diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 881b2f158..8962a8e70 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -222,6 +222,27 @@ def unset_external(name: str) -> None: os.environ.pop(name, None) +@contextmanager +def override_external(name: str, value: str | None) -> Iterator[None]: + """Set, or unset when ``value`` is None, a variable outside the registry, and restore it on exit. + + Tests use it instead of ``monkeypatch.setenv`` for variables that other + tools or the CI read. + """ + previous = os.environ.get(name) + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + try: + yield + finally: + if previous is None: + os.environ.pop(name, None) + else: + os.environ[name] = previous + + # ================== Paths ================== FASTVIDEO_CONFIG_ROOT = EnvPath( @@ -530,6 +551,133 @@ FASTVIDEO_TEST_KANDINSKY5_OVERFIT_OUTPUT_DIR = EnvStr("data/kandinsky5_overfit_p doc="Output directory for preprocess_kandinsky5_overfit.py.", deprecated_names=("KANDINSKY5_OVERFIT_OUTPUT_DIR", )) +# Switches and paths that only tests read. Each old name stays readable, with a +# warning, until the next minor release. +FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO = EnvStr("FastVideo/ssim-reference-videos", + category="test", + doc="Hugging Face repository that holds the SSIM reference videos.", + deprecated_names=("FASTVIDEO_SSIM_REFERENCE_HF_REPO", )) +FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE = EnvStr("dataset", + category="test", + doc="Repository type of FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.", + deprecated_names=("FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE", )) +FASTVIDEO_TEST_SSIM_SKIP_REFERENCE_DOWNLOAD = EnvBool(False, + category="test", + doc="SSIM tests use local reference videos without downloading.", + deprecated_names=("FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", )) +FASTVIDEO_TEST_SSIM_FULL_QUALITY = EnvBool(False, + category="test", + doc="SSIM tests use the full-quality sampling configurations.", + deprecated_names=("FASTVIDEO_SSIM_FULL_QUALITY", )) +FASTVIDEO_TEST_NIGHTLY = EnvBool(False, + category="test", + doc="Run the nightly end-to-end overfit tests.", + deprecated_names=("FASTVIDEO_NIGHTLY", )) +FASTVIDEO_TEST_ULYSSES_FAULT_RANK = EnvStr( + None, + category="test", + doc="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 = EnvStr( + None, + category="test", + doc="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 = EnvStr(None, + category="test", + doc="Local directory of golden-gate reference tensors.", + deprecated_names=("FASTVIDEO_GOLDEN_GATE_DIR", )) +FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY = EnvBool( + False, + category="test", + doc="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 = EnvStr(None, + category="test", + doc="Local Wan2.2 5B checkpoint for the MLX real-weights parity test.", + deprecated_names=("FASTVIDEO_WAN22_5B_ROOT", )) +FASTVIDEO_TEST_GRADNORM_UPDATE = EnvBool(False, + category="test", + doc="Gradient-norm regression tests update their references.", + deprecated_names=("FASTVIDEO_GRADNORM_UPDATE", )) +FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH = EnvStr("FastVideo/DreamX-World-5B-Cam-Diffusers", + category="test", + doc="Model for the DreamX-World camera SSIM test.", + deprecated_names=("DREAMX_WORLD_SSIM_MODEL_PATH", )) +FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH = EnvStr("FastVideo/DreamX-World-5B-Diffusers", + category="test", + doc="Model for the DreamX-World autoregressive SSIM test.", + deprecated_names=("DREAMX_WORLD_AR_SSIM_MODEL_PATH", )) +FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR = EnvStr("black-forest-labs/FLUX.1-dev", + category="test", + doc="Model for the Flux text-to-image SSIM test.", + deprecated_names=("FLUX_T2I_MODEL_DIR", )) +FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH = EnvStr(None, + category="test", + doc="Local Flux transformer for the Flux transformer test.", + deprecated_names=("FLUX_TRANSFORMER_PATH", )) +FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH = EnvStr("FastVideo/HunyuanGameCraft-Diffusers", + category="test", + doc="Model for the HunyuanGameCraft SSIM test.", + deprecated_names=("GAMECRAFT_MODEL_PATH", )) +FASTVIDEO_TEST_GEN3C_MODEL_PATH = EnvStr("FastVideo/GEN3C-Cosmos-7B-Diffusers", + category="test", + doc="Model for the GEN3C SSIM test.", + deprecated_names=("GEN3C_MODEL_PATH", )) +FASTVIDEO_TEST_GEN3C_IMAGE_PATH = EnvStr(None, + category="test", + doc="Input image for the GEN3C SSIM test.", + deprecated_names=("GEN3C_TEST_IMAGE_PATH", )) +FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR = EnvStr(None, + category="test", + doc="Local official GLM-Image weights for the GLM-Image SSIM test.", + deprecated_names=("GLM_IMAGE_LOCAL_WEIGHTS_DIR", )) +FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR = EnvStr(None, + category="test", + doc="Model for the GLM-Image SSIM test.", + deprecated_names=("GLM_IMAGE_MODEL_DIR", )) +FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS = EnvInt(1, + category="test", + doc="GPUs for the Kandinsky5 nightly end-to-end overfit test.", + deprecated_names=("KANDINSKY5_E2E_NUM_GPUS", )) +FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE = EnvBool( + False, + category="test", + doc="The Kandinsky5 nightly end-to-end test writes a missing reference video.", + deprecated_names=("KANDINSKY5_E2E_WRITE_REFERENCE", )) +FASTVIDEO_TEST_LONGCAT_MODEL_ROOT = EnvStr(None, + category="test", + doc="Local LongCat-Video checkpoint for the golden-gate test.", + deprecated_names=("LONGCAT_MODEL_ROOT", )) +FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR = EnvStr(None, + category="test", + doc="Local directory of MiniMax-H3 golden-gate tensors.", + deprecated_names=("MINIMAX_H3_GATE_GOLDEN_DIR", )) +FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER = EnvInt(0, + category="test", + doc="Transformer layer that the MiniMax-H3 golden-gate test checks.", + deprecated_names=("MINIMAX_H3_GATE_LAYER", )) +FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT = EnvStr(None, + category="test", + doc="Local MiniMax-H3 checkpoint for the golden-gate test.", + deprecated_names=("MINIMAX_H3_MODEL_ROOT", )) +FASTVIDEO_TEST_SD35_MODEL_DIR = EnvStr("stabilityai/stable-diffusion-3.5-medium", + category="test", + doc="Model for the Stable Diffusion 3.5 SSIM test.", + deprecated_names=("SD35_MODEL_DIR", )) +FASTVIDEO_TEST_TAEH3_REFERENCE_DIR = EnvStr(None, + category="test", + doc="Upstream taehv checkout for the MLX TAEH3 parity test.", + deprecated_names=("TAEH3_REFERENCE_DIR", )) +FASTVIDEO_TEST_ZIMAGE_MODEL_DIR = EnvStr("Tongyi-MAI/Z-Image-Turbo", + category="test", + doc="Model for the Z-Image SSIM test.", + deprecated_names=("ZIMAGE_MODEL_DIR", )) +FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION = EnvStr("f332072aa78be7aecdf3ee76d5c247082da564a6", + category="test", + doc="Hugging Face revision of the Z-Image model for its SSIM test.", + deprecated_names=("ZIMAGE_MODEL_REVISION", )) + # Variables that FastVideo no longer reads. Setting one logs a warning; delete # the entries in the next minor release. DEPRECATED_VARIABLES = { diff --git a/fastvideo/tests/api/test_attention_selector_resolution.py b/fastvideo/tests/api/test_attention_selector_resolution.py index c108ea43b..51086ec7f 100644 --- a/fastvideo/tests/api/test_attention_selector_resolution.py +++ b/fastvideo/tests/api/test_attention_selector_resolution.py @@ -27,6 +27,7 @@ import pytest import torch import fastvideo.attention.selector as selector +import fastvideo.envs as envs import fastvideo.platforms as platforms from fastvideo.platforms import AttentionBackendEnum @@ -68,10 +69,10 @@ class _FakePlatform: @pytest.fixture(autouse=True) -def _fake_platform(monkeypatch): +def _fake_platform(monkeypatch, env_overrides): monkeypatch.setattr(platforms, "_current_platform", _FakePlatform()) monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name) - monkeypatch.delenv("FASTVIDEO_ATTENTION_BACKEND", raising=False) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(None)) selector._cached_get_attn_backend.cache_clear() yield # Fake-platform resolutions must not outlive the test: mutations no @@ -94,9 +95,9 @@ def _oracle(requested, env, default): @pytest.mark.parametrize("scoped", [None, "unset", SAGE]) @pytest.mark.parametrize("env", [None, "TORCH_SDPA"]) @pytest.mark.parametrize("default", [None, SAGE]) -def test_precedence_matrix_matches_previous_semantics(monkeypatch, scoped, env, default) -> None: +def test_precedence_matrix_matches_previous_semantics(env_overrides, scoped, env, default) -> None: if env is not None: - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", env) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(env)) if scoped == "unset": # No scope and no explicit request: env is consulted. @@ -113,21 +114,21 @@ def test_precedence_matrix_matches_previous_semantics(monkeypatch, scoped, env, @pytest.mark.parametrize("requested", [None, SAGE]) @pytest.mark.parametrize("env", [None, "TORCH_SDPA"]) @pytest.mark.parametrize("default", [None, SAGE]) -def test_explicit_request_matches_the_same_oracle(monkeypatch, requested, env, default) -> None: +def test_explicit_request_matches_the_same_oracle(env_overrides, requested, env, default) -> None: """An explicit `requested=` resolves exactly as the same request would via a scope, and never consults the environment behind it.""" if env is not None: - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", env) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(env)) got = selector.get_attn_backend(default_backend=default, requested=requested, **KWARGS) assert got == _oracle(requested, None, default) -def test_scope_none_ignores_env_request(monkeypatch) -> None: +def test_scope_none_ignores_env_request(env_overrides) -> None: """The dense teacher/critic case: auto-select, ignore the process-wide request — previously implemented by popping the env var and flushing the selector cache around the build.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" with selector._component_attention_backend_scope(None, component="transformer"): @@ -135,10 +136,10 @@ def test_scope_none_ignores_env_request(monkeypatch) -> None: assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" -def test_interleaved_component_scopes_need_no_cache_clear(monkeypatch) -> None: +def test_interleaved_component_scopes_need_no_cache_clear(env_overrides) -> None: """Per-role/per-component builds interleave freely; identical shape/dtype keys resolve per scope with zero cache management.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) for _ in range(2): with selector._component_attention_backend_scope(SDPA, component="student"): @@ -148,8 +149,8 @@ def test_interleaved_component_scopes_need_no_cache_clear(monkeypatch) -> None: assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" -def test_scope_is_exception_safe(monkeypatch) -> None: - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") +def test_scope_is_exception_safe(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) with pytest.raises(RuntimeError, match="boom"): with selector._component_attention_backend_scope(SDPA, component="dit"): raise RuntimeError("boom") @@ -157,11 +158,11 @@ def test_scope_is_exception_safe(monkeypatch) -> None: assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" -def test_explicit_none_means_auto_not_absence(monkeypatch) -> None: +def test_explicit_none_means_auto_not_absence(env_overrides) -> None: """`requested=None` is an answer ("this component resolved to automatic selection"), so the env var must NOT be consulted behind it. Omitting the argument is the absence, and still falls back to the env.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) # Explicit "automatic selection" -> the fake platform's auto answer. assert selector.get_attn_backend(requested=None, **KWARGS) == "FLASH_ATTN" @@ -170,22 +171,22 @@ def test_explicit_none_means_auto_not_absence(monkeypatch) -> None: assert selector.get_attn_backend(requested=selector.NO_REQUEST, **KWARGS) == "SAGE_ATTN" -def test_explicit_request_outranks_the_construction_scope(monkeypatch) -> None: +def test_explicit_request_outranks_the_construction_scope(env_overrides) -> None: """A caller that knows its component's decision is not overridden by an ambient scope that happens to be active.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) with selector._component_attention_backend_scope(SDPA, component="dit"): assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA" assert selector.get_attn_backend(requested=SAGE, **KWARGS) == "SAGE_ATTN" -def test_env_change_takes_effect_without_cache_clear(monkeypatch) -> None: +def test_env_change_takes_effect_without_cache_clear(env_overrides) -> None: """Selection inputs live in the cache key, so an env change is simply a different key. (Previously a changed env var was silently ignored until someone remembered to flush the cache.)""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA" @@ -195,8 +196,8 @@ def test_scope_typo_fails_fast() -> None: pass -def test_consult_env_scope_keeps_env_visible(monkeypatch) -> None: - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") +def test_consult_env_scope_keeps_env_visible(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) with selector._component_attention_backend_scope(None, component="dit", consult_env=True): assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" @@ -231,22 +232,22 @@ def test_active_device_is_part_of_the_cache_key(monkeypatch) -> None: # --------------------------------------------------------------------------- -def test_env_is_folded_into_the_typed_request_once(monkeypatch): +def test_env_is_folded_into_the_typed_request_once(env_overrides): """``FastVideoArgs.attention_backend`` is the parse-once adapter.""" from fastvideo.fastvideo_args import FastVideoArgs - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) assert FastVideoArgs(model_path="x").attention_backend == "SAGE_ATTN" # An explicit request always wins over the environment. assert FastVideoArgs(model_path="x", attention_backend="TORCH_SDPA").attention_backend == "TORCH_SDPA" -def test_unparseable_env_falls_through_instead_of_raising(monkeypatch): +def test_unparseable_env_falls_through_instead_of_raising(env_overrides): """The env var keeps its permissive parse; only explicit requests raise.""" from fastvideo.fastvideo_args import FastVideoArgs - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "flash_atn") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("flash_atn")) assert FastVideoArgs(model_path="x").attention_backend is None with pytest.raises(ValueError, match="Unknown attention backend"): @@ -312,7 +313,7 @@ def test_component_attention_backend_reads_the_recorded_decision(): assert selector.component_attention_backend(_Component(config)) == SAGE -def test_component_with_no_concrete_decision_reports_no_request(monkeypatch): +def test_component_with_no_concrete_decision_reports_no_request(env_overrides): """"Stamped with no request" and "never stamped" are the SAME stored value. ``ModelConfig._resolved_attention_backend`` is a field defaulting to None, @@ -338,7 +339,7 @@ def test_component_with_no_concrete_decision_reports_no_request(monkeypatch): assert selector.component_attention_backend(_Component(untouched)) is selector.NO_REQUEST # (c) and it does NOT suppress the env var, which is the whole point - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) selector._cached_get_attn_backend.cache_clear() passthrough = selector.get_attn_backend(requested=selector.component_attention_backend(_Component(untouched)), **KWARGS) @@ -376,28 +377,28 @@ def test_component_without_a_recorded_decision_reports_no_request(): assert selector.component_attention_backend(_HFStyle()) is selector.NO_REQUEST -def test_effective_backend_prefers_the_recorded_decision(monkeypatch) -> None: +def test_effective_backend_prefers_the_recorded_decision(env_overrides) -> None: """A component built by a loader follows the decision recorded on its config, even when the env var asks for something else.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) config = SimpleNamespace(_resolved_attention_backend=SDPA) assert selector.effective_attention_backend(config) == SDPA -def test_effective_backend_follows_env_for_direct_construction(monkeypatch) -> None: +def test_effective_backend_follows_env_for_direct_construction(env_overrides) -> None: """A component constructed without a loader has no recorded decision; like the attention layers inside it, it follows the env var.""" config = SimpleNamespace(_resolved_attention_backend=None) assert selector.effective_attention_backend(config) is None - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) assert selector.effective_attention_backend(config) == SAGE assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN" -def test_effective_backend_follows_the_active_scope(monkeypatch) -> None: +def test_effective_backend_follows_the_active_scope(env_overrides) -> None: """While a loader scope is active, the scope's rule applies: a scope that ignores the env var (the dense teacher/critic case) keeps ignoring it.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) config = SimpleNamespace(_resolved_attention_backend=None) with selector._component_attention_backend_scope(None, component="teacher"): assert selector.effective_attention_backend(config) is None diff --git a/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py b/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py index a49883fb7..44549d08a 100644 --- a/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py +++ b/fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py @@ -25,6 +25,8 @@ import sys import pytest +import fastvideo.envs as envs + if importlib.util.find_spec("flash_attn") is None: pytest.skip( "flash_attn not installed; resolver tests require it for the unconditional top-level imports in flash_attn_no_pad", @@ -37,10 +39,10 @@ def _reload_resolver_module(): return importlib.import_module("fastvideo.attention.utils.flash_attn_no_pad") -def test_resolver_skips_cute_without_opt_in(monkeypatch) -> None: +def test_resolver_skips_cute_without_opt_in(monkeypatch, env_overrides) -> None: """Without ``FASTVIDEO_FA4=1`` the resolver must not even attempt the cute import.""" - monkeypatch.delenv("FASTVIDEO_FA4", raising=False) + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(None)) attempted: list[str] = [] real_import = builtins.__import__ @@ -58,7 +60,7 @@ def test_resolver_skips_cute_without_opt_in(monkeypatch) -> None: assert "fastvideo.attention.utils.flash_attn_cute" not in attempted -def test_resolver_raises_when_opted_in_but_cute_unavailable(monkeypatch) -> None: +def test_resolver_raises_when_opted_in_but_cute_unavailable(monkeypatch, env_overrides) -> None: """With ``FASTVIDEO_FA4=1`` an unimportable cute build fails loudly instead of silently falling through to FA3/FA2. @@ -67,7 +69,7 @@ def test_resolver_raises_when_opted_in_but_cute_unavailable(monkeypatch) -> None ImportError as "flash-attn not installed" (``bsa_attn.py``) cannot swallow the opted-in failure. """ - monkeypatch.setenv("FASTVIDEO_FA4", "1") + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(True)) real_import = builtins.__import__ def patched_import(name, globals=None, locals=None, fromlist=(), level=0): @@ -81,9 +83,9 @@ def test_resolver_raises_when_opted_in_but_cute_unavailable(monkeypatch) -> None _reload_resolver_module() -def test_resolver_returns_flash_attn_when_interface_unavailable(monkeypatch) -> None: +def test_resolver_returns_flash_attn_when_interface_unavailable(monkeypatch, env_overrides) -> None: """The terminal fallback is the plain ``flash_attn`` import.""" - monkeypatch.delenv("FASTVIDEO_FA4", raising=False) + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(None)) real_import = builtins.__import__ def patched_import(name, globals=None, locals=None, fromlist=(), level=0): diff --git a/fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py b/fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py index 5c6c07f25..5d28b4228 100644 --- a/fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py +++ b/fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py @@ -11,6 +11,8 @@ from unittest.mock import patch import pytest +import fastvideo.envs as envs + # The backend module needs a flash-attention package (FA2/FA3, or FA4 with # FASTVIDEO_FA4=1) at import. exc_type=ImportError also covers partial # installs (e.g. an FA4-only environment without the opt-in env set raises @@ -44,23 +46,23 @@ def fa4_runtime_available(): yield -def test_nvfp4_disabled_by_default(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False) +def test_nvfp4_disabled_by_default(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(None)) assert _build_impl().nvfp4_fa4 is False -def test_nvfp4_env_opt_in_enables_when_arg_absent(monkeypatch, fa4_runtime_available) -> None: - monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1") +def test_nvfp4_env_opt_in_enables_when_arg_absent(env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(True)) assert _build_impl().nvfp4_fa4 is True -def test_explicit_disable_wins_over_env_opt_in(monkeypatch) -> None: +def test_explicit_disable_wins_over_env_opt_in(env_overrides) -> None: """The H3 VAE constructs its impl with ``nvfp4_fa4=False``; the DiT env opt-in must not FP4-quantize Q/K inside the FP32-pinned VAE.""" - monkeypatch.setenv("FASTVIDEO_NVFP4_FA4", "1") + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(True)) assert _build_impl(nvfp4_fa4=False).nvfp4_fa4 is False -def test_explicit_enable_works_without_env(monkeypatch, fa4_runtime_available) -> None: - monkeypatch.delenv("FASTVIDEO_NVFP4_FA4", raising=False) +def test_explicit_enable_works_without_env(env_overrides, fa4_runtime_available) -> None: + env_overrides.enter_context(envs.FASTVIDEO_NVFP4_FA4.override(None)) assert _build_impl(nvfp4_fa4=True).nvfp4_fa4 is True diff --git a/fastvideo/tests/attention/test_vsa_h3_backward.py b/fastvideo/tests/attention/test_vsa_h3_backward.py index 7c9e45417..8d2d4395f 100644 --- a/fastvideo/tests/attention/test_vsa_h3_backward.py +++ b/fastvideo/tests/attention/test_vsa_h3_backward.py @@ -7,9 +7,12 @@ therefore poisons the graph, and the failure only appears once the VSA-256 CuTe path has a backward at all. These tests pin the composition. """ +import contextlib + import pytest import torch +import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder) _SPEC = dict(raw_latent_shape=(16, 16, 24), patch_size=(1, 2, 2), prefix_segments=(64, 32, 16)) @@ -28,18 +31,18 @@ def _build_meta(device, sparsity=0.5): ) -def _select_backend(monkeypatch, backend): +def _select_backend(env_overrides, backend): if backend == "cute": pytest.importorskip( "flash_attn.cute.block_sparsity", reason="optional FA4 CuTe build (flash_attn.cute) not installed", ) - monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1") - monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False) - monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_CUTEDSL", "1")) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", None)) + env_overrides.enter_context(envs.override_external("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", None)) else: - monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") - monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1")) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_CUTEDSL", None)) def _forward_backward(impl, meta, gate_compress, device): @@ -61,14 +64,14 @@ def _forward_backward(impl, meta, gate_compress, device): @pytest.mark.parametrize("backend", ["triton", "cute"]) @pytest.mark.parametrize("gate_compress", [False, True]) -def test_h3_vsa_backward_runs(monkeypatch, backend: str, gate_compress: bool) -> None: +def test_h3_vsa_backward_runs(env_overrides, backend: str, gate_compress: bool) -> None: """Regression: with the CuTe backend and a non-zero gate this used to die with "one of the variables needed for gradient computation has been modified by an inplace operation ... output 0 of FlashAttnFuncBackward". """ if not torch.cuda.is_available(): pytest.skip("CUDA is required") - _select_backend(monkeypatch, backend) + _select_backend(env_overrides, backend) device = torch.device("cuda") meta = _build_meta(device) @@ -84,7 +87,7 @@ def test_h3_vsa_backward_runs(monkeypatch, backend: str, gate_compress: bool) -> @pytest.mark.parametrize("gate_compress", [False, True]) -def test_h3_vsa_backward_cute_matches_triton(monkeypatch, gate_compress: bool) -> None: +def test_h3_vsa_backward_cute_matches_triton(gate_compress: bool) -> None: """CuTe and Triton take different routes to the same math; their gradients should agree to bf16 tolerance.""" if not torch.cuda.is_available(): @@ -95,8 +98,8 @@ def test_h3_vsa_backward_cute_matches_triton(monkeypatch, gate_compress: bool) - grads = {} for backend in ("triton", "cute"): - with monkeypatch.context() as m: - _select_backend(m, backend) + with contextlib.ExitStack() as backend_env: + _select_backend(backend_env, backend) meta = _build_meta(device) _, leaves = _forward_backward(impl, meta, gate_compress, device) grads[backend] = [leaf.grad.detach().float() for leaf in leaves] diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 5b03fabe9..9b26678c8 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -14,6 +14,7 @@ import pytest import torch import fastvideo.attention.backends.video_sparse_attn_h3 as vsa_h3 +import fastvideo.envs as envs import fastvideo.logger as fastvideo_logger from fastvideo.attention.backends.video_sparse_attn_h3 import (VSA_SM100A_ENV, MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder, _sm100a_unavailable_reason) @@ -126,10 +127,10 @@ def test_reason_covers_every_precondition(): assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None -def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch): +def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch, env_overrides): fake_sm = _FakeSm100a(supported=True) monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) @@ -146,10 +147,10 @@ def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch): assert impl._compile_layer_idx.item() == -1 -def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch): +def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch, env_overrides): fake_sm = _FakeSm100a(supported=True) monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - monkeypatch.delenv(VSA_SM100A_ENV, raising=False) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) unsupported = impl.prepare_for_regional_compile(torch.device("cpu")) @@ -161,7 +162,7 @@ def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch): assert fake_sm.support_calls == [] -def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch): +def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch, env_overrides): class _IndexOnlySm100a: @@ -169,7 +170,7 @@ def test_prepare_for_regional_compile_requires_mask_entry(monkeypatch): raise AssertionError("missing mask entry must be rejected before the support probe") monkeypatch.setattr(vsa_h3, "_sm100a", _IndexOnlySm100a()) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) @@ -210,11 +211,11 @@ def test_mask_dispatch_uses_local_compatibility_route_for_older_kernel_wheel(mon assert lse is None and calls == [(mask, vbs)] -def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch): +def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrides): """Capture must not revisit env/capability checks or raw map_to_index.""" fake_sm = _FakeSm100a(supported=True) monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) impl.prepare_for_regional_compile(torch.device("cpu")) @@ -258,11 +259,11 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch): torch._dynamo.reset() -def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch): +def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch, env_overrides): """Fifty H3 blocks must not specialize the shared graph on layer_idx.""" fake_sm = _FakeSm100a(supported=True) monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64)) assert meta.dense_layers_tensor.tolist() == [0, 17] @@ -308,11 +309,11 @@ def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_ov torch._dynamo.reset() -def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(monkeypatch): +def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(monkeypatch, env_overrides): """Pipeline compile must share one graph without selecting sm_100a.""" fake_sm = _FakeSm100a(supported=True) monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None) meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64)) q, k, v = _tiled_qkv(meta) @@ -363,17 +364,17 @@ def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(m assert len(compiled_graphs) == 1 -def test_default_off_routes_triton(routed, monkeypatch): +def test_default_off_routes_triton(routed, env_overrides): fake_sm, fake_triton, run, _ = routed - monkeypatch.delenv(VSA_SM100A_ENV, raising=False) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) run() assert fake_triton.calls == 1 assert fake_sm.calls == [] -def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeypatch): +def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) run(requires_grad=True) @@ -382,9 +383,9 @@ def test_unprepared_training_compile_keeps_existing_triton_route(routed, monkeyp assert fake_sm.calls == [] -def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch): +def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, meta = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) messages = [] monkeypatch.setattr(vsa_h3.logger, "info_once", messages.append) out = run() @@ -402,9 +403,9 @@ def test_env_on_routes_sm100a_with_index_metadata(routed, monkeypatch): assert out.shape == (1, n_tiles * 64, _HEADS, _DIM) -def test_sm100a_engagement_receipt_logs_once_without_stacklevel_conflict(routed, monkeypatch): +def test_sm100a_engagement_receipt_logs_once_without_stacklevel_conflict(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) records = [] def capture_log(level, msg, *args, **kwargs): @@ -425,9 +426,9 @@ def test_sm100a_engagement_receipt_logs_once_without_stacklevel_conflict(routed, {"stacklevel": 2})] -def test_env_on_grad_inputs_fall_back_to_triton(routed, monkeypatch): +def test_env_on_grad_inputs_fall_back_to_triton(routed, env_overrides): fake_sm, fake_triton, run, _ = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) run(requires_grad=True) assert fake_triton.calls == 1 assert fake_sm.calls == [] @@ -436,10 +437,10 @@ def test_env_on_grad_inputs_fall_back_to_triton(routed, monkeypatch): assert len(fake_sm.calls) == 1 -def test_env_on_unsupported_warns_once_and_falls_back(routed, monkeypatch): +def test_env_on_unsupported_warns_once_and_falls_back(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed fake_sm.supported = False - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) run() @@ -451,10 +452,10 @@ def test_env_on_unsupported_warns_once_and_falls_back(routed, monkeypatch): assert VSA_SM100A_ENV in warnings[0] and "is_supported" in warnings[0] -def test_env_on_missing_module_warns_and_falls_back(routed, monkeypatch): +def test_env_on_missing_module_warns_and_falls_back(routed, monkeypatch, env_overrides): fake_sm, fake_triton, run, _ = routed monkeypatch.setattr(vsa_h3, "_sm100a", None) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) warnings = [] monkeypatch.setattr(vsa_h3.logger, "warning_once", warnings.append) run() @@ -462,10 +463,10 @@ def test_env_on_missing_module_warns_and_falls_back(routed, monkeypatch): assert warnings and "not installed" in warnings[0] -def test_env_on_no_grad_context_detaches_route_from_leaf_flags(routed, monkeypatch): +def test_env_on_no_grad_context_detaches_route_from_leaf_flags(routed, env_overrides): """A requires_grad leaf under torch.no_grad() is still a no-grad forward.""" fake_sm, fake_triton, run, meta = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) q, k, v = _tiled_qkv(meta, requires_grad=True) with torch.no_grad(): @@ -483,15 +484,15 @@ def test_env_on_no_grad_context_detaches_route_from_leaf_flags(routed, monkeypat ], ) def test_preprocess_adds_partner_only_for_odd_no_grad_sm100a( - monkeypatch, + env_overrides, env_enabled, requires_grad, extra_tiles, ): if env_enabled: - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) else: - monkeypatch.delenv(VSA_SM100A_ENV, raising=False) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) meta = _build_meta() n_tiles = meta.variable_block_sizes.numel() assert n_tiles % 2 == 1 @@ -513,9 +514,9 @@ def test_preprocess_adds_partner_only_for_odd_no_grad_sm100a( assert torch.count_nonzero(tiled[:, n_tiles * 64:]) == 0 -def test_odd_partner_is_visible_only_to_sm100a_transport(routed, monkeypatch): +def test_odd_partner_is_visible_only_to_sm100a_transport(routed, env_overrides): fake_sm, fake_triton, _, meta = routed - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) raw_qkv = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) tiled_qkv = impl.preprocess_qkv(raw_qkv, meta) @@ -540,10 +541,10 @@ def test_odd_partner_is_visible_only_to_sm100a_transport(routed, monkeypatch): assert output.shape == (1, logical_seq_len, _HEADS, _DIM) -def test_unsupported_odd_route_strips_partner_before_triton(routed, monkeypatch): +def test_unsupported_odd_route_strips_partner_before_triton(routed, monkeypatch, env_overrides): fake_sm, _, _, meta = routed fake_sm.supported = False - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) monkeypatch.setattr(vsa_h3.logger, "warning_once", lambda *_args, **_kwargs: None) calls = [] @@ -568,8 +569,8 @@ def test_unsupported_odd_route_strips_partner_before_triton(routed, monkeypatch) assert torch.equal(sizes, meta.variable_block_sizes) -def test_geometry_change_clears_reused_partner_tile(monkeypatch): - monkeypatch.setenv(VSA_SM100A_ENV, "1") +def test_geometry_change_clears_reused_partner_tile(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) builder = MiniMaxH3VSAMetadataBuilder() even_meta = _build_meta(prefix_segments=(70, 70), builder=builder) odd_meta = _build_meta(prefix_segments=(70, 30), builder=builder) @@ -605,7 +606,7 @@ def test_forward_rejects_non_contract_transport_shapes(routed): impl.forward(logical, logical, logical, partner, meta) -def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): +def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch, env_overrides): """Exercise the padded route through the actual data-center Blackwell extension.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") @@ -620,7 +621,7 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): raw_qkv = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) with torch.inference_mode(): - monkeypatch.delenv(VSA_SM100A_ENV, raising=False) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) triton_tiled = impl.preprocess_qkv(raw_qkv, meta) triton_query, triton_key, triton_value = triton_tiled.chunk(3, dim=0) triton_output = impl.postprocess_output( @@ -630,7 +631,7 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): raise AssertionError("odd no-grad VSA-H3 unexpectedly fell back to Triton-64") monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", reject_triton) - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) sm100a_tiled = impl.preprocess_qkv(raw_qkv, meta) sm100a_query, sm100a_key, sm100a_value = sm100a_tiled.chunk(3, dim=0) sm100a_output = impl.postprocess_output( @@ -642,7 +643,7 @@ def test_real_sm100a_odd_route_matches_triton_oracle(monkeypatch): torch.testing.assert_close(sm100a_output.float(), triton_output.float(), atol=0.04, rtol=0.02) -def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(monkeypatch): +def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(env_overrides): """Capture #1745's odd transport buffer mutation with real Inductor.""" if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: pytest.skip("requires a GB200/B300 (sm_100a/sm_103a) compute node") @@ -650,7 +651,7 @@ def test_real_sm100a_odd_preprocess_forward_postprocess_fullgraph(monkeypatch): pytest.skip("requires a fastvideo_kernel build containing data-center Blackwell VSA") device = torch.device("cuda") - monkeypatch.setenv(VSA_SM100A_ENV, "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) meta = _build_meta(device=device, sparsity=0.5) assert meta.variable_block_sizes.numel() % 2 == 1 impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5) diff --git a/fastvideo/tests/conftest.py b/fastvideo/tests/conftest.py index 66c1e71ce..a010e5ecb 100644 --- a/fastvideo/tests/conftest.py +++ b/fastvideo/tests/conftest.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import contextlib + import pytest import torch @@ -20,3 +22,15 @@ def distributed_setup(): yield cleanup_dist_env_and_memory() + + +@pytest.fixture +def env_overrides(): + """Keep environment overrides until the end of the test. + + Enter ``envs.NAME.override(...)`` or ``envs.override_external(...)`` with + ``env_overrides.enter_context(...)``; the previous values come back at + teardown. See docs/contributing/env_vars.md. + """ + with contextlib.ExitStack() as stack: + yield stack diff --git a/fastvideo/tests/contract/test_env_policy.py b/fastvideo/tests/contract/test_env_policy.py index 99f882ed1..c84eb44f2 100644 --- a/fastvideo/tests/contract/test_env_policy.py +++ b/fastvideo/tests/contract/test_env_policy.py @@ -5,7 +5,8 @@ docs/contributing/env_vars.md. The test parses every Python file under fastvideo/ (except fastvideo/third_party/ and the registry fastvideo/envs.py) with ``ast`` and reports code that: -- reads the environment directly for a name outside EXTERNAL_ALLOWLIST; +- reads the environment directly for a name outside EXTERNAL_ALLOWLIST (test + code under fastvideo/tests/ may also read CI_ONLY_VARIABLES); - writes the environment directly (os.environ, os.putenv, monkeypatch.setenv), or through the envs.*_external helpers for a name outside EXTERNAL_WRITE_ALLOWLIST; @@ -42,6 +43,7 @@ REGISTRY_PATH = PACKAGE_ROOT / "envs.py" DOC_PATH = REPO_ROOT / "docs" / "contributing" / "env_vars.md" POLICY_DOC = "docs/contributing/env_vars.md" EXCLUDED_DIRS = (PACKAGE_ROOT / "third_party", ) +TESTS_ROOT = PACKAGE_ROOT / "tests" DOC_TABLE_BEGIN = "" DOC_TABLE_END = "" @@ -68,6 +70,51 @@ EXTERNAL_ALLOWLIST = { "CEREBRAS_API_KEY": "Cerebras API key for the streaming prompt provider.", "GROQ_API_KEY": "Groq API key for the streaming prompt provider.", "RAY_USAGE_STATS_ENABLED": "Ray usage-statistics switch.", + "BUILDKITE_*": "Set by the Buildkite agent in each CI job.", + "GITHUB_STEP_SUMMARY": "Set by GitHub Actions.", + "WANDB_MODE": "wandb setting.", + "WANDB_API_KEY": "wandb API key.", + "UV_TORCH_BACKEND": "uv setting that selects the PyTorch build.", + "CFLAGS": "C compiler flags.", + "CXXFLAGS": "C++ compiler flags.", + "LDFLAGS": "Linker flags.", + "CMAKE_ARGS": "Extra CMake arguments.", + "CC": "C compiler.", + "CXX": "C++ compiler.", + "LD": "Linker.", + "CUDACXX": "CUDA compiler.", + "CUBLAS_WORKSPACE_CONFIG": "cuBLAS workspace setting.", + "TORCHDYNAMO_DISABLE": "PyTorch Dynamo switch.", + "PYTORCH_MPS_*": "PyTorch MPS memory settings.", + "USER": "Login name of the current user.", + "PYTHONPATH": "Python module search path.", + "FASTVIDEO_VSA_CUTEDSL": "Read by fastvideo-kernel, which is outside this policy.", + "FASTVIDEO_VSA_TRITON": "Read by fastvideo-kernel, which is outside this policy.", + "FASTVIDEO_KERNEL_VSA_FORCE_TRITON": "Read by fastvideo-kernel, which is outside this policy.", +} + +# Variables that FastVideo's own CI and CI tooling define. They keep their names +# and test code under fastvideo/tests/ reads them directly; the value names what sets them. +CI_ONLY_VARIABLES = { + "TEST_SCOPE": ".github/workflows/ci-slash-commands.yml, ci-trigger-full-suite.yml, ci-scheduled-ssim.yml", + "PERF_RUN_SOURCE": ".buildkite/scripts/lanes/performance.sh", + "PERF_UPLOAD_POLICY": ".buildkite/scripts/lanes/performance.sh", + "PERF_REPORTS_DIR": ".buildkite/scripts/lanes/performance.sh", + "PERF_PYTEST_RC": ".buildkite/scripts/lanes/performance.sh", + "IMAGE_VERSION": ".buildkite/pipeline.yml", + "PERFORMANCE_TRACKING_ROOT": ".buildkite/scripts/lanes/performance.sh", + "HF_REPO_ID": "Performance tracking tooling (fastvideo/tests/modal/pr_test.py).", + "PERFORMANCE_RESEED_STAGING_ROOT": "Performance reseed tooling (.agents/skills/reseed-performance-baseline).", + "DASHBOARD_DAYS": "Performance dashboard tooling (fastvideo/tests/performance/dashboard.py).", + "GPU_BACKEND": "fastvideo-kernel/CMakeLists.txt (kernel build).", + "FASTVIDEO_PERFORMANCE_PROFILE_VERSION": "fastvideo/tests/modal/launch_l40s_job.py; recorded in performance identity.", + "FASTVIDEO_CONTAINER_IMAGE_REF": "Modal CI launchers; recorded in performance identity.", + "FASTVIDEO_SSIM_MODEL_ID": "fastvideo/tests/ssim/ci_runner.py (SSIM CI scheduler).", + "FASTVIDEO_MODAL_IMAGE": "Modal launch tooling in fastvideo/tests/modal/.", + "FASTVIDEO_MODAL_VOLUME": "Modal launch tooling in fastvideo/tests/modal/.", + "FASTVIDEO_KERNEL_CACHE_ROOT": "Kernel build-cache tooling (fastvideo/tests/modal/kernel_build_cache.py).", + "FASTVIDEO_KERNEL_PREBUILT_INFO": "Kernel build-cache tooling (fastvideo/tests/modal/kernel_build_cache.py).", + "FASTVIDEO_SSIM_BOOTSTRAP_MODE": ".buildkite/scripts/lanes/ssim.sh; --ssim-bootstrap-mode in fastvideo/tests/ssim/conftest.py.", } # Variables that FastVideo sets for other tools. Code writes them only through @@ -87,8 +134,9 @@ EXTERNAL_WRITE_ALLOWLIST = { "VIDEO_MAX_PIXELS": "Pixel budget that qwen_vl_utils reads for Qwen2.5-Omni.", "RAY_USAGE_STATS_ENABLED": "Turns off Ray usage statistics unless the user turned them on.", "CUTE_DSL_ENABLE_TVM_FFI": "CUTLASS DSL switch that the NVFP4 FlashAttention-4 path needs.", + "CUBLAS_WORKSPACE_CONFIG": "cuBLAS workspace size; golden-gate tests set it before cuBLAS initializes.", } -EXTERNAL_WRITE_HELPERS = {"set_external", "setdefault_external", "unset_external"} +EXTERNAL_WRITE_HELPERS = {"set_external", "setdefault_external", "unset_external", "override_external"} # Methods that registry fields expose; ``get`` and ``is_set`` count as reads. REGISTRY_METHODS = {"get", "set", "override", "is_set", "clear"} @@ -110,365 +158,36 @@ KNOWN_VIOLATIONS: dict[str, int] = { 'fastvideo/performance/hf_store.py: read HF_REPO_ID': 1, 'fastvideo/performance/hf_store.py: read PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS': 1, 'fastvideo/performance_dashboard/api.py: read PERFORMANCE_TRACKING_ROOT': 1, - 'fastvideo/tests/api/test_attention_selector_resolution.py: write FASTVIDEO_ATTENTION_BACKEND': 17, - 'fastvideo/tests/attention/test_flash_attn_no_pad_resolver.py: write FASTVIDEO_FA4': 3, - 'fastvideo/tests/attention/test_flash_attn_nvfp4_opt_out.py: write FASTVIDEO_NVFP4_FA4': 4, - 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_KERNEL_VSA_FORCE_TRITON': 1, - 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_VSA_CUTEDSL': 2, - 'fastvideo/tests/attention/test_vsa_h3_backward.py: write FASTVIDEO_VSA_TRITON': 2, - 'fastvideo/tests/attention/test_vsa_h3_sm100a_route.py: write ': 22, 'fastvideo/tests/contract/test_profiler_regions.py: whole-environ': 2, - 'fastvideo/tests/contract/test_profiler_regions.py: write FASTVIDEO_NVTX_PROFILE': 4, - 'fastvideo/tests/contract/test_ssim_ci_runner.py: write CUDA_VISIBLE_DEVICES': 1, - 'fastvideo/tests/contract/test_tensor_golden_isolation.py: read FASTVIDEO_ATTENTION_BACKEND': 2, - 'fastvideo/tests/contract/test_tensor_golden_isolation.py: write FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/contract/test_wan_validation_order.py: whole-environ': 1, 'fastvideo/tests/distributed/test_sp_hunyuanvideo.py: whole-environ': 1, 'fastvideo/tests/distributed/test_sp_lingbot_video.py: whole-environ': 1, 'fastvideo/tests/distributed/test_sp_ltx2.py: whole-environ': 1, 'fastvideo/tests/distributed/test_sp_wan.py: whole-environ': 1, 'fastvideo/tests/distributed/test_ulysses_a2a_parity.py: whole-environ': 1, - 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: read FASTVIDEO_ULYSSES_FAULT_RANK': 1, - 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: read FASTVIDEO_ULYSSES_FAULT_STAGE': 1, 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: whole-environ': 1, - 'fastvideo/tests/distributed/test_ulysses_fault_injection.py: write FASTVIDEO_ULYSSES_A2A': 1, - 'fastvideo/tests/encoders/test_clip_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_clip_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_hyt5_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_hyt5_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_llama_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_llama_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_qwen2_5_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_qwen2_5_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_siglip_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_siglip_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/encoders/test_t5_encoder.py: write MASTER_ADDR': 1, - 'fastvideo/tests/encoders/test_t5_encoder.py: write MASTER_PORT': 1, - 'fastvideo/tests/entrypoints/streaming/test_prompt_providers.py: write CEREBRAS_API_KEY': 3, - 'fastvideo/tests/entrypoints/streaming/test_prompt_providers.py: write GROQ_API_KEY': 2, 'fastvideo/tests/entrypoints/test_openai_api_integration.py: whole-environ': 1, 'fastvideo/tests/entrypoints/test_openai_video_client.py: whole-environ': 1, - 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_FFMPEG_BIN': 1, - 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_OUTPUT_PIX_FMT': 1, - 'fastvideo/tests/entrypoints/test_video_generator.py: write FASTVIDEO_VIDEO_CODEC': 1, - 'fastvideo/tests/golden_gate/_harness.py: read ': 1, - 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_FA4': 2, - 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_GOLDEN_GATE_DIR': 1, - 'fastvideo/tests/golden_gate/_harness.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, - 'fastvideo/tests/golden_gate/_harness.py: write CUBLAS_WORKSPACE_CONFIG': 1, - 'fastvideo/tests/golden_gate/_harness.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/golden_gate/_harness.py: write LOCAL_RANK': 1, - 'fastvideo/tests/golden_gate/_harness.py: write MASTER_ADDR': 1, - 'fastvideo/tests/golden_gate/_harness.py: write MASTER_PORT': 1, - 'fastvideo/tests/golden_gate/_harness.py: write RANK': 1, - 'fastvideo/tests/golden_gate/_harness.py: write WORLD_SIZE': 1, - 'fastvideo/tests/golden_gate/_tensor_golden.py: read FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/golden_gate/_tensor_golden.py: write FASTVIDEO_ATTENTION_BACKEND': 3, - 'fastvideo/tests/golden_gate/test_kandinsky5.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/golden_gate/test_kandinsky5.py: write FASTVIDEO_FA4': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_ATTENTION_BACKEND': 2, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_GATE_GOLDEN_DIR': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_GATE_LAYER': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: read MINIMAX_H3_MODEL_ROOT': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write CUBLAS_WORKSPACE_CONFIG': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write LOCAL_RANK': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write MASTER_ADDR': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write MASTER_PORT': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write RANK': 1, - 'fastvideo/tests/golden_gate/test_minimax_h3_t2v.py: write WORLD_SIZE': 1, - 'fastvideo/tests/golden_gate/test_wan_denoising.py: write FASTVIDEO_CFG_GATE_STEP': 1, - 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_ACTIVATIONS': 6, - 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_LAYERS': 5, - 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_OUTPUT': 5, - 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_STATS': 1, - 'fastvideo/tests/hooks/test_activation_trace.py: write FASTVIDEO_TRACE_STEPS': 1, - 'fastvideo/tests/inference/bsa/test_bsa_inference.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write MASTER_ADDR': 1, - 'fastvideo/tests/inference/lora/test_lora_inference_similarity.py: write MASTER_PORT': 1, 'fastvideo/tests/inference/test_basic_fasth3_profile.py: write ': 2, - 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_DISABLE_ATTENTION_COMPILE': 9, - 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_H3_VSA_PROBE': 4, - 'fastvideo/tests/inference/test_inference_regional_compile.py: write FASTVIDEO_VSA_SM100A': 3, - 'fastvideo/tests/inference/vmoba/test_vmoba_inference.py: write FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/layers/test_rmsnorm_forward_dispatch.py: whole-environ': 1, - 'fastvideo/tests/mlx/test_memory_limits.py: read PYTORCH_MPS_HIGH_WATERMARK_RATIO': 1, - 'fastvideo/tests/mlx/test_memory_limits.py: read PYTORCH_MPS_LOW_WATERMARK_RATIO': 1, - 'fastvideo/tests/mlx/test_memory_limits.py: write PYTORCH_MPS_HIGH_WATERMARK_RATIO': 1, - 'fastvideo/tests/mlx/test_memory_limits.py: write PYTORCH_MPS_LOW_WATERMARK_RATIO': 1, - 'fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py: read FASTVIDEO_MLX_DQ_GEMM': 1, - 'fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py: write FASTVIDEO_MLX_DQ_GEMM': 11, - 'fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py: write MASTER_ADDR': 1, - 'fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py: write MASTER_PORT': 1, - 'fastvideo/tests/mlx/test_mlx_taeh3.py: read TAEH3_REFERENCE_DIR': 1, - 'fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py: write HOME': 1, - 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: read FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY': 1, - 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: read FASTVIDEO_WAN22_5B_ROOT': 1, - 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: write MASTER_ADDR': 1, - 'fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py: write MASTER_PORT': 1, - 'fastvideo/tests/mlx/tiny_h3.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/mlx/tiny_h3.py: write MASTER_ADDR': 1, - 'fastvideo/tests/mlx/tiny_h3.py: write MASTER_PORT': 1, - 'fastvideo/tests/mlx/tiny_wan.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/mlx/tiny_wan.py: write MASTER_ADDR': 1, - 'fastvideo/tests/mlx/tiny_wan.py: write MASTER_PORT': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read ': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read CFLAGS': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read CMAKE_ARGS': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read CXXFLAGS': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read FASTVIDEO_KERNEL_CACHE_ROOT': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read FASTVIDEO_KERNEL_PREBUILT_INFO': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read GPU_BACKEND': 1, - 'fastvideo/tests/modal/kernel_build_cache.py: read LDFLAGS': 1, 'fastvideo/tests/modal/kernel_build_cache.py: whole-environ': 1, - 'fastvideo/tests/modal/kernel_cache_smoke.py: read BUILDKITE_COMMIT': 1, - 'fastvideo/tests/modal/kernel_cache_smoke.py: read BUILDKITE_REPO': 1, - 'fastvideo/tests/modal/kernel_cache_smoke.py: read FASTVIDEO_MODAL_IMAGE': 1, - 'fastvideo/tests/modal/kernel_cache_smoke.py: read IMAGE_VERSION': 1, - 'fastvideo/tests/modal/kernel_cache_smoke.py: read UV_TORCH_BACKEND': 1, 'fastvideo/tests/modal/launch_l40s_job.py: read ': 2, - 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_COMMIT': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_PULL_REQUEST': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read BUILDKITE_REPO': 1, 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_ATTENTION_BACKEND': 1, 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_MODAL_IMAGE': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_MODAL_VOLUME': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 2, - 'fastvideo/tests/modal/launch_l40s_job.py: read IMAGE_VERSION': 1, - 'fastvideo/tests/modal/launch_l40s_job.py: read UV_TORCH_BACKEND': 1, 'fastvideo/tests/modal/launch_l40s_job.py: whole-environ': 1, - 'fastvideo/tests/modal/modal_image_utils.py: read UV_TORCH_BACKEND': 1, 'fastvideo/tests/modal/pr_test.py: read ': 2, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BRANCH': 1, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BUILD_ID': 1, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_BUILD_URL': 1, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_COMMIT': 2, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_JOB_ID': 1, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_PULL_REQUEST': 2, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_REPO': 2, - 'fastvideo/tests/modal/pr_test.py: read BUILDKITE_SOURCE': 1, 'fastvideo/tests/modal/pr_test.py: read FASTVIDEO_FA4': 1, 'fastvideo/tests/modal/pr_test.py: read HF_API_KEY': 1, - 'fastvideo/tests/modal/pr_test.py: read IMAGE_VERSION': 1, - 'fastvideo/tests/modal/pr_test.py: read TEST_SCOPE': 1, - 'fastvideo/tests/modal/pr_test.py: read UV_TORCH_BACKEND': 1, - 'fastvideo/tests/modal/pr_test.py: read WANDB_API_KEY': 1, 'fastvideo/tests/modal/ssim_test.py: read ': 1, - 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_COMMIT': 1, - 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_PULL_REQUEST': 1, - 'fastvideo/tests/modal/ssim_test.py: read BUILDKITE_REPO': 1, 'fastvideo/tests/modal/ssim_test.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/modal/ssim_test.py: read IMAGE_VERSION': 1, - 'fastvideo/tests/modal/ssim_test.py: read UV_TORCH_BACKEND': 1, 'fastvideo/tests/modal/ssim_test.py: whole-environ': 2, - 'fastvideo/tests/modal/test_kernel_build_cache.py: write ': 2, - 'fastvideo/tests/modal/test_kernel_build_cache.py: write TORCH_CUDA_ARCH_LIST': 2, - 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_COMMIT': 2, - 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_PULL_REQUEST': 2, - 'fastvideo/tests/modal/test_pr_test.py: write BUILDKITE_REPO': 2, - 'fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py: write WANDB_MODE': 1, - 'fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py: write WANDB_MODE': 1, - 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read KANDINSKY5_E2E_NUM_GPUS': 1, - 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: read KANDINSKY5_E2E_WRITE_REFERENCE': 1, 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: whole-environ': 3, - 'fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py: write WANDB_MODE': 1, - 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: read FASTVIDEO_NIGHTLY': 1, 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: whole-environ': 2, - 'fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py: write FASTVIDEO_NIGHTLY': 1, - 'fastvideo/tests/nightly/test_e2e_overfit_single_sample.py: write WANDB_MODE': 1, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BRANCH': 2, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BUILD_ID': 1, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_BUILD_URL': 1, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_COMMIT': 3, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_JOB_ID': 1, - 'fastvideo/tests/performance/compare_baseline.py: read BUILDKITE_PULL_REQUEST': 2, - 'fastvideo/tests/performance/compare_baseline.py: read GITHUB_STEP_SUMMARY': 1, - 'fastvideo/tests/performance/compare_baseline.py: read PERFORMANCE_TRACKING_ROOT': 1, - 'fastvideo/tests/performance/compare_baseline.py: read PERF_PYTEST_RC': 4, - 'fastvideo/tests/performance/compare_baseline.py: read PERF_REPORTS_DIR': 1, - 'fastvideo/tests/performance/compare_baseline.py: read PERF_RUN_SOURCE': 1, - 'fastvideo/tests/performance/compare_baseline.py: read PERF_UPLOAD_POLICY': 1, - 'fastvideo/tests/performance/compare_baseline.py: read TEST_SCOPE': 2, - 'fastvideo/tests/performance/dashboard.py: read BUILDKITE_COMMIT': 1, - 'fastvideo/tests/performance/dashboard.py: read DASHBOARD_DAYS': 1, - 'fastvideo/tests/performance/dashboard.py: read PERFORMANCE_TRACKING_ROOT': 1, - 'fastvideo/tests/performance/dashboard.py: read PERF_REPORTS_DIR': 1, - 'fastvideo/tests/performance/identity.py: read FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/performance/identity.py: whole-environ': 1, - 'fastvideo/tests/performance/seed_baseline.py: read PERFORMANCE_RESEED_STAGING_ROOT': 1, - 'fastvideo/tests/performance/seed_baseline.py: read PERFORMANCE_TRACKING_ROOT': 1, - 'fastvideo/tests/performance/seed_baseline.py: read USER': 1, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BRANCH': 3, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BUILD_ID': 3, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_BUILD_URL': 2, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_JOB_ID': 3, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write BUILDKITE_PULL_REQUEST': 3, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write PERF_PYTEST_RC': 8, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write PERF_RUN_SOURCE': 23, - 'fastvideo/tests/performance/test_compare_baseline_policy.py: write TEST_SCOPE': 3, - 'fastvideo/tests/performance/test_dashboard_service.py: write ': 1, - 'fastvideo/tests/performance/test_dashboard_service.py: write HF_TOKEN': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BRANCH': 2, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BUILD_ID': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_BUILD_URL': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_COMMIT': 2, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_JOB_ID': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read BUILDKITE_PULL_REQUEST': 2, - 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read FASTVIDEO_STAGE_LOGGING': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read IMAGE_VERSION': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read PERF_RUN_SOURCE': 1, - 'fastvideo/tests/performance/test_inference_performance.py: read TEST_SCOPE': 2, - 'fastvideo/tests/performance/test_inference_performance.py: write FASTVIDEO_STAGE_LOGGING': 3, - 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_ATTENTION_BACKEND': 5, - 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_CONTAINER_IMAGE_REF': 2, - 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_FA4': 2, - 'fastvideo/tests/performance/test_inference_performance_identity.py: write FASTVIDEO_PERFORMANCE_PROFILE_VERSION': 1, - 'fastvideo/tests/performance/test_inference_performance_identity.py: write IMAGE_VERSION': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BRANCH': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BUILD_ID': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_BUILD_URL': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_COMMIT': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_JOB_ID': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write BUILDKITE_PULL_REQUEST': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write PERF_RUN_SOURCE': 1, - 'fastvideo/tests/performance/test_inference_performance_result_schema.py: write TEST_SCOPE': 1, - 'fastvideo/tests/ssim/bootstrap_references.py: read ': 3, 'fastvideo/tests/ssim/ci_runner.py: whole-environ': 1, - 'fastvideo/tests/ssim/conftest.py: read ': 1, - 'fastvideo/tests/ssim/conftest.py: read FASTVIDEO_SSIM_MODEL_ID': 1, - 'fastvideo/tests/ssim/conftest.py: read FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD': 1, - 'fastvideo/tests/ssim/conftest.py: write ': 3, - 'fastvideo/tests/ssim/inference_similarity_utils.py: read FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/ssim/inference_similarity_utils.py: write FASTVIDEO_ATTENTION_BACKEND': 3, - 'fastvideo/tests/ssim/reference_utils.py: read FASTVIDEO_SSIM_FULL_QUALITY': 1, 'fastvideo/tests/ssim/reference_videos_cli.py: read ': 1, - 'fastvideo/tests/ssim/reference_videos_cli.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO': 1, - 'fastvideo/tests/ssim/reference_videos_cli.py: read FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE': 1, - 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_AR_SSIM_MODEL_PATH': 1, - 'fastvideo/tests/ssim/test_dreamx_world_similarity.py: read DREAMX_WORLD_SSIM_MODEL_PATH': 1, - 'fastvideo/tests/ssim/test_flux_t2i_similarity.py: read FLUX_T2I_MODEL_DIR': 1, - 'fastvideo/tests/ssim/test_gamecraft_similarity.py: read GAMECRAFT_MODEL_PATH': 1, - 'fastvideo/tests/ssim/test_gamecraft_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 2, - 'fastvideo/tests/ssim/test_gen3c_similarity.py: read GEN3C_MODEL_PATH': 1, - 'fastvideo/tests/ssim/test_gen3c_similarity.py: read GEN3C_TEST_IMAGE_PATH': 1, - 'fastvideo/tests/ssim/test_gen3c_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/ssim/test_glm_image_similarity.py: read GLM_IMAGE_LOCAL_WEIGHTS_DIR': 1, - 'fastvideo/tests/ssim/test_glm_image_similarity.py: read GLM_IMAGE_MODEL_DIR': 1, - 'fastvideo/tests/ssim/test_kandinsky5_similarity.py: read FASTVIDEO_FA4': 1, - 'fastvideo/tests/ssim/test_kandinsky5_similarity.py: write FASTVIDEO_FA4': 3, - 'fastvideo/tests/ssim/test_lingbot_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/ssim/test_longcat_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 3, - 'fastvideo/tests/ssim/test_matrixgame2_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/ssim/test_matrixgame3_similarity.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/ssim/test_sd35_similarity.py: read SD35_MODEL_DIR': 1, - 'fastvideo/tests/ssim/test_zimage_similarity.py: read ZIMAGE_MODEL_DIR': 1, - 'fastvideo/tests/ssim/test_zimage_similarity.py: read ZIMAGE_MODEL_REVISION': 1, - 'fastvideo/tests/stages/_denoising_fixtures.py: write FASTVIDEO_CFG_GATE_STEP': 2, - 'fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/stages/test_minimax_h3_encoding_offload.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/train/methods/grad_norm_regression.py: read FASTVIDEO_GRADNORM_UPDATE': 1, - 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_cosmos_finetune.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_ltx2_finetune.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_ltx2_finetune.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_matrixgame2_finetune.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_matrixgame2_finetune.py: write MASTER_PORT': 1, 'fastvideo/tests/train/methods/test_minimax_h3_finetune.py: whole-environ': 2, - 'fastvideo/tests/train/methods/test_streaming_long_tuning.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_streaming_long_tuning.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_wan_causal_cd.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_wan_causal_cd.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_wan_causal_dfsft.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_wan_causal_dfsft.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_wan_causal_tfsft.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_wan_causal_tfsft.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/methods/test_wan_finetune.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/methods/test_wan_finetune.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_cosmos.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_cosmos.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_hunyuan.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_hunyuan.py: write MASTER_PORT': 1, 'fastvideo/tests/train/models/test_load_kandinsky5.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/train/models/test_load_kandinsky5.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_kandinsky5.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_longcat.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_longcat.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_ltx2.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_ltx2.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_matrixgame2.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_matrixgame2.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_wan.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_wan.py: write MASTER_PORT': 1, - 'fastvideo/tests/train/models/test_load_wan_causal.py: write MASTER_ADDR': 1, - 'fastvideo/tests/train/models/test_load_wan_causal.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/VSA/test_training_loss_VSA.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/training/VSA/test_training_loss_VSA.py: write WANDB_MODE': 1, - 'fastvideo/tests/training/Vanilla/mfu_calculation.py: read PYTHONPATH': 1, - 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write MASTER_ADDR': 1, - 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write PYTHONPATH': 1, - 'fastvideo/tests/training/Vanilla/mfu_calculation.py: write WANDB_MODE': 1, - 'fastvideo/tests/training/Vanilla/test_training_loss.py: write MASTER_ADDR': 1, - 'fastvideo/tests/training/Vanilla/test_training_loss.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/Vanilla/test_training_loss.py: write WANDB_MODE': 1, 'fastvideo/tests/training/distill/test_anyflow_smoke.py: whole-environ': 2, - 'fastvideo/tests/training/distill/test_distill_dmd.py: write MASTER_ADDR': 1, - 'fastvideo/tests/training/distill/test_distill_dmd.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/distill/test_distill_dmd.py: write WANDB_MODE': 1, - 'fastvideo/tests/training/lora/test_lora_training.py: write MASTER_ADDR': 1, - 'fastvideo/tests/training/lora/test_lora_training.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/lora/test_lora_training.py: write WANDB_MODE': 1, - 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write MASTER_ADDR': 1, - 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write MASTER_PORT': 1, - 'fastvideo/tests/training/self-forcing/test_self_forcing.py: write WANDB_MODE': 1, - 'fastvideo/tests/transformers/test_cosmos.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_cosmos.py: write MASTER_PORT': 1, - 'fastvideo/tests/transformers/test_cosmos2_5.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_cosmos2_5.py: write MASTER_PORT': 1, - 'fastvideo/tests/transformers/test_flux.py: read FLUX_TRANSFORMER_PATH': 1, - 'fastvideo/tests/transformers/test_flux.py: write FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/transformers/test_flux.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_flux.py: write MASTER_PORT': 1, - 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write DISABLE_SP': 1, - 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write MASTER_PORT': 1, - 'fastvideo/tests/transformers/test_hunyuangamecraft.py: write TORCHDYNAMO_DISABLE': 1, - 'fastvideo/tests/transformers/test_hyworld.py: read FASTVIDEO_ATTENTION_BACKEND': 1, - 'fastvideo/tests/transformers/test_hyworld.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_hyworld.py: write MASTER_PORT': 1, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write FASTVIDEO_ATTENTION_BACKEND': 2, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write FASTVIDEO_FA4': 1, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write LOCAL_RANK': 2, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write MASTER_ADDR': 2, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write MASTER_PORT': 2, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write RANK': 2, - 'fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py: write WORLD_SIZE': 2, - 'fastvideo/tests/transformers/test_wanvideo.py: write MASTER_ADDR': 1, - 'fastvideo/tests/transformers/test_wanvideo.py: write MASTER_PORT': 1, - 'fastvideo/tests/vaes/test_hunyuan15_vae.py: write MASTER_ADDR': 1, - 'fastvideo/tests/vaes/test_hunyuan15_vae.py: write MASTER_PORT': 1, - 'fastvideo/tests/vaes/test_wan_vae.py: write MASTER_ADDR': 1, - 'fastvideo/tests/vaes/test_wan_vae.py: write MASTER_PORT': 1, - 'fastvideo/tests/worker/test_gpu_worker.py: write FASTVIDEO_NVTX_PROFILE': 2, - 'fastvideo/tests/worker/test_gpu_worker.py: write LOCAL_RANK': 1, 'fastvideo/train/entrypoint/train.py: write FASTVIDEO_ATTENTION_BACKEND': 2, 'fastvideo/utils.py: read ': 5, 'fastvideo/utils.py: write ': 1, @@ -490,6 +209,18 @@ def load_registry() -> ModuleType: return registry +def is_external_write_allowed(helper: str, name: str, in_tests: bool) -> bool: + """Whether an envs.*_external call may write ``name``. + + override_external restores the previous value, so test code under fastvideo/tests/ + may use it for any variable that code reads directly; library code and the other + helpers need the write allowlist. + """ + if name in EXTERNAL_WRITE_ALLOWLIST: + return True + return in_tests and helper == "override_external" and (is_allowlisted(name) or name in CI_ONLY_VARIABLES) + + def is_allowlisted(name: str) -> bool: for pattern in EXTERNAL_ALLOWLIST: if pattern.endswith("*") and name.startswith(pattern[:-1]): @@ -506,8 +237,9 @@ class EnvAccessScanner: ``envs.X.get()`` and ``envs.X.is_set()`` calls by variable name. """ - def __init__(self, tree: ast.Module, registry_names: set[str]) -> None: + def __init__(self, tree: ast.Module, registry_names: set[str], in_tests: bool) -> None: self.registry_names = registry_names + self.in_tests = in_tests self.violations: list[tuple[str, str, int]] = [] self.registry_reads: Counter[str] = Counter() self.parents: dict[ast.AST, ast.AST] = {} @@ -628,7 +360,7 @@ class EnvAccessScanner: or (isinstance(func, ast.Attribute) and func.attr in ("setenv", "delenv"))): self.violations.append(("write", name, node.lineno)) elif (isinstance(func, ast.Attribute) and func.attr in EXTERNAL_WRITE_HELPERS - and self._is_envs_module(func.value) and name not in EXTERNAL_WRITE_ALLOWLIST): + and self._is_envs_module(func.value) and not is_external_write_allowed(func.attr, name, self.in_tests)): self.violations.append(("write", name, node.lineno)) def _check_registry_use(self, node: ast.Attribute, parent: ast.AST | None, in_function: bool) -> None: @@ -667,10 +399,12 @@ def collect_violations(registry: ModuleType) -> tuple[dict[str, list[int]], Coun reads: Counter[str] = registry_internal_reads(registry_names) for path in scanned_files(): relative = path.relative_to(REPO_ROOT).as_posix() - scanner = EnvAccessScanner(ast.parse(path.read_text(encoding="utf-8"), filename=relative), registry_names) + in_tests = TESTS_ROOT in path.parents + scanner = EnvAccessScanner(ast.parse(path.read_text(encoding="utf-8"), filename=relative), registry_names, + in_tests) reads.update(scanner.registry_reads) for kind, name, line in scanner.violations: - if kind == "read" and is_allowlisted(name): + if kind == "read" and (is_allowlisted(name) or (in_tests and name in CI_ONLY_VARIABLES)): continue key = f"{relative}: {kind} {name}".rstrip() found.setdefault(key, []).append(line) @@ -709,7 +443,7 @@ def test_registry_method_counts_only_when_called(): "def f():\n" " getter = envs.FASTVIDEO_FA4.get\n" " return envs.FASTVIDEO_FA4.get()\n") - scanner = EnvAccessScanner(ast.parse(source), {"FASTVIDEO_FA4"}) + scanner = EnvAccessScanner(ast.parse(source), {"FASTVIDEO_FA4"}, in_tests=False) assert scanner.violations == [("bare-field", "FASTVIDEO_FA4", 3)] assert scanner.registry_reads["FASTVIDEO_FA4"] == 1 diff --git a/fastvideo/tests/contract/test_profiler_regions.py b/fastvideo/tests/contract/test_profiler_regions.py index e9b706eba..adf7d5ae3 100644 --- a/fastvideo/tests/contract/test_profiler_regions.py +++ b/fastvideo/tests/contract/test_profiler_regions.py @@ -20,6 +20,7 @@ from unittest.mock import Mock import pytest import torch +import fastvideo.envs as envs from fastvideo.profiler import nvtx_range # Five-window child: ops before any region, inside a region, between regions, @@ -113,11 +114,11 @@ def test_noop_without_profiler_dir(tmp_path): assert proc.returncode == 0, proc.stderr -def test_nvtx_range_disabled_is_noop(monkeypatch): +def test_nvtx_range_disabled_is_noop(monkeypatch, env_overrides): """Keep CUDA NVTX untouched when external profiling is disabled.""" range_push = Mock() range_pop = Mock() - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "0") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(False)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda.nvtx, "range_push", range_push) monkeypatch.setattr(torch.cuda.nvtx, "range_pop", range_pop) @@ -130,11 +131,11 @@ def test_nvtx_range_disabled_is_noop(monkeypatch): range_pop.assert_not_called() -def test_nvtx_range_without_cuda_is_noop(monkeypatch): +def test_nvtx_range_without_cuda_is_noop(monkeypatch, env_overrides): """Keep NVTX untouched when profiling is enabled on a CPU-only process.""" range_push = Mock() range_pop = Mock() - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(True)) monkeypatch.setattr(torch.cuda, "is_available", lambda: False) monkeypatch.setattr(torch.cuda.nvtx, "range_push", range_push) monkeypatch.setattr(torch.cuda.nvtx, "range_pop", range_pop) @@ -147,10 +148,10 @@ def test_nvtx_range_without_cuda_is_noop(monkeypatch): range_pop.assert_not_called() -def test_nvtx_range_enabled_orders_push_body_pop(monkeypatch): +def test_nvtx_range_enabled_orders_push_body_pop(monkeypatch, env_overrides): """Place the profiled body between one matching NVTX push and pop.""" events = [] - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(True)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda.nvtx, "range_push", lambda name: events.append(("push", name))) monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: events.append(("pop", None))) @@ -165,10 +166,10 @@ def test_nvtx_range_enabled_orders_push_body_pop(monkeypatch): ] -def test_nvtx_range_body_exception_pops_and_propagates(monkeypatch): +def test_nvtx_range_body_exception_pops_and_propagates(monkeypatch, env_overrides): """Balance the NVTX stack while preserving a body exception.""" events = [] - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(True)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda.nvtx, "range_push", lambda name: events.append(("push", name))) monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: events.append(("pop", None))) diff --git a/fastvideo/tests/contract/test_ssim_ci_runner.py b/fastvideo/tests/contract/test_ssim_ci_runner.py index 4c9765f48..744199855 100644 --- a/fastvideo/tests/contract/test_ssim_ci_runner.py +++ b/fastvideo/tests/contract/test_ssim_ci_runner.py @@ -6,6 +6,8 @@ import sys import pytest +import fastvideo.envs as envs + SSIM_DIR = Path(__file__).resolve().parents[1] / "ssim" sys.path.insert(0, str(SSIM_DIR)) @@ -76,10 +78,9 @@ def test_discovery_rejects_missing_or_unsafe_test_file_selections(tmp_path: Path discover_tasks(tmp_path, ["../test_alpha.py"]) -def test_visible_gpu_ids_preserve_the_slurm_lease(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3,2") - - assert visible_gpu_ids() == ["2", "3"] +def test_visible_gpu_ids_preserve_the_slurm_lease() -> None: + with envs.override_external("CUDA_VISIBLE_DEVICES", "2,3,2"): + assert visible_gpu_ids() == ["2", "3"] def test_parallel_tasks_get_offsets_in_the_runner_assigned_port_range() -> None: diff --git a/fastvideo/tests/contract/test_tensor_golden_isolation.py b/fastvideo/tests/contract/test_tensor_golden_isolation.py index 34665c6d6..651fe3639 100644 --- a/fastvideo/tests/contract/test_tensor_golden_isolation.py +++ b/fastvideo/tests/contract/test_tensor_golden_isolation.py @@ -1,22 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 """Golden recipes must not inherit or leak the model loader's compute policy.""" -import os - import pytest import torch +import fastvideo.envs as envs from fastvideo import utils from fastvideo.tests.golden_gate._tensor_golden import deterministic_forward @pytest.mark.parametrize("existing_policy", [False, True]) @pytest.mark.parametrize("fails", [False, True]) -def test_tensor_golden_restores_runtime_state(monkeypatch, existing_policy, fails): +def test_tensor_golden_restores_runtime_state(monkeypatch, env_overrides, existing_policy, fails): # No CUDA work is needed to check context ownership. monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: True) - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) previous_policy = utils.MixedPrecisionState(param_dtype=torch.float16) if existing_policy: monkeypatch.setattr(utils._mixed_precision_state, "state", previous_policy, raising=False) @@ -32,7 +31,7 @@ def test_tensor_golden_restores_runtime_state(monkeypatch, existing_policy, fail try: with deterministic_forward("FLASH_ATTN"): assert utils.get_compute_dtype() == torch.get_default_dtype() - assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "FLASH_ATTN" + assert envs.FASTVIDEO_ATTENTION_BACKEND.get() == "FLASH_ATTN" assert torch.are_deterministic_algorithms_enabled() assert not torch.backends.cudnn.benchmark # Same process-state mutation performed by TransformerLoader. @@ -46,7 +45,7 @@ def test_tensor_golden_restores_runtime_state(monkeypatch, existing_policy, fail assert hasattr(utils._mixed_precision_state, "state") == existing_policy if existing_policy: assert utils.get_mixed_precision_state() is previous_policy - assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "TORCH_SDPA" + assert envs.FASTVIDEO_ATTENTION_BACKEND.get() == "TORCH_SDPA" assert torch.are_deterministic_algorithms_enabled() == deterministic assert torch.is_deterministic_algorithms_warn_only_enabled() == warn_only assert torch.backends.cudnn.benchmark == benchmark diff --git a/fastvideo/tests/distributed/test_ulysses_fault_injection.py b/fastvideo/tests/distributed/test_ulysses_fault_injection.py index 6af14b421..dbc7efc33 100644 --- a/fastvideo/tests/distributed/test_ulysses_fault_injection.py +++ b/fastvideo/tests/distributed/test_ulysses_fault_injection.py @@ -41,6 +41,7 @@ FAULT_CASES = [ def _worker() -> None: + import fastvideo.envs as envs from fastvideo.distributed import ( cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel, @@ -52,8 +53,8 @@ def _worker() -> None: world = int(os.environ["WORLD_SIZE"]) rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) - fault_rank = int(os.environ["FASTVIDEO_ULYSSES_FAULT_RANK"]) - fault_stage = os.environ["FASTVIDEO_ULYSSES_FAULT_STAGE"] + fault_rank = int(envs.FASTVIDEO_TEST_ULYSSES_FAULT_RANK.get()) + fault_stage = envs.FASTVIDEO_TEST_ULYSSES_FAULT_STAGE.get() torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") @@ -65,7 +66,7 @@ def _worker() -> None: UlyssesA2AHelper._can_attempt = ( # type: ignore[method-assign] lambda self: (False, "injected capability failure")) elif rank == fault_rank and fault_stage == "configuration": - os.environ["FASTVIDEO_ULYSSES_A2A"] = "off" + envs.FASTVIDEO_ULYSSES_A2A.set("off") maybe_init_distributed_environment_and_model_parallel(1, world) communicator = get_sp_group().device_communicator @@ -168,8 +169,8 @@ def test_rank_local_setup_failure_falls_back_group_wide(world: int, fault_rank: environment = dict( os.environ, FASTVIDEO_ULYSSES_A2A="auto", - FASTVIDEO_ULYSSES_FAULT_RANK=str(fault_rank), - FASTVIDEO_ULYSSES_FAULT_STAGE=fault_stage, + FASTVIDEO_TEST_ULYSSES_FAULT_RANK=str(fault_rank), + FASTVIDEO_TEST_ULYSSES_FAULT_STAGE=fault_stage, ) try: process = subprocess.run( diff --git a/fastvideo/tests/encoders/test_clip_encoder.py b/fastvideo/tests/encoders/test_clip_encoder.py index f97ba7ace..e988175d8 100644 --- a/fastvideo/tests/encoders/test_clip_encoder.py +++ b/fastvideo/tests/encoders/test_clip_encoder.py @@ -6,6 +6,7 @@ import pytest import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextModel import gc +import fastvideo.envs as envs from fastvideo.configs.pipelines import HunyuanConfig, PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.fastvideo_args import FastVideoArgs @@ -17,8 +18,8 @@ from torch.testing import assert_close logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py index 9597f89b7..568fce3c5 100644 --- a/fastvideo/tests/encoders/test_hyt5_encoder.py +++ b/fastvideo/tests/encoders/test_hyt5_encoder.py @@ -8,6 +8,7 @@ from torch.distributed.tensor import DTensor from torch.testing import assert_close from transformers import AutoConfig, AutoTokenizer, T5EncoderModel +import fastvideo.envs as envs from fastvideo.configs.pipelines import Hunyuan15T2V480PConfig, PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger @@ -18,8 +19,8 @@ from fastvideo.configs.models.encoders import T5Config logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") @pytest.fixture diff --git a/fastvideo/tests/encoders/test_llama_encoder.py b/fastvideo/tests/encoders/test_llama_encoder.py index 22a08a26e..6b940160f 100644 --- a/fastvideo/tests/encoders/test_llama_encoder.py +++ b/fastvideo/tests/encoders/test_llama_encoder.py @@ -6,6 +6,7 @@ import pytest import torch from transformers import AutoConfig, AutoTokenizer, LlamaModel import gc +import fastvideo.envs as envs from fastvideo.configs.pipelines import HunyuanConfig from fastvideo.forward_context import set_forward_context from fastvideo.fastvideo_args import FastVideoArgs @@ -18,8 +19,8 @@ from torch.testing import assert_close logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py index cd42434c7..4437db862 100644 --- a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_fp8.py @@ -2,13 +2,14 @@ from __future__ import annotations import json -import os import pytest import torch -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29514") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29514") import fastvideo.models.encoders.minimax_h3_checkpoint_fp8 as h3_fp8 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py index 74be1a9e8..a92bb7eb7 100644 --- a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py @@ -2,13 +2,14 @@ from __future__ import annotations import json -import os import pytest import torch -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29515") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29515") import fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 as h3_nvfp4 from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import ( diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py index b0c03c0e4..c0d10681e 100644 --- a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_truncation.py @@ -3,13 +3,14 @@ from __future__ import annotations import inspect -import os import pytest import torch -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29513") from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import ( MiniMaxH3Qwen3VLArchConfig, diff --git a/fastvideo/tests/encoders/test_qwen2_5_encoder.py b/fastvideo/tests/encoders/test_qwen2_5_encoder.py index 72b4cd0a5..b6a09bb5d 100644 --- a/fastvideo/tests/encoders/test_qwen2_5_encoder.py +++ b/fastvideo/tests/encoders/test_qwen2_5_encoder.py @@ -6,6 +6,7 @@ from torch.distributed.tensor import DTensor from torch.testing import assert_close from transformers import AutoConfig, AutoTokenizer, Qwen2_5_VLTextModel +import fastvideo.envs as envs from fastvideo.configs.pipelines import Hunyuan15T2V480PConfig, PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger @@ -16,8 +17,8 @@ from fastvideo.configs.models.encoders import Qwen2_5_VLConfig logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29505") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29505") @pytest.fixture diff --git a/fastvideo/tests/encoders/test_siglip_encoder.py b/fastvideo/tests/encoders/test_siglip_encoder.py index 9bd8d9033..3900e6292 100644 --- a/fastvideo/tests/encoders/test_siglip_encoder.py +++ b/fastvideo/tests/encoders/test_siglip_encoder.py @@ -3,7 +3,6 @@ import gc import json -import os import pytest import torch @@ -13,6 +12,7 @@ from torch.distributed.tensor import DTensor from torch.testing import assert_close from transformers import SiglipVisionModel as HFSiglipVisionModel +import fastvideo.envs as envs from fastvideo.configs.models.encoders import SiglipVisionConfig from fastvideo.configs.models.encoders.siglip import SiglipVisionArchConfig from fastvideo.forward_context import set_forward_context @@ -20,8 +20,8 @@ from fastvideo.logger import init_logger logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29505") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29505") # HYWorld model path - SigLIP image encoder MODEL_ID = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v" diff --git a/fastvideo/tests/encoders/test_t5_encoder.py b/fastvideo/tests/encoders/test_t5_encoder.py index 5183d5f48..01d40348c 100644 --- a/fastvideo/tests/encoders/test_t5_encoder.py +++ b/fastvideo/tests/encoders/test_t5_encoder.py @@ -8,6 +8,7 @@ from torch.distributed.tensor import DTensor from torch.testing import assert_close from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel, T5EncoderModel +import fastvideo.envs as envs from fastvideo.configs.pipelines import CosmosConfig, PipelineConfig, WanT2V480PConfig from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger @@ -19,8 +20,8 @@ from fastvideo.configs.models.encoders import T5Config, T5LargeConfig logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") @pytest.fixture diff --git a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py index 66288ebef..dc5349bea 100644 --- a/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py +++ b/fastvideo/tests/entrypoints/streaming/test_prompt_providers.py @@ -15,6 +15,7 @@ from typing import Any import pytest +import fastvideo.envs as envs from fastvideo.entrypoints.streaming.prompt.providers import ( CerebrasProvider, GroqProvider, @@ -98,19 +99,19 @@ class TestCerebrasProvider: def test_is_llm_provider(self): assert isinstance(CerebrasProvider(api_key="x"), LLMProvider) - def test_requires_api_key(self, monkeypatch): - monkeypatch.delenv("CEREBRAS_API_KEY", raising=False) + def test_requires_api_key(self, env_overrides): + env_overrides.enter_context(envs.override_external("CEREBRAS_API_KEY", None)) provider = CerebrasProvider() with pytest.raises(LLMProviderError, match="CEREBRAS_API_KEY"): asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) - def test_api_key_from_env(self, monkeypatch): - monkeypatch.setenv("CEREBRAS_API_KEY", "from-env") + def test_api_key_from_env(self, env_overrides): + env_overrides.enter_context(envs.override_external("CEREBRAS_API_KEY", "from-env")) provider = CerebrasProvider() assert provider.api_key == "from-env" - def test_explicit_api_key_wins(self, monkeypatch): - monkeypatch.setenv("CEREBRAS_API_KEY", "from-env") + def test_explicit_api_key_wins(self, env_overrides): + env_overrides.enter_context(envs.override_external("CEREBRAS_API_KEY", "from-env")) provider = CerebrasProvider(api_key="explicit") assert provider.api_key == "explicit" @@ -260,14 +261,14 @@ class TestGroqProvider: def test_is_llm_provider(self): assert isinstance(GroqProvider(api_key="x"), LLMProvider) - def test_requires_api_key(self, monkeypatch): - monkeypatch.delenv("GROQ_API_KEY", raising=False) + def test_requires_api_key(self, env_overrides): + env_overrides.enter_context(envs.override_external("GROQ_API_KEY", None)) provider = GroqProvider() with pytest.raises(LLMProviderError, match="GROQ_API_KEY"): asyncio.run(provider.complete(LLMRequest(messages=[], model="m"))) - def test_api_key_from_env(self, monkeypatch): - monkeypatch.setenv("GROQ_API_KEY", "from-env") + def test_api_key_from_env(self, env_overrides): + env_overrides.enter_context(envs.override_external("GROQ_API_KEY", "from-env")) provider = GroqProvider() assert provider.api_key == "from-env" diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index dd1f0185b..e446b4389 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -11,6 +11,7 @@ import torch from einops import rearrange import fastvideo.entrypoints.video_generator as video_generator_module +import fastvideo.envs as envs import fastvideo.utils as fastvideo_utils from fastvideo.api import ( ConfigValidationError, @@ -1025,14 +1026,14 @@ def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch): @pytest.mark.parametrize("audio_seconds", [1, 2, 3]) -def test_ffmpeg_pipe_preserves_video_duration(tmp_path, monkeypatch, audio_seconds): +def test_ffmpeg_pipe_preserves_video_duration(tmp_path, env_overrides, audio_seconds): ffmpeg = shutil.which("ffmpeg") ffprobe = shutil.which("ffprobe") if ffmpeg is None or ffprobe is None: pytest.skip("ffmpeg and ffprobe are required") - monkeypatch.setenv("FASTVIDEO_FFMPEG_BIN", ffmpeg) - monkeypatch.setenv("FASTVIDEO_VIDEO_CODEC", "libx264") - monkeypatch.setenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p") + env_overrides.enter_context(envs.FASTVIDEO_FFMPEG_BIN.override(ffmpeg)) + env_overrides.enter_context(envs.FASTVIDEO_VIDEO_CODEC.override("libx264")) + env_overrides.enter_context(envs.FASTVIDEO_OUTPUT_PIX_FMT.override("yuv420p")) fps, frame_count, sample_rate = 24, 48, 32000 frames = [np.full((64, 64, 3), index * 4, dtype=np.uint8) for index in range(frame_count)] audio = np.zeros((2, audio_seconds * sample_rate), dtype=np.float32) diff --git a/fastvideo/tests/golden_gate/_harness.py b/fastvideo/tests/golden_gate/_harness.py index 3bc1969b8..f2104c63d 100644 --- a/fastvideo/tests/golden_gate/_harness.py +++ b/fastvideo/tests/golden_gate/_harness.py @@ -24,17 +24,17 @@ import pytest import torch from torch.testing import assert_close -GOLDEN_REPO_ID = os.environ.get("FASTVIDEO_SSIM_REFERENCE_HF_REPO", "FastVideo/ssim-reference-videos") -GOLDEN_ROOT = Path(os.environ.get("FASTVIDEO_GOLDEN_GATE_DIR", Path(__file__).resolve().parent / "goldens")) +import fastvideo.envs as envs + DEFAULT_SEED = 20260805 # Must be set before cuBLAS initializes for use_deterministic_algorithms. -os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") -os.environ.setdefault("MASTER_ADDR", "127.0.0.1") -os.environ.setdefault("MASTER_PORT", "29661") -os.environ.setdefault("RANK", "0") -os.environ.setdefault("WORLD_SIZE", "1") -os.environ.setdefault("LOCAL_RANK", "0") +envs.setdefault_external("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +envs.setdefault_external("MASTER_ADDR", "127.0.0.1") +envs.setdefault_external("MASTER_PORT", "29661") +envs.setdefault_external("RANK", "0") +envs.setdefault_external("WORLD_SIZE", "1") +envs.setdefault_external("LOCAL_RANK", "0") @dataclass @@ -55,7 +55,7 @@ class GateSpec: dtype: torch.dtype = torch.bfloat16 # some blocks return tuples (e.g. double-stream img/txt) — reduce to one tensor postprocess: Callable[[Any], torch.Tensor] = field(default=lambda out: out) - model_root_env: str | None = None # env var naming a local checkpoint root + model_root: Callable[[], str | None] | None = None # returns a local checkpoint root, or None to download weight_file: str = "diffusion_pytorch_model.safetensors" # single-file fallback revision: str | None = None # Immutable checkpoint revision for pinned gates. @@ -74,7 +74,7 @@ def env_fingerprint(spec: GateSpec | str | None) -> dict[str, str]: "flash_attn": fa_version, "device_name": torch.cuda.get_device_name(0), "attention_backend": backend or "NONE", - "fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", ""), + "fastvideo_fa4": str(int(envs.FASTVIDEO_FA4.get())) if envs.FASTVIDEO_FA4.is_set() else "", "tf32_matmul": str(torch.backends.cuda.matmul.allow_tf32), "tf32_cudnn": str(torch.backends.cudnn.allow_tf32), } @@ -84,17 +84,18 @@ def env_fingerprint(spec: GateSpec | str | None) -> dict[str, str]: result.pop("flash_attn") result.pop("fastvideo_fa4") elif not isinstance(spec, GateSpec): - # New tensor gates record the effective switch. Unset and "0" both - # select FA2; retain the old block-gate metadata contract unchanged. - result["fastvideo_fa4"] = str(int(os.environ.get("FASTVIDEO_FA4", "0") != "0")) + # Tensor gates record the effective switch, so unset and "0" both record + # "0". Block gates record "" when FASTVIDEO_FA4 is unset, and "1" + # or "0" when it is set. + result["fastvideo_fa4"] = str(int(envs.FASTVIDEO_FA4.get())) return result def _component_dir(spec: GateSpec) -> Path: """Local checkout wins; otherwise fetch the index + only the shards holding this layer (or the single weight file for index-less checkpoints).""" - if spec.model_root_env: - root = os.environ.get(spec.model_root_env) + if spec.model_root is not None: + root = spec.model_root() if root: return Path(root) / spec.subfolder @@ -157,14 +158,22 @@ def _resolve_golden(spec: GateSpec) -> tuple[Path, bool]: return resolve_golden_path(_golden_relpath(spec)) +def _golden_repo_id() -> str: + return envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.get() + + +def _golden_root() -> Path: + return Path(envs.FASTVIDEO_TEST_GOLDEN_GATE_DIR.get() or Path(__file__).resolve().parent / "goldens") + + def resolve_golden_path(relative_path: str) -> tuple[Path, bool]: - local = GOLDEN_ROOT / relative_path + local = _golden_root() / relative_path if local.exists(): return local, True from huggingface_hub import hf_hub_download from huggingface_hub.errors import EntryNotFoundError try: - remote = hf_hub_download(GOLDEN_REPO_ID, f"golden_gates/{relative_path}", repo_type="dataset") + remote = hf_hub_download(_golden_repo_id(), f"golden_gates/{relative_path}", repo_type="dataset") return Path(remote), True except EntryNotFoundError: return local, False @@ -174,7 +183,12 @@ def run_gate(spec: GateSpec) -> None: if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): pytest.skip("golden-gate requires a bf16-capable CUDA GPU") - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = spec.attention_backend + with envs.FASTVIDEO_ATTENTION_BACKEND.override(spec.attention_backend): + _run_block_against_golden(spec) + + +def _run_block_against_golden(spec: GateSpec) -> None: + """Run the gated block on the fixed seeded batch and compare the output bitwise with the golden.""" torch.use_deterministic_algorithms(True) torch.backends.cudnn.benchmark = False @@ -209,7 +223,7 @@ def run_gate(spec: GateSpec) -> None: torch.save({"output": out, "metadata": meta}, golden_path) pytest.fail( f"golden latent seeded at {golden_path} ({meta}). Verify with a second run, " - f"then upload to {GOLDEN_REPO_ID} at golden_gates/{_golden_relpath(spec)}", + f"then upload to {_golden_repo_id()} at golden_gates/{_golden_relpath(spec)}", pytrace=False, ) diff --git a/fastvideo/tests/golden_gate/_tensor_golden.py b/fastvideo/tests/golden_gate/_tensor_golden.py index 4c9d04470..bd9c3594c 100644 --- a/fastvideo/tests/golden_gate/_tensor_golden.py +++ b/fastvideo/tests/golden_gate/_tensor_golden.py @@ -5,12 +5,12 @@ Keep weight loading and forward calls in each test. This helper only controls determinism and compares device/runtime-matched tensors; it is not a model runner. """ -import os from contextlib import contextmanager import pytest import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, device_folder, env_fingerprint, resolve_golden_path @@ -23,7 +23,6 @@ def deterministic_forward(attention_backend=None): deterministic = torch.are_deterministic_algorithms_enabled() warn_only = torch.is_deterministic_algorithms_warn_only_enabled() benchmark = torch.backends.cudnn.benchmark - previous_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") had_precision_state = hasattr(_mixed_precision_state, "state") previous_precision_state = getattr(_mixed_precision_state, "state", None) try: @@ -33,11 +32,12 @@ def deterministic_forward(attention_backend=None): # backend despite their unchanged torch default dtype and environment. if had_precision_state: del _mixed_precision_state.state - if attention_backend is not None: - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = attention_backend - torch.use_deterministic_algorithms(True) - torch.backends.cudnn.benchmark = False - yield torch.device("cuda:0") + # Restore the caller's FASTVIDEO_ATTENTION_BACKEND on exit, also when attention_backend is None. + with envs.FASTVIDEO_ATTENTION_BACKEND.override( + envs.FASTVIDEO_ATTENTION_BACKEND.get() if attention_backend is None else attention_backend): + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.benchmark = False + yield torch.device("cuda:0") finally: if had_precision_state: _mixed_precision_state.state = previous_precision_state @@ -45,10 +45,6 @@ def deterministic_forward(attention_backend=None): del _mixed_precision_state.state torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) torch.backends.cudnn.benchmark = benchmark - if previous_backend is None: - os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) - else: - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous_backend def assert_tensor_golden(name, outputs, *, identity, attention_backend=None, seed=DEFAULT_SEED): diff --git a/fastvideo/tests/golden_gate/test_flux.py b/fastvideo/tests/golden_gate/test_flux.py index 22babd762..19020f177 100644 --- a/fastvideo/tests/golden_gate/test_flux.py +++ b/fastvideo/tests/golden_gate/test_flux.py @@ -9,13 +9,14 @@ are float64 [S, 128]. temb is [B, 3072] (SD3AdaLayerNormZero 6-way chunk, sd3.py:394-395). Forward returns the TUPLE (encoder_hidden_states, hidden_states) (flux.py:329) — postprocess cats both streams along seq into one golden. The flux SSIM test pins TORCH_SDPA. black-forest-labs/FLUX.1-dev -is a GATED HF repo; FLUX_T2I_MODEL_DIR can point at a local checkout. +is a GATED HF repo; FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR can point at a local checkout. """ from __future__ import annotations import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -75,7 +76,9 @@ SPEC = GateSpec( attention_backend="TORCH_SDPA", # forward returns (encoder_hidden_states, hidden_states) — both [1, S, 3072] postprocess=lambda out: torch.cat(out, dim=1), - model_root_env="FLUX_T2I_MODEL_DIR", + # The registry default is the Hugging Face repo id, so only an explicitly set value names a local checkout. + model_root=lambda: (envs.FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR.get() + if envs.FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR.is_set() else None), ) diff --git a/fastvideo/tests/golden_gate/test_gamecraft.py b/fastvideo/tests/golden_gate/test_gamecraft.py index 6457071b7..d9c5f36ac 100644 --- a/fastvideo/tests/golden_gate/test_gamecraft.py +++ b/fastvideo/tests/golden_gate/test_gamecraft.py @@ -17,6 +17,7 @@ from __future__ import annotations import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -76,7 +77,9 @@ SPEC = GateSpec( renames=(), # checkpoint keys already match the block's state_dict 1:1 attention_backend="FLASH_ATTN", postprocess=lambda out: torch.cat(out, dim=1), # (img, txt) tuple -> one tensor - model_root_env="GAMECRAFT_MODEL_PATH", + # The registry default is the Hugging Face repo id, so only an explicitly set value names a local checkout. + model_root=lambda: (envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() + if envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.is_set() else None), ) diff --git a/fastvideo/tests/golden_gate/test_glm_image.py b/fastvideo/tests/golden_gate/test_glm_image.py index 48b567dc5..1ecc65ab1 100644 --- a/fastvideo/tests/golden_gate/test_glm_image.py +++ b/fastvideo/tests/golden_gate/test_glm_image.py @@ -10,13 +10,14 @@ returns a TUPLE (hidden_states, encoder_hidden_states) (glm_image.py:502) — postprocess cats encoder first along seq, matching the attention layout (glm_image.py:343). attn1.to_out is nn.ModuleList so "attn1.to_out.0.*" matches natively; only the two ff.net renames apply. The SSIM test pins -TORCH_SDPA and loads local weights via GLM_IMAGE_MODEL_DIR. +TORCH_SDPA and loads local weights via FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR. """ from __future__ import annotations import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -77,7 +78,7 @@ SPEC = GateSpec( # forward returns (hidden_states, encoder_hidden_states); encoder first, # like the attention concat layout (glm_image.py:343) postprocess=lambda out: torch.cat([out[1], out[0]], dim=1), - model_root_env="GLM_IMAGE_MODEL_DIR", + model_root=lambda: envs.FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR.get(), ) diff --git a/fastvideo/tests/golden_gate/test_kandinsky5.py b/fastvideo/tests/golden_gate/test_kandinsky5.py index 43b6d2fca..a40ee077a 100644 --- a/fastvideo/tests/golden_gate/test_kandinsky5.py +++ b/fastvideo/tests/golden_gate/test_kandinsky5.py @@ -18,10 +18,9 @@ FLASH_ATTN; mirrored here before any fastvideo import. from __future__ import annotations -import os - import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -97,11 +96,11 @@ SPEC = GateSpec( ) -def test_kandinsky5_golden_gate(distributed_runtime, monkeypatch) -> None: +def test_kandinsky5_golden_gate(distributed_runtime, env_overrides) -> None: # Pin FA4 off for THIS test only (test_kandinsky5_similarity.py:84). A # module-level os.environ.setdefault here leaked into every other # golden-gate test's env fingerprint at pytest collection time, breaking # full-directory runs in environments where FASTVIDEO_FA4 is unset # (observed on the self-hosted CI runner lane). - monkeypatch.setenv("FASTVIDEO_FA4", os.environ.get("FASTVIDEO_FA4", "0")) + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(envs.FASTVIDEO_FA4.get())) run_gate(SPEC) diff --git a/fastvideo/tests/golden_gate/test_longcat.py b/fastvideo/tests/golden_gate/test_longcat.py index 5ea910bde..004950aa6 100644 --- a/fastvideo/tests/golden_gate/test_longcat.py +++ b/fastvideo/tests/golden_gate/test_longcat.py @@ -16,8 +16,10 @@ plain tensor, so no postprocess. Backend FLASH_ATTN per the ssim test from __future__ import annotations + import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -69,7 +71,7 @@ SPEC = GateSpec( prefix_template="blocks.{N}.", renames=(), attention_backend="FLASH_ATTN", - model_root_env="LONGCAT_MODEL_ROOT", + model_root=lambda: envs.FASTVIDEO_TEST_LONGCAT_MODEL_ROOT.get(), weight_file="model.safetensors", # single 54.3 GB fp32 file, no index.json ) diff --git a/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py index 675b5a508..f7c35a57e 100644 --- a/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py +++ b/fastvideo/tests/golden_gate/test_minimax_h3_t2v.py @@ -21,7 +21,7 @@ never as unexplained bit drift. Goldens are device-keyed like SSIM references and live in the same HF dataset (``FastVideo/ssim-reference-videos`` under ``golden_gates//``); a -local ``goldens/`` directory or ``MINIMAX_H3_GATE_GOLDEN_DIR`` wins over the +local ``goldens/`` directory or ``FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR`` wins over the download. A missing golden is seeded locally and the test fails with upload instructions, mirroring the SSIM missing-reference convention. """ @@ -37,21 +37,19 @@ import pytest import torch from torch.testing import assert_close -GOLDEN_REPO_ID = os.environ.get("FASTVIDEO_SSIM_REFERENCE_HF_REPO", "FastVideo/ssim-reference-videos") +import fastvideo.envs as envs + MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3" -LAYER = int(os.environ.get("MINIMAX_H3_GATE_LAYER", "0")) SEED = 20260805 -# Pin the same attention path the H3 SSIM test runs, before any fastvideo -# import can cache a backend decision. CUBLAS workspace must be set before -# cuBLAS initializes for use_deterministic_algorithms to hold. -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") -os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") -os.environ.setdefault("MASTER_ADDR", "127.0.0.1") -os.environ.setdefault("MASTER_PORT", "29661") -os.environ.setdefault("RANK", "0") -os.environ.setdefault("WORLD_SIZE", "1") -os.environ.setdefault("LOCAL_RANK", "0") +# CUBLAS workspace must be set before cuBLAS initializes for +# use_deterministic_algorithms to hold. +envs.setdefault_external("CUBLAS_WORKSPACE_CONFIG", ":4096:8") +envs.setdefault_external("MASTER_ADDR", "127.0.0.1") +envs.setdefault_external("MASTER_PORT", "29661") +envs.setdefault_external("RANK", "0") +envs.setdefault_external("WORLD_SIZE", "1") +envs.setdefault_external("LOCAL_RANK", "0") # Checkpoint -> FastVideo renames, the block-level subset of # MiniMaxH3Config.param_names_mapping (fastvideo/configs/models/dits/minimax_h3.py). @@ -62,6 +60,14 @@ _RENAMES = ( ) +def _golden_repo_id() -> str: + return envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.get() + + +def _layer() -> int: + return envs.FASTVIDEO_TEST_MINIMAX_H3_GATE_LAYER.get() + + @pytest.fixture(scope="module", autouse=True) def _distributed_runtime(): if not torch.cuda.is_available(): @@ -91,8 +97,8 @@ def _env_fingerprint() -> dict[str, str]: "cudnn": str(torch.backends.cudnn.version()), "flash_attn": fa_version, "device_name": torch.cuda.get_device_name(0), - "attention_backend": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", ""), - "fastvideo_fa4": os.environ.get("FASTVIDEO_FA4", ""), + "attention_backend": envs.FASTVIDEO_ATTENTION_BACKEND.get() or "", + "fastvideo_fa4": str(int(envs.FASTVIDEO_FA4.get())) if envs.FASTVIDEO_FA4.is_set() else "", "tf32_matmul": str(torch.backends.cuda.matmul.allow_tf32), "tf32_cudnn": str(torch.backends.cudnn.allow_tf32), } @@ -101,18 +107,18 @@ def _env_fingerprint() -> dict[str, str]: def _component_dir() -> Path: """Local checkout wins; otherwise fetch only the index + the shards that hold this layer's tensors (~4.5 GB of the 62 GB checkpoint).""" - root = os.environ.get("MINIMAX_H3_MODEL_ROOT") + root = envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() if root: component = Path(root) / "transformer" if not (component / "diffusion_pytorch_model.safetensors.index.json").is_file(): - pytest.fail(f"MINIMAX_H3_MODEL_ROOT set but index missing under {component}", pytrace=False) + pytest.fail(f"FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT set but index missing under {component}", pytrace=False) return component from huggingface_hub import hf_hub_download index_path = Path(hf_hub_download(MODEL_REPO_ID, "transformer/diffusion_pytorch_model.safetensors.index.json")) weight_map = json.loads(index_path.read_text())["weight_map"] - prefix = f"transformer_blocks.{LAYER}." + prefix = f"transformer_blocks.{_layer()}." shards = sorted({shard for name, shard in weight_map.items() if name.startswith(prefix)}) if not shards: pytest.fail(f"no shards found for prefix {prefix} in {MODEL_REPO_ID}", pytrace=False) @@ -207,13 +213,14 @@ def _device_slug() -> str: def _golden_filename() -> str: - backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "default") - return f"minimax_h3_t2v_layer{LAYER}_{backend}_seed{SEED}.pt" + backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() or "default" + return f"minimax_h3_t2v_layer{_layer()}_{backend}_seed{SEED}.pt" def _resolve_golden() -> tuple[Path, bool]: """Return (path, exists). Local dir wins; falls back to the HF dataset.""" - local_root = Path(os.environ.get("MINIMAX_H3_GATE_GOLDEN_DIR", Path(__file__).resolve().parent / "goldens")) + local_root = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_GATE_GOLDEN_DIR.get() + or Path(__file__).resolve().parent / "goldens") local = local_root / _device_slug() / _golden_filename() if local.exists(): return local, True @@ -222,7 +229,7 @@ def _resolve_golden() -> tuple[Path, bool]: from huggingface_hub.errors import EntryNotFoundError try: remote = hf_hub_download( - GOLDEN_REPO_ID, + _golden_repo_id(), f"golden_gates/{_device_slug()}/{_golden_filename()}", repo_type="dataset", ) @@ -231,10 +238,14 @@ def _resolve_golden() -> tuple[Path, bool]: return local, False -def test_minimax_h3_t2v_golden_gate() -> None: +def test_minimax_h3_t2v_golden_gate(env_overrides) -> None: if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): pytest.skip("golden-gate requires a bf16-capable CUDA GPU") + # Pin the same attention path the H3 SSIM test runs unless the caller chose a backend. + if not envs.FASTVIDEO_ATTENTION_BACKEND.is_set(): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) + # Tripwire: any nondeterministic op added to the block errors immediately # instead of flaking the gate. benchmark=False kills cuDNN autotune # variance (no-op today — the block has no convs — but the gate may grow). @@ -246,8 +257,8 @@ def test_minimax_h3_t2v_golden_gate() -> None: from fastvideo.forward_context import set_forward_context - block = _build_block(LAYER) - state = _load_layer_state(component_dir, LAYER) + block = _build_block(_layer()) + state = _load_layer_state(component_dir, _layer()) missing, unexpected = block.load_state_dict(state, strict=True) assert not missing and not unexpected block = block.to(device=device, dtype=torch.bfloat16).eval() @@ -259,7 +270,7 @@ def test_minimax_h3_t2v_golden_gate() -> None: golden_path, golden_exists = _resolve_golden() meta = { - "layer": LAYER, + "layer": _layer(), "seed": SEED, "shape": list(out.shape), "env": _env_fingerprint(), @@ -269,7 +280,7 @@ def test_minimax_h3_t2v_golden_gate() -> None: torch.save({"output": out, "metadata": meta}, golden_path) pytest.fail( f"golden latent seeded at {golden_path} ({meta}). Verify with a second " - f"run, then upload to {GOLDEN_REPO_ID} at " + f"run, then upload to {_golden_repo_id()} at " f"golden_gates/{_device_slug()}/{_golden_filename()}", pytrace=False, ) diff --git a/fastvideo/tests/golden_gate/test_wan_denoising.py b/fastvideo/tests/golden_gate/test_wan_denoising.py index 45d8ee710..de332ff70 100644 --- a/fastvideo/tests/golden_gate/test_wan_denoising.py +++ b/fastvideo/tests/golden_gate/test_wan_denoising.py @@ -7,6 +7,7 @@ per-step tensor reference catches stage/scheduler drift after the block gate. import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import DEFAULT_SEED, distributed_runtime from fastvideo.tests.golden_gate._tensor_golden import assert_tensor_golden, deterministic_forward from fastvideo.tests.golden_gate._wan_checkpoint import checkpoint_identity, component_path @@ -53,8 +54,8 @@ def denoising_outputs(device): } -def test_wan_denoising_golden_gate(distributed_runtime, monkeypatch): - monkeypatch.setenv("FASTVIDEO_CFG_GATE_STEP", "1.0") +def test_wan_denoising_golden_gate(distributed_runtime, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_CFG_GATE_STEP.override(1.0)) with deterministic_forward("FLASH_ATTN") as device: outputs, identity = denoising_outputs(device) assert_tensor_golden("wan_denoising", outputs, identity=identity, attention_backend="FLASH_ATTN") diff --git a/fastvideo/tests/golden_gate/test_zimage.py b/fastvideo/tests/golden_gate/test_zimage.py index cd27b6774..9271a5b83 100644 --- a/fastvideo/tests/golden_gate/test_zimage.py +++ b/fastvideo/tests/golden_gate/test_zimage.py @@ -17,6 +17,7 @@ from __future__ import annotations import torch +import fastvideo.envs as envs from fastvideo.tests.golden_gate._harness import GateSpec, distributed_runtime, run_gate __all__ = ["distributed_runtime"] @@ -78,7 +79,9 @@ SPEC = GateSpec( prefix_template="layers.{N}.", renames=(), attention_backend="TORCH_SDPA", - model_root_env="ZIMAGE_MODEL_DIR", + # The registry default is the Hugging Face repo id, so only an explicitly set value names a local checkout. + model_root=lambda: (envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() + if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else None), ) diff --git a/fastvideo/tests/hooks/test_activation_trace.py b/fastvideo/tests/hooks/test_activation_trace.py index e5da540ef..cce426b03 100644 --- a/fastvideo/tests/hooks/test_activation_trace.py +++ b/fastvideo/tests/hooks/test_activation_trace.py @@ -6,6 +6,7 @@ from pathlib import Path import torch from torch import nn +import fastvideo.envs as envs from fastvideo.hooks.activation_trace import ( attach_activation_trace, detach_activation_trace, @@ -50,8 +51,8 @@ def _read_jsonl(path: Path) -> list[dict]: return [json.loads(line) for line in path.read_text().splitlines()] -def test_attach_activation_trace_off_returns_none(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_TRACE_ACTIVATIONS", raising=False) +def test_attach_activation_trace_off_returns_none(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(None)) model = ToyModel() manager = attach_activation_trace(model) @@ -62,12 +63,12 @@ def test_attach_activation_trace_off_returns_none(monkeypatch) -> None: def test_attach_activation_trace_on_respects_layer_filter( - monkeypatch, + env_overrides, tmp_path, ) -> None: - monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") - monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0.*") - monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_LAYERS.override(r"block\.0.*")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_OUTPUT.override(str(tmp_path / "trace.jsonl"))) model = ToyModel() manager = attach_activation_trace(model) @@ -81,12 +82,12 @@ def test_attach_activation_trace_on_respects_layer_filter( detach_activation_trace(manager) -def test_activation_trace_writes_configured_stats(monkeypatch, tmp_path) -> None: +def test_activation_trace_writes_configured_stats(env_overrides, tmp_path) -> None: path = tmp_path / "trace.jsonl" - monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") - monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") - monkeypatch.setenv("FASTVIDEO_TRACE_STATS", "abs_mean,sum,shape,dtype") - monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_LAYERS.override(r"block\.0")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_STATS.override("abs_mean,sum,shape,dtype")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_OUTPUT.override(str(path))) model = ToyModel() manager = attach_activation_trace(model) @@ -107,12 +108,12 @@ def test_activation_trace_writes_configured_stats(monkeypatch, tmp_path) -> None assert record["dtype"] == "torch.float32" -def test_activation_trace_step_filter(monkeypatch, tmp_path) -> None: +def test_activation_trace_step_filter(env_overrides, tmp_path) -> None: path = tmp_path / "trace.jsonl" - monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") - monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") - monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) - monkeypatch.setenv("FASTVIDEO_TRACE_STEPS", "0,2") + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_LAYERS.override(r"block\.0")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_OUTPUT.override(str(path))) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_STEPS.override("0,2")) model = ToyModel() manager = attach_activation_trace(model) @@ -126,11 +127,11 @@ def test_activation_trace_step_filter(monkeypatch, tmp_path) -> None: assert [record["step"] for record in _read_jsonl(path)] == [0, 2] -def test_activation_trace_flattens_tuple_outputs(monkeypatch, tmp_path) -> None: +def test_activation_trace_flattens_tuple_outputs(env_overrides, tmp_path) -> None: path = tmp_path / "trace.jsonl" - monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") - monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", "tuple$") - monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(path)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_LAYERS.override("tuple$")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_OUTPUT.override(str(path))) model = TupleOutputModel() manager = attach_activation_trace(model) @@ -143,10 +144,10 @@ def test_activation_trace_flattens_tuple_outputs(monkeypatch, tmp_path) -> None: assert [record["tensor"] for record in records] == ["out[0]", "out[1]"] -def test_detach_activation_trace_removes_hooks(monkeypatch, tmp_path) -> None: - monkeypatch.setenv("FASTVIDEO_TRACE_ACTIVATIONS", "1") - monkeypatch.setenv("FASTVIDEO_TRACE_LAYERS", r"block\.0") - monkeypatch.setenv("FASTVIDEO_TRACE_OUTPUT", str(tmp_path / "trace.jsonl")) +def test_detach_activation_trace_removes_hooks(env_overrides, tmp_path) -> None: + env_overrides.enter_context(envs.FASTVIDEO_TRACE_ACTIVATIONS.override(True)) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_LAYERS.override(r"block\.0")) + env_overrides.enter_context(envs.FASTVIDEO_TRACE_OUTPUT.override(str(tmp_path / "trace.jsonl"))) model = ToyModel() manager = attach_activation_trace(model) diff --git a/fastvideo/tests/inference/bsa/test_bsa_inference.py b/fastvideo/tests/inference/bsa/test_bsa_inference.py index 199b50f45..fd5795ebb 100644 --- a/fastvideo/tests/inference/bsa/test_bsa_inference.py +++ b/fastvideo/tests/inference/bsa/test_bsa_inference.py @@ -7,13 +7,15 @@ import sys import tempfile from pathlib import Path +import fastvideo.envs as envs -def test_inference_bsa(): + +def test_inference_bsa(env_overrides): """Test FastVideo BSA_ATTN inference pipeline""" output_dir = Path("outputs_video/bsa_1.3B/") - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "BSA_ATTN" + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("BSA_ATTN")) config = { "generator": { diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index 97e642bf2..4e2ec6800 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -5,6 +5,7 @@ import re import pytest +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.logger import init_logger from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results @@ -18,8 +19,8 @@ from fastvideo.worker import MultiprocExecutor import torch logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29500") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29500") # Base parameters for LoRA inference tests WAN_LORA_PARAMS = { @@ -132,12 +133,12 @@ def test_merge_lora_weights(model_id): @pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_lora_inference_similarity(ATTENTION_BACKEND, model_id): +def test_lora_inference_similarity(ATTENTION_BACKEND, model_id, env_overrides): """ Test that runs LoRA inference with LoRA switching and compares the output to reference videos using SSIM. """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py index e3aa224e0..b1b9c825c 100644 --- a/fastvideo/tests/inference/test_inference_regional_compile.py +++ b/fastvideo/tests/inference/test_inference_regional_compile.py @@ -28,6 +28,7 @@ import pytest import torch from torch import nn +import fastvideo.envs as envs from fastvideo.attention import layer as attention_layer from fastvideo.attention.layer import DistributedAttention from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention, MiniMaxH3Transformer3DModel @@ -45,8 +46,8 @@ def _init_params_for(backend_name: str | None) -> dict: return {"config": SimpleNamespace(_resolved_attention_backend=resolved)} -def test_legacy_vsa_backend_degrades_to_eager(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) +def test_legacy_vsa_backend_degrades_to_eager(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) backend_name = "VIDEO_SPARSE_ATTN" reason = _regional_compile_unsupported_reason(_init_params_for(backend_name)) assert reason is not None @@ -55,16 +56,16 @@ def test_legacy_vsa_backend_degrades_to_eager(monkeypatch) -> None: @pytest.mark.parametrize("backend_name", [None, "TORCH_SDPA"]) -def test_supported_backends_allow_compile(backend_name, monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) - monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) +def test_supported_backends_allow_compile(backend_name, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override(None)) assert _regional_compile_unsupported_reason(_init_params_for(backend_name)) is None -def test_h3_vsa_sm100a_tile64_allows_compile(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) - monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) - monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") +def test_h3_vsa_sm100a_tile64_allows_compile(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) reason = _regional_compile_unsupported_reason( _init_params_for("VIDEO_SPARSE_ATTN_H3"), @@ -75,13 +76,13 @@ def test_h3_vsa_sm100a_tile64_allows_compile(monkeypatch) -> None: @pytest.mark.parametrize(("sm100a", "tile_size"), [(False, 64), (True, 256), (True, None)]) -def test_h3_vsa_unsupported_compile_route_degrades_to_eager(sm100a, tile_size, monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) - monkeypatch.delenv("FASTVIDEO_H3_VSA_PROBE", raising=False) +def test_h3_vsa_unsupported_compile_route_degrades_to_eager(sm100a, tile_size, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override(None)) if sm100a: - monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(True)) else: - monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) reason = _regional_compile_unsupported_reason( _init_params_for("VIDEO_SPARSE_ATTN_H3"), @@ -92,9 +93,9 @@ def test_h3_vsa_unsupported_compile_route_degrades_to_eager(sm100a, tile_size, m assert "eager" in reason -def test_h3_vsa_probe_degrades_regional_compile_to_eager(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) - monkeypatch.setenv("FASTVIDEO_H3_VSA_PROBE", "/tmp/h3-vsa-probe") +def test_h3_vsa_probe_degrades_regional_compile_to_eager(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_PROBE.override("/tmp/h3-vsa-probe")) reason = _regional_compile_unsupported_reason(_init_params_for("VIDEO_SPARSE_ATTN_H3")) @@ -251,16 +252,16 @@ def test_prepared_zero_vsa_gate_still_runs_in_grad_enabled_training() -> None: assert attention._gate_compress_active is False -def test_attention_compile_escape_hatch_degrades_to_eager(monkeypatch) -> None: - monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "1") +def test_attention_compile_escape_hatch_degrades_to_eager(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(True)) reason = _regional_compile_unsupported_reason(_init_params_for("TORCH_SDPA")) assert reason is not None assert "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in reason @pytest.mark.parametrize("fa_version", ["2", "3", "4"]) -def test_dense_flash_attention_inference_allows_compile(fa_version, monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) +def test_dense_flash_attention_inference_allows_compile(fa_version, monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default") fake_module.fa_version = fa_version monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) @@ -268,8 +269,8 @@ def test_dense_flash_attention_inference_allows_compile(fa_version, monkeypatch) assert _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN")) is None -def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None: - monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) +def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(None)) disabled_calls: list[str] = [] def _fake_disable(fn): @@ -297,7 +298,7 @@ def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None # Preserve the existing process-wide escape hatch for callers that opt in # before constructing their attention modules. - monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "0") + env_overrides.enter_context(envs.FASTVIDEO_DISABLE_ATTENTION_COMPILE.override(False)) explicit = _ToyAttention() assert explicit.forward() == "forward" assert disabled_calls == ["disabled"] diff --git a/fastvideo/tests/inference/vmoba/test_vmoba_inference.py b/fastvideo/tests/inference/vmoba/test_vmoba_inference.py index 333a3a9f8..c543a55ff 100644 --- a/fastvideo/tests/inference/vmoba/test_vmoba_inference.py +++ b/fastvideo/tests/inference/vmoba/test_vmoba_inference.py @@ -6,14 +6,16 @@ import subprocess import tempfile from pathlib import Path +import fastvideo.envs as envs -def test_inference_vmoba(): + +def test_inference_vmoba(env_overrides): """Test FastVideo VMOBA_ATTN inference pipeline""" output_dir = Path("outputs_video/vmoba_1.3B/") moba_config = "fastvideo/configs/backend/vmoba/wan_1.3B_77_480_832.json" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VMOBA_ATTN" + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VMOBA_ATTN")) config = { "generator": { diff --git a/fastvideo/tests/mlx/test_memory_limits.py b/fastvideo/tests/mlx/test_memory_limits.py index 0f6a02935..ec91fcc3a 100644 --- a/fastvideo/tests/mlx/test_memory_limits.py +++ b/fastvideo/tests/mlx/test_memory_limits.py @@ -4,6 +4,7 @@ import os import pytest +import fastvideo.envs as envs from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, gib_to_bytes @@ -48,9 +49,9 @@ def test_gib_to_bytes_rejects_non_positive_values() -> None: gib_to_bytes(0) -def test_apply_memory_limits_sets_mlx_limits_and_metrics(monkeypatch) -> None: - monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False) - monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False) +def test_apply_memory_limits_sets_mlx_limits_and_metrics(env_overrides) -> None: + env_overrides.enter_context(envs.override_external("PYTORCH_MPS_HIGH_WATERMARK_RATIO", None)) + env_overrides.enter_context(envs.override_external("PYTORCH_MPS_LOW_WATERMARK_RATIO", None)) fake_mlx = _FakeMLX() applied = apply_memory_limits( diff --git a/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py b/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py index 79df58b16..9c8d3813a 100644 --- a/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py +++ b/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py @@ -3,13 +3,12 @@ from __future__ import annotations -import os - import numpy as np import pytest mx = pytest.importorskip("mlx.core", reason="MLX is required for affine dq-GEMM tests") +import fastvideo.envs as envs # noqa: E402 from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 MLXQuantizationSpec, affine_dq_gemm_min_m, @@ -57,8 +56,8 @@ def _rel_l2(a, b) -> float: @pytest.mark.parametrize("bits", AFFINE_BITS) @pytest.mark.parametrize("group_size", GROUP_SIZES) -def test_dq_gemm_matches_qmm_for_supported_bit_widths(bits: int, group_size: int, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "8") +def test_dq_gemm_matches_qmm_for_supported_bit_widths(bits: int, group_size: int, env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("8")) reset_dq_gemm_telemetry() in_features = group_size * 4 out_features = group_size * 2 @@ -78,8 +77,8 @@ def test_dq_gemm_matches_qmm_for_supported_bit_widths(bits: int, group_size: int assert float(np.max(np.abs(got_np - ref_np))) / scale < 0.08 -def test_dq_gemm_with_bias_and_batched_rows(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "4") +def test_dq_gemm_with_bias_and_batched_rows(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("4")) reset_dq_gemm_telemetry() quantized = _try_quantize(64, 128, bits=6, group_size=64) x = mx.random.normal((2, 8, 128)).astype(mx.bfloat16) @@ -94,8 +93,8 @@ def test_dq_gemm_with_bias_and_batched_rows(monkeypatch: pytest.MonkeyPatch) -> assert got.shape == (2, 8, 64) -def test_dq_gemm_stays_on_qmm_below_threshold(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "768") +def test_dq_gemm_stays_on_qmm_below_threshold(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("768")) reset_dq_gemm_telemetry() quantized = _try_quantize(64, 128, bits=6, group_size=64) x = mx.random.normal((32, 128)).astype(mx.bfloat16) @@ -107,8 +106,8 @@ def test_dq_gemm_stays_on_qmm_below_threshold(monkeypatch: pytest.MonkeyPatch) - np.testing.assert_array_equal(np.asarray(got.astype(mx.float32)), np.asarray(ref.astype(mx.float32))) -def test_dq_gemm_env_zero_disables_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "0") +def test_dq_gemm_env_zero_disables_dispatch(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("0")) assert affine_dq_gemm_min_m() is None reset_dq_gemm_telemetry() quantized = _try_quantize(64, 128, bits=6, group_size=64) @@ -121,8 +120,8 @@ def test_dq_gemm_env_zero_disables_dispatch(monkeypatch: pytest.MonkeyPatch) -> np.testing.assert_array_equal(np.asarray(got.astype(mx.float32)), np.asarray(ref.astype(mx.float32))) -def test_shared_linear_stays_on_qmm_at_wide_m(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "1") +def test_shared_linear_stays_on_qmm_at_wide_m(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("1")) reset_dq_gemm_telemetry() quantized = _try_quantize(64, 128, bits=6, group_size=64) x = mx.random.normal((1024, 128)).astype(mx.bfloat16) @@ -133,8 +132,8 @@ def test_shared_linear_stays_on_qmm_at_wide_m(monkeypatch: pytest.MonkeyPatch) - np.testing.assert_array_equal(np.asarray(got.astype(mx.float32)), np.asarray(ref.astype(mx.float32))) -def test_non_affine_weights_stay_on_quantized_matmul(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "1") +def test_non_affine_weights_stay_on_quantized_matmul(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("1")) reset_dq_gemm_telemetry() spec = MLXQuantizationSpec(mode="mxfp8") weight = mx.random.normal((64, 64)).astype(mx.bfloat16) @@ -151,15 +150,13 @@ def test_non_affine_weights_stay_on_quantized_matmul(monkeypatch: pytest.MonkeyP assert got.shape == (1024, 64) -def test_default_floor_is_measured_768(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "1") +def test_default_floor_is_measured_768(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("1")) assert affine_dq_gemm_min_m() == 768 - monkeypatch.setenv("FASTVIDEO_MLX_DQ_GEMM", "2048") + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override("2048")) assert affine_dq_gemm_min_m() == 2048 - monkeypatch.delenv("FASTVIDEO_MLX_DQ_GEMM", raising=False) - os.environ.pop("FASTVIDEO_MLX_DQ_GEMM", None) + env_overrides.enter_context(envs.FASTVIDEO_MLX_DQ_GEMM.override(None)) # Default with unset env is on at the measured floor. - monkeypatch.delenv("FASTVIDEO_MLX_DQ_GEMM", raising=False) - if "FASTVIDEO_MLX_DQ_GEMM" in os.environ: + if envs.FASTVIDEO_MLX_DQ_GEMM.is_set(): pytest.skip("parent environment pinned FASTVIDEO_MLX_DQ_GEMM") assert affine_dq_gemm_min_m() == 768 diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py index e2a4e3778..ca3fe3d11 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py @@ -15,6 +15,7 @@ import pytest mx = pytest.importorskip("mlx.core") +import fastvideo.envs as envs from fastvideo.mlx_runtime import minimax_h3 as h3 from fastvideo.mlx_runtime import minimax_h3_vsa as vsa from fastvideo.mlx_runtime import minimax_h3_vsa_simd as simd @@ -121,12 +122,12 @@ def test_configuration_is_transactional_and_invalidates_geometry(): @pytest.fixture -def tiny_h3_model(monkeypatch, request): +def tiny_h3_model(env_overrides, request): from fastvideo.tests.mlx.tiny_h3 import build_hf_config, build_tiny_h3_config, build_torch_model from fastvideo.tests.mlx.tiny_h3 import mlx_dit_from_torch_model - monkeypatch.setenv("MASTER_ADDR", "localhost") - monkeypatch.setenv("MASTER_PORT", "29513") + env_overrides.enter_context(envs.override_external("MASTER_ADDR", "localhost")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", "29513")) request.getfixturevalue("distributed_setup") return mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_h3_config())) diff --git a/fastvideo/tests/mlx/test_mlx_taeh3.py b/fastvideo/tests/mlx/test_mlx_taeh3.py index 788379997..898a1021f 100644 --- a/fastvideo/tests/mlx/test_mlx_taeh3.py +++ b/fastvideo/tests/mlx/test_mlx_taeh3.py @@ -1,26 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 """TAEH3 port parity. Requires Apple MLX and an explicit upstream checkout. -TAEH3_REFERENCE_DIR=/path/to/taehv python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py +FASTVIDEO_TEST_TAEH3_REFERENCE_DIR=/path/to/taehv python -m pytest fastvideo/tests/mlx/test_mlx_taeh3.py No weights or source code are downloaded by this test. """ import importlib.util -import os from pathlib import Path import numpy as np import pytest mx = pytest.importorskip('mlx.core') +import fastvideo.envs as envs from fastvideo.mlx_runtime.minimax_h3_taeh3 import MLXTAEH3Decoder, ensure_taeh3_checkpoint @pytest.fixture(scope='module') def models(): torch = pytest.importorskip('torch') - location = os.environ.get('TAEH3_REFERENCE_DIR') + location = envs.FASTVIDEO_TEST_TAEH3_REFERENCE_DIR.get() if not location: - pytest.skip('Set TAEH3_REFERENCE_DIR to an upstream taehv checkout') + pytest.skip('Set FASTVIDEO_TEST_TAEH3_REFERENCE_DIR to an upstream taehv checkout') root = Path(location) spec = importlib.util.spec_from_file_location('taeh3_reference', root / 'taehv.py') module = importlib.util.module_from_spec(spec) diff --git a/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py b/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py index 1a2e5495f..e25b9ed0a 100644 --- a/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py +++ b/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py @@ -12,6 +12,7 @@ from pathlib import Path import numpy as np import pytest +import fastvideo.envs as envs from fastvideo.mlx_runtime.prompt_cache import load_prompt_cache, save_prompt_cache @@ -268,6 +269,7 @@ def test_wan22_default_cache_path_tracks_encoder_files(tmp_path: Path) -> None: ) def test_wan22_main_resolves_prompt_cache( monkeypatch, + env_overrides, tmp_path: Path, cache_args: list[str], expected: str | None, @@ -293,7 +295,7 @@ def test_wan22_main_resolves_prompt_cache( raise CacheProbe resolved_args = [str(explicit) if arg == "explicit.npy" else arg for arg in cache_args] - monkeypatch.setenv("HOME", str(tmp_path)) + env_overrides.enter_context(envs.override_external("HOME", str(tmp_path))) monkeypatch.setattr(module, "load_prompt_cache", probe) monkeypatch.setattr( sys, diff --git a/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py index 027baf204..07c3724a6 100644 --- a/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py +++ b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py @@ -11,26 +11,27 @@ t≈900), and asserts allclose. Gated on Metal + local weights so Linux from __future__ import annotations import json -import os from pathlib import Path import numpy as np import pytest +import fastvideo.envs as envs + # The torch reference side brings up FastVideo's single-process distributed # environment, which rendezvouses over env://. Match the other torch-backed # tests in this repo and provide a local rendezvous before that happens. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29513") mx = pytest.importorskip("mlx.core", reason="MLX is required for the Wan2.2 real-weight parity test") -_ROOT = Path(os.environ.get("FASTVIDEO_WAN22_5B_ROOT", str(Path.home() / "models" / "fastwan22_5b"))) -_CHECKPOINT = _ROOT / "transformer" / "diffusion_pytorch_model.safetensors" -_CONFIG = _ROOT / "transformer" / "config.json" - _HAS_METAL = bool(getattr(mx, "metal", None) and mx.metal.is_available()) -_HAS_WEIGHTS = _CHECKPOINT.exists() and _CONFIG.exists() and _CHECKPOINT.stat().st_size > 1_000_000_000 + + +def _wan22_root() -> Path: + return Path(envs.FASTVIDEO_TEST_WAN22_5B_ROOT.get() or str(Path.home() / "models" / "fastwan22_5b")) + # The real-weight test loads ~10 GB fp16 weights + activations. Gate on the # documented 36 GB requirement unless explicitly opted in. @@ -41,7 +42,7 @@ def _check_memory_requirement() -> bool: bool: `true` if the system has at least 36 GiB of memory or low-memory execution is explicitly enabled, `false` otherwise. """ - if os.environ.get("FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY") == "1": + if envs.FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY.get(): return True try: # On macOS, check unified memory via sysctl. @@ -63,20 +64,8 @@ def _check_memory_requirement() -> bool: return False -_HAS_SUFFICIENT_MEMORY = _check_memory_requirement() - pytestmark = [ pytest.mark.skipif(not _HAS_METAL, reason="Metal required for real 5B fp16 parity"), - pytest.mark.skipif( - not _HAS_WEIGHTS, - reason=f"Wan2.2-5B checkpoint not found/incomplete under {_ROOT} " - "(set FASTVIDEO_WAN22_5B_ROOT; expected ~10 GB safetensors)", - ), - pytest.mark.skipif( - not _HAS_SUFFICIENT_MEMORY, - reason="Real-weight 5B test requires >= 36 GB unified memory " - "(set FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY=1 to override)", - ), ] @@ -151,6 +140,15 @@ def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: """ Verify that the MLX and PyTorch Wan2.2 implementations produce equivalent outputs for per-token timesteps using real checkpoint weights. """ + root = _wan22_root() + checkpoint = root / "transformer" / "diffusion_pytorch_model.safetensors" + config_path = root / "transformer" / "config.json" + if not (checkpoint.exists() and config_path.exists() and checkpoint.stat().st_size > 1_000_000_000): + pytest.skip(f"Wan2.2-5B checkpoint not found/incomplete under {root} " + "(set FASTVIDEO_TEST_WAN22_5B_ROOT; expected ~10 GB safetensors)") + if not _check_memory_requirement(): + pytest.skip("Real-weight 5B test requires >= 36 GB unified memory " + "(set FASTVIDEO_TEST_WAN22_5B_ALLOW_LOW_MEMORY=1 to override)") import torch from fastvideo.forward_context import set_forward_context @@ -158,7 +156,7 @@ def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from examples.inference.basic.mlx_wan_prompt_to_video import make_rotary_embeddings - config = json.loads(_CONFIG.read_text()) + config = json.loads(config_path.read_text()) in_ch = int(config["in_channels"]) text_dim = int(config["text_dim"]) # Small latent so activations fit alongside the ~10 GB fp16 weights on 36 GB. @@ -177,7 +175,7 @@ def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: text_np = (rng.standard_normal((1, 32, text_dim)) * 0.1).astype(np.float32) # --- torch reference (CPU fp16, same weight dtype as MLX deploy path) --- - torch_model = _load_torch_wan22_from_diffusers(_CHECKPOINT, _CONFIG, dtype=torch.float16) + torch_model = _load_torch_wan22_from_diffusers(checkpoint, config_path, dtype=torch.float16) hidden_t = torch.from_numpy(hidden_np).to(torch.float16) text_t = torch.from_numpy(text_np).to(torch.float16) with torch.no_grad(), set_forward_context( @@ -191,7 +189,7 @@ def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: gc.collect() # --- MLX (fp16 weights, fp16 compute) --- - mlx_model = mlx_wan22_dit_from_diffusers_safetensors(_CHECKPOINT, _CONFIG, dtype="fp16") + mlx_model = mlx_wan22_dit_from_diffusers_safetensors(checkpoint, config_path, dtype="fp16") freqs_cis = make_rotary_embeddings( config, latent_frames=frames, latent_height=height, latent_width=width) out = mlx_model( diff --git a/fastvideo/tests/mlx/tiny_h3.py b/fastvideo/tests/mlx/tiny_h3.py index afdf71740..556c35732 100644 --- a/fastvideo/tests/mlx/tiny_h3.py +++ b/fastvideo/tests/mlx/tiny_h3.py @@ -11,14 +11,13 @@ inside the head dim, as in the real config (96 <= 128). from __future__ import annotations -import os - import numpy as np import torch -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29513") from fastvideo.configs.models.dits.minimax_h3 import ( # noqa: E402 MiniMaxH3ArchConfig, @@ -116,7 +115,9 @@ def initialize_model_parameters(model: torch.nn.Module) -> None: def build_torch_model() -> MiniMaxH3Transformer3DModel: config = build_tiny_h3_config() - model = MiniMaxH3Transformer3DModel(config=config, hf_config=build_hf_config(config)) + # The CPU float32 reference model uses TORCH_SDPA unless FASTVIDEO_ATTENTION_BACKEND is already set. + with envs.FASTVIDEO_ATTENTION_BACKEND.override(envs.FASTVIDEO_ATTENTION_BACKEND.get() or "TORCH_SDPA"): + model = MiniMaxH3Transformer3DModel(config=config, hf_config=build_hf_config(config)) model = model.to(device="cpu", dtype=torch.float32) initialize_model_parameters(model) model.eval() diff --git a/fastvideo/tests/mlx/tiny_wan.py b/fastvideo/tests/mlx/tiny_wan.py index 2770ca31f..363f09930 100644 --- a/fastvideo/tests/mlx/tiny_wan.py +++ b/fastvideo/tests/mlx/tiny_wan.py @@ -8,14 +8,13 @@ the conversions into ``fastvideo.mlx_runtime.fastwan.MLXWanDiT``. from __future__ import annotations -import os - import numpy as np import torch -os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29513") from fastvideo.models.wan.config import ( # noqa: E402 WanVideoArchConfig, @@ -114,7 +113,9 @@ def build_torch_model() -> WanTransformer3DModel: WanTransformer3DModel: The initialized model in evaluation mode. """ config = build_tiny_wan_config() - model = WanTransformer3DModel(config=config, hf_config=build_hf_config(config)) + # The CPU float32 reference model uses TORCH_SDPA unless FASTVIDEO_ATTENTION_BACKEND is already set. + with envs.FASTVIDEO_ATTENTION_BACKEND.override(envs.FASTVIDEO_ATTENTION_BACKEND.get() or "TORCH_SDPA"): + model = WanTransformer3DModel(config=config, hf_config=build_hf_config(config)) model = model.to(device="cpu", dtype=torch.float32) initialize_model_parameters(model) model.eval() diff --git a/fastvideo/tests/modal/kernel_build_cache.py b/fastvideo/tests/modal/kernel_build_cache.py index 9dd170f13..7f109e102 100644 --- a/fastvideo/tests/modal/kernel_build_cache.py +++ b/fastvideo/tests/modal/kernel_build_cache.py @@ -182,8 +182,8 @@ def _torch_metadata() -> dict[str, object]: } -def _selected_command_metadata(environment_name: str, default_command: str) -> dict[str, str]: - raw = os.environ.get(environment_name, "").strip() or default_command +def _selected_command_metadata(configured_command: str, default_command: str) -> dict[str, str]: + raw = configured_command.strip() or default_command try: command = shlex.split(raw) except ValueError as error: @@ -207,13 +207,13 @@ def _selected_command_metadata(environment_name: str, default_command: str) -> d def _compiler_libc_metadata() -> dict[str, object]: return { "compiler": { - "cc": _selected_command_metadata("CC", "cc"), - "cxx": _selected_command_metadata("CXX", "c++"), + "cc": _selected_command_metadata(os.environ.get("CC", ""), "cc"), + "cxx": _selected_command_metadata(os.environ.get("CXX", ""), "c++"), }, "build_tools": { "cmake_version": _run_optional(["cmake", "--version"]), "ninja_version": _run_optional(["ninja", "--version"]), - "linker": _selected_command_metadata("LD", "ld"), + "linker": _selected_command_metadata(os.environ.get("LD", ""), "ld"), }, "libc": { "platform_libc": list(platform.libc_ver()), @@ -252,7 +252,7 @@ def _build_metadata(repo_root: Path) -> dict[str, object]: "torch": torch_cache_metadata, "cuda": { "cuda_home": os.environ.get("CUDA_HOME", ""), - "nvcc": _selected_command_metadata("CUDACXX", "nvcc"), + "nvcc": _selected_command_metadata(os.environ.get("CUDACXX", ""), "nvcc"), }, "abi": _compiler_libc_metadata(), "build": cache_key_build, diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 15ccbb693..04a558d7f 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -418,7 +418,7 @@ def seed_grad_norm_references(): dormant Modal rollback runtime), so this function only seeds the ``L40S`` key in ``fastvideo/tests/train/methods/grad_norm_refs.json``. - ``FASTVIDEO_GRADNORM_UPDATE=1`` makes ``check_grad_norm_regression`` record + ``FASTVIDEO_TEST_GRADNORM_UPDATE=1`` makes ``check_grad_norm_regression`` record the measured norm instead of asserting; ``-rs`` surfaces the recorded value in the log so it can be copied into the JSON. @@ -428,7 +428,7 @@ def seed_grad_norm_references(): the local command and the ``_DEVICE_MAPPINGS`` table. """ run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && " - "FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs") + "FASTVIDEO_FA4=0 FASTVIDEO_TEST_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs") @app.function(gpu="L40S:1", diff --git a/fastvideo/tests/modal/test_kernel_build_cache.py b/fastvideo/tests/modal/test_kernel_build_cache.py index 3345518f6..077918564 100644 --- a/fastvideo/tests/modal/test_kernel_build_cache.py +++ b/fastvideo/tests/modal/test_kernel_build_cache.py @@ -12,6 +12,8 @@ from pathlib import Path import pytest import yaml +import fastvideo.envs as envs + def _load_kernel_build_cache(): module_path = Path(__file__).with_name("kernel_build_cache.py") @@ -86,18 +88,15 @@ def _write_prebuilt_artifact(prebuilt_root: Path, name: str, cache_key: str) -> return wheel -def _patch_stable_metadata(monkeypatch) -> None: - for name in ( - "GPU_BACKEND", - "CMAKE_ARGS", - "CFLAGS", - "CXXFLAGS", - "LDFLAGS", - "TORCH_CUDA_ARCH_LIST", - "FASTVIDEO_CONTAINER_IMAGE_REF", - "CUDACXX", - ): - monkeypatch.delenv(name, raising=False) +def _patch_stable_metadata(monkeypatch, env_overrides) -> None: + env_overrides.enter_context(envs.override_external("GPU_BACKEND", None)) + env_overrides.enter_context(envs.override_external("CMAKE_ARGS", None)) + env_overrides.enter_context(envs.override_external("CFLAGS", None)) + env_overrides.enter_context(envs.override_external("CXXFLAGS", None)) + env_overrides.enter_context(envs.override_external("LDFLAGS", None)) + env_overrides.enter_context(envs.override_external("TORCH_CUDA_ARCH_LIST", None)) + env_overrides.enter_context(envs.override_external("FASTVIDEO_CONTAINER_IMAGE_REF", None)) + env_overrides.enter_context(envs.override_external("CUDACXX", None)) monkeypatch.setattr(kernel_build_cache, "_kernel_source_hash", lambda repo_root: {"source_hash": "source"}) monkeypatch.setattr(kernel_build_cache, "_read_kernel_version", lambda repo_root: "0.3.2") monkeypatch.setattr( @@ -130,7 +129,7 @@ def _patch_stable_metadata(monkeypatch) -> None: monkeypatch.setattr( kernel_build_cache, "_selected_command_metadata", - lambda environment_name, default_command: { + lambda configured_command, default_command: { "raw": default_command, "command": json.dumps([default_command]), "resolved_executable": f"/usr/bin/{default_command}", @@ -191,14 +190,14 @@ def test_docker_image_bakes_modal_apt_and_rust_layer() -> None: assert "ENV PATH=/root/.cargo/bin:${PATH}" in dockerfile -def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None: - _patch_stable_metadata(monkeypatch) +def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, env_overrides, tmp_path) -> None: + _patch_stable_metadata(monkeypatch, env_overrides) monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") - monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "9.0a") + env_overrides.enter_context(envs.override_external("TORCH_CUDA_ARCH_LIST", "9.0a")) explicit_hopper = kernel_build_cache._build_metadata(tmp_path) - monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + env_overrides.enter_context(envs.override_external("TORCH_CUDA_ARCH_LIST", None)) detected_hopper = kernel_build_cache._build_metadata(tmp_path) monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "8.9") @@ -210,8 +209,8 @@ def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None assert detected_l40s["cache_key"] != detected_hopper["cache_key"] -def test_cache_key_ignores_runtime_only_torch_config(monkeypatch, tmp_path) -> None: - _patch_stable_metadata(monkeypatch) +def test_cache_key_ignores_runtime_only_torch_config(monkeypatch, env_overrides, tmp_path) -> None: + _patch_stable_metadata(monkeypatch, env_overrides) build_host = kernel_build_cache._build_metadata(tmp_path) torch_metadata = kernel_build_cache._torch_metadata() monkeypatch.setattr( @@ -230,19 +229,25 @@ def test_cache_key_ignores_runtime_only_torch_config(monkeypatch, tmp_path) -> N @pytest.mark.parametrize("environment_name", ["CFLAGS", "CXXFLAGS", "LDFLAGS"]) -def test_cache_key_changes_with_build_flags(monkeypatch, tmp_path, environment_name) -> None: - _patch_stable_metadata(monkeypatch) +def test_cache_key_changes_with_build_flags(monkeypatch, env_overrides, tmp_path, environment_name) -> None: + _patch_stable_metadata(monkeypatch, env_overrides) monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") baseline = kernel_build_cache._build_metadata(tmp_path) - monkeypatch.setenv(environment_name, "-DFASTVIDEO_ABI_VARIANT=1") + # The env policy contract test requires a literal name in override_external, so each flag has its own override. + flag_overrides = { + "CFLAGS": envs.override_external("CFLAGS", "-DFASTVIDEO_ABI_VARIANT=1"), + "CXXFLAGS": envs.override_external("CXXFLAGS", "-DFASTVIDEO_ABI_VARIANT=1"), + "LDFLAGS": envs.override_external("LDFLAGS", "-DFASTVIDEO_ABI_VARIANT=1"), + } + env_overrides.enter_context(flag_overrides[environment_name]) changed = kernel_build_cache._build_metadata(tmp_path) assert baseline["cache_key"] != changed["cache_key"] -def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) -> None: - _patch_stable_metadata(monkeypatch) +def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, env_overrides, tmp_path) -> None: + _patch_stable_metadata(monkeypatch, env_overrides) monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") baseline = kernel_build_cache._build_metadata(tmp_path) @@ -263,7 +268,7 @@ def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) -> ) compiler_changed = kernel_build_cache._build_metadata(tmp_path) - _patch_stable_metadata(monkeypatch) + _patch_stable_metadata(monkeypatch, env_overrides) torch_metadata = kernel_build_cache._torch_metadata() monkeypatch.setattr(kernel_build_cache, "_torch_metadata", lambda: {**torch_metadata, "cxx11_abi": False}) torch_abi_changed = kernel_build_cache._build_metadata(tmp_path) diff --git a/fastvideo/tests/modal/test_pr_test.py b/fastvideo/tests/modal/test_pr_test.py index 76100a77d..3bd8f090e 100644 --- a/fastvideo/tests/modal/test_pr_test.py +++ b/fastvideo/tests/modal/test_pr_test.py @@ -10,6 +10,8 @@ from pathlib import Path import pytest +import fastvideo.envs as envs + class _FakeImage: @@ -247,14 +249,14 @@ def test_checkout_repository_rejects_invalid_buildkite_values(monkeypatch, git_r (False, ""), ], ) -def test_run_test_command_composes_valid_post_checkout_shell(monkeypatch, build_kernel, install_command): +def test_run_test_command_composes_valid_post_checkout_shell(monkeypatch, env_overrides, build_kernel, install_command): module = _load_pr_test_module(monkeypatch) real_run = subprocess.run events = [] - monkeypatch.setenv("BUILDKITE_REPO", "https://github.com/hao-ai-lab/FastVideo.git") - monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef0123456789abcdef01234567") - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") + env_overrides.enter_context(envs.override_external("BUILDKITE_REPO", "https://github.com/hao-ai-lab/FastVideo.git")) + env_overrides.enter_context(envs.override_external("BUILDKITE_COMMIT", "0123456789abcdef0123456789abcdef01234567")) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "false")) monkeypatch.setattr( module, "_checkout_repository", @@ -290,7 +292,7 @@ def test_run_test_command_composes_valid_post_checkout_shell(monkeypatch, build_ real_run(["/bin/bash", "-n"], input=shell_command, text=True, check=True) -def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch): +def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch, env_overrides): module = _load_pr_test_module(monkeypatch) commands = [] @@ -298,9 +300,9 @@ def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch commands.append(args[-1]) return types.SimpleNamespace(returncode=0) - monkeypatch.setenv("BUILDKITE_REPO", "https://example.com/FastVideo.git") - monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef") - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") + env_overrides.enter_context(envs.override_external("BUILDKITE_REPO", "https://example.com/FastVideo.git")) + env_overrides.enter_context(envs.override_external("BUILDKITE_COMMIT", "0123456789abcdef")) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "false")) monkeypatch.setattr(module, "_checkout_repository", lambda *_args: None) monkeypatch.setattr(subprocess, "run", fake_run) diff --git a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py index b851d6de5..1bd98fa8e 100644 --- a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py +++ b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py @@ -4,6 +4,7 @@ from huggingface_hub import snapshot_download import subprocess import sys import shutil +import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline @@ -210,11 +211,10 @@ def run_training(): def test_e2e_overfit_single_sample(): - os.environ["WANDB_MODE"] = "online" - - download_data() - run_preprocessing() - run_training() + with envs.override_external("WANDB_MODE", "online"): + download_data() + run_preprocessing() + run_training() reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4") print(f"reference_video_file: {reference_video_file}") diff --git a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py index 4f82a6fd9..6bd5a8380 100644 --- a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -4,6 +4,7 @@ from huggingface_hub import snapshot_download import shutil import subprocess import sys +import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline @@ -192,11 +193,10 @@ def run_training(): def test_e2e_overfit_single_sample(): - os.environ["WANDB_MODE"] = "online" - - download_data() - run_preprocessing() - run_training() + with envs.override_external("WANDB_MODE", "online"): + download_data() + run_preprocessing() + run_training() reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4") print(f"reference_video_file: {reference_video_file}") diff --git a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py index bd9f06b18..afa6f7b90 100644 --- a/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py +++ b/fastvideo/tests/nightly/test_e2e_kandinsky5_dmd_t2v_overfit.py @@ -38,7 +38,7 @@ and go green on stale artifacts. Reference bootstrap: the committed reference (``reference_video_kandinsky5_dmd_v0.mp4`` next to this file) is produced by running this test once on a sanctioned GPU box with -``KANDINSKY5_E2E_WRITE_REFERENCE=1``, reviewing the written video by eye, +``FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE=1``, reviewing the written video by eye, and committing it. Review bar: expect blurry-but-structured content loosely matching the prompt (4-step no-CFG sampling of near-base weights cannot look good) -- soft shapes, warm sunflower-ish colors, motion. @@ -62,6 +62,7 @@ this test fully controls. from __future__ import annotations +import contextlib import json import os import shutil @@ -73,10 +74,11 @@ import cv2 import numpy as np import pytest +import fastvideo.envs as envs + REPO_ROOT = Path(__file__).resolve().parents[3] THIS_FILE = Path(__file__).resolve() -NUM_GPUS = os.environ.get("KANDINSKY5_E2E_NUM_GPUS", "1") # Every artifact this test reads or writes lives under this single # test-owned root. It must NEVER point at (or contain) the documented # default dataset paths -- data/kandinsky5_overfit and @@ -111,7 +113,6 @@ STAGE2_CONFIG = (REPO_ROOT / "examples" / "train" / "configs" / "distribution_ma # committed next to this file -- mirroring reference_video_1_sample_v0.mp4 # in the Wan e2e test. REFERENCE_VIDEO = THIS_FILE.parent / "reference_video_kandinsky5_dmd_v0.mp4" -WRITE_REFERENCE_ENV = "KANDINSKY5_E2E_WRITE_REFERENCE" # Deliberately semantic even though the training clip is random noise: the # caption's relation to the clip content is irrelevant for plumbing, but # the caption doubles as the generation prompt, and a semantic prompt makes @@ -201,15 +202,15 @@ def _synthesize_single_sample() -> None: writer.write(frame) writer.release() - # One physical clip, but max(2, NUM_GPUS) manifest rows all pointing at + # One physical clip, but max(2, FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS) manifest rows all pointing at # it: DP_SP_BatchSampler (fastvideo/dataset/parquet_dataset_map_style.py) # floor-divides the number of batches by the number of data-parallel # groups with drop_last=True, so a dataset with fewer rows than GPUs # yields an EMPTY dataloader on every rank. Duplicate rows keep the # single-sample-overfit semantics (identical latents/caption) while - # making KANDINSKY5_E2E_NUM_GPUS=2+ a working escape hatch when one GPU + # making FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS=2+ a working escape hatch when one GPU # doesn't have the memory headroom for stage 2 (see _run_stage). - num_rows = max(2, int(NUM_GPUS)) + num_rows = max(2, envs.FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS.get()) with open(RAW_DATA_DIR / "videos2caption.json", "w") as f: json.dump([{ "path": "sample_0.mp4", @@ -269,6 +270,7 @@ def _latest_checkpoint(output_dir: Path) -> Path: def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_overrides: list[str], env_overrides: dict[str, str]) -> None: + num_gpus = str(envs.FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS.get()) cmd = [ sys.executable, "-m", @@ -276,7 +278,7 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov "--nnodes", "1", "--nproc_per_node", - NUM_GPUS, + num_gpus, "-m", "fastvideo.train.entrypoint.train", "--config", @@ -284,9 +286,9 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov "--training.data.data_path", str(PREPROCESSED_DIR), "--training.distributed.num_gpus", - NUM_GPUS, + num_gpus, "--training.distributed.hsdp_shard_dim", - NUM_GPUS, + num_gpus, "--training.loop.max_train_steps", str(max_train_steps), "--training.checkpoint.output_dir", @@ -312,7 +314,7 @@ def _run_stage(config: Path, output_dir: Path, *, max_train_steps: int, extra_ov # to 1 above -- a recorded run peaked at ~76 GiB and died allocating # the last MiBs on an 80 GiB device with ~1.9 GiB lost to # fragmentation). Expandable segments reclaims that headroom; if a - # single device still can't fit, shard with KANDINSKY5_E2E_NUM_GPUS=2+ + # single device still can't fit, shard with FASTVIDEO_TEST_KANDINSKY5_E2E_NUM_GPUS=2+ # (the synthesized manifest guarantees >= one row per rank). env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") env.update(env_overrides) @@ -432,7 +434,7 @@ def _assert_video_not_degenerate(video_path: Path) -> None: def _assert_matches_reference(video_path: Path) -> None: if not REFERENCE_VIDEO.exists(): - if os.environ.get(WRITE_REFERENCE_ENV) == "1": + if envs.FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE.get(): shutil.copy2(video_path, REFERENCE_VIDEO) print(f"\nWrote new reference video to {REFERENCE_VIDEO} -- " "review it by eye and commit it. This bootstrap run " @@ -440,7 +442,7 @@ def _assert_matches_reference(video_path: Path) -> None: else: pytest.fail(f"reference video missing at {REFERENCE_VIDEO}. This test " "must not pass on artifact existence alone -- run once on a " - f"sanctioned GPU box with {WRITE_REFERENCE_ENV}=1, review " + "sanctioned GPU box with FASTVIDEO_TEST_KANDINSKY5_E2E_WRITE_REFERENCE=1, review " "the written video, and commit it (see module docstring).") from fastvideo.tests.utils import compute_video_ssim_torchvision @@ -461,13 +463,13 @@ def _assert_matches_reference(video_path: Path) -> None: @pytest.mark.nightly -def test_e2e_kandinsky5_dmd_overfit_single_sample(): +def test_e2e_kandinsky5_dmd_overfit_single_sample(env_overrides): if not STAGE1_CONFIG.exists() or not STAGE2_CONFIG.exists(): pytest.skip( "Kandinsky5 QAT configs not found -- see examples/train/configs/{fine_tuning,distribution_matching}/kandinsky5/" ) - os.environ.setdefault("WANDB_MODE", "offline") + env_overrides.enter_context(envs.override_external("WANDB_MODE", os.environ.get("WANDB_MODE", "offline"))) _clean_previous_artifacts() _synthesize_single_sample() @@ -543,4 +545,5 @@ if __name__ == "__main__": if len(sys.argv) >= 2 and sys.argv[1] == "--generate": _generate_main(sys.argv[2], sys.argv[3]) else: - test_e2e_kandinsky5_dmd_overfit_single_sample() + with contextlib.ExitStack() as env_overrides: + test_e2e_kandinsky5_dmd_overfit_single_sample(env_overrides) diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py index 476bc55c3..6eeb0daf5 100644 --- a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py +++ b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py @@ -14,7 +14,7 @@ Requires the FastVideo LTX-2 distilled checkpoints (set HF_HOME to a cache that contains them, or allow ~60GB of downloads per version). The NVFP4-QAT case additionally requires FlashInfer FP4 support and SM100+. -Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast. +Guarded by FASTVIDEO_TEST_NIGHTLY=1 so the default test suite stays fast. """ import glob @@ -32,6 +32,7 @@ from huggingface_hub import snapshot_download sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) +import fastvideo.envs as envs # noqa: E402 from fastvideo.tests.utils import compute_video_ssim_torchvision # noqa: E402 NUM_GPUS_TRAINING = "4" @@ -136,12 +137,10 @@ def _validation_videos_by_step(out_dir: Path) -> dict[int, str]: return videos -@pytest.mark.skipif( - os.environ.get("FASTVIDEO_NIGHTLY") != "1", - reason="nightly e2e overfit test; set FASTVIDEO_NIGHTLY=1 to run", -) @pytest.mark.parametrize("case_id", _CASES.keys()) def test_e2e_ltx2_overfit_new_stack(case_id: str): + if not envs.FASTVIDEO_TEST_NIGHTLY.get(): + pytest.skip("nightly e2e overfit test; set FASTVIDEO_TEST_NIGHTLY=1 to run") case = _CASES[case_id] if case.get("requires_nvfp4"): if (not torch.cuda.is_available() or torch.cuda.get_device_capability(0) < (10, 0)): @@ -186,6 +185,6 @@ def test_e2e_ltx2_overfit_new_stack(case_id: str): if __name__ == "__main__": - os.environ.setdefault("FASTVIDEO_NIGHTLY", "1") - for _case_id in _CASES: - test_e2e_ltx2_overfit_new_stack(_case_id) + with envs.FASTVIDEO_TEST_NIGHTLY.override(True): + for _case_id in _CASES: + test_e2e_ltx2_overfit_new_stack(_case_id) diff --git a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py index 1d1f8f9c7..ebdc4c1b9 100644 --- a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py @@ -4,6 +4,7 @@ from huggingface_hub import snapshot_download import shutil import subprocess import sys +import fastvideo.envs as envs from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline @@ -197,11 +198,10 @@ def run_training(): def test_e2e_overfit_single_sample(): - os.environ["WANDB_MODE"] = "online" - - download_data() - run_preprocessing() - run_training() + with envs.override_external("WANDB_MODE", "online"): + download_data() + run_preprocessing() + run_training() reference_video_file = os.path.join(os.path.dirname(__file__), "reference_video_1_sample_v0.mp4") print(f"reference_video_file: {reference_video_file}") diff --git a/fastvideo/tests/performance/identity.py b/fastvideo/tests/performance/identity.py index 4732620c3..7a19d00ff 100644 --- a/fastvideo/tests/performance/identity.py +++ b/fastvideo/tests/performance/identity.py @@ -17,6 +17,8 @@ from typing import Any import torch +import fastvideo.envs as envs + RECIPE_SCHEMA_VERSION = 2 PROFILE_ID_LENGTH = 16 _PATH_EXCLUDED_GENERATION_KEYS = {"output_path", "output_video_name"} @@ -138,7 +140,7 @@ def build_recipe_from_benchmark_config( prompts = list(measured_prompts) if measured_prompts is not None else list( cfg.get("test_prompts") or ["A cinematic video."]) if attention_backend is None: - attention_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") + attention_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get() negative_prompt = generation_kwargs.pop("negative_prompt", generation_kwargs.pop("neg_prompt", None)) generation_recipe = { @@ -272,6 +274,17 @@ def software_profile( } +def _process_profile_env() -> dict[str, str | None]: + """Read the variables that _PROFILE_ENV_VARS names from the environment of this process.""" + return { + "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"), + "FASTVIDEO_ATTENTION_BACKEND": envs.FASTVIDEO_ATTENTION_BACKEND.get(), + "IMAGE_VERSION": os.environ.get("IMAGE_VERSION"), + "UV_TORCH_BACKEND": os.environ.get("UV_TORCH_BACKEND"), + "FASTVIDEO_CONTAINER_IMAGE_REF": os.environ.get("FASTVIDEO_CONTAINER_IMAGE_REF"), + } + + def environment_metadata( *, env: Mapping[str, str] | None = None, @@ -281,7 +294,7 @@ def environment_metadata( ) -> dict[str, Any]: """Return audit metadata kept separate from comparable identity keys.""" - source_env = env if env is not None else os.environ + source_env = env if env is not None else _process_profile_env() full_package_versions = dict(package_versions) if package_versions is not None else _installed_package_versions() return { "python": { diff --git a/fastvideo/tests/performance/test_compare_baseline_policy.py b/fastvideo/tests/performance/test_compare_baseline_policy.py index 5de5b7316..336bfbdb9 100644 --- a/fastvideo/tests/performance/test_compare_baseline_policy.py +++ b/fastvideo/tests/performance/test_compare_baseline_policy.py @@ -6,6 +6,7 @@ from pathlib import Path import pytest +import fastvideo.envs as envs from fastvideo.performance import hf_store from fastvideo.tests.performance import compare_baseline from fastvideo.tests.performance import seed_baseline @@ -19,9 +20,9 @@ from fastvideo.performance.metric_policy import resolve_metric_policies @pytest.fixture(autouse=True) -def _stub_seed_baseline_remote_sync(monkeypatch): - monkeypatch.delenv("BUILDKITE_BUILD_ID", raising=False) - monkeypatch.delenv("BUILDKITE_JOB_ID", raising=False) +def _stub_seed_baseline_remote_sync(monkeypatch, env_overrides): + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_ID", None)) + env_overrides.enter_context(envs.override_external("BUILDKITE_JOB_ID", None)) def sync(local_dir, *, strict=False, revision=None): assert strict is True @@ -129,25 +130,25 @@ def test_sync_marker_reuse_requires_the_exact_revision(tmp_path): assert hf_store._sync_marker_matches_request(str(marker), "pinned-old-revision") is False -def test_detect_run_source_prefers_explicit_env(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "local") - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "123") +def test_detect_run_source_prefers_explicit_env(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "local")) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "123")) assert compare_baseline._detect_run_source() == "local" -def test_detect_run_source_infers_pr(monkeypatch): - monkeypatch.delenv("PERF_RUN_SOURCE", raising=False) - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "123") +def test_detect_run_source_infers_pr(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", None)) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "123")) assert compare_baseline._detect_run_source() == "pr" -def test_detect_run_source_infers_scheduled_main(monkeypatch): - monkeypatch.delenv("PERF_RUN_SOURCE", raising=False) - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") - monkeypatch.setenv("BUILDKITE_BRANCH", "main") - monkeypatch.setenv("TEST_SCOPE", "full") +def test_detect_run_source_infers_scheduled_main(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", None)) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "false")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BRANCH", "main")) + env_overrides.enter_context(envs.override_external("TEST_SCOPE", "full")) assert compare_baseline._detect_run_source() == "scheduled_main" @@ -165,13 +166,13 @@ def test_upload_policy_always_uploads_failures(monkeypatch): assert compare_baseline._upload_allowed({"success": False}) is True -def test_normalized_record_includes_source_metadata(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") - monkeypatch.setenv("BUILDKITE_BRANCH", "feature/perf") - monkeypatch.setenv("TEST_SCOPE", "direct") - monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/build") - monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-1") - monkeypatch.setenv("BUILDKITE_JOB_ID", "job-1") +def test_normalized_record_includes_source_metadata(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BRANCH", "feature/perf")) + env_overrides.enter_context(envs.override_external("TEST_SCOPE", "direct")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_URL", "https://buildkite.example/build")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_ID", "build-1")) + env_overrides.enter_context(envs.override_external("BUILDKITE_JOB_ID", "job-1")) record = compare_baseline.normalize_performance_result(_raw_result()) @@ -185,8 +186,8 @@ def test_normalized_record_includes_source_metadata(monkeypatch): assert record["job_id"] == "job-1" -def test_normalized_record_includes_effective_regression_thresholds(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_normalized_record_includes_effective_regression_thresholds(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) raw = _raw_result() raw["regression_thresholds"] = { "latency": { @@ -232,8 +233,8 @@ def test_boolean_regression_threshold_values_are_ignored(): assert latency.gated is False -def test_normalized_record_preserves_identity_metadata(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_normalized_record_preserves_identity_metadata(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) raw = _raw_result() raw.update({ "result_schema_version": 2, @@ -280,13 +281,13 @@ def test_normalized_record_preserves_identity_metadata(monkeypatch): assert record["quality_metadata"] == {"quality_status": "canonical"} -def test_normalized_record_prefers_raw_v2_provenance(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") - monkeypatch.setenv("BUILDKITE_BRANCH", "feature/from-env") - monkeypatch.setenv("TEST_SCOPE", "direct") - monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/env") - monkeypatch.setenv("BUILDKITE_BUILD_ID", "env-build") - monkeypatch.setenv("BUILDKITE_JOB_ID", "env-job") +def test_normalized_record_prefers_raw_v2_provenance(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BRANCH", "feature/from-env")) + env_overrides.enter_context(envs.override_external("TEST_SCOPE", "direct")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_URL", "https://buildkite.example/env")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_ID", "env-build")) + env_overrides.enter_context(envs.override_external("BUILDKITE_JOB_ID", "env-job")) raw = _raw_result() raw.update({ "result_schema_version": 2, @@ -311,21 +312,21 @@ def test_normalized_record_prefers_raw_v2_provenance(monkeypatch): assert record["pr_number"] == "" -def test_v1_normalized_record_has_no_result_schema_version(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_v1_normalized_record_has_no_result_schema_version(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = compare_baseline.normalize_performance_result(_raw_result()) assert "result_schema_version" not in record -def test_main_writes_v1_current_artifact_without_comparison_identity(monkeypatch, tmp_path, capsys): +def test_main_writes_v1_current_artifact_without_comparison_identity(monkeypatch, env_overrides, tmp_path, capsys): results_dir = tmp_path / "results" reports_dir = tmp_path / "reports" tracking_root = tmp_path / "tracking" results_dir.mkdir() - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) raw_result = perf_test._build_result_record( cfg={"benchmark_id": "wan-t2v-1.3b-2gpu"}, model_info={}, @@ -379,8 +380,8 @@ def test_main_writes_v1_current_artifact_without_comparison_identity(monkeypatch assert uploaded_records[0]["baseline_eligible"] is False -def test_normalized_record_reads_identity_labels_from_recipe(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_normalized_record_reads_identity_labels_from_recipe(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) raw = _raw_result() raw["recipe"] = { "benchmark": { @@ -540,8 +541,8 @@ def test_informational_metric_remains_visible_without_failing(): assert row["failing_metrics"] == [] -def test_normalized_record_preserves_v2_comparison_identity(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_normalized_record_preserves_v2_comparison_identity(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record() @@ -553,8 +554,8 @@ def test_normalized_record_preserves_v2_comparison_identity(monkeypatch): assert record["software_profile_id"] == "sw-cuda" -def test_exact_comparable_baseline_without_regression_reports_pass(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_exact_comparable_baseline_without_regression_reports_pass(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record() baseline_records = [{ "latency": 10.0, @@ -576,8 +577,8 @@ def test_exact_comparable_baseline_without_regression_reports_pass(monkeypatch): assert "no gated regressions" in reason -def test_slower_gated_metric_reports_regression(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_slower_gated_metric_reports_regression(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record(avg_generation_time_s=11.0) baseline_records = [{"latency": 10.0}] @@ -596,8 +597,8 @@ def test_slower_gated_metric_reports_regression(monkeypatch): assert reason == failures[0] -def test_missing_exact_v2_baseline_reports_calibration_needed(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_missing_exact_v2_baseline_reports_calibration_needed(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record() failures, status, reason = compare_baseline._evaluate_record_comparison( @@ -616,6 +617,7 @@ def test_missing_exact_v2_baseline_reports_calibration_needed(monkeypatch): def test_seeded_v2_calibration_artifact_enables_next_compare_pass( monkeypatch, + env_overrides, tmp_path, ): tracking_root = tmp_path / "tracking" @@ -657,13 +659,13 @@ def test_seeded_v2_calibration_artifact_enables_next_compare_pass( timestamp="2026-06-18T00:00:00+00:00", ), f) - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root)) monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir)) monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never") monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir) - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) assert compare_baseline.main() == 0 @@ -676,8 +678,8 @@ def test_seeded_v2_calibration_artifact_enables_next_compare_pass( assert normalized["baseline_eligible"] is False -def test_baseline_seed_rejects_non_calibration_sources(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") +def test_baseline_seed_rejects_non_calibration_sources(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) record = _v2_record() record.update({ "comparison_status": compare_baseline.STATUS_PASS, @@ -726,8 +728,8 @@ def test_baseline_seed_rejects_non_calibration_sources(monkeypatch): }, "main-branch full-suite"), ], ) -def test_baseline_seed_rejects_untrusted_calibration_sources(monkeypatch, overrides, match): - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") +def test_baseline_seed_rejects_untrusted_calibration_sources(env_overrides, overrides, match): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) record = _v2_record() record.update({ "comparison_status": compare_baseline.STATUS_CALIBRATION_NEEDED, @@ -770,8 +772,8 @@ def test_baseline_seed_rejects_invalid_calibration_schema(overrides, match): seed_baseline.build_baseline_seed_record(record, reason="invalid source") -def test_baseline_seed_rejects_mixed_exact_identities(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") +def test_baseline_seed_rejects_mixed_exact_identities(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) first = _v2_record() second = _v2_record(variant_id="different-variant") for record in (first, second): @@ -1415,8 +1417,8 @@ def test_baseline_seed_rejects_inconsistent_batch_before_persistence( assert not tracking_root.exists() -def test_same_variant_changed_recipe_reports_recipe_mismatch(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_same_variant_changed_recipe_reports_recipe_mismatch(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record(recipe_fingerprint="recipe-2") recipe_mismatch_records = [{ "recipe_fingerprint": "recipe-1", @@ -1437,8 +1439,8 @@ def test_same_variant_changed_recipe_reports_recipe_mismatch(monkeypatch): assert "recipe-1" in reason -def test_same_recipe_calibration_record_does_not_report_recipe_mismatch(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") +def test_same_recipe_calibration_record_does_not_report_recipe_mismatch(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) record = _v2_record() assert compare_baseline._recipe_mismatch_records( @@ -1452,6 +1454,7 @@ def test_same_recipe_calibration_record_does_not_report_recipe_mismatch(monkeypa def test_main_recipe_mismatch_takes_precedence_over_static_regression( monkeypatch, + env_overrides, tmp_path, ): tracking_root = tmp_path / "tracking" @@ -1494,7 +1497,7 @@ def test_main_recipe_mismatch_takes_precedence_over_static_regression( monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never") monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir) - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) assert compare_baseline.main() == 1 @@ -1510,6 +1513,7 @@ def test_main_recipe_mismatch_takes_precedence_over_static_regression( def test_prior_pr_calibration_record_does_not_report_recipe_mismatch( monkeypatch, + env_overrides, tmp_path, ): tracking_root = tmp_path / "tracking" @@ -1545,7 +1549,7 @@ def test_prior_pr_calibration_record_does_not_report_recipe_mismatch( monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never") monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir) - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) assert compare_baseline.main() == 0 @@ -1616,6 +1620,7 @@ def test_v2_identity_lookup_ignores_model_id_and_gpu_display_string(tmp_path): def test_main_v2_identity_lookup_ignores_model_id_and_gpu_display_string( monkeypatch, + env_overrides, tmp_path, ): tracking_root = tmp_path / "tracking" @@ -1644,13 +1649,13 @@ def test_main_v2_identity_lookup_ignores_model_id_and_gpu_display_string( with open(results_dir / "perf_current.json", "w", encoding="utf-8") as f: json.dump(_v2_raw_result(), f) - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root)) monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir)) monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never") monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir) - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) assert compare_baseline.main() == 0 @@ -1665,6 +1670,7 @@ def test_main_v2_identity_lookup_ignores_model_id_and_gpu_display_string( def test_scheduled_main_static_regression_does_not_contaminate_passing_record( monkeypatch, + env_overrides, tmp_path, ): tracking_root = tmp_path / "tracking" @@ -1722,8 +1728,8 @@ def test_scheduled_main_static_regression_does_not_contaminate_passing_record( f, ) - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") - monkeypatch.setenv("PERF_PYTEST_RC", "1") + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", "1")) monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root)) monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir)) monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) @@ -1748,8 +1754,8 @@ def test_scheduled_main_static_regression_does_not_contaminate_passing_record( assert normalized["passing-benchmark"]["baseline_eligible"] is True -def test_unattributed_performance_pytest_failure_reports_infra_error(monkeypatch): - monkeypatch.setenv("PERF_PYTEST_RC", "2") +def test_unattributed_performance_pytest_failure_reports_infra_error(env_overrides): + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", "2")) record = _v2_record() failures, status, reason = compare_baseline._evaluate_record_comparison( @@ -1768,7 +1774,7 @@ def test_unattributed_performance_pytest_failure_reports_infra_error(monkeypatch assert "without an attributable static-threshold regression" in reason -def test_main_partial_v2_identity_takes_precedence_over_static_regression(monkeypatch, tmp_path): +def test_main_partial_v2_identity_takes_precedence_over_static_regression(monkeypatch, env_overrides, tmp_path): results_dir = tmp_path / "results" reports_dir = tmp_path / "reports" tracking_root = tmp_path / "tracking" @@ -1782,13 +1788,13 @@ def test_main_partial_v2_identity_takes_precedence_over_static_regression(monkey with open(results_dir / "perf_current.json", "w", encoding="utf-8") as f: json.dump(raw, f) - monkeypatch.setenv("PERF_RUN_SOURCE", "pr") + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "pr")) monkeypatch.setattr(compare_baseline, "TRACKING_ROOT", str(tracking_root)) monkeypatch.setattr(compare_baseline, "RESULTS_DIR", str(results_dir)) monkeypatch.setattr(compare_baseline, "PERF_REPORTS_DIR", str(reports_dir)) monkeypatch.setattr(compare_baseline, "UPLOAD_POLICY", "never") monkeypatch.setattr(compare_baseline, "sync_from_hf", lambda local_dir, strict=False: local_dir) - monkeypatch.delenv("PERF_PYTEST_RC", raising=False) + env_overrides.enter_context(envs.override_external("PERF_PYTEST_RC", None)) assert compare_baseline.main() == 1 diff --git a/fastvideo/tests/performance/test_dashboard_service.py b/fastvideo/tests/performance/test_dashboard_service.py index cb76b2ad0..cae52fdd0 100644 --- a/fastvideo/tests/performance/test_dashboard_service.py +++ b/fastvideo/tests/performance/test_dashboard_service.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import fastvideo.envs as envs from fastvideo.performance import hf_store from fastvideo.performance_dashboard.service import build_latest_summary, build_trends, filter_records @@ -477,11 +478,10 @@ def test_trends_include_source_metadata_with_legacy_defaults(): assert trends[0]["points"][1]["baseline_eligible"] is True -def test_hf_token_resolution_accepts_standard_env_names(monkeypatch): - for env_var in hf_store.HF_TOKEN_ENV_VARS: - monkeypatch.delenv(env_var, raising=False) +def test_hf_token_resolution_accepts_standard_env_names(env_overrides): + env_overrides.enter_context(envs.override_external("HUGGING_FACE_HUB_TOKEN", None)) - monkeypatch.setenv("HF_TOKEN", "hf_local") + env_overrides.enter_context(envs.override_external("HF_TOKEN", "hf_local")) assert hf_store.resolve_hf_token() == "hf_local" diff --git a/fastvideo/tests/performance/test_inference_performance.py b/fastvideo/tests/performance/test_inference_performance.py index 99ceea70e..3bc0c1a97 100644 --- a/fastvideo/tests/performance/test_inference_performance.py +++ b/fastvideo/tests/performance/test_inference_performance.py @@ -17,6 +17,7 @@ from typing import Any import torch import pytest +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.logger import init_logger from fastvideo.tests.performance.identity import ( @@ -383,8 +384,8 @@ def _build_identity_fields(cfg, init_kwargs, prompt, runtime_identity): hw_profile = hardware_profile(num_gpus=num_gpus) sw_profile = software_profile() sw_profile.update({ - "attention_backend": os.environ.get("FASTVIDEO_ATTENTION_BACKEND") or "auto", - "flash_attention_4_enabled": os.environ.get("FASTVIDEO_FA4", "0") != "0", + "attention_backend": envs.FASTVIDEO_ATTENTION_BACKEND.get() or "auto", + "flash_attention_4_enabled": envs.FASTVIDEO_FA4.get(), }) performance_profile_version = os.environ.get("FASTVIDEO_PERFORMANCE_PROFILE_VERSION") if performance_profile_version: @@ -633,12 +634,5 @@ def test_inference_performance(cfg): (text encoder, DiT, VAE decode). Assert each against device-aware thresholds. """ - original_env = os.environ.get("FASTVIDEO_STAGE_LOGGING") - os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" - try: + with envs.FASTVIDEO_STAGE_LOGGING.override(True): _run_benchmark(cfg) - finally: - if original_env is None: - os.environ.pop("FASTVIDEO_STAGE_LOGGING", None) - else: - os.environ["FASTVIDEO_STAGE_LOGGING"] = original_env diff --git a/fastvideo/tests/performance/test_inference_performance_identity.py b/fastvideo/tests/performance/test_inference_performance_identity.py index f8eda3340..128406912 100644 --- a/fastvideo/tests/performance/test_inference_performance_identity.py +++ b/fastvideo/tests/performance/test_inference_performance_identity.py @@ -2,6 +2,7 @@ import copy +import fastvideo.envs as envs from fastvideo.tests.performance import compare_baseline from fastvideo.tests.performance import test_inference_performance as perf @@ -45,8 +46,8 @@ def _benchmark_config(): } -def test_performance_producer_emits_v2_identity_from_raw_result_shape(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") +def test_performance_producer_emits_v2_identity_from_raw_result_shape(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) cfg = _benchmark_config() init_kwargs = dict(cfg["init_kwargs"]) cfg["generation_kwargs"]["output_path"] = "/tmp/generated" @@ -95,8 +96,8 @@ def test_performance_producer_emits_v2_identity_from_raw_result_shape(monkeypatc assert "output_path" not in record["recipe"]["generation_kwargs"] -def test_display_benchmark_id_rename_preserves_producer_fingerprint_and_cohort(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") +def test_display_benchmark_id_rename_preserves_producer_fingerprint_and_cohort(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) original = _benchmark_config() renamed = copy.deepcopy(original) renamed["benchmark_id"] = "wan-t2v-renamed-display-id" @@ -122,8 +123,8 @@ def test_display_benchmark_id_rename_preserves_producer_fingerprint_and_cohort(m original_identity) == compare_baseline._comparison_identity_filters(renamed_identity) -def test_v2_identity_tolerates_null_run_config(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") +def test_v2_identity_tolerates_null_run_config(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) cfg = _benchmark_config() cfg["run_config"] = None init_kwargs = dict(cfg["init_kwargs"]) @@ -141,8 +142,8 @@ def test_v2_identity_tolerates_null_run_config(monkeypatch): assert identity_fields["hardware_profile"]["gpu_count"] == 2 -def test_effective_gpu_count_is_recorded_in_recipe_and_hardware(monkeypatch): - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") +def test_effective_gpu_count_is_recorded_in_recipe_and_hardware(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) cfg = _benchmark_config() del cfg["init_kwargs"]["num_gpus"] init_kwargs = dict(cfg["init_kwargs"]) @@ -161,7 +162,7 @@ def test_effective_gpu_count_is_recorded_in_recipe_and_hardware(monkeypatch): assert identity_fields["hardware_profile"]["gpu_count"] == 2 -def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypatch): +def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypatch, env_overrides): base_profile = { "python": "3.12", "pytorch": "2.7", @@ -172,14 +173,15 @@ def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypat }, } monkeypatch.setattr(perf, "software_profile", lambda: dict(base_profile)) - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN") - monkeypatch.setenv("FASTVIDEO_FA4", "1") - monkeypatch.setenv("FASTVIDEO_PERFORMANCE_PROFILE_VERSION", "perf-profile-v2") - monkeypatch.setenv("IMAGE_VERSION", "py3.12-cuda13.0") - monkeypatch.setenv( - "FASTVIDEO_CONTAINER_IMAGE_REF", - "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:abc", - ) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("SAGE_ATTN")) + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(True)) + env_overrides.enter_context(envs.override_external("FASTVIDEO_PERFORMANCE_PROFILE_VERSION", "perf-profile-v2")) + env_overrides.enter_context(envs.override_external("IMAGE_VERSION", "py3.12-cuda13.0")) + env_overrides.enter_context( + envs.override_external( + "FASTVIDEO_CONTAINER_IMAGE_REF", + "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:abc", + )) cfg = _benchmark_config() identity_fields = perf._build_identity_fields( @@ -204,10 +206,11 @@ def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypat assert identity_fields["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith("sha256:abc") first_profile_id = identity_fields["software_profile_id"] - monkeypatch.setenv( - "FASTVIDEO_CONTAINER_IMAGE_REF", - "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:def", - ) + env_overrides.enter_context( + envs.override_external( + "FASTVIDEO_CONTAINER_IMAGE_REF", + "ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev@sha256:def", + )) changed_audit = perf._build_identity_fields( cfg, dict(cfg["init_kwargs"]), @@ -220,7 +223,7 @@ def test_producer_tracks_runtime_software_identity_and_container_audit(monkeypat assert changed_audit["software_profile_id"] == first_profile_id assert changed_audit["environment_metadata"]["env"]["FASTVIDEO_CONTAINER_IMAGE_REF"].endswith("sha256:def") - monkeypatch.setenv("FASTVIDEO_FA4", "0") + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(False)) changed_runtime = perf._build_identity_fields( cfg, dict(cfg["init_kwargs"]), diff --git a/fastvideo/tests/performance/test_inference_performance_result_schema.py b/fastvideo/tests/performance/test_inference_performance_result_schema.py index 30c6191a8..2a45f626a 100644 --- a/fastvideo/tests/performance/test_inference_performance_result_schema.py +++ b/fastvideo/tests/performance/test_inference_performance_result_schema.py @@ -2,6 +2,7 @@ import pytest +import fastvideo.envs as envs from fastvideo.tests.performance import test_inference_performance as perf_test @@ -60,15 +61,15 @@ def _identity_fields(): } -def test_build_result_record_emits_v2_wan_shape(monkeypatch): - monkeypatch.setenv("PERF_RUN_SOURCE", "scheduled_main") - monkeypatch.setenv("BUILDKITE_COMMIT", "a" * 40) - monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") - monkeypatch.setenv("BUILDKITE_BRANCH", "main") - monkeypatch.setenv("TEST_SCOPE", "full") - monkeypatch.setenv("BUILDKITE_BUILD_URL", "https://buildkite.example/build") - monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-1") - monkeypatch.setenv("BUILDKITE_JOB_ID", "job-1") +def test_build_result_record_emits_v2_wan_shape(monkeypatch, env_overrides): + env_overrides.enter_context(envs.override_external("PERF_RUN_SOURCE", "scheduled_main")) + env_overrides.enter_context(envs.override_external("BUILDKITE_COMMIT", "a" * 40)) + env_overrides.enter_context(envs.override_external("BUILDKITE_PULL_REQUEST", "false")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BRANCH", "main")) + env_overrides.enter_context(envs.override_external("TEST_SCOPE", "full")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_URL", "https://buildkite.example/build")) + env_overrides.enter_context(envs.override_external("BUILDKITE_BUILD_ID", "build-1")) + env_overrides.enter_context(envs.override_external("BUILDKITE_JOB_ID", "job-1")) monkeypatch.setattr(perf_test, "_build_identity_fields", lambda *_args: _identity_fields()) record = perf_test._build_result_record( diff --git a/fastvideo/tests/ssim/README.md b/fastvideo/tests/ssim/README.md index 922615ba7..25f708195 100644 --- a/fastvideo/tests/ssim/README.md +++ b/fastvideo/tests/ssim/README.md @@ -12,7 +12,7 @@ reference layout: `/_reference_videos/...` is still read as fallback.) Before SSIM tests run, missing reference videos are auto-downloaded from a -public HF repo (configured by `FASTVIDEO_SSIM_REFERENCE_HF_REPO`, default: +public HF repo (configured by `FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO`, default: `FastVideo/ssim-reference-videos`). Use the CLI: diff --git a/fastvideo/tests/ssim/bootstrap_references.py b/fastvideo/tests/ssim/bootstrap_references.py index e9c380369..73bd44a92 100644 --- a/fastvideo/tests/ssim/bootstrap_references.py +++ b/fastvideo/tests/ssim/bootstrap_references.py @@ -7,21 +7,16 @@ from pathlib import Path import pytest +import fastvideo.envs as envs from fastvideo.tests.ssim.reference_utils import get_output_quality_tier from fastvideo.tests.ssim.reference_videos_cli import ( - BOOTSTRAP_ENV_KEY, - DEFAULT_REPO_ID, - DEFAULT_REPO_TYPE, - HF_REPO_ENV_KEY, - HF_REPO_TYPE_ENV_KEY, - upload_draft_reference_artifact, -) + upload_draft_reference_artifact, ) TRUE_VALUES = {"1", "true", "yes", "on"} def bootstrap_mode_enabled() -> bool: - return os.environ.get(BOOTSTRAP_ENV_KEY, "").strip().lower() in TRUE_VALUES + return os.environ.get("FASTVIDEO_SSIM_BOOTSTRAP_MODE", "").strip().lower() in TRUE_VALUES def xfail_missing_reference_in_bootstrap_mode( @@ -38,8 +33,8 @@ def xfail_missing_reference_in_bootstrap_mode( raise FileNotFoundError( f"SSIM bootstrap mode is enabled, but generated {artifact_kind} artifact is missing: {generated_path}") - repo_id = os.environ.get(HF_REPO_ENV_KEY, DEFAULT_REPO_ID) - repo_type = os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE) + repo_id = envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.get() + repo_type = envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE.get() draft_path = upload_draft_reference_artifact( repo_id=repo_id, repo_type=repo_type, diff --git a/fastvideo/tests/ssim/conftest.py b/fastvideo/tests/ssim/conftest.py index d93a59eb1..bbfc5caab 100644 --- a/fastvideo/tests/ssim/conftest.py +++ b/fastvideo/tests/ssim/conftest.py @@ -1,11 +1,10 @@ +import contextlib import os +import fastvideo.envs as envs from fastvideo.tests.ssim.reference_utils import ( - FULL_QUALITY_ENV_VAR, - get_output_quality_tier, -) + get_output_quality_tier, ) from fastvideo.tests.ssim.reference_videos_cli import ( - BOOTSTRAP_ENV_KEY, HF_REPO_ENV_KEY, ensure_reference_videos_available, ) @@ -21,7 +20,7 @@ def pytest_addoption(parser): ) parser.addoption( "--ssim-reference-repo", - default=os.environ.get(HF_REPO_ENV_KEY, ""), + default="", help=("HF repo id for SSIM reference videos " f"(overrides {HF_REPO_ENV_KEY})."), ) @@ -41,18 +40,20 @@ def pytest_addoption(parser): def pytest_configure(config): if config.getoption("--ssim-full-quality"): - os.environ[FULL_QUALITY_ENV_VAR] = "1" + envs.FASTVIDEO_TEST_SSIM_FULL_QUALITY.set(True) repo_id = config.getoption("--ssim-reference-repo") if repo_id: - os.environ[HF_REPO_ENV_KEY] = repo_id + envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.set(repo_id) if config.getoption("--ssim-bootstrap-mode"): - os.environ[BOOTSTRAP_ENV_KEY] = "1" + # Kept for the whole pytest session and restored when pytest exits. + session_env = contextlib.ExitStack() + config.add_cleanup(session_env.close) + session_env.enter_context(envs.override_external("FASTVIDEO_SSIM_BOOTSTRAP_MODE", "1")) skip_download = config.getoption("--skip-ssim-reference-download") - skip_download = skip_download or os.environ.get("FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", - "").strip().lower() in {"1", "true", "yes", "on"} + skip_download = skip_download or envs.FASTVIDEO_TEST_SSIM_SKIP_REFERENCE_DOWNLOAD.get() if not skip_download: ensure_reference_videos_available( diff --git a/fastvideo/tests/ssim/inference_similarity_utils.py b/fastvideo/tests/ssim/inference_similarity_utils.py index 9ae9b0d46..68670816e 100644 --- a/fastvideo/tests/ssim/inference_similarity_utils.py +++ b/fastvideo/tests/ssim/inference_similarity_utils.py @@ -7,6 +7,7 @@ from collections.abc import Iterator from contextlib import contextmanager from logging import Logger +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.tests.ssim.bootstrap_references import ( xfail_missing_reference_in_bootstrap_mode, ) @@ -35,15 +36,8 @@ DEVICE_MAPPINGS = ( @contextmanager def attention_backend(backend: str) -> Iterator[None]: - previous = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend - try: + with envs.FASTVIDEO_ATTENTION_BACKEND.override(backend): yield - finally: - if previous is None: - os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) - else: - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous def shutdown_executor(generator: VideoGenerator | None) -> None: diff --git a/fastvideo/tests/ssim/latent_similarity_utils.py b/fastvideo/tests/ssim/latent_similarity_utils.py index 9dbb15670..abe0c71e9 100644 --- a/fastvideo/tests/ssim/latent_similarity_utils.py +++ b/fastvideo/tests/ssim/latent_similarity_utils.py @@ -188,7 +188,7 @@ def load_latent_reference(path: str) -> dict[str, Any]: # ``weights_only=False`` is required because the payload is a dict of # tensors + plain-Python metadata (slice_spec, prompt, ...). The trust # boundary is the controlled HF dataset configured via - # FASTVIDEO_SSIM_REFERENCE_HF_REPO (default + # FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO (default # FastVideo/ssim-reference-videos), which is org-write-gated. payload = torch.load(path, map_location="cpu", weights_only=False) fmt = payload.get("format_version") if isinstance(payload, dict) else None diff --git a/fastvideo/tests/ssim/reference_utils.py b/fastvideo/tests/ssim/reference_utils.py index e7390fadb..33e2a24c5 100644 --- a/fastvideo/tests/ssim/reference_utils.py +++ b/fastvideo/tests/ssim/reference_utils.py @@ -9,8 +9,9 @@ from typing import TypeVar import torch +import fastvideo.envs as envs + DEVICE_REFERENCE_FOLDER_SUFFIX = "_reference_videos" -FULL_QUALITY_ENV_VAR = "FASTVIDEO_SSIM_FULL_QUALITY" DEFAULT_OUTPUT_QUALITY_TIER = "default" FULL_OUTPUT_QUALITY_TIER = "full_quality" REFERENCE_VIDEOS_DIRNAME = "reference_videos" @@ -96,8 +97,7 @@ def build_reference_folder_path( def use_full_quality_configs() -> bool: - value = os.environ.get(FULL_QUALITY_ENV_VAR, "").strip().lower() - return value in {"1", "true", "yes", "on"} + return envs.FASTVIDEO_TEST_SSIM_FULL_QUALITY.get() def get_output_quality_tier() -> str: @@ -126,3 +126,8 @@ def select_ssim_params(default_params: T, full_quality_params: T) -> T: if use_full_quality_configs(): return full_quality_params return default_params + + +def with_model_path(params_map: dict[str, dict[str, object]], model_path: str) -> dict[str, dict[str, object]]: + """Return a copy of ``params_map`` in which every entry loads ``model_path``.""" + return {model_id: {**params, "model_path": model_path} for model_id, params in params_map.items()} diff --git a/fastvideo/tests/ssim/reference_videos_cli.py b/fastvideo/tests/ssim/reference_videos_cli.py index 4c71c46f2..02d8cacf9 100644 --- a/fastvideo/tests/ssim/reference_videos_cli.py +++ b/fastvideo/tests/ssim/reference_videos_cli.py @@ -4,12 +4,32 @@ from __future__ import annotations import argparse import contextlib +import importlib.util import os import shutil import sys import tempfile from collections.abc import Iterable, Sequence from pathlib import Path +from types import ModuleType + + +def _load_envs_registry() -> ModuleType: + """Load fastvideo/envs.py as a standalone module, without importing the fastvideo package. + + Importing fastvideo.envs runs fastvideo/__init__.py, which imports torch and the + generation stack. This CLI runs as a plain script that needs only the standard + library, so it loads the registry file by path. + """ + registry_path = Path(__file__).resolve().parents[2] / "envs.py" + spec = importlib.util.spec_from_file_location("_fastvideo_envs_registry", registry_path) + assert spec is not None and spec.loader is not None + registry = importlib.util.module_from_spec(spec) + spec.loader.exec_module(registry) + return registry + + +envs = _load_envs_registry() VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv") # Additional artefact types stored under the same reference folders. Latent @@ -19,9 +39,8 @@ VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv") LATENT_EXTENSIONS = (".pt", ) REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png", ) HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN") -HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO" -HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE" -BOOTSTRAP_ENV_KEY = "FASTVIDEO_SSIM_BOOTSTRAP_MODE" +HF_REPO_ENV_KEY = "FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO" +HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE" DEFAULT_REPO_ID = "FastVideo/ssim-reference-videos" DEFAULT_REPO_TYPE = "dataset" @@ -42,11 +61,11 @@ def _ssim_dir() -> Path: def _default_repo_id() -> str: - return os.environ.get(HF_REPO_ENV_KEY, DEFAULT_REPO_ID) + return envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.get() def _default_repo_type() -> str: - return os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE) + return envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE.get() def _iter_reference_files(root: Path) -> Iterable[Path]: diff --git a/fastvideo/tests/ssim/test_dreamx_world_similarity.py b/fastvideo/tests/ssim/test_dreamx_world_similarity.py index 24f5ffd60..e8648382e 100644 --- a/fastvideo/tests/ssim/test_dreamx_world_similarity.py +++ b/fastvideo/tests/ssim/test_dreamx_world_similarity.py @@ -7,12 +7,14 @@ from pathlib import Path import pytest from PIL import Image, ImageDraw +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( resolve_inference_device_reference_folder, run_image_to_video_similarity_test, run_text_to_video_similarity_test, ) +from fastvideo.tests.ssim.reference_utils import with_model_path logger = init_logger(__name__) @@ -21,18 +23,10 @@ pytestmark = pytest.mark.skip(reason="Disabled pending removal of DreamX World s device_reference_folder = resolve_inference_device_reference_folder(logger) -_MODEL_PATH = os.getenv( - "DREAMX_WORLD_SSIM_MODEL_PATH", - "FastVideo/DreamX-World-5B-Cam-Diffusers", -) -_AR_MODEL_PATH = os.getenv( - "DREAMX_WORLD_AR_SSIM_MODEL_PATH", - "FastVideo/DreamX-World-5B-Diffusers", -) - +# The tests add "model_path" from FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH and +# FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH. DREAMX_WORLD_PARAMS = { "num_gpus": 1, - "model_path": _MODEL_PATH, "height": 64, "width": 64, "num_frames": 9, @@ -53,7 +47,6 @@ DREAMX_WORLD_FULL_QUALITY_PARAMS = { DREAMX_WORLD_AR_PARAMS = { "num_gpus": 1, - "model_path": _AR_MODEL_PATH, "height": 192, "width": 192, "num_frames": 81, @@ -129,6 +122,7 @@ def test_dreamx_world_inference_similarity( ) -> None: image_path = tmp_path / "dreamx_world_ssim_input.png" _write_deterministic_reference_image(image_path) + model_path = envs.FASTVIDEO_TEST_DREAMX_WORLD_SSIM_MODEL_PATH.get() run_image_to_video_similarity_test( logger=logger, @@ -138,8 +132,8 @@ def test_dreamx_world_inference_similarity( image_path=str(image_path), attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=DREAMX_WORLD_MODEL_TO_PARAMS, - full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS, + default_params_map=with_model_path(DREAMX_WORLD_MODEL_TO_PARAMS, model_path), + full_quality_params_map=with_model_path(FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS, model_path), min_acceptable_ssim=0.98, init_kwargs_override={ "use_fsdp_inference": False, @@ -166,6 +160,7 @@ def test_dreamx_world_ar_inference_similarity( attention_backend_name: str, model_id: str, ) -> None: + model_path = envs.FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH.get() run_text_to_video_similarity_test( logger=logger, script_dir=os.path.dirname(os.path.abspath(__file__)), @@ -173,8 +168,8 @@ def test_dreamx_world_ar_inference_similarity( prompt=prompt, attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=DREAMX_WORLD_AR_MODEL_TO_PARAMS, - full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS, + default_params_map=with_model_path(DREAMX_WORLD_AR_MODEL_TO_PARAMS, model_path), + full_quality_params_map=with_model_path(FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS, model_path), min_acceptable_ssim=0.98, init_kwargs_override={ "use_fsdp_inference": False, diff --git a/fastvideo/tests/ssim/test_flux_t2i_similarity.py b/fastvideo/tests/ssim/test_flux_t2i_similarity.py index 5d78feab4..4205f610e 100644 --- a/fastvideo/tests/ssim/test_flux_t2i_similarity.py +++ b/fastvideo/tests/ssim/test_flux_t2i_similarity.py @@ -7,6 +7,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo.api.flux import FluxSamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( @@ -14,6 +15,7 @@ from fastvideo.tests.ssim.inference_similarity_utils import ( from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, + with_model_path, ) logger = init_logger(__name__) @@ -23,11 +25,6 @@ REQUIRED_GPUS = 1 # MS-SSIM gate (see module docstring). FLUX_T2I_MIN_SSIM = 0.98 -FLUX_MODEL_PATH = os.getenv( - "FLUX_T2I_MODEL_DIR", - "black-forest-labs/FLUX.1-dev", -) - device_reference_folder = resolve_device_reference_folder( ( ("A40", "A40"), @@ -51,9 +48,9 @@ TEST_PROMPTS = [ "a photo of a cat", ] +# The test adds "model_path" from FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR. FLUX_DEFAULT_PARAMS: dict[str, object] = { "num_gpus": 1, - "model_path": FLUX_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": 256, @@ -68,7 +65,6 @@ FLUX_DEFAULT_PARAMS: dict[str, object] = { _flux_full_defaults = FluxSamplingParam() FLUX_FULL_QUALITY_PARAMS: dict[str, object] = { "num_gpus": 1, - "model_path": FLUX_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": _flux_full_defaults.height, @@ -100,10 +96,11 @@ def test_flux_t2i_similarity( attention_backend_name: str, model_id: str, ) -> None: - is_hf_repo = "/" in FLUX_MODEL_PATH and not FLUX_MODEL_PATH.startswith("/") - if not is_hf_repo and not os.path.isdir(FLUX_MODEL_PATH): - pytest.skip(f"FLUX weights not found at {FLUX_MODEL_PATH} " - f"(set FLUX_T2I_MODEL_DIR to override)") + flux_model_path = envs.FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR.get() + is_hf_repo = "/" in flux_model_path and not flux_model_path.startswith("/") + if not is_hf_repo and not os.path.isdir(flux_model_path): + pytest.skip(f"FLUX weights not found at {flux_model_path} " + f"(set FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR to override)") run_text_to_video_similarity_test( logger=logger, @@ -112,8 +109,8 @@ def test_flux_t2i_similarity( prompt=prompt, attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=FLUX_MODEL_TO_PARAMS, - full_quality_params_map=FLUX_FULL_QUALITY_MODEL_TO_PARAMS, + default_params_map=with_model_path(FLUX_MODEL_TO_PARAMS, flux_model_path), + full_quality_params_map=with_model_path(FLUX_FULL_QUALITY_MODEL_TO_PARAMS, flux_model_path), min_acceptable_ssim=FLUX_T2I_MIN_SSIM, media_extension=".png", init_kwargs_override={ diff --git a/fastvideo/tests/ssim/test_gamecraft_similarity.py b/fastvideo/tests/ssim/test_gamecraft_similarity.py index 97d1e4934..e415eaaea 100644 --- a/fastvideo/tests/ssim/test_gamecraft_similarity.py +++ b/fastvideo/tests/ssim/test_gamecraft_similarity.py @@ -17,6 +17,7 @@ import os import torch import pytest +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger @@ -26,6 +27,7 @@ from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, select_ssim_params, + with_model_path, ) from fastvideo.tests.utils import ( compute_video_ssim_torchvision, @@ -71,15 +73,10 @@ def _shutdown_executor(generator: VideoGenerator | None) -> None: # --------------------------------------------------------------------------- # Model parameters # --------------------------------------------------------------------------- -# Same as basic_gamecraft.py: default HF path; set GAMECRAFT_MODEL_PATH for local weights. -_GAMECRAFT_MODEL_PATH = os.environ.get( - "GAMECRAFT_MODEL_PATH", - "FastVideo/HunyuanGameCraft-Diffusers", -) - +# Same as basic_gamecraft.py: default HF path; set FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH for local weights. +# The tests add "model_path" from FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH. GAMECRAFT_T2V_PARAMS = { "num_gpus": 1, - "model_path": _GAMECRAFT_MODEL_PATH, "height": 480, "width": 832, "num_frames": 33, @@ -91,44 +88,46 @@ GAMECRAFT_T2V_PARAMS = { "negative_prompt": "", } -_GAMECRAFT_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(_GAMECRAFT_MODEL_PATH) -GAMECRAFT_T2V_FULL_QUALITY_PARAMS = { - "num_gpus": GAMECRAFT_T2V_PARAMS["num_gpus"], - "model_path": GAMECRAFT_T2V_PARAMS["model_path"], - "height": _GAMECRAFT_FULL_QUALITY_DEFAULTS.height, - "width": _GAMECRAFT_FULL_QUALITY_DEFAULTS.width, - "num_frames": GAMECRAFT_T2V_PARAMS["num_frames"], # default num_frames: 33 - "num_inference_steps": _GAMECRAFT_FULL_QUALITY_DEFAULTS.num_inference_steps, - "guidance_scale": _GAMECRAFT_FULL_QUALITY_DEFAULTS.guidance_scale, - "seed": _GAMECRAFT_FULL_QUALITY_DEFAULTS.seed, - "action": GAMECRAFT_T2V_PARAMS["action"], - "action_speed": GAMECRAFT_T2V_PARAMS["action_speed"], - "negative_prompt": _GAMECRAFT_FULL_QUALITY_DEFAULTS.negative_prompt, -} + +def _gamecraft_t2v_full_quality_params(model_path: str) -> dict[str, object]: + """Return the full-quality T2V params, taken from the default sampling params of ``model_path``.""" + full_quality_defaults = SamplingParam.from_pretrained(model_path) + return { + "num_gpus": GAMECRAFT_T2V_PARAMS["num_gpus"], + "model_path": model_path, + "height": full_quality_defaults.height, + "width": full_quality_defaults.width, + "num_frames": GAMECRAFT_T2V_PARAMS["num_frames"], # default num_frames: 33 + "num_inference_steps": full_quality_defaults.num_inference_steps, + "guidance_scale": full_quality_defaults.guidance_scale, + "seed": full_quality_defaults.seed, + "action": GAMECRAFT_T2V_PARAMS["action"], + "action_speed": GAMECRAFT_T2V_PARAMS["action_speed"], + "negative_prompt": full_quality_defaults.negative_prompt, + } + GAMECRAFT_I2V_PARAMS = { **GAMECRAFT_T2V_PARAMS, "image_path": ("https://huggingface.co/datasets/huggingface/documentation-images/" "resolve/main/diffusers/astronaut.jpg"), } -GAMECRAFT_I2V_FULL_QUALITY_PARAMS = { - **GAMECRAFT_T2V_FULL_QUALITY_PARAMS, - "image_path": GAMECRAFT_I2V_PARAMS["image_path"], -} + + +def _gamecraft_i2v_full_quality_params(model_path: str) -> dict[str, object]: + return { + **_gamecraft_t2v_full_quality_params(model_path), + "image_path": GAMECRAFT_I2V_PARAMS["image_path"], + } + MODEL_TO_PARAMS = { "HunyuanGameCraft-T2V": GAMECRAFT_T2V_PARAMS, } -FULL_QUALITY_MODEL_TO_PARAMS = { - "HunyuanGameCraft-T2V": GAMECRAFT_T2V_FULL_QUALITY_PARAMS, -} I2V_MODEL_TO_PARAMS = { "HunyuanGameCraft-I2V": GAMECRAFT_I2V_PARAMS, } -FULL_QUALITY_I2V_MODEL_TO_PARAMS = { - "HunyuanGameCraft-I2V": GAMECRAFT_I2V_FULL_QUALITY_PARAMS, -} TEST_PROMPTS = [ "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.", @@ -146,9 +145,9 @@ I2V_TEST_PROMPTS = [ @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): +def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): """Generate a T2V video with GameCraft and compare to reference via SSIM.""" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) output_dir = build_generated_output_dir( @@ -161,9 +160,10 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): os.makedirs(output_dir, exist_ok=True) + model_path = envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() params_map = select_ssim_params( - MODEL_TO_PARAMS, - FULL_QUALITY_MODEL_TO_PARAMS, + with_model_path(MODEL_TO_PARAMS, model_path), + {model_id: _gamecraft_t2v_full_quality_params(model_path)}, ) BASE_PARAMS = params_map[model_id] num_inference_steps = BASE_PARAMS["num_inference_steps"] @@ -268,9 +268,9 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id): @pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) @pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys())) -def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): +def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): """Generate an I2V video with GameCraft and compare to reference via SSIM.""" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) output_dir = build_generated_output_dir( @@ -283,9 +283,10 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id): os.makedirs(output_dir, exist_ok=True) + model_path = envs.FASTVIDEO_TEST_GAMECRAFT_MODEL_PATH.get() params_map = select_ssim_params( - I2V_MODEL_TO_PARAMS, - FULL_QUALITY_I2V_MODEL_TO_PARAMS, + with_model_path(I2V_MODEL_TO_PARAMS, model_path), + {model_id: _gamecraft_i2v_full_quality_params(model_path)}, ) BASE_PARAMS = params_map[model_id] num_inference_steps = BASE_PARAMS["num_inference_steps"] diff --git a/fastvideo/tests/ssim/test_gen3c_similarity.py b/fastvideo/tests/ssim/test_gen3c_similarity.py index 07f41761c..a264efd61 100644 --- a/fastvideo/tests/ssim/test_gen3c_similarity.py +++ b/fastvideo/tests/ssim/test_gen3c_similarity.py @@ -10,9 +10,9 @@ Usage: pytest fastvideo/tests/ssim/test_gen3c_similarity.py -v Environment variables: - GEN3C_MODEL_PATH - Diffusers-format GEN3C model path/repo id. - Default: FastVideo/GEN3C-Cosmos-7B-Diffusers - (local converted path also supported) + FASTVIDEO_TEST_GEN3C_MODEL_PATH - Diffusers-format GEN3C model path/repo id. + Default: FastVideo/GEN3C-Cosmos-7B-Diffusers + (local converted path also supported) """ import os @@ -22,6 +22,7 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.logger import init_logger from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results @@ -35,10 +36,10 @@ def _resolve_gen3c_test_image_path() -> str: Resolve image path for GEN3C I2V SSIM tests. Priority: - 1) GEN3C_TEST_IMAGE_PATH env var + 1) FASTVIDEO_TEST_GEN3C_IMAGE_PATH env var 2) Repo asset image """ - env_image = os.getenv("GEN3C_TEST_IMAGE_PATH") + env_image = envs.FASTVIDEO_TEST_GEN3C_IMAGE_PATH.get() if env_image: return env_image @@ -69,9 +70,10 @@ else: # GEN3C generation parameters # --------------------------------------------------------------------------- +# The test reads the model path from FASTVIDEO_TEST_GEN3C_MODEL_PATH and the input image from +# _resolve_gen3c_test_image_path(). GEN3C_T2V_PARAMS = { "num_gpus": 1, - "model_path": os.getenv("GEN3C_MODEL_PATH", "FastVideo/GEN3C-Cosmos-7B-Diffusers"), "height": 720, "width": 1280, "num_frames": 121, @@ -80,7 +82,6 @@ GEN3C_T2V_PARAMS = { "embedded_cfg_scale": 6, "flow_shift": 1.0, "seed": 1024, - "image_path": _resolve_gen3c_test_image_path(), "sp_size": 1, "tp_size": 1, "fps": 24, @@ -110,11 +111,11 @@ CANDIDATE_VIDEO_NAME = "gen3c_ssim_candidate.mp4" @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): +def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): """ Generate a GEN3C video and compare against the reference using MS-SSIM. """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) base_output_dir = os.path.join(script_dir, "generated_videos", model_id) @@ -124,22 +125,24 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): BASE_PARAMS = MODEL_TO_PARAMS[model_id] num_inference_steps = BASE_PARAMS["num_inference_steps"] - model_path = BASE_PARAMS["model_path"] + model_path = envs.FASTVIDEO_TEST_GEN3C_MODEL_PATH.get() # Guard common misconfigurations to keep CI behavior explicit. if model_path.lower() == "nvidia/gen3c-cosmos-7b": pytest.skip("nvidia/GEN3C-Cosmos-7B is the official raw checkpoint repo, not Diffusers format. " - "Use GEN3C_MODEL_PATH=FastVideo/GEN3C-Cosmos-7B-Diffusers or a local converted path.") + "Use FASTVIDEO_TEST_GEN3C_MODEL_PATH=FastVideo/GEN3C-Cosmos-7B-Diffusers " + "or a local converted path.") local_like = model_path.startswith(("/", "./", "../")) if local_like and not os.path.exists(model_path): pytest.skip(f"Local GEN3C model path not found: {model_path}. " - "Set GEN3C_MODEL_PATH to a valid local path or HF Diffusers repo id.") + "Set FASTVIDEO_TEST_GEN3C_MODEL_PATH to a valid local path or HF Diffusers repo id.") if os.path.exists(model_path): model_index_path = os.path.join(model_path, "model_index.json") if not os.path.exists(model_index_path): - pytest.skip(f"GEN3C_MODEL_PATH is not Diffusers-format (missing model_index.json): {model_path}") + pytest.skip("FASTVIDEO_TEST_GEN3C_MODEL_PATH is not Diffusers-format " + f"(missing model_index.json): {model_path}") init_kwargs = { "num_gpus": BASE_PARAMS["num_gpus"], @@ -158,13 +161,13 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id): "guidance_scale": BASE_PARAMS["guidance_scale"], "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"], "seed": BASE_PARAMS["seed"], - "image_path": BASE_PARAMS["image_path"], + "image_path": _resolve_gen3c_test_image_path(), "fps": BASE_PARAMS["fps"], } if not os.path.exists(generation_kwargs["image_path"]): pytest.skip(f"GEN3C test image not found: {generation_kwargs['image_path']}. " - "Set GEN3C_TEST_IMAGE_PATH to a valid local image.") + "Set FASTVIDEO_TEST_GEN3C_IMAGE_PATH to a valid local image.") # Keep local reruns deterministic: remove prior candidate outputs so # VideoGenerator does not auto-suffix (_1, _2, ...). diff --git a/fastvideo/tests/ssim/test_glm_image_similarity.py b/fastvideo/tests/ssim/test_glm_image_similarity.py index feaaa3e12..31834733e 100644 --- a/fastvideo/tests/ssim/test_glm_image_similarity.py +++ b/fastvideo/tests/ssim/test_glm_image_similarity.py @@ -8,12 +8,14 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( run_text_to_video_similarity_test, ) from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, + with_model_path, ) logger = init_logger(__name__) @@ -22,8 +24,6 @@ REQUIRED_GPUS = 1 pytestmark = pytest.mark.skip(reason="Disabled pending removal of GLM-Image support.") REPO_ROOT = Path(__file__).resolve().parents[3] -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "glm_image")) -GLM_IMAGE_MODEL_PATH = os.getenv("GLM_IMAGE_MODEL_DIR", str(LOCAL_WEIGHTS_DIR)) device_reference_folder = resolve_device_reference_folder( ( @@ -50,9 +50,9 @@ TEST_PROMPTS = [ "Warm golden light, photorealistic style.", ] +# The test adds "model_path" from FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR. GLM_IMAGE_PARAMS = { "num_gpus": 1, - "model_path": GLM_IMAGE_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": 256, @@ -67,7 +67,6 @@ GLM_IMAGE_PARAMS = { GLM_IMAGE_FULL_QUALITY_PARAMS = { "num_gpus": 1, - "model_path": GLM_IMAGE_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": 1024, @@ -88,9 +87,13 @@ GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS = { } +def _local_weights_dir() -> Path: + return Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or REPO_ROOT / "official_weights" / "glm_image") + + def _has_weights() -> bool: required = ["transformer", "vae", "text_encoder", "vision_language_encoder", "processor", "tokenizer", "scheduler"] - return all((LOCAL_WEIGHTS_DIR / r).exists() for r in required) + return all((_local_weights_dir() / r).exists() for r in required) def _upstream_glm_image_available() -> bool: @@ -105,10 +108,6 @@ def _upstream_glm_image_available() -> bool: not torch.cuda.is_available(), reason="GLM-Image SSIM test requires CUDA", ) -@pytest.mark.skipif( - not _has_weights(), - reason=f"GLM-Image full weights not found at {LOCAL_WEIGHTS_DIR}.", -) @pytest.mark.skipif( not _upstream_glm_image_available(), reason="GLM-Image needs transformers>=5.0.0rc0 (ships the AR encoder).", @@ -121,6 +120,10 @@ def test_glm_image_similarity( attention_backend_name: str, model_id: str, ) -> None: + if not _has_weights(): + pytest.skip(f"GLM-Image full weights not found at {_local_weights_dir()}.") + model_path = envs.FASTVIDEO_TEST_GLM_IMAGE_MODEL_DIR.get() or str(_local_weights_dir()) + run_text_to_video_similarity_test( logger=logger, script_dir=os.path.dirname(os.path.abspath(__file__)), @@ -128,8 +131,8 @@ def test_glm_image_similarity( prompt=prompt, attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=GLM_IMAGE_MODEL_TO_PARAMS, - full_quality_params_map=GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS, + default_params_map=with_model_path(GLM_IMAGE_MODEL_TO_PARAMS, model_path), + full_quality_params_map=with_model_path(GLM_IMAGE_FULL_QUALITY_MODEL_TO_PARAMS, model_path), min_acceptable_ssim=0.98, init_kwargs_override={ "trust_remote_code": True, diff --git a/fastvideo/tests/ssim/test_kandinsky5_similarity.py b/fastvideo/tests/ssim/test_kandinsky5_similarity.py index 36641c66a..a0944feb7 100644 --- a/fastvideo/tests/ssim/test_kandinsky5_similarity.py +++ b/fastvideo/tests/ssim/test_kandinsky5_similarity.py @@ -11,6 +11,7 @@ import os import pytest +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.pipelines.basic.kandinsky5.presets import KANDINSKY5_T2V_LITE_5S from fastvideo.tests.ssim.inference_similarity_utils import ( @@ -80,9 +81,7 @@ def test_kandinsky5_t2v_inference_similarity( # and restored afterwards: the other SSIM references are seeded with FA4 # on, so a module-level pin would corrupt every test collected in the # same pytest process. - saved_fa4 = os.environ.get("FASTVIDEO_FA4") - os.environ["FASTVIDEO_FA4"] = "0" - try: + with envs.FASTVIDEO_FA4.override(False): run_text_to_video_similarity_test( logger=logger, script_dir=os.path.dirname(os.path.abspath(__file__)), @@ -102,8 +101,3 @@ def test_kandinsky5_t2v_inference_similarity( "pin_cpu_memory": True, }, ) - finally: - if saved_fa4 is None: - os.environ.pop("FASTVIDEO_FA4", None) - else: - os.environ["FASTVIDEO_FA4"] = saved_fa4 diff --git a/fastvideo/tests/ssim/test_lingbot_similarity.py b/fastvideo/tests/ssim/test_lingbot_similarity.py index a0163d839..2324b14ab 100644 --- a/fastvideo/tests/ssim/test_lingbot_similarity.py +++ b/fastvideo/tests/ssim/test_lingbot_similarity.py @@ -13,6 +13,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger @@ -132,8 +133,8 @@ TEST_PROMPTS = [ @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND +def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) params = select_ssim_params(LINGBOT_PARAMS, LINGBOT_FULL_QUALITY_PARAMS) diff --git a/fastvideo/tests/ssim/test_longcat_similarity.py b/fastvideo/tests/ssim/test_longcat_similarity.py index beefe40cd..c68eddb6e 100644 --- a/fastvideo/tests/ssim/test_longcat_similarity.py +++ b/fastvideo/tests/ssim/test_longcat_similarity.py @@ -19,6 +19,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger @@ -192,13 +193,13 @@ def _resolve_asset_path(asset_path: str) -> str: @pytest.mark.parametrize("prompt", T2V_TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): +def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): """ Test LongCat T2V inference and compare output to reference videos using SSIM. Parameters derived from examples/inference/basic/basic_longcat_t2v.py """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) params = select_ssim_params(LONGCAT_T2V_PARAMS, LONGCAT_T2V_FULL_QUALITY_PARAMS) @@ -280,13 +281,13 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): @pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): +def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): """ Test LongCat I2V inference and compare output to reference videos using SSIM. Parameters derived from examples/inference/basic/basic_longcat_i2v.py """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) params = select_ssim_params(LONGCAT_I2V_PARAMS, LONGCAT_I2V_FULL_QUALITY_PARAMS) @@ -373,13 +374,13 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): @pytest.mark.parametrize("prompt", VC_TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): +def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): """ Test LongCat VC (Video Continuation) inference and compare output to reference videos using SSIM. Parameters derived from examples/inference/basic/basic_longcat_vc.py """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) params = select_ssim_params(LONGCAT_VC_PARAMS, LONGCAT_VC_FULL_QUALITY_PARAMS) diff --git a/fastvideo/tests/ssim/test_matrixgame2_similarity.py b/fastvideo/tests/ssim/test_matrixgame2_similarity.py index 9f9dc002d..6f03594db 100644 --- a/fastvideo/tests/ssim/test_matrixgame2_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame2_similarity.py @@ -4,6 +4,7 @@ import os import torch import pytest +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger @@ -83,12 +84,12 @@ TEST_IMAGE_PATHS = [ @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_matrixgame2_similarity(prompt, ATTENTION_BACKEND, model_id): +def test_matrixgame2_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): """ Test that runs inference with different parameters and compares the output to reference videos using SSIM. """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/fastvideo/tests/ssim/test_matrixgame3_similarity.py b/fastvideo/tests/ssim/test_matrixgame3_similarity.py index b383698b1..256f431cc 100644 --- a/fastvideo/tests/ssim/test_matrixgame3_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame3_similarity.py @@ -3,6 +3,7 @@ import os import pytest +import fastvideo.envs as envs from fastvideo import VideoGenerator from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger @@ -78,12 +79,12 @@ TEST_IMAGE_PATHS = [ @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_matrixgame3_similarity(prompt, ATTENTION_BACKEND, model_id): +def test_matrixgame3_similarity(prompt, ATTENTION_BACKEND, model_id, env_overrides): """ Test that runs MG3 inference (action conditions auto-generated from seed) and compares the output to reference videos using SSIM. """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) script_dir = os.path.dirname(os.path.abspath(__file__)) diff --git a/fastvideo/tests/ssim/test_sd35_similarity.py b/fastvideo/tests/ssim/test_sd35_similarity.py index ba64fa029..d11b3f41b 100644 --- a/fastvideo/tests/ssim/test_sd35_similarity.py +++ b/fastvideo/tests/ssim/test_sd35_similarity.py @@ -6,6 +6,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo.api.sampling_param import SamplingParam from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( @@ -13,17 +14,13 @@ from fastvideo.tests.ssim.inference_similarity_utils import ( from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, + with_model_path, ) logger = init_logger(__name__) REQUIRED_GPUS = 1 -SD35_MODEL_PATH = os.getenv( - "SD35_MODEL_DIR", - "stabilityai/stable-diffusion-3.5-medium", -) - device_reference_folder = resolve_device_reference_folder( ( ("A40", "A40"), @@ -45,9 +42,9 @@ TEST_PROMPTS = [ "a photo of a cat", ] +# The test adds "model_path" from FASTVIDEO_TEST_SD35_MODEL_DIR. SD35_PARAMS = { "num_gpus": 1, - "model_path": SD35_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": 256, @@ -60,28 +57,29 @@ SD35_PARAMS = { "neg_prompt": "", } -_SD35_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(SD35_MODEL_PATH) -SD35_FULL_QUALITY_PARAMS = { - "num_gpus": 1, - "model_path": SD35_MODEL_PATH, - "sp_size": 1, - "tp_size": 1, - "height": _SD35_FULL_QUALITY_DEFAULTS.height, - "width": _SD35_FULL_QUALITY_DEFAULTS.width, - "num_frames": SD35_PARAMS["num_frames"], - "fps": _SD35_FULL_QUALITY_DEFAULTS.fps, - "num_inference_steps": (_SD35_FULL_QUALITY_DEFAULTS.num_inference_steps), - "guidance_scale": _SD35_FULL_QUALITY_DEFAULTS.guidance_scale, - "seed": _SD35_FULL_QUALITY_DEFAULTS.seed, - "neg_prompt": _SD35_FULL_QUALITY_DEFAULTS.negative_prompt, -} + +def _sd35_full_quality_params(model_path: str) -> dict[str, object]: + """Return the full-quality SD3.5 params, taken from the default sampling params of ``model_path``.""" + full_quality_defaults = SamplingParam.from_pretrained(model_path) + return { + "num_gpus": 1, + "model_path": model_path, + "sp_size": 1, + "tp_size": 1, + "height": full_quality_defaults.height, + "width": full_quality_defaults.width, + "num_frames": SD35_PARAMS["num_frames"], + "fps": full_quality_defaults.fps, + "num_inference_steps": (full_quality_defaults.num_inference_steps), + "guidance_scale": full_quality_defaults.guidance_scale, + "seed": full_quality_defaults.seed, + "neg_prompt": full_quality_defaults.negative_prompt, + } + SD35_MODEL_TO_PARAMS = { MODEL_ID: SD35_PARAMS, } -SD35_FULL_QUALITY_MODEL_TO_PARAMS = { - MODEL_ID: SD35_FULL_QUALITY_PARAMS, -} @pytest.mark.skipif( @@ -96,9 +94,10 @@ def test_sd35_similarity( attention_backend_name: str, model_id: str, ) -> None: - is_hf_repo = "/" in SD35_MODEL_PATH and not SD35_MODEL_PATH.startswith("/") - if not is_hf_repo and not os.path.isdir(SD35_MODEL_PATH): - pytest.skip(f"SD3.5 weights not found at {SD35_MODEL_PATH} (set SD35_MODEL_DIR to override)") + sd35_model_path = envs.FASTVIDEO_TEST_SD35_MODEL_DIR.get() + is_hf_repo = "/" in sd35_model_path and not sd35_model_path.startswith("/") + if not is_hf_repo and not os.path.isdir(sd35_model_path): + pytest.skip(f"SD3.5 weights not found at {sd35_model_path} (set FASTVIDEO_TEST_SD35_MODEL_DIR to override)") run_text_to_video_similarity_test( logger=logger, @@ -107,8 +106,8 @@ def test_sd35_similarity( prompt=prompt, attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=SD35_MODEL_TO_PARAMS, - full_quality_params_map=SD35_FULL_QUALITY_MODEL_TO_PARAMS, + default_params_map=with_model_path(SD35_MODEL_TO_PARAMS, sd35_model_path), + full_quality_params_map={model_id: _sd35_full_quality_params(sd35_model_path)}, min_acceptable_ssim=0.98, init_kwargs_override={ "workload_type": "t2v", diff --git a/fastvideo/tests/ssim/test_zimage_similarity.py b/fastvideo/tests/ssim/test_zimage_similarity.py index aaf488fd1..f739e9cbf 100644 --- a/fastvideo/tests/ssim/test_zimage_similarity.py +++ b/fastvideo/tests/ssim/test_zimage_similarity.py @@ -8,12 +8,14 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo.logger import init_logger from fastvideo.tests.ssim.inference_similarity_utils import ( run_text_to_video_similarity_test, ) from fastvideo.tests.ssim.reference_utils import ( get_cuda_device_name, resolve_device_reference_folder, + with_model_path, ) logger = init_logger(__name__) @@ -22,15 +24,6 @@ REQUIRED_GPUS = 1 pytestmark = pytest.mark.skip(reason="Disabled pending removal of Z-Image support.") ZIMAGE_MIN_SSIM = 0.98 -ZIMAGE_MODEL_PATH = os.getenv( - "ZIMAGE_MODEL_DIR", - "Tongyi-MAI/Z-Image-Turbo", -) -ZIMAGE_MODEL_REVISION = os.getenv( - "ZIMAGE_MODEL_REVISION", - "f332072aa78be7aecdf3ee76d5c247082da564a6", -) - TEST_PROMPTS = [ "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. " "Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. " @@ -55,10 +48,10 @@ device_reference_folder = resolve_device_reference_folder( logger=logger, ) +# The test adds "model_path" from FASTVIDEO_TEST_ZIMAGE_MODEL_DIR. ZIMAGE_MODEL_TO_PARAMS: dict[str, dict[str, object]] = { "Tongyi-MAI__Z-Image-Turbo": { "num_gpus": 1, - "model_path": ZIMAGE_MODEL_PATH, "sp_size": 1, "tp_size": 1, "height": 1024, @@ -86,10 +79,11 @@ def test_zimage_similarity( attention_backend_name: str, model_id: str, ) -> None: - is_hf_repo = "/" in ZIMAGE_MODEL_PATH and not ZIMAGE_MODEL_PATH.startswith("/") - if not is_hf_repo and not os.path.isdir(ZIMAGE_MODEL_PATH): - pytest.skip(f"Z-Image weights not found at {ZIMAGE_MODEL_PATH} " - "(set ZIMAGE_MODEL_DIR to override)") + zimage_model_path = envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() + is_hf_repo = "/" in zimage_model_path and not zimage_model_path.startswith("/") + if not is_hf_repo and not os.path.isdir(zimage_model_path): + pytest.skip(f"Z-Image weights not found at {zimage_model_path} " + "(set FASTVIDEO_TEST_ZIMAGE_MODEL_DIR to override)") run_text_to_video_similarity_test( logger=logger, @@ -98,12 +92,12 @@ def test_zimage_similarity( prompt=prompt, attention_backend_name=attention_backend_name, model_id=model_id, - default_params_map=ZIMAGE_MODEL_TO_PARAMS, - full_quality_params_map=ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS, + default_params_map=with_model_path(ZIMAGE_MODEL_TO_PARAMS, zimage_model_path), + full_quality_params_map=with_model_path(ZIMAGE_FULL_QUALITY_MODEL_TO_PARAMS, zimage_model_path), min_acceptable_ssim=ZIMAGE_MIN_SSIM, media_extension=".png", init_kwargs_override={ - "revision": ZIMAGE_MODEL_REVISION, + "revision": envs.FASTVIDEO_TEST_ZIMAGE_MODEL_REVISION.get(), "workload_type": "t2i", "trust_remote_code": True, "use_fsdp_inference": False, diff --git a/fastvideo/tests/stages/_denoising_fixtures.py b/fastvideo/tests/stages/_denoising_fixtures.py index 6ad590248..dbbde882a 100644 --- a/fastvideo/tests/stages/_denoising_fixtures.py +++ b/fastvideo/tests/stages/_denoising_fixtures.py @@ -6,6 +6,8 @@ from types import SimpleNamespace import torch +import fastvideo.envs as envs + class RecordingLogger: @@ -102,19 +104,19 @@ def _tiny_batch(): ) -def _patch_denoising_module(monkeypatch, cfg_gate_step): +def _patch_denoising_module(monkeypatch, env_overrides, cfg_gate_step): if cfg_gate_step is None: - monkeypatch.delenv("FASTVIDEO_CFG_GATE_STEP", raising=False) + env_overrides.enter_context(envs.FASTVIDEO_CFG_GATE_STEP.override(None)) expected_gate_step = 1.0 else: - monkeypatch.setenv("FASTVIDEO_CFG_GATE_STEP", str(cfg_gate_step)) + env_overrides.enter_context(envs.FASTVIDEO_CFG_GATE_STEP.override(float(cfg_gate_step))) expected_gate_step = float(cfg_gate_step) import fastvideo.pipelines.stages.denoising as denoising from fastvideo.pipelines.stages.base import PipelineStage # envs.FASTVIDEO_CFG_GATE_STEP.get() reads the environment on each call, so - # the stage sees monkeypatched values without reloading the module. + # the stage sees overridden values without reloading the module. assert denoising.envs.FASTVIDEO_CFG_GATE_STEP.get() == expected_gate_step logger = RecordingLogger() diff --git a/fastvideo/tests/stages/test_cfg_gating.py b/fastvideo/tests/stages/test_cfg_gating.py index 9e20d0a30..4ebb1b048 100644 --- a/fastvideo/tests/stages/test_cfg_gating.py +++ b/fastvideo/tests/stages/test_cfg_gating.py @@ -9,8 +9,8 @@ from fastvideo.tests.stages._denoising_fixtures import ( ) -def _run_stage(monkeypatch, cfg_gate_step): - denoising, logger = _patch_denoising_module(monkeypatch, cfg_gate_step) +def _run_stage(monkeypatch, env_overrides, cfg_gate_step): + denoising, logger = _patch_denoising_module(monkeypatch, env_overrides, cfg_gate_step) model = TinyDenoiser() stage = denoising.DenoisingStage(model, TinyScheduler()) stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() @@ -39,8 +39,8 @@ def _run_legacy_two_pass(): @pytest.mark.parametrize("cfg_gate_step", [None, "1.0"]) -def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, cfg_gate_step): - out, model, logger = _run_stage(monkeypatch, cfg_gate_step) +def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, env_overrides, cfg_gate_step): + out, model, logger = _run_stage(monkeypatch, env_overrides, cfg_gate_step) legacy_out = _run_legacy_two_pass() assert torch.equal(out, legacy_out) @@ -49,8 +49,8 @@ def test_cfg_gating_default_off_matches_legacy_two_pass(monkeypatch, cfg_gate_st assert any("gate_step=-1/4" in msg and "reused=0" in msg for msg in logger.infos) -def test_cfg_gating_reuses_cached_delta_after_gate(monkeypatch): - out, model, logger = _run_stage(monkeypatch, "0.5") +def test_cfg_gating_reuses_cached_delta_after_gate(monkeypatch, env_overrides): + out, model, logger = _run_stage(monkeypatch, env_overrides, "0.5") legacy_out = _run_legacy_two_pass() assert model.calls == ["cond", "uncond", "cond", "uncond", "cond", "cond"] diff --git a/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py b/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py index 075530022..19a102752 100644 --- a/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py +++ b/fastvideo/tests/stages/test_kandinsky5_dmd_stage_backend_engages.py @@ -30,6 +30,7 @@ import types import pytest import torch +import fastvideo.envs as envs from fastvideo.attention import LocalAttention from fastvideo.attention.selector import _cached_get_attn_backend from fastvideo.configs.models.dits.kandinsky5 import ( @@ -63,14 +64,14 @@ def _tiny_arch() -> Kandinsky5ArchConfig: ) -def _build_stage_and_spy(monkeypatch): +def _build_stage_and_spy(monkeypatch, env_overrides): """Tiny transformer + real DMD stage on CPU, with every resolved backend impl wrapped in a call counter.""" # Pin the backend so resolution is identical on CPU-only and CUDA # machines, and clear the process-wide selector cache keyed on # (head_size, dtype, supported_backends) -- NOT on the env var (same # defensive pattern as test_kandinsky5_qat_attention_engages.py). - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) _cached_get_attn_backend.cache_clear() from fastvideo.models.dits.kandinsky5 import Kandinsky5Transformer3DModel @@ -131,8 +132,8 @@ def _make_fastvideo_args(arch: Kandinsky5ArchConfig) -> types.SimpleNamespace: ) -def test_dmd_stage_invokes_selected_backend_every_step(monkeypatch): - stage, _, arch, backend_calls = _build_stage_and_spy(monkeypatch) +def test_dmd_stage_invokes_selected_backend_every_step(monkeypatch, env_overrides): + stage, _, arch, backend_calls = _build_stage_and_spy(monkeypatch, env_overrides) try: batch = _make_batch() @@ -152,11 +153,11 @@ def test_dmd_stage_invokes_selected_backend_every_step(monkeypatch): _cached_get_attn_backend.cache_clear() -def test_raw_transformer_without_context_bypasses_backend(monkeypatch): +def test_raw_transformer_without_context_bypasses_backend(monkeypatch, env_overrides): """Documents the trap the stage wrapper exists to prevent: without a forward context the transformer still returns output, but the selected backend is never invoked.""" - stage, transformer, _, backend_calls = _build_stage_and_spy(monkeypatch) + stage, transformer, _, backend_calls = _build_stage_and_spy(monkeypatch, env_overrides) del stage # only the spied transformer is needed here try: with torch.no_grad(): @@ -180,11 +181,11 @@ def test_raw_transformer_without_context_bypasses_backend(monkeypatch): _cached_get_attn_backend.cache_clear() -def test_dmd_stage_with_context_and_raw_without_share_one_spy(monkeypatch): +def test_dmd_stage_with_context_and_raw_without_share_one_spy(monkeypatch, env_overrides): """Same spy, both paths, one process: proves the counter difference between the two tests above is the forward-context wrapper itself, not some environmental difference.""" - stage, transformer, arch, backend_calls = _build_stage_and_spy(monkeypatch) + stage, transformer, arch, backend_calls = _build_stage_and_spy(monkeypatch, env_overrides) try: with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None): transformer( diff --git a/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py b/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py index bce953796..9ec4d6d05 100644 --- a/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py +++ b/fastvideo/tests/stages/test_minimax_h3_encoding_offload.py @@ -9,6 +9,7 @@ import pytest import torch from PIL import Image +import fastvideo.envs as envs from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig, MiniMaxH3AudioVAEConfig from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig from fastvideo.distributed import get_local_torch_device @@ -25,9 +26,9 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires @pytest.fixture -def stage(distributed_setup, monkeypatch): +def stage(distributed_setup, env_overrides): """Construct finite seeded VAEs with the tiny geometries used in parity tests.""" - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) torch.manual_seed(20260921) video_arch = MiniMaxH3VideoVAEArchConfig( latent_channels=4, diff --git a/fastvideo/tests/stages/test_wan_causal_denoising.py b/fastvideo/tests/stages/test_wan_causal_denoising.py index 4e00fec30..b17c390e1 100644 --- a/fastvideo/tests/stages/test_wan_causal_denoising.py +++ b/fastvideo/tests/stages/test_wan_causal_denoising.py @@ -41,8 +41,8 @@ class RecordingUniPC(FlowUniPCMultistepScheduler): return super().set_timesteps(*args, **kwargs) -def test_causal_standard_resets_scheduler_per_block_and_caches_per_request(monkeypatch): - _patch_denoising_module(monkeypatch, "1.0") +def test_causal_standard_resets_scheduler_per_block_and_caches_per_request(monkeypatch, env_overrides): + _patch_denoising_module(monkeypatch, env_overrides, "1.0") from fastvideo.pipelines.basic.wan.stages import causal_denoising monkeypatch.setattr(causal_denoising, "get_local_torch_device", lambda: torch.device("cpu")) monkeypatch.setattr(causal_denoising, "set_forward_context", lambda **kwargs: nullcontext()) diff --git a/fastvideo/tests/stages/test_wan_denoising.py b/fastvideo/tests/stages/test_wan_denoising.py index 083e09e0d..a1b55e2e2 100644 --- a/fastvideo/tests/stages/test_wan_denoising.py +++ b/fastvideo/tests/stages/test_wan_denoising.py @@ -14,8 +14,8 @@ from fastvideo.tests.stages._denoising_fixtures import ( ) -def _stage(monkeypatch, model, scheduler, *, second=None, gate="1.0"): - _, logger = _patch_denoising_module(monkeypatch, gate) +def _stage(monkeypatch, env_overrides, model, scheduler, *, second=None, gate="1.0"): + _, logger = _patch_denoising_module(monkeypatch, env_overrides, gate) from fastvideo.pipelines.basic.wan.stages import denoising monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) stage = denoising.WanDenoisingStage(model, scheduler, transformer_2=second) @@ -25,7 +25,7 @@ def _stage(monkeypatch, model, scheduler, *, second=None, gate="1.0"): @pytest.mark.parametrize("steps", [3, 50]) @pytest.mark.parametrize("cfg", [False, True]) -def test_wan_all_scheduler_steps_match_explicit_loop(monkeypatch, steps, cfg): +def test_wan_all_scheduler_steps_match_explicit_loop(monkeypatch, env_overrides, steps, cfg): batch = _batch(steps, cfg) scheduler = FlowUniPCMultistepScheduler(shift=3.0) scheduler.set_timesteps(steps, device="cpu") @@ -45,7 +45,7 @@ def test_wan_all_scheduler_steps_match_explicit_loop(monkeypatch, steps, cfg): trajectory.append(expected) model = RecordingDenoiser() - stage, _ = _stage(monkeypatch, model, scheduler) + stage, _ = _stage(monkeypatch, env_overrides, model, scheduler) result = stage.forward(batch, _args()) torch.testing.assert_close(result.latents, expected, atol=0, rtol=0) @@ -55,9 +55,9 @@ def test_wan_all_scheduler_steps_match_explicit_loop(monkeypatch, steps, cfg): assert scheduler.step_index == steps -def test_wan_expert_boundary_invalidates_cfg_cache(monkeypatch): +def test_wan_expert_boundary_invalidates_cfg_cache(monkeypatch, env_overrides): primary, secondary = RecordingDenoiser(), RecordingDenoiser(offset=0.125) - stage, logger = _stage(monkeypatch, primary, TinyScheduler(), second=secondary, gate="0.0") + stage, logger = _stage(monkeypatch, env_overrides, primary, TinyScheduler(), second=secondary, gate="0.0") args, batch = _args(), _batch() args.pipeline_config.dit_config.boundary_ratio = 0.9 batch.boundary_ratio = 0.5 # Per-request override wins; equality still uses the high-noise expert. @@ -70,7 +70,7 @@ def test_wan_expert_boundary_invalidates_cfg_cache(monkeypatch): @pytest.mark.parametrize("kind,channels", [("t2v", 2), ("i2v", 4), ("v2v", 6), ("lucy", 4), ("ti2v", 2)]) -def test_wan_conditioning_layout_and_first_frame(monkeypatch, kind, channels): +def test_wan_conditioning_layout_and_first_frame(monkeypatch, env_overrides, kind, channels): args, batch = _args(), _batch(cfg=False) model, vae = RecordingDenoiser(), TinyVAE() if kind == "i2v": @@ -81,7 +81,7 @@ def test_wan_conditioning_layout_and_first_frame(monkeypatch, kind, channels): args.pipeline_config.ti2v_task = kind == "ti2v" if kind == "ti2v": batch.pil_image = torch.zeros(1, 3, 1, 16, 32) - stage, _ = _stage(monkeypatch, model, TinyScheduler()) + stage, _ = _stage(monkeypatch, env_overrides, model, TinyScheduler()) from fastvideo.pipelines.basic.wan.stages import conditioning monkeypatch.setattr(conditioning, "get_local_torch_device", lambda: torch.device("cpu")) conditioning.WanFirstFrameEncodingStage(vae).forward(batch, args) @@ -103,9 +103,9 @@ def test_wan_conditioning_layout_and_first_frame(monkeypatch, kind, channels): assert vae.encode_calls == 0 -def test_wan_request_does_not_reuse_previous_cfg_delta(monkeypatch): +def test_wan_request_does_not_reuse_previous_cfg_delta(monkeypatch, env_overrides): model = RecordingDenoiser() - stage, _ = _stage(monkeypatch, model, TinyScheduler(), gate="0.0") + stage, _ = _stage(monkeypatch, env_overrides, model, TinyScheduler(), gate="0.0") first = stage.forward(_batch(), _args()).latents.clone() second = stage.forward(_batch(), _args()).latents.clone() torch.testing.assert_close(first, second, atol=0, rtol=0) diff --git a/fastvideo/tests/stages/test_wan_dmd_denoising.py b/fastvideo/tests/stages/test_wan_dmd_denoising.py index 7d92834d0..c927e443b 100644 --- a/fastvideo/tests/stages/test_wan_dmd_denoising.py +++ b/fastvideo/tests/stages/test_wan_dmd_denoising.py @@ -11,8 +11,8 @@ from fastvideo.tests.stages._denoising_fixtures import NullProgressBar, _patch_d from fastvideo.tests.stages._denoising_fixtures import RecordingDenoiser, _args, _batch -def test_wan_dmd_uses_full_training_table_and_preserves_rng_order(monkeypatch): - _patch_denoising_module(monkeypatch, "1.0") +def test_wan_dmd_uses_full_training_table_and_preserves_rng_order(monkeypatch, env_overrides): + _patch_denoising_module(monkeypatch, env_overrides, "1.0") from fastvideo.pipelines.basic.wan.stages import dmd monkeypatch.setattr(dmd, "get_local_torch_device", lambda: torch.device("cpu")) monkeypatch.setattr(dmd, "set_forward_context", lambda **kwargs: nullcontext()) diff --git a/fastvideo/tests/stages/test_wan_pipeline_wiring.py b/fastvideo/tests/stages/test_wan_pipeline_wiring.py index 62d2040ab..900f9aa59 100644 --- a/fastvideo/tests/stages/test_wan_pipeline_wiring.py +++ b/fastvideo/tests/stages/test_wan_pipeline_wiring.py @@ -23,8 +23,9 @@ from fastvideo.tests.stages._denoising_fixtures import RecordingDenoiser, TinyVA ("turbodiffusion.turbodiffusion_pipeline", "WanDenoisingStage", False), ("turbodiffusion.turbodiffusion_i2v_pipeline", "WanDenoisingStage", False), ]) -def test_pipeline_wires_family_sampler_and_owns_scheduler(monkeypatch, module_name, sampler, first_frame): - _patch_denoising_module(monkeypatch, "1.0") +def test_pipeline_wires_family_sampler_and_owns_scheduler(monkeypatch, env_overrides, module_name, sampler, + first_frame): + _patch_denoising_module(monkeypatch, env_overrides, "1.0") from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler module = importlib.import_module(f"fastvideo.pipelines.basic.{module_name}") diff --git a/fastvideo/tests/train/methods/grad_norm_regression.py b/fastvideo/tests/train/methods/grad_norm_regression.py index 4a7b40081..8c4d1f771 100644 --- a/fastvideo/tests/train/methods/grad_norm_regression.py +++ b/fastvideo/tests/train/methods/grad_norm_regression.py @@ -24,7 +24,7 @@ Seeding a reference for the current device: copy the recorded value from the log into ``grad_norm_refs.json``. - **Local / other GPUs** — on that workstation:: - FASTVIDEO_GRADNORM_UPDATE=1 \\ + FASTVIDEO_TEST_GRADNORM_UPDATE=1 \\ pytest fastvideo/tests/train/methods -vs -rs The harness writes the measured norm into ``grad_norm_refs.json`` under the @@ -35,14 +35,15 @@ Seeding a reference for the current device: from __future__ import annotations import json -import os from pathlib import Path import pytest import torch +import fastvideo.envs as envs + _REFS_PATH = Path(__file__).resolve().parent / "grad_norm_refs.json" -_UPDATE_ENV = "FASTVIDEO_GRADNORM_UPDATE" +_UPDATE_ENV = "FASTVIDEO_TEST_GRADNORM_UPDATE" # bf16 single-step smoke: catch gross breakage (wrong wiring, dead grads, # scale regressions), not micro-drift from reduction nondeterminism. @@ -147,13 +148,13 @@ def check_grad_norm_regression( - Skips when the current GPU has no reference (unsupported device, or not yet seeded) so a new runner never hard-fails before its golden exists. - - With ``FASTVIDEO_GRADNORM_UPDATE=1`` records/updates the reference for the + - With ``FASTVIDEO_TEST_GRADNORM_UPDATE=1`` records/updates the reference for the current device instead of asserting. """ norm = layer0_grad_norm(transformer) device_key = resolve_device_key() - if os.environ.get(_UPDATE_ENV) == "1": + if envs.FASTVIDEO_TEST_GRADNORM_UPDATE.get(): if device_key is None: pytest.skip(f"{_UPDATE_ENV}=1 but GPU '{_device_name()}' has no reference " "key; add it to _DEVICE_MAPPINGS first") diff --git a/fastvideo/tests/train/methods/test_cosmos_finetune.py b/fastvideo/tests/train/methods/test_cosmos_finetune.py index 3029a4567..f68ae0954 100644 --- a/fastvideo/tests/train/methods/test_cosmos_finetune.py +++ b/fastvideo/tests/train/methods/test_cosmos_finetune.py @@ -12,10 +12,10 @@ the fixture (``training_cfg_rate=0`` is required, and from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29521") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29521") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_ltx2_finetune.py b/fastvideo/tests/train/methods/test_ltx2_finetune.py index 5a9bc3467..951b7626c 100644 --- a/fastvideo/tests/train/methods/test_ltx2_finetune.py +++ b/fastvideo/tests/train/methods/test_ltx2_finetune.py @@ -12,10 +12,10 @@ memory (e.g. the L40S CI runner). from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29522") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29522") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py index 9a591b983..5c5f33d8d 100644 --- a/fastvideo/tests/train/methods/test_matrixgame2_finetune.py +++ b/fastvideo/tests/train/methods/test_matrixgame2_finetune.py @@ -19,10 +19,10 @@ by a loading-only smoke for the same reason. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29522") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29522") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_streaming_long_tuning.py b/fastvideo/tests/train/methods/test_streaming_long_tuning.py index d58a36bbc..05c4244c4 100644 --- a/fastvideo/tests/train/methods/test_streaming_long_tuning.py +++ b/fastvideo/tests/train/methods/test_streaming_long_tuning.py @@ -16,10 +16,10 @@ bottom — they need no GPU, no mocks, and run wherever this file is collected. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29523") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29523") from pathlib import Path from types import SimpleNamespace @@ -174,7 +174,7 @@ def test_streaming_long_tuning_multi_stage_train_steps(monkeypatch: pytest.Monke # Grad-norm regression on the post-reset chunk (grads zeroed at the top # of this iteration, so the reference captures exactly one step). Keep - # this last: under FASTVIDEO_GRADNORM_UPDATE=1 it records and skips. + # this last: under FASTVIDEO_TEST_GRADNORM_UPDATE=1 it records and skips. check_grad_norm_regression( "test_streaming_long_tuning", student.transformer, diff --git a/fastvideo/tests/train/methods/test_wan_causal_cd.py b/fastvideo/tests/train/methods/test_wan_causal_cd.py index 3e4b3b47f..0aa3147b1 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_cd.py +++ b/fastvideo/tests/train/methods/test_wan_causal_cd.py @@ -11,10 +11,10 @@ gradients, frozen teacher, and a post-step EMA update. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29519") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py index 7692e9b24..43f5d1ad5 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_dfsft.py +++ b/fastvideo/tests/train/methods/test_wan_causal_dfsft.py @@ -13,10 +13,10 @@ in the fixture) and is the natural training counterpart of the from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29517") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29517") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py b/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py index 77e21732c..43d026da4 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py +++ b/fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py @@ -10,10 +10,10 @@ nonzero gradients reaching the first transformer block. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29520") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29520") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_wan_causal_tfsft.py b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py index 590a95ece..e64617292 100644 --- a/fastvideo/tests/train/methods/test_wan_causal_tfsft.py +++ b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py @@ -10,10 +10,10 @@ gradients reaching the first transformer block. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29518") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29518") from pathlib import Path diff --git a/fastvideo/tests/train/methods/test_wan_finetune.py b/fastvideo/tests/train/methods/test_wan_finetune.py index 531579384..64d1c660a 100644 --- a/fastvideo/tests/train/methods/test_wan_finetune.py +++ b/fastvideo/tests/train/methods/test_wan_finetune.py @@ -19,12 +19,12 @@ device-keyed grad-norm regression on top of this same harness. from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29516") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29516") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py index e3db14651..f34276470 100644 --- a/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py +++ b/fastvideo/tests/train/models/test_kandinsky5_qat_attention_engages.py @@ -18,9 +18,9 @@ Two silent-fallback traps this guards against: This test itself must not become that stale-selection source for whichever model test runs next in the same pytest process: ``FASTVIDEO_ATTENTION_BACKEND`` -is set via ``monkeypatch`` (auto-restored at teardown, including on +is set via ``env_overrides`` (auto-restored at teardown, including on failure) and the cache is cleared again in a ``finally`` block, since -``monkeypatch`` only restores the env var, not the separate process-wide +``env_overrides`` only restores the env var, not the separate process-wide cache keyed on it. Unlike ``fastvideo.platforms.cuda``'s other backends, ATTN_QAT_TRAIN has no @@ -44,10 +44,10 @@ sidesteps both problems. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29519") from pathlib import Path @@ -63,7 +63,7 @@ _FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "kandinsky5 @pytest.mark.usefixtures("distributed_setup") -def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): +def test_kandinsky5_attn_qat_train_engages_and_backprops(env_overrides): if not torch.cuda.is_available(): pytest.skip("requires CUDA") @@ -72,7 +72,7 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): if not is_attn_qat_train_available(): pytest.skip("fastvideo_kernel ATTN_QAT_TRAIN kernel is not built") - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "ATTN_QAT_TRAIN") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("ATTN_QAT_TRAIN")) _cached_get_attn_backend.cache_clear() try: from fastvideo.train.models.kandinsky5 import Kandinsky5Model @@ -156,7 +156,7 @@ def test_kandinsky5_attn_qat_train_engages_and_backprops(monkeypatch): assert torch.isfinite(attn.to_query.weight.grad.to_local() if hasattr(attn.to_query.weight.grad, "to_local") else attn.to_query.weight.grad).all().item(), "weight grad contains NaN/Inf" finally: - # monkeypatch restores FASTVIDEO_ATTENTION_BACKEND itself, but the + # env_overrides restores FASTVIDEO_ATTENTION_BACKEND itself, but the # selector cache is a separate process-wide functools.cache keyed on # (head_size, dtype, supported_backends) -- not on the env var -- so # it must be cleared again here or a later model test in this same diff --git a/fastvideo/tests/train/models/test_load_cosmos.py b/fastvideo/tests/train/models/test_load_cosmos.py index 444aa5caf..8c7e316b0 100644 --- a/fastvideo/tests/train/models/test_load_cosmos.py +++ b/fastvideo/tests/train/models/test_load_cosmos.py @@ -15,12 +15,12 @@ kwargs in ``CosmosModel._build_distill_input_kwargs``. from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29518") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29518") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_hunyuan.py b/fastvideo/tests/train/models/test_load_hunyuan.py index 90d71d1c3..0d2479e7c 100644 --- a/fastvideo/tests/train/models/test_load_hunyuan.py +++ b/fastvideo/tests/train/models/test_load_hunyuan.py @@ -11,10 +11,10 @@ forward signature than Wan (no ``encoder_attention_mask``, no from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29517") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29517") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_kandinsky5.py b/fastvideo/tests/train/models/test_load_kandinsky5.py index 2eff39833..beda9603d 100644 --- a/fastvideo/tests/train/models/test_load_kandinsky5.py +++ b/fastvideo/tests/train/models/test_load_kandinsky5.py @@ -15,8 +15,10 @@ from __future__ import annotations import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29518") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29518") os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_longcat.py b/fastvideo/tests/train/models/test_load_longcat.py index f0fcdffd9..7bf0f8a47 100644 --- a/fastvideo/tests/train/models/test_load_longcat.py +++ b/fastvideo/tests/train/models/test_load_longcat.py @@ -16,12 +16,12 @@ same runner). from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29519") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_ltx2.py b/fastvideo/tests/train/models/test_load_ltx2.py index 5fa4c02c6..4d2603357 100644 --- a/fastvideo/tests/train/models/test_load_ltx2.py +++ b/fastvideo/tests/train/models/test_load_ltx2.py @@ -17,12 +17,12 @@ the denoised x0 prediction. This mirrors the kwargs in from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29519") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29519") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_matrixgame2.py b/fastvideo/tests/train/models/test_load_matrixgame2.py index ac182db84..b2b31012d 100644 --- a/fastvideo/tests/train/models/test_load_matrixgame2.py +++ b/fastvideo/tests/train/models/test_load_matrixgame2.py @@ -17,12 +17,12 @@ mirrors ``MatrixGame2Model._build_distill_input_kwargs``. from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29520") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29520") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_wan.py b/fastvideo/tests/train/models/test_load_wan.py index 5b6c63a62..c8633d58a 100644 --- a/fastvideo/tests/train/models/test_load_wan.py +++ b/fastvideo/tests/train/models/test_load_wan.py @@ -10,12 +10,12 @@ or forward-signature regressions in from __future__ import annotations -import os +import fastvideo.envs as envs # Required by the ``distributed_setup`` fixture pulled from -# ``fastvideo/tests/conftest.py``. Set before any fastvideo import. -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29515") +# ``fastvideo/tests/conftest.py``. Set at import, before the fixture runs. +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29515") from pathlib import Path diff --git a/fastvideo/tests/train/models/test_load_wan_causal.py b/fastvideo/tests/train/models/test_load_wan_causal.py index ec8321de1..3978f2a87 100644 --- a/fastvideo/tests/train/models/test_load_wan_causal.py +++ b/fastvideo/tests/train/models/test_load_wan_causal.py @@ -14,10 +14,10 @@ streaming forward path end-to-end. from __future__ import annotations -import os +import fastvideo.envs as envs -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29516") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29516") from pathlib import Path diff --git a/fastvideo/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/tests/training/VSA/test_training_loss_VSA.py index e576c3f5c..214ee73c7 100644 --- a/fastvideo/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/tests/training/VSA/test_training_loss_VSA.py @@ -6,13 +6,12 @@ import json from huggingface_hub import snapshot_download import torch -# Ensure backend selection happens during import-time initialization -os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" # Force VSA to use Triton implementation even on H100 / when CUDA extension is available # os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1" # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) +import fastvideo.envs as envs from fastvideo.training.wan_training_pipeline import main from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.utils import FlexibleArgumentParser @@ -56,8 +55,6 @@ def run_worker(): def test_distributed_training(): """Test the distributed training setup""" - os.environ.setdefault("WANDB_MODE", "offline") - data_dir = Path("data/mini_dataset_i2v_VSA") if not data_dir.exists(): @@ -73,7 +70,12 @@ def test_distributed_training(): # Run torchrun command cmd = ["torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, str(current_file)] - process = subprocess.run(cmd, check=True) + # The torchrun workers select VSA when the training pipeline initializes attention. + with ( + envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN"), + envs.override_external("WANDB_MODE", "offline"), + ): + process = subprocess.run(cmd, check=True) summary_file = 'data/wan_finetune_test_VSA/tracker/wandb/latest-run/files/wandb-summary.json' diff --git a/fastvideo/tests/training/Vanilla/mfu_calculation.py b/fastvideo/tests/training/Vanilla/mfu_calculation.py index d45f8e0e8..bf78e9a34 100644 --- a/fastvideo/tests/training/Vanilla/mfu_calculation.py +++ b/fastvideo/tests/training/Vanilla/mfu_calculation.py @@ -6,12 +6,12 @@ from pathlib import Path current_dir = str(Path(__file__).parent.parent.parent.parent.parent) if current_dir not in sys.path: sys.path.insert(0, current_dir) -os.environ["PYTHONPATH"] = current_dir + ":" + os.environ.get("PYTHONPATH", "") import subprocess import torch import json from huggingface_hub import snapshot_download +import fastvideo.envs as envs from fastvideo.utils import logger # Import the training pipeline from fastvideo.training.wan_training_pipeline import main @@ -31,8 +31,8 @@ NUM_GPUS_PER_NODE = "2" GRAD_ACCUM = "1" MASTER_PORT = os.environ.get("MASTER_PORT", "29504") -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", MASTER_PORT) +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", MASTER_PORT) def run_worker(): @@ -131,8 +131,6 @@ def run_worker(): def test_distributed_training(): """Test the distributed training setup""" - os.environ.setdefault("WANDB_MODE", "offline") - data_dir = Path("data/crush-smol_processed_t2v") if not data_dir.exists(): @@ -150,7 +148,10 @@ def test_distributed_training(): "torchrun", "--nnodes", NUM_NODES, "--nproc_per_node", NUM_GPUS_PER_NODE, "--master_port", MASTER_PORT, str(current_file) ] - process = subprocess.run(cmd, capture_output=True, text=True) + # The torchrun workers run this file and import fastvideo from the repository root. + worker_pythonpath = current_dir + ":" + os.environ.get("PYTHONPATH", "") + with envs.override_external("WANDB_MODE", "offline"), envs.override_external("PYTHONPATH", worker_pythonpath): + process = subprocess.run(cmd, capture_output=True, text=True) # Print stdout and stderr for debugging if process.stdout: diff --git a/fastvideo/tests/training/Vanilla/test_training_loss.py b/fastvideo/tests/training/Vanilla/test_training_loss.py index 59adae0a5..fddbc928c 100644 --- a/fastvideo/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/tests/training/Vanilla/test_training_loss.py @@ -1,7 +1,9 @@ import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29512") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29512") import sys import subprocess from pathlib import Path @@ -58,8 +60,6 @@ def run_worker(): def test_distributed_training(): """Test the distributed training setup""" - os.environ.setdefault("WANDB_MODE", "offline") - data_dir = Path("data/crush-smol_processed_t2v") if not data_dir.exists(): @@ -78,7 +78,8 @@ def test_distributed_training(): os.environ["MASTER_PORT"], str(current_file) ] - process = subprocess.run(cmd, capture_output=True, text=True) + with envs.override_external("WANDB_MODE", "offline"): + process = subprocess.run(cmd, capture_output=True, text=True) # Print stdout and stderr for debugging if process.stdout: diff --git a/fastvideo/tests/training/distill/test_distill_dmd.py b/fastvideo/tests/training/distill/test_distill_dmd.py index 8d80811e8..240dcadab 100644 --- a/fastvideo/tests/training/distill/test_distill_dmd.py +++ b/fastvideo/tests/training/distill/test_distill_dmd.py @@ -1,7 +1,9 @@ import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29512") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29512") import sys import subprocess from pathlib import Path @@ -57,8 +59,6 @@ def run_worker(): def test_distributed_training(): """Test the distributed training setup""" - os.environ["WANDB_MODE"] = "offline" - data_dir = Path("data/crush-smol_processed_t2v") if not data_dir.exists(): @@ -77,7 +77,8 @@ def test_distributed_training(): os.environ["MASTER_PORT"], str(current_file) ] - process = subprocess.run(cmd, capture_output=True, text=True) + with envs.override_external("WANDB_MODE", "offline"): + process = subprocess.run(cmd, capture_output=True, text=True) # Print stdout and stderr for debugging if process.stdout: diff --git a/fastvideo/tests/training/lora/test_lora_training.py b/fastvideo/tests/training/lora/test_lora_training.py index e5664d1ad..67a5170e7 100644 --- a/fastvideo/tests/training/lora/test_lora_training.py +++ b/fastvideo/tests/training/lora/test_lora_training.py @@ -1,7 +1,9 @@ import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29514") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29514") import sys import subprocess from pathlib import Path @@ -18,8 +20,6 @@ NUM_GPUS_PER_NODE = "2" def test_lora_training(): """Test the LoRA training setup""" - os.environ.setdefault("WANDB_MODE", "offline") - data_dir = Path("data/crush-smol_processed_t2v") if not data_dir.exists(): @@ -50,7 +50,8 @@ def test_lora_training(): "--lora_training", "True", "--seed", "42" ] - process = subprocess.run(cmd, check=True) + with envs.override_external("WANDB_MODE", "offline"): + process = subprocess.run(cmd, check=True) summary_file = '/workspace/tracker/wandb/latest-run/files/wandb-summary.json' diff --git a/fastvideo/tests/training/self-forcing/test_self_forcing.py b/fastvideo/tests/training/self-forcing/test_self_forcing.py index 1dc64ea93..defa3e3a3 100644 --- a/fastvideo/tests/training/self-forcing/test_self_forcing.py +++ b/fastvideo/tests/training/self-forcing/test_self_forcing.py @@ -1,7 +1,9 @@ import os -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29513") +import fastvideo.envs as envs + +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29513") import sys import subprocess from pathlib import Path @@ -156,8 +158,6 @@ def run_worker(): def test_distributed_training(): """Test the distributed self-forcing training setup""" - os.environ["WANDB_MODE"] = "offline" - data_dir = Path("data/crush-smol_processed_t2v") if not data_dir.exists(): @@ -176,7 +176,8 @@ def test_distributed_training(): os.environ["MASTER_PORT"], str(current_file) ] - process = subprocess.run(cmd, capture_output=True, text=True) + with envs.override_external("WANDB_MODE", "offline"): + process = subprocess.run(cmd, capture_output=True, text=True) # Print stdout and stderr for debugging if process.stdout: diff --git a/fastvideo/tests/transformers/test_cosmos.py b/fastvideo/tests/transformers/test_cosmos.py index bf8455099..eb2bdaae6 100644 --- a/fastvideo/tests/transformers/test_cosmos.py +++ b/fastvideo/tests/transformers/test_cosmos.py @@ -6,6 +6,7 @@ import pytest import torch from diffusers.models.transformers.transformer_cosmos import CosmosTransformer3DModel +import fastvideo.envs as envs from fastvideo.configs.pipelines import PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.fastvideo_args import FastVideoArgs @@ -18,8 +19,8 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29504") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29504") BASE_MODEL_PATH = "nvidia/Cosmos-Predict2-2B-Video2World" diff --git a/fastvideo/tests/transformers/test_cosmos2_5.py b/fastvideo/tests/transformers/test_cosmos2_5.py index b1e6d2e4e..70161f91d 100644 --- a/fastvideo/tests/transformers/test_cosmos2_5.py +++ b/fastvideo/tests/transformers/test_cosmos2_5.py @@ -17,6 +17,7 @@ COSMOS_PREDICT2_5_PATH = os.path.normpath(COSMOS_PREDICT2_5_PATH) if os.path.exists(COSMOS_PREDICT2_5_PATH) and COSMOS_PREDICT2_5_PATH not in sys.path: sys.path.insert(0, COSMOS_PREDICT2_5_PATH) +import fastvideo.envs as envs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.tests.utils import skip_if_gated_repo_inaccessible @@ -33,8 +34,8 @@ if os.path.exists(COSMOS_PREDICT2_5_PATH): else: logger.warning(f"cosmos-predict2.5 not found at: {COSMOS_PREDICT2_5_PATH}") -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29505") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29505") # COSMOS 2.5 model path - update this based on the actual HuggingFace model ID # The model has subdirectories: base/pre-trained, base/post-trained, auto/multiview, robot/action-cond diff --git a/fastvideo/tests/transformers/test_flux.py b/fastvideo/tests/transformers/test_flux.py index 902cb6599..5917fd481 100644 --- a/fastvideo/tests/transformers/test_flux.py +++ b/fastvideo/tests/transformers/test_flux.py @@ -10,6 +10,7 @@ import torch from diffusers import FluxTransformer2DModel as HFFluxTransformer2DModel from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.models.dits.flux import FluxDiTConfig from fastvideo.configs.pipelines.base import PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs @@ -17,8 +18,8 @@ from fastvideo.forward_context import set_forward_context from fastvideo.models.loader.component_loader import TransformerLoader from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29517") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29517") _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) _DEFAULT_FLUX_TRANSFORMER = os.path.join( @@ -30,7 +31,7 @@ _DEFAULT_FLUX_TRANSFORMER = os.path.join( def _flux_transformer_path() -> str: - return os.environ.get("FLUX_TRANSFORMER_PATH", _DEFAULT_FLUX_TRANSFORMER) + return envs.FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH.get() or _DEFAULT_FLUX_TRANSFORMER def _prepare_latent_image_ids( @@ -49,8 +50,8 @@ def _prepare_latent_image_ids( @pytest.fixture -def torch_sdpa_attention_backend(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +def torch_sdpa_attention_backend(env_overrides) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) requires_cuda = pytest.mark.skipif( @@ -58,21 +59,17 @@ requires_cuda = pytest.mark.skipif( reason="FLUX DiT parity test requires CUDA", ) -requires_weights = pytest.mark.skipif( - not glob.glob(os.path.join(_flux_transformer_path(), "*.safetensors")), - reason=(f"No safetensors under {_flux_transformer_path()} — download FLUX.1-dev " - "transformer or set FLUX_TRANSFORMER_PATH"), -) - @requires_cuda -@requires_weights @pytest.mark.usefixtures("distributed_setup", "torch_sdpa_attention_backend") def test_flux_transformer_parity_vs_diffusers() -> None: """Single forward: FastVideo DiT vs Diffusers ``FluxTransformer2DModel``.""" + transformer_path = _flux_transformer_path() + if not glob.glob(os.path.join(transformer_path, "*.safetensors")): + pytest.skip(f"No safetensors under {transformer_path} — download FLUX.1-dev " + "transformer or set FASTVIDEO_TEST_FLUX_TRANSFORMER_PATH") device = torch.device("cuda:0") precision = torch.bfloat16 - transformer_path = _flux_transformer_path() args = FastVideoArgs( model_path=transformer_path, diff --git a/fastvideo/tests/transformers/test_hunyuangamecraft.py b/fastvideo/tests/transformers/test_hunyuangamecraft.py index 5660b7a3b..66864a79f 100644 --- a/fastvideo/tests/transformers/test_hunyuangamecraft.py +++ b/fastvideo/tests/transformers/test_hunyuangamecraft.py @@ -9,6 +9,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig from fastvideo.configs.pipelines import PipelineConfig from fastvideo.distributed.parallel_state import ( @@ -23,8 +24,8 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29507") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29507") # Path to converted weights (local path, not HuggingFace) TRANSFORMER_PATH = "official_weights/hunyuan-gamecraft/transformer" @@ -37,18 +38,17 @@ REFERENCE_LATENT = 42351.12903189659 @pytest.mark.usefixtures("distributed_setup") -def test_hunyuangamecraft_transformer(monkeypatch: pytest.MonkeyPatch): +def test_hunyuangamecraft_transformer(env_overrides): """Test HunyuanGameCraft transformer regression.""" if not os.path.exists(TRANSFORMER_PATH): pytest.skip(f"Weights not found at {TRANSFORMER_PATH}") - # Keep these model-specific overrides local to this test. Setting + # Keep this model-specific override local to this test. Setting # TORCHDYNAMO_DISABLE during module collection disables torch.compile for # every later transformer test in the same pytest process; FakeTensor # compile checks then execute Triton kernels against fake CUDA pointers. - monkeypatch.setenv("TORCHDYNAMO_DISABLE", "1") - monkeypatch.setenv("DISABLE_SP", "1") + env_overrides.enter_context(envs.override_external("TORCHDYNAMO_DISABLE", "1")) sp_rank = get_sp_parallel_rank() sp_world_size = get_sp_world_size() diff --git a/fastvideo/tests/transformers/test_hyworld.py b/fastvideo/tests/transformers/test_hyworld.py index d2887ee66..14ffea7b9 100644 --- a/fastvideo/tests/transformers/test_hyworld.py +++ b/fastvideo/tests/transformers/test_hyworld.py @@ -4,6 +4,7 @@ import os import pytest import torch +import fastvideo.envs as envs from fastvideo.configs.models.dits import HYWorldConfig from fastvideo.configs.pipelines import PipelineConfig from fastvideo.distributed.parallel_state import ( @@ -20,8 +21,8 @@ from fastvideo.utils import maybe_download_model logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") MODEL_PATH = maybe_download_model("FastVideo/HY-WorldPlay-Bidirectional-Diffusers") TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") @@ -43,7 +44,7 @@ def test_hyworld_transformer(): # An env-forced backend is knowable before the multi-GB download and GPU # load; the authoritative check on the layer's resolved backend still runs # after construction below. - forced_backend_name = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") + forced_backend_name = envs.FASTVIDEO_ATTENTION_BACKEND.get() if forced_backend_name: forced_backend = getattr(AttentionBackendEnum, forced_backend_name, None) if forced_backend is not None and forced_backend not in REFERENCE_LATENTS: diff --git a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py index 2ba2e932f..63de68b04 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py +++ b/fastvideo/tests/transformers/test_minimax_h3_fusion_routing.py @@ -6,6 +6,7 @@ from __future__ import annotations import pytest import torch +import fastvideo.envs as envs from fastvideo.platforms import AttentionBackendEnum @@ -135,7 +136,7 @@ def test_swiglu_fusion_stays_on_eager_path_with_grad(monkeypatch: pytest.MonkeyP def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4( - monkeypatch: pytest.MonkeyPatch, + env_overrides, ) -> None: """Exercise the real block wiring without loading any H3 checkpoint.""" if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): @@ -145,13 +146,13 @@ def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4( if "fa4" not in getattr(flash_attn, "__version__", "").lower(): pytest.skip("the focused integration test requires the FA4 environment") - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN") - monkeypatch.setenv("FASTVIDEO_FA4", "1") - monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") - monkeypatch.setenv("MASTER_PORT", "29573") - monkeypatch.setenv("RANK", "0") - monkeypatch.setenv("WORLD_SIZE", "1") - monkeypatch.setenv("LOCAL_RANK", "0") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) + env_overrides.enter_context(envs.FASTVIDEO_FA4.override(True)) + env_overrides.enter_context(envs.override_external("MASTER_ADDR", "127.0.0.1")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", "29573")) + env_overrides.enter_context(envs.override_external("RANK", "0")) + env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1")) + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel from fastvideo.forward_context import set_forward_context @@ -223,7 +224,7 @@ def test_all_minimax_h3_fusions_match_one_eager_block_under_fa4( cleanup_dist_env_and_memory() -def test_minimax_h3_fusions_engage_on_cuda_inference(monkeypatch: pytest.MonkeyPatch) -> None: +def test_minimax_h3_fusions_engage_on_cuda_inference(monkeypatch: pytest.MonkeyPatch, env_overrides) -> None: """Pin the positive side of the routing guard. The parity test above still passes if ``_can_run_minimax_h3_fusion`` @@ -238,12 +239,12 @@ def test_minimax_h3_fusions_engage_on_cuda_inference(monkeypatch: pytest.MonkeyP pytest.skip("BF16 CUDA is required") pytest.importorskip("triton") - monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") - monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") - monkeypatch.setenv("MASTER_PORT", "29574") - monkeypatch.setenv("RANK", "0") - monkeypatch.setenv("WORLD_SIZE", "1") - monkeypatch.setenv("LOCAL_RANK", "0") + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + env_overrides.enter_context(envs.override_external("MASTER_ADDR", "127.0.0.1")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", "29574")) + env_overrides.enter_context(envs.override_external("RANK", "0")) + env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1")) + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) import fastvideo.models.dits.minimax_h3 as h3 from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index c8790c5f1..1827cd667 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """Full-checkpoint Wan parity against Diffusers; requires one CUDA GPU and weights.""" -import os - import pytest import torch from diffusers import WanTransformer3DModel from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.pipelines import PipelineConfig from fastvideo.forward_context import set_forward_context from fastvideo.fastvideo_args import FastVideoArgs @@ -16,8 +15,8 @@ from fastvideo.tests.golden_gate._wan_checkpoint import component_path from fastvideo.models.wan.config import WanVideoConfig from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Wan transformer parity requires one CUDA GPU") diff --git a/fastvideo/tests/vaes/test_hunyuan15_vae.py b/fastvideo/tests/vaes/test_hunyuan15_vae.py index 000f8c702..707139426 100644 --- a/fastvideo/tests/vaes/test_hunyuan15_vae.py +++ b/fastvideo/tests/vaes/test_hunyuan15_vae.py @@ -7,6 +7,7 @@ import pytest import torch from diffusers import AutoencoderKLHunyuanVideo15 +import fastvideo.envs as envs from fastvideo.configs.pipelines import PipelineConfig from fastvideo.logger import init_logger # from fastvideo.models.vaes.hunyuanvae import ( @@ -18,8 +19,8 @@ from fastvideo.utils import maybe_download_model logger = init_logger(__name__) -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) diff --git a/fastvideo/tests/vaes/test_wan_vae.py b/fastvideo/tests/vaes/test_wan_vae.py index 532060a99..b12d7ecd9 100644 --- a/fastvideo/tests/vaes/test_wan_vae.py +++ b/fastvideo/tests/vaes/test_wan_vae.py @@ -1,21 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 """Wan VAE encode/decode parity against Diffusers (one CUDA GPU, FP32).""" -import os - import pytest import torch from diffusers import AutoencoderKLWan from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.pipelines import PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import VAELoader from fastvideo.models.wan.vae_config import WanVAEConfig from fastvideo.tests.golden_gate._wan_checkpoint import component_path -os.environ.setdefault("MASTER_ADDR", "localhost") -os.environ.setdefault("MASTER_PORT", "29503") +envs.setdefault_external("MASTER_ADDR", "localhost") +envs.setdefault_external("MASTER_PORT", "29503") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Wan VAE parity requires one CUDA GPU") diff --git a/fastvideo/tests/worker/test_gpu_worker.py b/fastvideo/tests/worker/test_gpu_worker.py index 56cfd9e1e..8acc2c3ea 100644 --- a/fastvideo/tests/worker/test_gpu_worker.py +++ b/fastvideo/tests/worker/test_gpu_worker.py @@ -5,15 +5,16 @@ from unittest.mock import Mock import pytest import torch +import fastvideo.envs as envs from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines import ForwardBatch from fastvideo.worker.gpu_worker import Worker, _log_cuda_device_uuid -def test_cuda_device_uuid_receipt_is_disabled_without_nvtx_profiling(monkeypatch) -> None: +def test_cuda_device_uuid_receipt_is_disabled_without_nvtx_profiling(monkeypatch, env_overrides) -> None: """Avoid NVIDIA property access during ordinary worker initialization.""" get_device_properties = Mock() - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "0") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(False)) monkeypatch.setattr(torch.cuda, "get_device_properties", get_device_properties) _log_cuda_device_uuid(0, torch.device("cuda:0")) @@ -21,10 +22,10 @@ def test_cuda_device_uuid_receipt_is_disabled_without_nvtx_profiling(monkeypatch get_device_properties.assert_not_called() -def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> None: +def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch, env_overrides) -> None: """Bind one profiled worker rank to its NVIDIA device UUID in logs.""" log_info = Mock() - monkeypatch.setenv("FASTVIDEO_NVTX_PROFILE", "1") + env_overrides.enter_context(envs.FASTVIDEO_NVTX_PROFILE.override(True)) monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: SimpleNamespace(uuid="device-uuid")) monkeypatch.setattr("fastvideo.worker.gpu_worker.logger.info", log_info) @@ -39,14 +40,15 @@ def test_cuda_device_uuid_receipt_identifies_profiled_worker(monkeypatch) -> Non @pytest.mark.parametrize("executor_backend", ["mp", "ray"]) -def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypatch, executor_backend: str) -> None: +def test_init_device_applies_offload_policy_after_binding_worker_device(monkeypatch, env_overrides, + executor_backend: str) -> None: """The runtime probe must see this worker's device, never driver device 0.""" events = [] args = FastVideoArgs(model_path="test", num_gpus=1, distributed_executor_backend=executor_backend) args.finalize_device_offload_policy = Mock(side_effect=lambda device_id: events.append(("policy", device_id))) worker = Worker(args, local_rank=3, rank=3, distributed_init_method="env://") - monkeypatch.setenv("LOCAL_RANK", "0") + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda_alike", lambda: True) monkeypatch.setattr("fastvideo.platforms.current_platform.is_cuda", lambda: False) monkeypatch.setattr(torch.cuda, "set_device", lambda device: events.append(("set_device", device.index))) diff --git a/tests/local_tests/dreamx_world/PORT_STATUS.md b/tests/local_tests/dreamx_world/PORT_STATUS.md index 012d9c9bf..986a33e31 100644 --- a/tests/local_tests/dreamx_world/PORT_STATUS.md +++ b/tests/local_tests/dreamx_world/PORT_STATUS.md @@ -56,7 +56,7 @@ | ar_pipeline_config | `PYTHONPATH=/workspace/FastVideo python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -q -rs` | 10_passed_0_skipped | 2026-07-02: `DreamXWorld5BARPipelineConfig`, `DreamXWorldARPipeline`, registry config selection for `GD-ML/DreamX-World-5B`, and `dreamx_world_5b_ar` preset pass. | | ar_full_generation | `PYTHONPATH=/workspace/FastVideo python /tmp/run_dreamx_ar_video_smoke.py` | generated_video_pass | 2026-07-02: A40 short full-generation smoke passed from `/tmp/converted_dreamx_world_ar` with 64x64, 9 frames, 4 denoise steps, `output_type=pil`, `save_video=True`; MP4 saved at `outputs_video/dreamx_world_ar_video_smoke/a quiet road through a futuristic coastal city at sunrise.mp4` and decoded as 9 frames of `(64, 64, 3)` uint8. | | ar_long_horizon | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /tmp/run_dreamx_ar_long_horizon.py` | generated_video_pass | 2026-07-02: A40 long-horizon AR generation passed from `/tmp/converted_dreamx_world_ar` with 64x64, 1005 frames, 4 denoise steps, seed 4096; MP4 saved at `outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4` and decoded as 1005 frames of `(64, 64, 3)`; end-to-end generation latency was 231.68s after load. | -| ar_ssim_default | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download` | 1_passed_0_skipped | 2026-07-02: A40 default AR SSIM reference seeded locally at `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/`; helper removes stale generated base MP4 before generation; default params are 192x192, 81 frames, 4 steps, seed 2048, min SSIM 0.98. | +| ar_ssim_default | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download` | 1_passed_0_skipped | 2026-07-02: A40 default AR SSIM reference seeded locally at `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/`; helper removes stale generated base MP4 before generation; default params are 192x192, 81 frames, 4 steps, seed 2048, min SSIM 0.98. | | ar_ssim_modal_l40s | `modal run /tmp/modal_dreamx_ar_ssim_git.py` | 1_passed_0_skipped | 2026-07-02: Modal L40S run checked out `post-fix dreamx-world-5b-cam branch commit`, downloaded default `L40S_reference_videos` via `reference_videos_cli.py download`, used cached converted AR weights under `/root/data/dreamx_world_ar_converted`, seeded the missing AR L40S reference from generated output, reran a fresh generated-vs-reference compare successfully (`mean_ssim=1.0`), and exported the reference to Modal volume `hf-model-weights:dreamx_ar_ssim_l40s`. Downloaded local reference decodes as 81 frames of `(192, 192, 3)`. A40-vs-L40S reference spot check after deterministic AR noise fix: mean SSIM 0.7770, min 0.6139, max 0.9944. | | pipeline_smoke | `python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs` | 4_passed_0_skipped | 2026-06-30: combined smoke/parity passed. 2026-07-01: smoke alone passed with real `image_path` TI2V coverage (`3 passed`), validating image load, TI2V preprocessing, VAE first-frame encode under CPU offload, camera conditioning, and 1-step latent generation from `converted_weights/dreamx_world`. Tests force `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` to avoid the local FlashAttention-4 cute ABI mismatch. | | basic_example | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA DREAMX_WORLD_MODEL_DIR=converted_weights/dreamx_world DREAMX_WORLD_IMAGE_PATH= DREAMX_WORLD_HEIGHT=64 DREAMX_WORLD_WIDTH=64 DREAMX_WORLD_NUM_FRAMES=9 DREAMX_WORLD_STEPS=1 DREAMX_WORLD_GUIDANCE=1.0 DREAMX_WORLD_OUTPUT_PATH=outputs_video/dreamx_world_example_smoke python examples/inference/basic/basic_dreamx_world.py` | generated_video_pass | 2026-06-30: example saved an MP4 under `outputs_video/dreamx_world_example_smoke`; imageio/ffmpeg decoded frame 0 as `(64, 64, 3)` uint8, fps 16, duration 0.56s. | diff --git a/tests/local_tests/dreamx_world/README.md b/tests/local_tests/dreamx_world/README.md index 3febd5215..2a836b17b 100644 --- a/tests/local_tests/dreamx_world/README.md +++ b/tests/local_tests/dreamx_world/README.md @@ -205,7 +205,7 @@ PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /t # output: outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4 # decoded: 1005 frames, (64, 64, 3) -PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download +PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B FASTVIDEO_TEST_DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download # 2026-07-02: DreamX-World-5B AR default SSIM passed: 1 passed, 0 skipped # local A40 reference: fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/ diff --git a/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py b/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py index 4de60de79..7005b96b9 100644 --- a/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py +++ b/tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py @@ -24,6 +24,7 @@ import torch from PIL import Image from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel from fastvideo.models.loader.component_loader import TextEncoderLoader @@ -38,9 +39,9 @@ def _require_assets() -> tuple[torch.device, Path]: if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): pytest.fail("MiniMax-H3 Qwen3-VL parity requires a bf16-capable CUDA GPU", pytrace=False) - model_root = os.environ.get("MINIMAX_H3_MODEL_ROOT") + model_root = envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() if not model_root: - pytest.fail("set MINIMAX_H3_MODEL_ROOT to the downloaded MiniMax-H3 checkpoint", pytrace=False) + pytest.fail("set FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT to the downloaded MiniMax-H3 checkpoint", pytrace=False) root = Path(model_root) required = (root / "text_encoder", root / "tokenizer", root / "processor") missing = [str(path) for path in required if not path.is_dir()] diff --git a/tests/local_tests/flux/README.md b/tests/local_tests/flux/README.md index bf023834c..250e43eb6 100644 --- a/tests/local_tests/flux/README.md +++ b/tests/local_tests/flux/README.md @@ -66,7 +66,7 @@ Reference images must be seeded first via `reference_videos_cli.py`. ```bash FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ -FLUX_T2I_MODEL_DIR=official_weights/FLUX.1-dev \ +FASTVIDEO_TEST_FLUX_T2I_MODEL_DIR=official_weights/FLUX.1-dev \ pytest fastvideo/tests/ssim/test_flux_t2i_similarity.py -vs ``` diff --git a/tests/local_tests/glm_image/test_ar_parity.py b/tests/local_tests/glm_image/test_ar_parity.py index c4ebf4795..3c4f4f0cf 100644 --- a/tests/local_tests/glm_image/test_ar_parity.py +++ b/tests/local_tests/glm_image/test_ar_parity.py @@ -7,13 +7,16 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) AR_DIR = LOCAL_WEIGHTS_DIR / "vision_language_encoder" PROCESSOR_DIR = LOCAL_WEIGHTS_DIR / "processor" diff --git a/tests/local_tests/glm_image/test_edit_pipeline_parity.py b/tests/local_tests/glm_image/test_edit_pipeline_parity.py index 598e1619a..a727d7516 100644 --- a/tests/local_tests/glm_image/test_edit_pipeline_parity.py +++ b/tests/local_tests/glm_image/test_edit_pipeline_parity.py @@ -21,6 +21,8 @@ import pytest import torch from PIL import Image +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29520") os.environ.setdefault("DISABLE_SP", "1") @@ -28,7 +30,8 @@ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) CONDITION_IMAGE = REPO_ROOT / "assets" / "images" / "couple.jpg" diff --git a/tests/local_tests/glm_image/test_pipeline_parity.py b/tests/local_tests/glm_image/test_pipeline_parity.py index f6f0fc85c..2ce23a6e6 100644 --- a/tests/local_tests/glm_image/test_pipeline_parity.py +++ b/tests/local_tests/glm_image/test_pipeline_parity.py @@ -8,6 +8,8 @@ import numpy as np import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") @@ -15,7 +17,8 @@ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" diff --git a/tests/local_tests/glm_image/test_t5_parity.py b/tests/local_tests/glm_image/test_t5_parity.py index c25772ab5..ef189eca3 100644 --- a/tests/local_tests/glm_image/test_t5_parity.py +++ b/tests/local_tests/glm_image/test_t5_parity.py @@ -18,13 +18,16 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) TEXT_ENCODER_DIR = LOCAL_WEIGHTS_DIR / "text_encoder" TOKENIZER_DIR = LOCAL_WEIGHTS_DIR / "tokenizer" diff --git a/tests/local_tests/glm_image/test_transformer_parity.py b/tests/local_tests/glm_image/test_transformer_parity.py index 23e5b22cc..95aac33e0 100644 --- a/tests/local_tests/glm_image/test_transformer_parity.py +++ b/tests/local_tests/glm_image/test_transformer_parity.py @@ -16,6 +16,8 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") @@ -23,7 +25,8 @@ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" diff --git a/tests/local_tests/glm_image/test_transformer_strict_load.py b/tests/local_tests/glm_image/test_transformer_strict_load.py index 93299aa92..0aab4d32b 100644 --- a/tests/local_tests/glm_image/test_transformer_strict_load.py +++ b/tests/local_tests/glm_image/test_transformer_strict_load.py @@ -20,6 +20,8 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") @@ -27,7 +29,8 @@ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) TRANSFORMER_DIR = LOCAL_WEIGHTS_DIR / "transformer" diff --git a/tests/local_tests/glm_image/test_vae_parity.py b/tests/local_tests/glm_image/test_vae_parity.py index 614514e08..673e724c8 100644 --- a/tests/local_tests/glm_image/test_vae_parity.py +++ b/tests/local_tests/glm_image/test_vae_parity.py @@ -18,13 +18,16 @@ from pathlib import Path import pytest import torch +import fastvideo.envs as envs + os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29519") os.environ.setdefault("DISABLE_SP", "1") REPO_ROOT = Path(__file__).resolve().parents[3] FAMILY = "glm_image" -LOCAL_WEIGHTS_DIR = Path(os.getenv("GLM_IMAGE_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / FAMILY)) +LOCAL_WEIGHTS_DIR = Path(envs.FASTVIDEO_TEST_GLM_IMAGE_LOCAL_WEIGHTS_DIR.get() or + REPO_ROOT / "official_weights" / FAMILY) VAE_DIR = LOCAL_WEIGHTS_DIR / "vae" diff --git a/tests/local_tests/minimax_h3/README.md b/tests/local_tests/minimax_h3/README.md index ffdf94155..d6b0d3dfa 100644 --- a/tests/local_tests/minimax_h3/README.md +++ b/tests/local_tests/minimax_h3/README.md @@ -45,14 +45,14 @@ PYTHONPATH="${MINIMAX_H3_OFFICIAL_REF_DIR:-$PWD/DiffusersMiniMaxH3}/src:$PWD" py ## Checkpoint component parity ```bash -export MINIMAX_H3_MODEL_ROOT=/path/to/MiniMax-H3 +export FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT=/path/to/MiniMax-H3 export MINIMAX_H3_OFFICIAL_REF_DIR=/path/to/DiffusersMiniMaxH3 PYTHONPATH="$MINIMAX_H3_OFFICIAL_REF_DIR/src:$PWD" \ MINIMAX_H3_RUN_ENCODER_PARITY=1 \ pytest tests/local_tests/encoders/test_minimax_h3_qwen3_vl_parity.py -v -s -MINIMAX_H3_RUN_NVFP4_PARITY=1 MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ +MINIMAX_H3_RUN_NVFP4_PARITY=1 FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ MINIMAX_H3_NVFP4_TEXT_ENCODER=/path/to/FastH3-text-encoder-nvfp4 \ pytest tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py -s @@ -116,12 +116,12 @@ sum of rank-local maxima, not a simultaneous node peak. ```bash python tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \ - --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \ + --source-root "$PWD" --model-root "$FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT" \ --revision-label candidate --operation encode python -m torch.distributed.run --nproc_per_node=4 \ tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \ - --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \ + --source-root "$PWD" --model-root "$FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT" \ --revision-label candidate-sp4 --operation decode ``` diff --git a/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py b/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py index b47132898..e98a8f2cc 100644 --- a/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py +++ b/tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py @@ -11,7 +11,7 @@ noise stays above the thresholds. Run on a Blackwell GPU:: MINIMAX_H3_RUN_NVFP4_PARITY=1 \ - MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ + FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT=/path/to/FastH3 \ MINIMAX_H3_NVFP4_TEXT_ENCODER=/path/to/FastH3-nvfp4/text_encoder \ pytest tests/local_tests/minimax_h3/test_minimax_h3_text_encoder_nvfp4_parity.py -s """ @@ -26,6 +26,7 @@ from types import SimpleNamespace import pytest import torch +import fastvideo.envs as envs from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConfig from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4LinearMethod @@ -60,10 +61,10 @@ def _require_assets() -> tuple[torch.device, Path, Path]: pytest.skip("set MINIMAX_H3_RUN_NVFP4_PARITY=1 on a Blackwell GPU node") if not torch.cuda.is_available(): pytest.fail("MiniMax-H3 NVFP4 parity requires a CUDA GPU", pytrace=False) - root = os.environ.get("MINIMAX_H3_MODEL_ROOT") + root = envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() converted = os.environ.get("MINIMAX_H3_NVFP4_TEXT_ENCODER") if not root or not converted: - pytest.fail("set MINIMAX_H3_MODEL_ROOT and MINIMAX_H3_NVFP4_TEXT_ENCODER", pytrace=False) + pytest.fail("set FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT and MINIMAX_H3_NVFP4_TEXT_ENCODER", pytrace=False) root_path, converted_path = Path(root), Path(converted) missing = [str(p) for p in (root_path / "text_encoder", root_path / "tokenizer", converted_path) if not p.is_dir()] if missing: diff --git a/tests/local_tests/minimax_h3/test_mlx_audio_vae_parity.py b/tests/local_tests/minimax_h3/test_mlx_audio_vae_parity.py index f550c3374..f9940291e 100644 --- a/tests/local_tests/minimax_h3/test_mlx_audio_vae_parity.py +++ b/tests/local_tests/minimax_h3/test_mlx_audio_vae_parity.py @@ -9,13 +9,12 @@ case validates the released FP32 checkpoint end to end. pytest tests/local_tests/minimax_h3/test_mlx_audio_vae_parity.py -v The real-weight case activates automatically when the released snapshot is -present; point MINIMAX_H3_MODEL_ROOT at another checkout otherwise. +present; point FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT at another checkout otherwise. """ from __future__ import annotations import json -import os from pathlib import Path import numpy as np @@ -24,6 +23,7 @@ import pytest torch = pytest.importorskip("torch", reason="PyTorch reference needed for audio VAE parity") mx = pytest.importorskip("mlx.core", reason="MLX needed for audio VAE parity") +import fastvideo.envs as envs # noqa: E402 from fastvideo.configs.models.vaes.minimax_h3_audio import ( # noqa: E402 MiniMaxH3AudioVAEArchConfig, MiniMaxH3AudioVAEConfig, @@ -256,7 +256,7 @@ def test_audio_loader_validates_decoder_weights_before_construction(tmp_path) -> # Real released weights (bounded segment) # --------------------------------------------------------------------------- -DEFAULT_MODEL_ROOT = Path(os.environ.get("MINIMAX_H3_MODEL_ROOT", Path.home() / "models/FastH3-Preview-v0.2")) +DEFAULT_MODEL_ROOT = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() or Path.home() / "models/FastH3-Preview-v0.2") REAL_WEIGHTS_PRESENT = (DEFAULT_MODEL_ROOT / "audio_vae").is_dir() and any( (DEFAULT_MODEL_ROOT / "audio_vae").glob("*.safetensors")) diff --git a/tests/local_tests/minimax_h3/test_mlx_video_vae_parity.py b/tests/local_tests/minimax_h3/test_mlx_video_vae_parity.py index 23027e2e1..e6c213198 100644 --- a/tests/local_tests/minimax_h3/test_mlx_video_vae_parity.py +++ b/tests/local_tests/minimax_h3/test_mlx_video_vae_parity.py @@ -11,12 +11,11 @@ Runs in any environment with torch + mlx (CPU is fine): pytest tests/local_tests/minimax_h3/test_mlx_video_vae_parity.py -v The real-weight case activates automatically when the released snapshot is -present; point MINIMAX_H3_MODEL_ROOT at another checkout otherwise. +present; point FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT at another checkout otherwise. """ from __future__ import annotations -import os from pathlib import Path import numpy as np @@ -25,6 +24,7 @@ import pytest torch = pytest.importorskip("torch", reason="PyTorch reference needed for video VAE parity") mx = pytest.importorskip("mlx.core", reason="MLX needed for video VAE parity") +import fastvideo.envs as envs # noqa: E402 from fastvideo.configs.models.vaes.minimax_h3_video import ( # noqa: E402 MiniMaxH3VideoVAEArchConfig, MiniMaxH3VideoVAEConfig, @@ -228,7 +228,7 @@ def test_tiled_decode_matches_untiled() -> None: # Real released weights (bounded tile) # --------------------------------------------------------------------------- -DEFAULT_MODEL_ROOT = Path(os.environ.get("MINIMAX_H3_MODEL_ROOT", Path.home() / "models/FastH3-Preview-v0.2")) +DEFAULT_MODEL_ROOT = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() or Path.home() / "models/FastH3-Preview-v0.2") REAL_WEIGHTS_PRESENT = (DEFAULT_MODEL_ROOT / "vae").is_dir() and any( (DEFAULT_MODEL_ROOT / "vae").glob("*.safetensors")) diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_parity.py b/tests/local_tests/pipelines/test_zimage_pipeline_parity.py index 1b9ae21f9..b3aadddb2 100644 --- a/tests/local_tests/pipelines/test_zimage_pipeline_parity.py +++ b/tests/local_tests/pipelines/test_zimage_pipeline_parity.py @@ -16,6 +16,7 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs from fastvideo.configs.pipelines.zimage import ZImagePipelineConfig from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -32,7 +33,8 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch REPO_ROOT = Path(__file__).resolve().parents[3] REFERENCE_REVISION = "26f23eda626ffadda020b04ff79488e1d72004cd" REFERENCE_REPO = Path(os.getenv("ZIMAGE_REFERENCE_REPO", REPO_ROOT / "Z-Image")) -MODEL_DIR = Path(os.getenv("ZIMAGE_MODEL_DIR", REPO_ROOT / "official_weights" / "Z-Image")) +MODEL_DIR = Path(envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else + REPO_ROOT / "official_weights" / "Z-Image") class _ConstantTransformer: diff --git a/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py b/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py index 7a99e9829..61f4ccc49 100644 --- a/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py +++ b/tests/local_tests/pipelines/test_zimage_pipeline_smoke.py @@ -3,7 +3,6 @@ from __future__ import annotations -import os from contextlib import nullcontext from pathlib import Path from typing import Any, cast @@ -11,7 +10,9 @@ from typing import Any, cast import pytest import torch -MODEL_DIR_ENV = os.getenv("ZIMAGE_MODEL_DIR") +import fastvideo.envs as envs + +MODEL_DIR_ENV = envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.get() if envs.FASTVIDEO_TEST_ZIMAGE_MODEL_DIR.is_set() else None def test_zimage_typed_surface_preflight() -> None: @@ -106,7 +107,7 @@ def test_model_download_honors_pinned_revision(monkeypatch: pytest.MonkeyPatch, @pytest.mark.skipif(not torch.cuda.is_available(), reason="Z-Image load/generate smoke requires CUDA") def test_zimage_pipeline_load_generate_smoke() -> None: if MODEL_DIR_ENV is None: - pytest.skip("Set ZIMAGE_MODEL_DIR to activate the real load/generate smoke") + pytest.skip("Set FASTVIDEO_TEST_ZIMAGE_MODEL_DIR to activate the real load/generate smoke") model_dir = Path(MODEL_DIR_ENV) if not (model_dir / "model_index.json").is_file(): pytest.skip(f"Z-Image model_index.json not found under {model_dir}") diff --git a/tests/local_tests/sd35/README.md b/tests/local_tests/sd35/README.md index bd5f61f8d..802927c16 100644 --- a/tests/local_tests/sd35/README.md +++ b/tests/local_tests/sd35/README.md @@ -16,7 +16,7 @@ checkpoint. Skipped in CI; CUDA required. | Official commit/version | `` | | HF weights | `stabilityai/stable-diffusion-3.5-medium` | | HF revision | `` | -| Local weights dir | `official_weights/stabilityai__stable-diffusion-3.5-medium` (env: `SD35_MODEL_DIR`) | +| Local weights dir | `official_weights/stabilityai__stable-diffusion-3.5-medium` (env: `FASTVIDEO_TEST_SD35_MODEL_DIR`) | | Source layout | `diffusers` | | Needs conversion | `` | diff --git a/tests/local_tests/sd35/test_sd35_component_parity.py b/tests/local_tests/sd35/test_sd35_component_parity.py index 8deb50e38..3689cf3b6 100644 --- a/tests/local_tests/sd35/test_sd35_component_parity.py +++ b/tests/local_tests/sd35/test_sd35_component_parity.py @@ -21,6 +21,8 @@ from transformers import T5EncoderModel as RefT5EncoderModel from diffusers import FlowMatchEulerDiscreteScheduler as RefScheduler from safetensors.torch import safe_open +import fastvideo.envs as envs + pytestmark = [ pytest.mark.skipif( not torch.cuda.is_available(), @@ -32,10 +34,8 @@ pytestmark = [ os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") hf_logging.set_verbosity_error() -MODEL_DIR = Path(os.getenv( - "SD35_MODEL_DIR", - "/FastVideo/official_weights/stabilityai__stable-diffusion-3.5-medium", -)) +MODEL_DIR = Path(envs.FASTVIDEO_TEST_SD35_MODEL_DIR.get() if envs.FASTVIDEO_TEST_SD35_MODEL_DIR.is_set() else + "/FastVideo/official_weights/stabilityai__stable-diffusion-3.5-medium") def _load_json(path: Path) -> dict: diff --git a/tests/local_tests/transformers/test_minimax_h3_transformer_parity.py b/tests/local_tests/transformers/test_minimax_h3_transformer_parity.py index 1a14f29c1..05491790f 100644 --- a/tests/local_tests/transformers/test_minimax_h3_transformer_parity.py +++ b/tests/local_tests/transformers/test_minimax_h3_transformer_parity.py @@ -20,10 +20,12 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs + REPO_ROOT = Path(__file__).resolve().parents[3] OFFICIAL_REF_DIR = Path(os.environ.get("MINIMAX_H3_OFFICIAL_REF_DIR", REPO_ROOT / "DiffusersMiniMaxH3")) OFFICIAL_SRC = OFFICIAL_REF_DIR / "src" -MODEL_ROOT = Path(os.environ.get("MINIMAX_H3_MODEL_ROOT", REPO_ROOT / "official_weights" / "MiniMax-H3")) +MODEL_ROOT = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() or REPO_ROOT / "official_weights" / "MiniMax-H3") RUN_ENV = "MINIMAX_H3_RUN_DIT_PARITY" PARITY_SCOPE = "production_loader" diff --git a/tests/local_tests/vaes/test_minimax_h3_audio_vae_parity.py b/tests/local_tests/vaes/test_minimax_h3_audio_vae_parity.py index d6f1d11bb..2848b3d7a 100644 --- a/tests/local_tests/vaes/test_minimax_h3_audio_vae_parity.py +++ b/tests/local_tests/vaes/test_minimax_h3_audio_vae_parity.py @@ -14,10 +14,12 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs + REPO_ROOT = Path(__file__).resolve().parents[3] OFFICIAL_REF_DIR = Path(os.environ.get("MINIMAX_H3_OFFICIAL_REF_DIR", REPO_ROOT / "DiffusersMiniMaxH3")) OFFICIAL_SRC = OFFICIAL_REF_DIR / "src" -MODEL_ROOT = Path(os.environ.get("MINIMAX_H3_MODEL_ROOT", REPO_ROOT / "official_weights" / "MiniMax-H3")) +MODEL_ROOT = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() or REPO_ROOT / "official_weights" / "MiniMax-H3") RUN_ENV = "MINIMAX_H3_RUN_AUDIO_VAE_PARITY" PARITY_SCOPE = "production_loader" diff --git a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py index 98308a441..3893386cf 100644 --- a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py +++ b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py @@ -14,10 +14,12 @@ import pytest import torch from torch.testing import assert_close +import fastvideo.envs as envs + REPO_ROOT = Path(__file__).resolve().parents[3] OFFICIAL_REF_DIR = Path(os.environ.get("MINIMAX_H3_OFFICIAL_REF_DIR", REPO_ROOT / "DiffusersMiniMaxH3")) OFFICIAL_SRC = OFFICIAL_REF_DIR / "src" -MODEL_ROOT = Path(os.environ.get("MINIMAX_H3_MODEL_ROOT", REPO_ROOT / "official_weights" / "MiniMax-H3")) +MODEL_ROOT = Path(envs.FASTVIDEO_TEST_MINIMAX_H3_MODEL_ROOT.get() or REPO_ROOT / "official_weights" / "MiniMax-H3") RUN_ENV = "MINIMAX_H3_RUN_VIDEO_VAE_PARITY" PARITY_SCOPE = "production_loader" diff --git a/tests/local_tests/zimage/PORT_STATUS.md b/tests/local_tests/zimage/PORT_STATUS.md index b66f631ac..1a3b7d085 100644 --- a/tests/local_tests/zimage/PORT_STATUS.md +++ b/tests/local_tests/zimage/PORT_STATUS.md @@ -87,7 +87,7 @@ intermediate size 9728, 32 attention heads, 8 KV heads, and head dimension | Complete component closure | `ZIMAGE_OFFICIAL_REF_DIR= ZIMAGE_TRANSFORMER_DIR= pytest tests/local_tests/zimage/ -v -s` | 34 PASSED, no skips, GB200, 2026-07-15, `f310aacb1` | Official source `26f23eda...`, HF snapshot `f332072a...`; production/native encoder, scheduler, tokenizer, transformer, and VAE all passed. | | Pipeline typed surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS, hardware-free, 2026-07-14 | Covers exact `EntryClass`, config, preset, registry, model detection, and official defaults. | | Native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS, hardware-free, 2026-07-14 | Covers conditioning trim, endpoint-preserving timestep setup, denoising sign/timestep math, and VAE scale/shift. | -| Full pipeline parity | `ZIMAGE_REFERENCE_REPO= ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS, RTX 5090, `fec2baa6c` | Real FastVideo one-step latents match pinned native Tongyi exactly (`max=0`, `mean=0`). Later commits through `f310aacb1` change only API inventory/tests, not inference math. | +| Full pipeline parity | `ZIMAGE_REFERENCE_REPO= FASTVIDEO_TEST_ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS, RTX 5090, `fec2baa6c` | Real FastVideo one-step latents match pinned native Tongyi exactly (`max=0`, `mean=0`). Later commits through `f310aacb1` change only API inventory/tests, not inference math. | | PNG SSIM quality | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | L40S REFERENCE UPLOADED; END-TO-END SSIM CI PENDING | Dataset revision `7b27d00004b723ef0cf4923e77ae57128573f01c`; remote SHA-256 `6c2be648...f7a6`; upload scoped to `Tongyi-MAI__Z-Image-Turbo/TORCH_SDPA`. | ## Open Questions diff --git a/tests/local_tests/zimage/README.md b/tests/local_tests/zimage/README.md index 223b2b0f9..e7c24ff3c 100644 --- a/tests/local_tests/zimage/README.md +++ b/tests/local_tests/zimage/README.md @@ -127,7 +127,7 @@ pytest tests/local_tests/zimage/ -v -s |---|---|---| | Hardware-free pipeline surface | `pytest tests/local_tests/pipelines/test_zimage_pipeline_smoke.py::test_zimage_typed_surface_preflight -v` | PASS | | Hardware-free native stage contract | `pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_native_default_stage_math -v` | PASS | -| Pinned native full GPU pipeline parity | `ZIMAGE_REFERENCE_REPO= ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS on RTX 5090; one-step latents bit-exact (`max=0`, `mean=0`) | +| Pinned native full GPU pipeline parity | `ZIMAGE_REFERENCE_REPO= FASTVIDEO_TEST_ZIMAGE_MODEL_DIR= DISABLE_SP=1 pytest tests/local_tests/pipelines/test_zimage_pipeline_parity.py::test_zimage_pipeline_latents_match_pinned_native_repo -v -s` | PASS on RTX 5090; one-step latents bit-exact (`max=0`, `mean=0`) | | 8-step 1024x1024 PNG SSIM | `pytest fastvideo/tests/ssim/test_zimage_similarity.py -v -s` | Reviewed L40S PNG uploaded; remote SHA-256 `6c2be648...f7a6`; end-to-end SSIM CI pending | ## Component contracts From 8322b01815969b5e9f5a31565a4cf0930f3b9367 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Thu, 1 Oct 2026 22:17:32 -0700 Subject: [PATCH 043/151] [bugfix]: restore the NVFP4 module after test_nvfp4_config re-imports it (#1908) Co-authored-by: aryan5v --- .../ops/quantization/test_nvfp4_config.py | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_config.py b/fastvideo/tests/ops/quantization/test_nvfp4_config.py index a07fb7eb5..977528189 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_config.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_config.py @@ -8,12 +8,21 @@ clear error when flashinfer is missing. """ from __future__ import annotations +import importlib import sys import types import pytest +def _clear_nvfp4_module(monkeypatch): + # Load the original first, even when this test runs in isolation. + module = importlib.import_module("fastvideo.layers.quantization.nvfp4_config") + # A fresh import replaces both bindings; restore both at teardown. + monkeypatch.delattr(sys.modules[module.__package__], "nvfp4_config") + monkeypatch.delitem(sys.modules, module.__name__) + + def test_nvfp4config_imports_without_flashinfer(monkeypatch): """Importing the module on a host without flashinfer must succeed. @@ -24,7 +33,7 @@ def test_nvfp4config_imports_without_flashinfer(monkeypatch): # Hide flashinfer from sys.modules and the import path. monkeypatch.setitem(sys.modules, "flashinfer", None) # Force a re-import of the target module. - sys.modules.pop("fastvideo.layers.quantization.nvfp4_config", None) + _clear_nvfp4_module(monkeypatch) from fastvideo.layers.quantization.nvfp4_config import NVFP4Config config = NVFP4Config() assert config.get_name() == "nvfp4" @@ -45,7 +54,7 @@ def test_nvfp4_kernel_call_raises_clear_error_without_flashinfer(monkeypatch): AttributeError or NameError.""" # Stage a fake flashinfer that fails on import. monkeypatch.setitem(sys.modules, "flashinfer", _raise_module_on_import("flashinfer")) - sys.modules.pop("fastvideo.layers.quantization.nvfp4_config", None) + _clear_nvfp4_module(monkeypatch) from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer with pytest.raises(ImportError, match="flashinfer"): _require_flashinfer() @@ -86,3 +95,22 @@ def test_coerce_fp4_input_dtype_casts_and_rejects(): for bad in (torch.int32, torch.int64, torch.bool): with pytest.raises(TypeError, match="floating-point"): _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=bad)) + + +@pytest.mark.parametrize("import_test", [ + test_nvfp4config_imports_without_flashinfer, + test_nvfp4_kernel_call_raises_clear_error_without_flashinfer, +]) +def test_nvfp4_import_tests_restore_module_identity(import_test): + original = importlib.import_module("fastvideo.layers.quantization.nvfp4_config") + + with pytest.MonkeyPatch.context() as monkeypatch: + import_test(monkeypatch) + assert sys.modules[original.__name__] is not original + + from fastvideo.layers.quantization import nvfp4_config + from fastvideo.layers.quantization.nvfp4_config import NVFP4Config + + assert sys.modules[original.__name__] is original + assert nvfp4_config is original + assert NVFP4Config is original.NVFP4Config From 6809a751fb3a6a66d95367875f4fa93ef57d7a9c Mon Sep 17 00:00:00 2001 From: Ishan <34205320+Ishxn20@users.noreply.github.com> Date: Thu, 1 Oct 2026 22:17:46 -0700 Subject: [PATCH 044/151] [bugfix] Fixed the MiniMax-H3 pin-fallback decode test setup (#1909) --- .../tests/stages/test_minimax_h3_vae_streaming.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 74510139f..c7f0f7bbd 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -156,9 +156,15 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None: 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) - MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward( + MiniMaxH3VideoDecodingStage(VAE()).forward( batch, - SimpleNamespace(output_type="pil", pin_cpu_memory=True, vae_cpu_offload=False, vae_parallel_decode=False), + SimpleNamespace( + output_type="pil", + pin_cpu_memory=True, + vae_cpu_offload=False, + vae_parallel_decode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ), ) assert observed == { From 9491c8638adeb266dd37656ba6224b2200db91ea Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Thu, 1 Oct 2026 22:22:17 -0700 Subject: [PATCH 045/151] [bugfix]: keep the full audio track when saving through the ffmpeg pipe (#1910) --- fastvideo/entrypoints/video_generator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 2f8468a68..5bdc819b7 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -1256,10 +1256,10 @@ class VideoGenerator: envs.FASTVIDEO_OUTPUT_PIX_FMT.get(), # Audio and video decoders may produce different durations. # Preserve every video frame and pad/trim audio to match. + # No -frames:v: ffmpeg 4.4 closes every stream once the video + # reaches that count, which cuts the padded audio short. "-af", "apad", - "-frames:v", - str(len(frames)), "-t", str(len(frames) / fps), "-movflags", From 8444c0897a8b96848eb85b6e5750ef486f79fc92 Mon Sep 17 00:00:00 2001 From: William Lin Date: Fri, 2 Oct 2026 11:28:07 -0700 Subject: [PATCH 046/151] [feat] FastH3 Ref2VA PDD inference with reference-video VSA sparsity (#1907) Co-authored-by: Davids048 Co-authored-by: Claude Opus 5.5 --- .buildkite/scripts/unit_test.sh | 4 + .../inference_schema_parity_inventory.yaml | 2 + docs/inference/fasth3-distilled.md | 124 ++++ docs/inference/support_matrix.md | 5 +- examples/inference/basic/README.md | 18 + .../basic/basic_fasth3_omniref_pdd.py | 268 ++++++++ .../backends/video_sparse_attn_h3.py | 288 ++++++--- fastvideo/configs/models/dits/minimax_h3.py | 8 + fastvideo/configs/pipelines/base.py | 26 +- fastvideo/configs/pipelines/minimax_h3.py | 229 ++++++- fastvideo/entrypoints/video_generator.py | 8 +- fastvideo/fastvideo_args.py | 19 +- fastvideo/layers/pdd.py | 312 ++++++++++ fastvideo/models/dits/minimax_h3.py | 65 +- .../basic/minimax_h3/minimax_h3_pipeline.py | 126 +++- .../pipelines/basic/minimax_h3/packing.py | 16 + .../minimax_h3/stages/minimax_h3_denoising.py | 138 ++++- .../tests/attention/test_vsa_h3_backward.py | 3 +- .../tests/attention/test_vsa_h3_metadata.py | 29 +- .../attention/test_vsa_h3_ref2va_regions.py | 572 ++++++++++++++++++ .../attention/test_vsa_h3_sm100a_route.py | 6 +- .../attention/test_vsa_h3_tile_grad_safety.py | 3 +- .../tests/entrypoints/test_video_generator.py | 34 +- .../test_basic_fasth3_omniref_pdd.py | 122 ++++ fastvideo/tests/layers/test_pdd_linear.py | 209 +++++++ .../loader/test_minimax_h3_pdd_loading.py | 231 +++++++ .../tests/mlx/test_mlx_minimax_h3_vsa.py | 5 +- .../stages/test_minimax_h3_pdd_contract.py | 472 +++++++++++++++ .../stages/test_minimax_h3_pdd_denoising.py | 327 ++++++++++ fastvideo/tests/test_partial_hf_download.py | 39 +- fastvideo/utils.py | 26 + 31 files changed, 3582 insertions(+), 152 deletions(-) create mode 100644 examples/inference/basic/basic_fasth3_omniref_pdd.py create mode 100644 fastvideo/layers/pdd.py create mode 100644 fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py create mode 100644 fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py create mode 100644 fastvideo/tests/layers/test_pdd_linear.py create mode 100644 fastvideo/tests/loader/test_minimax_h3_pdd_loading.py create mode 100644 fastvideo/tests/stages/test_minimax_h3_pdd_contract.py create mode 100644 fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 7fa4b3799..aacba0fcc 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -15,8 +15,12 @@ exec pytest \ ./fastvideo/tests/ops/ \ ./fastvideo/tests/worker/ \ ./fastvideo/tests/training/test_trackers.py \ + ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \ ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \ ./fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py \ + ./fastvideo/tests/attention/test_vsa_h3_metadata.py \ + ./fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py \ + ./fastvideo/tests/layers/test_pdd_linear.py \ ./fastvideo/tests/modal/test_kernel_build_cache.py \ ./fastvideo/tests/modal/test_pr_test.py \ ./fastvideo/tests/modal/test_ssim_test.py \ diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 463977b4c..4309e9851 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -470,12 +470,14 @@ surfaces: coords_style: "MagiHuman internal data-proxy coordinate convention." frame_receptive_field: "MagiHuman internal data-proxy receptive-field setting." image_conditioning: "MagiHuman preset variant marker for reference-image conditioning." + pdd_step_indices: "MiniMax-H3 PDD fine-grid partition; resolve_checkpoint_settings sets it from fastvideo_inference.json." ref_audio_offset: "MagiHuman internal data-proxy audio alignment offset." scheduler_sigma_min: "Z-Image scheduler parity invariant; not part of the public typed inference API." scheduler_use_reference_discrete_timesteps: "Z-Image scheduler parity invariant; not part of the public typed inference API." sr_local_attn_layers: "MagiHuman SR internal sparse-attention layer selection." text_offset: "MagiHuman internal data-proxy text alignment offset." vae_stride: "MagiHuman internal VAE/data-proxy stride setting." + vsa_ref_keep_rate: "MiniMax-H3 PDD reference-video VSA keep rate; defaults to the checkpoint's trained value." z_dim: "MagiHuman internal VAE latent channel setting." vocoder_config: "Legacy internal component config object." vocoder_precision: "Precision override pending dedicated component precision design." diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index 27035b92f..d77f68bd5 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -74,6 +74,130 @@ This documents execution support for the published checkpoint. It is not a quality claim: compare video/audio output against base MiniMax-H3 on your own prompts before adopting it. +## Ref2VA PDD students + +A Parallel Decoding Distillation (PDD) student widens the transformer's two +output projections to `pdd_steps` heads, one per interval of a fixed fine +time grid on `[0, 0.999]`. Each transformer forward fuses a block of +consecutive heads into their integration-weighted mean, and an ordinary Euler +step over the block's two node sigmas applies it. The FastH3 OmniRef PDD-8 +student is a Ref2VA (`transformer_ref`) student with 32 heads, sampled in +eight blocks of four. + +Its export carries only what distillation changed: + +| Path | Contents | +| --- | --- | +| `transformer_ref/` | Widened `proj_out` and `audio_proj_out`, trained VSA compression gates; `config.json` records `pdd_steps` | +| `scheduler/`, `audio_scheduler/` | Video and audio shifts (12 and 3) | +| `fastvideo_inference.json` | The sampling contract below | +| `modular_model_index.json` | The Diffusers manifest | + +The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's. +`basic_fasth3_omniref_pdd.py` composes the two into one local directory of +symlinks and runs the recipe the contract records. `--model-path` is the +export, as a local directory or a Hugging Face repo id. The base comes from +the revision the contract pins in `base_model_revision`, or from +`--base-model-path`: + +```bash +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path \ + --video reference.mp4 --image character.png \ + --prompt 'The dancer from the video performs the routine in the pictured outfit.' \ + --height 480 --width 832 --num-frames 124 \ + --output outputs/fasth3-omniref-pdd +``` + +References are ordered: pass `--image`, `--video`, and `--audio` in the order +the prompt refers to them. At least one image or video is required. + +A PDD export uses `fasth3-inference-contract-v1` with PDD fields in place of +the DMD ladder: + +```json +{ + "schema_version": "fasth3-inference-contract-v1", + "model_type": "ref2va", + "transformer_component": "transformer_ref", + "pdd_steps": 32, + "pdd_step_indices": [0, 4, 8, 12, 16, 20, 24, 28, 32], + "num_inference_steps": 8, + "transformer_forwards": 8, + "grid_max_t": 0.999, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, + "vsa_ref_policy": "p2_multi_region", + "vsa_ref_keep_rate": 0.1 +} +``` + +The export also records `schema`, `conditioning`, and `base_model_revision`, +the base snapshot it was distilled against, as `hf://@` +(for example `hf://MiniMaxAI/MiniMax-H3@`); any other form is an +error. FastVideo reads the file once, when the run's arguments are built and +before any weights load. The file must carry exactly these fields: a missing +or unknown field is an error. Fields that repeat a value stored elsewhere must +equal it: `pdd_steps` equals `transformer_ref/config.json`, the shifts equal +the scheduler configs, and `num_inference_steps` and `transformer_forwards` +equal the block count of `pdd_step_indices`, a strictly increasing partition +of the grid from 0 to `pdd_steps`. Only `MiniMaxH3Ref2VAModularPipeline` runs +the export. A `transformer_ref` whose `config.json` sets `pdd_steps` without +this file is an error. + +Each sampling setting has one source, the file: + +| Setting | Run leaves it unset | Run sets a different value | +| ------------------------------- | ---------------------- | ---------------------------------------------- | +| `pdd_step_indices` | Taken from the file | Error | +| `num_inference_steps` (request) | Set to the block count | Error | +| `attention_backend` | Taken from the file | Error, including `FASTVIDEO_ATTENTION_BACKEND` | +| `VSA_tile_size` | Taken from the file | Error | +| `VSA_sparsity` | Taken from the file | The run's value, with a warning | +| `vsa_ref_keep_rate` | Taken from the file | The run's value, with a warning | + +A request leaves `num_inference_steps` unset only when it is parsed from a +mapping or a config file. A `GenerationRequest` built in Python counts every +field as set, so it must pass the block count. The trained compression gates +of `transformer_ref` load only under `VIDEO_SPARSE_ATTN_H3`, so the attention +backend cannot change. `dmd_denoising_steps` must stay unset. + +A MiniMax-H3 checkpoint without PDD fields, such as a DMD export, whose +transformer carries VSA compression gates (`to_gate_compress` weights) also +runs only with `VIDEO_SPARSE_ATTN_H3`. The pipeline reads the shard index (or +the safetensors headers) and rejects any other backend, including automatic +selection, before it loads any component. + +### Reference-video sparsity + +Every PDD contract sets `"vsa_ref_policy": "p2_multi_region"`, the only value FastVideo accepts, which means: +`VIDEO_SPARSE_ATTN_H3` tiles every reference video as its own sparse region, +in place in the packed sequence. Each video query keeps `vsa_ref_keep_rate` of +every reference video's tiles and `1 - VSA_sparsity` of the target video's +tiles. Text, audio, and image references stay dense. With `VSA_sparsity` 0, +every region is dense and the reference keep rate has no effect. Checkpoints +other than PDD students keep every conditioning row dense. + +### Hardware + +The contract's 128-token tiles, `(4, 4, 8)`, run only on the sm_100a/sm_103a +CUDA block-sparse kernel (B200, B300, GB200, GB300) of a fastvideo-kernel +build with the Blackwell VSA extension. There is no Triton fallback for tile +128: the backend raises instead. Tile 128 runs eagerly. Regional compile +requires 64-token tiles. With `--num-gpus` above 1 the example shards the DiT +across the GPUs (FSDP) and splits the sequence across them. + +On GB200, a 480x832, 124-frame request fits on one GPU: the eight forwards +take about 54 s, and device memory in use peaks near 103 GiB. A 768x1344, +345-frame request with `--num-gpus 4` takes about 36 s for the eight +forwards. With the DiT sharded, peak allocated memory is about 66 GiB per GPU +(about 95 GiB in use), against 83 GiB (103 GiB) with a full DiT copy on each +GPU; the output is identical. Model loading and decoding come on top of this. + ## Apple Silicon `mlx_fasth3.py` stays on FastH3 V1 and its uniform AdaLN cache. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 9b2bc90ec..5fb510219 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -101,7 +101,10 @@ to convert the official weights locally and set `MMAUDIO_MODEL_PATH`. **Note (MiniMax H3)**: T2VA, FL2VA, and Ref2VA all generate video with stereo audio. Use the Ref2VA example when passing ordered image, video, or audio -references. +references. Distilled Ref2VA PDD students (eight transformer forwards, with +reference videos as sparse VSA regions) run through +[basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py); +see [FastH3 distilled checkpoint schedules](fasth3-distilled.md#ref2va-pdd-students). **Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered model ID exists on `main` diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 7c948bd9b..fcc5b0696 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -167,6 +167,24 @@ launchers still use sparse attention at strength `0` and require FastVideo's tile-64 VSA kernel; the dense launcher selects FA4. Each launcher writes to its own variant directory by default so comparison outputs do not collide. +### FastH3 OmniRef PDD (Ref2VA) + +[basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py) +runs a Parallel Decoding Distillation Ref2VA student in eight transformer +forwards. The export carries only its `transformer_ref`, scheduler configs, +and `fastvideo_inference.json`; the script links them with the base +MiniMax-H3 components into one local model directory. Its 128-token VSA tiles +need the sm_100a/sm_103a kernel: + +```bash +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path \ + --video reference.mp4 --image character.png --prompt "your prompt" +``` + +See [FastH3 distilled checkpoint schedules](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/inference/fasth3-distilled.md#ref2va-pdd-students) +for the contract and the reference-video sparsity policy. + ## Basic Walkthrough All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines! diff --git a/examples/inference/basic/basic_fasth3_omniref_pdd.py b/examples/inference/basic/basic_fasth3_omniref_pdd.py new file mode 100644 index 000000000..8f9673c82 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_omniref_pdd.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Eight-forward Ref2VA video+audio generation with a FastH3 OmniRef PDD student. + +A FastH3 OmniRef PDD export is a Parallel Decoding Distillation (PDD) student of +MiniMax-H3's reference-conditioned ``transformer_ref`` partition. Its directory +(local, or a Hugging Face repo) carries only what distillation changed: +``transformer_ref/`` (output heads widened to the fine grid, trained VSA +compression gates), the two scheduler configs, and ``fastvideo_inference.json``. +The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's. + +This script composes the two into one local model directory of symlinks and +runs the recipe the contract records: fused blocks of the fine grid (one +transformer forward each), the video/audio shifts, and VIDEO_SPARSE_ATTN_H3 +with its sparsity and tile size, where every reference video is its own +sparse region. The contract's 128-token tiles run only on the sm_100a/sm_103a +CUDA kernel (B200/B300/GB200/GB300) of a fastvideo-kernel build with the +Blackwell VSA extension; the VSA-H3 backend raises at the first attention call +when either is missing. +With ``--num-gpus`` above 1 the DiT is sharded across the GPUs (FSDP) and the +sequence is split across them (sequence parallelism). + +References are ordered. Pass them in order with --image / --video / --audio, +for example ``--video dance.mp4 --image outfit.png --audio voice.wav``. +""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +from fastvideo import VideoGenerator +from fastvideo.api import ( + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) +from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision +from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference + +CONTRACT = "fastvideo_inference.json" +# Components the distilled export replaces, and the ones it shares with base MiniMax-H3. +EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler") +BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae") +MANIFESTS = ("modular_model_index.json", "model_index.json") +# Components the composed directory's manifest must declare, besides its transformer_ref. +MANIFEST_COMPONENTS = ("scheduler", "audio_scheduler", *BASE_COMPONENTS) + + +class _AppendReference(argparse.Action): + """Collect --image/--video/--audio into one list that keeps command-line order.""" + + def __call__(self, parser, namespace, value, option_string=None): + references = list(getattr(namespace, self.dest, None) or []) + references.append((self.const, value)) + setattr(namespace, self.dest, references) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", + required=True, + help="FastH3 OmniRef PDD export: a local directory or a Hugging Face repo id") + parser.add_argument("--revision", default=None, help="revision of --model-path when it is a Hugging Face repo") + parser.add_argument("--base-model-path", + default=None, + help="base MiniMax-H3 snapshot (local directory or repo id). Default: the repo and revision " + "pinned by the export's base_model_revision") + parser.add_argument("--base-revision", + default=None, + help="revision of --base-model-path when it is a repo id (the export's pinned revision " + "applies only to its own base repo); without --base-model-path, overrides that pin") + parser.add_argument("--composed-dir", + default=None, + help="where to write the composed model directory of symlinks (default: " + "/fasth3_omniref_model)") + for flag in ("image", "video", "audio"): + parser.add_argument(f"--{flag}", + dest="references", + action=_AppendReference, + const=flag, + metavar="PATH", + help=f"an ordered {flag} reference (repeatable)") + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", default="outputs/fasth3_omniref_pdd") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--num-gpus", type=int, default=1, help="sequence-parallel GPUs (must divide 56 heads)") + return parser + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = build_parser() + args = parser.parse_args(argv) + if not args.references: + parser.error("pass at least one ordered reference with --image, --video, or --audio") + if not any(kind != "audio" for kind, _ in args.references): + parser.error("Ref2VA needs at least one image or video reference") + return args + + +def _snapshot(path_or_repo: str, revision: str | None, allow_patterns: list[str]) -> Path: + local = Path(path_or_repo).expanduser() + if local.is_dir(): + return local + from huggingface_hub import snapshot_download + + return Path(snapshot_download(repo_id=path_or_repo, revision=revision, allow_patterns=allow_patterns)) + + +def load_contract(export_dir: Path) -> dict[str, Any]: + path = export_dir / CONTRACT + if not path.is_file(): + raise FileNotFoundError(f"{export_dir} has no {CONTRACT}; it is not a FastH3 distilled export.") + contract = json.loads(path.read_text(encoding="utf-8")) + if "pdd_steps" not in contract or contract.get("model_type") != "ref2va": + raise ValueError(f"{path} is not a Ref2VA PDD contract; use the example matching the checkpoint.") + return contract + + +def base_model_source(args: argparse.Namespace, contract: dict[str, Any]) -> tuple[str, str | None]: + """The base snapshot and revision: --base-model-path as given, else the export's pin.""" + base_repo, base_revision = base_model_from_contract(contract) + if args.base_model_path: + return args.base_model_path, args.base_revision + return base_repo, args.base_revision or base_revision + + +def base_model_from_contract(contract: dict[str, Any]) -> tuple[str, str]: + """The base repo id and revision that the export pins as ``hf://@``.""" + return parse_base_model_revision(contract.get("base_model_revision")) + + +def _link(destination: Path, source: Path) -> None: + if not source.exists(): + raise FileNotFoundError(f"Missing checkpoint component: {source}") + source = source.resolve() + if destination.is_symlink(): + if destination.resolve() == source: + return + destination.unlink() + elif destination.exists(): + raise FileExistsError(f"{destination} exists and is not a symlink; choose another --composed-dir.") + destination.symlink_to(source, target_is_directory=source.is_dir()) + + +def _declares_components(manifest: Path) -> bool: + """Whether a Diffusers manifest declares every component of the composed directory.""" + declared = { + name + for name, spec in json.loads(manifest.read_text(encoding="utf-8")).items() + if isinstance(spec, list) and spec and spec[0] is not None + } + return set(MANIFEST_COMPONENTS) <= declared and bool({"transformer", "transformer_ref"} & declared) + + +def select_manifest(export_dir: Path, base_dir: Path) -> Path: + """The export's manifest when it declares every composed component, else the base's.""" + for root in (export_dir, base_dir): + for name in MANIFESTS: + if (root / name).is_file() and _declares_components(root / name): + return root / name + raise FileNotFoundError("Neither the export nor the base snapshot has a Diffusers model manifest that declares " + f"{', '.join(MANIFEST_COMPONENTS)} and transformer_ref.") + + +def compose_model_dir(export_dir: Path, base_dir: Path, composed_dir: Path) -> Path: + """One MiniMax-H3 model directory: distilled components from the export, the rest from the base.""" + composed_dir.mkdir(parents=True, exist_ok=True) + for name in (*EXPORT_COMPONENTS, CONTRACT): + _link(composed_dir / name, export_dir / name) + manifest_source = select_manifest(export_dir, base_dir) + for name in MANIFESTS: + stale = composed_dir / name + if name != manifest_source.name and stale.is_symlink(): + stale.unlink() + _link(composed_dir / manifest_source.name, manifest_source) + for name in BASE_COMPONENTS: + _link(composed_dir / name, base_dir / name) + return composed_dir + + +def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]: + export_dir = _snapshot(args.model_path, args.revision, + [CONTRACT, *MANIFESTS, *(f"{name}/**" for name in EXPORT_COMPONENTS)]) + contract = load_contract(export_dir) + base_source, base_revision = base_model_source(args, contract) + base_dir = _snapshot(base_source, base_revision, [*MANIFESTS, *(f"{name}/**" for name in BASE_COMPONENTS)]) + composed_dir = Path(args.composed_dir) if args.composed_dir else Path(args.output) / "fasth3_omniref_model" + return compose_model_dir(export_dir, base_dir, composed_dir), contract + + +def build_generator_config(model_dir: Path, num_gpus: int) -> GeneratorConfig: + return GeneratorConfig( + model_path=str(model_dir), + engine=EngineConfig( + num_gpus=num_gpus, + # Shard the DiT across the GPUs rather than holding a full copy on each. + use_fsdp_inference=num_gpus > 1, + parallelism=ParallelismConfig(tp_size=1, sp_size=num_gpus), + offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False), + ), + pipeline=PipelineSelection( + workload_type="i2v", + # FastVideo reads the trained fused-block partition and attention + # settings from the composed directory's fastvideo_inference.json. + components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"), + ), + ) + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + model_dir, contract = resolve_model(args) + print(f"Composed model directory: {model_dir}") + print(f"Contract: {contract['num_inference_steps']} fused blocks of a {contract['pdd_steps']}-interval grid, " + f"{contract['attention_backend']} with sparsity {contract['vsa_sparsity']}, " + f"{contract['vsa_tile_size']}-token tiles, reference keep rate {contract['vsa_ref_keep_rate']}") + references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in args.references] + + generator = VideoGenerator.from_config(build_generator_config(model_dir, args.num_gpus)) + try: + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + negative_prompt="", + inputs=InputConfig(references=references), + sampling=SamplingConfig( + height=args.height, + width=args.width, + num_frames=args.num_frames, + fps=24, + # PDD counts fused blocks: one transformer forward each. + num_inference_steps=contract["num_inference_steps"], + guidance_scale=1.0, + batch_cfg=False, + seed=args.seed, + ), + output=OutputConfig( + output_path=str(output_dir / "fasth3_omniref_pdd.mp4"), + save_video=True, + return_frames=False, + ), + )) + print(f"Output written to: {result.video_path}") + if result.generation_time is not None: + print(f"Generation time: {result.generation_time:.1f} s") + if result.peak_memory_mb is not None: + print(f"Peak GPU memory allocated (rank 0): {result.peak_memory_mb / 1024:.1f} GiB") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index a6f195ece..65401519b 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -7,8 +7,18 @@ backend differs from the Wan-tuned ``video_sparse_attn``: - Tiles are ``[segment-pure prefix chunks] + [3D video tiles]``; prefix tiles never straddle segment boundaries. The tile size is selectable at - metadata build time: 256 tokens ``(4,8,8)`` (default) or 64 tokens - ``(4,4,4)`` (see ``VSA_H3_TILE_SHAPES``). + metadata build time: 256 tokens ``(4,8,8)`` (default), 128 tokens + ``(4,4,8)``, or 64 tokens ``(4,4,4)`` (see ``VSA_H3_TILE_SHAPES``). +- The builder takes the packed sequence as ordered segments: dense row + counts (text, audio, image references) and one 3-D sparse region per + video, wherever it sits. The generated video is the last region; a Ref2VA + PDD run adds one earlier region per reference video. Rows are permuted to + ``[dense chunks][region 0 tiles][region 1 tiles]...`` and + ``untile_combined_index`` inverts the permutation, so the kernels see one + dense prefix followed by video tiles. Every video query keeps its own + top-k of EACH region (``video_tile_spans`` / ``span_sparsities``); the + reference-video regions may use a different keep rate than the generated + video. - Selection is pure Python on pooled tile scores; the block-sparse kernel consumes an explicit bool mask, so no kernel changes are needed. - The compression branch is gated by ``to_gate_compress``, which the base @@ -32,15 +42,21 @@ backward run the Triton block-sparse kernels directly (no expansion, ``FASTVIDEO_VSA_CUTEDSL`` does not apply). A third, opt-in route exists for the tile-64 FORWARD only: ``FASTVIDEO_VSA_SM100A=1`` sends no-grad forwards through the data-center Blackwell CUDA block-sparse kernel -(``fastvideo_kernel.block_sparse_attn_sm100a``, upstream PR #1719 plus -our per-q-tile ``q2k_num`` fix) when the extension is built, the device +(``fastvideo_kernel.block_sparse_attn_sm100a``, which reads a separate +``q2k_num`` key-tile count per query tile) when the extension is built, the device is sm_100 or sm_103, and the geometry qualifies. The CUDA kernel assigns adjacent pairs of query tiles to CTAs, so an odd logical tile count receives one internal, zero-valid partner tile for the no-grad call only. Score search, the trained mask, gate-compress, and the returned packed sequence remain on the original logical tiles. Grad-tracking forwards and every backward stay -on Triton unchanged. If the env is set but a precondition fails, the route +on the Triton kernels. If the env is set but a precondition fails, the route logs one warning and falls back. + +Tile 128 has exactly one implementation: the same sm_100a/sm_103a CUDA +forward, which carries a 128-token block instantiation. It needs no opt-in, +pads an odd tile count with the same zero-valid partner tile, runs no-grad +forwards only, and fails closed (no Triton fallback) when the extension or +device cannot run it. """ import functools @@ -61,7 +77,7 @@ except ImportError: try: # Optional: only present in fastvideo_kernel builds that carry the - # sm_100a/sm_103a CUDA block-sparse forward (upstream PR #1719). The module itself imports + # sm_100a/sm_103a CUDA block-sparse forward. The module itself imports # fine without the compiled symbols (`_HAS_VSA_SM100A` is then False and # `is_supported` says no), so this only guards *module* availability. from fastvideo_kernel import block_sparse_attn_sm100a as _sm100a @@ -91,7 +107,12 @@ _TILE_ELEMS = math.prod(VSA_H3_TILE_SIZE) VSA_H3_TILE_SHAPES: dict[int, tuple[int, int, int]] = { _TILE_ELEMS: VSA_H3_TILE_SIZE, 64: (4, 4, 4), + # sm_100a/sm_103a CUDA forward only (128-token block instantiation). + 128: (4, 4, 8), } +# Tile sizes the sm_100a/sm_103a CUDA kernel runs natively; its CTAs own +# adjacent pairs of query tiles, so an odd count gets one zero-valid partner. +_SM100A_TILE_ELEMS = (64, 128) @torch.library.custom_op( @@ -168,15 +189,17 @@ def token_tile_and_valid(variable_block_sizes: torch.Tensor, return token_tile, token_valid -def _validate_h3_tile_geometry( - prefix_segments: tuple[int, ...], - dit_seq_shape: tuple[int, int, int], +def _validate_h3_segment_geometry( + segments: tuple[int | tuple[int, int, int], ...], variable_block_sizes: torch.Tensor, untile_combined_index: torch.Tensor, tile_elems: int = _TILE_ELEMS, ) -> None: """Fail synchronously on out-of-bounds tile geometry. + ``segments`` is the packed-order interleave: an ``int`` is a dense + (prefix) segment's row count, a triple is one video region's token grid. + Invariants the block-sparse kernel trusts without checking: every tile's valid size is in (0, tile_elems]; the sizes sum to the packed sequence length; and ``untile_combined_index`` maps each packed @@ -185,82 +208,99 @@ def _validate_h3_tile_geometry( collective (e.g. an FSDP all-gather), which is unattributable — so raise here, once per cached geometry, with the numbers in hand. """ - total = sum(prefix_segments) + math.prod(dit_seq_shape) + total = sum(math.prod(segment) if isinstance(segment, tuple) else segment for segment in segments) n_pad = variable_block_sizes.numel() * tile_elems sizes_min = int(variable_block_sizes.min()) sizes_max = int(variable_block_sizes.max()) sizes_sum = int(variable_block_sizes.sum()) if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total: - raise ValueError(f"VSA-H3 tile sizes out of bounds for prefix={prefix_segments}, video={dit_seq_shape}, " + raise ValueError(f"VSA-H3 tile sizes out of bounds for segments={segments}, " f"tile_elems={tile_elems}: min={sizes_min}, max={sizes_max}, sum={sizes_sum}, " f"expected sum={total}.") if untile_combined_index.numel() != total: raise ValueError(f"VSA-H3 untile index has {untile_combined_index.numel()} entries for a packed " - f"sequence of {total} rows (prefix={prefix_segments}, video={dit_seq_shape}).") + f"sequence of {total} rows (segments={segments}).") idx_min = int(untile_combined_index.min()) idx_max = int(untile_combined_index.max()) if idx_min < 0 or idx_max >= n_pad: # Range first: the pad-slot gather below would itself index out of # bounds (the very async fault this guard exists to preempt). raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: range " - f"[{idx_min}, {idx_max}] vs padded length {n_pad} " - f"(prefix={prefix_segments}, video={dit_seq_shape}).") + f"[{idx_min}, {idx_max}] vs padded length {n_pad} (segments={segments}).") in_tile_offset = untile_combined_index % tile_elems maps_into_pad = bool((in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any()) if maps_into_pad or int(torch.unique(untile_combined_index).numel()) != total: raise ValueError(f"VSA-H3 untile index is not an injective map into non-pad slots: " - f"pad-slot hit={maps_into_pad} " - f"(prefix={prefix_segments}, video={dit_seq_shape}).") + f"pad-slot hit={maps_into_pad} (segments={segments}).") @functools.lru_cache(maxsize=10) -def _h3_tile_geometry( - prefix_segments: tuple[int, ...], - dit_seq_shape: tuple[int, int, int], +def _h3_segment_tile_geometry( + segments: tuple[int | tuple[int, int, int], ...], device: torch.device, tile_shape: tuple[int, int, int] = VSA_H3_TILE_SIZE, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: - """Tile the packed sequence: segment-pure prefix chunks, then video tiles. +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, tuple[tuple[int, int], ...]]: + """Tile a packed sequence of interleaved dense segments and video regions. + + ``segments`` lists, in packed order, dense segment row counts (``int``) + and video-region token grids (``(t, h, w)`` triples). Rows are permuted so + every dense chunk lands in the leading prefix tiles and each video region + contributes one contiguous run of 3-D tiles, wherever the dense segments + sit in the packed sequence. Dense chunks never straddle a segment + boundary. Returns (tile_partition_indices, variable_block_sizes, - untile_combined_index, num_prefix_tiles, num_video_tiles). + untile_combined_index, num_prefix_tiles, num_video_tiles, + video_tile_spans), where ``video_tile_spans`` holds one ``[start, end)`` + tile range per video region (prefix tiles included in the numbering). """ tile_elems = math.prod(tile_shape) - prefix_len = sum(prefix_segments) prefix_sizes: list[int] = [] - for segment in prefix_segments: + prefix_rows: list[torch.Tensor] = [] + video_entries: list[tuple[int, tuple[int, int, int]]] = [] + cursor = 0 + for segment in segments: + if isinstance(segment, tuple): + video_entries.append((cursor, segment)) + cursor += math.prod(segment) + continue full, rem = divmod(segment, tile_elems) prefix_sizes.extend([tile_elems] * full) if rem: prefix_sizes.append(rem) + prefix_rows.append(torch.arange(cursor, cursor + segment, device=device, dtype=torch.long)) + cursor += segment num_prefix_tiles = len(prefix_sizes) ts_t, ts_h, ts_w = tile_shape - t, h, w = dit_seq_shape - num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w)) - video_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device, tile_shape) - num_video_tiles = int(video_sizes.numel()) + video_index_parts: list[torch.Tensor] = [] + video_size_parts: list[torch.Tensor] = [] + video_tile_spans: list[tuple[int, int]] = [] + tile_cursor = num_prefix_tiles + for start, grid in video_entries: + t, h, w = grid + num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w)) + region_sizes = construct_variable_block_sizes(grid, num_tiles, device, tile_shape) + video_index_parts.append(get_tile_partition_indices(grid, tile_shape, device) + start) + video_size_parts.append(region_sizes) + video_tile_spans.append((tile_cursor, tile_cursor + int(region_sizes.numel()))) + tile_cursor += int(region_sizes.numel()) + num_video_tiles = tile_cursor - num_prefix_tiles - video_indices = get_tile_partition_indices(dit_seq_shape, tile_shape, device) + prefix_len - tile_partition_indices = torch.cat([ - torch.arange(prefix_len, device=device, dtype=torch.long), - video_indices, - ]) + tile_partition_indices = torch.cat(prefix_rows + video_index_parts) # cat promotes the int32 helper output to int64 alongside the prefix sizes - variable_block_sizes = torch.cat([ - torch.tensor(prefix_sizes, dtype=torch.long, device=device), - video_sizes, - ]) + variable_block_sizes = torch.cat([torch.tensor(prefix_sizes, dtype=torch.long, device=device)] + video_size_parts) # get_non_pad_index is lru-cached on tensor identity; variable_block_sizes # is itself cached by this function, so the identity stays stable. non_pad_index = get_non_pad_index(variable_block_sizes, tile_elems) untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)] - # One-time (lru-cached) synchronous bounds check; see _validate_h3_tile_geometry. - _validate_h3_tile_geometry(prefix_segments, dit_seq_shape, variable_block_sizes, untile_combined_index, tile_elems) - return (tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles) + # One-time (lru-cached) synchronous bounds check; see _validate_h3_segment_geometry. + _validate_h3_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems) + return (tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles, + tuple(video_tile_spans)) class MiniMaxH3VSABackend(AttentionBackend): @@ -308,8 +348,13 @@ class MiniMaxH3VSAMetadata(AttentionMetadata): # this tensor with each implementation's tensor-valued layer index so the # shared block code does not specialize once per Python ``layer_idx``. dense_layers_tensor: torch.Tensor - # tokens per tile (256 or 64); selects the tile geometry AND the kernel - # route in forward() (256 -> VSA-256 CuTe/Triton, 64 -> native Triton) + # One [start, end) tile range per sparse video region, in packed order + # with the generated video last, and the sparsity of that region's top-k. + video_tile_spans: tuple[tuple[int, int], ...] + span_sparsities: tuple[float, ...] + # tokens per tile (256, 128 or 64); selects the tile geometry AND the + # kernel route in forward() (256 -> VSA-256 CuTe/Triton, 128 -> sm_100a + # CUDA, 64 -> native Triton or opt-in sm_100a CUDA) tile_elems: int = _TILE_ELEMS # layers forced dense regardless of sparsity (probe-guided opt-outs) dense_layers: tuple[int, ...] = () @@ -330,26 +375,61 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): def build( # type: ignore self, current_timestep: int, - raw_latent_shape: tuple[int, int, int], patch_size: tuple[int, int, int], VSA_sparsity: float, - prefix_segments: tuple[int, ...], + packed_segments: tuple[int | tuple[int, int, int], ...], device: torch.device, exempt: bool = True, dense_layers: tuple[int, ...] = (), tile_size: int = _TILE_ELEMS, + ref_keep_rate: float | None = None, **kwargs: dict[str, Any], ) -> MiniMaxH3VSAMetadata: + """Build per-step metadata for one packed H3 sequence. + + ``packed_segments`` lists the sequence in packed order: an ``int`` is + a dense segment's row count, and a ``(t, h, w)`` triple is the raw + latent shape of one sparse video region. The last region is the + generated video and follows ``VSA_sparsity``; every earlier region is + a reference video and keeps ``ref_keep_rate`` of its tiles, which a + build with reference-video regions requires. ``exempt=False`` + (compete mode) supports only the generated-video region. + """ tile_shape = VSA_H3_TILE_SHAPES.get(int(tile_size)) if tile_shape is None: raise ValueError(f"VSA-H3 tile_size must be one of {sorted(VSA_H3_TILE_SHAPES)}, got {tile_size!r}") - dit_seq_shape = (raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1], - raw_latent_shape[2] // patch_size[2]) - prefix_segments = tuple(int(s) for s in prefix_segments if s > 0) - total_seq_length = sum(prefix_segments) + math.prod(dit_seq_shape) + # Video regions become token grids under the patch size; empty dense segments are dropped. + token_segments: list[int | tuple[int, int, int]] = [] + for segment in packed_segments: + if isinstance(segment, tuple): + if any(int(v) <= 0 or int(v) % p for v, p in zip(segment, patch_size, strict=True)): + raise ValueError(f"VSA-H3 video region latent shape {tuple(segment)} is not a positive multiple " + f"of patch {tuple(patch_size)}.") + t, h, w = (int(v) // p for v, p in zip(segment, patch_size, strict=True)) + token_segments.append((t, h, w)) + elif segment > 0: + token_segments.append(int(segment)) + num_regions = sum(isinstance(segment, tuple) for segment in token_segments) + if num_regions == 0: + raise ValueError("VSA-H3 needs at least one video region.") + reference_sparsities: tuple[float, ...] = () + if num_regions > 1: + if ref_keep_rate is None: + raise ValueError(f"A VSA-H3 build with {num_regions - 1} reference-video region(s) needs " + "ref_keep_rate.") + if not exempt: + raise ValueError(f"vsa_mode='compete' supports only the generated-video region; this build has " + f"{num_regions - 1} reference-video region(s). Use vsa_mode='exempt'.") + reference_sparsities = (1.0 - float(ref_keep_rate), ) * (num_regions - 1) + total_seq_length = sum(math.prod(s) if isinstance(s, tuple) else s for s in token_segments) - (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, - num_video_tiles) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device, tile_shape) + (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles, + video_tile_spans) = _h3_segment_tile_geometry(tuple(token_segments), device, tile_shape) + + # A dense build (sparsity 0: dense steps) keeps every region dense, + # whatever the reference keep rate. + VSA_sparsity = float(VSA_sparsity) + span_sparsities = (0.0, ) * num_regions if VSA_sparsity <= 0.0 else (*reference_sparsities, VSA_sparsity) dense_layers = tuple(int(layer) for layer in dense_layers) return MiniMaxH3VSAMetadata( @@ -361,9 +441,11 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): exempt=exempt, variable_block_sizes=variable_block_sizes, untile_combined_index=untile_combined_index, + dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64), + video_tile_spans=video_tile_spans, + span_sparsities=span_sparsities, tile_elems=int(tile_size), dense_layers=dense_layers, - dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64), tile_buf_holder=self._tile_buf_holder, ) @@ -386,23 +468,46 @@ def _pool_tiles(x: torch.Tensor, variable_block_sizes: torch.Tensor, tile_elems: def _build_block_mask( scores: torch.Tensor, num_prefix_tiles: int, - num_video_tiles: int, VSA_sparsity: float, exempt: bool, + video_tile_spans: tuple[tuple[int, int], ...], + span_sparsities: tuple[float, ...], ) -> torch.Tensor: - """scores: [B, H, n_tiles, n_tiles] -> bool mask, same shape.""" + """scores: [B, H, n_tiles, n_tiles] -> bool mask, same shape. + + Prefix rows (non-video queries) stay dense. With ``exempt``, every query + keeps every prefix column, and every video query keeps its own top-k of + EACH video region, with the budget from that region's ``span_sparsities`` + entry. Without ``exempt`` (compete mode, generated-video region only), + prefix and video columns compete for one FLOP-matched top-k. + ``VSA_sparsity <= 0`` (dense steps and layers) keeps every tile. + ``video_tile_spans`` come from the cached geometry, which lays them out + contiguously after the prefix tiles, with one ``span_sparsities`` entry + each. + """ n_tiles = scores.shape[-1] - k_vid = compute_topk(VSA_sparsity, num_video_tiles) - if k_vid == num_video_tiles: + if VSA_sparsity <= 0.0: + span_sparsities = tuple(VSA_sparsity for _ in video_tile_spans) + span_topk = [ + compute_topk(sparsity, end - start) + for (start, end), sparsity in zip(video_tile_spans, span_sparsities, strict=True) + ] + if all(k == end - start for k, (start, end) in zip(span_topk, video_tile_spans, strict=True)): return torch.ones_like(scores, dtype=torch.bool) mask = torch.zeros_like(scores, dtype=torch.bool) if exempt or num_prefix_tiles == 0: - video_cols = scores[..., num_prefix_tiles:] - idx = video_cols.topk(k_vid, dim=-1).indices + num_prefix_tiles - mask.scatter_(-1, idx, True) + for (start, end), k_span in zip(video_tile_spans, span_topk, strict=True): + if k_span == end - start: + mask[..., start:end] = True + continue + idx = scores[..., start:end].topk(k_span, dim=-1).indices + start + mask.scatter_(-1, idx, True) mask[..., :num_prefix_tiles] = True else: - k_total = min(k_vid + num_prefix_tiles, n_tiles) + # compete: prefix keys enter the top-k under a FLOP-matched budget, + # which is defined for the generated-video region alone (the builder + # rejects compete mode with reference-video regions) + k_total = min(span_topk[0] + num_prefix_tiles, n_tiles) idx = scores.topk(k_total, dim=-1).indices mask.scatter_(-1, idx, True) mask[:, :, :num_prefix_tiles, :] = True @@ -452,6 +557,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): # request-time env/probe/fallback behavior; only Dynamo capture reads # the prepared, static route. self._regional_compile_sm100a_enabled: bool | None = None + # Tile-128 route per (device, dtype, head size): None when the CUDA + # kernel can run it, else why not. The kernel's predicate otherwise + # depends only on the tile-128 buffer contract (contiguous BHSD, + # 128-token blocks, an even tile count with the partner tile, integer + # tile sizes), which every call meets, so it is evaluated once. + self._tile128_route: dict[tuple[torch.device, torch.dtype, int], str | None] = {} def prepare_for_compile(self, device: torch.device) -> None: """Tensorize per-layer state shared by every torch.compile route.""" @@ -511,9 +622,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): buffer; callers must consume it before the next ``tile()`` (both call sites in ``forward()`` read it immediately). A grad-tracking forward instead receives a fresh buffer and leaves the holder untouched, so - the builder never retains autograd state across steps. Odd tile-64 - no-grad sm100a requests carry one additional all-zero tile internally; - metadata and all observable outputs retain the logical geometry. + the builder never retains autograd state across steps. Odd no-grad + sm100a requests (tile 128 always, tile 64 when opted in) carry one + additional all-zero tile internally; metadata and all observable + outputs retain the logical geometry. """ if x.shape[1] != attn_metadata.total_seq_length: raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, " @@ -530,7 +642,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): sm100a_requested = False else: sm100a_requested = envs.FASTVIDEO_VSA_SM100A.get() - needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode and sm100a_requested) + # Tile 128 has no route other than the sm100a CUDA kernel. + sm100a_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64 and sm100a_requested) + needs_sm100a_pair = n_tiles % 2 != 0 and not grad_mode and sm100a_route kernel_tiles = n_tiles + int(needs_sm100a_pair) target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1]) @@ -584,11 +698,11 @@ class MiniMaxH3VSAImpl(AttentionImpl): tile_elems = attn_metadata.tile_elems if regional_compiling and tile_elems != 64: raise RuntimeError("VSA-H3 regional fullgraph compile requires 64-token tiles; disable " - "inference_torch_compile for tile-256/CuTe runs.") + "inference_torch_compile for tile-128/256 runs.") if tile_elems == 64: if block_sparse_attn_64_bhsd is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn is not installed") - elif block_sparse_attn_256_bshd is None: + elif tile_elems == 256 and block_sparse_attn_256_bshd is None: raise NotImplementedError("fastvideo_kernel.block_sparse_attn_256 is not installed") # Probe recording performs filesystem writes and host synchronizations, @@ -596,12 +710,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): # environment switch while Dynamo captures a regional full graph. # The metadata always describes the trained logical geometry. # ``tile()`` may append exactly one transport-only partner for an odd - # tile-64 sm100a call. Keep score selection and the gate branch on the - # logical prefix, and reject every other shape before a kernel sees it. + # tile-64/128 sm100a call. Keep score selection and the gate branch on + # the logical prefix, and reject every other shape before a kernel sees it. n_tiles = attn_metadata.variable_block_sizes.numel() logical_seq_len = n_tiles * tile_elems pair_pad_seq_len = logical_seq_len + tile_elems - pair_pad_is_valid = tile_elems == 64 and n_tiles % 2 != 0 + pair_pad_is_valid = tile_elems in _SM100A_TILE_ELEMS and n_tiles % 2 != 0 allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, ) if query.shape[1] not in allowed_seq_lengths: expected = (f"the logical length {logical_seq_len} or one sm100a partner tile " @@ -649,18 +763,19 @@ class MiniMaxH3VSAImpl(AttentionImpl): mask = _build_block_mask( scores, attn_metadata.num_prefix_tiles, - attn_metadata.num_video_tiles, layer_sparsity, attn_metadata.exempt, + attn_metadata.video_tile_spans, + attn_metadata.span_sparsities, ) if force_dense is not None: # A scalar bool tensor broadcasts over the block map. This exactly # preserves the eager dense-layer contract without a Python branch. mask = mask | force_dense - if tile_elems == 64: - # Native 64-token path: the block map is already at the kernels' - # granularity. Both 64-token entries take BHSD ([B, H, S_pad, D]); + if tile_elems in _SM100A_TILE_ELEMS: + # Native 64/128-token path: the block map is already at the + # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose # around the call. q_bhsd = query.transpose(1, 2).contiguous() @@ -680,14 +795,30 @@ class MiniMaxH3VSAImpl(AttentionImpl): value=0, ) - # Opt-in sm_100a CUDA forward (upstream PR #1719 + per-q-tile - # q2k_num fix). Forward-only: grad-tracking calls stay on Triton + # sm_100a CUDA forward, opt-in at tile 64 and the only route at + # tile 128. Forward-only: tile-64 grad-tracking calls stay on Triton # so autograd keeps the Triton fwd+bwd pairing untouched. The # kernel does return an LSE in Triton's M format, so a future # fwd/bwd pairing is possible, but it is not built here. grad_mode = torch.is_grad_enabled() and (query.requires_grad or key.requires_grad or value.requires_grad) use_sm100a = False - if regional_compiling: + if tile_elems == 128: + # The CUDA kernel's 128-token instantiation is the only tile-128 + # implementation, so it needs no opt-in and never falls back. + if grad_mode: + raise NotImplementedError("VSA-H3 tile 128 supports no-grad inference forwards only.") + route_key = (q_bhsd.device, q_bhsd.dtype, int(q_bhsd.shape[-1])) + if route_key not in self._tile128_route: + reason = _sm100a_unavailable_reason(_sm100a, q_bhsd, sm100a_variable_block_sizes, grad_mode) + if reason is None and map_to_index is None: + reason = "fastvideo_kernel.triton_kernels.index (map_to_index) is not importable" + self._tile128_route[route_key] = reason + reason = self._tile128_route[route_key] + if reason is not None: + raise RuntimeError(f"VSA-H3 tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel: " + f"{reason}. Use VSA_tile_size=64 or 256 elsewhere.") + use_sm100a = True + elif regional_compiling: if self._regional_compile_sm100a_enabled is None: raise RuntimeError( "VSA-H3 sm_100a routing was not resolved before torch.compile; " @@ -712,11 +843,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): logger.warning_once(f"{VSA_SM100A_ENV}=1 but falling back to the Triton-64 kernels: {reason}") if use_sm100a: - # Regional preparation emits the compile-route receipt before + # prepare_for_regional_compile logs the selected mask route before # capture. Logging from this branch would itself break a # ``fullgraph=True`` forward. if not compiling: - logger.info_once("MiniMax-H3 VSA tile-64 forward: using the sm100a/sm103a CUDA block-sparse kernel") + logger.info_once( + f"MiniMax-H3 VSA tile-{tile_elems} forward: using the sm100a/sm103a CUDA block-sparse kernel") if regional_compiling: # The compile-safe wrapper keeps both Triton mask # compaction and the raw pybind launch behind one @@ -732,7 +864,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): # Preserve the established eager/index-native route and # compatibility with older kernel wheels. Per-row counts # are non-uniform (prefix queries are dense; video queries - # run prefix+top-k), which the fixed kernel supports. + # run prefix+top-k); the kernel reads each query tile's own count. q2k_idx, q2k_num = map_to_index(sm100a_mask) out_bhsd, _ = _sm100a.block_sparse_attn_sm100a( q_bhsd, diff --git a/fastvideo/configs/models/dits/minimax_h3.py b/fastvideo/configs/models/dits/minimax_h3.py index 464f75a69..b67f2993f 100644 --- a/fastvideo/configs/models/dits/minimax_h3.py +++ b/fastvideo/configs/models/dits/minimax_h3.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import dataclass, field from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.layers.pdd import validate_pdd_steps from fastvideo.platforms import AttentionBackendEnum @@ -59,6 +60,11 @@ class MiniMaxH3ArchConfig(DiTArchConfig): norm_eps: float = 1e-5 qk_norm_eps: float = 1e-5 final_norm_eps: float = 1e-5 + # Parallel Decoding Distillation (PDD) students widen both output + # projections to this many fine-grid heads (``fastvideo.layers.pdd``). + # ``None`` is the released single-head transformer. A PDD export records + # the key in its transformer ``config.json``. + pdd_steps: int | None = None def __post_init__(self) -> None: super().__post_init__() @@ -67,6 +73,8 @@ class MiniMaxH3ArchConfig(DiTArchConfig): self.patch_size = (self.patch_size[0], self.patch_size[1], self.patch_size[2]) self.num_channels_latents = self.in_channels self.out_channels = self.in_channels + if self.pdd_steps is not None: + validate_pdd_steps(self.pdd_steps, "MiniMax H3 pdd_steps") if self.adaln_rank is not None and not 0 < self.adaln_rank <= self.time_embed_dim: raise ValueError(f"MiniMax H3 adaln_rank must be in (0, time_embed_dim={self.time_embed_dim}], " f"got {self.adaln_rank}.") diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 271d06763..17c29bae3 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -2,7 +2,7 @@ import json from collections.abc import Callable from dataclasses import asdict, dataclass, field, fields -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import torch @@ -12,6 +12,11 @@ from fastvideo.configs.utils import update_config_from_args from fastvideo.logger import init_logger from fastvideo.utils import FlexibleArgumentParser, StoreBoolean, shallow_asdict +if TYPE_CHECKING: + from fastvideo.api.sampling_param import SamplingParam + from fastvideo.api.schema import GenerationRequest + from fastvideo.fastvideo_args import FastVideoArgs + logger = init_logger(__name__) @@ -281,6 +286,25 @@ class PipelineConfig: f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})" ) + def resolve_checkpoint_settings(self, fastvideo_args: "FastVideoArgs") -> None: + """Apply the run settings that the checkpoint at ``fastvideo_args.model_path`` fixes. + + ``FastVideoArgs.__post_init__`` calls this once in the main process, before + argument validation and before any weights load. Most checkpoints fix no + run settings, so the base implementation does nothing. + """ + + def apply_request_constraints(self, request: "GenerationRequest", + sampling_param: "SamplingParam") -> "SamplingParam": + """Apply this model's constraints on a generation request to its sampling parameters. + + ``sampling_param`` holds the preset defaults with the request's explicit + fields applied. A model sets the values it fixes when the request leaves + them unset, and raises when the request sets a value it does not accept. + The base PipelineConfig accepts every request unchanged. + """ + return sampling_param + def dump_to_json(self, file_path: str): output_dict = shallow_asdict(self) del_keys = [] diff --git a/fastvideo/configs/pipelines/minimax_h3.py b/fastvideo/configs/pipelines/minimax_h3.py index d2a0d5f14..366028600 100644 --- a/fastvideo/configs/pipelines/minimax_h3.py +++ b/fastvideo/configs/pipelines/minimax_h3.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any from fastvideo.configs.models import EncoderConfig, VAEConfig from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config @@ -11,6 +12,137 @@ from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3 from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEConfig from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.logger import init_logger +from fastvideo.utils import read_optional_model_json + +if TYPE_CHECKING: + from fastvideo.api.sampling_param import SamplingParam + from fastvideo.api.schema import GenerationRequest + from fastvideo.fastvideo_args import FastVideoArgs + +logger = init_logger(__name__) + +# The checkpoint file that records how a distilled FastH3 export must be sampled. +FASTH3_INFERENCE_FILE = "fastvideo_inference.json" +FASTH3_INFERENCE_SCHEMA = "fasth3-inference-contract-v1" +# The fields a FastH3 Ref2VA PDD export writes to FASTH3_INFERENCE_FILE. Missing +# and unknown fields are both errors, so no field is silently ignored. +_PDD_INFERENCE_FIELDS = frozenset({ + "schema_version", + "schema", + "model_type", + "transformer_component", + "conditioning", + "base_model_revision", + "pdd_steps", + "pdd_step_indices", + "num_inference_steps", + "transformer_forwards", + "grid_max_t", + "video_scheduler_shift", + "audio_scheduler_shift", + "guidance_scale", + "attention_backend", + "vsa_sparsity", + "vsa_tile_size", + "vsa_ref_policy", + "vsa_ref_keep_rate", +}) +# Fields that repeat a value FastVideo fixes in code: the Ref2VA pipeline, its +# conditioning, sampling without classifier-free guidance, and the per-reference +# sparse regions that only VIDEO_SPARSE_ATTN_H3 builds. +_PDD_FIXED_FIELDS: dict[str, Any] = { + "schema_version": FASTH3_INFERENCE_SCHEMA, + "schema": FASTH3_INFERENCE_SCHEMA, + "model_type": "ref2va", + "transformer_component": "transformer_ref", + "conditioning": "fixed_ordered_references_target_only_flow", + "guidance_scale": 1.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_ref_policy": "p2_multi_region", +} + +# The base snapshot an export was distilled against: "hf://@". +_BASE_MODEL_REVISION_PREFIX = "hf://" + + +def parse_base_model_revision(value: Any) -> tuple[str, str]: + """Split a contract's ``base_model_revision``, ``hf://@``, into repo id and revision.""" + if isinstance(value, str) and value.startswith(_BASE_MODEL_REVISION_PREFIX): + repo, separator, revision = value[len(_BASE_MODEL_REVISION_PREFIX):].partition("@") + if separator and repo.strip() and revision.strip() and "@" not in revision: + return repo, revision + raise ValueError(f"FastH3 base_model_revision={value!r} must be hf://@.") + + +def _require_fraction(name: str, value: Any, *, allow_zero: bool) -> None: + """Raise unless *value* is a real number in [0, 1) (``allow_zero``) or in (0, 1).""" + is_real = isinstance(value, int | float) and not isinstance(value, bool) + if not (is_real and (value >= 0.0 if allow_zero else value > 0.0) and value < 1.0): + raise ValueError(f"{name} must be in {'[0, 1)' if allow_zero else '(0, 1)'}, got {value!r}.") + + +def _read_pdd_inference_file(model_path: str, revision: str | None) -> dict[str, Any] | None: + """Read and validate a FastH3 Ref2VA PDD checkpoint's fastvideo_inference.json; None for other checkpoints. + + The file must carry exactly the PDD fields. Fields that FastVideo fixes in + code must hold that value, and fields that repeat a value stored in another + checkpoint file must equal that file: ``pdd_steps`` the transformer config, + the shifts the scheduler configs, and the step counts the partition, which + must increase strictly from 0 to ``pdd_steps``. + """ + inference_file = read_optional_model_json(model_path, FASTH3_INFERENCE_FILE, revision) + if inference_file is None or "pdd_steps" not in inference_file: + return None + source = f"{model_path}/{FASTH3_INFERENCE_FILE}" + missing = sorted(_PDD_INFERENCE_FIELDS - inference_file.keys()) + unknown = sorted(inference_file.keys() - _PDD_INFERENCE_FIELDS) + if missing or unknown: + raise ValueError(f"{source} must carry exactly the FastH3 PDD fields; missing {missing}, unknown {unknown}.") + for key, expected in _PDD_FIXED_FIELDS.items(): + if inference_file[key] != expected: + raise ValueError(f"{source} {key}={inference_file[key]!r} is unsupported; expected {expected!r}.") + parse_base_model_revision(inference_file["base_model_revision"]) + from fastvideo.layers.pdd import PDD_GRID_MAX_T + if inference_file["grid_max_t"] != PDD_GRID_MAX_T: + raise ValueError(f"{source} grid_max_t={inference_file['grid_max_t']!r} is unsupported; the PDD fine grid " + f"ends at {PDD_GRID_MAX_T}.") + + owner_files = {} + for owner in ("transformer_ref/config.json", "scheduler/scheduler_config.json", + "audio_scheduler/scheduler_config.json"): + owner_files[owner] = read_optional_model_json(model_path, owner, revision) + if owner_files[owner] is None: + raise ValueError(f"{model_path} has {FASTH3_INFERENCE_FILE} but no {owner}.") + pdd_steps = owner_files["transformer_ref/config.json"].get("pdd_steps") + if inference_file["pdd_steps"] != pdd_steps: + raise ValueError(f"{source} pdd_steps={inference_file['pdd_steps']!r} disagrees with " + f"transformer_ref/config.json pdd_steps={pdd_steps!r}.") + for key, owner in (("video_scheduler_shift", "scheduler/scheduler_config.json"), + ("audio_scheduler_shift", "audio_scheduler/scheduler_config.json")): + shift = owner_files[owner].get("shift") + if inference_file[key] != shift: + raise ValueError(f"{source} {key}={inference_file[key]!r} disagrees with {owner} shift={shift!r}.") + indices = inference_file["pdd_step_indices"] + if (not isinstance(indices, list) or len(indices) < 2 or any(type(index) is not int for index in indices) + or indices[0] != 0 or indices[-1] != pdd_steps + or any(left >= right for left, right in zip(indices, indices[1:], strict=False))): + raise ValueError(f"{source} pdd_step_indices={indices!r} must increase strictly from 0 to " + f"pdd_steps={pdd_steps}.") + for key in ("num_inference_steps", "transformer_forwards"): + if inference_file[key] != len(indices) - 1: + raise ValueError(f"{source} {key}={inference_file[key]!r} disagrees with the {len(indices) - 1} blocks " + "of pdd_step_indices.") + + # Trained VSA settings. A run's override of a tunable one is checked where it is applied. + from fastvideo.attention.backends.video_sparse_attn_h3 import VSA_H3_TILE_SHAPES + tile_size = inference_file["vsa_tile_size"] + if type(tile_size) is not int or tile_size not in VSA_H3_TILE_SHAPES: + raise ValueError(f"{source} vsa_tile_size={tile_size!r} must be one of {sorted(VSA_H3_TILE_SHAPES)}.") + _require_fraction(f"{source} vsa_sparsity", inference_file["vsa_sparsity"], allow_zero=True) + # A keep rate of 1 would leave reference-video attention dense. + _require_fraction(f"{source} vsa_ref_keep_rate", inference_file["vsa_ref_keep_rate"], allow_zero=False) + return inference_file @dataclass @@ -28,6 +160,15 @@ class MiniMaxH3PipelineConfig(PipelineConfig): vae_precision: str = "fp32" text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) vae_sp: bool = False + # Parallel Decoding Distillation (PDD) students: fine-grid node indices of + # the fused blocks, one transformer forward each. resolve_checkpoint_settings + # sets it from the checkpoint's fastvideo_inference.json; None means the + # checkpoint is not a PDD student. + pdd_step_indices: tuple[int, ...] | None = None + # Fraction of each reference video's VSA tiles that every video query keeps, + # for PDD students. None takes the checkpoint's trained value; another value + # overrides it with a warning. + vsa_ref_keep_rate: float | None = None def check_pipeline_config(self) -> None: super().check_pipeline_config() @@ -35,5 +176,91 @@ class MiniMaxH3PipelineConfig(PipelineConfig): raise ValueError("MiniMax-H3 uses separate checkpoint-defined video/audio scheduler shifts; " "flow_shift must remain unset.") + def resolve_checkpoint_settings(self, fastvideo_args: FastVideoArgs) -> None: + """Apply a FastH3 Ref2VA PDD checkpoint's fastvideo_inference.json to this run. -__all__ = ["MiniMaxH3PipelineConfig"] + The file is the one source of the PDD partition and the trained VSA + settings, and a setting the run leaves unset takes the file's value. + The partition, attention backend and tile size are fixed by training, + so a different run value raises; VSA sparsity and the reference keep + rate may differ, with a warning. For checkpoints without PDD fields + (base MiniMax-H3, DMD exports) this method sets nothing. + """ + inference_file = _read_pdd_inference_file(fastvideo_args.model_path, fastvideo_args.revision) + if inference_file is None: + for name in ("pdd_step_indices", "vsa_ref_keep_rate"): + if getattr(self, name) is not None: + raise ValueError(f"{name} applies only to FastH3 PDD checkpoints, whose {FASTH3_INFERENCE_FILE} " + f"sets it; {fastvideo_args.model_path} has no PDD fields.") + return + + # Training-fixed settings: an unset run value takes the file's value; a different one raises. + from fastvideo.attention.selector import coerce_attn_backend + trained_backend = inference_file["attention_backend"] + if fastvideo_args.attention_backend is None: + fastvideo_args.attention_backend = trained_backend + elif coerce_attn_backend(fastvideo_args.attention_backend) != coerce_attn_backend(trained_backend): + raise ValueError(f"This FastH3 PDD checkpoint was trained with attention_backend={trained_backend}; " + f"this run requests {fastvideo_args.attention_backend} (argument or " + "FASTVIDEO_ATTENTION_BACKEND). Leave it unset or pass the trained value.") + trained_tile = inference_file["vsa_tile_size"] + if fastvideo_args.VSA_tile_size is None: + fastvideo_args.VSA_tile_size = trained_tile + elif fastvideo_args.VSA_tile_size != trained_tile: + raise ValueError(f"This FastH3 PDD checkpoint was trained with VSA_tile_size={trained_tile}; this run " + f"requests {fastvideo_args.VSA_tile_size}. Leave it unset or pass the trained value.") + indices = tuple(inference_file["pdd_step_indices"]) + if self.pdd_step_indices is not None and tuple(self.pdd_step_indices) != indices: + raise ValueError(f"pdd_step_indices comes from {FASTH3_INFERENCE_FILE} ({list(indices)}); this run sets " + f"{list(self.pdd_step_indices)}.") + if self.dmd_denoising_steps is not None: + raise ValueError(f"{fastvideo_args.model_path} is a FastH3 PDD checkpoint; dmd_denoising_steps must be " + "unset.") + self.pdd_step_indices = indices + + # Tunable settings: an unset run value takes the trained value; a different one wins, with a warning. + trained_sparsity = inference_file["vsa_sparsity"] + if fastvideo_args.VSA_sparsity is None: + fastvideo_args.VSA_sparsity = trained_sparsity + elif fastvideo_args.VSA_sparsity != trained_sparsity: + _require_fraction("VSA_sparsity", fastvideo_args.VSA_sparsity, allow_zero=True) + logger.warning("FastH3 PDD checkpoint was trained with VSA_sparsity=%s; this run uses %s.", + trained_sparsity, fastvideo_args.VSA_sparsity) + trained_keep_rate = inference_file["vsa_ref_keep_rate"] + if self.vsa_ref_keep_rate is None: + self.vsa_ref_keep_rate = trained_keep_rate + elif self.vsa_ref_keep_rate != trained_keep_rate: + _require_fraction("vsa_ref_keep_rate", self.vsa_ref_keep_rate, allow_zero=False) + logger.warning("FastH3 PDD checkpoint was trained with vsa_ref_keep_rate=%s; this run uses %s.", + trained_keep_rate, self.vsa_ref_keep_rate) + + logger.info( + "FastH3 PDD checkpoint %s: %d fused blocks %s of a %d-interval grid; attention_backend=%s, " + "VSA_tile_size=%s, VSA_sparsity=%s, vsa_ref_keep_rate=%s", fastvideo_args.model_path, + len(indices) - 1, list(indices), indices[-1], fastvideo_args.attention_backend, + fastvideo_args.VSA_tile_size, fastvideo_args.VSA_sparsity, self.vsa_ref_keep_rate) + + def apply_request_constraints(self, request: GenerationRequest, sampling_param: SamplingParam) -> SamplingParam: + """Fix ``num_inference_steps`` to the block count of a PDD checkpoint. + + A PDD checkpoint runs exactly one transformer forward per trained fused + block: a request that leaves ``num_inference_steps`` unset gets the block + count, and a request that sets another count raises. Checkpoints without + a PDD partition return ``sampling_param`` unchanged. + """ + if self.pdd_step_indices is None: + return sampling_param + from fastvideo.api.compat import explicit_request_updates + + num_blocks = len(self.pdd_step_indices) - 1 + requested_steps = explicit_request_updates(request).get("num_inference_steps", num_blocks) + if requested_steps != num_blocks: + raise ValueError(f"This FastH3 PDD checkpoint runs exactly {num_blocks} transformer forwards; the request " + f"sets num_inference_steps={requested_steps}. Pass num_inference_steps={num_blocks}. Only " + "a request parsed from a mapping or a config file can leave it unset; a " + "GenerationRequest built in Python counts every field as set.") + sampling_param.num_inference_steps = num_blocks + return sampling_param + + +__all__ = ["FASTH3_INFERENCE_FILE", "FASTH3_INFERENCE_SCHEMA", "MiniMaxH3PipelineConfig", "parse_base_model_revision"] diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 5bdc819b7..90d454aa4 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -458,9 +458,9 @@ class VideoGenerator: raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) - resolved_sampling_param = request_to_sampling_param( + resolved_sampling_param = fastvideo_args.pipeline_config.apply_request_constraints( request, - model_path=self.fastvideo_args.model_path, + request_to_sampling_param(request, model_path=self.fastvideo_args.model_path), ) return self._generate_video_impl( prompt=request.prompt, @@ -508,9 +508,9 @@ class VideoGenerator: raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) - sampling_param = request_to_sampling_param( + sampling_param = fastvideo_args.pipeline_config.apply_request_constraints( request, - model_path=self.fastvideo_args.model_path, + request_to_sampling_param(request, model_path=self.fastvideo_args.model_path), ) batch_extra = request_to_batch_extra(request) result = self._generate_video_impl( diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 8cb4ef920..d57a179a7 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -231,9 +231,10 @@ class FastVideoArgs: disable_autocast: bool = False - # VSA parameters - VSA_sparsity: float = 0.0 # inference/validation sparsity - VSA_tile_size: int = 256 # VSA-H3 tile size (256 or 64); 64 = native Triton path + # VSA parameters. None means unset: __post_init__ fills the checkpoint's + # trained value when the checkpoint fixes one, else 0.0 and 256. + VSA_sparsity: float | None = None # inference/validation sparsity + VSA_tile_size: int | None = None # VSA-H3 tokens per tile (256, 128, 64); 128 = sm_100a CUDA only # V-MoBA parameters moba_config_path: str | None = None @@ -361,6 +362,13 @@ class FastVideoArgs: if env_backend is not None and backend_name_to_enum(env_backend) is not None: self.attention_backend = env_backend self._fold_vae_parallel_env() + # Runs after FASTVIDEO_ATTENTION_BACKEND is copied into attention_backend, + # so a backend chosen by that env var counts as the run's request. + self.pipeline_config.resolve_checkpoint_settings(self) + if self.VSA_sparsity is None: + self.VSA_sparsity = 0.0 + if self.VSA_tile_size is None: + self.VSA_tile_size = 256 import fastvideo.envs as envs envs.warn_deprecated_variables() self.check_fastvideo_args() @@ -800,13 +808,14 @@ class FastVideoArgs: "--VSA-sparsity", type=float, default=FastVideoArgs.VSA_sparsity, - help="Validation sparsity for VSA", + help="Validation sparsity for VSA (default: the checkpoint's trained value, else 0.0)", ) parser.add_argument( "--VSA-tile-size", type=int, default=FastVideoArgs.VSA_tile_size, - help="VSA-H3 tile size in tokens (256 or 64); 64 runs the native Triton block-sparse path", + help="VSA-H3 tile size in tokens (256, 128 or 64; default: the checkpoint's trained value, else 256); " + "64 runs the native Triton block-sparse path, 128 requires the sm_100a/sm_103a CUDA kernel", ) # Master port for distributed training/inference diff --git a/fastvideo/layers/pdd.py b/fastvideo/layers/pdd.py new file mode 100644 index 000000000..bcfb15fe4 --- /dev/null +++ b/fastvideo/layers/pdd.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Parallel Decoding Distillation (PDD) sampling primitives. + +A PDD student widens a diffusion transformer's final output projection so one +forward emits ``pdd_steps`` head predictions, one per interval of a fixed fine +time grid. Head ``j`` is the modality's rectified-flow direction frozen on +interval ``j``, so a block ``[start, end)`` of heads advances the state by +``sum_j (sigma(u_{j+1}) - sigma(u_j)) * q_j``. + +The sampler never materializes every head. For each runtime step it fuses one +block of head parameters inside the linear layer (:meth:`PDDReplicatedLinear.fuse`) +into their integration-weighted mean, and an ordinary Euler step over the +block's two node sigmas applies it: the block's total weight is exactly its +node-sigma increment. + +The fine grid lives on the base (unshifted) clock ``[0, 0.999]``; every +modality reaches its own noise level through its rational time shift. The +arithmetic matches the rectified-flow schedule the students are trained with. +""" + +from __future__ import annotations + +import contextlib +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod +from fastvideo.layers.quantization.base_config import QuantizationConfig + +# Rectified-flow students end their finite time domain at 0.999 and build the +# PDD fine grid on ``[0, 0.999]``. +PDD_GRID_MAX_T = 0.999 + + +def validate_finite_nonzero(value: torch.Tensor, name: str) -> None: + if not bool(torch.all(torch.isfinite(value) & (value != 0))): + raise ValueError(f"{name} must be finite and non-zero") + + +def validate_pdd_steps(value: Any, name: str = "pdd_steps") -> int: + """A PDD grid needs at least two intervals to have a block structure.""" + if isinstance(value, bool) or not isinstance(value, int) or value < 2: + raise ValueError(f"{name} must be an int >= 2, got {value!r}") + return value + + +# ---------------------------------------------------------------------------- +# Shifted rectified-flow arithmetic on the base clock +# ---------------------------------------------------------------------------- + + +def shifted_noise_amount(base: torch.Tensor, shift: float, *, max_t: float = PDD_GRID_MAX_T) -> torch.Tensor: + """Rational time shift ``f_s(u) = s * u * M / (u * (s - 1) + M)`` on ``[0, M]``. + + ``M = max_t`` is a fixed point of every shift, so the first grid node has + the same noise level in every modality. + """ + if shift == 1.0: + return base + return shift * base * max_t / (base * (shift - 1.0) + max_t) + + +def shifted_noise_delta( + base_start: torch.Tensor, + base_end: torch.Tensor, + shift: float, + *, + max_t: float = PDD_GRID_MAX_T, +) -> torch.Tensor: + """Return ``f_s(base_end) - f_s(base_start)`` without subtractive cancellation. + + The rational difference simplifies to + ``s * M^2 * (b - a) / (D(a) * D(b))`` with ``D(u) = M + (s - 1) * u``; the + identity shift keeps the direct difference bit-exact. + """ + if shift == 1.0: + return base_end - base_start + denominator_start = max_t + (shift - 1.0) * base_start + denominator_end = max_t + (shift - 1.0) * base_end + return shift * max_t**2 * (base_end - base_start) / (denominator_start * denominator_end) + + +@dataclass(frozen=True) +class PDDModalitySchedule: + """One modality's noise level as a function of the base clock: ``sigma(u) = f_shift(u)``.""" + + shift: float + max_t: float = PDD_GRID_MAX_T + + def sigma(self, base: torch.Tensor) -> torch.Tensor: + return shifted_noise_amount(base, self.shift, max_t=self.max_t) + + def integration_weights(self, fine_grid: torch.Tensor) -> torch.Tensor: + """``sigma(u_{j+1}) - sigma(u_j)`` for every fine-grid interval ``j``. + + Freezing head ``j`` on its interval advances the state by exactly this + coefficient. Negative for a descending grid. + """ + return shifted_noise_delta(fine_grid[:-1], fine_grid[1:], self.shift, max_t=self.max_t) + + +# ---------------------------------------------------------------------------- +# Fine grid and runtime partitions +# ---------------------------------------------------------------------------- + + +def pdd_fine_grid( + pdd_steps: int, + *, + max_t: float = PDD_GRID_MAX_T, + device: torch.device | str | None = None, +) -> torch.Tensor: + """The descending base-clock grid ``u_0 = max_t > ... > u_N = 0`` in float64.""" + grid = torch.linspace(max_t, 0.0, pdd_steps + 1, dtype=torch.float64, device=device) + return grid.clamp(max=max_t) + + +@dataclass(frozen=True) +class PDDSamplingPlan: + """One runtime partition of the fine grid plus every modality's coefficients.""" + + fine_grid: torch.Tensor + indices: torch.Tensor + node_sigmas: dict[str, torch.Tensor] + integration_weights: dict[str, torch.Tensor] + + @property + def pdd_steps(self) -> int: + return int(self.fine_grid.numel()) - 1 + + @property + def num_steps(self) -> int: + return int(self.indices.numel()) - 1 + + @property + def block_sizes(self) -> list[int]: + return (self.indices[1:] - self.indices[:-1]).tolist() + + def block(self, step: int) -> tuple[int, int]: + """Half-open fine-grid block ``[start, end)`` of runtime step *step*.""" + return int(self.indices[step]), int(self.indices[step + 1]) + + +def build_pdd_sampling_plan( + step_indices: Sequence[int], + schedules: Mapping[str, PDDModalitySchedule], + *, + device: torch.device | str | None = None, +) -> PDDSamplingPlan: + """Partition the fine grid into fused blocks for every modality. + + *step_indices* is the checkpoint's trained partition: fine-grid nodes that + increase strictly from 0 to the grid size, one fused block per consecutive + pair. ``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` validates it. + """ + max_ts = {schedule.max_t for schedule in schedules.values()} + if len(max_ts) != 1: + raise ValueError(f"Every PDD modality schedule must share max_t, got {sorted(max_ts)}") + fine_grid = pdd_fine_grid(int(step_indices[-1]), max_t=next(iter(max_ts)), device=device) + node_indices = torch.as_tensor(list(step_indices), dtype=torch.long, device=fine_grid.device) + weights = {} + for name, schedule in schedules.items(): + weights[name] = schedule.integration_weights(fine_grid) + validate_finite_nonzero(weights[name], f"PDD {name} integration weight") + return PDDSamplingPlan( + fine_grid=fine_grid, + indices=node_indices, + node_sigmas={ + name: schedule.sigma(fine_grid[node_indices]) + for name, schedule in schedules.items() + }, + integration_weights=weights, + ) + + +# ---------------------------------------------------------------------------- +# Widened output projection +# ---------------------------------------------------------------------------- + + +class PDDReplicatedLinear(ReplicatedLinear): + """Final linear layer widened to ``grid_size`` PDD heads. + + Output features are laid out head-major, ``(grid, output_size)``, so after + the architecture's own unpatchify the output channel axis reads + ``(grid, C)``. An ordinary ``forward(input)`` returns every head. + + Inside :meth:`fuse`, ``forward(input)`` instead returns the normalized + weighted average over the heads of one contiguous block ``[start, end)``, + computed by collapsing that block's parameter slices into one effective + linear. The architecture's inline final projection therefore needs no + seam, and the fused parameters are built inside ``forward``, where a + sharded weight is already gathered. + """ + + def __init__( + self, + input_size: int, + output_size: int, + *, + grid_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + input_size, + output_size * grid_size, + bias=bias, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + ) + if not isinstance(self.quant_method, UnquantizedLinearMethod): + raise NotImplementedError("PDDReplicatedLinear supports only unquantized weights") + self.head_output_size = int(output_size) + self.grid_size = grid_size + # Transient fusion state is a plain tuple, never a buffer, so it stays + # out of state_dict and can be restored as one value. + self._fusion_state: tuple[int, int, torch.Tensor, torch.dtype] | None = None + + @contextlib.contextmanager + def fuse( + self, + start: int, + end: int, + weights: torch.Tensor, + precision_decoding: torch.dtype, + ) -> Iterator[PDDReplicatedLinear]: + """Temporarily make ``forward(input)`` return the fused block output. + + *weights* are this modality's per-interval integration weights over the + whole grid (``[grid_size]``); the block ``[start, end)`` is fused with + their normalized form, so the layer emits the weighted mean and the + caller scales by the block's total weight. ``precision_decoding`` is + used only for the fused-parameter accumulation; the fused parameters + return to the layer's dtype before the projection runs. The previous + fusion state is restored on exit. + """ + previous = self._fusion_state + self._fusion_state = (start, end, weights, precision_decoding) + try: + yield self + finally: + self._fusion_state = previous + + def _fused_params( + self, + start: int, + end: int, + weights: torch.Tensor, + precision_decoding: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """The ``[O, I]`` fused weight (and bias) of block ``[start, end)``.""" + # Normalize in the incoming schedule precision, then cast only the + # fused-parameter accumulation. The layer emits the weighted mean and + # the caller multiplies by the block total, so together they represent + # ``sum_j w_j pred_j``. + block_weights = weights[start:end].to(device=self.weight.device) + block_scale = block_weights.sum() + normalized = (block_weights / block_scale).to(dtype=precision_decoding) + + def fuse_parameter(parameter: torch.Tensor) -> torch.Tensor: + heads = parameter.unflatten(0, (self.grid_size, -1))[start:end].to(dtype=precision_decoding) + return torch.einsum("n,n...->...", normalized, heads).to(dtype=parameter.dtype) + + return fuse_parameter(self.weight), None if self.bias is None else fuse_parameter(self.bias) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: # type: ignore[override] + fusion_state = self._fusion_state + if fusion_state is None: + return super().forward(x) + fused_weight, fused_bias = self._fused_params(*fusion_state) + output = F.linear(x, fused_weight, None if self.skip_bias_add else fused_bias) + return output, fused_bias if self.skip_bias_add else None + + +@contextlib.contextmanager +def fuse_pdd_heads( + linears: Mapping[str, PDDReplicatedLinear], + start: int, + end: int, + integration_weights: Mapping[str, torch.Tensor], + precision_decoding: torch.dtype, +) -> Iterator[None]: + """Fuse block ``[start, end)`` on every modality head at once.""" + with contextlib.ExitStack() as stack: + for name, linear in linears.items(): + stack.enter_context(linear.fuse(start, end, integration_weights[name], precision_decoding)) + yield + + +__all__ = [ + "PDD_GRID_MAX_T", + "PDDModalitySchedule", + "PDDReplicatedLinear", + "PDDSamplingPlan", + "build_pdd_sampling_plan", + "fuse_pdd_heads", + "pdd_fine_grid", + "shifted_noise_amount", + "shifted_noise_delta", + "validate_finite_nonzero", + "validate_pdd_steps", +] diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 1cd1d16d0..e25650466 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -3,7 +3,9 @@ from __future__ import annotations +import contextlib import math +from collections.abc import Iterator, Mapping from typing import Any import torch @@ -23,6 +25,7 @@ from fastvideo.distributed.parallel_state import get_sp_world_size, model_parall from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod from fastvideo.layers.mlp import MLP +from fastvideo.layers.pdd import PDDReplicatedLinear, fuse_pdd_heads from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps from fastvideo.logger import init_logger @@ -731,22 +734,56 @@ class MiniMaxH3Transformer3DModel(BaseDiT): prefix=f"{config.prefix}.norm_out", apply_silu=self.adaln_rank is None, ) - self.proj_out = ReplicatedLinear( - arch.hidden_size, - video_patch_dim, - bias=True, - quant_config=config.quant_config, - prefix=f"{config.prefix}.proj_out", - ) - self.audio_proj_out = ReplicatedLinear( - arch.hidden_size, - arch.audio_in_channels, - bias=True, - quant_config=config.quant_config, - prefix=f"{config.prefix}.audio_proj_out", - ) + # Parallel Decoding Distillation (PDD) students widen both output + # projections to ``pdd_steps`` heads, one per interval of a fixed fine + # time grid. The sampler fuses one block of heads per forward + # (``fuse_pdd_block``), so the fused output keeps the released width. + self.pdd_steps: int | None = arch.pdd_steps + self.proj_out = self._output_projection(arch.hidden_size, video_patch_dim, config, "proj_out") + self.audio_proj_out = self._output_projection(arch.hidden_size, arch.audio_in_channels, config, + "audio_proj_out") self.__post_init__() + def _output_projection(self, hidden_size: int, output_size: int, config: MiniMaxH3Config, + name: str) -> ReplicatedLinear: + if self.pdd_steps is None: + return ReplicatedLinear( + hidden_size, + output_size, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.{name}", + ) + return PDDReplicatedLinear( + hidden_size, + output_size, + grid_size=self.pdd_steps, + bias=True, + quant_config=config.quant_config, + prefix=f"{config.prefix}.{name}", + ) + + @property + def pdd_linears(self) -> dict[str, PDDReplicatedLinear]: + """The widened ``{"video": proj_out, "audio": audio_proj_out}`` heads.""" + if self.pdd_steps is None: + # Not AttributeError: nn.Module.__getattr__ would swallow it and + # report a misleading missing-attribute message. + raise RuntimeError("This MiniMax H3 transformer has no PDD heads (arch_config.pdd_steps is None)") + return {"video": self.proj_out, "audio": self.audio_proj_out} + + @contextlib.contextmanager + def fuse_pdd_block( + self, + start: int, + end: int, + integration_weights: Mapping[str, torch.Tensor], + precision_decoding: torch.dtype, + ) -> Iterator[None]: + """Fuse fine-grid block ``[start, end)`` on both widened heads for the enclosed forwards.""" + with fuse_pdd_heads(self.pdd_linears, start, end, integration_weights, precision_decoding): + yield + @staticmethod def _compile_setup_device(attention: MiniMaxH3Attention) -> torch.device: """Return the loaded device even when FP8 replaced the query weight.""" diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index e044822b5..b39f5c443 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -12,9 +12,12 @@ from typing import Any import torch +from fastvideo.attention.selector import (_active_component_attention_backend_scope, coerce_attn_backend, + get_env_variable_attn_backend) from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig -from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.configs.pipelines.minimax_h3 import (FASTH3_INFERENCE_FILE, FASTH3_INFERENCE_SCHEMA, + MiniMaxH3PipelineConfig) from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.models.hf_transformer_utils import get_diffusers_config @@ -29,6 +32,7 @@ from fastvideo.pipelines.basic.minimax_h3.stages import ( from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.platforms import AttentionBackendEnum logger = init_logger(__name__) @@ -38,6 +42,41 @@ logger = init_logger(__name__) _DENOISE_MODULE_NAMES = ("vae", "audio_vae", "transformer") +def _requested_attention_backend(fastvideo_args: FastVideoArgs) -> AttentionBackendEnum | None: + """The backend the transformer loader builds with; None means automatic selection. + + Follows ``PipelineComponentLoader.load_module``: an active per-component + request wins, then the run's ``attention_backend``, then + ``FASTVIDEO_ATTENTION_BACKEND``. Automatic selection never picks + VIDEO_SPARSE_ATTN_H3. + """ + scope = _active_component_attention_backend_scope() + if scope is not None: + if scope.backend is not None or not scope.consult_env: + return scope.backend + return get_env_variable_attn_backend() + requested = coerce_attn_backend(getattr(fastvideo_args, "attention_backend", None)) + return requested if requested is not None else get_env_variable_attn_backend() + + +def _checkpoint_has_vsa_gates(transformer_dir: Path) -> bool: + """Whether a transformer checkpoint ships VSA-H3 compression gates (``to_gate_compress``). + + Reads only the shard index, or the safetensors headers when there is none. + """ + indexes = sorted(transformer_dir.glob("*.safetensors.index.json")) + if indexes: + return any(".to_gate_compress." in name for index in indexes + for name in json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})) + from safetensors import safe_open + + for path in sorted(transformer_dir.glob("*.safetensors")): + with safe_open(str(path), framework="pt") as handle: + if any(".to_gate_compress." in name for name in handle.keys()): # noqa: SIM118 + return True + return False + + @dataclass(frozen=True) class _H3VideoGeometry: spatial_compression_ratio: int @@ -148,7 +187,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): assert patterns is not None # Keep the optional distilled schedule even when downloading only the # selected transformer partition. Otherwise Hub and local loads differ. - return [*patterns, "fastvideo_inference.json"] + return [*patterns, FASTH3_INFERENCE_FILE] def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: _apply_h3_checkpoint_arch_configs(self.model_path, fastvideo_args, self._extra_config_module_map) @@ -161,14 +200,85 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): raise ValueError(f"MiniMax-H3 {modality} scheduler must expose a positive finite shift, got {shift}.") self._load_checkpoint_schedule(fastvideo_args) + def _checkpoint_facts(self) -> tuple[Any, bool]: + """The checkpoint's parsed ``fastvideo_inference.json`` (None without one) and whether its transformer + carries VSA compression gates. + + Read once per resolved local path: ``_load_config`` replaces a Hub + repo id in ``model_path`` with its downloaded snapshot. + """ + root = Path(self.model_path) + cached = getattr(self, "_checkpoint_facts_cache", None) + if cached is None or cached[0] != root: + path = root / FASTH3_INFERENCE_FILE + contract = json.loads(path.read_text(encoding="utf-8")) if path.is_file() else None + cached = (root, contract, _checkpoint_has_vsa_gates(self._transformer_dir())) + self._checkpoint_facts_cache = cached + return cached[1], cached[2] + + def _checkpoint_contract(self) -> Any: + """The checkpoint's ``fastvideo_inference.json``, or None when there is none.""" + return self._checkpoint_facts()[0] + + def _transformer_dir(self) -> Path: + return Path(self.model_path) / self._extra_config_module_map.get("transformer", "transformer") + + def _require_checkpoint_attention_backend(self, requested: AttentionBackendEnum | None) -> None: + """Reject a backend other than VIDEO_SPARSE_ATTN_H3 for a transformer that carries VSA compression gates. + + MiniMax-H3 builds its compression gates only under that backend and loads + strictly, so any other backend would otherwise fail later, while loading + weights, as a missing ``to_gate_compress`` parameter. + """ + if requested == AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 or not self._checkpoint_facts()[1]: + return + raise ValueError( + f"This MiniMax-H3 checkpoint needs attention_backend=VIDEO_SPARSE_ATTN_H3: its " + f"{self._transformer_dir().name} carries VSA compression gates (to_gate_compress), and the gates are " + f"built only under that backend. This run requests " + f"{'automatic selection' if requested is None else requested.name}. Select VIDEO_SPARSE_ATTN_H3.") + + def _check_pdd_transformer(self, fastvideo_args: FastVideoArgs) -> None: + """Check this pipeline's transformer against the resolved PDD settings. + + A transformer whose config.json sets ``pdd_steps`` samples only with its + trained partition, which comes from the checkpoint's PDD fields; and + those fields describe the Ref2VA ``transformer_ref`` student. + """ + indices = getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) + config_path = self._transformer_dir() / "config.json" + widened = config_path.is_file() and json.loads(config_path.read_text(encoding="utf-8")).get("pdd_steps") + if widened and indices is None: + raise ValueError(f"{config_path} sets pdd_steps, but {self.model_path} has no PDD " + f"{FASTH3_INFERENCE_FILE}; a PDD student samples only with its trained partition.") + if indices is not None and not self._ref2va: + raise ValueError(f"{self.model_path} is a FastH3 Ref2VA PDD checkpoint; select " + "MiniMaxH3Ref2VAModularPipeline.") + + def _load_config(self, model_path: str) -> dict[str, Any]: + config = super()._load_config(model_path) + # super()._load_config resolved self.model_path to a local directory, + # downloading a Hub repo id. Check the transformer before any component + # loads, unless the caller supplied the transformer already built. + if getattr(self, "_check_transformer_before_loading", True): + self._check_pdd_transformer(self.fastvideo_args) + # A PDD checkpoint's backend comes from its fastvideo_inference.json + # (MiniMaxH3PipelineConfig.resolve_checkpoint_settings). + if getattr(self.fastvideo_args.pipeline_config, "pdd_step_indices", None) is None: + self._require_checkpoint_attention_backend(_requested_attention_backend(self.fastvideo_args)) + return config + def _load_checkpoint_schedule(self, fastvideo_args: FastVideoArgs) -> None: """A distilled export's schedule is explicit; never silently use a uniform grid.""" - path = Path(self.model_path) / "fastvideo_inference.json" - if not path.is_file(): + if getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) is not None: + # MiniMaxH3PipelineConfig.resolve_checkpoint_settings already applied + # this PDD checkpoint's file when the run's FastVideoArgs were built. return - contract = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(contract, dict) or contract.get("schema_version") != "fasth3-inference-contract-v1": - raise ValueError("Unsupported FastH3 fastvideo_inference.json schema.") + contract = self._checkpoint_contract() + if contract is None: + return + if not isinstance(contract, dict) or contract.get("schema_version") != FASTH3_INFERENCE_SCHEMA: + raise ValueError(f"Unsupported FastH3 {FASTH3_INFERENCE_FILE} schema.") steps = contract.get("dmd_denoising_steps") if (not isinstance(steps, list) or not steps or any(type(step) is not int or not 0 < step <= 1000 for step in steps) @@ -237,6 +347,8 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): fastvideo_args: FastVideoArgs, loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]: """Load the Qwen3-VL conditioner first; defer DiT and VAEs until after encode.""" + # _load_config checks the transformer against the run's settings before any component loads. + self._check_transformer_before_loading = loaded_modules is None or "transformer" not in loaded_modules if not self._defer_denoise_modules(fastvideo_args): if _use_taeh3_t2va(fastvideo_args, ref2va=self._ref2va): saved = list(self.required_config_modules) diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 9d273e48a..bef2a6513 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -89,6 +89,14 @@ class MiniMaxH3PackedLayout: latent_height: int latent_width: int num_audio_latents: int + # Per-reference conditioning spans in packed order (Ref2VA only; empty + # elsewhere): (kind, row_count, latent_shape) covering exactly the rows + # between the text tokens and the target audio. kind is "audio", "image", + # or "video"; latent_shape is the reference's raw latent (frames, height, + # width) for visual kinds and (0, 0, 0) for audio. A video reference with + # sound contributes two entries, its audio rows first. VSA-H3 uses them to + # tile (and sparsify) reference-video regions in place. + reference_segments: tuple[tuple[str, int, tuple[int, int, int]], ...] = () def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: @@ -395,10 +403,13 @@ def build_ref2va_packed_sequence( video_indices: list[torch.Tensor] = [] audio_indices: list[torch.Tensor] = [] + reference_segments: list[tuple[str, int, tuple[int, int, int]]] = [] cursor = num_text_tokens rotary_time = float(num_text_tokens) for reference, visual_row_count in zip(references, visual_row_counts, strict=True): + visual_latent_shape = (reference.num_latent_frames, reference.latent_height, reference.latent_width) if reference.media_type == "image": + reference_segments.append(("image", visual_row_count, visual_latent_shape)) rows = slice(cursor, cursor + visual_row_count) cursor = rows.stop video_indices.append(torch.arange(rows.start, rows.stop)) @@ -413,6 +424,7 @@ def build_ref2va_packed_sequence( rotary_time += 1.0 elif reference.media_type == "audio": count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + reference_segments.append(("audio", count, (0, 0, 0))) rows = slice(cursor, cursor + count) cursor = rows.stop audio_indices.append(torch.arange(rows.start, rows.stop)) @@ -420,6 +432,9 @@ def build_ref2va_packed_sequence( rotary_time += float(reference.num_audio_latents) elif reference.media_type == "video": audio_count = reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS if reference.has_audio else 0 + if audio_count: + reference_segments.append(("audio", audio_count, (0, 0, 0))) + reference_segments.append(("video", visual_row_count, visual_latent_shape)) audio_rows = slice(cursor, cursor + audio_count) video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count) cursor = video_rows.stop @@ -483,6 +498,7 @@ def build_ref2va_packed_sequence( latent_height=latent_height, latent_width=latent_width, num_audio_latents=num_audio_latents, + reference_segments=tuple(reference_segments), ) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index bce4ea978..c59ac6f11 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -3,6 +3,7 @@ from __future__ import annotations +import contextlib from typing import Any import torch @@ -12,6 +13,8 @@ from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.hooks.activation_trace import trace_step +from fastvideo.layers.pdd import PDDModalitySchedule, PDDSamplingPlan, build_pdd_sampling_plan +from fastvideo.logger import init_logger from fastvideo.profiler import nvtx_range, profiler_region from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_KEYFRAME_NOISE_AUG, @@ -25,6 +28,18 @@ from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.utils import get_compute_dtype +logger = init_logger(__name__) + + +def scale_pdd_initial_noise(video_rows: torch.Tensor, audio_rows: torch.Tensor, layout: MiniMaxH3PackedLayout, + plan: PDDSamplingPlan) -> None: + """Start PDD targets at the first node's noise level (``0.999 * eps``), as the students are trained. + + Latent preparation draws unit noise; conditioning rows are left untouched. + """ + video_rows[layout.num_condition_video_rows:] *= float(plan.node_sigmas["video"][0]) + audio_rows[layout.num_condition_audio_rows:] *= float(plan.node_sigmas["audio"][0]) + def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> Any: """Builder instance when the transformer resolved to VSA-H3, else None. @@ -45,8 +60,16 @@ def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> return backend.get_builder_cls()() -def _h3_vsa_prefix_segments(layout: MiniMaxH3PackedLayout, patch_size: tuple[int, int, int]) -> tuple[int, ...]: - """Segment sizes preceding the generated-video tail, validated against the layout.""" +def _h3_vsa_single_region_segments( + layout: MiniMaxH3PackedLayout, + patch_size: tuple[int, int, int], +) -> tuple[int | tuple[int, int, int], ...]: + """VSA-H3 packed segments with the generated video as the only sparse region. + + Returns ``(text rows, condition rows, audio rows, target video latent + shape)`` for ``MiniMaxH3VSAMetadataBuilder.build``, validated against the + layout; every conditioning row stays dense. + """ n_text = int(layout.text_indices.numel()) n_cond = int(layout.num_condition_video_rows) n_audio = int(layout.audio_indices.numel()) @@ -56,7 +79,45 @@ def _h3_vsa_prefix_segments(layout: MiniMaxH3PackedLayout, patch_size: tuple[int raise ValueError("VSA-H3 supports the standard [text|cond|audio|video] packing only; " f"segments ({n_text}, {n_cond}, {n_audio}) + video {n_video} do not sum to " f"sequence length {layout.sequence_length}.") - return n_text, n_cond, n_audio + return n_text, n_cond, n_audio, (layout.num_video_latent_frames, layout.latent_height, layout.latent_width) + + +def _h3_vsa_ref2va_segments( + layout: MiniMaxH3PackedLayout, + patch_size: tuple[int, int, int], +) -> tuple[int | tuple[int, int, int], ...]: + """VSA-H3 packed segments of a Ref2VA layout with every reference video as a sparse region. + + Returns the segments for ``MiniMaxH3VSAMetadataBuilder.build`` in true + packed order, ``[text | reference spans... | target audio | target + video]``: each reference VIDEO is a region (its latent shape); text, + every reference's audio rows, image references, and the target audio are + dense row counts; the target video is the last region. + """ + if not layout.reference_segments: + raise ValueError("Ref2VA VSA sparsification needs per-reference spans on the layout; " + "this layout carries none (T2VA/FL2VA packing).") + cursor = int(layout.text_indices.numel()) + segments: list[int | tuple[int, int, int]] = [cursor] + for kind, rows, latent_shape in layout.reference_segments: + rows = int(rows) + if kind == "video": + segments.append((int(latent_shape[0]), int(latent_shape[1]), int(latent_shape[2]))) + else: + # audio rows and single-frame image latents stay dense + segments.append(rows) + cursor += rows + n_target_audio = int(layout.audio_indices.numel()) - int(layout.num_condition_audio_rows) + segments.append(n_target_audio) + cursor += n_target_audio + segments.append((layout.num_video_latent_frames, layout.latent_height, layout.latent_width)) + n_target_video = ((layout.num_video_latent_frames // patch_size[0]) * (layout.latent_height // patch_size[1]) * + (layout.latent_width // patch_size[2])) + if cursor + n_target_video != layout.sequence_length: + raise ValueError(f"Ref2VA reference spans {layout.reference_segments} + target segments do not tile the " + f"packed sequence: reached row {cursor} + video {n_target_video} != " + f"{layout.sequence_length}.") + return tuple(segments) class MiniMaxH3DenoisingStage(PipelineStage): @@ -86,6 +147,30 @@ class MiniMaxH3DenoisingStage(PipelineStage): # clean time (1 - sigma); passing integer rungs to step() is wrong. scheduler.set_timesteps(sigmas=sigmas, device=device) + def _pdd_sampling_plan(self, fastvideo_args: FastVideoArgs, device: torch.device) -> PDDSamplingPlan | None: + """The fused-block plan of a PDD checkpoint, or None for any other checkpoint. + + ``pdd_step_indices`` is the checkpoint's trained partition, which + ``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` read and checked + against the transformer's ``pdd_steps``. Each modality's node sigmas + come from its own scheduler shift on the shared base clock. + """ + indices = getattr(fastvideo_args.pipeline_config, "pdd_step_indices", None) + if indices is None: + return None + schedules = { + "video": PDDModalitySchedule(shift=float(self.scheduler.shift)), + "audio": PDDModalitySchedule(shift=float(self.audio_scheduler.shift)), + } + return build_pdd_sampling_plan(indices, schedules, device=device) + + def _pdd_head_fusion(self, plan: PDDSamplingPlan | None, index: int) -> contextlib.AbstractContextManager[None]: + """Fuse runtime step *index*'s block of heads inside both widened output projections.""" + if plan is None: + return contextlib.nullcontext() + block_start, block_end = plan.block(index) + return self.transformer.fuse_pdd_block(block_start, block_end, plan.integration_weights, torch.float32) + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) @@ -117,7 +202,16 @@ class MiniMaxH3DenoisingStage(PipelineStage): device = get_local_torch_device() dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps - if dmd_steps is None: + pdd_plan = self._pdd_sampling_plan(fastvideo_args, device) + if pdd_plan is not None: + # One runtime step fuses one block of heads into their weighted + # mean velocity; the ordinary Euler step over the block's two node + # sigmas then applies the block's total increment per modality. + self.scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["video"].to(torch.float32), device=device) + self.audio_scheduler.set_timesteps(sigmas=pdd_plan.node_sigmas["audio"].to(torch.float32), device=device) + logger.info("MiniMax-H3 PDD denoising: %d fused block(s) %s over a %d-interval fine grid.", + pdd_plan.num_steps, pdd_plan.block_sizes, pdd_plan.pdd_steps) + elif dmd_steps is None: self.scheduler.set_timesteps(batch.num_inference_steps, device=device) self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) else: @@ -153,25 +247,44 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args) if vsa_metadata_builder is not None: vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size - vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size) # Per-request knobs (sweeps flip these between generate_video calls # without respawning workers); mode None defers to the env default. vsa_mode = batch.extra.get("vsa_mode", "exempt") if vsa_mode not in ("exempt", "compete"): raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.") vsa_exempt = vsa_mode == "exempt" + # PDD checkpoints set a reference keep rate: each reference video is + # then its own sparse region keeping that fraction of its tiles. + # Without one, every conditioning row stays dense. + vsa_ref_keep_rate = getattr(fastvideo_args.pipeline_config, "vsa_ref_keep_rate", None) + if vsa_ref_keep_rate is None: + vsa_packed_segments = _h3_vsa_single_region_segments(layout, vsa_patch_size) + else: + vsa_packed_segments = _h3_vsa_ref2va_segments(layout, vsa_patch_size) vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ())) vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) - # Run-level tile geometry (256 default, 64 = native Triton path), - # plumbed like the run-level sparsity; the builder validates the - # value against VSA_H3_TILE_SHAPES. + # Run-level tile geometry (256 default, 64 = native Triton path, + # 128 = sm_100a CUDA), plumbed like the run-level sparsity; the + # builder validates the value against VSA_H3_TILE_SHAPES. vsa_tile_size = int(fastvideo_args.VSA_tile_size) + if vsa_ref_keep_rate is not None: + num_reference_regions = sum(isinstance(segment, tuple) for segment in vsa_packed_segments) - 1 + run_sparsity = float(batch.VSA_sparsity) + if run_sparsity <= 0.0: + effect = "VSA_sparsity=0 keeps every region dense" + else: + effect = (f"each reference keeps {vsa_ref_keep_rate:g} of its tiles and the target keeps " + f"{1.0 - run_sparsity:g}") + logger.info("MiniMax-H3 VSA-H3: %d reference video region(s); %s; %d-token tiles.", + num_reference_regions, effect, vsa_tile_size) try: if full_cpu_offload: self.transformer.to(device) batch.latents = batch.latents.to(device) batch.audio_latents = batch.audio_latents.to(device) + if pdd_plan is not None: + scale_pdd_initial_noise(batch.latents, batch.audio_latents, layout, pdd_plan) # The stage range groups the complete denoising loop while the # indexed model ranges retain timing detail for every H3 block. @@ -187,15 +300,14 @@ class MiniMaxH3DenoisingStage(PipelineStage): vsa_sparsity = 0.0 if index < vsa_dense_first_n else float(batch.VSA_sparsity) attn_metadata = vsa_metadata_builder.build( current_timestep=index, - raw_latent_shape=(layout.num_video_latent_frames, layout.latent_height, - layout.latent_width), patch_size=vsa_patch_size, VSA_sparsity=vsa_sparsity, - prefix_segments=vsa_prefix_segments, + packed_segments=vsa_packed_segments, device=device, exempt=vsa_exempt, dense_layers=vsa_dense_layers, tile_size=vsa_tile_size, + ref_keep_rate=vsa_ref_keep_rate, ) # Under torch.compile(mode="reduce-overhead") each denoising # step must be marked, or cudagraph trees flag cross-step @@ -203,7 +315,7 @@ class MiniMaxH3DenoisingStage(PipelineStage): # CUDAGraphs that has been overwritten" (surfaces at sp=1; # sp>1 is masked by collective-induced graph breaks). torch.compiler.cudagraph_mark_step_begin() - with trace_step(index), set_forward_context( + with self._pdd_head_fusion(pdd_plan, index), trace_step(index), set_forward_context( current_timestep=index, attn_metadata=attn_metadata, forward_batch=batch, @@ -247,4 +359,4 @@ class MiniMaxH3DenoisingStage(PipelineStage): return batch -__all__ = ["MiniMaxH3DenoisingStage"] +__all__ = ["MiniMaxH3DenoisingStage", "scale_pdd_initial_noise"] diff --git a/fastvideo/tests/attention/test_vsa_h3_backward.py b/fastvideo/tests/attention/test_vsa_h3_backward.py index 8d2d4395f..fed125f75 100644 --- a/fastvideo/tests/attention/test_vsa_h3_backward.py +++ b/fastvideo/tests/attention/test_vsa_h3_backward.py @@ -23,10 +23,9 @@ _DIM = 128 def _build_meta(device, sparsity=0.5): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=sparsity, - prefix_segments=_SPEC["prefix_segments"], + packed_segments=(*_SPEC["prefix_segments"], _SPEC["raw_latent_shape"]), device=device, ) diff --git a/fastvideo/tests/attention/test_vsa_h3_metadata.py b/fastvideo/tests/attention/test_vsa_h3_metadata.py index 24292f26b..7823ccc9c 100644 --- a/fastvideo/tests/attention/test_vsa_h3_metadata.py +++ b/fastvideo/tests/attention/test_vsa_h3_metadata.py @@ -11,7 +11,7 @@ import torch.nn.functional as F from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder, _build_block_mask, - _pool_tiles, _validate_h3_tile_geometry, + _pool_tiles, _validate_h3_segment_geometry, token_tile_and_valid) _720P = dict(raw_latent_shape=(30, 44, 80), patch_size=(1, 2, 2), prefix_segments=(512, 1760, 400)) @@ -30,10 +30,9 @@ _CPU = torch.device("cpu") def _build(spec, sparsity=0.0, device=_CPU, tile_size=_TILE_ELEMS): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=spec["raw_latent_shape"], patch_size=spec["patch_size"], VSA_sparsity=sparsity, - prefix_segments=spec["prefix_segments"], + packed_segments=(*spec["prefix_segments"], spec["raw_latent_shape"]), device=device, tile_size=tile_size, ) @@ -92,16 +91,16 @@ def test_mask_policy(): k_vid = math.ceil(0.1 * V) scores = torch.randn(1, 2, n, n) - exempt = _build_block_mask(scores, P, V, 0.9, exempt=True) + exempt = _build_block_mask(scores, P, 0.9, True, meta.video_tile_spans, meta.span_sparsities) assert exempt[:, :, :P].all(), "prefix queries must be dense" assert exempt[..., :P].all(), "prefix keys must be visible to every query" assert (exempt[:, :, P:, P:].sum(-1) == k_vid).all(), "video rows select exactly k_vid video tiles" - compete = _build_block_mask(scores, P, V, 0.9, exempt=False) + compete = _build_block_mask(scores, P, 0.9, False, meta.video_tile_spans, meta.span_sparsities) assert compete[:, :, :P].all() assert (compete[:, :, P:].sum(-1) == min(k_vid + P, n)).all(), "budget-matched top-k" - dense = _build_block_mask(scores, P, V, 0.0, exempt=True) + dense = _build_block_mask(scores, P, 0.0, True, meta.video_tile_spans, meta.span_sparsities) assert dense.all(), "sparsity 0 must select everything" @@ -115,7 +114,8 @@ def test_sparsity_zero_matches_dense_sdpa(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.0, exempt=True) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.0, True, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -134,7 +134,8 @@ def test_prefix_queries_stay_dense_at_high_sparsity(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) for exempt in (True, False): - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.75, exempt=exempt) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.75, exempt, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -220,7 +221,8 @@ def test_sparsity_zero_matches_dense_sdpa_tile64(): scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes, meta.tile_elems), _pool_tiles(tk, meta.variable_block_sizes, meta.tile_elems).transpose(-2, -1)) - mask = _build_block_mask(scores, meta.num_prefix_tiles, meta.num_video_tiles, 0.0, exempt=True) + mask = _build_block_mask(scores, meta.num_prefix_tiles, 0.0, True, meta.video_tile_spans, + meta.span_sparsities) sparse_out = impl.postprocess_output(reference_sparse_attention(tq, tk, tv, mask, meta), meta) dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) @@ -230,18 +232,17 @@ def test_sparsity_zero_matches_dense_sdpa_tile64(): def test_geometry_guard_enforces_tile64_bound(): """A 65-token tile passes the 256 bound but must fail the 64 one.""" meta = _build(_TINY64, tile_size=64) - prefix = tuple(s for s in _TINY64["prefix_segments"] if s > 0) - dit_shape = (9, 10, 13) + segments = (*(s for s in _TINY64["prefix_segments"] if s > 0), (9, 10, 13)) sizes = meta.variable_block_sizes.clone() sizes[0] = 65 with pytest.raises(ValueError, match="tile sizes out of bounds"): - _validate_h3_tile_geometry(prefix, dit_shape, sizes, meta.untile_combined_index, 64) + _validate_h3_segment_geometry(segments, sizes, meta.untile_combined_index, 64) # the untampered tile-64 geometry passes its own bound - _validate_h3_tile_geometry(prefix, dit_shape, meta.variable_block_sizes, meta.untile_combined_index, 64) + _validate_h3_segment_geometry(segments, meta.variable_block_sizes, meta.untile_combined_index, 64) def test_builder_rejects_unknown_tile_size(): - for bad in (0, 128, 512): + for bad in (0, 96, 512): with pytest.raises(ValueError, match="tile_size"): _build(_TINY, tile_size=bad) diff --git a/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py new file mode 100644 index 000000000..ec48bbc84 --- /dev/null +++ b/fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for VSA-H3 reference-video regions and 128-token tiles. + +With vsa_ref_policy="p2_multi_region" every Ref2VA reference video is its own +sparse region, tiled in place at its packed offset, while text, audio, and +image references stay dense. Each video query keeps its own top-k of every +region: the reference keep rate may differ from the target's. The kernel +route for 128-token tiles is exercised through a stand-in for the sm_100a +extension; real kernels are not needed. +""" + +import hashlib +import math + +import pytest +import torch +import torch.nn.functional as F + +import fastvideo.attention.backends.video_sparse_attn_h3 as vsa_h3 +import fastvideo.envs as envs +from fastvideo.attention.backends.video_sparse_attn import compute_topk +from fastvideo.attention.backends.video_sparse_attn_h3 import (_TILE_ELEMS, VSA_H3_TILE_SHAPES, MiniMaxH3VSAImpl, + MiniMaxH3VSAMetadataBuilder, _build_block_mask, + _h3_segment_tile_geometry, + _pool_tiles, token_tile_and_valid) +from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_ref2va_packed_sequence +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import _h3_vsa_ref2va_segments + +_CPU = torch.device("cpu") +_PATCH = (1, 2, 2) + +# packed order: [text 37 | ref_audio 9 | ref_video (5,4,6)=120 | tgt_audio 11 +# | tgt_video (8,4,6)=192] -> S=369; raw latents under patch (1,2,2); the +# reference video keeps 25% of its tiles +_R2V = dict( + packed_segments=(37, 9, (5, 8, 12), 11, (8, 8, 12)), + ref_keep_rate=0.25, +) +_R2V_SEQ = 369 +_R2V_DENSE_RUNS = ((0, 37), (37, 46), (166, 177)) # packed [start, end) of text/ref_audio/tgt_audio + + +def _build_r2v(sparsity=0.0, tile_size=_TILE_ELEMS, device=_CPU, **overrides): + call = dict(_R2V, **overrides) + return MiniMaxH3VSAMetadataBuilder().build( + current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=sparsity, + device=device, + tile_size=tile_size, + **call, + ) + + +def _impl(head_size=8): + return MiniMaxH3VSAImpl(num_heads=2, head_size=head_size, causal=False, softmax_scale=head_size**-0.5) + + +def _mask(meta, scores, sparsity): + return _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans, + meta.span_sparsities) + + +def _assert_multi_region_metadata(metadata, *, expected_reference_video_regions, target_sparsity, ref_keep_rate): + """Test oracle: metadata holds one sparse span per reference video plus the target, with the policy's sparsities. + + The last span is the generated target and every earlier span is a + reference video; references keep ``ref_keep_rate`` of their tiles and the + target keeps ``1 - target_sparsity``. Text, reference audio and images, + and target audio are dense prefix tiles. + """ + expected_span_count = expected_reference_video_regions + 1 + if len(metadata.video_tile_spans) != expected_span_count: + raise ValueError("Multi-region metadata must contain one sparse span per reference video plus the target; " + f"expected {expected_span_count} spans, got {metadata.video_tile_spans}. Metadata whose only " + "sparse region is the generated video keeps reference-video conditioning dense.") + if len(metadata.span_sparsities) != expected_span_count: + raise ValueError(f"{len(metadata.span_sparsities)} span sparsities for {expected_span_count} spans.") + cursor = metadata.num_prefix_tiles + for index, (start, end) in enumerate(metadata.video_tile_spans): + if start != cursor or end <= start: + raise ValueError(f"span {index} is {(start, end)} after tile {cursor}; spans must be contiguous.") + cursor = end + if cursor != metadata.num_prefix_tiles + metadata.num_video_tiles: + raise ValueError(f"the spans end at tile {cursor} instead of covering every video tile.") + if target_sparsity <= 0.0: + expected = (0.0, ) * expected_span_count + else: + reference_sparsity = target_sparsity if ref_keep_rate is None else 1.0 - ref_keep_rate + if expected_reference_video_regions and reference_sparsity <= 0.0: + raise ValueError("sparse target attention would leave reference-video conditioning dense " + f"(ref_keep_rate={ref_keep_rate!r}).") + expected = (reference_sparsity, ) * expected_reference_video_regions + (target_sparsity, ) + if any(not math.isclose(actual, wanted, rel_tol=0.0, abs_tol=1e-12) + for actual, wanted in zip(metadata.span_sparsities, expected, strict=True)): + raise ValueError(f"per-region sparsities {metadata.span_sparsities} differ from the policy's {expected}.") + + +def _token_allow(meta, mask, b=0, h=0): + """Block mask mapped to packed-token coordinates: allow[i, j].""" + tile_of = meta.untile_combined_index // meta.tile_elems + return mask[b, h][tile_of][:, tile_of] + + +def _sparse_attention_oracle(query, key, value, mask, meta): + """SDPA over the padded tile buffer with the block mask expanded to tokens.""" + token_tile, token_valid = token_tile_and_valid(meta.variable_block_sizes, meta.tile_elems) + out = torch.empty_like(query) + for b in range(query.shape[0]): + for h in range(query.shape[2]): + allow = mask[b, h][token_tile][:, token_tile] & token_valid[None, :] + bias = torch.zeros(allow.shape, dtype=query.dtype, device=query.device) + bias.masked_fill_(~allow, float("-inf")) + out[b, :, h] = F.scaled_dot_product_attention(query[b, :, h][None], key[b, :, h][None], + value[b, :, h][None], attn_mask=bias[None])[0] + return out + + +# --------------------------------------------------------------------------- +# Packed-order segments for representative Ref2VA layouts +# --------------------------------------------------------------------------- + + +def _reference(kind, *, frames=1, height=4, width=6, audio=0): + return MiniMaxH3PreparedReference(media_type=kind, + has_audio=audio > 0, + num_latent_frames=frames, + latent_height=height, + latent_width=width, + num_audio_latents=audio) + + +def _layout(references, *, text=7, target=(3, 4, 6), audio=5): + tags = torch.full((text, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long) + return build_ref2va_packed_sequence(tags, references, *target, audio, _PATCH) + + +# (references, expected packed segments, expected packed row of each video region). +# Rows: an image (1, 4, 6) latent is 1*2*3 = 6 rows; a (5, 4, 6) video is 30; the +# (3, 4, 6) target is 18; audio rows are 2 per latent; text is 7. +_LAYOUTS = { + "images_only": ([_reference("image"), _reference("image", height=8, width=4)], (7, 6, 8, 10, (3, 4, 6)), (31, )), + "image_and_silent_video": ([_reference("image"), _reference("video", frames=5)], (7, 6, (5, 4, 6), 10, + (3, 4, 6)), (13, 53)), + "video_with_audio_and_audio": ([_reference("video", frames=5, audio=4), + _reference("audio", audio=3)], (7, 8, (5, 4, 6), 6, 10, (3, 4, 6)), (15, 61)), + "two_videos": ([_reference("video", frames=2, height=8, width=4), + _reference("video", frames=5, audio=2)], (7, (2, 8, 4), 4, (5, 4, 6), 10, (3, 4, 6)), + (7, 27, 67)), +} + + +@pytest.mark.parametrize("name", sorted(_LAYOUTS)) +@pytest.mark.parametrize("tile_size", [64, 128, 256]) +def test_h3_vsa_ref2va_segments_tile_layouts_in_packed_order(name, tile_size): + references, segments, offsets = _LAYOUTS[name] + layout = _layout(references) + assert _h3_vsa_ref2va_segments(layout, _PATCH) == segments + videos = [segment for segment in segments if isinstance(segment, tuple)] + + meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.9, + packed_segments=segments, + device=_CPU, + tile_size=tile_size, + ref_keep_rate=0.25) + assert meta.total_seq_length == layout.sequence_length + # References keep 25% of their tiles, distinct from the target's 10%. + _assert_multi_region_metadata(meta, + expected_reference_video_regions=len(videos) - 1, + target_sparsity=0.9, + ref_keep_rate=0.25) + # Dense rows (text, audio, image references) map to prefix tiles only; + # each video region's rows map to its own span. + tile_of = meta.untile_combined_index // meta.tile_elems + region_rows = torch.zeros(layout.sequence_length, dtype=torch.bool) + for (t, h, w), offset, (start, end) in zip(videos, offsets, meta.video_tile_spans, strict=True): + rows = slice(offset, offset + (t // _PATCH[0]) * (h // _PATCH[1]) * (w // _PATCH[2])) + region_rows[rows] = True + assert bool(((tile_of[rows] >= start) & (tile_of[rows] < end)).all()) + assert bool((tile_of[~region_rows] < meta.num_prefix_tiles).all()) + x = torch.randn(1, layout.sequence_length, 2, 4) + assert torch.equal(_impl().tile(x, meta)[:, meta.untile_combined_index], x) + + +def test_h3_vsa_ref2va_segments_missing_reference_spans(): + layout = _layout([_reference("image")]) + stripped = type(layout)(**{**layout.__dict__, "reference_segments": ()}) + with pytest.raises(ValueError, match="per-reference spans"): + _h3_vsa_ref2va_segments(stripped, _PATCH) + + +# --------------------------------------------------------------------------- +# Multi-region geometry and per-region top-k +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("tile_size,width", [(64, 4), (128, 8)]) +def test_build_multi_region_native_tile_geometry(tile_size, width): + """Hand-computed per-region (4,4,4)/(4,4,8) tile ids and sizes at true offsets.""" + meta = _build_r2v(tile_size=tile_size) + assert meta.total_seq_length == _R2V_SEQ + assert meta.num_prefix_tiles == 3 # 37, 9, 11 each fit one tile + per_region = 4 if tile_size == 64 else 2 + assert meta.num_video_tiles == 2 * per_region + assert meta.video_tile_spans == ((3, 3 + per_region), (3 + per_region, 3 + 2 * per_region)) + assert meta.variable_block_sizes[:3].tolist() == [37, 9, 11] + + idx = meta.untile_combined_index + for (t, h, w), start, first_tile in (((5, 4, 6), 46, 3), ((8, 4, 6), 177, 3 + per_region)): + n_t, n_h, n_w = -(-t // 4), -(-h // 4), -(-w // width) + expected_sizes = torch.tensor([ + min(4, t - 4 * tt) * min(4, h - 4 * hh) * min(width, w - width * ww) for tt in range(n_t) + for hh in range(n_h) for ww in range(n_w) + ]) + n_region = n_t * n_h * n_w + assert torch.equal(meta.variable_block_sizes[first_tile:first_tile + n_region], expected_sizes) + row = torch.arange(t * h * w) + row_t, row_h, row_w = row // (h * w), (row // w) % h, row % w + expected_tile = first_tile + ((row_t // 4) * n_h + row_h // 4) * n_w + row_w // width + assert torch.equal(idx[start:start + t * h * w] // tile_size, expected_tile) + # Each prefix tile stays inside ONE dense segment even though the target + # audio sits between the two video regions. + for seg_start, seg_end in _R2V_DENSE_RUNS: + for tile in torch.unique(idx[seg_start:seg_end] // tile_size).tolist(): + assert tile < meta.num_prefix_tiles + rows = torch.nonzero(idx // tile_size == tile).flatten() + assert rows.min() >= seg_start and rows.max() < seg_end + + +# SHA-256 of the single-region tile geometry (partition indices, block sizes, +# untile index, prefix and video tile counts) recorded at commit 9edc8adf5, +# where MiniMaxH3VSAMetadataBuilder supported only one sparse video region. The +# hashes are fixed so the geometry is compared with that recorded output +# instead of with the code under test. +_SINGLE_REGION_GEOMETRY_SHA256 = { + ((37, 9, 120, 11), (8, 4, 6), 64): "e246d39c1be206be5b9ad938438d0594fea013d5e8c4fcb08bce05face5d4220", + ((37, 9, 120, 11), (8, 4, 6), 128): "bce9def60b23ef5dc71d790e367480c2c22e0e6f414678b72844b808d7f96bd9", + ((37, 9, 120, 11), (8, 4, 6), 256): "a710f91b80cb02bf8bbf2d37f06b612c2a4e805356c81f55c744dd066417a535", + ((512, 1200, 216), (16, 24, 40), 64): "22d1581bd63b2dfb77951e31e788aa93cd53b50c89290312b1340605889f21a9", + ((512, 1200, 216), (16, 24, 40), 128): "23b19d850777a4a6063164c5ceea33991fca5d0e7ded67d22be0a2d24c9ff0a0", + ((512, 1200, 216), (16, 24, 40), 256): "93d86b693252631c8f35efd2fdbdc216e2979c574d10c5d8d5b4fb8288750462", + ((5, 3, 250), (6, 10, 14), 64): "3226fc23f45bec409992d2bf0f42794c609d53bac401c5162e7c37615904cc41", + ((5, 3, 250), (6, 10, 14), 128): "2645e5a4c49b30de3f355ee0163e8703de2676a9a21e6d04b6b1a210a01242ab", + ((5, 3, 250), (6, 10, 14), 256): "220fa91d488c88832d08177e28d8bb57f2477b7be36677302ecfefb9cef7c56c", +} + + +def _geometry_sha256(outputs): + digest = hashlib.sha256() + for item in outputs: + if torch.is_tensor(item): + tensor = item.detach().cpu().contiguous() + digest.update(str(tensor.dtype).encode()) + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(tensor.numpy().tobytes()) + else: + digest.update(repr(int(item)).encode()) + return digest.hexdigest() + + +@pytest.mark.parametrize("prefix_segments,dit_seq_shape,tile_size", sorted(_SINGLE_REGION_GEOMETRY_SHA256)) +def test_h3_segment_tile_geometry_single_region_matches_frozen_hashes(prefix_segments, dit_seq_shape, tile_size): + geometry = _h3_segment_tile_geometry((*prefix_segments, dit_seq_shape), _CPU, VSA_H3_TILE_SHAPES[tile_size])[:5] + assert _geometry_sha256(geometry) == _SINGLE_REGION_GEOMETRY_SHA256[prefix_segments, dit_seq_shape, tile_size] + + +def test_build_single_region_one_generated_video_span(): + single = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.5, + packed_segments=(37, 9, 120, 11, (8, 8, 12)), + device=_CPU, + tile_size=128) + # The 37, 9, 120, and 11 dense rows each fit one 128-token tile; the (8, 4, 6) video grid + # makes 2 * 1 * 1 (4, 4, 8) tiles. + assert (single.num_prefix_tiles, single.num_video_tiles) == (4, 2) + assert single.video_tile_spans == ((4, 6), ) + assert single.span_sparsities == (0.5, ) + + +def test_build_block_mask_per_region_topk_and_reference_keep_rate(): + """Every video query keeps exactly k_i tiles of EACH region; prefix stays dense.""" + torch.manual_seed(3) + meta = _build_r2v(sparsity=0.75, tile_size=64, ref_keep_rate=0.5) + P = meta.num_prefix_tiles + n = P + meta.num_video_tiles + scores = torch.randn(1, 2, n, n) + mask = _mask(meta, scores, 0.75) + assert mask[:, :, :P].all(), "prefix queries dense" + assert mask[..., :P].all(), "prefix keys visible to every query" + for (start, end), sparsity in zip(meta.video_tile_spans, meta.span_sparsities, strict=True): + assert (mask[:, :, P:, start:end].sum(-1) == compute_topk(sparsity, end - start)).all(), (start, end) + assert meta.span_sparsities == pytest.approx((0.5, 0.75)) + + # A span at sparsity 0 keeps every column; the other span keeps its own top-k. + ref_span, tgt_span = meta.video_tile_spans + mask = _build_block_mask(scores, P, 0.75, meta.exempt, meta.video_tile_spans, (0.0, 0.75)) + assert mask[..., ref_span[0]:ref_span[1]].all(), "a sparsity-0 span's columns are dense" + assert (mask[:, :, P:, tgt_span[0]:tgt_span[1]].sum(-1) == compute_topk(0.75, tgt_span[1] - tgt_span[0])).all() + + keep_quarter = _build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.25) + assert keep_quarter.span_sparsities == pytest.approx((0.75, 0.9)) + # A dense build (sparsity 0: dense steps) ignores the reference override. + assert _build_r2v(sparsity=0.0, ref_keep_rate=0.25).span_sparsities == (0.0, 0.0) + # Dense layers pass sparsity 0 to the mask builder, which overrides the spans. + assert _build_block_mask(scores, P, 0.0, True, meta.video_tile_spans, meta.span_sparsities).all() + + +def test_build_reference_video_region_sparsifies_reference_keys(): + """A reference-video region sparsifies that video's keys and queries; target-to-target selection is unchanged.""" + torch.manual_seed(4) + q = torch.randn(1, _R2V_SEQ, 2, 8) + k = torch.randn(1, _R2V_SEQ, 2, 8) + allows = [] + # The reference video's 120 rows as a dense segment, then as its own region. + for build in (dict(packed_segments=(37, 9, 120, 11, (8, 8, 12)), ref_keep_rate=None), {}): + meta = _build_r2v(sparsity=0.75, **build) + impl = _impl() + tq, tk = (impl.tile(t, meta).clone() for t in (q, k)) + scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes), + _pool_tiles(tk, meta.variable_block_sizes).transpose(-2, -1)) + allows.append(torch.stack([_token_allow(meta, _mask(meta, scores, 0.75), 0, h) for h in range(2)])) + dense_reference_allow, multi_region_allow = allows + assert dense_reference_allow[..., 46:166].all(), "a dense reference video is visible to every query" + assert not torch.equal(dense_reference_allow, multi_region_allow) + for start, end in _R2V_DENSE_RUNS: + assert multi_region_allow[:, start:end].all() and multi_region_allow[..., start:end].all(), (start, end) + assert not multi_region_allow[:, 200, 46:166].all(), "a reference-video region sparsifies its keys" + assert torch.equal(dense_reference_allow[:, 177:, 177:], multi_region_allow[:, 177:, 177:]) + + +@pytest.mark.parametrize("tile_size", [_TILE_ELEMS, 64, 128]) +def test_build_multi_region_sparsity_zero_matches_dense_sdpa(tile_size): + torch.manual_seed(5) + meta = _build_r2v(tile_size=tile_size) + q, k, v = (torch.randn(1, _R2V_SEQ, 2, 8) for _ in range(3)) + impl = _impl() + # An odd 128-token tile count carries one transport-only partner tile; the + # attention problem is the logical prefix. + logical = meta.variable_block_sizes.numel() * meta.tile_elems + tq, tk, tv = (impl.tile(t, meta)[:, :logical].clone() for t in (q, k, v)) + scores = torch.matmul(_pool_tiles(tq, meta.variable_block_sizes, meta.tile_elems), + _pool_tiles(tk, meta.variable_block_sizes, meta.tile_elems).transpose(-2, -1)) + mask = _mask(meta, scores, 0.0) + assert mask.all() + sparse_out = impl.postprocess_output(_sparse_attention_oracle(tq, tk, tv, mask, meta), meta) + dense_out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2) + assert torch.allclose(sparse_out, dense_out, atol=1e-5), (sparse_out - dense_out).abs().max() + + +def test_build_region_shape_not_multiple_of_patch(): + with pytest.raises(ValueError, match="not a positive multiple of patch"): + _build_r2v(packed_segments=(37, 9, (5, 7, 12), 11, (8, 8, 12))) + + +def test_build_no_video_region(): + with pytest.raises(ValueError, match="at least one video region"): + _build_r2v(packed_segments=(37, 9, 11)) + + +def test_build_compete_with_reference_video_regions(): + with pytest.raises(ValueError, match="compete' supports only the generated-video region"): + _build_r2v(sparsity=0.75, exempt=False) + + +def test_build_compete_with_generated_video_only(): + meta = _build_r2v(sparsity=0.75, exempt=False, packed_segments=(37, 20, (8, 8, 12))) + assert meta.exempt is False and len(meta.video_tile_spans) == 1 + + +def test_assert_multi_region_metadata_rejects_mismatched_metadata(): + single = MiniMaxH3VSAMetadataBuilder().build(current_timestep=0, + patch_size=_PATCH, + VSA_sparsity=0.9, + packed_segments=(37, 9, 120, 11, (8, 8, 12)), + device=_CPU, + tile_size=64) + with pytest.raises(ValueError, match="only sparse region is the generated video"): + _assert_multi_region_metadata(single, expected_reference_video_regions=1, target_sparsity=0.9, + ref_keep_rate=0.1) + with pytest.raises(ValueError, match="leave reference-video conditioning dense"): + _assert_multi_region_metadata(_build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.9, + ref_keep_rate=1.0) + with pytest.raises(ValueError, match="differ from the policy"): + _assert_multi_region_metadata(_build_r2v(sparsity=0.9, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.9, + ref_keep_rate=0.2) + _assert_multi_region_metadata(_build_r2v(sparsity=0.0, tile_size=64, ref_keep_rate=0.1), + expected_reference_video_regions=1, + target_sparsity=0.0, + ref_keep_rate=0.1) + + +def test_build_multi_region_missing_ref_keep_rate(): + with pytest.raises(ValueError, match="needs ref_keep_rate"): + _build_r2v(sparsity=0.75, ref_keep_rate=None) + + +# --------------------------------------------------------------------------- +# 128-token tiles: sm_100a/sm_103a CUDA only +# --------------------------------------------------------------------------- + +_HEADS, _DIM = 2, 128 + + +class _FakeSm100a: + """Stands in for fastvideo_kernel.block_sparse_attn_sm100a.""" + + def __init__(self, supported=True): + self.supported = supported + self.calls = [] + + def is_supported(self, q, variable_block_sizes): + return self.supported + + def block_sparse_attn_sm100a(self, q, k, v, q2k_idx, q2k_num, variable_block_sizes, need_lse=True): + self.calls.append(dict(q=q, q2k_num=q2k_num, vbs=variable_block_sizes, need_lse=need_lse)) + return q.clone(), None + + +def _fake_map_to_index(block_map): + num = block_map.sum(dim=-1, dtype=torch.int32) + idx = torch.full(block_map.shape, -1, dtype=torch.int32) + for position in torch.nonzero(block_map.flatten(0, -2).any(-1)).flatten().tolist(): + cols = torch.nonzero(block_map.flatten(0, -2)[position]).flatten() + idx.flatten(0, -2)[position, :cols.numel()] = cols.to(torch.int32) + return idx, num + + +@pytest.fixture() +def tile128(monkeypatch, env_overrides): + fake = _FakeSm100a() + + def forbid_triton(*args, **kwargs): + raise AssertionError("tile 128 must never fall back to the Triton-64 kernel") + + monkeypatch.setattr(vsa_h3, "_sm100a", fake) + monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index) + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", forbid_triton) + # Tile 128 has no opt-in: it routes to the CUDA kernel with the switch off. + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + return fake + + +def _run_tile128(meta, requires_grad=False): + impl = _impl(_DIM) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16, requires_grad=requires_grad) + tiled = impl.preprocess_qkv(raw, meta) + query, key, value = tiled.chunk(3, dim=0) + return tiled, impl.forward(query, key, value, None, meta) + + +def test_forward_tile128_odd_tile_count_adds_partner_tile(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128, ref_keep_rate=0.25) + n_tiles = meta.variable_block_sizes.numel() + assert n_tiles % 2 == 1 + tiled, out = _run_tile128(meta) + assert tiled.shape[1] == (n_tiles + 1) * 128 + assert torch.count_nonzero(tiled[:, n_tiles * 128:]) == 0 + assert len(tile128.calls) == 1 + call = tile128.calls[0] + assert call["need_lse"] is False + assert call["vbs"].dtype == torch.int32 and call["vbs"].numel() == n_tiles + 1 and call["vbs"][-1] == 0 + assert (call["q2k_num"][..., n_tiles] == 0).all() + assert out.shape == (1, n_tiles * 128, _HEADS, _DIM) + + +def test_forward_tile128_even_tile_count_no_partner_tile(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128, packed_segments=(37, 9, (5, 8, 12), 11, (8, 8, 12), 64)) + assert meta.variable_block_sizes.numel() % 2 == 0 + tiled, _ = _run_tile128(meta) + assert tiled.shape[1] == meta.variable_block_sizes.numel() * 128 + assert len(tile128.calls) == 1 + + +def test_forward_tile128_unavailable_cuda_kernel(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + tile128.supported = False + with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + _run_tile128(meta) + monkeypatch.setattr(vsa_h3, "_sm100a", None) + with pytest.raises(RuntimeError, match="not installed"): + _run_tile128(meta) + + +def test_forward_tile128_route_resolved_once_per_device(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128) + probes = [] + tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) or True + impl = _impl(_DIM) + for _ in range(3): + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) + impl.forward(query, key, value, None, meta) + assert len(tile128.calls) == 3 and probes == [_CPU] + # A failed probe is cached too, and every later call still fails closed. + tile128.is_supported = lambda q, variable_block_sizes: probes.append(q.device) and False + impl = _impl(_DIM) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16) + query, key, value = impl.preprocess_qkv(raw, meta).chunk(3, dim=0) + for _ in range(2): + with pytest.raises(RuntimeError, match="tile 128 requires the sm_100a/sm_103a CUDA block-sparse kernel"): + impl.forward(query, key, value, None, meta) + assert probes == [_CPU, _CPU] + + +def test_forward_tile128_requires_no_grad(tile128): + meta = _build_r2v(sparsity=0.5, tile_size=128) + with pytest.raises(NotImplementedError, match="no-grad inference"): + _run_tile128(meta, requires_grad=True) + assert tile128.calls == [] + + +def test_forward_tile128_rejects_regional_fullgraph_compile(tile128, monkeypatch): + meta = _build_r2v(sparsity=0.5, tile_size=128) + impl = _impl(_DIM) + impl._regional_compile_sm100a_enabled = True + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + q = torch.zeros(1, meta.variable_block_sizes.numel() * 128, _HEADS, _DIM, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="requires 64-token tiles"): + impl.forward(q, q, q, None, meta) + + +def test_build_tile128_uses_4x4x8_tiles(): + meta = _build_r2v(tile_size=128) + assert meta.tile_elems == 128 + assert int(meta.variable_block_sizes.max()) <= 128 + assert int(meta.variable_block_sizes.sum()) == _R2V_SEQ + # (5, 4, 6) and (8, 4, 6) token grids -> ceil(t/4) * ceil(4/4) * ceil(6/8) tiles each + assert meta.video_tile_spans == ((3, 5), (5, 7)) + assert math.prod((4, 4, 8)) == 128 + + +@pytest.mark.parametrize("sparsity,ref_keep_rate", [(0.5, 0.25), (0.9, 0.1)]) +def test_forward_tile128_real_sm100a_kernel_matches_oracle(env_overrides, sparsity, ref_keep_rate): + """The real 128-token CUDA forward on a multi-region layout against an FP32 masked-SDPA oracle.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)}: + pytest.skip("requires an sm_100a/sm_103a GPU (B200, B300, GB200, GB300)") + forwards = getattr(vsa_h3._sm100a, "_FWD_BY_BLOCK", {}) if vsa_h3._sm100a is not None else {} + if forwards.get(128) is None or vsa_h3.map_to_index is None: + pytest.skip("requires a fastvideo_kernel build with the 128-token sm_100a forward") + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(None)) + device = torch.device("cuda") + meta = _build_r2v(sparsity=sparsity, tile_size=128, ref_keep_rate=ref_keep_rate, device=device) + n_tiles = meta.variable_block_sizes.numel() + logical = n_tiles * 128 + impl = _impl(_DIM) + torch.manual_seed(19) + raw = torch.randn(3, meta.total_seq_length, _HEADS, _DIM, device=device, dtype=torch.bfloat16) + with torch.inference_mode(): + tiled = impl.preprocess_qkv(raw, meta) + assert tiled.shape[1] == (n_tiles + n_tiles % 2) * 128 + query, key, value = tiled.chunk(3, dim=0) + actual = impl.postprocess_output(impl.forward(query, key, value, None, meta), meta) + # The mask the backend selected: pooled scores of the logical tiles. + logical_query, logical_key, logical_value = (t[:, :logical] for t in (query, key, value)) + scores = torch.matmul(_pool_tiles(logical_query, meta.variable_block_sizes, 128), + _pool_tiles(logical_key, meta.variable_block_sizes, 128).transpose(-2, -1)) / _DIM**0.5 + mask = _mask(meta, scores, sparsity) + assert not mask.all(), "the oracle must exercise a sparse mask" + expected = impl.postprocess_output( + _sparse_attention_oracle(logical_query.float(), logical_key.float(), logical_value.float(), mask, meta), + meta) + torch.cuda.synchronize() + assert actual.shape == (1, meta.total_seq_length, _HEADS, _DIM) + torch.testing.assert_close(actual.float(), expected, atol=0.04, rtol=0.02) diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 9b26678c8..0eaa79130 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -35,10 +35,9 @@ def _build_meta( builder = builder or MiniMaxH3VSAMetadataBuilder() return builder.build( current_timestep=0, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=sparsity, - prefix_segments=prefix_segments, + packed_segments=(*prefix_segments, _SPEC["raw_latent_shape"]), device=device, dense_layers=dense_layers, tile_size=64, @@ -224,10 +223,9 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch, env_overrid # adjacent-query-block contract. meta = MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=0.0, - prefix_segments=(64, 64), + packed_segments=(64, 64, _SPEC["raw_latent_shape"]), device=torch.device("cpu"), tile_size=64, ) diff --git a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py index 30c9aa4b4..93ecf8dd3 100644 --- a/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py +++ b/fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py @@ -36,10 +36,9 @@ def _build(builder=None, step=0): # A fresh metadata per step, all sharing the builder's holder. return builder.build( current_timestep=step, - raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], VSA_sparsity=0.9, - prefix_segments=_SPEC["prefix_segments"], + packed_segments=(*_SPEC["prefix_segments"], _SPEC["raw_latent_shape"]), device=torch.device("cpu"), tile_size=_TILE_ELEMS, ), builder diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index e446b4389..82fbaff4b 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -1,3 +1,4 @@ +import dataclasses import json import shutil import subprocess @@ -41,6 +42,7 @@ def _new_runtime_video_generator() -> VideoGenerator: model_path="test-model", prompt_txt=None, workload_type=SimpleNamespace(value="t2v"), + pipeline_config=SimpleNamespace(apply_request_constraints=lambda request, sampling_param: sampling_param), ) generator.executor = SimpleNamespace( set_log_queue=lambda queue: None, @@ -114,7 +116,9 @@ def _single_video_args(output_type="pil"): output_type=output_type, pin_cpu_memory=False, VSA_sparsity=0.0, - pipeline_config=SimpleNamespace(flow_shift=1.0, embedded_cfg_scale=1.0), + pipeline_config=SimpleNamespace(flow_shift=1.0, + embedded_cfg_scale=1.0, + apply_request_constraints=lambda request, sampling_param: sampling_param), workload_type=SimpleNamespace(value="t2v"), ) @@ -704,6 +708,31 @@ def test_generate_uses_typed_request_path(monkeypatch): assert result.video_path == "outputs/test.mp4" +@pytest.mark.parametrize("entry", ["generate", "generate_video"]) +def test_generate_applies_pipeline_config_request_constraints(monkeypatch, entry): + """Both entry points run the sampling params that the loaded model's request constraints return.""" + generator = _new_runtime_video_generator() + _patch_sampling_param_from_pretrained(monkeypatch) + + def apply_request_constraints(request, sampling_param): + return dataclasses.replace(sampling_param, num_inference_steps=8) + + generator.fastvideo_args.pipeline_config = SimpleNamespace(apply_request_constraints=apply_request_constraints) + captured = {} + + def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs): + captured["sampling_param"] = sampling_param + return {"prompts": prompt, "video_path": "outputs/test.mp4"} + + monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl) + if entry == "generate": + generator.generate(GenerationRequest(prompt="hello world")) + else: + with pytest.warns(DeprecationWarning): + generator.generate_video(prompt="hello world") + assert captured["sampling_param"].num_inference_steps == 8 + + def test_generate_rejects_stage_override_outside_registered_stage(monkeypatch) -> None: """Validate typed stage overrides at the public generation entrypoint.""" generator = _new_runtime_video_generator() @@ -923,7 +952,8 @@ def test_generate_video_legacy_call_uses_legacy_impl(monkeypatch): def test_generate_video_legacy_call_routes_compat_kwargs(monkeypatch): generator = _new_runtime_video_generator() - generator.fastvideo_args.pipeline_config = SimpleNamespace(embedded_cfg_scale=1.0) + generator.fastvideo_args.pipeline_config = SimpleNamespace( + embedded_cfg_scale=1.0, apply_request_constraints=lambda request, sampling_param: sampling_param) captured = {} def fake_generate_video_impl( diff --git a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py new file mode 100644 index 000000000..26db7a4c8 --- /dev/null +++ b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of examples/inference/basic/basic_fasth3_omniref_pdd.py (no weights, no GPU).""" +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3_omniref_pdd.py" +BASE_PIN = "9bfb6693f2cf6de171db46d1aa586f67d773a1da" +CONTRACT = { + "schema_version": "fasth3-inference-contract-v1", + "base_model_revision": f"hf://MiniMaxAI/MiniMax-H3@{BASE_PIN}", + "pdd_steps": 32, + "pdd_step_indices": list(range(0, 33, 4)), + "num_inference_steps": 8, + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, + "vsa_ref_policy": "p2_multi_region", + "vsa_ref_keep_rate": 0.1, +} + + +def _load_example(): + spec = importlib.util.spec_from_file_location("basic_fasth3_omniref_pdd", EXAMPLE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +example = _load_example() + + +def _args(*overrides: str): + return example.parse_args(["--model-path", "export", "--image", "a.png", "--prompt", "p", *overrides]) + + +@pytest.mark.parametrize("overrides,expected", [ + ((), ("MiniMaxAI/MiniMax-H3", BASE_PIN)), + (("--base-revision", "abc"), ("MiniMaxAI/MiniMax-H3", "abc")), + # The export's pin names a revision of its own base repo; it never applies to another one. + (("--base-model-path", "someone/mirror"), ("someone/mirror", None)), + (("--base-model-path", "someone/mirror", "--base-revision", "abc"), ("someone/mirror", "abc")), +]) +def test_base_model_source_command_line_overrides(overrides, expected): + assert example.base_model_source(_args(*overrides), CONTRACT) == expected + + +@pytest.mark.parametrize("pin", [ + BASE_PIN, f"MiniMaxAI/MiniMax-H3@{BASE_PIN}", "hf://MiniMaxAI/MiniMax-H3", "hf://@abc", "", None, +]) +@pytest.mark.parametrize("overrides", [(), ("--base-model-path", "someone/mirror")]) +def test_base_model_source_unparsable_base_pin(pin, overrides): + """Never a silent fall back to the base repo's latest revision; the pipeline rejects the same values.""" + with pytest.raises(ValueError, match="must be hf://@"): + example.base_model_source(_args(*overrides), {**CONTRACT, "base_model_revision": pin}) + + +_COMPONENT = ["diffusers", "Component"] +_FULL_MANIFEST = { + "_class_name": "MiniMaxH3ModularPipeline", + **{name: _COMPONENT for name in ("transformer", "transformer_ref", "scheduler", "audio_scheduler", + "text_encoder", "tokenizer", "processor", "vae", "audio_vae")}, +} +_EXPORT_ONLY_MANIFEST = { + "_class_name": "MiniMaxH3ModularPipeline", + **{name: _COMPONENT for name in ("transformer_ref", "scheduler", "audio_scheduler")}, +} + + +def _snapshot_dir(root: Path, components, manifest): + root.mkdir() + for name in components: + (root / name).mkdir() + if manifest is not None: + (root / "modular_model_index.json").write_text(json.dumps(manifest)) + return root + + +@pytest.mark.parametrize("export_manifest,chosen", [ + (_FULL_MANIFEST, "export"), + (_EXPORT_ONLY_MANIFEST, "base"), + (None, "base"), +]) +def test_compose_model_dir_selects_complete_manifest(tmp_path, export_manifest, chosen): + export = _snapshot_dir(tmp_path / "export", example.EXPORT_COMPONENTS, export_manifest) + (export / example.CONTRACT).write_text(json.dumps(CONTRACT)) + base = _snapshot_dir(tmp_path / "base", example.BASE_COMPONENTS, _FULL_MANIFEST) + composed = example.compose_model_dir(export, base, tmp_path / "composed") + manifest = composed / "modular_model_index.json" + assert manifest.resolve() == ((export if chosen == "export" else base) / "modular_model_index.json").resolve() + for name in (*example.EXPORT_COMPONENTS, *example.BASE_COMPONENTS, example.CONTRACT): + assert (composed / name).exists() + + +def test_compose_model_dir_no_complete_manifest(tmp_path): + export = _snapshot_dir(tmp_path / "export", example.EXPORT_COMPONENTS, _EXPORT_ONLY_MANIFEST) + (export / example.CONTRACT).write_text(json.dumps(CONTRACT)) + base = _snapshot_dir(tmp_path / "base", example.BASE_COMPONENTS, _EXPORT_ONLY_MANIFEST) + with pytest.raises(FileNotFoundError, match="manifest that declares"): + example.compose_model_dir(export, base, tmp_path / "composed") + + +def test_build_generator_config_leaves_trained_settings_to_checkpoint(tmp_path): + # FastVideo reads the trained attention settings from the composed directory's fastvideo_inference.json. + assert example.build_generator_config(tmp_path, 1).pipeline.experimental == {} + + +@pytest.mark.parametrize("num_gpus,sharded", [(1, False), (4, True)]) +def test_build_generator_config_multi_gpu_shards_dit(tmp_path, num_gpus, sharded): + config = example.build_generator_config(tmp_path, num_gpus) + assert config.engine.use_fsdp_inference is sharded + assert config.engine.parallelism.sp_size == num_gpus + assert config.pipeline.components.override_pipeline_cls_name == "MiniMaxH3Ref2VAModularPipeline" diff --git a/fastvideo/tests/layers/test_pdd_linear.py b/fastvideo/tests/layers/test_pdd_linear.py new file mode 100644 index 000000000..1de8ce50a --- /dev/null +++ b/fastvideo/tests/layers/test_pdd_linear.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for the PDD sampling primitives in ``fastvideo.layers.pdd``. + +References are spelled out locally: the rational time shift on the +``[0, 0.999]`` base clock, its cancellation-free difference, the fused blocks +of an explicit fine-grid partition, and the fused-head forward as the +integration-weighted mean of the materialized heads. +""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.layers.pdd import ( + PDD_GRID_MAX_T, + PDDModalitySchedule, + PDDReplicatedLinear, + build_pdd_sampling_plan, + fuse_pdd_heads, + pdd_fine_grid, + shifted_noise_amount, + shifted_noise_delta, +) +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler + +# The OmniRef PDD-8 export: a 32-interval grid sampled in eight 4-head blocks. +GRID32_BLOCKS8 = list(range(0, 33, 4)) + + +def _time_shift(value: torch.Tensor, shift: float, max_t: float = PDD_GRID_MAX_T) -> torch.Tensor: + """``s * t * M / (t * (s - 1) + M)`` in float64.""" + value = value.to(torch.float64) + return value * shift * max_t / (value * (shift - 1.0) + max_t) + + +@pytest.mark.parametrize("shift", [1.0, 0.25, 3.0, 12.0]) +def test_shifted_noise_amount_matches_rational_time_shift(shift: float) -> None: + base = torch.linspace(0.0, PDD_GRID_MAX_T, 17, dtype=torch.float64) + torch.testing.assert_close(shifted_noise_amount(base, shift), _time_shift(base, shift), rtol=0.0, atol=1e-15) + # max_t is a fixed point of every shift, so the first node has the same + # noise level in every modality. + fixed_point = float(shifted_noise_amount(torch.tensor(PDD_GRID_MAX_T, dtype=torch.float64), shift)) + assert fixed_point == pytest.approx(PDD_GRID_MAX_T, abs=1e-15) + + +@pytest.mark.parametrize("shift", [1.0, 0.25, 12.0]) +def test_shifted_noise_delta_matches_direct_difference(shift: float) -> None: + grid = pdd_fine_grid(256) + direct = shifted_noise_amount(grid[1:], shift) - shifted_noise_amount(grid[:-1], shift) + delta = shifted_noise_delta(grid[:-1], grid[1:], shift) + torch.testing.assert_close(delta, direct, rtol=1e-12, atol=1e-15) + assert bool((delta < 0).all()) + if shift == 1.0: + assert torch.equal(delta, grid[1:] - grid[:-1]) + + +def test_pdd_fine_grid_descending_float64_from_max_t_to_zero() -> None: + grid = pdd_fine_grid(8) + assert grid.dtype == torch.float64 + assert grid.shape == (9, ) + assert float(grid[0]) == PDD_GRID_MAX_T + assert float(grid[-1]) == 0.0 + assert bool((grid[:-1] > grid[1:]).all()) + + +@pytest.mark.parametrize(("step_indices", "block_sizes"), [ + ([0, 64, 128, 192, 256], [64, 64, 64, 64]), + ([0, 2, 5, 8], [2, 3, 3]), +], ids=["uniform", "uneven"]) +def test_build_pdd_sampling_plan_blocks_and_node_sigmas(step_indices: list[int], block_sizes: list[int]) -> None: + schedules = {"video": PDDModalitySchedule(shift=12.0), "audio": PDDModalitySchedule(shift=3.0)} + plan = build_pdd_sampling_plan(step_indices, schedules) + + grid_size = step_indices[-1] + assert plan.pdd_steps == grid_size + assert plan.num_steps == len(block_sizes) + assert plan.block_sizes == block_sizes + assert plan.block(1) == (step_indices[1], step_indices[2]) + # Node i of an N-interval grid sits at base time max_t * (1 - i / N). + node_base = PDD_GRID_MAX_T * (1.0 - torch.tensor(step_indices, dtype=torch.float64) / grid_size) + for name, schedule in schedules.items(): + node_sigmas = plan.node_sigmas[name] + torch.testing.assert_close(node_sigmas, _time_shift(node_base, schedule.shift), rtol=0.0, atol=1e-12) + assert float(node_sigmas[0]) == pytest.approx(PDD_GRID_MAX_T, abs=1e-15) + assert float(node_sigmas[-1]) == 0.0 + # A fused block's total weight is exactly its node-sigma increment, + # which is what makes one Euler step over the two nodes the block advance. + for step in range(plan.num_steps): + start, end = plan.block(step) + total = plan.integration_weights[name][start:end].sum() + torch.testing.assert_close(total, node_sigmas[step + 1] - node_sigmas[step], rtol=1e-12, atol=1e-15) + + +def test_build_pdd_sampling_plan_mismatched_max_t() -> None: + schedules = {"a": PDDModalitySchedule(1.0), "b": PDDModalitySchedule(1.0, max_t=1.0)} + with pytest.raises(ValueError, match="share max_t"): + build_pdd_sampling_plan([0, 4, 8], schedules) + + +def test_build_pdd_sampling_plan_grid32_eight_blocks_feed_h3_schedulers() -> None: + """The exported contract (Grid32, blocks of 4, shifts 12/3) as the stage builds it.""" + plan = build_pdd_sampling_plan(GRID32_BLOCKS8, { + "video": PDDModalitySchedule(shift=12.0), + "audio": PDDModalitySchedule(shift=3.0) + }) + assert plan.block_sizes == [4] * 8 + for name, shift in (("video", 12.0), ("audio", 3.0)): + expected = _time_shift(pdd_fine_grid(32)[GRID32_BLOCKS8], shift) + torch.testing.assert_close(plan.node_sigmas[name], expected, rtol=0.0, atol=1e-15) + scheduler = MiniMaxH3Scheduler(shift=shift) + scheduler.set_timesteps(sigmas=plan.node_sigmas[name].to(torch.float32)) + assert scheduler.num_inference_steps == 8 + torch.testing.assert_close(scheduler.timesteps, 1.0 - plan.node_sigmas[name][:-1].to(torch.float32)) + + +def _randomized(linear: PDDReplicatedLinear) -> PDDReplicatedLinear: + # ReplicatedLinear storage is uninitialized until a checkpoint loads. + with torch.no_grad(): + linear.weight.normal_() + linear.bias.normal_() + return linear + + +def test_pdd_replicated_linear_head_major_and_state_dict_compatible() -> None: + linear = PDDReplicatedLinear(5, 3, grid_size=4, params_dtype=torch.float32) + assert linear.weight.shape == (12, 5) and linear.bias.shape == (12, ) + assert linear.head_output_size == 3 and linear.grid_size == 4 + # Fusion state is transient: the state dict is a plain linear's. + assert set(linear.state_dict()) == {"weight", "bias"} + + +def test_fused_params_use_normalized_integration_weights() -> None: + torch.manual_seed(2) + grid, channels, in_features = 5, 3, 7 + linear = _randomized(PDDReplicatedLinear(in_features, channels, grid_size=grid, params_dtype=torch.float32)) + weights = PDDModalitySchedule(shift=0.25).integration_weights(pdd_fine_grid(grid)) + start, end = 1, 5 + + actual_weight, actual_bias = linear._fused_params(start, end, weights, torch.float32) + alpha = weights[start:end] + alpha = (alpha / alpha.sum()).float() + expected_weight = torch.einsum("n,nci->ci", alpha, linear.weight.reshape(grid, channels, in_features)[start:end]) + expected_bias = torch.einsum("n,nc->c", alpha, linear.bias.reshape(grid, channels)[start:end]) + torch.testing.assert_close(actual_weight, expected_weight, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(actual_bias, expected_bias, rtol=1e-6, atol=1e-6) + + +def test_forward_fused_block_weighted_mean_of_materialized_heads() -> None: + torch.manual_seed(3) + grid, out_features, in_features = 6, 4, 5 + linear = _randomized(PDDReplicatedLinear(in_features, out_features, grid_size=grid, params_dtype=torch.float32)) + x = torch.randn(2, 3, in_features) + weights = PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(grid)) + + heads, extra_bias = linear(x) + assert extra_bias is None + heads = heads.unflatten(-1, (grid, out_features)) + for start, end in ((0, grid), (2, 5), (4, 5)): + alpha = (weights[start:end] / weights[start:end].sum()).float() + expected = torch.einsum("n,btnc->btc", alpha, heads[..., start:end, :]) + with linear.fuse(start, end, weights, torch.float32): + fused, fused_extra = linear(x) + assert fused_extra is None + torch.testing.assert_close(fused, expected, rtol=1e-5, atol=1e-5) + # Leaving the context restores the materialized-head forward. + torch.testing.assert_close(linear(x)[0].unflatten(-1, (grid, out_features)), heads) + + +def test_fused_params_bf16_heads_accumulate_in_decoding_precision() -> None: + torch.manual_seed(5) + linear = _randomized(PDDReplicatedLinear(8, 3, grid_size=4, params_dtype=torch.bfloat16)) + weights = PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(4)) + fused_weight, fused_bias = linear._fused_params(0, 4, weights, torch.float32) + alpha = (weights / weights.sum()).float() + expected = torch.einsum("n,nci->ci", alpha, linear.weight.float().reshape(4, 3, 8)).to(torch.bfloat16) + assert fused_weight.dtype == torch.bfloat16 and fused_bias is not None and fused_bias.dtype == torch.bfloat16 + assert torch.equal(fused_weight, expected) + + +def test_fuse_pdd_heads_fuses_every_modality_head_at_once() -> None: + torch.manual_seed(4) + grid = 4 + linears = { + "video": _randomized(PDDReplicatedLinear(3, 2, grid_size=grid, params_dtype=torch.float32)), + "audio": _randomized(PDDReplicatedLinear(3, 1, grid_size=grid, params_dtype=torch.float32)), + } + weights = { + "video": PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(grid)), + "audio": PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(grid)), + } + x = torch.randn(2, 3) + with fuse_pdd_heads(linears, 1, 3, weights, torch.float32): + fused = {name: linear(x)[0] for name, linear in linears.items()} + for name, linear in linears.items(): + with linear.fuse(1, 3, weights[name], torch.float32): + torch.testing.assert_close(fused[name], linear(x)[0]) + assert linear._fusion_state is None + + +def test_fuse_nested_contexts_restore_outer_state() -> None: + linear = _randomized(PDDReplicatedLinear(3, 2, grid_size=4, params_dtype=torch.float32)) + weights = PDDModalitySchedule(shift=1.0).integration_weights(pdd_fine_grid(4)) + with linear.fuse(0, 4, weights, torch.float32): + outer = linear._fusion_state + with linear.fuse(1, 3, weights, torch.float32): + assert linear._fusion_state[:2] == (1, 3) + assert linear._fusion_state is outer + assert linear._fusion_state is None diff --git a/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py new file mode 100644 index 000000000..808a0a2ed --- /dev/null +++ b/fastvideo/tests/loader/test_minimax_h3_pdd_loading.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Loading a PDD-widened MiniMax-H3 ``transformer_ref`` export through the transformer loader. + +A PDD export records ``pdd_steps`` in its transformer ``config.json`` and +stores both output projections widened to that many heads, next to its trained +VSA compression gates. The loader must build the widened heads from the +config, load every tensor exactly, and fail closed when the config and the +weights disagree. Runs on CPU with a tiny synthetic checkpoint. +""" +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +import fastvideo.envs as envs +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.layers.pdd import PDDModalitySchedule, PDDReplicatedLinear, pdd_fine_grid +from fastvideo.models.dits import minimax_h3 +from fastvideo.models.loader import component_loader +from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_TEXT_TAG, build_packed_sequence, build_row_timesteps +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.platforms import AttentionBackendEnum + +PDD_STEPS = 4 +TINY_ARCH = dict( + num_attention_heads=4, + attention_head_dim=16, + hidden_size=32, + num_layers=2, + num_refiner_layers=1, + ffn_dim=48, + in_channels=4, + audio_in_channels=4, + patch_size=[1, 2, 2], + text_dim=16, + freq_dim=16, + time_embed_hidden_dim=32, + time_embed_dim=24, + rope_freq_dim=2, + rope_theta=10000.0, + norm_eps=1e-5, + qk_norm_eps=1e-5, + final_norm_eps=1e-5, +) +VIDEO_PATCH_DIM = 4 * 1 * 2 * 2 + + +def _released_name(name: str) -> str: + """FastVideo parameter name -> released (diffusers) checkpoint name.""" + for fastvideo_name, released_name in (("time_embedder.fc_in.", "time_embedder.linear_1."), + ("time_embedder.fc_out.", "time_embedder.linear_2."), + (".attn.to_out.", ".attn.to_out.0."), (".ff.fc_in.", ".ff.net.0.proj."), + (".ff.fc_out.", ".ff.net.2.")): + name = name.replace(fastvideo_name, released_name) + return name + + +class _VSAH3Selection: + """What the selector returns when VIDEO_SPARSE_ATTN_H3 resolves on a CUDA host.""" + + @staticmethod + def get_name() -> str: + return "VIDEO_SPARSE_ATTN_H3" + + +@pytest.fixture() +def cpu_loader(monkeypatch, env_overrides): + """Load onto the CPU with SDPA attention, on CPU-only and CUDA hosts alike.""" + monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu")) + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + + +@pytest.fixture() +def vsa_gates(monkeypatch, cpu_loader): + """Construct H3 attention as the VSA-H3 backend does, so the compression gates exist on a CPU host.""" + select = minimax_h3.get_attn_backend + + def select_vsa_h3(head_size, dtype, supported_attention_backends=None, **kwargs): + if supported_attention_backends and AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3 in supported_attention_backends: + return _VSAH3Selection + return select(head_size, dtype, supported_attention_backends=supported_attention_backends, **kwargs) + + monkeypatch.setattr(minimax_h3, "get_attn_backend", select_vsa_h3) + + +@pytest.fixture() +def single_process_group(env_overrides): + """A 1x1 sequence/tensor-parallel group, which DistributedAttention needs even on one CPU process.""" + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + + # Keep a launcher- or CI-assigned rendezvous: packed CI lanes share the tray's network namespace, and each + # lease gets its own port range. + env_overrides.enter_context(envs.override_external("MASTER_ADDR", os.environ.get("MASTER_ADDR") or "127.0.0.1")) + env_overrides.enter_context(envs.override_external("MASTER_PORT", os.environ.get("MASTER_PORT") or "29591")) + env_overrides.enter_context(envs.override_external("RANK", "0")) + env_overrides.enter_context(envs.override_external("WORLD_SIZE", "1")) + env_overrides.enter_context(envs.override_external("LOCAL_RANK", "0")) + maybe_init_distributed_environment_and_model_parallel(1, 1) + try: + yield + finally: + cleanup_dist_env_and_memory() + + +def _write_transformer_ref(root, *, config_pdd_steps, weight_pdd_steps=PDD_STEPS, with_gates=True): + """A synthetic BF16 export: every tensor the (widened) model defines, in the released naming.""" + arch = MiniMaxH3ArchConfig(**{**TINY_ARCH, "patch_size": tuple(TINY_ARCH["patch_size"])}, + pdd_steps=weight_pdd_steps) + with torch.device("meta"): + reference = minimax_h3.MiniMaxH3Transformer3DModel(MiniMaxH3Config(arch_config=arch), hf_config={}) + generator = torch.Generator().manual_seed(0) + tensors = {} + for name, parameter in reference.named_parameters(): + if with_gates or "to_gate_compress" not in name: + tensors[_released_name(name)] = (torch.randn(parameter.shape, generator=generator) * 0.1).to(torch.bfloat16) + path = root / "transformer_ref" + path.mkdir() + save_file(tensors, str(path / "diffusion_pytorch_model.safetensors")) + config = {"_class_name": "MiniMaxH3Transformer3DModel", "_diffusers_version": "0.36.0.dev0", **TINY_ARCH} + if config_pdd_steps is not None: + config["pdd_steps"] = config_pdd_steps + (path / "config.json").write_text(json.dumps(config)) + return path, tensors + + +def _load(path): + args = SimpleNamespace( + pipeline_config=MiniMaxH3PipelineConfig(), + override_transformer_cls_name=None, + model_paths={}, + init_weights_from_safetensors=None, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + dit_cpu_offload=True, + pin_cpu_memory=False, + use_fsdp_inference=False, + training_mode=False, + enable_torch_compile=False, + torch_compile_kwargs=None, + inference_torch_compile=False, + VSA_tile_size=128, + lora_path=None, + lora_strength=1.0, + inference_mode=True, + dit_layerwise_offload=False, + ) + return component_loader.TransformerLoader().load(str(path), args) + + +def test_load_widened_heads_and_trained_gates_exactly(tmp_path, vsa_gates): + path, tensors = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS) + model = _load(path) + + assert model.pdd_steps == PDD_STEPS + for head, rows in (("proj_out", VIDEO_PATCH_DIM), ("audio_proj_out", TINY_ARCH["audio_in_channels"])): + linear = getattr(model, head) + assert isinstance(linear, PDDReplicatedLinear) + assert linear.grid_size == PDD_STEPS and linear.head_output_size == rows + assert linear.weight.shape == (PDD_STEPS * rows, TINY_ARCH["hidden_size"]) + # Output projections stay FP32 at inference, like the released heads. + assert linear.weight.dtype == torch.float32 + gate_names = [name for name, _ in model.named_parameters() if "to_gate_compress" in name] + assert gate_names == [f"transformer_blocks.{index}.attn.to_gate_compress.weight" for index in range(2)] + + loaded = dict(model.named_parameters()) + assert {_released_name(name) for name in loaded} == set(tensors) + for name, parameter in loaded.items(): + assert torch.equal(parameter, tensors[_released_name(name)].to(parameter.dtype)), name + with torch.no_grad(): + assert all(block.attn._gate_active() for block in model.transformer_blocks) + + +def test_load_export_without_gates_disables_gate_branch(tmp_path, vsa_gates): + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS, with_gates=False) + model = _load(path) + with torch.no_grad(): + assert not any(block.attn._gate_active() for block in model.transformer_blocks) + + +@pytest.mark.parametrize("config_pdd_steps,weight_pdd_steps", [(None, PDD_STEPS), (PDD_STEPS, None), (8, PDD_STEPS)]) +def test_load_mismatched_config_and_weight_grids(tmp_path, vsa_gates, config_pdd_steps, weight_pdd_steps): + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=config_pdd_steps, weight_pdd_steps=weight_pdd_steps) + # ReplicatedLinear.weight_loader rejects the shape; never a silent partial load. + with pytest.raises((AssertionError, RuntimeError), match="size"): + _load(path) + + +def test_fuse_pdd_block_weighted_mean_of_widened_heads(tmp_path, cpu_loader, single_process_group): + """End to end through the DiT: a fused block equals the weighted mean of its heads' outputs.""" + path, _ = _write_transformer_ref(tmp_path, config_pdd_steps=PDD_STEPS, with_gates=False) + model = _load(path) + layout = build_packed_sequence(torch.full((3, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long), 2, 4, 4, 2, (1, 2, 2)) + unique, inverse = build_row_timesteps(layout, 0.3, 0.4, 0.999, 1.0) + generator = torch.Generator().manual_seed(1) + inputs = dict( + hidden_states=torch.randn(1, int(layout.video_indices.numel()), VIDEO_PATCH_DIM, generator=generator), + audio_hidden_states=torch.randn(1, + int(layout.audio_indices.numel()), + TINY_ARCH["audio_in_channels"], + generator=generator), + encoder_hidden_states=torch.randn(1, 3, TINY_ARCH["text_dim"], generator=generator), + timestep=unique, + timestep_indices=inverse, + token_tags=layout.token_tags, + position_ids=layout.position_ids, + video_indices=layout.video_indices, + audio_indices=layout.audio_indices, + text_indices=layout.text_indices, + ) + weights = { + "video": PDDModalitySchedule(shift=12.0).integration_weights(pdd_fine_grid(PDD_STEPS)), + "audio": PDDModalitySchedule(shift=3.0).integration_weights(pdd_fine_grid(PDD_STEPS)), + } + with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=None, + forward_batch=ForwardBatch(data_type="dummy")): + heads = model(**inputs) + with model.fuse_pdd_block(1, 3, weights, torch.float32): + fused = model(**inputs) + for name, all_heads, block in zip(("video", "audio"), heads, fused, strict=True): + all_heads = all_heads.unflatten(-1, (PDD_STEPS, -1)) + alpha = weights[name][1:3] / weights[name][1:3].sum() + expected = torch.einsum("n,brnc->brc", alpha.float(), all_heads[..., 1:3, :].float()) + assert block.shape == expected.shape + torch.testing.assert_close(block.float(), expected, rtol=1e-4, atol=1e-5) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py index 439da0c0e..611e04fc6 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py @@ -62,10 +62,9 @@ def _dit_shape(spec: dict) -> tuple[int, int, int]: def _torch_meta(spec: dict, tile_size: int = 256, sparsity: float = 0.0): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, - raw_latent_shape=spec["raw_latent_shape"], patch_size=spec["patch_size"], VSA_sparsity=sparsity, - prefix_segments=spec["prefix_segments"], + packed_segments=(*spec["prefix_segments"], spec["raw_latent_shape"]), device=_CPU, tile_size=tile_size, ) @@ -112,7 +111,7 @@ def test_routing_mask_parity_exempt_and_compete() -> None: k_vid = compute_topk(0.9, v) for exempt in (True, False): mlx_mask = build_block_mask(scores, p, v, 0.9, exempt=exempt) - torch_mask = _build_block_mask(torch.tensor(scores), p, v, 0.9, exempt=exempt).numpy() + torch_mask = _build_block_mask(torch.tensor(scores), p, 0.9, exempt, ((p, p + v), ), (0.9, )).numpy() assert mlx_mask[:, :, :p].all() if exempt: assert mlx_mask[..., :p].all() diff --git a/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py b/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py new file mode 100644 index 000000000..22f892919 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_pdd_contract.py @@ -0,0 +1,472 @@ +# SPDX-License-Identifier: Apache-2.0 +"""How a FastH3 Ref2VA PDD checkpoint's ``fastvideo_inference.json`` sets a run's settings. + +``MiniMaxH3PipelineConfig.resolve_checkpoint_settings`` runs inside +``FastVideoArgs.__post_init__``. It fills each unset run setting from the file, +rejects an explicit value that differs for a setting fixed by training (attention +backend, VSA tile size, PDD partition), keeps an explicit value that differs for +a tunable setting (VSA sparsity, reference keep rate) with a warning, and +rejects a file that is incomplete or disagrees with the checkpoint file that owns +a value. ``MiniMaxH3PipelineConfig.apply_request_constraints`` then fills an +unset request step count with the block count and rejects a different one. In +the worker, ``MiniMaxH3BasePipeline`` checks its transformer before loading +weights: against those settings, and, when the transformer carries VSA +compression gates, against the run's attention backend. A DMD export's schedule +comes from its ``dmd_denoising_steps`` (see test_minimax_h3_distilled_schedule.py). +""" +from __future__ import annotations + +import json +import logging +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo import envs +from fastvideo.api.compat import normalize_generation_request +from fastvideo.api.parser import parse_config +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.api.schema import GenerationRequest +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import ( + MiniMaxH3ModularPipeline, + MiniMaxH3Ref2VAModularPipeline, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase + +GRID32_BLOCKS8 = [0, 4, 8, 12, 16, 20, 24, 28, 32] +# The OmniRef PDD-8 export's fastvideo_inference.json, verbatim. +PDD_FILE = { + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "audio_scheduler_shift": 3.0, + "base_model_revision": "hf://MiniMaxAI/MiniMax-H3@9bfb6693f2cf6de171db46d1aa586f67d773a1da", + "conditioning": "fixed_ordered_references_target_only_flow", + "grid_max_t": 0.999, + "guidance_scale": 1.0, + "model_type": "ref2va", + "num_inference_steps": 8, + "pdd_step_indices": GRID32_BLOCKS8, + "pdd_steps": 32, + "schema": "fasth3-inference-contract-v1", + "schema_version": "fasth3-inference-contract-v1", + "transformer_component": "transformer_ref", + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "vsa_ref_keep_rate": 0.1, + "vsa_ref_policy": "p2_multi_region", + "vsa_sparsity": 0.9, + "vsa_tile_size": 128, +} +# A DMD export's fastvideo_inference.json, which has no PDD fields. +DMD_FILE = { + "schema_version": "fasth3-inference-contract-v1", + "dmd_denoising_steps": [999, 874, 749, 624, 500, 375, 250, 125], + "num_inference_steps": 9, + "transformer_forwards": 8, + "video_scheduler_shift": 12.0, + "audio_scheduler_shift": 3.0, +} + + +@pytest.fixture(autouse=True) +def _unset_backend_env_var(env_overrides): + """Start every test without FASTVIDEO_ATTENTION_BACKEND, which FastVideoArgs reads as an explicit backend.""" + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(None)) + + +def _checkpoint(root: Path, inference_file: dict | None, *, pdd_steps: int | None = 32) -> str: + """Write a Ref2VA checkpoint's JSON files: the inference file and the configs that own its repeated values. + + The transformer_ref config.json sets ``pdd_steps`` unless it is None; the + scheduler configs hold the video/audio shifts 12.0/3.0. No weights are written. + """ + files = { + "transformer_ref/config.json": {} if pdd_steps is None else {"pdd_steps": pdd_steps}, + "scheduler/scheduler_config.json": {"shift": 12.0}, + "audio_scheduler/scheduler_config.json": {"shift": 3.0}, + } + if inference_file is not None: + files["fastvideo_inference.json"] = inference_file + for relative_path, content in files.items(): + path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(content)) + return str(root) + + +def _run_args(model_path: str, config_settings: dict | None = None, **run_settings) -> FastVideoArgs: + """Build a run's FastVideoArgs; construction resolves the checkpoint's settings.""" + pipeline_config = MiniMaxH3PipelineConfig(**(config_settings or {})) + return FastVideoArgs(model_path=model_path, pipeline_config=pipeline_config, **run_settings) + + +def _pipeline(model_path: str, cls=MiniMaxH3Ref2VAModularPipeline): + """A pipeline without loaded weights: the checkpoint path, its schedulers, and the Ref2VA flag __init__ sets.""" + pipeline = object.__new__(cls) + pipeline.model_path = model_path + pipeline._ref2va = cls._ref2va_default + pipeline.modules = { + "scheduler": MiniMaxH3Scheduler(shift=12.0), + "audio_scheduler": MiniMaxH3Scheduler(shift=3.0), + } + return pipeline + + +# ------------------------------------------------- resolve_checkpoint_settings: run settings on a PDD checkpoint + + +def test_resolve_checkpoint_settings_unset_run_settings(tmp_path): + args = _run_args(_checkpoint(tmp_path, PDD_FILE)) + assert args.attention_backend == "VIDEO_SPARSE_ATTN_H3" + assert args.VSA_tile_size == 128 + assert args.VSA_sparsity == 0.9 + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert args.pipeline_config.vsa_ref_keep_rate == 0.1 + + +def test_resolve_checkpoint_settings_matching_explicit_settings(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), { + "pdd_step_indices": GRID32_BLOCKS8, + "vsa_ref_keep_rate": 0.1 + }, + attention_backend="VIDEO_SPARSE_ATTN_H3", + VSA_tile_size=128, + VSA_sparsity=0.9) + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert "trained with" not in caplog.text + + +def test_resolve_checkpoint_settings_repeated_call(tmp_path): + args = _run_args(_checkpoint(tmp_path, PDD_FILE)) + args.pipeline_config.resolve_checkpoint_settings(args) + assert (args.attention_backend, args.VSA_tile_size, args.VSA_sparsity) == ("VIDEO_SPARSE_ATTN_H3", 128, 0.9) + assert args.pipeline_config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert args.pipeline_config.vsa_ref_keep_rate == 0.1 + + +@pytest.mark.parametrize("config_settings,run_settings,match", [ + ({}, {"attention_backend": "FLASH_ATTN"}, + "trained with attention_backend=VIDEO_SPARSE_ATTN_H3; this run requests FLASH_ATTN"), + ({}, {"VSA_tile_size": 64}, "trained with VSA_tile_size=128; this run requests 64"), + ({"pdd_step_indices": (0, 8, 16, 24, 32)}, {}, re.escape("this run sets [0, 8, 16, 24, 32]")), +], ids=["backend", "tile_size", "partition"]) +def test_resolve_checkpoint_settings_explicit_training_fixed_conflict(tmp_path, config_settings, run_settings, + match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, PDD_FILE), config_settings, **run_settings) + + +def test_resolve_checkpoint_settings_env_backend_conflict(tmp_path, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("FLASH_ATTN")) + with pytest.raises(ValueError, match="this run requests FLASH_ATTN"): + _run_args(_checkpoint(tmp_path, PDD_FILE)) + + +def test_resolve_checkpoint_settings_explicit_sparsity_override(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), VSA_sparsity=0.8) + assert args.VSA_sparsity == 0.8 + assert "trained with VSA_sparsity=0.9; this run uses 0.8" in caplog.text + + +def test_resolve_checkpoint_settings_explicit_keep_rate_override(tmp_path, caplog): + with caplog.at_level(logging.WARNING): + args = _run_args(_checkpoint(tmp_path, PDD_FILE), {"vsa_ref_keep_rate": 0.2}) + assert args.pipeline_config.vsa_ref_keep_rate == 0.2 + assert "trained with vsa_ref_keep_rate=0.1; this run uses 0.2" in caplog.text + + +@pytest.mark.parametrize("config_settings,run_settings,match", [ + ({}, {"VSA_sparsity": 1.0}, re.escape("VSA_sparsity must be in [0, 1), got 1.0")), + ({"vsa_ref_keep_rate": 1.0}, {}, re.escape("vsa_ref_keep_rate must be in (0, 1), got 1.0")), +], ids=["sparsity", "keep_rate"]) +def test_resolve_checkpoint_settings_tunable_out_of_range(tmp_path, config_settings, run_settings, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, PDD_FILE), config_settings, **run_settings) + + +def test_resolve_checkpoint_settings_dmd_steps_on_pdd_checkpoint(tmp_path): + with pytest.raises(ValueError, match="dmd_denoising_steps must be unset"): + _run_args(_checkpoint(tmp_path, PDD_FILE), {"dmd_denoising_steps": [999, 749, 500, 250]}) + + +# ------------------------------------------------- resolve_checkpoint_settings: PDD file validation + + +def test_resolve_checkpoint_settings_missing_field(tmp_path): + inference_file = {key: value for key, value in PDD_FILE.items() if key != "vsa_sparsity"} + with pytest.raises(ValueError, match=re.escape("missing ['vsa_sparsity']")): + _run_args(_checkpoint(tmp_path, inference_file)) + + +@pytest.mark.parametrize("pin", ["", None, "9bfb6693", "MiniMaxAI/MiniMax-H3@9bfb6693", "hf://MiniMaxAI/MiniMax-H3"]) +def test_resolve_checkpoint_settings_unparsable_base_model_revision(tmp_path, pin): + with pytest.raises(ValueError, match="must be hf://@"): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "base_model_revision": pin})) + + +def test_resolve_checkpoint_settings_unknown_field(tmp_path): + with pytest.raises(ValueError, match=re.escape("unknown ['dmd_denoising_steps']")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "dmd_denoising_steps": [999, 500]})) + + +@pytest.mark.parametrize("field,value", [ + ("schema_version", "fasth3-inference-contract-v2"), + ("schema", "fasth3-inference-contract-v2"), + ("model_type", "t2va"), + ("transformer_component", "transformer"), + ("conditioning", "noised_references"), + ("guidance_scale", 2.0), + ("attention_backend", "FLASH_ATTN"), + ("vsa_ref_policy", "dense"), + ("grid_max_t", 1.0), +]) +def test_resolve_checkpoint_settings_fixed_field_other_value(tmp_path, field, value): + with pytest.raises(ValueError, match=re.escape(f"{field}={value!r} is unsupported")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +@pytest.mark.parametrize("field,value,match", [ + ("pdd_steps", 16, "pdd_steps=16 disagrees with transformer_ref/config.json pdd_steps=32"), + ("video_scheduler_shift", 10.0, + "video_scheduler_shift=10.0 disagrees with scheduler/scheduler_config.json shift=12.0"), + ("audio_scheduler_shift", 4.0, + "audio_scheduler_shift=4.0 disagrees with audio_scheduler/scheduler_config.json shift=3.0"), + ("num_inference_steps", 9, "num_inference_steps=9 disagrees with the 8 blocks of pdd_step_indices"), + ("transformer_forwards", 4, "transformer_forwards=4 disagrees with the 8 blocks of pdd_step_indices"), +]) +def test_resolve_checkpoint_settings_duplicate_disagrees_with_owner(tmp_path, field, value, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +@pytest.mark.parametrize("owner", ["transformer_ref/config.json", "audio_scheduler/scheduler_config.json"]) +def test_resolve_checkpoint_settings_missing_owner_file(tmp_path, owner): + model_path = _checkpoint(tmp_path, PDD_FILE) + (tmp_path / owner).unlink() + with pytest.raises(ValueError, match=f"has fastvideo_inference.json but no {owner}"): + _run_args(model_path) + + +@pytest.mark.parametrize("field,value,match", [ + ("vsa_tile_size", 96, "vsa_tile_size=96 must be one of"), + ("vsa_sparsity", "0.9", "vsa_sparsity must be in \\[0, 1\\), got '0.9'"), + ("vsa_ref_keep_rate", 1.0, "vsa_ref_keep_rate must be in \\(0, 1\\), got 1.0"), +]) +def test_resolve_checkpoint_settings_invalid_trained_vsa_value(tmp_path, field, value, match): + with pytest.raises(ValueError, match=match): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, field: value})) + + +def test_resolve_checkpoint_settings_transformer_config_without_pdd_steps(tmp_path): + with pytest.raises(ValueError, match="pdd_steps=32 disagrees with transformer_ref/config.json pdd_steps=None"): + _run_args(_checkpoint(tmp_path, PDD_FILE, pdd_steps=None)) + + +@pytest.mark.parametrize("indices", [ + [0, 8, 4, 12, 16, 20, 24, 28, 32], + [1, 4, 8, 12, 16, 20, 24, 28, 32], + [0, 4, 8, 12, 16, 20, 24, 28, 31], + [0, 4.0, 8, 12, 16, 20, 24, 28, 32], + [32], +], ids=["decreasing", "starts_after_0", "ends_before_pdd_steps", "float_node", "single_node"]) +def test_resolve_checkpoint_settings_malformed_partition(tmp_path, indices): + with pytest.raises(ValueError, match=re.escape(f"pdd_step_indices={indices!r} must increase strictly from 0 to " + "pdd_steps=32")): + _run_args(_checkpoint(tmp_path, {**PDD_FILE, "pdd_step_indices": indices})) + + +# ------------------------------------------------- resolve_checkpoint_settings: checkpoints without PDD fields + + +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +def test_resolve_checkpoint_settings_non_pdd_checkpoint(tmp_path, inference_file): + args = _run_args(_checkpoint(tmp_path, inference_file, pdd_steps=None)) + assert (args.attention_backend, args.VSA_tile_size, args.VSA_sparsity) == (None, 256, 0.0) + config = args.pipeline_config + assert (config.pdd_step_indices, config.vsa_ref_keep_rate, config.dmd_denoising_steps) == (None, None, None) + + +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +@pytest.mark.parametrize("config_settings", [{"pdd_step_indices": (0, 16, 32)}, {"vsa_ref_keep_rate": 0.1}], + ids=["partition", "keep_rate"]) +def test_resolve_checkpoint_settings_pdd_setting_on_non_pdd_checkpoint(tmp_path, inference_file, config_settings): + with pytest.raises(ValueError, match="applies only to FastH3 PDD checkpoints"): + _run_args(_checkpoint(tmp_path, inference_file, pdd_steps=None), config_settings) + + +# ------------------------------------------------- apply_request_constraints: the step count of a request + + +def _pdd_config(): + return MiniMaxH3PipelineConfig(pdd_step_indices=tuple(GRID32_BLOCKS8)) + + +@pytest.mark.parametrize("sampling", [{}, {"num_inference_steps": 8}]) +def test_apply_request_constraints_pdd_unset_or_matching_steps(sampling): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": sampling}) + sampling_param = _pdd_config().apply_request_constraints(request, SamplingParam(num_inference_steps=50)) + assert sampling_param.num_inference_steps == 8 + + +def test_apply_request_constraints_pdd_conflicting_steps(): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": {"num_inference_steps": 50}}) + with pytest.raises(ValueError, match="runs exactly 8 transformer forwards.*num_inference_steps=50"): + _pdd_config().apply_request_constraints(request, SamplingParam()) + + +def test_apply_request_constraints_pdd_python_request_counts_schema_default_as_set(): + """A GenerationRequest built in Python carries the schema default 50 as a set value.""" + request = normalize_generation_request(GenerationRequest(prompt="fox")) + with pytest.raises(ValueError, match="num_inference_steps=50. Pass num_inference_steps=8"): + _pdd_config().apply_request_constraints(request, SamplingParam()) + + +def test_apply_request_constraints_non_pdd_checkpoint_keeps_request_steps(): + request = parse_config(GenerationRequest, {"prompt": "fox", "sampling": {"num_inference_steps": 50}}) + sampling_param = SamplingParam(num_inference_steps=50) + assert MiniMaxH3PipelineConfig().apply_request_constraints(request, sampling_param) is sampling_param + assert sampling_param.num_inference_steps == 50 + + +# ------------------------------------------------- worker pipeline: initialize_pipeline and the transformer check + + +def test_initialize_pipeline_pdd_checkpoint_keeps_resolved_settings(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + args = _run_args(model_path) + _pipeline(model_path).initialize_pipeline(args) + config = args.pipeline_config + assert config.dit_config.arch_config.pdd_steps == 32 + assert config.pdd_step_indices == tuple(GRID32_BLOCKS8) + assert config.dmd_denoising_steps is None + + +def test_initialize_pipeline_dmd_checkpoint_applies_rungs(tmp_path): + model_path = _checkpoint(tmp_path, DMD_FILE, pdd_steps=None) + args = _run_args(model_path) + _pipeline(model_path).initialize_pipeline(args) + assert args.pipeline_config.dmd_denoising_steps == DMD_FILE["dmd_denoising_steps"] + assert args.pipeline_config.pdd_step_indices is None + + +def test_check_pdd_transformer_widened_transformer_without_pdd_file(tmp_path): + model_path = _checkpoint(tmp_path, None, pdd_steps=32) + with pytest.raises(ValueError, match="sets pdd_steps, but .* has no PDD fastvideo_inference.json"): + _pipeline(model_path)._check_pdd_transformer(_run_args(model_path)) + + +def test_check_pdd_transformer_pdd_checkpoint_on_t2va_pipeline(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + with pytest.raises(ValueError, match="select MiniMaxH3Ref2VAModularPipeline"): + _pipeline(model_path, MiniMaxH3ModularPipeline)._check_pdd_transformer(_run_args(model_path)) + + +def test_check_pdd_transformer_pdd_checkpoint_on_ref2va_pipeline(tmp_path): + model_path = _checkpoint(tmp_path, PDD_FILE) + _pipeline(model_path)._check_pdd_transformer(_run_args(model_path)) + + +# ------------------------------------------------- worker pipeline: _load_config checks, before any component loads + +_GATE = "transformer_blocks.0.attn.to_gate_compress.weight" + + +def _write_gates(transformer_dir: Path, *, indexed: bool) -> None: + """Give a transformer checkpoint one VSA compression gate, listed in a shard index or in a safetensors file.""" + if indexed: + (transformer_dir / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {_GATE: "diffusion_pytorch_model-00001-of-00001.safetensors"}})) + else: + save_file({_GATE: torch.zeros(2, 2)}, str(transformer_dir / "diffusion_pytorch_model.safetensors")) + + +@pytest.fixture +def hub_snapshot(tmp_path, monkeypatch): + """Resolve the Hub repo id "org/fasth3" to tmp_path, as maybe_download_model does after downloading.""" + from fastvideo.pipelines import composed_pipeline_base + + downloads = [] + + def download(model_path, **kwargs): + downloads.append(model_path) + return str(tmp_path) if model_path == "org/fasth3" else model_path + + monkeypatch.setattr(composed_pipeline_base, "maybe_download_model", download) + monkeypatch.setattr(composed_pipeline_base, "verify_model_config_and_directory", + lambda model_path, **kwargs: {"_class_name": "MiniMaxH3ModularPipeline"}) + return downloads + + +@pytest.mark.parametrize("indexed", [True, False], ids=["shard_index", "safetensors_header"]) +@pytest.mark.parametrize("inference_file", [None, DMD_FILE], ids=["base", "dmd"]) +def test_load_config_gated_non_pdd_transformer_requires_vsa(tmp_path, hub_snapshot, inference_file, indexed): + """The gates are read from the resolved Hub snapshot. A PDD checkpoint's backend is checked by + resolve_checkpoint_settings instead (test_resolve_checkpoint_settings_explicit_training_fixed_conflict).""" + model_path = _checkpoint(tmp_path, inference_file, pdd_steps=None) + _write_gates(tmp_path / "transformer_ref", indexed=indexed) + pipeline = _pipeline("org/fasth3") + pipeline.fastvideo_args = _run_args(model_path, attention_backend="FLASH_ATTN") + with pytest.raises(ValueError, match="its transformer_ref carries VSA compression gates.*requests FLASH_ATTN"): + pipeline._load_config(pipeline.model_path) + assert hub_snapshot == ["org/fasth3"] and pipeline.model_path == model_path + pipeline.fastvideo_args.attention_backend = "VIDEO_SPARSE_ATTN_H3" + assert pipeline._load_config(pipeline.model_path) == {"_class_name": "MiniMaxH3ModularPipeline"} + + +@pytest.mark.parametrize("pdd_steps,gated,match", [ + (32, False, "has no PDD fastvideo_inference.json"), + (None, True, "carries VSA compression gates"), +], ids=["widened_without_pdd_file", "gated_with_other_backend"]) +def test_load_modules_transformer_check_before_loading_unless_supplied(tmp_path, monkeypatch, hub_snapshot, + pdd_steps, gated, match): + """A transformer that fails its check raises before any component loads; a transformer the caller already + built is not rebuilt, so its checkpoint is not checked.""" + loads = [] + + def base_load_modules(self, fastvideo_args, loaded_modules=None): + self._load_config(self.model_path) + loads.append(loaded_modules) + return {} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", base_load_modules) + model_path = _checkpoint(tmp_path, None, pdd_steps=pdd_steps) + if gated: + _write_gates(tmp_path / "transformer_ref", indexed=True) + pipeline = _pipeline(model_path) + args = _run_args(model_path, attention_backend="FLASH_ATTN", h3_sequential_load=False) + pipeline.fastvideo_args = args + with pytest.raises(ValueError, match=match): + pipeline.load_modules(args) + assert loads == [] + supplied = {"transformer": object()} + pipeline.load_modules(args, supplied) + assert loads == [supplied] + + +def test_checkpoint_facts_read_once_per_resolved_path(tmp_path, monkeypatch, hub_snapshot): + """_load_config (gate check) and initialize_pipeline (DMD rungs) share one parse of the inference file.""" + from fastvideo.pipelines.basic.minimax_h3 import minimax_h3_pipeline + + parsed = [] + + def recording_loads(text): + parsed.append(json.loads(text)) + return parsed[-1] + + monkeypatch.setattr(minimax_h3_pipeline, "json", SimpleNamespace(loads=recording_loads)) + model_path = _checkpoint(tmp_path, DMD_FILE, pdd_steps=None) + args = _run_args(model_path) + pipeline = _pipeline("org/fasth3") + pipeline.fastvideo_args = args + pipeline._load_config(pipeline.model_path) + pipeline.initialize_pipeline(args) + assert args.pipeline_config.dmd_denoising_steps == DMD_FILE["dmd_denoising_steps"] + assert parsed.count(DMD_FILE) == 1 diff --git a/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py b/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py new file mode 100644 index 000000000..d49706abb --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_pdd_denoising.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PDD sampling and Ref2VA reference-video VSA in the MiniMax-H3 denoising stage. + +A PDD-widened transformer samples fused blocks of its fine grid: the stage +partitions the grid (the checkpoint contract's ``pdd_step_indices``), hands +each scheduler its modality's node sigmas, starts the targets at the first +node's noise level, and fuses one block of heads per forward. The ordinary +Euler step over two node sigmas then equals the block advance. With +``vsa_ref_keep_rate`` set, every reference video is its own sparse VSA region. +All checks are CPU-only; the transformer is a tiny stand-in with real widened +heads. +""" +from __future__ import annotations + +import contextlib +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.forward_context import get_forward_context +from fastvideo.layers.pdd import ( + PDD_GRID_MAX_T, + PDDModalitySchedule, + PDDReplicatedLinear, + fuse_pdd_heads, + pdd_fine_grid, + shifted_noise_amount, +) +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_KEYFRAME_NOISE_AUG, + MINIMAX_H3_TEXT_TAG, + build_ref2va_packed_sequence, +) +from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference +from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_denoising as denoising +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +PATCH = (1, 2, 2) +VIDEO_CHANNELS = 96 # 24 latent channels x (1, 2, 2) patch +AUDIO_CHANNELS = 32 +GRID32_BLOCKS8 = tuple(range(0, 33, 4)) +# A 4-block partition of an 8-interval grid, for checks that need a short plan. +GRID8_BLOCKS4 = (0, 2, 4, 6, 8) + + +def _reference(kind: str, *, frames: int = 1, height: int = 4, width: int = 4, audio: int = 0): + return MiniMaxH3PreparedReference(media_type=kind, + has_audio=audio > 0, + num_latent_frames=frames, + latent_height=height, + latent_width=width, + num_audio_latents=audio) + + +def _ref2va_layout(references, *, text=5, frames=3, height=4, width=6, audio=3): + tags = torch.full((text, ), MINIMAX_H3_TEXT_TAG, dtype=torch.long) + return build_ref2va_packed_sequence(tags, references, frames, height, width, audio, PATCH) + + +def _stage() -> denoising.MiniMaxH3DenoisingStage: + return denoising.MiniMaxH3DenoisingStage( + transformer=SimpleNamespace(), + scheduler=MiniMaxH3Scheduler(shift=12.0), + audio_scheduler=MiniMaxH3Scheduler(shift=3.0), + ) + + +def _args(**pipeline_fields): + return SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig(**pipeline_fields), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + VSA_tile_size=128) + + +def _plan(stage, args): + return stage._pdd_sampling_plan(args, torch.device("cpu")) + + +def test_pdd_sampling_plan_unset_step_indices(): + assert _plan(_stage(), _args()) is None + + +def test_pdd_sampling_plan_uses_contract_partition_and_modality_shifts(): + plan = _plan(_stage(), _args(pdd_step_indices=GRID32_BLOCKS8)) + assert plan.indices.tolist() == list(GRID32_BLOCKS8) + base = pdd_fine_grid(32)[list(GRID32_BLOCKS8)] + torch.testing.assert_close(plan.node_sigmas["video"], shifted_noise_amount(base, 12.0), rtol=0, atol=0) + torch.testing.assert_close(plan.node_sigmas["audio"], shifted_noise_amount(base, 3.0), rtol=0, atol=0) + + +def test_step_pdd_block_equals_fused_block_advance(): + """x_end = x_start + sum_j w_j * dx/dsigma_j with the raw H3 output m = x0 - eps = -dx/dsigma.""" + torch.manual_seed(0) + stage = _stage() + plan = _plan(stage, _args(pdd_step_indices=GRID8_BLOCKS4)) + stage.scheduler.set_timesteps(sigmas=plan.node_sigmas["video"].to(torch.float32)) + stage.audio_scheduler.set_timesteps(sigmas=plan.node_sigmas["audio"].to(torch.float32)) + video, audio = torch.randn(5, 6), torch.randn(4, 3) + for step in range(plan.num_steps): + start, end = plan.block(step) + fused_video, fused_audio = torch.randn(5, 6), torch.randn(4, 3) + stepped_video = stage.scheduler.step(fused_video, stage.scheduler.timesteps[step], video, return_dict=False)[0] + stepped_audio = stage.audio_scheduler.step(fused_audio, + stage.audio_scheduler.timesteps[step], + audio, + return_dict=False)[0] + for name, state, fused, stepped in (("video", video, fused_video, stepped_video), + ("audio", audio, fused_audio, stepped_audio)): + total = plan.integration_weights[name][start:end].sum() + expected = state.double() - total * fused.double() + torch.testing.assert_close(stepped.double(), expected, rtol=1e-5, atol=1e-5) + video, audio = stepped_video, stepped_audio + + +def test_scale_pdd_initial_noise_starts_targets_at_first_node_sigma(): + plan = _plan(_stage(), _args(pdd_step_indices=GRID8_BLOCKS4)) + layout = SimpleNamespace(num_condition_video_rows=2, num_condition_audio_rows=1) + video, audio = torch.randn(5, 6), torch.randn(4, 3) + expected_video, expected_audio = video.clone(), audio.clone() + expected_video[2:] *= PDD_GRID_MAX_T + expected_audio[1:] *= PDD_GRID_MAX_T + + denoising.scale_pdd_initial_noise(video, audio, layout, plan) + + assert float(plan.node_sigmas["video"][0]) == pytest.approx(PDD_GRID_MAX_T) + assert float(plan.node_sigmas["audio"][0]) == pytest.approx(PDD_GRID_MAX_T) + torch.testing.assert_close(video, expected_video, rtol=0, atol=0) + torch.testing.assert_close(audio, expected_audio, rtol=0, atol=0) + + +class _TinyPDDTransformer: + """Stand-in DiT with real widened heads: each row's output depends on its own value and clock.""" + + def __init__(self, pdd_steps: int, hidden: int = 8): + self.pdd_steps = pdd_steps + generator = torch.Generator().manual_seed(7) + self.video_in = torch.randn(VIDEO_CHANNELS, hidden, generator=generator) / VIDEO_CHANNELS**0.5 + self.audio_in = torch.randn(AUDIO_CHANNELS, hidden, generator=generator) / AUDIO_CHANNELS**0.5 + self.proj_out = PDDReplicatedLinear(hidden, VIDEO_CHANNELS, grid_size=pdd_steps, params_dtype=torch.float32) + self.audio_proj_out = PDDReplicatedLinear(hidden, + AUDIO_CHANNELS, + grid_size=pdd_steps, + params_dtype=torch.float32) + with torch.no_grad(): + for linear in (self.proj_out, self.audio_proj_out): + linear.weight.copy_(torch.randn(linear.weight.shape, generator=generator) * 0.1) + linear.bias.copy_(torch.randn(linear.bias.shape, generator=generator) * 0.1) + self.blocks: list[tuple[int, int]] = [] + self.calls: list[dict] = [] + + @contextlib.contextmanager + def fuse_pdd_block(self, start, end, integration_weights, precision_decoding): + self.blocks.append((start, end)) + with fuse_pdd_heads({"video": self.proj_out, "audio": self.audio_proj_out}, start, end, integration_weights, + precision_decoding): + yield + + def features(self, rows, weights, row_times): + return torch.tanh(rows.float() @ weights + row_times[:, None]) + + def __call__(self, **kwargs): + row_times = kwargs["timestep"][kwargs["timestep_indices"]] + self.calls.append({ + **{key: value.clone() for key, value in kwargs.items() if torch.is_tensor(value)}, + "attn_metadata": get_forward_context().attn_metadata, + }) + video = self.features(kwargs["hidden_states"][0], self.video_in, row_times[kwargs["video_indices"]]) + audio = self.features(kwargs["audio_hidden_states"][0], self.audio_in, row_times[kwargs["audio_indices"]]) + return self.proj_out(video)[0][None], self.audio_proj_out(audio)[0][None] + + +def _run(monkeypatch, layout, transformer, *, args, num_inference_steps=8, sparsity=0.0, extra=None, builder=None): + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(denoising, "_h3_vsa_metadata_builder", lambda *_: builder) + stage = denoising.MiniMaxH3DenoisingStage(transformer, MiniMaxH3Scheduler(shift=12.0), + MiniMaxH3Scheduler(shift=3.0)) + generator = torch.Generator().manual_seed(3) + latents = torch.randn(int(layout.video_indices.numel()), VIDEO_CHANNELS, generator=generator) + audio_latents = torch.randn(int(layout.audio_indices.numel()), AUDIO_CHANNELS, generator=generator) + batch = ForwardBatch(data_type="video", + prompt_embeds=[torch.zeros(1, int(layout.text_indices.numel()), 8)], + latents=latents.clone(), + audio_latents=audio_latents.clone(), + num_inference_steps=num_inference_steps, + VSA_sparsity=sparsity, + extra={ + MINIMAX_H3_LAYOUT_KEY: layout, + **(extra or {}) + }) + return stage, stage.forward(batch, args), latents, audio_latents + + +def _reference_video_and_audio_references(): + return [ + _reference("image", height=4, width=6), + _reference("video", frames=5, height=4, width=6, audio=2), + _reference("audio", audio=3), + ] + + +def test_forward_pdd_runs_eight_fused_blocks_from_contract(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8) + stage, result, latents, audio_latents = _run(monkeypatch, layout, transformer, args=args) + + assert len(transformer.calls) == 8 and result.step_index == 7 + assert transformer.blocks == list(zip(GRID32_BLOCKS8[:-1], GRID32_BLOCKS8[1:], strict=True)) + for proj in (transformer.proj_out, transformer.audio_proj_out): + assert proj._fusion_state is None + + # Independent recomputation: node sigmas on each modality's clock, targets + # start at 0.999 * noise, conditions stay fixed, and every block advances + # x by -(sigma_end - sigma_start) * (weighted mean of its raw heads). + base = pdd_fine_grid(32)[list(GRID32_BLOCKS8)] + sigmas = {"video": shifted_noise_amount(base, 12.0), "audio": shifted_noise_amount(base, 3.0)} + torch.testing.assert_close(stage.scheduler.sigmas, sigmas["video"].float(), rtol=0, atol=0) + torch.testing.assert_close(stage.audio_scheduler.sigmas, sigmas["audio"].float(), rtol=0, atol=0) + weights = {name: PDDModalitySchedule(shift).integration_weights(pdd_fine_grid(32)) + for name, shift in (("video", 12.0), ("audio", 3.0))} + video_start, audio_start = layout.num_condition_video_rows, layout.num_condition_audio_rows + video, audio = latents.double().clone(), audio_latents.double().clone() + video[video_start:] *= PDD_GRID_MAX_T + audio[audio_start:] *= PDD_GRID_MAX_T + for step, call in enumerate(transformer.calls): + start, end = GRID32_BLOCKS8[step], GRID32_BLOCKS8[step + 1] + torch.testing.assert_close(call["hidden_states"][0].double(), video, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(call["audio_hidden_states"][0].double(), audio, rtol=1e-5, atol=1e-5) + row_times = call["timestep"][call["timestep_indices"]] + video_time = 1.0 - float(sigmas["video"][step].float()) + audio_time = 1.0 - float(sigmas["audio"][step].float()) + video_rows, audio_rows = call["video_indices"], call["audio_indices"] + for rows, value in ((video_rows[video_start:], video_time), (video_rows[:video_start], + max(video_time, MINIMAX_H3_KEYFRAME_NOISE_AUG)), + (audio_rows[audio_start:], audio_time), (audio_rows[:audio_start], 1.0), + (call["text_indices"], video_time)): + torch.testing.assert_close(row_times[rows], torch.full_like(row_times[rows], value), rtol=0, atol=1e-7) + for name, state, rows, first, heads, features_in in ( + ("video", video, video_rows, video_start, transformer.proj_out, transformer.video_in), + ("audio", audio, audio_rows, audio_start, transformer.audio_proj_out, transformer.audio_in), + ): + features = transformer.features(state, features_in, row_times[rows].double().float()) + materialized = heads(features.float())[0].double().unflatten(-1, (32, -1)) + alpha = weights[name][start:end] + mean = torch.einsum("n,rnc->rc", alpha / alpha.sum(), materialized[:, start:end]) + state[first:] -= (sigmas[name][step + 1] - sigmas[name][step]) * mean[first:] + torch.testing.assert_close(result.latents.double(), video, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(result.audio_latents.double(), audio, rtol=1e-4, atol=1e-4) + # Reference conditions never move. + torch.testing.assert_close(result.latents[:video_start], latents[:video_start], rtol=0, atol=0) + torch.testing.assert_close(result.audio_latents[:audio_start], audio_latents[:audio_start], rtol=0, atol=0) + + +@pytest.mark.parametrize("dense_first_n", [0, 1]) +def test_forward_ref_keep_rate_builds_one_sparse_region_per_reference_video(monkeypatch, dense_first_n): + references = [ + _reference("video", frames=5, height=4, width=6), + _reference("image", height=4, width=6), + _reference("video", frames=2, height=8, width=4, audio=2), + ] + layout = _ref2va_layout(references) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.25) + _run(monkeypatch, + layout, + transformer, + args=args, + sparsity=0.9, + extra={"vsa_dense_first_n_steps": dense_first_n}, + builder=MiniMaxH3VSAMetadataBuilder()) + + assert len(transformer.calls) == 8 + for step, call in enumerate(transformer.calls): + metadata = call["attn_metadata"] + assert metadata.tile_elems == 128 + assert metadata.total_seq_length == layout.sequence_length + assert len(metadata.video_tile_spans) == 3 + expected = (0.0, 0.0, 0.0) if step < dense_first_n else (0.75, 0.75, 0.9) + assert metadata.span_sparsities == pytest.approx(expected) + + +@pytest.mark.parametrize("sparsity,effect", [ + (0.9, "each reference keeps 0.25 of its tiles and the target keeps 0.1"), + (0.0, "VSA_sparsity=0 keeps every region dense"), +]) +def test_forward_ref_keep_rate_logs_applied_sparsity(monkeypatch, sparsity, effect): + infos = [] + monkeypatch.setattr(denoising.logger, "info", lambda message, *args: infos.append(message % args)) + layout = _ref2va_layout(_reference_video_and_audio_references()) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.25) + _run(monkeypatch, + layout, + _TinyPDDTransformer(pdd_steps=32), + args=args, + sparsity=sparsity, + builder=MiniMaxH3VSAMetadataBuilder()) + assert f"MiniMax-H3 VSA-H3: 1 reference video region(s); {effect}; 128-token tiles." in infos + + +def test_forward_unset_ref_keep_rate_keeps_references_in_dense_prefix(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + transformer = _TinyPDDTransformer(pdd_steps=32) + args = _args(pdd_step_indices=GRID32_BLOCKS8) + _run(monkeypatch, layout, transformer, args=args, sparsity=0.9, builder=MiniMaxH3VSAMetadataBuilder()) + for call in transformer.calls: + metadata = call["attn_metadata"] + # The generated video is the only sparse region. + assert len(metadata.video_tile_spans) == 1 and metadata.span_sparsities == (0.9, ) + + +def test_forward_ref_keep_rate_requires_exempt_prefix_keys(monkeypatch): + layout = _ref2va_layout(_reference_video_and_audio_references()) + args = _args(pdd_step_indices=GRID32_BLOCKS8, vsa_ref_keep_rate=0.1) + with pytest.raises(ValueError, match="compete' supports only the generated-video region"): + _run(monkeypatch, + layout, + _TinyPDDTransformer(pdd_steps=32), + args=args, + sparsity=0.9, + extra={"vsa_mode": "compete"}, + builder=MiniMaxH3VSAMetadataBuilder()) diff --git a/fastvideo/tests/test_partial_hf_download.py b/fastvideo/tests/test_partial_hf_download.py index 6cb9df925..b83e2a503 100644 --- a/fastvideo/tests/test_partial_hf_download.py +++ b/fastvideo/tests/test_partial_hf_download.py @@ -175,7 +175,7 @@ def test_direct_pipeline_load_forwards_partial_patterns(tmp_path: Path, monkeypa monkeypatch.setattr("fastvideo.pipelines.composed_pipeline_base.maybe_download_model", fake_download) pipeline = object.__new__(MiniMaxH3Ref2VAModularPipeline) pipeline.model_path = "MiniMaxAI/MiniMax-H3" - pipeline.fastvideo_args = SimpleNamespace(revision="test-revision") + pipeline.fastvideo_args = SimpleNamespace(revision="test-revision", pipeline_config=SimpleNamespace(pdd_step_indices=None)) pipeline._load_config(pipeline.model_path) @@ -250,3 +250,40 @@ def test_model_index_supports_umbrella_repo_paths(tmp_path: Path, monkeypatch) - assert captured["filename"] == "minimax-h3/model_index.json" assert captured["revision"] == "test-revision" assert config["pipeline_name"] == "MiniMaxH3ModularPipeline" + + +def test_read_optional_model_json_local_directory(tmp_path: Path) -> None: + (tmp_path / "scheduler").mkdir() + (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps({"shift": 12.0})) + + assert utils.read_optional_model_json(str(tmp_path), "scheduler/scheduler_config.json") == {"shift": 12.0} + assert utils.read_optional_model_json(str(tmp_path), "fastvideo_inference.json") is None + + +def test_read_optional_model_json_umbrella_repo_file(tmp_path: Path, monkeypatch) -> None: + captured: dict = {} + + def fake_hf_hub_download(**kwargs): + captured.update(kwargs) + path = tmp_path / kwargs["filename"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"pdd_steps": 32})) + return str(path) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download) + + config = utils.read_optional_model_json("org/repo/student", "transformer_ref/config.json", revision="rev") + + assert config == {"pdd_steps": 32} + assert captured == {"repo_id": "org/repo", "filename": "student/transformer_ref/config.json", "revision": "rev"} + + +def test_read_optional_model_json_absent_hub_file(monkeypatch) -> None: + from huggingface_hub.utils import EntryNotFoundError + + def fake_hf_hub_download(**kwargs): + raise EntryNotFoundError(f"{kwargs['filename']} is not in the repo") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_hf_hub_download) + + assert utils.read_optional_model_json("org/repo", "fastvideo_inference.json") is None diff --git a/fastvideo/utils.py b/fastvideo/utils.py index b42c3f606..10d142b76 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -747,6 +747,32 @@ def maybe_download_model_index(model_name_or_path: str, revision: str | None = N raise ValueError(f"Failed to download or parse a Diffusers manifest for {model_name_or_path}: {e}") from e +def read_optional_model_json(model_name_or_path: str, + relative_path: str, + revision: str | None = None) -> dict[str, Any] | None: + """Read one JSON file of a local model directory or Hub repo; None when the model has no such file. + + For a Hub repo only that file is downloaded, into the shared HF cache, so + a run can read small checkpoint metadata before any component weights. + """ + if os.path.exists(model_name_or_path): + path = os.path.join(model_name_or_path, relative_path) + if not os.path.isfile(path): + return None + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + repo_id, subfolder = _split_hf_repo_subfolder(model_name_or_path) + filename = f"{subfolder}/{relative_path}" if subfolder else relative_path + try: + path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) + except EntryNotFoundError: + return None + with open(path, encoding="utf-8") as f: + return cast(dict[str, Any], json.load(f)) + + # Token variables that `huggingface_hub` itself reads, in its priority order. HF_TOKEN_ENV_VARS = ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN") # FastVideo-specific token names, kept as deprecated aliases until the next From 0cc41a22dc651a7b77a559225f0cbdfe384e8c98 Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Fri, 2 Oct 2026 12:37:34 -0700 Subject: [PATCH 047/151] [bugfix] sm_100a VSA: fence the CLC response read before the slot release; forward PV descriptor lifetime (#1904) Co-authored-by: Claude Opus 5.5 --- .../attention/block_sparse_kernel_sm100a.cuh | 8 +- .../csrc/attention/primitives.cuh | 5 +- .../tests/test_block_sparse_sm100a_clc.py | 75 +++++++++++++++++++ 3 files changed, 80 insertions(+), 8 deletions(-) create mode 100644 fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py diff --git a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh index fa17cf251..d2af53a1a 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh @@ -2,8 +2,6 @@ // data-center Blackwell sm_100a/sm_103a. The filename is retained for API compatibility. // Warp-specialized: load / MMA (tcgen05) / softmax / correction / epilogue / scheduler. // Writes O and, when asked, the log-sum-exp the backward consumes. -// -// Generated (comments stripped). Do not edit by hand. #ifndef BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH #define BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH @@ -506,9 +504,11 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q, const uint32_t o_tmem_addr = tmem_base + (uint32_t)(2 * S_COLS + i * O_COLS); int slot = 0; + // Declared outside the loop: with BLK128 only p == 0 initializes it and p == 1 continues + // the p == 0 descriptor walk. + SmemDescPair desc_bv; #pragma unroll for (int p = 0; p < V_SUBTILES; ++p) { - SmemDescPair desc_bv; if (!BLK128 || p == 0) { slot = kv_ph.get_stage(); mbarrier_wait_parity(smem_ptr_u32(&full_bar[slot]), kv_ph.get_phase()); @@ -518,8 +518,6 @@ fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q, desc_bv.u64 = desc_v0; desc_bv.w.x += (uint32_t)(slot * (int)KV_DESC_DELTA); } - const uint64_t dbV = desc_kv0 + (uint64_t)slot * KV_DESC_DELTA - + (BLK128 ? (uint64_t)p * (V_BLK_BYTES >> 4) : 0u); #pragma unroll for (int ki = 0; ki < K_ATOMS_PER_TILE; ++ki) { const int a = p * K_ATOMS_PER_TILE + ki; diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh index 895519dda..3ffa17e56 100644 --- a/fastvideo-kernel/csrc/attention/primitives.cuh +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -1,8 +1,6 @@ // primitives.cuh -- device primitives for the sm_100a/sm_103a VSA block-sparse attention // forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store / // tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers. -// -// Generated and pruned to what the kernel reaches -- do not edit by hand. #pragma once #include #include @@ -218,8 +216,8 @@ template __device__ __forceinline__ ClcTileInfo clc_parse_response(uint32_t resp_smem_addr) { uint32_t d0, d1, d2, d3; - fence_proxy_async_shared_cta(); clc_load_response(resp_smem_addr, d0, d1, d2, d3); + fence_proxy_async_shared_cta(); const int ctaid_x = static_cast(d0); const int ctaid_y = static_cast(d1 & 0xFFFFu); const bool valid = (d2 & 1u) != 0u; @@ -251,6 +249,7 @@ ClcTileInfo clc_fetch_next_tile( __cvta_generic_to_shared(&clc_response[clc_cons_stage * 4])); ClcTileInfo t = clc_parse_response< CLUSTER_SHAPE_M, CLUSTER_SHAPE_N, ORDER>(resp_addr); + __syncwarp(); if (do_release) { uint32_t empty_local = static_cast( __cvta_generic_to_shared(&clc_empty_bar[clc_cons_stage])); diff --git a/fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py b/fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py new file mode 100644 index 000000000..2a22aee23 --- /dev/null +++ b/fastvideo-kernel/tests/test_block_sparse_sm100a_clc.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""sm_100a VSA forward: CLC work-item scheduler under slot reuse. + +The persistent forward steals not-yet-launched CTAs with clusterlaunchcontrol.try_cancel and +hands the 16-byte responses through a two-slot ring. The ring is only exercised when the grid +is several waves deep: with every CTA resident no cancel succeeds and no slot is ever reused, +which is the case for the small correctness tests. Here the grid is >= 3 waves, every call's +output and LSE are compared bitwise against a golden call of the same binary, and the buffers +the kernel writes are pre-filled with NaN, so a tile the kernel never wrote shows up as NaN +and a tile written from the wrong work item as a value mismatch. +""" + +import pytest +import torch + +from fastvideo_kernel import block_sparse_attn_sm100a as vsa + +HEAD_DIM = 128 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} + or not vsa._HAS_VSA_SM100A, + reason="requires data-center Blackwell (sm_100a/sm_103a) and a built fastvideo_kernel extension", +) + + +def make_case(block, num_blocks, heads, topk, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + S = num_blocks * block + shape = (1, heads, S, HEAD_DIM) if vsa.BHSD else (1, S, heads, HEAD_DIM) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=g) for _ in range(3)) + scores = torch.rand((heads * num_blocks, num_blocks), device="cuda", generator=g) + idx = scores.topk(topk, dim=1).indices.sort(dim=1).values.to(torch.int32) + idx = idx.view(1, heads, num_blocks, topk).contiguous() + num = torch.randint(1, topk + 1, (1, heads, num_blocks), device="cuda", generator=g, dtype=torch.int32) + vbs = torch.randint(block // 2, block + 1, (num_blocks, ), device="cuda", generator=g, dtype=torch.int32) + return q, k, v, idx, num, vbs + + +@pytest.mark.parametrize("block,num_blocks", [(128, 128), (64, 256)], ids=["blk128", "blk64"]) +def test_clc_slot_reuse_never_drops_or_misplaces_a_tile(block, num_blocks): + if vsa._FWD_BY_BLOCK.get(block) is None: + pytest.skip(f"no {block}-token forward in this build") + heads, topk, calls = 16, 4, 300 + q, k, v, idx, num, vbs = make_case(block, num_blocks, heads, topk) + assert vsa.is_supported(q, vbs) + ctas = heads * num_blocks // 2 + assert ctas >= 3 * torch.cuda.get_device_properties(0).multi_processor_count, ctas + + out_g, lse_g = (t.clone() for t in vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True)) + assert torch.isfinite(out_g).all() and torch.isfinite(lse_g).all() + for _ in range(3): + out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True) + assert torch.equal(out, out_g) and torch.equal(lse, lse_g), "forward is not bitwise deterministic" + del out, lse + + S = num_blocks * block + tile_shape = (1, heads, num_blocks, block * HEAD_DIM) + nan = float("nan") + failures = [] + for call in range(calls): + # Pre-fill the blocks the op is about to allocate: same sizes, freed in reverse order. + poison_lse = torch.full((1, heads, S), nan, device="cuda", dtype=torch.float32) + poison_out = torch.full_like(q, nan) + out_ptr, lse_ptr = poison_out.data_ptr(), poison_lse.data_ptr() + del poison_out, poison_lse + out, lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs, need_lse=True) + assert out.data_ptr() == out_ptr and lse.data_ptr() == lse_ptr, "allocator did not reuse the poisoned blocks" + wrong = (out != out_g).view(tile_shape).any(-1) | (lse != lse_g).view(1, heads, num_blocks, block).any(-1) + if bool(wrong.any()): + never_written = int(out.isnan().view(tile_shape).all(-1).sum()) + failures.append((call, int(wrong.sum()), never_written)) + del out, lse + assert not failures, (f"{len(failures)} of {calls} calls had wrong tiles; first (call, wrong tiles, " + f"tiles never written): {failures[:5]}") From a4d6416c526b2af8493bb617a76ec1c412d3d747 Mon Sep 17 00:00:00 2001 From: YZJF Date: Sun, 4 Oct 2026 12:20:06 -0700 Subject: [PATCH 048/151] [docs] Add OpenAI serving recipes for Wan cookbook pages (#1906) Co-authored-by: SolitaryThinker Co-authored-by: Claude Sonnet 5.5 --- docs/assets/cookbook-recipes.json | 37 ++++++++-- docs/assets/cookbook.js | 43 ++++++++---- docs/assets/custom.css | 3 +- docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 10 +-- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 53 +++++++++++++- docs/cookbook/z-image.md | 2 +- docs/generate_examples.py | 13 +++- examples/serving/openai_wan22_t2v_a14b.yaml | 1 + examples/serving/openai_wan22_ti2v_5b.yaml | 9 ++- tests/local_tests/test_cookbook_serving.py | 76 +++++++++++++++++++++ 22 files changed, 228 insertions(+), 43 deletions(-) diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 8875361dc..e89e58d12 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 11, + "version": 12, "recipes": [ { "id": "fastwan21-t2v", @@ -10,6 +10,11 @@ "summary": "Generate a video in three denoising steps with the distilled FastWan2.1 1.3B checkpoint and video sparse attention.", "model": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", "source": "scripts/inference/inference_wan_VSA_DMD_1_3B.yaml", + "serving": { + "source": "examples/serving/openai_fastwan21_1_3b.yaml", + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", + "env": "FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN" + }, "command": "FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN fastvideo generate --config scripts/inference/inference_wan_VSA_DMD_1_3B.yaml", "gpu_types": ["NVIDIA"], "hardware": {"gpu_count": 1, "evidence": "source-configured"}, @@ -25,6 +30,13 @@ "summary": "The maintained high-capacity Wan2.2 text-to-video example with CPU offload settings encoded in its checked-in Python source.", "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "source": "examples/inference/basic/basic_wan2_2.py", + "serving": { + "source": "examples/serving/openai_wan22_t2v_a14b.yaml", + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", + "limitations": [ + "Generation at 1280 × 720 on two GPUs with DiT CPU offload is slow. The client examples stop polling after 30 minutes while the job keeps running on the server; retrieve it by ID or raise the deadline." + ] + }, "command": "python examples/inference/basic/basic_wan2_2.py", "gpu_types": ["NVIDIA"], "hardware": {"gpu_count": 2, "evidence": "source-configured"}, @@ -54,6 +66,14 @@ "summary": "Use one maintained 5B checkpoint for text-to-video or add an image input to switch the same recipe to image-to-video.", "model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "source": "examples/inference/basic/basic_wan2_2_ti2v.py", + "serving": { + "source": "examples/serving/openai_wan22_ti2v_5b.yaml", + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", + "task": "Text to video", + "limitations": [ + "This server walkthrough sends a text prompt only. The server API also accepts an image reference (input_reference) for this checkpoint, but the checked-in clients do not upload one. For the cookbook's image-to-video workflow, use the Python command." + ] + }, "command": "python examples/inference/basic/basic_wan2_2_ti2v.py", "gpu_types": ["NVIDIA"], "hardware": {"gpu_count": 1, "evidence": "source-configured"}, @@ -456,7 +476,8 @@ "source": "examples/inference/basic/basic_fasth3.py", "serving": { "source": "examples/serving/openai_fasth3.yaml", - "install": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"" + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"", + "audio": true }, "command": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"\npython examples/inference/basic/basic_fasth3.py --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --profile all", "gpu_types": ["NVIDIA"], @@ -481,7 +502,8 @@ "serving": { "source": "examples/serving/mlx_fasth3.yaml", "install": "uv pip install -e \".[mlx]\"", - "prepare": "hf download FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2 --local-dir ./FastH3-Preview-v0.2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-Preview-v0.2/transformer --out ./FastH3-MLX --formats \"int6\"" + "prepare": "hf download FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2 --local-dir ./FastH3-Preview-v0.2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-Preview-v0.2/transformer --out ./FastH3-MLX --formats \"int6\"", + "audio": true }, "group": "fasth3-preview", "group_label": "FastH3 V1", @@ -528,7 +550,8 @@ "serving": { "source": "examples/serving/openai_fasth3_spark.yaml", "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", - "env": "FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3" + "env": "FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3", + "audio": true }, "command": "FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_fasth3_spark.yaml", "gpu_types": ["NVIDIA"], @@ -592,7 +615,8 @@ "source": "examples/inference/basic/basic_fasth3_8step.py", "serving": { "source": "examples/serving/openai_fasth3_8step.yaml", - "install": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"" + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"", + "audio": true }, "command": "UV_TORCH_BACKEND=cu130 uv pip install -e \".[fasth3]\"\npython examples/inference/basic/basic_fasth3_8step.py --prompt \"(S1) A presenter says [English] FastVideo runs FastH3.\" --profile strict --no-inference-torch-compile --no-compile-vae", "gpu_types": ["NVIDIA"], @@ -621,7 +645,8 @@ "serving": { "source": "examples/serving/mlx_fasth3_8step.yaml", "install": "uv pip install -e \".[mlx]\"", - "prepare": "hf download FastVideo/FastVideo-FastH3-8-Step-V2 --local-dir ./FastH3-8-Step-V2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-8-Step-V2/transformer --out ./FastH3-8-Step-V2-MLX --formats \"int8\" --include-vsa" + "prepare": "hf download FastVideo/FastVideo-FastH3-8-Step-V2 --local-dir ./FastH3-8-Step-V2\npython scripts/checkpoint_conversion/convert_minimax_h3_mlx.py --model-root ./FastH3-8-Step-V2/transformer --out ./FastH3-8-Step-V2-MLX --formats \"int8\" --include-vsa", + "audio": true }, "group": "fasth3-8step-v2", "group_label": "FastH3 V2", diff --git a/docs/assets/cookbook.js b/docs/assets/cookbook.js index 5b00e95ae..689552355 100644 --- a/docs/assets/cookbook.js +++ b/docs/assets/cookbook.js @@ -506,6 +506,10 @@ const runtime = runtimeFor(recipe); const profile = servingPanel && servingProfiles[recipe.id]; const useServer = Boolean(profile && usagePreference === "server"); + // Only some servers expose the browser playground or return audio; both come from the serving profile. + const hasPlayground = Boolean(profile && profile.playground_url); + const hasAudio = Boolean(profile && profile.audio); + const playgroundAny = familyRecipes.some((item) => servingProfiles[item.id] && servingProfiles[item.id].playground_url); // The measured local profile and the server config have separate evidence. const activeRecipe = useServer ? { ...recipe, hardware: profile.hardware, evidence: "Source-backed" } : recipe; const knobs = knobsFor(recipe); @@ -515,14 +519,19 @@ usage.querySelectorAll("[data-cookbook-mode]").forEach((option) => { const selected = option.dataset.cookbookMode === (useServer ? "server" : "python"); option.disabled = option.dataset.cookbookMode === "server" && !profile; + const hint = option.querySelector("[data-cookbook-server-hint]"); + if (hint) hint.textContent = (profile ? hasPlayground : playgroundAny) ? "Playground, cURL, or an API client" : "cURL or an API client"; option.classList.toggle("cookbook-option--selected", selected); option.setAttribute("aria-pressed", String(selected)); }); + const servedNames = [...new Set(familyRecipes.filter((item) => servingProfiles[item.id]).map((item) => item.group_label || item.label))]; servingAvailability.textContent = profile - ? "The playground and the OpenAI Python client share one server process. Both workflows can run on your own machine." + ? `${hasPlayground ? "The playground" : "cURL"} and the OpenAI Python client share one server process. Both workflows can run on your own machine.` : servingLoadFailed - ? "Server examples could not be loaded. Open the H3 server guide below, or use Python directly." - : "This recipe uses Python directly. FastH3 V1 and FastH3 V2 can also run a local server for the playground and the OpenAI Python client."; + ? "Server examples could not be loaded. Use Python directly." + : servedNames.length + ? `This recipe uses Python directly. ${new Intl.ListFormat("en").format(servedNames)} can also run a local server for ${playgroundAny ? "the playground and " : "cURL and "}the OpenAI Python client.` + : "This recipe uses Python directly."; servingPanel.hidden = !useServer; commandBlock.hidden = useServer; root.querySelector("[data-cookbook-python-note]").hidden = useServer; @@ -531,11 +540,12 @@ if (useServer) { const isMLX = profile.runtime === "mlx"; const isSpark = runtime.id === "spark"; + const where = hasPlayground ? "in the playground or your app" : "in your app"; servingPanel.querySelector("[data-cookbook-server-lifetime]").textContent = isMLX - ? "Start once, then change prompts in the playground or your app. MLX reuses its pipeline and prompt cache, but loads and releases model components between phases to limit unified-memory use. It does not keep all weights resident." + ? `Start once, then change prompts ${where}. MLX reuses its pipeline and prompt cache, but loads and releases model components between phases to limit unified-memory use. It does not keep all weights resident.` : isSpark - ? "Start once, then change prompts in the playground or your app. On a DGX Spark, lazy module load still reloads Qwen3-VL and the DiT between phases of each request, so later prompts are not a free hot cache." - : "Start once, then change prompts in the playground or your app. CUDA requests reuse the loaded model. The Python SDK can also reuse a generator within one process."; + ? `Start once, then change prompts ${where}. On a DGX Spark, lazy module load still reloads Qwen3-VL and the DiT between phases of each request, so later prompts are not a free hot cache.` + : `Start once, then change prompts ${where}. CUDA requests reuse the loaded model. The Python SDK can also reuse a generator within one process.`; servingPanel.querySelector("[data-cookbook-install-guide]").href = isMLX ? "../../getting_started/installation/mlx/" : isSpark @@ -546,7 +556,10 @@ servingPanel.querySelector("[data-cookbook-server-install]").textContent = profile.install; servingPanel.querySelector("[data-cookbook-server-command]").textContent = profile.command; servingPanel.querySelector("[data-cookbook-health-command]").textContent = profile.health_command; - servingPanel.querySelector("[data-cookbook-playground]").href = profile.playground_url; + servingPanel.querySelectorAll("[data-cookbook-playground-only]").forEach((element) => { + element.hidden = !hasPlayground; + }); + if (hasPlayground) servingPanel.querySelector("[data-cookbook-playground]").href = profile.playground_url; const client = profile.clients[selectedClient]; const filename = client.source.split("/").pop(); servingPanel.querySelector("[data-cookbook-client-install]").textContent = client.install; @@ -574,13 +587,17 @@ }); description.textContent = useServer - ? `${recipe.group_label || recipe.label} generates video with audio. Start the local server, then use the playground or the OpenAI Python client. This profile uses the checked-in ${runtime.label} configuration.` + ? `${recipe.group_label || recipe.label} generates video${hasAudio ? " with audio" : ""}. Start the local server, then use ${hasPlayground ? "the playground or " : "cURL or "}the OpenAI Python client. This profile uses the checked-in ${runtime.label} configuration.` : recipe.summary; label.textContent = useServer ? `${recipe.group_label || recipe.label} · Server` : recipe.label; model.textContent = recipe.model; - task.textContent = recipe.task; + task.textContent = (useServer && recipe.serving && recipe.serving.task) || recipe.task; hardwareValue.textContent = runtimeSummary(activeRecipe); - if (artifact) artifact.textContent = useServer ? "MP4 with audio" : recipe.expected_artifact || "Not yet documented for this recipe."; + if (artifact) { + artifact.textContent = useServer + ? (hasAudio ? "MP4 with audio" : "MP4 video") + : recipe.expected_artifact || "Not yet documented for this recipe."; + } if (evidenceCell) { evidenceCell.textContent = activeRecipe.evidence || "Source-backed"; evidenceCell.classList.toggle("cookbook-badge--verified", activeRecipe.evidence === "Verified"); @@ -602,10 +619,12 @@ ? "This MLX server config has no recorded hardware run. Measurements from the Python recipe are not server memory requirements. Only text-to-video/audio is wired; reference inputs and fast modes are not exposed here." : runtime.id === "spark" ? "This Spark server config has no recorded serving benchmark. Lazy module load reloads Qwen3-VL and the DiT between phases of each request. Compilation of the DiT is disabled." - : "This server config has no recorded serving benchmark. Compilation is disabled, unlike the measured Python performance profile.", + : `This server config has no recorded serving benchmark.${profile.compile_enabled === false ? " Compilation is disabled, unlike the measured Python performance profile." : ""}`, `${profile.sampling.width} × ${profile.sampling.height} · ${profile.sampling.num_frames} frames · ${profile.sampling.fps} fps. The server supplies these defaults; the client sends the model and prompt.`, "Generation is serialized. Job metadata is held in memory and is lost when the server restarts.", - ] : recipe.limitations || []), ...knobCaveats]; + ] : recipe.limitations || []), + ...(useServer && recipe.serving && Array.isArray(recipe.serving.limitations) ? recipe.serving.limitations : []), + ...knobCaveats]; notes.replaceChildren(); notes.hidden = limitations.length === 0; if (limitations.length) { diff --git a/docs/assets/custom.css b/docs/assets/custom.css index 03a6eee7a..f0a30131f 100644 --- a/docs/assets/custom.css +++ b/docs/assets/custom.css @@ -1515,7 +1515,8 @@ img { .cookbook-serving[hidden], .cookbook-command[hidden], -[data-cookbook-python-note][hidden] { +[data-cookbook-python-note][hidden], +[data-cookbook-playground-only][hidden] { display: none; } diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 2ee864e5b..fbb813fae 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
      +
      All model families
      diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index b8bbb5905..f2f159aed 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
      +
      All model families
      diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index b4a2c3967..ed88351a1 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
      +
      All model families
      diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index 8408dad62..66acf7f28 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
      +
      All model families
      diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index e07993f24..119296eae 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
      +
      All model families
      diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index f89806052..ce2a81574 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
      +
      All model families
      diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index de1bf1c98..a9354f30d 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
      +
      All model families
      diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index 8bcb544f5..26160e865 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
      +
      All model families
      diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 8f3b4f33b..01f77a712 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -11,7 +11,7 @@ a full model, not a demo. **V2** is the eight-step checkpoint. More forwards is why V2 is the higher-quality FastH3. The V2 schedule contract is in [FastH3 distilled checkpoint schedules](../inference/fasth3-distilled.md). -
      +
      All model families
      @@ -154,7 +154,7 @@ is why V2 is the higher-quality FastH3. The V2 schedule contract is in Both can run locally
      - +
      @@ -201,17 +201,17 @@ is why V2 is the higher-quality FastH3. The V2 schedule contract is in

      Start the server

      -

      Keep this terminal running while you use the playground or API clients.

      +

      Keep this terminal running while you use the playground or API clients.

      GPU machine · Terminal
      Check that the server is ready

      In another terminal, this returns {"status":"ok"} after startup.

      Generate and download a video

      -
      +
      Try prompts in your browser

      Edit a prompt, generate, and watch the result. The playground uses the same server as cURL and your app.

      Open playground
      -

      Open after the server is ready. On a remote GPU machine, forward port 8000 to your computer first. This opens a local page, not a hosted demo.

      +

      Open after the server is ready. On a remote GPU machine, forward port 8000 to your computer first. This opens a local page, not a hosted demo.

      Use cURL or an SDK

      Each example submits a job, checks its status, and saves the MP4. The Python and JavaScript examples use OpenAI-compatible clients; no OpenAI account is needed.

      diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 8bdd93515..5e127c9b2 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
      +
      All model families
      diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index a2ca588ef..20dc7bfd7 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
      +
      All model families
      diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index a3c8c62ab..07c87fdb5 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index 92b259fec..14820b563 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index 4554e0966..c36c32508 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
      +
      All model families
      @@ -122,6 +122,17 @@ hide:

      Loading recipe details...

      +
      +
      + Workflow + Both can run locally +
      +
      + + +
      +
      +

      Exact device and memory details appear only when a recorded run supports them.

      @@ -150,7 +161,45 @@ hide:
      Terminal
      - 
      Loading...
      + 
      Loading...
      +
      +

      Running this script again starts a new process and reloads the model. To iterate in Python, create the generator once and reuse it for multiple prompts.

      + +
      diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index 8f4301f64..b53c79ef1 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
      +
      All model families
      diff --git a/docs/generate_examples.py b/docs/generate_examples.py index 690084081..cf0900b58 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -293,12 +293,16 @@ def cookbook_serving_profile(recipe: dict) -> dict: for language, (filename, install) in COOKBOOK_CLIENTS.items(): path = ROOT_DIR / "examples/serving/clients" / filename text = path.read_text(encoding="utf-8") - # Keep displayed snippets identical to executable sources for the - # checked-in H3 profile, with only endpoint and alias substitutions. + # Keep displayed snippets identical to the executable sources, with + # only endpoint and alias substitutions. text = text.replace("http://127.0.0.1:8000/v1", base_url) text = text.replace('"fasth3"', json.dumps(model)).replace("${FASTVIDEO_MODEL:-fasth3}", "${FASTVIDEO_MODEL:-" + model + "}") clients[language] = {"source": path.relative_to(ROOT_DIR).as_posix(), "code": text, "install": install} + # The playground router only serves H3 servers (see require_h3 in + # fastvideo/entrypoints/openai/playground.py), so other families must not link to it. + has_playground = recipe["family"] == "minimax_h3" + compile_enabled = ((generator.get("engine") or {}).get("compile") or {}).get("enabled") return { "source": serving["source"], "install": serving["install"], @@ -307,7 +311,10 @@ def cookbook_serving_profile(recipe: dict) -> dict: "prepare": serving.get("prepare", ""), "model": model, "base_url": base_url, - "playground_url": f"http://127.0.0.1:{port}/playground/", + "playground_url": f"http://127.0.0.1:{port}/playground/" if has_playground else None, + "audio": bool(serving.get("audio")), + # None when the config leaves compilation at its default. + "compile_enabled": compile_enabled, "health_command": f"curl --fail-with-body http://127.0.0.1:{port}/health", "hardware": hardware, "sampling": config["default_request"]["sampling"], diff --git a/examples/serving/openai_wan22_t2v_a14b.yaml b/examples/serving/openai_wan22_t2v_a14b.yaml index de7e65799..88ce2672f 100644 --- a/examples/serving/openai_wan22_t2v_a14b.yaml +++ b/examples/serving/openai_wan22_t2v_a14b.yaml @@ -25,5 +25,6 @@ default_request: num_frames: 81 height: 720 width: 1280 + fps: 16 output: return_frames: false diff --git a/examples/serving/openai_wan22_ti2v_5b.yaml b/examples/serving/openai_wan22_ti2v_5b.yaml index 1505d1421..e8afbb18d 100644 --- a/examples/serving/openai_wan22_ti2v_5b.yaml +++ b/examples/serving/openai_wan22_ti2v_5b.yaml @@ -1,6 +1,8 @@ # OpenAI-compatible server for Wan2.2 TI2V 5B. Supports both text-to-video # (omit the image reference) and image-to-video (set input_reference or -# image_reference) from the same model. +# image_reference) from the same model. The checked-in cookbook clients send +# text prompts only. The sampling block pins the values the cookbook displays; +# they match the registered preset, and a test keeps them in sync. # fastvideo serve --config examples/serving/openai_wan22_ti2v_5b.yaml generator: model_path: Wan-AI/Wan2.2-TI2V-5B-Diffusers @@ -23,5 +25,10 @@ server: served_model_name: wan22-ti2v-5b default_request: + sampling: + height: 704 + width: 1280 + num_frames: 121 + fps: 24 output: return_frames: false diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index 833c2d869..2903f6b9c 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -1,5 +1,6 @@ """Check serving snippets against their executable sources, without model imports.""" +import ast import copy import json from pathlib import Path @@ -157,6 +158,81 @@ def test_example_docs_exclude_installed_client_dependencies(tmp_path): assert Example(tmp_path).other_files == [tmp_path / "client.mjs"] +@pytest.mark.parametrize( + ("recipe_id", "config_name", "fps", "env_prefix"), + [ + ("fastwan21-t2v", "openai_fastwan21_1_3b.yaml", 16, "FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN "), + ("wan22-t2v", "openai_wan22_t2v_a14b.yaml", 16, ""), + ("wan22-ti2v", "openai_wan22_ti2v_5b.yaml", 24, ""), + ], +) +def test_wan_text_recipes_publish_server_profiles(recipe_id, config_name, fps, env_prefix): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + recipe = next(item for item in recipes if item["id"] == recipe_id) + assert "audio" not in recipe["serving"] + profile = cookbook_serving_profile(recipe) + # The playground router rejects non-H3 servers, so Wan pages must not link to it. + assert profile["playground_url"] is None + assert profile["audio"] is False + assert profile["command"].startswith(f"{env_prefix}fastvideo serve --config examples/serving/{config_name}") + assert "--server.host 127.0.0.1" in profile["command"] + assert profile["sampling"]["fps"] == fps + if recipe_id == "wan22-ti2v": + assert recipe["serving"]["task"] == "Text to video" + assert any("text prompt only" in item for item in recipe["serving"]["limitations"]) + + +def test_h3_serving_recipes_advertise_audio_and_playground(): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + h3 = [item for item in recipes if item["family"] == "minimax_h3" and "serving" in item] + assert h3 + for recipe in h3: + profile = cookbook_serving_profile(recipe) + assert recipe["serving"]["audio"] is True, recipe["id"] + assert profile["audio"] is True, recipe["id"] + assert profile["playground_url"] == "http://127.0.0.1:8000/playground/", recipe["id"] + + +def test_serving_profile_reports_explicit_compile_setting(): + recipes = {item["id"]: item for item in json.loads(COOKBOOK_DATA.read_text())["recipes"]} + assert cookbook_serving_profile(recipes["fasth3-preview-cuda"])["compile_enabled"] is False + assert cookbook_serving_profile(recipes["fastwan21-t2v"])["compile_enabled"] is True + assert cookbook_serving_profile(recipes["wan22-t2v"])["compile_enabled"] is None + + +@pytest.mark.parametrize( + ("config_name", "preset_name"), + [ + ("openai_wan22_t2v_a14b.yaml", "WAN_2_2_T2V_A14B"), + ("openai_wan22_ti2v_5b.yaml", "WAN_2_2_TI2V_5B"), + ], +) +def test_wan_serving_sampling_matches_registered_preset(config_name, preset_name): + # Read the literal defaults from source so the test needs no model imports (torch). + tree = ast.parse((ROOT / "fastvideo/pipelines/basic/wan/presets.py").read_text()) + call = next( + node.value for node in ast.walk(tree) + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call) + and any(isinstance(target, ast.Name) and target.id == preset_name for target in node.targets)) + defaults_node = next(keyword.value for keyword in call.keywords if keyword.arg == "defaults") + defaults = { + key.value: value.value + for key, value in zip(defaults_node.keys, defaults_node.values) + if isinstance(key, ast.Constant) and isinstance(value, ast.Constant) + } + assert defaults + config = yaml.safe_load((ROOT / "examples/serving" / config_name).read_text()) + for key, value in config["default_request"]["sampling"].items(): + if key in defaults: + assert defaults[key] == value, f"{config_name} {key} drifted from {preset_name}" + + +def test_wan21_i2v_stays_python_only(): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + recipe = next(item for item in recipes if item["id"] == "wan21-i2v") + assert "serving" not in recipe + + def test_h3_command_blocks_have_unique_copy_targets(): class CodeBlocks(HTMLParser): def __init__(self): From 6ded84ee0f60cb8081d6fc228671e52c7c71c749 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Sun, 4 Oct 2026 12:35:12 -0700 Subject: [PATCH 049/151] [ci]: improve CI efficiency by reordering lanes (#1914) --- .buildkite/pipeline.yml | 183 ++++++++++++++++++++-------------------- 1 file changed, 93 insertions(+), 90 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 3ec87fb79..2d96a3cf2 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -23,7 +23,100 @@ notify: # dispatcher, and every test payload executes inside the Slinky Slurm tray. # fastvideo/tests/modal remains available only for an explicit manual rollback; # no active pipeline or slash-command route invokes it. +# Buildkite hands jobs to free agents in the order they appear here. Golden-gate comes first +# because every later merge lane waits for it; the fastcheck lanes follow from longest to +# shortest measured runtime, so the longest lane never starts last and stretches the build. steps: + - label: ":test_tube: Golden-Gate Tests" + key: "golden-gate" + if: | + build.env("TEST_SCOPE") == "full" || + (build.env("TEST_SCOPE") == "merge" && + build.env("MERGE_TEST_PLAN") =~ /,golden-gate,/) || + (build.env("TEST_SCOPE") == "direct" && + (build.env("TEST_TYPE") == "golden_gate" || build.env("TEST_TYPE") == "golden_gate_ci")) + command: "/opt/fastvideo-ci-runner/run-ci" + timeout_in_minutes: 90 + env: + TEST_TYPE: "golden_gate_ci" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "ci-runner" + + - label: ":microscope: Unit Tests" + key: "unit" + if: | + build.env("TEST_SCOPE") == "full" || + (build.env("TEST_SCOPE") == "merge" && + build.env("MERGE_TEST_PLAN") =~ /,unit,/) || + build.env("TEST_SCOPE") == "fastcheck" || + build.env("TEST_SCOPE") == null || + (build.env("TEST_SCOPE") == "direct" && + (build.env("TEST_TYPE") == "unit_test" || build.env("TEST_TYPE") == "unit_test_ci")) + command: "/opt/fastvideo-ci-runner/run-unit" + timeout_in_minutes: 90 + env: + TEST_TYPE: "unit_test_ci" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "ci-runner" + + - label: ":microscope: Kernel Tests" + key: "kernel-tests" + if: | + build.env("TEST_SCOPE") == "full" || + (build.env("TEST_SCOPE") == "merge" && + build.env("MERGE_TEST_PLAN") =~ /,kernel-tests,/) || + build.env("TEST_SCOPE") == "fastcheck" || + build.env("TEST_SCOPE") == null || + (build.env("TEST_SCOPE") == "direct" && + (build.env("TEST_TYPE") == "kernel_tests" || build.env("TEST_TYPE") == "kernel_tests_ci")) + command: "/opt/fastvideo-ci-runner/run-ci" + timeout_in_minutes: 90 + env: + TEST_TYPE: "kernel_tests_ci" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "ci-runner" + + - label: ":microscope: DreamVerse App Tests" + key: "dreamverse" + if: | + build.env("TEST_SCOPE") == "full" || + (build.env("TEST_SCOPE") == "merge" && + build.env("MERGE_TEST_PLAN") =~ /,dreamverse,/) || + build.env("TEST_SCOPE") == "fastcheck" || + build.env("TEST_SCOPE") == null || + (build.env("TEST_SCOPE") == "direct" && + (build.env("TEST_TYPE") == "dreamverse_app" || build.env("TEST_TYPE") == "dreamverse_app_ci")) + command: "/opt/fastvideo-ci-runner/run-ci" + timeout_in_minutes: 90 + env: + TEST_TYPE: "dreamverse_app_ci" + retry: + automatic: + - exit_status: 128 + limit: 3 + - exit_status: -1 + limit: 2 + agents: + queue: "ci-runner" + - label: ":microscope: Encoder Tests" key: "encoder" if: | @@ -93,96 +186,6 @@ steps: agents: queue: "ci-runner" - - label: ":microscope: Kernel Tests" - key: "kernel-tests" - if: | - build.env("TEST_SCOPE") == "full" || - (build.env("TEST_SCOPE") == "merge" && - build.env("MERGE_TEST_PLAN") =~ /,kernel-tests,/) || - build.env("TEST_SCOPE") == "fastcheck" || - build.env("TEST_SCOPE") == null || - (build.env("TEST_SCOPE") == "direct" && - (build.env("TEST_TYPE") == "kernel_tests" || build.env("TEST_TYPE") == "kernel_tests_ci")) - command: "/opt/fastvideo-ci-runner/run-ci" - timeout_in_minutes: 90 - env: - TEST_TYPE: "kernel_tests_ci" - retry: - automatic: - - exit_status: 128 - limit: 3 - - exit_status: -1 - limit: 2 - agents: - queue: "ci-runner" - - - label: ":microscope: Unit Tests" - key: "unit" - if: | - build.env("TEST_SCOPE") == "full" || - (build.env("TEST_SCOPE") == "merge" && - build.env("MERGE_TEST_PLAN") =~ /,unit,/) || - build.env("TEST_SCOPE") == "fastcheck" || - build.env("TEST_SCOPE") == null || - (build.env("TEST_SCOPE") == "direct" && - (build.env("TEST_TYPE") == "unit_test" || build.env("TEST_TYPE") == "unit_test_ci")) - command: "/opt/fastvideo-ci-runner/run-unit" - timeout_in_minutes: 90 - env: - TEST_TYPE: "unit_test_ci" - retry: - automatic: - - exit_status: 128 - limit: 3 - - exit_status: -1 - limit: 2 - agents: - queue: "ci-runner" - - - label: ":microscope: DreamVerse App Tests" - key: "dreamverse" - if: | - build.env("TEST_SCOPE") == "full" || - (build.env("TEST_SCOPE") == "merge" && - build.env("MERGE_TEST_PLAN") =~ /,dreamverse,/) || - build.env("TEST_SCOPE") == "fastcheck" || - build.env("TEST_SCOPE") == null || - (build.env("TEST_SCOPE") == "direct" && - (build.env("TEST_TYPE") == "dreamverse_app" || build.env("TEST_TYPE") == "dreamverse_app_ci")) - command: "/opt/fastvideo-ci-runner/run-ci" - timeout_in_minutes: 90 - env: - TEST_TYPE: "dreamverse_app_ci" - retry: - automatic: - - exit_status: 128 - limit: 3 - - exit_status: -1 - limit: 2 - agents: - queue: "ci-runner" - - - label: ":test_tube: Golden-Gate Tests" - key: "golden-gate" - if: | - build.env("TEST_SCOPE") == "full" || - (build.env("TEST_SCOPE") == "merge" && - build.env("MERGE_TEST_PLAN") =~ /,golden-gate,/) || - (build.env("TEST_SCOPE") == "direct" && - (build.env("TEST_TYPE") == "golden_gate" || build.env("TEST_TYPE") == "golden_gate_ci")) - command: "/opt/fastvideo-ci-runner/run-ci" - timeout_in_minutes: 90 - env: - TEST_TYPE: "golden_gate_ci" - retry: - automatic: - - exit_status: 128 - limit: 3 - - exit_status: -1 - limit: 2 - agents: - queue: "ci-runner" - - label: ":bar_chart: SSIM Tests" key: "ssim" depends_on: "golden-gate" From f06d824054034e6b8d54a69286b16ee6b28c94c0 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 6 Oct 2026 10:50:17 -0700 Subject: [PATCH 050/151] [feat]: add Kandinsky-6.0 TI2VA and VSR pipelines (#1924) Co-authored-by: leffff Co-authored-by: Kirill Co-authored-by: shaoxiongduan Co-authored-by: Claude Opus 5.5 --- docs/assets/cookbook-recipes.json | 66 +- docs/assets/logos/SOURCES.md | 2 +- docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/index.md | 23 + docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/kandinsky6.md | 144 +++ docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 8 +- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 2 +- docs/cookbook/z-image.md | 2 +- .../inference_schema_parity_inventory.yaml | 10 + docs/generate_examples.py | 1 + docs/inference/kandinsky6.md | 107 ++ docs/inference/kandinsky6_sr.md | 83 ++ docs/inference/support_matrix.md | 8 + .../inference/basic/basic_kandinsky6_sr.py | 53 + .../inference/basic/basic_kandinsky6_ti2va.py | 51 + fastvideo/api/compat.py | 17 + fastvideo/api/sampling_param.py | 2 +- fastvideo/configs/models/audio/__init__.py | 6 + .../models/audio/kandinsky6_audio_vae.py | 30 + fastvideo/configs/models/dits/__init__.py | 5 +- fastvideo/configs/models/dits/kandinsky6.py | 167 +++ .../configs/models/dits/kandinsky6_sr.py | 139 +++ .../configs/models/upsamplers/__init__.py | 7 +- .../models/upsamplers/kandinsky6_sr.py | 207 ++++ fastvideo/configs/models/vaes/__init__.py | 3 + .../configs/models/vaes/kandinsky6_sr.py | 81 ++ fastvideo/configs/pipelines/__init__.py | 7 +- fastvideo/configs/pipelines/kandinsky6.py | 157 +++ fastvideo/configs/pipelines/kandinsky6_sr.py | 53 + .../pipelines/kandinsky6_sr_options.py | 25 + fastvideo/entrypoints/video_generator.py | 81 +- .../models/audio/kandinsky6_audio_vae.py | 134 +++ fastvideo/models/dits/kandinsky6.py | 1067 +++++++++++++++++ fastvideo/models/dits/kandinsky6_sr.py | 273 +++++ fastvideo/models/encoders/clip.py | 12 + fastvideo/models/loader/component_loader.py | 111 +- fastvideo/models/registry.py | 2 + .../models/schedulers/scheduling_piflow.py | 319 +++++ fastvideo/models/upsamplers/kandinsky6_sr.py | 252 ++++ fastvideo/models/vaes/kandinsky6_sr.py | 461 +++++++ .../pipelines/basic/kandinsky6/__init__.py | 1 + .../kandinsky6/kandinsky6_ti2va_pipeline.py | 112 ++ .../pipelines/basic/kandinsky6/presets.py | 65 + .../pipelines/basic/kandinsky6_sr/__init__.py | 1 + .../kandinsky6_sr/kandinsky6_sr_pipeline.py | 34 + .../pipelines/basic/kandinsky6_sr/presets.py | 57 + .../pipelines/basic/kandinsky6_sr/sr_io.py | 192 +++ .../pipelines/basic/kandinsky6_sr/tiling.py | 146 +++ fastvideo/pipelines/composed_pipeline_base.py | 11 + fastvideo/pipelines/stages/__init__.py | 8 + fastvideo/pipelines/stages/kandinsky6.py | 707 +++++++++++ fastvideo/pipelines/stages/kandinsky6_sr.py | 381 ++++++ fastvideo/pipelines/stages/text_encoding.py | 36 +- fastvideo/registry.py | 114 +- .../tests/api/test_extra_overrides_routing.py | 12 + .../test_video_frame_quantize_oom_fallback.py | 77 ++ .../tests/ssim/test_kandinsky6_similarity.py | 98 ++ .../tests/stages/kandinsky6_sr/conftest.py | 89 ++ .../stages/kandinsky6_sr/k6_sr_release.py | 33 + .../tests/stages/kandinsky6_sr/k6_sr_tiny.py | 211 ++++ .../test_kandinsky6_sr_arch_config.py | 68 ++ .../test_kandinsky6_sr_components.py | 278 +++++ ..._kandinsky6_sr_official_release_configs.py | 159 +++ .../test_kandinsky6_sr_registry_routing.py | 203 ++++ .../test_kandinsky6_sr_stages_pipeline.py | 211 ++++ .../test_kandinsky6_sr_tiling.py | 65 + .../test_kandinsky6_sr_transformer.py | 133 ++ ...est_kandinsky6_sr_video_generator_hooks.py | 157 +++ .../test_kandinsky6_sr_video_io.py | 117 ++ .../tests/stages/kandinsky6_ti2va_tiny.py | 62 + .../stages/test_kandinsky6_arch_config.py | 99 ++ .../stages/test_kandinsky6_audio_decoding.py | 208 ++++ .../test_kandinsky6_audio_latent_duration.py | 45 + .../stages/test_kandinsky6_cfg_resolution.py | 68 ++ .../test_kandinsky6_denoising_precision.py | 215 ++++ .../stages/test_kandinsky6_empty_prompt.py | 36 + .../test_kandinsky6_image_conditioning.py | 205 ++++ .../test_kandinsky6_latent_preparation.py | 123 ++ .../test_kandinsky6_max_sequence_length.py | 187 +++ .../stages/test_kandinsky6_parameter_dtype.py | 86 ++ .../stages/test_kandinsky6_piflow_guidance.py | 60 + .../tests/stages/test_kandinsky6_registry.py | 166 +++ .../stages/test_kandinsky6_target_dtype.py | 96 ++ .../test_kandinsky6_ti2va_transformer.py | 149 +++ fastvideo/utils.py | 8 +- mkdocs.yml | 3 + tests/local_tests/README.md | 2 + tests/local_tests/kandinsky6/README.md | 79 ++ .../test_kandinsky6_transformer_parity.py | 188 +++ tests/local_tests/kandinsky6_sr/README.md | 41 + tests/local_tests/kandinsky6_sr/__init__.py | 1 + .../kandinsky6_sr/lu_parity_common.py | 94 ++ tests/local_tests/kandinsky6_sr/reference.py | 196 +++ .../kandinsky6_sr/sr_pipeline_capture.py | 44 + .../test_kandinsky6_sr_dit_parity.py | 69 ++ ...st_kandinsky6_sr_latent_upscaler_parity.py | 125 ++ .../test_kandinsky6_sr_real_weights.py | 118 ++ .../test_kandinsky6_sr_vae_parity.py | 141 +++ .../test_kandinsky6_sr_pipeline_parity.py | 73 ++ tests/local_tests/test_cookbook_serving.py | 27 + 111 files changed, 10914 insertions(+), 78 deletions(-) create mode 100644 docs/cookbook/kandinsky6.md create mode 100644 docs/inference/kandinsky6.md create mode 100644 docs/inference/kandinsky6_sr.md create mode 100644 examples/inference/basic/basic_kandinsky6_sr.py create mode 100644 examples/inference/basic/basic_kandinsky6_ti2va.py create mode 100644 fastvideo/configs/models/audio/kandinsky6_audio_vae.py create mode 100644 fastvideo/configs/models/dits/kandinsky6.py create mode 100644 fastvideo/configs/models/dits/kandinsky6_sr.py create mode 100644 fastvideo/configs/models/upsamplers/kandinsky6_sr.py create mode 100644 fastvideo/configs/models/vaes/kandinsky6_sr.py create mode 100644 fastvideo/configs/pipelines/kandinsky6.py create mode 100644 fastvideo/configs/pipelines/kandinsky6_sr.py create mode 100644 fastvideo/configs/pipelines/kandinsky6_sr_options.py create mode 100644 fastvideo/models/audio/kandinsky6_audio_vae.py create mode 100644 fastvideo/models/dits/kandinsky6.py create mode 100644 fastvideo/models/dits/kandinsky6_sr.py create mode 100644 fastvideo/models/schedulers/scheduling_piflow.py create mode 100644 fastvideo/models/upsamplers/kandinsky6_sr.py create mode 100644 fastvideo/models/vaes/kandinsky6_sr.py create mode 100644 fastvideo/pipelines/basic/kandinsky6/__init__.py create mode 100644 fastvideo/pipelines/basic/kandinsky6/kandinsky6_ti2va_pipeline.py create mode 100644 fastvideo/pipelines/basic/kandinsky6/presets.py create mode 100644 fastvideo/pipelines/basic/kandinsky6_sr/__init__.py create mode 100644 fastvideo/pipelines/basic/kandinsky6_sr/kandinsky6_sr_pipeline.py create mode 100644 fastvideo/pipelines/basic/kandinsky6_sr/presets.py create mode 100644 fastvideo/pipelines/basic/kandinsky6_sr/sr_io.py create mode 100644 fastvideo/pipelines/basic/kandinsky6_sr/tiling.py create mode 100644 fastvideo/pipelines/stages/kandinsky6.py create mode 100644 fastvideo/pipelines/stages/kandinsky6_sr.py create mode 100644 fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py create mode 100644 fastvideo/tests/ssim/test_kandinsky6_similarity.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/conftest.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/k6_sr_release.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/k6_sr_tiny.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_arch_config.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_components.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_official_release_configs.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_registry_routing.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_stages_pipeline.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_tiling.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_transformer.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_generator_hooks.py create mode 100644 fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_io.py create mode 100644 fastvideo/tests/stages/kandinsky6_ti2va_tiny.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_arch_config.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_audio_decoding.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_audio_latent_duration.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_cfg_resolution.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_denoising_precision.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_empty_prompt.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_image_conditioning.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_latent_preparation.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_max_sequence_length.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_parameter_dtype.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_piflow_guidance.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_registry.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_target_dtype.py create mode 100644 fastvideo/tests/stages/test_kandinsky6_ti2va_transformer.py create mode 100644 tests/local_tests/kandinsky6/README.md create mode 100644 tests/local_tests/kandinsky6/test_kandinsky6_transformer_parity.py create mode 100644 tests/local_tests/kandinsky6_sr/README.md create mode 100644 tests/local_tests/kandinsky6_sr/__init__.py create mode 100644 tests/local_tests/kandinsky6_sr/lu_parity_common.py create mode 100644 tests/local_tests/kandinsky6_sr/reference.py create mode 100644 tests/local_tests/kandinsky6_sr/sr_pipeline_capture.py create mode 100644 tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_dit_parity.py create mode 100644 tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py create mode 100644 tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_real_weights.py create mode 100644 tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py create mode 100644 tests/local_tests/pipelines/test_kandinsky6_sr_pipeline_parity.py diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index e89e58d12..c92c95e22 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 12, + "version": 13, "recipes": [ { "id": "fastwan21-t2v", @@ -332,6 +332,70 @@ "expected_artifact": "MP4 videos under video_samples_kandinsky5_i2v/", "related": ["kandinsky5-t2v-lite-sft"] }, + { + "id": "kandinsky6-ti2va-base", + "family": "kandinsky6", + "stage": "inference", + "task": "Text or image to video with audio", + "label": "Kandinsky 6 Pro TI2VA", + "summary": "Generate a five-second, 512x768 video with synchronized audio from text, or set IMAGE_PATH in the maintained example to condition on an image.", + "model": "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_ti2va.py", + "command": "python examples/inference/basic/basic_kandinsky6_ti2va.py", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "MP4 video with synchronized audio under video_samples_kandinsky6_ti2va/", + "related": ["kandinsky6-ti2va-piflow", "kandinsky6-vsr"] + }, + { + "id": "kandinsky6-ti2va-piflow", + "family": "kandinsky6", + "stage": "inference", + "task": "Distilled text or image to video with audio", + "label": "Kandinsky 6 Pro pi-Flow", + "summary": "Generate a five-second video with synchronized audio using the distilled pi-Flow checkpoint in ten inference steps and guidance 1.0.", + "model": "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_ti2va.py", + "command": "KANDINSKY6_MODEL_PATH=kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers python examples/inference/basic/basic_kandinsky6_ti2va.py", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "MP4 video with synchronized audio under video_samples_kandinsky6_ti2va/", + "related": ["kandinsky6-ti2va-base", "kandinsky6-vsr-distilled"] + }, + { + "id": "kandinsky6-vsr", + "family": "kandinsky6", + "stage": "inference", + "task": "Video super-resolution", + "label": "Kandinsky 6 VSR", + "summary": "Upscale an existing clip by 2.25x with the four-step flow-matching VSR checkpoint while preserving its source audio.", + "model": "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_sr.py", + "command": ": \"${INPUT_VIDEO:?Set INPUT_VIDEO to an existing video}\"\npython examples/inference/basic/basic_kandinsky6_sr.py --video-path \"$INPUT_VIDEO\" --scale 2.25", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "Upscaled MP4 under outputs_video/kandinsky6_sr/", + "related": ["kandinsky6-vsr-distilled", "kandinsky6-ti2va-base"] + }, + { + "id": "kandinsky6-vsr-distilled", + "family": "kandinsky6", + "stage": "inference", + "task": "Distilled video super-resolution", + "label": "Kandinsky 6 VSR distilled", + "summary": "Upscale an existing clip by 2.25x with the distilled two-step VSR checkpoint while preserving its source audio.", + "model": "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_sr.py", + "command": ": \"${INPUT_VIDEO:?Set INPUT_VIDEO to an existing video}\"\npython examples/inference/basic/basic_kandinsky6_sr.py --model-path kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers --video-path \"$INPUT_VIDEO\" --scale 2.25", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "Upscaled MP4 under outputs_video/kandinsky6_sr/", + "related": ["kandinsky6-vsr", "kandinsky6-ti2va-piflow"] + }, { "id": "flux2-klein-t2i", "family": "flux", diff --git a/docs/assets/logos/SOURCES.md b/docs/assets/logos/SOURCES.md index a3edc8a56..21247b395 100644 --- a/docs/assets/logos/SOURCES.md +++ b/docs/assets/logos/SOURCES.md @@ -11,7 +11,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo. | `ltx.webp` | [Official Lightricks Hugging Face organization avatar](https://huggingface.co/Lightricks) | LTX family card and page header | | `tencent-hunyuan.webp` | [Official Tencent Hunyuan Hugging Face organization avatar](https://huggingface.co/Tencent-Hunyuan) | Hunyuan and GameCraft cards, Hunyuan page header | | `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header | -| `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 family card and page header | +| `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers | | `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header | | `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header | | `tongyi.webp` | [Official Tongyi MAI Hugging Face organization avatar](https://huggingface.co/Tongyi-MAI) | Z-Image family card and page header | diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index fbb813fae..89e39e845 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
      +
      All model families
      diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index f2f159aed..82508bd81 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
      +
      All model families
      diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index ed88351a1..51f4a7f28 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
      +
      All model families
      diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index 66acf7f28..75caf3144 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
      +
      All model families
      diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md index be3a1f727..6b4b43497 100644 --- a/docs/cookbook/index.md +++ b/docs/cookbook/index.md @@ -141,6 +141,29 @@ hide: + + + + + + + + + + Kandinsky 6Video with audio, and video SR + 4 recipes + +
        +
      • T2VA
        • +
      • I2VA
        • +
      • VSR
        • +
      +       +       +
      +
      All model families
      diff --git a/docs/cookbook/kandinsky6.md b/docs/cookbook/kandinsky6.md new file mode 100644 index 000000000..3aac661fa --- /dev/null +++ b/docs/cookbook/kandinsky6.md @@ -0,0 +1,144 @@ +--- +hide: +- toc +--- + +# Kandinsky 6 recipes + +
      +
      + All model families +
      + +
      +

      Maintained family · Inference

      +

      Kandinsky 6 inference recipes

      +

      Kandinsky 6 from the Kandinsky Lab generates five-second video with synchronized audio from text or an image, with base and distilled pi-Flow checkpoints, and upscales existing clips with base or distilled video super-resolution.

      +
      +
      +
      + Inference live + Distillation planned + Fine-tuning planned + Training planned + Evaluation planned + Optimization planned + Deployment planned +
      +
      + + +
      +
      +

      Pick a recipe and runtime

      +

      Start with the result you want, then choose one of the runtimes FastVideo actually maintains for it.

      +
      + +
      +
      +
      +
      + Recipe + Task and checkpoint +
      +
      + +
      +
      + +
      +
      + Runtime + Maintained paths only +
      +
      + +
      +
      + +

      Loading recipe details...

      +

      Exact device and memory details appear only when a recorded run supports them.

      + +
      + Reading recipe evidence... +
      +
      + +
      +
      +

      Loading...

      +
      + Maintained + Source-backed + Source config +
      +
      + +
      +
      Model
      Loading...
      +
      Workload
      Loading...
      +
      Source configuration
      Loading...
      +
      Expected output
      Loading...
      +
      + +
      +
      + Terminal +
      + 
      Loading...
      +
      + +
      + Open example source + View model card +
      +

      +
      +
      + + +
      +
      + +
      + Setup +
      +

      The generated commands expect a local clone:

      + 
      git clone https://github.com/hao-ai-lab/FastVideo.git
+cd FastVideo
      +

      Use Configuration for supported Python and CLI settings, Optimizations for attention and memory tradeoffs, and the support matrix for the supported model and optimization surface.

      +

      For model-specific behavior, see Kandinsky 6 video with audio and Kandinsky 6 video super-resolution.

      +
      +
      + +
      + Troubleshooting +
      +
        +
      • Text and image conditioning use the same TI2VA pipeline and example. Set IMAGE_PATH in the script to condition on an image; both base and distilled outputs include synchronized audio.
        • +
      • For pi-Flow, set KANDINSKY6_MODEL_PATH to the distilled model ID when running basic_kandinsky6_ti2va.py. Its registered preset uses 10 inference steps, guidance 1.0, eps=1e-6, final_step_size_scale=0.5, and num_policy_substeps=128.
        • +
      • For VSR, set INPUT_VIDEO in the generated command. The pipeline accepts x2, x2.25 and x4 scales, processes up to 121 frames at 24 fps, and preserves source audio.
        • +
      • Checkpoint key-layout errors indicate that the checkpoint and FastVideo checkout target different Kandinsky 6 Diffusers revisions.
        • +
      • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
        • +
      +
      +
      + +
      + Evidence status +
      +

      All recipes on this page are Source-backed: their commands, model IDs, and flags were validated against the checked-in FastVideo sources listed above (static validation). No runtime GPU validation is recorded for these recipes, so GPU model fit, memory use, throughput, and runtime duration are Unknown and deliberately not claimed. Runtime buttons show only the GPU counts configured in checked-in sources.

      +
      +
      diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index ce2a81574..3b43788d4 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
      +
      All model families
      diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index a9354f30d..53bf0fa14 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
      +
      All model families
      diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index 26160e865..15b62ae34 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
      +
      All model families
      diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 01f77a712..0aabcb034 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -5,13 +5,7 @@ hide: # MiniMax H3 recipes -FastH3 is two distilled MiniMax-H3 checkpoints. **V1** is the four-step -launch. Some Hub repo names still say Preview. That name is historical. V1 is -a full model, not a demo. **V2** is the eight-step checkpoint. More forwards -is why V2 is the higher-quality FastH3. The V2 schedule contract is in -[FastH3 distilled checkpoint schedules](../inference/fasth3-distilled.md). - -
      +
      All model families
      diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 5e127c9b2..a6e5654e7 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
      +
      All model families
      diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index 20dc7bfd7..b2ccbc533 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
      +
      All model families
      diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index 07c87fdb5..653b63677 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index 14820b563..860abee1c 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
      +
      All model families
      diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index c36c32508..9e5eb3e2e 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
      +
      All model families
      diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index b53c79ef1..7f39034dc 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
      +
      All model families
      diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 4309e9851..d45543555 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -458,6 +458,14 @@ surfaces: sources: [fastvideo.pipelines.basic.magi_human.pipeline_configs.MagiHumanBaseConfig] video_txt_guidance_scale: sources: [fastvideo.pipelines.basic.magi_human.pipeline_configs.MagiHumanBaseConfig] + piflow_eps: + sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig] + piflow_final_step_size_scale: + sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig] + piflow_num_policy_substeps: + sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig] + sample_fps: + sources: [fastvideo.configs.pipelines.kandinsky6.Kandinsky6TI2VAConfig] compatibility_only: batch_size: "Gen3C inference-only tuning field pending typed batching design." gradient_checkpointing: "Gen3C inference-only compatibility field pending typed batching design." @@ -481,6 +489,8 @@ surfaces: z_dim: "MagiHuman internal VAE latent channel setting." vocoder_config: "Legacy internal component config object." vocoder_precision: "Precision override pending dedicated component precision design." + audio_sample_rate: "Kandinsky6 audio-latent length constant (audio VAE sample rate); tied to the checkpoint, not a request knob." + audio_downsample_factor: "Kandinsky6 audio-latent length constant (audio VAE downsample factor); tied to the checkpoint, not a request knob." sampling_param_base: moved: diff --git a/docs/generate_examples.py b/docs/generate_examples.py index cf0900b58..bc2d9dd63 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -43,6 +43,7 @@ COOKBOOK_FAMILIES = { "hunyuan", "cosmos", "kandinsky5", + "kandinsky6", "flux", "glm_image", "zimage", diff --git a/docs/inference/kandinsky6.md b/docs/inference/kandinsky6.md new file mode 100644 index 000000000..da4307f7b --- /dev/null +++ b/docs/inference/kandinsky6.md @@ -0,0 +1,107 @@ +# Kandinsky 6 Text/Image to Video with Audio (T2IVA) + +`Kandinsky6TI2VAPipeline` generates a video and a synchronized audio track from a text prompt, optionally conditioned +on an image: one pipeline serves both, so pass `image_path` to condition on an image and leave it out for text only. +The audio is decoded by the checkpoint's audio VAE (mel decoder plus vocoder in one component) and muxed into the saved +mp4 automatically. To upscale a generated clip, see [Kandinsky 6 Video SR](kandinsky6_sr.md). + +## Models + +Both variants are official Diffusers repos, loaded directly through their `model_index.json`: + +| Variant | Hub repo | Scheduler | Steps | Guidance | Example | +|---|---|---|---|---|---| +| T2IVA | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | `FlowMatchEulerDiscreteScheduler` (shift 5.0) | 50 | 5.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| T2IVA distilled | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | `PiflowScheduler` (`n_grid` 10, shift 5.0) | 10 | 1.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | + +The steps and guidance columns are the defaults of the preset the registry selects for each repo id. Everything else is +shared: 512x768, 121 frames (5 s at 24 fps) and the Diffusers default negative prompt (only used when +`guidance_scale > 1`). The distilled preset is named `kandinsky6_ti2va_distilled`. + +## Usage + +```bash +python examples/inference/basic/basic_kandinsky6_ti2va.py +KANDINSKY6_MODEL_PATH=kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers \ + python examples/inference/basic/basic_kandinsky6_ti2va.py +``` + +Set `IMAGE_PATH` in the script to condition on an image. + +```python +from fastvideo import VideoGenerator + +generator = VideoGenerator.from_pretrained( + "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers", + num_gpus=1, + dit_cpu_offload=False, + text_encoder_cpu_offload=True, +) +generator.generate_video( + "cinematic shot: a giant stone samurai on a stormy cliff above a neon city opens glowing golden eyes and " + "raises a katana. Blue lightning strikes the blade, creating a massive shockwave through the clouds. The " + "camera rapidly pulls back from a low angle. Photorealistic, epic scale, dark blue and gold lighting, rain, " + "sparks, volumetric lightning, blockbuster quality. Audio: heavy rain, deep thunder, metallic sword hum, " + "rising brass and choir, electrical crackle, perfectly synchronized lightning impact, sub-bass shockwave. " + "No dialogue, text, or logos.", + image_path=None, # or the path of a conditioning image + output_path="video_samples_kandinsky6_ti2va", + height=512, + width=768, + num_frames=121, +) +``` + +A local copy of a repo works the same way. A local directory is treated as the distilled variant only when its name is +a Kandinsky-6 name containing `distill` (for example `Kandinsky-6.0-Pro-distill-5s-Diffusers`); any other directory name +selects the base preset (see below). + +## Distilled (pi-Flow) checkpoint + +- The distilled repo replaces the flow-matching scheduler with `PiflowScheduler` (`n_grid` 10, `eps` 1e-6, + `final_step_size_scale` 0.5, `num_policy_substeps` 128). Its DiT emits `n_grid` predictions per latent channel + (`out_visual_dim` 160 = 16 x 10, `out_audio_dim` 400 = 40 x 10). +- pi-Flow runs without classifier-free guidance: `guidance_scale` must be exactly `1.0`. The official Diffusers + pipeline rejects any other guidance for a `PiflowScheduler` too; FastVideo raises a `ValueError` naming the + required value. `num_inference_steps` is not constrained by the checkpoint -- `PiflowScheduler.set_timesteps` + accepts any step count and ignores the scheduler's `nfe`; `10` is only the value the distilled checkpoint was + trained for and the `kandinsky6_ti2va_distilled` preset's default. +- The `kandinsky6_ti2va_distilled` preset (10 steps, guidance 1.0) is the default for the distilled repo id and for + local directories named like it. Set `KANDINSKY6_MODEL_PATH` to either one when running the shared + `basic_kandinsky6_ti2va.py` example. A distilled copy under another directory name selects the base preset, so + retain `Kandinsky-6.0-Pro-distill-5s-Diffusers` as the final directory name. +- The policy values can be overridden on `Kandinsky6TI2VAConfig` (`piflow_eps`, `piflow_final_step_size_scale`, + `piflow_num_policy_substeps`); `None` keeps the values from `scheduler_config.json`. + +## Differences from the Diffusers pipeline + +A few Diffusers pipeline options are not ported, and are surfaced here instead of as a knob that would silently do +nothing: + +- `sample_audio=False` (video-only, no audio stream) is not exposed; FastVideo's DiT raises `NotImplementedError` for + a partial-modality call instead of denoising video alone. +- `expand_prompts` (the Qwen prompt-beautifier pass) is not ported; FastVideo's `PromptEnhancerConfig` is a separate, + external (Cerebras/Groq streaming) feature, not this pipeline's built-in expansion. +- Of the Diffusers reference's `visual_cond_scheme` values, only `tail_cond_first_frame` (append one clean reference + frame to the end of the sequence) is implemented; `pretrain` and plain `i2v` are not. +- Precomputed `prompt_embeds`/`negative_prompt_embeds` are not accepted; every call encodes its own prompt text. +- MagCache is not ported: a `magcache` block in a checkpoint's `transformer/config.json` is parsed and dropped + by `update_model_arch`, not read automatically or exposed as an opt-in cache config. +- RNG differs: FastVideo draws video then audio noise from one per-request CPU generator seeded by `seed` (default + 1024); Diffusers seeds a device generator from a value drawn out of `generator` (audio uses `seed+1`). The same + numeric seed produces different noise on the two stacks -- pass `latents`/`audio_latents` directly for bit-level + comparisons. +- Qwen prompt tokens are unpadded and carry no attention mask (Diffusers pads to a fixed length and masks the + padding in text self/cross-attention). Mathematically equivalent for a single prompt (measured 2e-7 relative + difference), but FastVideo has no attention-mask plumbing, so a hand-built batch of unequal-length prompts is not + supported. +- Attention is dense (`LocalAttention`, flash/SDPA); NABLA sparse attention is wired but unverified against a real + NABLA-flagged checkpoint. This matches the Diffusers pipeline itself, which never enables NABLA either. +- VAE tiling is on by default (`vae_tiling=True`); Diffusers decodes untiled. + +## Memory + +The Pro DiT has 30.1B parameters, about 60 GB in bf16 (`dit_precision` defaults to `bf16`), and the Qwen2.5-VL text +encoder adds 16.6 GB. FastVideo enables `dit_cpu_offload` by default; the examples turn it off (`dit_cpu_offload=False`) +to keep the DiT resident on the GPU and offload the text encoder instead (`text_encoder_cpu_offload=True`). See +[Offloading](offloading.md) for the memory knobs. diff --git a/docs/inference/kandinsky6_sr.md b/docs/inference/kandinsky6_sr.md new file mode 100644 index 000000000..c79d2a1cb --- /dev/null +++ b/docs/inference/kandinsky6_sr.md @@ -0,0 +1,83 @@ +# Kandinsky 6 Video Super-Resolution + +`Kandinsky6SRPipeline` upscales a low-resolution video by x2, x2.25 or x4 (video-to-video, no text prompt). The clip is +encoded once with the SR VAE (KVAE); its latent is cut into overlapping tiles, each tile is enlarged by the *latent +upscaler*, refined by a text-free SR DiT in a few denoising steps and decoded, and the tiles are blended back together. +The source audio is kept. The clip can come from anywhere, for example from the [Kandinsky 6 T2IVA](kandinsky6.md) +pipeline. + +## Models + +| Variant | Hub repo | Scheduler | Default steps per tile | +|---|---|---|---| +| Flow matching | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers` | `FlowMatchEulerDiscreteScheduler` (shift 5.0) | 4 | +| Distilled | `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | `PiflowScheduler` (shift 3.5, `n_grid` 10) | 2 | + +Both repos have the same layout and share `vae/` and `latent_upscaler/`: + +```text +model_index.json _class_name = Kandinsky6SRPipeline +transformer/ Kandinsky6SRTransformer3DModel (sr_params: trained base resolution, RoPE scale, noise level) +vae/ Kandinsky6SRVAE +latent_upscaler/ Kandinsky6SRLatentUpscalerBank (x2 and x4 models) +scheduler/ FlowMatchEulerDiscreteScheduler or PiflowScheduler +``` + +The scheduler component drives the denoising loop, and each repo resolves to its own preset (4 or 2 steps). The +distilled transformer's head holds `n_grid` predictions per latent channel, which `PiflowScheduler` integrates. + +## Usage + +```bash +python examples/inference/basic/basic_kandinsky6_sr.py --video-path input.mp4 --scale 2.25 +python examples/inference/basic/basic_kandinsky6_sr.py --video-path input.mp4 \ + --model-path kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers +``` + +```python +from fastvideo import VideoGenerator + +generator = VideoGenerator.from_pretrained("kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", num_gpus=1) +result = generator.generate({ + "inputs": {"video_path": "input.mp4"}, + "output": {"output_path": "outputs_video/sr", "return_frames": False}, + "extensions": {"sr_resolution_scale": 2.25, "sr_target_resolution": "fullhd"}, +}) +``` + +The output geometry and frame rate follow the input clip; the request's `height` / `width` / `num_frames` are not used. +Clips are resampled to 24 fps by a fixed stride and only the first 121 frames (5 s, `1 + 8k`-aligned) are processed. +The source audio (mono, 44.1 kHz) is trimmed to the processed span and muxed into the output. + +## Request parameters + +The `sr_*` options belong only to Kandinsky6 SR. Pass them in `request.extensions` (as above), or under +`request.stage_overrides.sr`. The SR pipeline reads them from `ForwardBatch.extra`; they are not fields of the +shared `SamplingParam` or `ForwardBatch`. Other model families reject these options. Existing +`generate_video(..., sr_resolution_scale=...)` calls remain supported for Kandinsky6 SR; replace direct +`SamplingParam(sr_...=...)` construction with request extensions or these keyword arguments. +For the config-based CLI, use dotted overrides such as `--request.extensions.sr_resolution_scale 4`, rather than +shared `--sr-*` flags. The model-specific example above keeps its `--scale` and `--tiles-batch-size` flags. + +| Field | Default | Meaning | +|---|---|---| +| `num_inference_steps` | 4 / 2 | Denoising steps (DiT calls) per tile. The upstream Diffusers pipeline counts grid points instead (its 5 is 4 steps here). The distilled model was trained for 2; other values run with a warning. | +| `sr_resolution_scale` | `2.25` | Total upscale: `2`, `4` or `2.25` (x1.125 pixel pre-upscale, then x2). | +| `sr_tiles_batch_size` | `1` | Tiles denoised per DiT call (raise it only if memory allows). | +| `sr_tile_min_overlap` | `0.20` | Minimum overlap between neighbouring tiles, as a fraction of the tile. | +| `sr_target_resolution` | `None` | Downscale the result to `hd`, `fullhd`, `2k` or `WxH`. | +| `sr_target_resize_mode` | `fit` | `fit` keeps the aspect ratio, `exact` uses the bucket dimensions. | +| `seed` | `42` | Tile group *k* (of `sr_tiles_batch_size` tiles) is seeded with `seed + index of its first tile`. | + +Two Python-only inputs take raw tensors and are passed as keyword arguments of `generate_video()`: + +- `sr_lr_latent`: an unscaled KVAE latent `[T, C, H, W]` of the source, instead of `video_path` (skips decoding and + encoding the video; scale 2 or 4 only, since 2.25 needs the pixel pre-upscale). +- `sr_audio` / `sr_audio_sample_rate`: a mono waveform in `[-1, 1]` to mux instead of the source's own audio. + +## Limitations + +- NABLA block-sparse attention, which both repos request for 512-pixel tiles, is not wired; the DiT runs dense attention + and logs a warning. +- No sequence or tensor parallelism: the DiT runs on one GPU. Tiles are processed one group after another. +- One clip per request. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 5fb510219..c89be5ce8 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -54,6 +54,9 @@ column links a runnable script in `examples/inference/basic/` where one exists. | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | +| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`
      `kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| kandinsky6_sr | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`
      `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | — | [basic_kandinsky6_sr.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_sr.py) | | lingbot_video | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) | | lingbot_video | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) | | lingbotworld | `FastVideo/LingBot-World-Base-Cam-Diffusers` | I2V | [basic_lingbotworld_base_cam.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld_base_cam.py) | @@ -106,6 +109,11 @@ reference videos as sparse VSA regions) run through [basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py); see [FastH3 distilled checkpoint schedules](fasth3-distilled.md#ref2va-pdd-students). +**Note (Kandinsky 6)**: the two `kandinsky6` IDs generate video with audio from +text, optionally plus an image (see the [T2IVA guide](kandinsky6.md)); the +`kandinsky6_sr` IDs upscale an existing video (see the +[Video SR guide](kandinsky6_sr.md)). + **Note (Wan-VACE)**: not currently supported — no VACE pipeline or registered model ID exists on `main` ([#1435](https://github.com/hao-ai-lab/FastVideo/issues/1435)). The closest diff --git a/examples/inference/basic/basic_kandinsky6_sr.py b/examples/inference/basic/basic_kandinsky6_sr.py new file mode 100644 index 000000000..c1e200041 --- /dev/null +++ b/examples/inference/basic/basic_kandinsky6_sr.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 video super-resolution: upscale a low-resolution video x2, x2.25 or x4. + +The request needs no prompt and no height / width / num_frames: the geometry and frame rate follow the input video +(the first 121 frames, 5 s at 24 fps, are processed and the source audio is kept). Use +``kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers`` for the 2-step distilled model. +See docs/inference/kandinsky6_sr.md. +""" +import argparse + +from fastvideo import VideoGenerator + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", default="kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers") + parser.add_argument("--video-path", required=True, help="Low-resolution source video") + parser.add_argument("--output-path", default="outputs_video/kandinsky6_sr", + help="Output directory (the file is named after the input) or .mp4 path") + parser.add_argument("--scale", type=float, default=2.25, choices=[2.0, 4.0, 2.25], help="Total upscale factor") + parser.add_argument("--num-steps", type=int, default=None, + help="Denoising steps per tile (default: the checkpoint's preset, 4 or 2)") + parser.add_argument("--tiles-batch-size", type=int, default=1, help="Tiles denoised per DiT call") + parser.add_argument("--target-resolution", default=None, help="Optional final size: hd, fullhd, 2k or WxH") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + generator = VideoGenerator.from_pretrained(args.model_path, num_gpus=1) + sampling = {"seed": args.seed} + if args.num_steps is not None: + sampling["num_inference_steps"] = args.num_steps + result = generator.generate({ + "inputs": { + "video_path": args.video_path + }, + "sampling": sampling, + "output": { + "output_path": args.output_path, + "save_video": True, + "return_frames": False, + }, + "extensions": { + "sr_resolution_scale": args.scale, + "sr_tiles_batch_size": args.tiles_batch_size, + "sr_target_resolution": args.target_resolution, + }, + }) + print(f"Wrote {result.video_path}") + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_kandinsky6_ti2va.py b/examples/inference/basic/basic_kandinsky6_ti2va.py new file mode 100644 index 000000000..100d000fc --- /dev/null +++ b/examples/inference/basic/basic_kandinsky6_ti2va.py @@ -0,0 +1,51 @@ +import os + +from fastvideo import VideoGenerator + +DEFAULT_PROMPT = ( + "cinematic shot: a giant stone samurai on a stormy cliff above a neon city opens glowing golden eyes and " + "raises a katana. Blue lightning strikes the blade, creating a massive shockwave through the clouds. The " + "camera rapidly pulls back from a low angle. Photorealistic, epic scale, dark blue and gold lighting, rain, " + "sparks, volumetric lightning, blockbuster quality. Audio: heavy rain, deep thunder, metallic sword hum, " + "rising brass and choir, electrical crackle, perfectly synchronized lightning impact, sub-bass shockwave. " + "No dialogue, text, or logos." +) + +OUTPUT_PATH = "video_samples_kandinsky6_ti2va" + +# The base checkpoint uses 50 steps and guidance 5.0. Set KANDINSKY6_MODEL_PATH to run the +# distilled pi-Flow checkpoint with this same example: +# kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers +# or a local directory named ``Kandinsky-6.0-Pro-distill-5s-Diffusers``. +# That name selects the registered distilled defaults automatically: 10 steps, +# guidance 1.0, eps 1e-6, final_step_size_scale 0.5, and +# num_policy_substeps 128. +MODEL_PATH = os.environ.get("KANDINSKY6_MODEL_PATH", "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers") + +IMAGE_PATH = None # e.g. "assets/girl.png" to condition on an image + + +def main(): + generator = VideoGenerator.from_pretrained( + MODEL_PATH, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + _ = generator.generate_video( + DEFAULT_PROMPT, + image_path=IMAGE_PATH, + output_path=OUTPUT_PATH, + save_video=True, + height=512, + width=768, + num_frames=121, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/compat.py b/fastvideo/api/compat.py index 07d4762cb..56a54e282 100644 --- a/fastvideo/api/compat.py +++ b/fastvideo/api/compat.py @@ -26,6 +26,7 @@ from fastvideo.api.schema import ( SamplingConfig, ) from fastvideo.api.sampling_param import SamplingParam +from fastvideo.configs.pipelines.kandinsky6_sr_options import SR_REQUEST_FIELDS from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.pipelines.basic.ltx2.stage_overrides import ( refine_preset_override_fields, @@ -53,6 +54,8 @@ REQUEST_BATCH_EXTRA_PASSTHROUGH_FIELDS = ( "vsa_mode", "vsa_dense_first_n_steps", "vsa_dense_layers", + # Recognized transport keys; validate_request_batch_extra restricts them to Kandinsky6 SR. + *SR_REQUEST_FIELDS, ) # torch.compile kwargs that map to first-class CompileConfig fields. _COMPILE_TYPED_KEYS = ("backend", "fullgraph", "mode", "dynamic") @@ -383,6 +386,7 @@ def request_to_sampling_param( if request.output.return_state: sampling_param.return_continuation_state = True updates = explicit_request_updates(request) + validate_request_batch_extra(updates, model_path=model_path) for key, value in updates.items(): if hasattr(sampling_param, key): @@ -486,6 +490,18 @@ def request_to_pipeline_overrides(request: GenerationRequest) -> dict[str, Any]: return overrides +def validate_request_batch_extra(updates: Mapping[str, Any], *, model_path: str) -> None: + """Reject Kandinsky6 SR extensions for other model families, including the legacy API.""" + sr_keys = sorted(set(updates).intersection(SR_REQUEST_FIELDS)) + if not sr_keys: + return + from fastvideo.registry import get_preset_selection + + _, model_family = get_preset_selection(model_path) + if model_family != "kandinsky6_sr": + raise ValueError(f"Request fields {sr_keys} are only supported by Kandinsky6 SR, not {model_path}") + + def request_to_batch_extra(request: GenerationRequest) -> dict[str, Any]: """Extract typed-request extensions consumed through ``ForwardBatch.extra``.""" return { @@ -677,6 +693,7 @@ __all__ = [ "normalize_generator_config", "register_continuation_kind", "request_to_batch_extra", + "validate_request_batch_extra", "request_to_pipeline_overrides", "request_to_sampling_param", ] diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 464d7f0fa..8dfb430e4 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -190,7 +190,7 @@ class SamplingParam: def __post_init__(self) -> None: self.data_type = "video" if self.num_frames > 1 else "image" - def check_sampling_param(self): + def check_sampling_param(self) -> None: if self.prompt_path and not self.prompt_path.endswith(".txt"): raise ValueError("prompt_path must be a txt file") diff --git a/fastvideo/configs/models/audio/__init__.py b/fastvideo/configs/models/audio/__init__.py index 69b30005e..f04040ef3 100644 --- a/fastvideo/configs/models/audio/__init__.py +++ b/fastvideo/configs/models/audio/__init__.py @@ -13,6 +13,10 @@ from fastvideo.configs.models.audio.bigvgan import ( BigVGANV2ArchConfig, BigVGANV2Config, ) +from fastvideo.configs.models.audio.kandinsky6_audio_vae import ( + Kandinsky6AudioVAEArchConfig, + Kandinsky6AudioVAEConfig, +) __all__ = [ "LTX2AudioEncoderConfig", @@ -22,4 +26,6 @@ __all__ = [ "MMAudioVAEConfig", "BigVGANV2ArchConfig", "BigVGANV2Config", + "Kandinsky6AudioVAEArchConfig", + "Kandinsky6AudioVAEConfig", ] diff --git a/fastvideo/configs/models/audio/kandinsky6_audio_vae.py b/fastvideo/configs/models/audio/kandinsky6_audio_vae.py new file mode 100644 index 000000000..1395f01b7 --- /dev/null +++ b/fastvideo/configs/models/audio/kandinsky6_audio_vae.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 audio VAE configuration. + +Mirrors the checkpoint's audio_vae/config.json: +{ + "_class_name": "MMAudioVAE", "mode": "44k", "sample_rate": 44100, + "downsample_factor": 1024, "scaling_factor": 0.5302, + "vocoder_config": {...BigVGAN-v2 hyperparams, ignored here -- redundant + backward-compat copy; the pipeline loads a separate "vocoder" component + instead, see fastvideo/models/audio/kandinsky6_audio_vae.py} +} +""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.base import ArchConfig, ModelConfig + + +@dataclass +class Kandinsky6AudioVAEArchConfig(ArchConfig): + architectures: list[str] = field(default_factory=lambda: ["Kandinsky6AudioVAE"]) + mode: str = "44k" + need_vae_decoder: bool = True + need_vae_encoder: bool = False + scaling_factor: float = 1.0 + + +@dataclass +class Kandinsky6AudioVAEConfig(ModelConfig): + arch_config: ArchConfig = field(default_factory=Kandinsky6AudioVAEArchConfig) diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 821cf944a..72b539f13 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -17,13 +17,16 @@ from fastvideo.models.wan.config import WanVideoConfig from fastvideo.configs.models.dits.zimage import ZImageDiTConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig +from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig +from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRArchConfig, Kandinsky6SRConfig from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", - "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", + "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig", + "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig", "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig" ] diff --git a/fastvideo/configs/models/dits/kandinsky6.py b/fastvideo/configs/models/dits/kandinsky6.py new file mode 100644 index 000000000..9f77710c7 --- /dev/null +++ b/fastvideo/configs/models/dits/kandinsky6.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_kandinsky6_transformer_block(n: str, m) -> bool: + return ("text_transformer_blocks" in n or "visual_transformer_blocks" in n) and n.split(".")[-1].isdigit() + + +@dataclass +class Kandinsky6ArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_kandinsky6_transformer_block]) + + # NABLA block-sparse attention for attention_engine="nabla" checkpoints, + # plus the dense backends every DiT supports. Kandinsky6 reuses the same + # NABLA backend as Kandinsky5 (fastvideo/attention/backends/nabla.py) for + # the video self-attention sub-layer only -- audio self-attention and the + # video<->audio cross-attention are always dense. + # + # Note: the diffusers reference's own `engine`/`attention_engine` doesn't + # actually gate NABLA at construction time -- every attention module can + # receive `sparse_params` at *call* time regardless of this value (NABLA + # dispatch is a per-call pipeline decision there, not a fixed model + # property). FastVideo's LocalAttention needs its backend bound at + # construction time, so this port keeps Kandinsky5's construction-time + # gate instead: only build the extra NABLA-backed attention module when + # attention_engine=="nabla". + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.NABLA_ATTN, + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + # Map the current Diffusers checkpoint names to native FastVideo layers. The loader applies only + # the FIRST matching rule: decoder renames must include the FFN rename in the same rule. + # Text-tower attn is self-attention; decoder self/cross-attention names are unchanged. + param_names_mapping: dict = field( + default_factory=lambda: { + r"^(visual_transformer_blocks\.\d+)\.video_dec_block\.feed_forward\.net\.0\.proj\.(weight|bias)$": + r"\1.videoT.feed_forward.mlp.fc_in.\2", + r"^(visual_transformer_blocks\.\d+)\.video_dec_block\.feed_forward\.net\.2\.(weight|bias)$": + r"\1.videoT.feed_forward.mlp.fc_out.\2", + r"^(visual_transformer_blocks\.\d+)\.video_dec_block\.(.*)$": r"\1.videoT.\2", + r"^(visual_transformer_blocks\.\d+)\.audio_dec_block\.feed_forward\.net\.0\.proj\.(weight|bias)$": + r"\1.audioT.feed_forward.mlp.fc_in.\2", + r"^(visual_transformer_blocks\.\d+)\.audio_dec_block\.feed_forward\.net\.2\.(weight|bias)$": + r"\1.audioT.feed_forward.mlp.fc_out.\2", + r"^(visual_transformer_blocks\.\d+)\.audio_dec_block\.(.*)$": r"\1.audioT.\2", + r"^((?:video|audio)_text_transformer_blocks\.\d+)\.attn\.(.*)$": r"\1.self_attention.\2", + r"^((?:video|audio)_text_transformer_blocks\.\d+)\.attn_norm\.(.*)$": r"\1.self_attention_norm.\2", + r"^((?:video|audio)_time_embeddings)\.timestep_embedder\.linear_1\.(weight|bias)$": r"\1.in_layer.\2", + r"^((?:video|audio)_time_embeddings)\.timestep_embedder\.linear_2\.(weight|bias)$": r"\1.out_layer.\2", + r"^(.*feed_forward)\.net\.0\.proj\.(weight|bias)$": r"\1.mlp.fc_in.\2", + r"^(.*feed_forward)\.net\.2\.(weight|bias)$": r"\1.mlp.fc_out.\2", + }) + + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Diffusers Kandinsky6Transformer3DModel config fields (mirror + # transformer/config.json 1:1). + in_visual_dim: int = 16 + out_visual_dim: int = 16 + in_text_dim: int = 3584 + in_text_dim2: int = 768 + time_dim: int = 1024 + patch_size: tuple[int, int, int] = (1, 2, 2) + model_dim: int = 4096 + ff_dim: int = 16384 + num_text_blocks: int = 4 + num_visual_blocks: int = 60 + axes_dims: tuple[int, int, int] = (32, 48, 48) + visual_cond: bool = True + + # RoPE (T, H, W) axis divisor from transformer/config.json. A fixed + # value, not derived from height/width. + scale_factor: tuple[float, float, float] = (1.0, 2.0, 2.0) + + # Always True: the DiT is the joint video+audio model. Kept so the + # transformer/config.json key round-trips; False is rejected below. + is_multimodal: bool = True + out_audio_dim: int | None = None + in_audio_dim: int = 20 + model_dim_a: int | None = None + time_dim_a: int | None = None + ff_dim_a: int | None = None + axes_dims_a: tuple[int, int, int] | None = None + audio_freqs_scaling: float = 1.0 + + # transformer/config.json's "attention_engine" key, kept as a plain string + # (not AttentionBackendEnum) so it round-trips. Only the value "nabla" is + # treated specially, per the note above. + attention_engine: str = "auto" + attention_causal: bool | None = None + attention_local: bool | None = None + attention_glob: bool | None = None + attention_window: int | None = None + attention_P: float | None = None + attention_wT: int | None = None + attention_wW: int | None = None + attention_wH: int | None = None + attention_add_sta: bool | None = None + attention_method: str | None = None + + # Parsed so the transformer/config.json key round-trips, but not read. + # Diffusers pads text tokens to a fixed length and masks the padding in + # attention; this port feeds the DiT the unpadded tokens with no mask + # (as Kandinsky5 does), which is numerically equivalent. + text_token_padding: bool = False + + # Video<->audio fused block knobs (Kandinsky6FusedTransformerDecoderBlock). + ca_rope: bool = False + cross_gates: bool = False + fix_modulation: bool = False + # >0 adds a learned embedding distinguishing generated vs. reference + # visual tokens, consumed by IT2VA's tail_cond_first_frame conditioning. + # Defaults to 2 (generated/reference) rather than diffusers' bare-DiT + # default of 0: FastVideo ships one merged pipeline that serves both + # T2VA (no image -> layer exists but unused) and IT2VA (image -> + # tail_cond_first_frame needs it) from the same checkpoint/config. + visual_token_type_num_embeddings: int = 2 + + def __post_init__(self): + super().__post_init__() + if not self.is_multimodal: + raise ValueError("Kandinsky6 only supports is_multimodal=True: the DiT is always the joint " + "video+audio model.") + head_dim = sum(self.axes_dims) + if self.model_dim % head_dim != 0: + raise ValueError(f"model_dim ({self.model_dim}) must be divisible by head_dim ({head_dim})") + self.hidden_size = self.model_dim + self.num_attention_heads = self.model_dim // head_dim + self.in_channels = self.in_visual_dim + self.out_channels = self.out_visual_dim + self.num_channels_latents = self.in_visual_dim + + # Resolve the audio-tower dims once here (mirrors the `x or default` + # resolution diffusers' Kandinsky6Transformer3DModel.__init__ does + # inline) so the model constructor can read them unconditionally. + self.model_dim_a = self.model_dim_a or self.model_dim + self.time_dim_a = self.time_dim_a or self.time_dim + self.ff_dim_a = self.ff_dim_a or self.ff_dim + self.axes_dims_a = self.axes_dims_a or self.axes_dims + head_dim_a = sum(self.axes_dims_a) + if self.model_dim_a % head_dim_a != 0: + raise ValueError(f"model_dim_a ({self.model_dim_a}) must be divisible by head_dim_a ({head_dim_a})") + + # Kandinsky6FusedTransformerDecoderBlock's cross-modal modulation is + # driven by the *other* modality's time embedding by default + # (fix_modulation=False, matching the diffusers reference): the + # video-conditioning-on-audio modulation is built from time_dim but + # invoked with the audio time embedding, and vice versa. That only + # type-checks when the two time embeddings are the same width, so + # reject the mismatched, fix_modulation=False combination here with a + # clear message instead of a cryptic matmul shape error deep inside + # the fused block. + if not self.fix_modulation and self.time_dim_a != self.time_dim: + raise ValueError(f"time_dim_a ({self.time_dim_a}) must equal time_dim ({self.time_dim}) unless " + "fix_modulation=True (Kandinsky6FusedTransformerDecoderBlock's cross-modal " + "modulation is driven by the other modality's time embedding by default).") + + +@dataclass +class Kandinsky6VideoAudioConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=Kandinsky6ArchConfig) + prefix: str = "Kandinsky6" diff --git a/fastvideo/configs/models/dits/kandinsky6_sr.py b/fastvideo/configs/models/dits/kandinsky6_sr.py new file mode 100644 index 000000000..09297a133 --- /dev/null +++ b/fastvideo/configs/models/dits/kandinsky6_sr.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Arch config of the Kandinsky6 video super-resolution (SR) DiT. + +Fields mirror ``transformer/config.json`` of the official ``Kandinsky6SRTransformer3DModel`` bundles, including the +nested ``sr_params`` (sampling parameters of the SR training setup). ``out_visual_dim`` is the full head width: the +flow-matching checkpoint predicts ``in_visual_dim`` channels, the pi-Flow checkpoint ``n_grid * in_visual_dim`` (``n_grid`` +lives in its ``PiflowScheduler`` config). ``update_model_arch`` drops undeclared keys, so configurations the port does +not implement are rejected by name instead of silently running a different model. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + +_METADATA_KEYS = frozenset({"_class_name", "_diffusers_version", "_name_or_path"}) +# The released checkpoints' sampling setup; the pipeline implements exactly this one. +_SUPPORTED_SR_PARAMS = { + "lq_noise_type": "ddpm", + "lq_channel_noise_scale": 0.0, + "cap_noise_timestep": False, +} + + +def _is_sr_visual_block(name: str, module: Any) -> bool: + return "visual_transformer_blocks" in name and name.split(".")[-1].isdigit() + + +@dataclass +class Kandinsky6SRArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_sr_visual_block]) + # NABLA sparse attention requested by ``attention_params`` is not wired; the DiT runs dense attention. + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + # Map the current Diffusers checkpoint names to native FastVideo layers. The loader applies only + # the FIRST matching rule: decoder renames must include the FFN rename in the same rule. + # Text-tower attn is self-attention; decoder self/cross-attention names are unchanged. + param_names_mapping: dict = field( + default_factory=lambda: { + r"^(time_embeddings)\.timestep_embedder\.linear_1\.(weight|bias)$": r"\1.in_layer.\2", + r"^(time_embeddings)\.timestep_embedder\.linear_2\.(weight|bias)$": r"\1.out_layer.\2", + r"^(.*feed_forward)\.net\.0\.proj\.(weight|bias)$": r"\1.mlp.fc_in.\2", + r"^(.*feed_forward)\.net\.2\.(weight|bias)$": r"\1.mlp.fc_out.\2", + }) + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + in_visual_dim: int = 64 + out_visual_dim: int = 64 + time_dim: int = 512 + patch_size: tuple[int, int, int] = (1, 1, 1) + model_dim: int = 1792 + ff_dim: int = 7168 + num_visual_blocks: int = 32 + axes_dims: tuple[int, int, int] = (16, 24, 24) + visual_cond: bool = True + instruct_type: str = "hybrid_anchor" + attention_params: dict | None = None + # Text-tower arguments of the shared constructor; the SR DiT is text-free. + use_text: bool = False + in_text_dim: int = 3584 + in_text_dim2: int = 768 + num_text_blocks: int = 2 + attribute_overrides: dict | None = None + sr_params: dict = field(default_factory=dict) + # New top-level fields the diffusers refactor registers on the checkpoint config (replacing the old + # attention_params/sr_params-nested scheme); the DiT does not read these yet (still dense-only), but they + # must be known so loading a checkpoint saved from the current diffusers code doesn't hard-fail on an + # "unsupported key". Default None/unset so an old-schema checkpoint (no such key in its config.json) keeps + # today's behavior exactly -- unlike diffusers' own constructor defaults, which would apply to every + # checkpoint that omits these keys, including ones never trained with NABLA. + scale_factor: tuple[float, float, float] | None = None + nabla_threshold: float | None = None + nabla_window: tuple[int, int, int] | None = None + tile_sizes: tuple[tuple[int, int], ...] | None = None + + def __post_init__(self) -> None: + super().__post_init__() + self.patch_size = tuple(int(p) for p in self.patch_size) # type: ignore[assignment] + self.axes_dims = tuple(int(a) for a in self.axes_dims) # type: ignore[assignment] + unsupported = { + "use_text": self.use_text, + "attribute_overrides": self.attribute_overrides, + } + for name, value in unsupported.items(): + if value: + raise NotImplementedError(f"Kandinsky6 SR: {name}={value!r} is not supported.") + if self.instruct_type != "hybrid_anchor" or not self.visual_cond: + raise NotImplementedError("Kandinsky6 SR supports instruct_type='hybrid_anchor' with visual_cond=True, got " + f"{self.instruct_type!r} / visual_cond={self.visual_cond}.") + if self.patch_size[0] != 1: + raise ValueError(f"Kandinsky6 SR requires patch_size[0] == 1, got {self.patch_size}") + head_dim = sum(self.axes_dims) + if self.model_dim % head_dim: + raise ValueError(f"model_dim ({self.model_dim}) must be divisible by head_dim ({head_dim})") + if self.sr_params: + for key, expected in _SUPPORTED_SR_PARAMS.items(): + if self.sr_params.get(key, expected) != expected: + raise NotImplementedError(f"Kandinsky6 SR: sr_params.{key}={self.sr_params[key]!r} is not " + f"supported (expected {expected!r}).") + + self.hidden_size = self.model_dim + self.num_attention_heads = self.model_dim // head_dim + self.in_channels = self.in_visual_dim + self.out_channels = self.out_visual_dim + self.num_channels_latents = self.in_visual_dim + + @property + def input_channels(self) -> int: + """``[noised latent | anchor | anchor mask]``.""" + return 2 * self.in_visual_dim + 1 + + @property + def visual_size(self) -> int: + size = self.sr_params.get("visual_size", 512) + return int(size[0] if isinstance(size, list | tuple) else size) + + @property + def rope_scale_factor(self) -> tuple[float, ...]: + """RoPE frequency scale per (t, h, w) axis for :attr:`visual_size`.""" + factors = self.sr_params.get("scale_factor", {}) + values = factors.get(str(self.visual_size), factors.get(self.visual_size, (1.0, 1.0, 1.0))) + return tuple(float(v) for v in values) + + @property + def lq_noise_scale(self) -> float: + return float(self.sr_params.get("lq_noise_scale", 0.7)) + + def known_config_keys(self) -> frozenset[str]: + return frozenset(self.__dataclass_fields__) | _METADATA_KEYS + + +@dataclass +class Kandinsky6SRConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=Kandinsky6SRArchConfig) + prefix: str = "Kandinsky6SR" diff --git a/fastvideo/configs/models/upsamplers/__init__.py b/fastvideo/configs/models/upsamplers/__init__.py index d4db478dd..1defc436f 100644 --- a/fastvideo/configs/models/upsamplers/__init__.py +++ b/fastvideo/configs/models/upsamplers/__init__.py @@ -1,4 +1,9 @@ from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig +from fastvideo.configs.models.upsamplers.kandinsky6_sr import (Kandinsky6SRLatentUpscalerConfig, + Kandinsky6SRLatentUpscalerEntryConfig) from fastvideo.configs.models.upsamplers.base import UpsamplerConfig -__all__ = ["SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"] +__all__ = [ + "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "Kandinsky6SRLatentUpscalerConfig", + "Kandinsky6SRLatentUpscalerEntryConfig", "UpsamplerConfig" +] diff --git a/fastvideo/configs/models/upsamplers/kandinsky6_sr.py b/fastvideo/configs/models/upsamplers/kandinsky6_sr.py new file mode 100644 index 000000000..dd570e7b2 --- /dev/null +++ b/fastvideo/configs/models/upsamplers/kandinsky6_sr.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config of the Kandinsky6 SR latent-upscaler bank (``latent_upscaler/config.json``). + +The component config is ``{"models": [{"target_scale": "4x" | "2x", "model": {...}}, ...], "scaling_factor": f}``. +Each ``model`` mapping describes one cascaded 2x+2x upsampler. Only the architecture of the released checkpoints is +implemented, so every key that would select a different architecture must carry its released value; a value that is +not supported raises a ``ValueError`` naming the key instead of silently building something else. +""" +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import Any + +from fastvideo.configs.models.upsamplers.base import UpsamplerConfig + +# key -> (required value, value a config that omits the key stands for). The omitted-key values are the training +# code's defaults, so e.g. a config without ``input_skip`` means ``input_skip=True`` and is rejected. +_FIXED_KEYS: dict[str, tuple[Any, Any]] = { + "upscale_factor": (4, 4), + "dims": (3, 2), + "temporal_padding": ("replicate", "zeros"), + "upsample_mode": ("pxs_v2", "pixel_shuffle"), + "upsample_padding_mode": ("zeros", "reflect"), + "modulated_norm": (True, False), + "modulated_output_proj": (True, False), + "bare_stem": (True, False), + "input_skip": (False, True), + "global_skip": (False, False), + "grn": (False, False), + "layer_scale_init": (None, None), + "depthwise": (False, False), + "bottleneck_channels": (None, None), + "motion_attention": (None, None), + # Only read by depthwise blocks; the residual convs are always 3x3x3. + "kernel_size": (3, 3), +} +# Fixed for entries built with the x2 entry (``enable_x2_entry=true``). +_X2_FIXED_KEYS: dict[str, tuple[Any, Any]] = { + "x2_tail_mode": ("private_full", "shared"), + "x2_finisher": ("pxs_residual", "none"), +} +# Accepted without effect at inference: training-time settings, and switches that only apply to the pixel-shuffle +# upsample modes (the pxs_v2 upsample has no temporal kernel and no ICNR init). +_IGNORED_KEYS = frozenset({ + "gradient_checkpointing", "stochastic_depth_rate", "loss_weight_2x", "loss_weight_4x", "temporal_mix", "icnr", + "upsample_position" +}) +_REQUIRED_INT_KEYS = ("in_channels", "hidden_channels", "num_pre_blocks", "num_mid_blocks", "num_post_blocks", + "expand_ratio") +_SUPPORTED_SCALES = (2, 4) + + +def _positive_int(where: str, key: str, value: Any) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{where}: `{key}` must be a positive integer, got {value!r}") + return value + + +def _parse_target_scale(where: str, value: Any) -> int: + text = str(value).removesuffix("x") + if isinstance(value, bool) or not text.isdigit() or int(text) not in _SUPPORTED_SCALES: + raise ValueError(f"{where}: `target_scale` must be one of '2x', '4x', got {value!r}") + return int(text) + + +@dataclass(frozen=True) +class Kandinsky6SRLatentUpscalerEntryConfig: + """One bank entry: a cascaded 2x+2x upsampler serving ``target_scale``. + + ``stage_channels`` are the widths at the 1x, 2x and 4x latent grids. With ``enable_x2_entry`` the entry also has + a private x2 path (its own input stem, ``x2_adapter_blocks`` residual blocks at 1x, and private copies of the mid + stage and second stage) that upsamples by 2 instead of 4. + """ + + target_scale: int + in_channels: int + hidden_channels: int + stage_channels: tuple[int, int, int] + num_pre_blocks: int + num_mid_blocks: int + num_post_blocks: int + expand_ratio: int + enable_x2_entry: bool = False + x2_adapter_blocks: int = 0 + + @classmethod + def from_dict(cls, spec: Any, index: int = 0) -> Kandinsky6SRLatentUpscalerEntryConfig: + """Validate one ``models[index]`` entry of ``latent_upscaler/config.json``.""" + where = f"latent upscaler models[{index}]" + if not isinstance(spec, Mapping) or set(spec) != {"target_scale", "model" + } or not isinstance(spec["model"], Mapping): + raise ValueError(f"{where} must be a mapping with exactly `target_scale` and a `model` mapping, " + f"got {spec!r}") + target_scale = _parse_target_scale(where, spec["target_scale"]) + model = dict(spec["model"]) + + if model.get("architecture") != "multi_scale": + raise ValueError(f"{where}: `architecture` must be 'multi_scale', got {model.get('architecture')!r}") + enable_x2_entry = model.get("enable_x2_entry", False) + if not isinstance(enable_x2_entry, bool): + raise ValueError(f"{where}: `enable_x2_entry` must be a boolean, got {enable_x2_entry!r}") + + fixed = dict(_FIXED_KEYS) + if enable_x2_entry: + fixed.update(_X2_FIXED_KEYS) + for key, (required, omitted) in fixed.items(): + value = model.get(key, omitted) + if value != required or type(value) is not type(required): + raise ValueError(f"{where}: unsupported `{key}`={value!r}" + f"{' (omitted)' if key not in model else ''}; only {required!r} is implemented") + + x2_defaults = { + "x2_adapter_blocks": 0, + "x2_adapter_sources": None, + "x2_tail_mode": "shared", + "x2_finisher": "none" + } + known = { + "architecture", "enable_x2_entry", "stage_channels", *_REQUIRED_INT_KEYS, *fixed, *x2_defaults, + *_IGNORED_KEYS + } + unknown = sorted(set(model) - known) + if unknown: + raise ValueError(f"{where}: unknown keys {unknown}") + if not enable_x2_entry: + present = sorted(key for key, default in x2_defaults.items() if model.get(key, default) != default) + if present: + raise ValueError(f"{where}: {present} require `enable_x2_entry`=true") + if target_scale == 2 and not enable_x2_entry: + raise ValueError(f"{where}: `target_scale`='2x' requires `enable_x2_entry`=true") + + missing = [key for key in _REQUIRED_INT_KEYS if key not in model] + if missing: + raise ValueError(f"{where}: missing keys {missing}") + ints = {key: _positive_int(where, key, model[key]) for key in _REQUIRED_INT_KEYS} + hidden = ints["hidden_channels"] + + raw_stages = model.get("stage_channels") + if raw_stages is None: + stage_channels = (hidden, hidden, hidden) + else: + if not isinstance(raw_stages, list | tuple) or len(raw_stages) != 3: + raise ValueError(f"{where}: `stage_channels` must list 3 widths, got {raw_stages!r}") + first, second, third = (_positive_int(where, "stage_channels", width) for width in raw_stages) + stage_channels = (first, second, third) + if stage_channels[0] != hidden: + raise ValueError(f"{where}: `hidden_channels` ({hidden}) must equal `stage_channels[0]` " + f"({stage_channels[0]})") + + x2_adapter_blocks = 0 + if enable_x2_entry: + x2_adapter_blocks = _positive_int(where, "x2_adapter_blocks", model.get("x2_adapter_blocks", 0)) + # Training warm-starts the adapter from these pre_blocks; only their consistency matters here. + sources = model.get("x2_adapter_sources") + if sources is not None and (not isinstance(sources, list | tuple) or len(sources) != x2_adapter_blocks + or any(not isinstance(i, int) or not 0 <= i < ints["num_pre_blocks"] + for i in sources)): + raise ValueError(f"{where}: `x2_adapter_sources` must list {x2_adapter_blocks} indices in " + f"[0, num_pre_blocks={ints['num_pre_blocks']}), got {sources!r}") + + return cls(target_scale=target_scale, + stage_channels=stage_channels, + enable_x2_entry=enable_x2_entry, + x2_adapter_blocks=x2_adapter_blocks, + **ints) + + +def _as_entry(entry: Any, index: int) -> Kandinsky6SRLatentUpscalerEntryConfig: + if isinstance(entry, Kandinsky6SRLatentUpscalerEntryConfig): + return entry + # ``PipelineConfig.dump_to_json`` writes entries flattened by ``dataclasses.asdict``; read them back as such. + if isinstance(entry, Mapping) and "model" not in entry and set(entry) == { + item.name + for item in fields(Kandinsky6SRLatentUpscalerEntryConfig) + }: + return Kandinsky6SRLatentUpscalerEntryConfig(**{**entry, "stage_channels": tuple(entry["stage_channels"])}) + return Kandinsky6SRLatentUpscalerEntryConfig.from_dict(entry, index) + + +@dataclass +class Kandinsky6SRLatentUpscalerConfig(UpsamplerConfig): + """``latent_upscaler/config.json``: one entry per served scale, plus the VAE latent ``scaling_factor``. + + ``models`` accepts the raw config.json entries and is normalised to + :class:`Kandinsky6SRLatentUpscalerEntryConfig` in ``__post_init__`` (also re-run by ``update_model_config``). + """ + + models: list[Kandinsky6SRLatentUpscalerEntryConfig] = field(default_factory=list) + scaling_factor: float = 1.0 + # Index order used by re-keyed ModuleList weights. The released re-keying + # script uses (2, 4) when the legacy models config has no scales field. + scales: tuple[int, ...] = (2, 4) + + def __post_init__(self) -> None: + if not isinstance(self.models, list | tuple): + raise ValueError(f"latent upscaler `models` must be a list of entries, got {self.models!r}") + entries = [_as_entry(entry, index) for index, entry in enumerate(self.models)] + scales = [entry.target_scale for entry in entries] + if len(set(scales)) != len(scales): + raise ValueError(f"latent upscaler `models` has duplicate target scales: {scales}") + self.models = entries + self.scales = tuple(self.scales) + if (not self.scales or len(set(self.scales)) != len(self.scales) + or any(type(scale) is not int or scale not in _SUPPORTED_SCALES for scale in self.scales)): + raise ValueError(f"latent upscaler `scales` must contain unique values from (2, 4), got {self.scales}") + self.scaling_factor = float(self.scaling_factor) diff --git a/fastvideo/configs/models/vaes/__init__.py b/fastvideo/configs/models/vaes/__init__.py index e341a058c..13210e3a2 100644 --- a/fastvideo/configs/models/vaes/__init__.py +++ b/fastvideo/configs/models/vaes/__init__.py @@ -5,6 +5,7 @@ from fastvideo.configs.models.vaes.gen3cvae import Gen3CVAEConfig from fastvideo.configs.models.vaes.glm_image import GlmImageVAEConfig from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig +from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEArchConfig, Kandinsky6SRVAEConfig from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig, MiniMaxH3AudioVAEConfig from fastvideo.configs.models.vaes.minimax_h3_video import ( @@ -23,6 +24,8 @@ __all__ = [ "Cosmos25VAEConfig", "Gen3CVAEConfig", "Hunyuan15VAEConfig", + "Kandinsky6SRVAEArchConfig", + "Kandinsky6SRVAEConfig", "LTX2VAEConfig", "MiniMaxH3AudioVAEArchConfig", "MiniMaxH3AudioVAEConfig", diff --git a/fastvideo/configs/models/vaes/kandinsky6_sr.py b/fastvideo/configs/models/vaes/kandinsky6_sr.py new file mode 100644 index 000000000..7a84a3716 --- /dev/null +++ b/fastvideo/configs/models/vaes/kandinsky6_sr.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Config of the Kandinsky6 SR causal video KVAE (``Kandinsky6SRVAE``). + +``vae/config.json`` comes in two layouts: an older one nesting the KVAE architecture under ``encoder_config`` / +``decoder_config`` dicts (the config stored the string ``"None"`` for their unused channel fields), and the current +one with those same knobs flattened to the top level instead -- ``decoder_ch`` / ``decoder_ch_mult`` are the only decoder-specific overrides, everything else +(``num_res_blocks``, ``z_channels``, the temporal-compression / norm / padding knobs) is shared between encoder and +decoder. ``update_model_arch`` copies whichever layout the checkpoint has; ``__post_init__`` synthesizes +``encoder_config`` / ``decoder_config`` from the flat fields when the nested ones weren't populated, so +``Kandinsky6SRVAE.__init__`` (and its ``_arch_kwargs`` validation) only ever has to handle the nested shape. +``scaling_factor`` multiplies the raw KVAE latent to get the latent space the SR DiT / latent upscaler work in. +""" +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class Kandinsky6SRVAEArchConfig(VAEArchConfig): + vae_type: str = "video-kvae" + encoder_config: dict[str, Any] = field(default_factory=dict) + decoder_config: dict[str, Any] = field(default_factory=dict) + scaling_factor: float = 1.0 + spatial_factor: int = 16 + temporal_factor: int = 4 + + # Flat top-level KVAE architecture fields (current official vae/config.json layout -- see module + # docstring). Left at their None default (and ignored) when the checkpoint instead uses the older + # nested encoder_config/decoder_config layout. + in_channels: int | None = None + out_channels: int | None = None + z_channels: int | None = None + ch: int | None = None + ch_mult: list[float] | None = None + decoder_ch: int | None = None + decoder_ch_mult: list[float] | None = None + num_res_blocks: int | None = None + padding_mode: str | None = None + temporal_compress_times: int | None = None + temporal_compress_start_level: int | None = None + norm_type: str | None = None + double_z: bool | None = None + downsample_version: int | None = None + fix_pxs: bool | None = None + + def __post_init__(self) -> None: + if self.vae_type != "video-kvae": + raise ValueError(f"Kandinsky6SRVAE supports only 'video-kvae', got {self.vae_type!r}") + self.spatial_compression_ratio = int(self.spatial_factor) + self.temporal_compression_ratio = int(self.temporal_factor) + + if not self.encoder_config and not self.decoder_config and self.ch is not None: + shared = dict(num_res_blocks=self.num_res_blocks, + z_channels=self.z_channels, + temporal_compress_times=self.temporal_compress_times, + temporal_compress_start_level=self.temporal_compress_start_level, + norm_type=self.norm_type, + padding_mode=self.padding_mode) + self.encoder_config = { + **shared, "ch": self.ch, + "ch_mult": self.ch_mult, + "in_channels": self.in_channels, + "double_z": self.double_z, + "downsample_version": self.downsample_version, + "fix_pxs": self.fix_pxs + } + self.decoder_config = { + **shared, "ch": self.decoder_ch, + "ch_mult": self.decoder_ch_mult, + "out_ch": self.out_channels + } + + +@dataclass +class Kandinsky6SRVAEConfig(VAEConfig): + arch_config: VAEArchConfig = field(default_factory=Kandinsky6SRVAEArchConfig) + # The KVAE has its own segment-cached encode / decode; the generic tiled-VAE knobs do not apply. + use_tiling: bool = False + use_temporal_tiling: bool = False + use_parallel_tiling: bool = False diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index fb97bbeb1..34461ebf6 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -7,6 +7,8 @@ from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyua from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig +from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig +from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig @@ -22,6 +24,7 @@ __all__ = [ "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", - "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", - "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name" + "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig", "Kandinsky6SRPipelineConfig", + "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", + "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/kandinsky6.py b/fastvideo/configs/pipelines/kandinsky6.py new file mode 100644 index 000000000..ddd61370c --- /dev/null +++ b/fastvideo/configs/pipelines/kandinsky6.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, ModelConfig, VAEConfig +from fastvideo.configs.models.audio import BigVGANV2Config, Kandinsky6AudioVAEConfig +from fastvideo.configs.models.dits import Kandinsky6VideoAudioConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput, CLIPTextConfig +from fastvideo.configs.models.encoders.reason1 import Reason1Config +from fastvideo.configs.models.vaes import HunyuanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text + +# Same Qwen2.5-VL "prompt engineer" system template and 129-token crop as +# Kandinsky5 (fastvideo/configs/pipelines/kandinsky5.py) -- Kandinsky6 reuses +# the identical text-conditioning stack (Qwen2.5-VL token embeddings + CLIP +# pooled embedding), byte-exact including its two misspelled words: the +# checkpoints were trained with this exact system prompt, and +# ENCODE_START_IDX is the tokenized length of everything before the user +# prompt. Kept as an independent copy (not imported from kandinsky5.py) so +# the two model families' pipeline configs stay self-contained. +KANDINSKY6_PROMPT_TEMPLATE = "\n".join([ + "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", # codespell:ignore promt + "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.", + "Describe the location of the video, main characters or objects and their action.", + "Describe the dynamism of the video and presented actions.", + "Name the visual style of the video: whether it is a professional footage, user generated content, some kind of animation, video game or scren content.", # codespell:ignore scren + "Describe the visual effects, postprocessing and transitions if they are presented in the video.", + "Pay attention to the order of key actions shown in the scene.<|im_end|>", + "<|im_start|>user\n{}<|im_end|>", +]) +KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX = 129 + + +def kandinsky6_qwen_preprocess_text(prompt: str) -> str: + # The diffusers reference encodes an empty prompt as an empty user turn + # (the template's system half is never empty, so the Qwen tokenizer + # never sees zero tokens) -- unlike the generic empty-prompt guard + # (`text_encoding.py`'s `treat_empty_as_dot`, not opted into here), + # substituting "." would tokenize differently and drift from the + # reference's embeddings for prompt="". + return KANDINSKY6_PROMPT_TEMPLATE.format(prompt) + + +def kandinsky6_qwen_postprocess_text(outputs: BaseEncoderOutput, + mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if outputs.hidden_states is None: + raise RuntimeError("Kandinsky6 Qwen prompt embeddings require hidden_states.") + hidden_states = outputs.hidden_states[-1] + prompt_embeds = hidden_states[:, KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX:] + mask = mask[:, KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX:] + if prompt_embeds.shape[1] == 0: + prompt_embeds = hidden_states[:, -1:] + mask = torch.ones((mask.shape[0], 1), dtype=mask.dtype, device=mask.device) + return prompt_embeds, mask + + +def kandinsky6_clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + if outputs.pooler_output is None: + raise RuntimeError("Kandinsky6 CLIP pooled output is required.") + return outputs.pooler_output + + +@dataclass +class Kandinsky6TI2VAConfig(PipelineConfig): + """Kandinsky6 TI2VA pipeline configuration. + + One pipeline (fastvideo.pipelines.basic.kandinsky6.Kandinsky6TI2VAPipeline) + serves text-to-video+audio generation with an optional conditioning image: + pure text when none is supplied at generation time, image+text (called + "I2VA" in the diffusers reference) when one is -- mirroring the diffusers + reference, which likewise has a single Kandinsky6TI2VAPipeline class + (pipeline_kandinsky6_ti2va.py) taking an optional image argument, not a + separate T2VA/I2VA subclass pair. See + fastvideo/pipelines/stages/kandinsky6.py's Kandinsky6ImageEncodingStage. + """ + + dit_config: DiTConfig = field(default_factory=Kandinsky6VideoAudioConfig) + vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig) + # Two separate checkpoint components, like LTX-2's audio_vae/vocoder pair: audio_vae_config (the + # mel-VAE decoder, a thin nesting wrapper around FastVideo's existing MMAudioVAE) and + # vocoder_config (BigVGAN-v2, reused unmodified from the existing mmaudio pipeline). + audio_vae_config: ModelConfig = field(default_factory=Kandinsky6AudioVAEConfig) + vocoder_config: ModelConfig = field(default_factory=BigVGANV2Config) + + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Reason1Config(), CLIPTextConfig())) + preprocess_text_funcs: tuple[Callable[[str], Any], ...] = field( + default_factory=lambda: (kandinsky6_qwen_preprocess_text, preprocess_text)) + postprocess_text_funcs: tuple[Callable[..., Any], ...] = field( + default_factory=lambda: (kandinsky6_qwen_postprocess_text, kandinsky6_clip_postprocess_text)) + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", "bf16")) + # Diffusers default: 1024 user tokens for Qwen (max_length = msl + 129, + # the template-prefix crop), 77 for CLIP (fixed, never overridden by a + # request -- see Kandinsky6TextEncodingStage._resolve_max_length). + text_encoder_max_lengths: tuple[int, ...] = field( + default_factory=lambda: (KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX + 1024, 77)) + + # Matches both official scheduler_config.json files (base and distilled + # alike): shift 5.0; SchedulerLoader also applies this to a checkpoint's + # own scheduler_config.json via set_shift(), so it is only a fallback for + # a config missing the key. + flow_shift: float | None = 5.0 + vae_tiling: bool = True + + # Scheduler constructor overrides; None keeps scheduler_config.json values. + piflow_eps: float | None = None + piflow_final_step_size_scale: float | None = None + piflow_num_policy_substeps: int | None = None + + # Audio<->video latent-length alignment, matching the diffusers + # reference's Kandinsky6TI2VAPipeline defaults exactly: + # audio_latent_frames = ceil(((T_lat-1)*4+1) / sample_fps * audio_sample_rate / audio_downsample_factor). + sample_fps: float = 24.0 + audio_sample_rate: int = 44100 + audio_downsample_factor: int = 1024 + + def __post_init__(self) -> None: + if len(self.text_encoder_configs) != 2: + raise ValueError(f"Kandinsky6 pipeline requires exactly 2 text encoders (qwen and clip), " + f"but got {len(self.text_encoder_configs)} encoder(s).") + if len(self.text_encoder_precisions) != 2: + raise ValueError("Kandinsky6 pipeline requires exactly 2 text encoder precisions, " + f"but got {len(self.text_encoder_precisions)}.") + if len(self.text_encoder_max_lengths) != 2: + raise ValueError("Kandinsky6 pipeline requires exactly 2 text encoder max lengths, " + f"but got {len(self.text_encoder_max_lengths)}.") + + # The merged pipeline can receive an optional conditioning image on + # any given call, so the VAE encoder must always be available (unlike + # Kandinsky5, which only flips this on for the dedicated I2V config). + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + + qwen_cfg = self.text_encoder_configs[0] + qwen_cfg.arch_config.output_hidden_states = True + qwen_cfg.arch_config.tokenizer_kwargs.update({ + "padding": True, + "truncation": True, + "return_tensors": "pt", + }) + + clip_cfg = self.text_encoder_configs[1] + clip_cfg.arch_config.tokenizer_kwargs.update({ + "padding": "max_length", + "max_length": 77, + "truncation": True, + "add_special_tokens": True, + "return_tensors": "pt", + }) diff --git a/fastvideo/configs/pipelines/kandinsky6_sr.py b/fastvideo/configs/pipelines/kandinsky6_sr.py new file mode 100644 index 000000000..c7cbafa1c --- /dev/null +++ b/fastvideo/configs/pipelines/kandinsky6_sr.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline config of the Kandinsky6 video super-resolution (VSR) pipeline. + +Video -> video with no text encoder (the SR DiT is text-free), so the text-encoder lists are empty. The bundle's +``scheduler`` component drives the denoising loop and the transformer config's ``sr_params`` hold the noise level and +RoPE scale; per-request knobs live in ``Kandinsky6SROptions`` and travel through ``ForwardBatch.extra``. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import ClassVar + +from fastvideo.configs.models import DiTConfig, UpsamplerConfig, VAEConfig +from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRConfig +from fastvideo.configs.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerConfig +from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class Kandinsky6SRPipelineConfig(PipelineConfig): + """Kandinsky6 SR: source video -> KVAE encode -> latent upscaler -> tiled SR DiT -> KVAE decode -> stitch.""" + + # Read by ``VideoGenerator`` (plain class attributes, deliberately not dataclass fields): + # prompt_optional -- a call without a text prompt is valid (V2V; the prompt is never used); + # output_shape_from_input -- the output geometry comes from the input video, so ``VideoGenerator`` must not + # pre-allocate the return buffer from the request's height / width / num_frames. + prompt_optional: ClassVar[bool] = True + output_shape_from_input: ClassVar[bool] = True + + dit_config: DiTConfig = field(default_factory=Kandinsky6SRConfig) + vae_config: VAEConfig = field(default_factory=Kandinsky6SRVAEConfig) + upsampler_config: UpsamplerConfig = field(default_factory=Kandinsky6SRLatentUpscalerConfig) + + # KVAE, latent upscaler and DiT run in bf16 (the latent state and the pi-Flow / Euler update maths stay fp32). + dit_precision: str = "bf16" + vae_precision: str = "bf16" + upsampler_precision: str = "bf16" + + # The KVAE tiles in time itself (segment cache) and the pipeline tiles in space: no generic tiled-VAE decode. + vae_tiling: bool = False + vae_sp: bool = False + + # No text encoders (see module docstring); ``check_pipeline_config`` wants the lists to have equal lengths. + text_encoder_configs: tuple = field(default_factory=tuple) + text_encoder_precisions: tuple = field(default_factory=tuple) + preprocess_text_funcs: tuple = field(default_factory=tuple) + postprocess_text_funcs: tuple = field(default_factory=tuple) + + def __post_init__(self) -> None: + # The SR pipeline always encodes (LU path) and decodes with the KVAE. + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True diff --git a/fastvideo/configs/pipelines/kandinsky6_sr_options.py b/fastvideo/configs/pipelines/kandinsky6_sr_options.py new file mode 100644 index 000000000..871a1dd72 --- /dev/null +++ b/fastvideo/configs/pipelines/kandinsky6_sr_options.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-request options owned by Kandinsky6 SR, carried in ``ForwardBatch.extra``.""" +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, fields +from typing import Any + + +@dataclass(frozen=True) +class Kandinsky6SROptions: + sr_resolution_scale: float = 2.25 + sr_tiles_batch_size: int = 1 + sr_tile_min_overlap: float = 0.20 + sr_target_resolution: str | None = None + sr_target_resize_mode: str = "fit" + + @classmethod + def from_extra(cls, extra: Mapping[str, Any]) -> Kandinsky6SROptions: + """Keep omitted options at their defaults; preserve explicit values, including None.""" + return cls(**{key: extra[key] for key in SR_OPTION_FIELDS if key in extra}) + + +SR_OPTION_FIELDS = tuple(field.name for field in fields(Kandinsky6SROptions)) +SR_REQUEST_FIELDS = (*SR_OPTION_FIELDS, "sr_lr_latent", "sr_audio", "sr_audio_sample_rate") diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 90d454aa4..69150fc44 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -37,6 +37,7 @@ from fastvideo.api.compat import ( normalize_generation_request, normalize_generator_config, request_to_batch_extra, + validate_request_batch_extra, request_to_pipeline_overrides, request_to_sampling_param, ) @@ -129,6 +130,35 @@ def _resolve_output_size( return fallback +def _quantize_video_frames_to_uint8(src: torch.Tensor) -> torch.Tensor: + """`[b, c, t, h, w]` decoded output -> CPU uint8 `[t, b, c, h, w]`. + + Quantizes on ``src``'s own device (typically CUDA) before the device->host copy: a full fp32 D->H copy scales with + resolution x frames x batch, while casting to uint8 first makes the transfer 4x smaller and moves the elementwise + work onto the GPU. ``clamp_()`` also fixes a latent overflow bug (VAE output slightly outside [0, 1] wrapped mod + 256 in the old unclamped cast). Equivalence with a CPU quantize is SSIM-gated, not bit-exact (float->uint8 differs + <=1 LSB CPU vs GPU). uint8 input is already quantized by the worker (MiniMax-H3 decode stage) and passes through + untouched (see ``pixels_to_uint8``). + + A large DiT run with ``dit_cpu_offload=False`` can leave a GPU with only a few hundred MB free once denoising and + decode are done, in a *different* process than the one that allocates ``src`` (the worker keeps the model + resident); the tiny elementwise allocation this needs can then be the single straw that OOMs. That failure is + recoverable -- the data itself is already on the device (``src`` is IPC-shared, not a fresh allocation) and only + needs a device->host copy, which does not need a device allocation -- so on ``torch.cuda.OutOfMemoryError`` this + falls back to copying ``src`` to CPU first and quantizing there instead of crashing the whole generation. + """ + try: + vid_u8 = pixels_to_uint8(src) + return rearrange(vid_u8, "b c t h w -> t b c h w").cpu() + except torch.cuda.OutOfMemoryError: + logger.warning("GPU out of memory quantizing the decoded video (the DiT/VAE likely still hold the device " + "resident); falling back to a CPU quantize. Pass dit_cpu_offload=True / " + "vae_cpu_offload=True to avoid this, at some speed cost.") + src_cpu = src.detach().to("cpu") + vid_u8 = pixels_to_uint8(src_cpu if src_cpu.dtype == torch.uint8 else src_cpu.float()) + return rearrange(vid_u8, "b c t h w -> t b c h w") + + def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None: """Validate typed stage overrides against the model's registered preset.""" if not request.stage_overrides: @@ -536,6 +566,8 @@ class VideoGenerator: if fastvideo_args is None: fastvideo_args = self.fastvideo_args + validate_request_batch_extra(kwargs, model_path=fastvideo_args.model_path) + # Handle batch processing from text file if sampling_param is None: sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path) @@ -600,15 +632,24 @@ class VideoGenerator: return results # Single prompt generation (original behavior) + output_name_hint: str | None = None if prompt is None: if fastvideo_args.workload_type is WorkloadType.V2A: # Video semantics are sufficient conditioning for V2A models; # model-specific text stages interpret the empty string using # their native tokenizer/empty-prompt contract. prompt = "" + elif getattr(fastvideo_args.pipeline_config, "prompt_optional", False): + # Video-to-video pipelines without a text tower (Kandinsky6 SR) + # are conditioned by the input video alone; name the output + # after that video instead of the (empty) prompt. + prompt = "" + if isinstance(sampling_param.video_path, str) and sampling_param.video_path: + output_name_hint = os.path.splitext(os.path.basename(sampling_param.video_path))[0] else: raise ValueError("Either prompt or prompt_txt must be provided") - output_path = self._prepare_output_path(sampling_param.output_path, prompt) + output_path = self._prepare_output_path(sampling_param.output_path, + output_name_hint if output_name_hint else prompt) kwargs["output_path"] = output_path if prompt_embeds is not None: kwargs["prompt_embeds"] = prompt_embeds @@ -816,7 +857,10 @@ class VideoGenerator: # *do* ask for the latent samples via ``return_frames=True``. # ``skip_pixel_prealloc`` also gates the slow-path warning. needs_samples_out = batch.return_frames - skip_pixel_prealloc = is_latent_output or not needs_samples_out + # ``output_shape_from_input`` pipelines (Kandinsky6 SR) decide the output geometry from the input video, so the + # request's height / width / num_frames say nothing about the buffer to pre-allocate. + output_shape_from_input = bool(getattr(fastvideo_args.pipeline_config, "output_shape_from_input", False)) + skip_pixel_prealloc = is_latent_output or not needs_samples_out or output_shape_from_input if skip_pixel_prealloc: samples = torch.empty(0, device='cpu') else: @@ -895,24 +939,10 @@ class VideoGenerator: elif not needs_frame_output: frames = None else: - # Quantize on the source device (typically CUDA) BEFORE the - # device->host copy. `samples` above is just the pinned-CPU - # mirror of `output_batch.output` (`samples.copy_(output)` or - # `output.cpu()`) with no intervening preprocessing, so reading - # `output_batch.output` here is the same data. The old path - # paid a full fp32 video D->H copy (which scales with - # resolution x frames x batch) and then a single-threaded - # per-frame CPU *255/cast loop. Casting to uint8 on-device - # makes the transfer 4x smaller, ships it in a single copy, - # and moves the elementwise work onto the GPU. clamp_() also - # fixes a latent overflow bug: VAE output slightly outside - # [0, 1] wrapped mod 256 in the old unclamped cast. - # (Equivalence is SSIM-gated, not bit-exact: float->uint8 - # differs <=1 LSB CPU vs GPU.) - # uint8 input is already quantized by the worker (MiniMax-H3 - # decode stage) and passes through untouched. - vid_u8 = pixels_to_uint8(output_batch.output) - vid_u8 = rearrange(vid_u8, "b c t h w -> t b c h w").cpu() + # `samples` above is just the pinned-CPU mirror of `output_batch.output` (`samples.copy_(output)` or + # `output.cpu()`) with no intervening preprocessing, so reading `output_batch.output` here is the same + # data. See `_quantize_video_frames_to_uint8` for why this quantizes on-device and its OOM fallback. + vid_u8 = _quantize_video_frames_to_uint8(output_batch.output) frames = [ torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy() for x in vid_u8 ] @@ -947,6 +977,9 @@ class VideoGenerator: logger.info("Saved image to %s", output_path) else: assert frames is not None # implied by save_to_disk and not audio_only + # ``batch`` is the request-side batch. A pipeline that derives the frame rate from its input (Kandinsky6 + # SR follows the source clip) mutates a worker-side copy, so it reports the rate via ``extra``. + save_fps = int(output_batch.extra.get("output_fps") or batch.fps) audio = output_batch.extra.get("audio") audio_sample_rate = output_batch.extra.get("audio_sample_rate") if audio is not None and audio_sample_rate is not None: @@ -957,7 +990,7 @@ class VideoGenerator: save_ok = self._save_video_with_audio_ffmpeg_pipe( output_path=output_path, frames=frames, - fps=batch.fps, + fps=save_fps, audio=audio, sample_rate=int(audio_sample_rate), ) @@ -966,7 +999,7 @@ class VideoGenerator: save_ok = self._save_video_with_audio_single_pass( output_path=output_path, frames=frames, - fps=batch.fps, + fps=save_fps, audio=audio, sample_rate=int(audio_sample_rate), ) @@ -976,7 +1009,7 @@ class VideoGenerator: else: logger.warning("Single-pass save failed; falling back to two-step save/mux.") save_start = time.perf_counter() - imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") + imageio.mimsave(output_path, frames, fps=save_fps, format="mp4") save_video_time = time.perf_counter() - save_start mux_start = time.perf_counter() mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate)) @@ -985,7 +1018,7 @@ class VideoGenerator: logger.warning("Audio mux failed; saved video without audio.") else: save_start = time.perf_counter() - imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") + imageio.mimsave(output_path, frames, fps=save_fps, format="mp4") save_video_time = time.perf_counter() - save_start audio_mux_time = 0.0 logger.info("Saved video to %s", output_path) diff --git a/fastvideo/models/audio/kandinsky6_audio_vae.py b/fastvideo/models/audio/kandinsky6_audio_vae.py new file mode 100644 index 000000000..80f12cab4 --- /dev/null +++ b/fastvideo/models/audio/kandinsky6_audio_vae.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 audio VAE: a thin nesting wrapper, not a new architecture. + +FastVideo's existing MMAudioVAE (the mel<->latent codec) nested directly under +`vae.*`, plus a `mel_converter` waveform->mel STFT frontend, matching the +checkpoint's audio_vae layout. Only the module nesting is new; the layers are +reused unmodified. + +The BigVGAN-v2 mel->waveform vocoder is a separate pipeline component +(`vocoder`, loaded by `VocoderLoader` as `BigVGANV2`; see +`Kandinsky6AudioDecodingStage`). The checkpoint's `audio_vae/*.safetensors` +also bundles a copy of the vocoder weights under `vocoder.*`; +`AudioDecoderLoader` drops those keys before the strict state-dict load, since +this class has no `vocoder` submodule. + +`mel_converter` is only needed to encode real audio, which T2VA/IT2VA +generation never does: it only decodes audio latents into a waveform. +""" +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.models.audio.mmaudio_vae import MMAudioVAE + + +class MelConverter(nn.Module): + """Waveform -> log-mel-spectrogram STFT frontend, matching the + diffusers reference's `MelConverter`/`get_mel_converter("44k")`. See + module docstring: not exercised by the current (decode-only) inference + path.""" + + def __init__( + self, + *, + sampling_rate: int = 44100, + n_fft: int = 2048, + num_mels: int = 128, + hop_size: int = 512, + win_size: int = 2048, + fmin: int = 0, + fmax: int | None = 22050, + ) -> None: + super().__init__() + self.n_fft = n_fft + self.hop_size = hop_size + self.win_size = win_size + self.register_buffer("hann_window", torch.hann_window(win_size), persistent=True) + try: + from librosa.filters import mel as librosa_mel_fn + mel_basis = torch.from_numpy( + librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)).float() + except Exception: + # librosa is an optional eval-only dependency here (see + # pyproject.toml's eval-audio extra), not a core one, and this + # buffer is only ever populated by the checkpoint's own loaded + # weights for the (currently unused, decode-only) encode path -- + # a zero filterbank still gives the loader a correctly-shaped + # buffer to load real weights into. + mel_basis = torch.zeros(num_mels, n_fft // 2 + 1) + self.register_buffer("mel_basis", mel_basis, persistent=True) + + def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: + waveform = waveform.clamp(min=-1.0, max=1.0) + pad = (self.n_fft - self.hop_size) // 2 + padded = F.pad(waveform.unsqueeze(1), (pad, pad), mode="reflect").squeeze(1) + spec = torch.stft( + padded, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.view_as_real(spec) + magnitude = torch.sqrt(spec.pow(2).sum(-1) + 1e-9).float() + mel = torch.matmul(self.mel_basis, magnitude) + return torch.log(torch.clamp(mel, min=1e-5)) + + +class Kandinsky6AudioVAE(nn.Module): + """Nests the reused MMAudioVAE to match the checkpoint's exact parameter + names. Decodes latents to a mel spectrogram only -- the separate + `vocoder` pipeline component turns that into a waveform (see module + docstring).""" + + def __init__(self, config: dict[str, Any]) -> None: + super().__init__() + config = dict(config) + config.pop("_class_name", None) + config.pop("_diffusers_version", None) + self.scaling_factor = float(config.get("scaling_factor", 1.0)) + # Audio-latent-frames-per-raw-audio-sample, matching the diffusers + # reference's MMAudioVAE.downsample_factor and the pipeline config's + # audio_downsample_factor default. + self.downsample_factor = 1024 + + mode = config.get("mode", "44k") + need_encoder = bool(config.get("need_vae_encoder", False)) + + self.mel_converter = MelConverter(sampling_rate=44100, n_fft=2048, num_mels=128, hop_size=512, + win_size=2048, fmin=0, fmax=22050) + self.vae = MMAudioVAE(mode=mode, need_encoder=need_encoder) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """latents: [B, embed_dim, A] (1D-conv channel-first) -> mel: [B, num_mels, T].""" + return self.vae.decode(latents, unnormalize_output=True) + + def encode_audio(self, waveform: torch.Tensor): + """waveform -> mel -> VAE posterior. Not used by generation (see the module docstring).""" + mel = self.mel_converter(waveform) + return self.vae.encode(mel) + + def wrapped_encode(self, waveform: torch.Tensor) -> torch.Tensor: + """waveform -> mean audio latent in one call, matching the + diffusers reference's MMAudioVAE.wrapped_encode.""" + return self.encode_audio(waveform).mean + + def remove_weight_norm(self) -> "Kandinsky6AudioVAE": + """Called by the loader after load_state_dict: MMAudioVAE's custom + post-load weight renormalization (needs real loaded values to + compute from).""" + self.vae.remove_weight_norm() + return self + + +EntryClass = Kandinsky6AudioVAE diff --git a/fastvideo/models/dits/kandinsky6.py b/fastvideo/models/dits/kandinsky6.py new file mode 100644 index 000000000..539474201 --- /dev/null +++ b/fastvideo/models/dits/kandinsky6.py @@ -0,0 +1,1067 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native FastVideo implementation of the Kandinsky6 T2VA/IT2VA transformer. + +Kandinsky6 extends Kandinsky5's DiT (fastvideo/models/dits/kandinsky5.py) +with a second, parallel "audio" tower and a fused video<->audio decoder +block, so this file mirrors kandinsky5.py's structure closely for every +shared piece (time/text/visual embeddings, RoPE, modulation, feed-forward, +NABLA sparse attention) and only adds what's new: the audio tower and +Kandinsky6FusedTransformerDecoderBlock's cross-modal attention, ported from +the reference diffusers.models.transformers.transformer_kandinsky6 module. + +Video and audio latents are ordinary batched tensors -- `[B, T, H, W, C]` for +video and `[B, A, D]` for audio -- consistent with Kandinsky5 and the rest of +FastVideo's pipeline/stage system. `Kandinsky6RoPE3D.forward` takes an explicit +`shape` tuple rather than deriving T/H/W from the position tensors' lengths. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.attention import LocalAttention +from fastvideo.attention.backends.nabla import CAN_USE_FLEX_ATTN, flex_attention, nablaT_v2 +from fastvideo.configs.models.dits import Kandinsky6VideoAudioConfig +from fastvideo.layers.layernorm import LayerNormScaleShift +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + +logger = init_logger(__name__) + +if not CAN_USE_FLEX_ATTN: + logger.warning("torch.nn.attention.flex_attention is unavailable in this PyTorch build; " + "Kandinsky6 NABLA sparse attention cannot be used.") + +FRACTAL_PIXEL_SIZE = 8 +_ARCH_CONFIG_DEFAULTS = Kandinsky6VideoAudioConfig().arch_config + +# Modulation submodules whose maths diffusers keeps fp32 in the distill checkpoint's +# `from_pretrained(torch_dtype=bf16)` load, matched by exact dotted-path component (not substring) against +# `visual_modulation`/`text_modulation` (the block-level AdaLN modules) or `modulation` (the plain attribute name of +# `Kandinsky6OutLayer`/`Kandinsky6OutLayerAudio`'s modulation submodule). The differently-named `va_modulation`/ +# `av_modulation` cross-modal gates and the `*_time_embeddings` towers are fp32 in the raw checkpoint file too, but +# diffusers' exact-component match does not keep them fp32, so they are deliberately not listed here. +_FP32_MODULATION_COMPONENTS = frozenset({"visual_modulation", "text_modulation", "modulation"}) + + +def _build_rotary_freqs(dim: int, max_period: float) -> torch.Tensor: + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + + +def local_patching(x: torch.Tensor, + shape: tuple[int, int, int, int], + group_size: tuple[int, int, int], + dim: int = 0) -> torch.Tensor: + batch_size, duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape( + *x.shape[:dim], + duration // g1, + g1, + height // g2, + g2, + width // g3, + g3, + *x.shape[dim + 3:], + ) + x = x.permute( + *range(len(x.shape[:dim])), + dim, + dim + 2, + dim + 4, + dim + 1, + dim + 3, + dim + 5, + *range(dim + 6, len(x.shape)), + ) + x = x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) + return x + + +def local_merge(x: torch.Tensor, + shape: tuple[int, int, int, int], + group_size: tuple[int, int, int], + dim: int = 0) -> torch.Tensor: + batch_size, duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape( + *x.shape[:dim], + duration // g1, + height // g2, + width // g3, + g1, + g2, + g3, + *x.shape[dim + 2:], + ) + x = x.permute( + *range(len(x.shape[:dim])), + dim, + dim + 3, + dim + 1, + dim + 4, + dim + 2, + dim + 5, + *range(dim + 6, len(x.shape)), + ) + x = x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) + return x + + +def fractal_flatten(x: torch.Tensor, rope: torch.Tensor, shape: tuple[int, int, int, int], block_mask: bool = False): + if block_mask: + pixel_size = FRACTAL_PIXEL_SIZE + x = local_patching(x, shape, (1, pixel_size, pixel_size), dim=1) + rope = local_patching(rope, shape, (1, pixel_size, pixel_size), dim=1) + x = x.flatten(1, 2) + rope = rope.flatten(1, 2) + else: + x = x.flatten(1, 3) + rope = rope.flatten(1, 3) + return x, rope + + +def fractal_unflatten(x: torch.Tensor, shape: tuple[int, int, int, int], block_mask: bool = False): + if block_mask: + pixel_size = FRACTAL_PIXEL_SIZE + x = x.reshape(x.shape[0], -1, pixel_size**2, *x.shape[2:]) + x = local_merge(x, shape, (1, pixel_size, pixel_size), dim=1) + else: + x = x.reshape(*shape, *x.shape[2:]) + return x + + +class Kandinsky6TimeEmbeddings(nn.Module): + + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + super().__init__() + assert model_dim % 2 == 0 + self.model_dim = model_dim + self.max_period = max_period + self.freqs = _build_rotary_freqs(self.model_dim // 2, self.max_period) + self.in_layer = ReplicatedLinear(model_dim, time_dim, bias=True) + self.activation = nn.SiLU() + self.out_layer = ReplicatedLinear(time_dim, time_dim, bias=True) + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, time: torch.Tensor) -> torch.Tensor: + args = torch.outer(time, self.freqs.to(device=time.device)) + time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + time_embed, _ = self.in_layer(time_embed) + time_embed = self.activation(time_embed) + time_embed, _ = self.out_layer(time_embed) + return time_embed + + +class Kandinsky6TextEmbeddings(nn.Module): + """Linear + LayerNorm projection, reused for text tokens and audio latents.""" + + def __init__(self, in_dim: int, model_dim: int): + super().__init__() + self.in_layer = ReplicatedLinear(in_dim, model_dim, bias=True) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.in_layer(x) + return self.norm(x).type_as(x) + + +class Kandinsky6VisualEmbeddings(nn.Module): + + def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]): + super().__init__() + self.patch_size = patch_size + self.in_layer = ReplicatedLinear(math.prod(patch_size) * visual_dim, model_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch_size, duration, height, width, dim = x.shape + x = (x.view( + batch_size, + duration // self.patch_size[0], + self.patch_size[0], + height // self.patch_size[1], + self.patch_size[1], + width // self.patch_size[2], + self.patch_size[2], + dim, + ).permute(0, 1, 3, 5, 2, 4, 6, 7).flatten(4, 7)) + x, _ = self.in_layer(x) + return x + + +class Kandinsky6RoPE1D(nn.Module): + """1D rotary embedding for text and audio sequences.""" + + def __init__(self, dim: int, max_pos: int = 2048, max_period: float = 10000.0, freqs_scaling: float = 1.0): + super().__init__() + self.max_period = max_period + self.dim = dim + self.max_pos = max_pos + self.freqs_scaling = freqs_scaling + freq = _build_rotary_freqs(dim // 2, max_period) * freqs_scaling + pos = torch.arange(max_pos, dtype=freq.dtype) + self.register_buffer("args", torch.outer(pos, freq), persistent=False) + + def forward(self, pos: torch.Tensor) -> torch.Tensor: + args = self.args[pos] + cosine = torch.cos(args) + sine = torch.sin(args) + rope = torch.stack([cosine, -sine, sine, cosine], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + +class Kandinsky6RoPE3D(nn.Module): + """3D rotary embedding for video spatial-temporal (T, H, W) tokens.""" + + def __init__(self, + axes_dims: tuple[int, int, int], + max_pos: tuple[int, int, int] = (128, 128, 128), + max_period: float = 10000.0): + super().__init__() + self.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + + for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos, strict=True)): + freq = _build_rotary_freqs(axes_dim // 2, max_period) + pos = torch.arange(ax_max_pos, dtype=freq.dtype) + self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) + + def forward(self, + shape: tuple[int, int, int, int], + pos: tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | list[torch.Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0)): + batch_size, duration, height, width = shape + args_t = self.args_0[pos[0]] / scale_factor[0] + args_h = self.args_1[pos[1]] / scale_factor[1] + args_w = self.args_2[pos[2]] / scale_factor[2] + + args = torch.cat( + [ + args_t.view(1, duration, 1, 1, -1).repeat(batch_size, 1, height, width, 1), + args_h.view(1, 1, height, 1, -1).repeat(batch_size, duration, 1, width, 1), + args_w.view(1, 1, 1, width, -1).repeat(batch_size, duration, height, 1, 1), + ], + dim=-1, + ) + cosine = torch.cos(args) + sine = torch.sin(args) + rope = torch.stack([cosine, -sine, sine, cosine], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + +class Kandinsky6Modulation(nn.Module): + + def __init__(self, time_dim: int, model_dim: int, num_params: int): + super().__init__() + self.activation = nn.SiLU() + self.out_layer = ReplicatedLinear(time_dim, num_params * model_dim, bias=True) + self.out_layer.weight.data.zero_() + self.out_layer.bias.data.zero_() + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.activation(x) + x, _ = self.out_layer(x) + return x + + +def _apply_rotary(x: torch.Tensor, rope: torch.Tensor) -> torch.Tensor: + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) + x_out = (rope * x_).sum(dim=-1) + return x_out.reshape(*x.shape).to(x.dtype) + + +class Kandinsky6Attention(nn.Module): + """Self- or cross-attention. ``kv_dim`` lets K/V come from a + differently-sized stream (the video<->audio cross-modal attentions).""" + + def __init__( + self, + num_channels: int, + head_dim: int, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None, + prefix: str = "", + kv_dim: int | None = None, + use_nabla: bool = False, + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + assert num_channels % head_dim == 0 + self.num_heads = num_channels // head_dim + kv_dim = kv_dim or num_channels + + self.to_query = ReplicatedLinear(num_channels, + num_channels, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_query") + self.to_key = ReplicatedLinear(kv_dim, + num_channels, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_key") + self.to_value = ReplicatedLinear(kv_dim, + num_channels, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_value") + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + self.out_layer = ReplicatedLinear(num_channels, + num_channels, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.out_layer") + self.local_attention = LocalAttention( + num_heads=self.num_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + ) + # Only the video self-attention gets a second NABLA-backed attention + # layer; audio self-attention and every cross-attention are dense. + self.nabla_attention = None + if use_nabla: + self.nabla_attention = LocalAttention( + num_heads=self.num_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + default_backend=AttentionBackendEnum.NABLA_ATTN, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + rotary_emb: torch.Tensor | None = None, + rotary_emb_kv: torch.Tensor | None = None, + ) -> torch.Tensor: + query, _ = self.to_query(hidden_states) + + kv_source = hidden_states if encoder_hidden_states is None else encoder_hidden_states + key, _ = self.to_key(kv_source) + value, _ = self.to_value(kv_source) + + shape, kv_shape = query.shape[:-1], key.shape[:-1] + query = query.reshape(*shape, self.num_heads, -1) + key = key.reshape(*kv_shape, self.num_heads, -1) + value = value.reshape(*kv_shape, self.num_heads, -1) + + query = self.query_norm(query.float()).type_as(query) + key = self.key_norm(key.float()).type_as(key) + + if rotary_emb is not None: + query = _apply_rotary(query, rotary_emb).type_as(query) + kv_rope = rotary_emb_kv if rotary_emb_kv is not None else (rotary_emb if encoder_hidden_states is None else None) + if kv_rope is not None: + key = _apply_rotary(key, kv_rope).type_as(key) + + if sparse_params is not None: + if self.nabla_attention is None: + raise RuntimeError("sparse_params passed to an attention layer built without use_nabla; " + "this checkpoint/config combination is inconsistent.") + try: + hidden_states = self.nabla_attention(query, key, value) + except AssertionError as exc: + # Standalone parity tests call the model without a pipeline + # forward context; run the NABLA kernel directly. + if "Forward context is not set" not in str(exc): + raise + attn_mask = nablaT_v2(query, key, sparse_params["sta_mask"], thr=sparse_params["P"]) + hidden_states = flex_attention( + query=query.transpose(1, 2), + key=key.transpose(1, 2), + value=value.transpose(1, 2), + block_mask=attn_mask, + ).transpose(1, 2) + else: + try: + hidden_states = self.local_attention(query, key, value) + except AssertionError as exc: + # LocalAttention requires pipeline forward context. Standalone + # parity tests call the model directly, so fall back to SDPA. + if "Forward context is not set" not in str(exc): + raise + + query_shape = query.shape[:-2] + key_shape = key.shape[:-2] + query = query.reshape(query_shape[0], -1, self.num_heads, query.shape[-1]).transpose(1, 2) + key = key.reshape(key_shape[0], -1, self.num_heads, key.shape[-1]).transpose(1, 2) + value = value.reshape(key_shape[0], -1, self.num_heads, value.shape[-1]).transpose(1, 2) + hidden_states = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=None, + is_causal=False, + ) + hidden_states = hidden_states.transpose(1, 2).reshape(*query_shape, self.num_heads, -1) + + hidden_states = hidden_states.flatten(-2, -1) + hidden_states, _ = self.out_layer(hidden_states) + return hidden_states + + +class Kandinsky6FeedForward(nn.Module): + + def __init__(self, dim: int, ff_dim: int, prefix: str = "", quant_config: QuantizationConfig | None = None): + super().__init__() + self.mlp = MLP(dim, ff_dim, bias=False, act_type="gelu", quant_config=quant_config, prefix=f"{prefix}.mlp") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +class Kandinsky6OutLayer(nn.Module): + """Projects visual hidden states back to packed latent patches.""" + + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]): + super().__init__() + self.patch_size = patch_size + self.modulation = Kandinsky6Modulation(time_dim, model_dim, 2) + self.norm = nn.LayerNorm(model_dim, eps=1e-5, elementwise_affine=False) + self.out_layer = ReplicatedLinear(model_dim, math.prod(patch_size) * visual_dim, bias=True) + + def forward(self, visual_embed: torch.Tensor, time_embed: torch.Tensor) -> torch.Tensor: + shift, scale = torch.chunk(self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + visual_embed = (self.norm(visual_embed.float()) * (scale.float()[:, None, None] + 1.0) + + shift.float()[:, None, None]).type_as(visual_embed) + + x, _ = self.out_layer(visual_embed) + + batch_size, duration, height, width, _ = x.shape + x = (x.view( + batch_size, + duration, + height, + width, + -1, + self.patch_size[0], + self.patch_size[1], + self.patch_size[2], + ).permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(1, 2).flatten(2, 3).flatten(3, 4)) + return x + + +class Kandinsky6OutLayerAudio(nn.Module): + """Projects audio hidden states back to audio latent channels.""" + + def __init__(self, model_dim: int, time_dim: int, audio_dim: int): + super().__init__() + self.modulation = Kandinsky6Modulation(time_dim, model_dim, 2) + self.norm = nn.LayerNorm(model_dim, eps=1e-5, elementwise_affine=False) + self.out_layer = ReplicatedLinear(model_dim, audio_dim, bias=True) + + def forward(self, audio_embed: torch.Tensor, time_embed: torch.Tensor) -> torch.Tensor: + shift, scale = torch.chunk(self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + x = (self.norm(audio_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as(audio_embed) + # Normalized a second time after the scale/shift affine, before the + # output projection, as in the Diffusers model; the weights expect it. + x = self.norm(x.float()).type_as(audio_embed) + out, _ = self.out_layer(x) + return out + + +class Kandinsky6TransformerEncoderBlock(nn.Module): + """Text-only self-attention + feed-forward block (video/audio text towers).""" + + def __init__(self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + quant_config: QuantizationConfig | None = None): + super().__init__() + self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) + + self.self_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.self_attention = Kandinsky6Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.self_attention", + quant_config=quant_config) + + self.feed_forward_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.feed_forward = Kandinsky6FeedForward(model_dim, + ff_dim, + prefix=f"{prefix}.feed_forward", + quant_config=quant_config) + + def forward(self, x: torch.Tensor, time_embed: torch.Tensor, rope: torch.Tensor) -> torch.Tensor: + self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) + out = self.self_attention_norm(x.float(), shift=shift, scale=scale, convert_modulation_dtype=True).type_as(x) + out = self.self_attention(out, rotary_emb=rope) + x = (x.float() + gate.float() * out.float()).type_as(x) + + ff_shift, ff_scale, ff_gate = torch.chunk(ff_params, 3, dim=-1) + out = self.feed_forward_norm(x.float(), shift=ff_shift, scale=ff_scale, + convert_modulation_dtype=True).type_as(x) + out = self.feed_forward(out) + x = (x.float() + ff_gate.float() * out.float()).type_as(x) + + return x + + +class Kandinsky6TransformerDecoderBlock(nn.Module): + """Self-attention + text cross-attention + feed-forward block. + + Used standalone for plain (non-multimodal) T2V/I2V-parity checkpoints, + and as the ``videoT``/``audioT`` sub-block inside + ``Kandinsky6FusedTransformerDecoderBlock`` for T2VA/IT2VA. + """ + + def __init__(self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + use_nabla: bool = False, + quant_config: QuantizationConfig | None = None): + super().__init__() + self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) + + self.self_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.self_attention = Kandinsky6Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.self_attention", + use_nabla=use_nabla, + quant_config=quant_config) + + self.cross_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.cross_attention = Kandinsky6Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.cross_attention", + quant_config=quant_config) + + self.feed_forward_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.feed_forward = Kandinsky6FeedForward(model_dim, + ff_dim, + prefix=f"{prefix}.feed_forward", + quant_config=quant_config) + + def forward(self, visual_embed: torch.Tensor, text_embed: torch.Tensor, time_embed: torch.Tensor, + rope: torch.Tensor | None, sparse_params: dict[str, Any] | None) -> torch.Tensor: + self_attn_params, cross_attn_params, ff_params = torch.chunk( + self.visual_modulation(time_embed).unsqueeze(dim=1), 3, dim=-1) + + self_shift, self_scale, self_gate = torch.chunk(self_attn_params, 3, dim=-1) + visual_out = self.self_attention_norm( + visual_embed.float(), + shift=self_shift, + scale=self_scale, + convert_modulation_dtype=True, + ).type_as(visual_embed) + visual_out = self.self_attention(visual_out, rotary_emb=rope, sparse_params=sparse_params) + visual_embed = (visual_embed.float() + self_gate.float() * visual_out.float()).type_as(visual_embed) + + cross_shift, cross_scale, cross_gate = torch.chunk(cross_attn_params, 3, dim=-1) + visual_out = self.cross_attention_norm( + visual_embed.float(), + shift=cross_shift, + scale=cross_scale, + convert_modulation_dtype=True, + ).type_as(visual_embed) + visual_out = self.cross_attention(visual_out, encoder_hidden_states=text_embed) + visual_embed = (visual_embed.float() + cross_gate.float() * visual_out.float()).type_as(visual_embed) + + ff_shift, ff_scale, ff_gate = torch.chunk(ff_params, 3, dim=-1) + visual_out = self.feed_forward_norm( + visual_embed.float(), + shift=ff_shift, + scale=ff_scale, + convert_modulation_dtype=True, + ).type_as(visual_embed) + visual_out = self.feed_forward(visual_out) + visual_embed = (visual_embed.float() + ff_gate.float() * visual_out.float()).type_as(visual_embed) + + return visual_embed + + +def _apply_gate_sum(x: torch.Tensor, out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + return (x.float() + gate.float() * out.float()).type_as(x) + + +def _apply_scale_shift(norm: nn.LayerNorm, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).type_as(x) + + +class Kandinsky6FusedTransformerDecoderBlock(nn.Module): + """Joint video+audio decoder block: per-modality self/text-cross + attention plus a dedicated bidirectional video<->audio cross-attention, + all independently AdaLN-modulated. Ported from + diffusers.models.transformers.transformer_kandinsky6.Kandinsky6FusedTransformerDecoderBlock. + """ + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + model_dim_a: int, + time_dim_a: int, + ff_dim_a: int, + head_dim_a: int, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = "", + use_nabla: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, + time_dim, + ff_dim, + head_dim, + supported_attention_backends, + prefix=f"{prefix}.videoT", + use_nabla=use_nabla, + quant_config=quant_config) + self.audioT = Kandinsky6TransformerDecoderBlock(model_dim_a, + time_dim_a, + ff_dim_a, + head_dim_a, + supported_attention_backends, + prefix=f"{prefix}.audioT", + use_nabla=False, + quant_config=quant_config) + self.va_cross_attention = Kandinsky6Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.va_cross_attention", + kv_dim=model_dim_a, + quant_config=quant_config) + self.av_cross_attention = Kandinsky6Attention(model_dim_a, + head_dim_a, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.av_cross_attention", + kv_dim=model_dim, + quant_config=quant_config) + self.va_modulation = Kandinsky6Modulation( + time_dim, + model_dim if not cross_gates else model_dim * 2 + model_dim_a, + 1 if cross_gates else 3, + ) + self.av_modulation = Kandinsky6Modulation( + time_dim_a, + model_dim_a if not cross_gates else model_dim_a * 2 + model_dim, + 1 if cross_gates else 3, + ) + self.va_normalization = nn.LayerNorm(model_dim, elementwise_affine=False) + self.av_normalization = nn.LayerNorm(model_dim_a, elementwise_affine=False) + self.ca_rope = ca_rope + self.cross_gates = cross_gates + self.fix_modulation = fix_modulation + self.model_dim = model_dim + self.model_dim_a = model_dim_a + + def forward( + self, + vis: torch.Tensor | None, + aud: torch.Tensor | None, + text_v: torch.Tensor, + text_a: torch.Tensor | None, + time_embed: tuple[torch.Tensor, torch.Tensor | None], + vis_rope: torch.Tensor | None, + aud_rope: torch.Tensor | None, + sparse_params: dict[str, Any] | None, + va_gate_scale: float = 1.0, + av_gate_scale: float = 1.0, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """``vis``/``aud`` may each be ``None`` (matches the diffusers reference's guarded + ``Kandinsky6FusedTransformerDecoderBlock``): every stage below only runs for a present + modality, and the video<->audio cross-modal mixing only runs when both are present.""" + t_v, t_a = time_embed + gate_v = vis_out_t = None + + if vis is not None: + sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v).unsqueeze(dim=1), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + vis = _apply_gate_sum( + vis, + self.videoT.self_attention( + self.videoT.self_attention_norm(vis.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(vis), + rotary_emb=vis_rope, + sparse_params=sparse_params, + ), + gate, + ) + shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) + vis_pre_ca = self.videoT.cross_attention_norm(vis.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(vis) + vis_out_t = self.videoT.cross_attention(vis_pre_ca, encoder_hidden_states=text_v) + + if aud is not None: + sa_p, ca_p, ff_p_a = torch.chunk(self.audioT.visual_modulation(t_a).unsqueeze(dim=1), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + aud = _apply_gate_sum( + aud, + self.audioT.self_attention( + self.audioT.self_attention_norm(aud.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(aud), + rotary_emb=aud_rope, + ), + gate, + ) + shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) + aud_pre_ca = self.audioT.cross_attention_norm(aud.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(aud) + aud_out_t = self.audioT.cross_attention(aud_pre_ca, encoder_hidden_states=text_a) + aud = _apply_gate_sum(aud, aud_out_t, gate_a) + + if vis is not None: + t_va_mod = t_a if not self.fix_modulation else t_v + t_av_mod = t_v if not self.fix_modulation else t_a + va_params = self.va_modulation(t_va_mod).unsqueeze(dim=1) + av_params = self.av_modulation(t_av_mod).unsqueeze(dim=1) + if self.cross_gates: + va_shift, va_scale, va_gate = torch.split( + va_params, [self.model_dim, self.model_dim, self.model_dim_a], dim=-1) + av_shift, av_scale, av_gate = torch.split( + av_params, [self.model_dim_a, self.model_dim_a, self.model_dim], dim=-1) + else: + va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) + av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) + + vis = _apply_gate_sum(vis, vis_out_t, gate_v) + vis_for_va = _apply_scale_shift(self.va_normalization, vis, va_scale, va_shift) + aud_for_av = _apply_scale_shift(self.av_normalization, aud, av_scale, av_shift) + rq_v = vis_rope if self.ca_rope else None + rk_a = aud_rope if self.ca_rope else None + vis_from_aud = self.va_cross_attention(vis_for_va, encoder_hidden_states=aud_pre_ca, rotary_emb=rq_v, + rotary_emb_kv=rk_a) + aud_from_vis = self.av_cross_attention(aud_for_av, encoder_hidden_states=vis_pre_ca, rotary_emb=rk_a, + rotary_emb_kv=rq_v) + vis = _apply_gate_sum(vis, vis_from_aud, (va_gate if not self.cross_gates else av_gate) * va_gate_scale) + aud = _apply_gate_sum(aud, aud_from_vis, (av_gate if not self.cross_gates else va_gate) * av_gate_scale) + elif vis is not None: + vis = _apply_gate_sum(vis, vis_out_t, gate_v) + + if vis is not None: + shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) + vis = _apply_gate_sum( + vis, + self.videoT.feed_forward( + self.videoT.feed_forward_norm(vis.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(vis)), + gate, + ) + if aud is not None: + shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) + aud = _apply_gate_sum( + aud, + self.audioT.feed_forward( + self.audioT.feed_forward_norm(aud.float(), shift=shift, scale=scale, + convert_modulation_dtype=True).type_as(aud)), + gate, + ) + return vis, aud + + +@dataclass +class Kandinsky6TransformerOutput: + sample: torch.Tensor | tuple[torch.Tensor, torch.Tensor] + + +class Kandinsky6Transformer3DModel(BaseDiT): + """Native FastVideo implementation of the Kandinsky6 T2VA/IT2VA transformer.""" + + _fsdp_shard_conditions = _ARCH_CONFIG_DEFAULTS._fsdp_shard_conditions + _compile_conditions = _ARCH_CONFIG_DEFAULTS._compile_conditions + param_names_mapping = _ARCH_CONFIG_DEFAULTS.param_names_mapping + reverse_param_names_mapping = _ARCH_CONFIG_DEFAULTS.reverse_param_names_mapping + lora_param_names_mapping = _ARCH_CONFIG_DEFAULTS.lora_param_names_mapping + _supported_attention_backends = _ARCH_CONFIG_DEFAULTS._supported_attention_backends + + def __init__(self, config: Kandinsky6VideoAudioConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + arch = config.arch_config + quant_config = config.quant_config + self.quant_config = quant_config + + head_dim = sum(arch.axes_dims) + head_dim_a = sum(arch.axes_dims_a) + self.in_visual_dim = arch.in_visual_dim + self.in_audio_dim = arch.in_audio_dim + self.model_dim = arch.model_dim + self.patch_size = arch.patch_size + self.visual_cond = arch.visual_cond + self.attention_engine = arch.attention_engine + self.visual_token_type_num_embeddings = arch.visual_token_type_num_embeddings + self.scale_factor = tuple(float(v) for v in arch.scale_factor) + + visual_embed_dim = (2 * arch.in_visual_dim + 1) if arch.visual_cond else arch.in_visual_dim + + self.visual_embeddings = Kandinsky6VisualEmbeddings(visual_embed_dim, arch.model_dim, arch.patch_size) + if self.visual_token_type_num_embeddings > 0: + self.visual_token_type_embeddings = nn.Embedding(self.visual_token_type_num_embeddings, arch.model_dim) + self.visual_rope_embeddings = Kandinsky6RoPE3D(arch.axes_dims) + self.out_layer = Kandinsky6OutLayer(arch.model_dim, arch.time_dim, arch.out_visual_dim, arch.patch_size) + + use_nabla = arch.attention_engine == "nabla" + + # The model is always the joint video+audio DiT; a video-only call passes + # hidden_states_audio=None instead of using a different model shape. + + # Registered before the (much smaller) text towers below: enable_layerwise_offload + # (fastvideo/hooks/layerwise_offload.py) hooks only the first top-level nn.ModuleList in + # registration order, and with dit_layerwise_offload=True (the default) that has to be these + # blocks rather than the 4-entry video_text_transformer_blocks. + self.visual_transformer_blocks = nn.ModuleList([ + Kandinsky6FusedTransformerDecoderBlock( + arch.model_dim, + arch.time_dim, + arch.ff_dim, + head_dim, + arch.model_dim_a, + arch.time_dim_a, + arch.ff_dim_a, + head_dim_a, + self._supported_attention_backends, + prefix=f"{config.prefix}.visual_transformer_blocks.{i}", + use_nabla=use_nabla, + ca_rope=arch.ca_rope, + cross_gates=arch.cross_gates, + fix_modulation=arch.fix_modulation, + quant_config=quant_config) for i in range(arch.num_visual_blocks) + ]) + + self.audio_embeddings = Kandinsky6TextEmbeddings(arch.in_audio_dim, arch.model_dim_a) + self.audio_rope_embeddings = Kandinsky6RoPE1D(head_dim_a, freqs_scaling=arch.audio_freqs_scaling) + self.audio_out_layer = Kandinsky6OutLayerAudio(arch.model_dim_a, arch.time_dim_a, arch.out_audio_dim or arch.in_audio_dim) + + for tower_prefix, model_dim, time_dim, hd in ( + ("video", arch.model_dim, arch.time_dim, head_dim), + ("audio", arch.model_dim_a, arch.time_dim_a, head_dim_a), + ): + setattr(self, f"{tower_prefix}_time_embeddings", Kandinsky6TimeEmbeddings(model_dim, time_dim)) + setattr(self, f"{tower_prefix}_text_embeddings", Kandinsky6TextEmbeddings(arch.in_text_dim, model_dim)) + setattr(self, f"{tower_prefix}_pooled_text_embeddings", + Kandinsky6TextEmbeddings(arch.in_text_dim2, time_dim)) + setattr(self, f"{tower_prefix}_text_rope_embeddings", Kandinsky6RoPE1D(hd)) + setattr( + self, f"{tower_prefix}_text_transformer_blocks", + nn.ModuleList([ + Kandinsky6TransformerEncoderBlock( + model_dim, + time_dim, + arch.ff_dim if tower_prefix == "video" else arch.ff_dim_a, + hd, + self._supported_attention_backends, + prefix=f"{config.prefix}.{tower_prefix}_text_transformer_blocks.{i}", + quant_config=quant_config) for i in range(arch.num_text_blocks) + ])) + + self.gradient_checkpointing = False + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + self.__post_init__() + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """The distill checkpoint's safetensors file stores its modulation and time-embedding tensors as F32; + diffusers' ``from_pretrained(torch_dtype=bf16)`` keeps only the exact-component modulation matches fp32 (see + ``_FP32_MODULATION_COMPONENTS``) and rounds the rest -- including ``va_modulation``/``av_modulation`` and the + ``*_time_embeddings`` towers -- to bf16. Mirrors ``Kandinsky6SRTransformer3DModel._get_parameter_dtype``. + """ + if any(component in _FP32_MODULATION_COMPONENTS for component in name.split(".")): + return torch.float32 + return default_dtype + + def _time_embed(self, prefix: str, time: torch.Tensor, pooled: torch.Tensor) -> torch.Tensor: + pooled_embeddings = getattr(self, f"{prefix}_pooled_text_embeddings") + time_embeddings = getattr(self, f"{prefix}_time_embeddings") + return time_embeddings(time) + pooled_embeddings(pooled) + + def _encode_text(self, prefix: str, text_embed: torch.Tensor, pooled: torch.Tensor, time: torch.Tensor, + text_rope_pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + text_embeddings = getattr(self, f"{prefix}_text_embeddings") + rope_embeddings = getattr(self, f"{prefix}_text_rope_embeddings") + blocks = getattr(self, f"{prefix}_text_transformer_blocks") + te = text_embeddings(text_embed) + tm = self._time_embed(prefix, time, pooled) + # Video and audio each own a text RoPE table sized to their own + # head_dim (they can differ), so the position indices are looked up + # per-tower rather than sharing one precomputed rope tensor. + text_rope = rope_embeddings(text_rope_pos).unsqueeze(dim=0) + for block in blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + te = torch.utils.checkpoint.checkpoint(block, te, tm, text_rope, use_reentrant=False) + else: + te = block(te, tm, text_rope) + return te, tm + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + pooled_projections: torch.Tensor | None = None, + hidden_states_audio: torch.Tensor | None = None, + audio_timestep: torch.Tensor | None = None, + visual_rope_pos: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | list[torch.Tensor] | None = None, + audio_rope_pos: torch.Tensor | None = None, + text_rope_pos: torch.Tensor | None = None, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + sparse_params: dict[str, Any] | None = None, + visual_token_type_ids: torch.Tensor | None = None, + va_gate_scale: float = 1.0, + av_gate_scale: float = 1.0, + return_dict: bool = False, + **kwargs: Any, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | Kandinsky6TransformerOutput: + if pooled_projections is None: + if encoder_hidden_states_image is None: + raise ValueError("pooled_projections must be provided for Kandinsky6.") + pooled_projections = encoder_hidden_states_image + if visual_rope_pos is None or text_rope_pos is None: + raise ValueError("visual_rope_pos and text_rope_pos are required for Kandinsky6.") + if visual_token_type_ids is not None and self.visual_token_type_num_embeddings == 0: + raise ValueError("visual_token_type_ids requires visual_token_type_num_embeddings > 0.") + + x_video = hidden_states + x_audio = hidden_states_audio + + # A video-only call (x_audio is None) runs the fused blocks with aud=None. + video_te, video_tm = self._encode_text("video", encoder_hidden_states, pooled_projections, timestep, + text_rope_pos) + + visual_embed = self.visual_embeddings(x_video) + if getattr(self, "visual_token_type_embeddings", None) is not None and visual_token_type_ids is not None: + type_embed = self.visual_token_type_embeddings(visual_token_type_ids) + visual_embed = visual_embed + type_embed[:, :, None, None, :] + visual_shape = visual_embed.shape[:-1] + visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) + # The NABLA sparse mask (inside the fused block's video self-attention) is built for + # fractal token order, so flatten/unflatten with the same `to_fractal` flag -- a no-op + # when to_fractal is False. + to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False + visual_embed, visual_rope = fractal_flatten(visual_embed, visual_rope, visual_shape, block_mask=to_fractal) + + audio_te = audio_tm = audio_embed = audio_rope = None + if x_audio is not None: + audio_timestep = audio_timestep if audio_timestep is not None else timestep + audio_rope_pos = audio_rope_pos if audio_rope_pos is not None else torch.arange( + x_audio.shape[1], device=x_audio.device) + audio_te, audio_tm = self._encode_text("audio", encoder_hidden_states, pooled_projections, + audio_timestep, text_rope_pos) + audio_embed = self.audio_embeddings(x_audio) + audio_rope = self.audio_rope_embeddings(audio_rope_pos).unsqueeze(dim=0) + + for block in self.visual_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + visual_embed, audio_embed = torch.utils.checkpoint.checkpoint( + block, + visual_embed, + audio_embed, + video_te, + audio_te, + (video_tm, audio_tm), + visual_rope, + audio_rope, + sparse_params, + va_gate_scale, + av_gate_scale, + use_reentrant=False, + ) + else: + visual_embed, audio_embed = block(visual_embed, audio_embed, video_te, audio_te, + (video_tm, audio_tm), visual_rope, audio_rope, sparse_params, + va_gate_scale, av_gate_scale) + + visual_embed = fractal_unflatten(visual_embed, visual_shape, block_mask=to_fractal) + video_out = self.out_layer(visual_embed, video_tm) + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor] + if audio_embed is not None: + audio_out = self.audio_out_layer(audio_embed, audio_tm) + result = (video_out, audio_out) + else: + result = video_out + + if return_dict: + return Kandinsky6TransformerOutput(sample=result) + return result + + def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None: + for i, (axes_dim, ax_max_pos) in enumerate( + zip(self.visual_rope_embeddings.axes_dims, self.visual_rope_embeddings.max_pos, strict=True)): + name = f"args_{i}" + buf = getattr(self.visual_rope_embeddings, name, None) + if isinstance(buf, torch.Tensor) and buf.is_meta: + freq = _build_rotary_freqs(axes_dim // 2, self.visual_rope_embeddings.max_period).to(device=device) + pos = torch.arange(ax_max_pos, dtype=freq.dtype, device=device) + self.visual_rope_embeddings._buffers[name] = torch.outer(pos, freq) + + rope1d_modules: list[Kandinsky6RoPE1D] = [ + self.video_text_rope_embeddings, + self.audio_text_rope_embeddings, + self.audio_rope_embeddings, + ] + time_embeds = [self.video_time_embeddings, self.audio_time_embeddings] + + for rope1d in rope1d_modules: + if isinstance(rope1d.args, torch.Tensor) and rope1d.args.is_meta: + freq = _build_rotary_freqs(rope1d.dim // 2, rope1d.max_period).to(device=device) * rope1d.freqs_scaling + pos = torch.arange(rope1d.max_pos, dtype=freq.dtype, device=device) + rope1d._buffers["args"] = torch.outer(pos, freq) + + for time_embed in time_embeds: + if isinstance(time_embed.freqs, torch.Tensor) and time_embed.freqs.is_meta: + time_embed.freqs = _build_rotary_freqs(time_embed.model_dim // 2, + time_embed.max_period).to(device=device) + + +EntryClass = Kandinsky6Transformer3DModel diff --git a/fastvideo/models/dits/kandinsky6_sr.py b/fastvideo/models/dits/kandinsky6_sr.py new file mode 100644 index 000000000..7e110da8e --- /dev/null +++ b/fastvideo/models/dits/kandinsky6_sr.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 video super-resolution (SR) DiT. + +A text-free Kandinsky6 DiT: no text tower and no cross-attention; a learned ``pooled_bias`` stands in for the pooled +empty-caption embedding the model was trained with. The blocks reuse the Kandinsky6 building blocks +(``models/dits/kandinsky6.py``). Input ``[B, T, H, W, 2C + 1]`` (noised latent | anchor | anchor mask), output +``[B, T, H, W, out_visual_dim]`` (``n_grid * C`` for the pi-Flow checkpoint). + +The residual stream stays in the model's ``compute_dtype`` between blocks, as in the Diffusers model. Modulation, +norm and time-embedding maths happen in fp32 internally and are cast back to ``compute_dtype`` right after each gated +residual add. Sequence parallelism is not supported (dense ``LocalAttention``), and NABLA sparse attention is not wired. +""" +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRArchConfig, Kandinsky6SRConfig +from fastvideo.layers.layernorm import LayerNormScaleShift +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.kandinsky6 import ( + Kandinsky6Attention, + Kandinsky6FeedForward, + Kandinsky6Modulation, + Kandinsky6OutLayer, + Kandinsky6RoPE3D, + Kandinsky6TimeEmbeddings, + Kandinsky6VisualEmbeddings, + _build_rotary_freqs, +) +from fastvideo.platforms import AttentionBackendEnum + +logger = init_logger(__name__) + +_ARCH_CONFIG_DEFAULTS = Kandinsky6SRConfig().arch_config + +# Parameters used in fp32 maths (modulation, time embedding, QK norms) stay fp32 when the DiT is loaded in bf16, so a +# fp32 checkpoint is not rounded there. +_FP32_PARAM_MARKERS = ("pooled_bias", "time_embeddings.", ".visual_modulation.", "out_layer.modulation.", + "query_norm.", "key_norm.") + + +class Kandinsky6SRDecoderBlock(nn.Module): + """Text-free decoder block: modulated self-attention + modulated feed-forward (6 modulation params). + + ``x = (x + gate * attention(LN(x) * (1 + scale) + shift)).to(compute_dtype)`` per sub-layer: the gated + residual add happens in fp32 (matching the norm/modulation maths) and is cast back to ``compute_dtype`` + immediately after, so the residual stream tracks the model's working dtype between blocks instead of staying + fp32 forever after the first block. + """ + + def __init__(self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = "", + quant_config: QuantizationConfig | None = None): + super().__init__() + self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) + self.self_attention_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.self_attention = Kandinsky6Attention(model_dim, + head_dim, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.self_attention", + quant_config=quant_config) + self.feed_forward_norm = LayerNormScaleShift(model_dim, + norm_type="layer", + eps=1e-5, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32) + self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim, prefix=f"{prefix}.feed_forward", + quant_config=quant_config) + + def forward(self, x: torch.Tensor, time_embed: torch.Tensor, rope: torch.Tensor, + compute_dtype: torch.dtype) -> torch.Tensor: + # `compute_dtype` comes from the model's never-quantized embedding layer (see + # Kandinsky6SRTransformer3DModel.compute_dtype), not a local linear here: under FP8, + # `convert_model_to_fp8` pops `.weight` off quantized ``ReplicatedLinear`` layers like + # ``self_attention.to_query``, so reading its dtype directly would raise AttributeError. + sa_params, ff_params = torch.chunk(self.visual_modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + + shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) + h = self.self_attention_norm(x.float(), shift=shift, scale=scale, convert_modulation_dtype=True) + h = self.self_attention(h.to(compute_dtype), rotary_emb=rope) + x = (x.float() + gate.float() * h.float()).to(compute_dtype) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + h = self.feed_forward_norm(x.float(), shift=shift, scale=scale, convert_modulation_dtype=True) + h = self.feed_forward(h.to(compute_dtype)) + return (x.float() + gate.float() * h.float()).to(compute_dtype) + + +class Kandinsky6SROutLayer(Kandinsky6OutLayer): + """``Kandinsky6OutLayer`` for an fp32 residual stream: modulate in fp32, cast to the parameter dtype, project.""" + + def forward(self, visual_embed: torch.Tensor, time_embed: torch.Tensor) -> torch.Tensor: + shift, scale = torch.chunk(self.modulation(time_embed).unsqueeze(dim=1), 2, dim=-1) + h = (self.norm(visual_embed.float()) * (scale.float()[:, None, None] + 1.0) + shift.float()[:, None, None]) + x, _ = self.out_layer(h.to(self.out_layer.weight.dtype)) + + batch_size, duration, height, width, _ = x.shape + x = (x.view( + batch_size, + duration, + height, + width, + -1, + self.patch_size[0], + self.patch_size[1], + self.patch_size[2], + ).permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(1, 2).flatten(2, 3).flatten(3, 4)) + return x + + +class Kandinsky6SRTransformer3DModel(BaseDiT): + """``forward(hidden_states [B, T, H, W, 2C + 1], timestep [B], visual_rope_pos, scale_factor)`` + -> ``[B, T, H, W, out_visual_dim]``.""" + + _fsdp_shard_conditions = _ARCH_CONFIG_DEFAULTS._fsdp_shard_conditions + _compile_conditions = _ARCH_CONFIG_DEFAULTS._compile_conditions + param_names_mapping = _ARCH_CONFIG_DEFAULTS.param_names_mapping + reverse_param_names_mapping = _ARCH_CONFIG_DEFAULTS.reverse_param_names_mapping + lora_param_names_mapping = _ARCH_CONFIG_DEFAULTS.lora_param_names_mapping + _supported_attention_backends = _ARCH_CONFIG_DEFAULTS._supported_attention_backends + + def __init__(self, config: Kandinsky6SRConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + arch = config.arch_config + assert isinstance(arch, Kandinsky6SRArchConfig) + self._reject_unknown_config_keys(arch, hf_config) + quant_config = config.quant_config + self.quant_config = quant_config + + head_dim = sum(arch.axes_dims) + self.in_visual_dim = arch.in_visual_dim + self.model_dim = arch.model_dim + self.patch_size = arch.patch_size + self.input_channels = arch.input_channels + + self.time_embeddings = Kandinsky6TimeEmbeddings(arch.model_dim, arch.time_dim) + self.pooled_bias = nn.Parameter(torch.zeros(arch.time_dim)) + + self.visual_embeddings = Kandinsky6VisualEmbeddings(self.input_channels, arch.model_dim, arch.patch_size) + self.visual_rope_embeddings = Kandinsky6RoPE3D(arch.axes_dims) + self.visual_transformer_blocks = nn.ModuleList([ + Kandinsky6SRDecoderBlock(arch.model_dim, + arch.time_dim, + arch.ff_dim, + head_dim, + self._supported_attention_backends, + prefix=f"{config.prefix}.visual_transformer_blocks.{i}", + quant_config=quant_config) for i in range(arch.num_visual_blocks) + ]) + self.out_layer = Kandinsky6SROutLayer(arch.model_dim, arch.time_dim, arch.out_visual_dim, arch.patch_size) + if arch.attention_params is not None: + self._warn_if_sparse_attention_requested(arch.attention_params) + if arch.nabla_threshold is not None: + logger.warning( + "Kandinsky6SR checkpoint requests NABLA sparse attention via nabla_threshold=%s; the SR DiT runs " + "dense attention (sparse NABLA is not wired), so results differ slightly from the reference.", + arch.nabla_threshold) + + self.gradient_checkpointing = False + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + self.__post_init__() + + @staticmethod + def _reject_unknown_config_keys(arch: Kandinsky6SRArchConfig, hf_config: dict[str, Any]) -> None: + """``update_model_arch`` drops keys the arch config does not declare; a real checkpoint flag we do not know + would then silently run a different model, so refuse it.""" + unknown = sorted(set(hf_config) - arch.known_config_keys()) + if unknown: + raise ValueError(f"Kandinsky6SR transformer/config.json has unsupported keys {unknown}; refusing to " + "silently ignore them.") + + @staticmethod + def _warn_if_sparse_attention_requested(attention_params: dict) -> None: + for size, params in attention_params.items(): + kind = params.get("type") if isinstance(params, dict) else getattr(params, "type", None) + if kind not in (None, "flash", "sdpa", "dense"): + logger.warning( + "Kandinsky6SR checkpoint requests attention type %r for visual size %s; the SR DiT runs dense " + "attention (sparse NABLA is not wired), so results differ slightly from the reference.", kind, + size) + return + + @property + def compute_dtype(self) -> torch.dtype: + """Parameter dtype of the linear stacks (the dtype the DiT input is cast to).""" + return self.visual_embeddings.in_layer.weight.dtype + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + if any(marker in name for marker in _FP32_PARAM_MARKERS): + return torch.float32 + return default_dtype + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + visual_rope_pos: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | list[torch.Tensor], + scale_factor: tuple[float, ...] = (1.0, 1.0, 1.0), + return_dict: bool = False, + **kwargs: Any, + ) -> torch.Tensor: + """Run the DiT. + + Args: + hidden_states: ``[B, T, H, W, 2C + 1]`` noised latent | anchor | anchor mask. + timestep: ``[B]`` scheduler timestep (sigma * 1000). + visual_rope_pos: ``(arange(T), arange(H // ph), arange(W // pw))`` position indices (shared by the batch). + scale_factor: RoPE frequency scale per axis (t, h, w). + """ + if kwargs: + raise TypeError(f"Kandinsky6SRTransformer3DModel.forward got unsupported arguments {sorted(kwargs)}; the " + "SR DiT has no text / sparse-attention / LQ inputs.") + if hidden_states.ndim != 5: + raise ValueError(f"hidden_states must be [B,T,H,W,C], got {tuple(hidden_states.shape)}") + if hidden_states.shape[-1] != self.input_channels: + raise ValueError( + f"hidden_states has {hidden_states.shape[-1]} channels but the trained input layer expects " + f"{self.input_channels} (noised latent | anchor | anchor mask).") + + time_embed = self.time_embeddings(timestep) + time_embed = time_embed + self.pooled_bias.to(time_embed.dtype) + + visual_embed = self.visual_embeddings(hidden_states.to(self.compute_dtype).contiguous()) + batch_size, duration, height, width, dim = visual_embed.shape + visual_rope = self.visual_rope_embeddings((batch_size, duration, height, width), visual_rope_pos, scale_factor) + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(1, 3) + + compute_dtype = self.compute_dtype + for block in self.visual_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + visual_embed = torch.utils.checkpoint.checkpoint(block, visual_embed, time_embed, visual_rope, + compute_dtype, use_reentrant=False) + else: + visual_embed = block(visual_embed, time_embed, visual_rope, compute_dtype) + + visual_embed = visual_embed.reshape(batch_size, duration, height, width, dim) + return self.out_layer(visual_embed, time_embed) + + def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None: + rope = self.visual_rope_embeddings + for i, (axes_dim, ax_max_pos) in enumerate(zip(rope.axes_dims, rope.max_pos, strict=True)): + name = f"args_{i}" + buf = getattr(rope, name, None) + if isinstance(buf, torch.Tensor) and buf.is_meta: + freq = _build_rotary_freqs(axes_dim // 2, rope.max_period).to(device=device) + pos = torch.arange(ax_max_pos, dtype=freq.dtype, device=device) + rope._buffers[name] = torch.outer(pos, freq) + + time_module = self.time_embeddings + if isinstance(time_module.freqs, torch.Tensor) and time_module.freqs.is_meta: + time_module.freqs = _build_rotary_freqs(time_module.model_dim // 2, time_module.max_period).to(device=device) + + +EntryClass = Kandinsky6SRTransformer3DModel diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/models/encoders/clip.py index 479e74477..9bf28014e 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/models/encoders/clip.py @@ -427,7 +427,19 @@ class CLIPTextModel(TextEncoder): ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + + weights = list(weights) + # Some checkpoints (Kandinsky6's, for one) save CLIP's + # inner CLIPTextTransformer's own state_dict() directly, without the + # `text_model.` prefix this class's `self.text_model` wrapper puts on + # every parameter -- unlike other checkpoints already loaded through + # this class (e.g. Kandinsky5's), whose keys already include it. + # Detect and normalize once, up front, rather than guessing per key. + needs_text_model_prefix = bool(weights) and not any(name.startswith("text_model.") for name, _ in weights) + for name, loaded_weight in weights: + if needs_text_model_prefix: + name = f"text_model.{name}" # Handle q_proj, k_proj, v_proj -> qkv_proj mapping for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name in name: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 148a6d1bb..7e969d2a4 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -25,7 +25,7 @@ from fastvideo.attention.selector import ( coerce_attn_backend, record_resolved_attention_backend, ) -from fastvideo.configs.models import EncoderConfig +from fastvideo.configs.models import EncoderConfig, UpsamplerConfig from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.layers.quantization import get_quantization_config @@ -122,6 +122,9 @@ class ComponentLoader(ABC): # generic config-only loader. "spatial_upsampler": (UpsamplerLoader, "diffusers"), "temporal_upsampler": (UpsamplerLoader, "diffusers"), + # Kandinsky6 SR latent-upscaler bank (x2 / x4 entries). + # (the official Diffusers save_pretrained writes the library as "kandinsky6", the Hub repo as "diffusers") + "latent_upscaler": (UpsamplerLoader, ("diffusers", "kandinsky6")), } if module_type in module_loaders: @@ -131,7 +134,8 @@ class ComponentLoader(ABC): is_fastvideo_module = transformers_or_diffusers.startswith("fastvideo.") if not is_fastvideo_module: # Assert that the library matches what's expected for this module type - assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + allowed = (expected_library, ) if isinstance(expected_library, str) else tuple(expected_library) + assert transformers_or_diffusers in allowed, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" return loader_cls() # For unknown module types, use a generic loader @@ -756,8 +760,10 @@ class VAELoader(ComponentLoader): 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) - class_name = config.pop("_class_name") + class_name = config.pop("_class_name", None) config.pop("_name_or_path", None) + if class_name is None: + class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("vae") assert class_name is not None, ( "Model config does not contain a _class_name attribute. Only diffusers format is supported.") fastvideo_args.model_paths["vae"] = model_path @@ -909,7 +915,8 @@ class VAELoader(ComponentLoader): # Diffusers-format AutoencoderKL checkpoints should match exactly; load # strictly so missing/unexpected keys are surfaced early. - strict_load = class_name in {"AutoencoderKL", "AutoencoderKLMiniMaxH3"} + # 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) if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False) and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")): @@ -923,7 +930,18 @@ class AudioDecoderLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) - class_name = config.pop("_class_name", None) or "LTX2AudioDecoder" + class_name = (config.pop("_class_name", None) + or getattr(fastvideo_args, "_model_index_class_names", {}).get("audio_vae") + or "LTX2AudioDecoder") + # The K6 Diffusers MMAudioVAE's config.json still nests a `vocoder_config` (checked here, not + # used for construction -- Kandinsky6AudioVAE itself has no vocoder submodule; the pipeline + # loads a separate `vocoder` component instead, see Kandinsky6AudioDecodingStage). The + # checkpoint's audio_vae/*.safetensors correspondingly still bundles a redundant `vocoder.*` + # copy of those weights for backward compatibility with pre-split code; filtered out below. + bundled_mmaudio = class_name == "MMAudioVAE" and "vocoder_config" in config + if bundled_mmaudio: + class_name = "Kandinsky6AudioVAE" + config["need_vae_encoder"] = True model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() @@ -954,8 +972,11 @@ class AudioDecoderLoader(ComponentLoader): # fp32 and only then casts the whole feature utility module to bf16. # Constructing/loading directly in bf16 quantizes the unnormalized # checkpoint first and changes the decoded mel trajectory. - construction_precision = "fp32" if class_name == "MMAudioVAE" else precision - construction_device = torch.device("cpu") if class_name == "MMAudioVAE" else target_device + # Kandinsky6AudioVAE wraps MMAudioVAE internally (self.vae), so the + # same precision constraint applies to it too. + _needs_fp32_construction = class_name in ("MMAudioVAE", "Kandinsky6AudioVAE") + construction_precision = "fp32" if _needs_fp32_construction else precision + construction_device = torch.device("cpu") if _needs_fp32_construction else target_device with set_default_torch_dtype(PRECISION_TO_TYPE[construction_precision]): audio_decoder = model_cls(config).to(construction_device) @@ -964,9 +985,15 @@ class AudioDecoderLoader(ComponentLoader): for sf_file in safetensors_list: loaded.update(safetensors_load_file(sf_file)) - if class_name == "MMAudioVAE": + if class_name in ("MMAudioVAE", "Kandinsky6AudioVAE"): + if bundled_mmaudio: + loaded = {k: v for k, v in loaded.items() if not k.startswith("vocoder.")} audio_decoder.load_state_dict(loaded, strict=True) - audio_decoder.remove_weight_norm() + if bundled_mmaudio: + # Diffusers removes MPConv normalization before saving this export. + audio_decoder.vae._weights_normalized = True + else: + audio_decoder.remove_weight_norm() return audio_decoder.to(device=target_device, dtype=PRECISION_TO_TYPE[precision]).eval() decoder_state = {} @@ -981,12 +1008,36 @@ class AudioDecoderLoader(ComponentLoader): return audio_decoder.eval() +def _mmaudio_vocoder_to_bigvgan(class_name: str, config: dict) -> tuple[str, dict]: + """Diffusers' standalone MMAudioVocoder is BigVGAN-v2. Its config.json omits a few hyperparameters + that Diffusers' fixed-architecture reimplementation hardcodes instead of exposing (resblock type "1", + snakebeta activation with log-scale, no output bias/tanh); fill them in. A no-op for any other + class_name. + """ + if class_name != "MMAudioVocoder": + return class_name, config + config = dict(config) + config.setdefault("resblock", "1") + config.setdefault("activation", "snakebeta") + config.setdefault("snake_logscale", True) + config.setdefault("use_bias_at_final", False) + config.setdefault("use_tanh_at_final", False) + # The checkpoint's vocoder/*.safetensors is saved with weight_norm already removed (plain + # `conv_pre.weight`, not `conv_pre.parametrizations.weight.original0/1`). BigVGANV2's constructor + # applies weight_norm parametrization by default; strip it during construction (before + # VocoderLoader's subsequent load_state_dict(strict=True) below) so the fresh module's keys match + # what's on disk. + config.setdefault("weight_norm_removed", True) + return "BigVGANV2", config + + class VocoderLoader(ComponentLoader): """Loader for native vocoders.""" def load(self, model_path: str, fastvideo_args: FastVideoArgs): config = get_diffusers_config(model=model_path) class_name = config.pop("_class_name", None) or "LTX2Vocoder" + class_name, config = _mmaudio_vocoder_to_bigvgan(class_name, config) model_cls, _ = ModelRegistry.resolve_model_cls(class_name) target_device = get_local_torch_device() @@ -1033,8 +1084,13 @@ class TransformerLoader(ComponentLoader): """Load the transformer based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) hf_config = deepcopy(config) - cls_name = config.pop("_class_name") + cls_name = config.pop("_class_name", None) config.pop("_name_or_path", None) + if cls_name is None: + # Some checkpoints' component config.json omits _class_name even + # though model_index.json declares it (see composed_pipeline_base + # .load_modules); fall back to that before giving up. + cls_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("transformer") if cls_name is None: raise ValueError("Model config does not contain a _class_name attribute. " "Only diffusers format is supported.") @@ -1180,12 +1236,23 @@ class SchedulerLoader(ComponentLoader): """Load the scheduler based on the model path, and inference args.""" config = get_diffusers_config(model=model_path) - class_name = config.pop("_class_name") + class_name = config.pop("_class_name", None) + if class_name is None: + class_name = getattr(fastvideo_args, "_model_index_class_names", {}).get("scheduler") assert class_name is not None, ( "Model config does not contain a _class_name attribute. Only diffusers format is supported.") scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) + if getattr(scheduler_cls, "is_piflow", False): + overrides = { + name: value + for name in ("eps", "final_step_size_scale", "num_policy_substeps") + if (value := getattr(fastvideo_args.pipeline_config, "piflow_" + name, None)) is not None + } + # Retain the checkpoint shift, not the base K6 pipeline default. + return scheduler_cls.from_config(config, **overrides) + scheduler = scheduler_cls(**config) if fastvideo_args.pipeline_config.flow_shift is not None: scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift) @@ -1251,7 +1318,11 @@ class UpsamplerLoader(ComponentLoader): def load(self, model_path: str, fastvideo_args: FastVideoArgs): """Load the upsampler based on the model path, and inference args.""" config_dict = get_diffusers_config(model=model_path) - class_name = config_dict.pop("_class_name", None) + # Some bundles (e.g. Kandinsky6 SR's latent_upscaler) declare the class only in model_index.json, whose + # entries are keyed by the component directory name. + component_name = os.path.basename(os.path.normpath(str(model_path))) + class_name = (config_dict.pop("_class_name", None) + or getattr(fastvideo_args, "_model_index_class_names", {}).get(component_name)) if class_name is None: raise ValueError("Model config does not contain a _class_name attribute. " @@ -1264,6 +1335,7 @@ class UpsamplerLoader(ComponentLoader): # otherwise the LTX-2 branch below handles the single-class # path that takes the diffusers config dict directly. upsampler_config_attr = getattr(fastvideo_args.pipeline_config, "upsampler_config", None) + single_config = False if isinstance(upsampler_config_attr, list | tuple): try: upsampler_cfg = deepcopy(upsampler_config_attr[0]) @@ -1276,15 +1348,24 @@ class UpsamplerLoader(ComponentLoader): # `LTX2LatentUpsampler` wrapper takes the raw diffusers config # dict directly via LatentUpsamplerConfigurator. upsampler_cfg = deepcopy(config_dict) + elif isinstance(upsampler_config_attr, UpsamplerConfig) and type(upsampler_config_attr) is not UpsamplerConfig: + # A pipeline that declares one concrete upsampler config: fill it from the component's config.json. + config_dict.pop("_name_or_path", None) + upsampler_cfg = deepcopy(upsampler_config_attr) + upsampler_cfg.update_model_config(config_dict) + single_config = True else: raise AttributeError("pipeline_config.upsampler_config is missing; cannot build " f"upsampler config for class {class_name}") model_cls, _ = ModelRegistry.resolve_model_cls(class_name) - model = model_cls(upsampler_cfg) + upsampler_precision = getattr(fastvideo_args.pipeline_config, "upsampler_precision", "bf16") + # Build directly in the target precision: large upsamplers (the Kandinsky6 SR bank has 3.7B parameters) + # would otherwise be materialised and initialised in fp32 first. + with (set_default_torch_dtype(PRECISION_TO_TYPE[upsampler_precision]) if single_config else nullcontext()): + model = model_cls(upsampler_cfg) target_device = get_local_torch_device() - upsampler_precision = getattr(fastvideo_args.pipeline_config, "upsampler_precision", "bf16") model = model.to(target_device, dtype=PRECISION_TO_TYPE[upsampler_precision]) safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) @@ -1310,6 +1391,8 @@ class UpsamplerLoader(ComponentLoader): else: target_module.load_state_dict(loaded, strict=True) + if single_config: + model.requires_grad_(False) return model.eval() diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 135730a83..e859f3e8c 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -47,6 +47,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = { ), "Gen3CTransformer3DModel": ("dits", "gen3c", "Gen3CTransformer3DModel"), "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"), + "Kandinsky6Transformer3DModel": ("dits", "kandinsky6", "Kandinsky6Transformer3DModel"), "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"), } @@ -128,6 +129,7 @@ _VAE_MODELS = { _AUDIO_MODELS = { "MMAudioVAE": ("audio", "mmaudio_vae", "MMAudioVAE"), "BigVGANV2": ("audio", "bigvgan", "BigVGANV2"), + "Kandinsky6AudioVAE": ("audio", "kandinsky6_audio_vae", "Kandinsky6AudioVAE"), "LTX2AudioEncoder": ("audio", "ltx2_audio_vae", "LTX2AudioEncoder"), "LTX2AudioDecoder": ("audio", "ltx2_audio_vae", "LTX2AudioDecoder"), "LTX2Vocoder": ("audio", "ltx2_audio_vae", "LTX2Vocoder"), diff --git a/fastvideo/models/schedulers/scheduling_piflow.py b/fastvideo/models/schedulers/scheduling_piflow.py new file mode 100644 index 000000000..633e1bf96 --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_piflow.py @@ -0,0 +1,319 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Diffusers scheduler for distilled Kandinsky 6 PiFlow checkpoints.""" + +from __future__ import annotations + +import torch +from diffusers.configuration_utils import register_to_config +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, + FlowMatchEulerDiscreteSchedulerOutput, +) + + +class DXPolicy: + """Network-free DX policy over one flow-matching segment.""" + + def __init__( # noqa: PLR0913 + self, + denoising_output: torch.Tensor, + x_t_src: torch.Tensor, + sigma_t_src: torch.Tensor, + segment_size: float | torch.Tensor = 1.0, + shift: float = 1.0, + mode: str = "grid", + eps: float = 1e-4, + ) -> None: + self.ndim = x_t_src.dim() + self.shift = shift + self.eps = eps + if mode not in ("grid", "polynomial"): + raise ValueError(f"Unknown mode: {mode}") + self.mode = mode + + sigma_t_src = sigma_t_src.reshape(*sigma_t_src.size(), *((self.ndim - sigma_t_src.dim()) * [1])) + self.raw_t_src = self._unwarp_t(sigma_t_src) + segment = segment_size + if isinstance(segment, torch.Tensor) and segment.dim() < self.raw_t_src.dim(): + segment = segment.reshape(*segment.size(), *((self.raw_t_src.dim() - segment.dim()) * [1])) + self.raw_t_dst = (self.raw_t_src - segment).clamp(min=0) + self.segment_size = (self.raw_t_src - self.raw_t_dst).clamp(min=eps) + self.denoising_output_x_0 = x_t_src.unsqueeze(1) - sigma_t_src.unsqueeze(1) * denoising_output + + @staticmethod + def _interpolate(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + n = x.size(1) + if n < 2: # noqa: PLR2004 + return x.squeeze(1) + t = t.clamp(min=0, max=1) * (n - 1) + t0 = t.floor().to(torch.long).clamp(min=0, max=n - 2) + t1 = t0 + 1 + indices = torch.stack([t0, t1], dim=1) + values = torch.gather(x, dim=1, index=indices.expand(-1, -1, *x.shape[2:])) + return (t1 - t) * values[:, 0] + (t - t0) * values[:, 1] + + def _unwarp_t(self, sigma_t: torch.Tensor) -> torch.Tensor: + return sigma_t / (self.shift + (1 - self.shift) * sigma_t) + + def pi(self, x_t: torch.Tensor, sigma_t: torch.Tensor) -> torch.Tensor: + sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1])) + raw_t = self._unwarp_t(sigma_t) + if self.mode == "grid": + x_0 = self._interpolate( + self.denoising_output_x_0, + (raw_t - self.raw_t_dst) / self.segment_size, + ) + else: + p_order = self.denoising_output_x_0.size(1) + diff_t = self.raw_t_src - raw_t + basis = torch.stack([diff_t**i for i in range(p_order)], dim=1) + x_0 = torch.sum(basis * self.denoising_output_x_0, dim=1) + return (x_t - x_0) / sigma_t.clamp(min=self.eps) + + +def shift_timesteps(t: torch.Tensor, shift: float) -> torch.Tensor: + """Map raw flow-matching time to the shifted DiT time.""" + return shift * t / (1 + (shift - 1) * t) + + +def policy_rollout_fm( # noqa: PLR0913 + x_t_start: torch.Tensor, + sigma_t_start: torch.Tensor, + raw_t_start: torch.Tensor, + raw_t_end: torch.Tensor, + total_substeps: int, + policy: DXPolicy, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Integrate ``policy.pi`` from ``raw_t_start`` to ``raw_t_end``.""" + num_batches = x_t_start.size(0) + ndim = x_t_start.dim() + shape = (num_batches, *((ndim - 1) * [1])) + raw_t_start = raw_t_start.reshape(shape) + raw_t_end = raw_t_end.reshape(shape) + sigma_t = sigma_t_start.reshape(shape) + + delta_raw_t = raw_t_start - raw_t_end + num_substeps = (delta_raw_t * total_substeps).round().to(torch.long).clamp(min=1) + substep_size = delta_raw_t / num_substeps + max_num_substeps = num_substeps.max() + + raw_t = raw_t_start + x_t = x_t_start + for substep_id in range(max_num_substeps.item()): + velocity = policy.pi(x_t, sigma_t) + raw_t_minus = (raw_t - substep_size).clamp(min=0) + sigma_t_minus = shift_timesteps(raw_t_minus, policy.shift) + x_t_minus = x_t + velocity * (sigma_t_minus - sigma_t) + + active_mask = num_substeps > substep_id + x_t = torch.where(active_mask, x_t_minus, x_t) + sigma_t = torch.where(active_mask, sigma_t_minus, sigma_t) + raw_t = torch.where(active_mask, raw_t_minus, raw_t) + + return x_t, sigma_t, sigma_t.flatten() * 1_000 + + +class PiflowScheduler(FlowMatchEulerDiscreteScheduler): + """Few-step PiFlow scheduler for widened-output diffusion transformers. + + PiFlow evaluates the denoising model at a small number of grid points and + integrates a network-free policy between those evaluations. The scheduler + is intended for distilled Kandinsky 6 checkpoints, including the main + video/audio model and the video super-resolution model. Their model output + contains ``n_grid`` predictions per sample channel. + + Args: + num_train_timesteps (`int`, *optional*, defaults to 1000): Number of + training diffusion steps. + shift (`float`, *optional*, defaults to 5.0): Flow-matching timestep shift. + n_grid (`int`, *optional*, defaults to 10): Number of predictions in the + widened model output. + nfe (`int`, *optional*): Number of model evaluations used at inference. + eps (`float`, *optional*, defaults to 1e-6): Minimum timestep and policy denominator. + final_step_size_scale (`float`, *optional*, defaults to 0.5): Relative + size of the final raw-timestep segment. + num_policy_substeps (`int`, *optional*, defaults to 128): Maximum policy + integration substeps per raw-timestep unit. + """ + + is_piflow = True + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 5.0, + n_grid: int = 10, + nfe: int | None = None, + eps: float = 1e-6, + final_step_size_scale: float = 0.5, + num_policy_substeps: int = 128, + ) -> None: + if n_grid < 2: + raise ValueError(f"PiflowScheduler requires n_grid >= 2, got {n_grid}") + if eps <= 0: + raise ValueError(f"PiflowScheduler requires eps > 0, got {eps}") + if not 0 < final_step_size_scale <= 1: + raise ValueError("PiflowScheduler requires 0 < final_step_size_scale <= 1") + if num_policy_substeps < 1: + raise ValueError("PiflowScheduler requires num_policy_substeps >= 1") + super().__init__( + num_train_timesteps=num_train_timesteps, + shift=shift, + ) + self.n_grid = int(n_grid) + self.eps = float(eps) + self.final_step_size_scale = float(final_step_size_scale) + self.num_policy_substeps = int(num_policy_substeps) + self._piflow_raw_timesteps = torch.empty(0) + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, + ) -> None: + """Set the distilled PiFlow timestep schedule. + + Args: + num_inference_steps (`int`): Number of model evaluations. + device (`str` or `torch.device`, *optional*): Device for the schedule. + sigmas (`list[float]`, *optional*): Unsupported custom sigma schedule. + mu (`float`, *optional*): Unsupported dynamic-shift parameter. + timesteps (`list[float]`, *optional*): Unsupported custom timestep schedule. + """ + if sigmas is not None or mu is not None or timesteps is not None: + raise ValueError("PiflowScheduler only supports its configured distilled timestep schedule") + if num_inference_steps is None or num_inference_steps < 1: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}") + one_minus_final = 1.0 - self.final_step_size_scale + segment = 1.0 / (num_inference_steps - one_minus_final) + raw = 1.0 - torch.arange(num_inference_steps, dtype=torch.float32, device=device) * segment + sigmas = shift_timesteps(raw, float(self.config.shift)) + self.num_inference_steps = int(num_inference_steps) + self._piflow_raw_timesteps = raw + self.timesteps = sigmas * self.config.num_train_timesteps + self.sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]) + self._step_index = None + self._begin_index = None + + def _to_grid(self, model_output: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + if model_output.ndim != sample.ndim or model_output.shape[:-1] != sample.shape[:-1]: + raise ValueError( + "Piflow model output must match sample shape except for the output channels: " + f"got {tuple(model_output.shape)} for sample {tuple(sample.shape)}" + ) + if model_output.shape[-1] % self.n_grid != 0: + raise ValueError( + f"Piflow model output channels {model_output.shape[-1]} are not divisible by n_grid={self.n_grid}" + ) + output_dim = model_output.shape[-1] // self.n_grid + if output_dim != sample.shape[-1]: + raise ValueError( + "Piflow model output channels do not match the sample: " + f"expected {sample.shape[-1] * self.n_grid}, got {model_output.shape[-1]}" + ) + return model_output.reshape(*model_output.shape[:-1], self.n_grid, output_dim).movedim(-2, 1) + + def _policy_step(self, model_output: torch.Tensor, sample: torch.Tensor, step_index: int) -> torch.Tensor: + model_output = self._to_grid(model_output, sample) + raw_src = self._piflow_raw_timesteps[step_index].to(device=sample.device) + raw_dst = ( + self._piflow_raw_timesteps[step_index + 1] + if step_index + 1 < self._piflow_raw_timesteps.numel() + else self._piflow_raw_timesteps.new_full((), self.eps) + ).to(device=sample.device) + sigma_src = self.sigmas[step_index].to(device=sample.device) + token_shape = (sample.shape[0], *((sample.ndim - 1) * [1])) + sigma = sigma_src.expand(sample.shape[0]).reshape(token_shape) + segment = (raw_src - raw_dst).expand(sample.shape[0]) + policy = DXPolicy( + model_output, + sample, + sigma, + segment, + shift=float(self.config.shift), + mode="grid", + eps=self.eps, + ) + updated, _, _ = policy_rollout_fm( + sample, + sigma, + raw_src.expand(sample.shape[0]), + raw_dst.expand(sample.shape[0]), + self.num_policy_substeps, + policy, + ) + return updated.to(dtype=sample.dtype) + + def _step_index_for(self, timestep: torch.Tensor | float) -> int: + if self.step_index is None: + if self.begin_index is not None: + self._step_index = self.begin_index + else: + schedule_timesteps = self.timesteps + timestep = torch.as_tensor( + timestep, + device=schedule_timesteps.device, + dtype=schedule_timesteps.dtype, + ) + indices = torch.nonzero(schedule_timesteps == timestep).flatten() + if not indices.numel(): + raise ValueError(f"timestep {timestep.item()} is not in the Piflow schedule") + position = 1 if indices.numel() > 1 else 0 + self._step_index = int(indices[position].item()) + if self.step_index is None or self.step_index >= self.num_inference_steps: + raise RuntimeError("PiflowScheduler.step called after the schedule was exhausted") + return int(self.step_index) + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + """Advance one step by integrating the PiFlow policy. + + Args: + model_output (`torch.FloatTensor`): Widened model output containing + ``n_grid`` predictions per sample channel. + timestep (`float` or `torch.FloatTensor`): Current scheduler timestep. + sample (`torch.FloatTensor`): Current noisy sample. + return_dict (`bool`, *optional*, defaults to True): Whether to return + a [`FlowMatchEulerDiscreteSchedulerOutput`]. + + Returns: + [`FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: Updated sample. + """ + if isinstance(timestep, int) or isinstance(timestep, (torch.IntTensor, torch.LongTensor)): + raise ValueError( + "Passing integer indices as timesteps to PiflowScheduler.step() is not supported; " + "pass a value from scheduler.timesteps instead" + ) + step_index = self._step_index_for(timestep) + # PiFlow's policy rollout performs its update in float32. Keep that + # precision across outer steps, matching the native sampler. + updated = self._policy_step(model_output, sample.to(torch.float32), step_index) + self._step_index += 1 + if return_dict: + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) + return (updated,) + + +EntryClass = PiflowScheduler diff --git a/fastvideo/models/upsamplers/kandinsky6_sr.py b/fastvideo/models/upsamplers/kandinsky6_sr.py new file mode 100644 index 000000000..0f8d27d41 --- /dev/null +++ b/fastvideo/models/upsamplers/kandinsky6_sr.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR latent upscaler (LU): a bank of convolutional upsamplers of scaled KVAE latents. + +Each bank entry is a cascade of two 2x stages over ``[B, C, T, h, w]`` latents:: + + x4: input_proj -> pre_blocks -> upsample_1 -> mid_blocks -> upsample_2 -> post_blocks -> output_proj + x2: mid_input_proj -> x2_branch (adapter -> finisher -> private mid stage -> private second stage) + +Every norm is an RMSNorm FiLM-modulated by the input latent itself (``zq``, nearest-resized to the feature grid), the +3x3x3 convs pad time by repeating the edge frame, and the 2x upsample is ``Conv1x1(up + Conv(1,3,3)(up))`` with +``up`` the nearest-neighbour 2x resize. The bank holds one entry per served scale under ``_models.`` in config.scales order. + +Module names match the re-keyed Diffusers checkpoint; no load-time compatibility hook is required. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from fastvideo.configs.models.upsamplers.kandinsky6_sr import (Kandinsky6SRLatentUpscalerConfig, + Kandinsky6SRLatentUpscalerEntryConfig) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class ReplicateTimeConv3d(nn.Conv3d): + """``Conv3d`` with 'same' padding that repeats the edge frame along T and zero-pads H and W. + + ``nn.Conv3d`` has a single padding mode for all axes, so T is padded explicitly before the convolution. + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> None: + pad = kernel_size // 2 + super().__init__(in_channels, out_channels, kernel_size, padding=(0, pad, pad)) + self.temporal_pad = pad + + def forward(self, x: Tensor) -> Tensor: + if self.temporal_pad: + x = F.pad(x, (0, 0, 0, 0, self.temporal_pad, self.temporal_pad), mode="replicate") + return super().forward(x) + + +class RMSNorm(nn.Module): + """Channel RMS norm of ``[B, C, T, H, W]`` features with a learnable per-channel gain.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim, 1, 1, 1)) + + def forward(self, x: Tensor) -> Tensor: + # Upcast to float32 around the normalize (matches the diffusers reference's + # Kandinsky6SRLatentUpscalerRMSNorm, added as a precision hardening fix): without it, bf16 activations + # drift from the reference across the many ResidualBlock/X2Branch calls that route through this norm. + normalized = F.normalize(x.float(), dim=1).to(x.dtype) + return normalized * self.scale * self.gamma + + +class ModulatedRMSNorm(nn.Module): + """``RMSNorm(x) * conv_y(zq) + conv_b(zq)`` with 1x1x1 convs of the conditioning latent ``zq``.""" + + def __init__(self, dim: int, zq_dim: int) -> None: + super().__init__() + self.norm = RMSNorm(dim) + self.conv_y = nn.Conv3d(zq_dim, dim, kernel_size=1) + self.conv_b = nn.Conv3d(zq_dim, dim, kernel_size=1) + + def forward(self, x: Tensor, zq: Tensor) -> Tensor: + if zq.shape[2:] != x.shape[2:]: + zq = F.interpolate(zq, size=x.shape[2:], mode="nearest") + return self.norm(x) * self.conv_y(zq) + self.conv_b(zq) + + +class ResidualBlock(nn.Module): + """Pre-activation block ``norm1 -> SiLU -> conv1 -> norm2 -> SiLU -> conv2`` plus a (1x1x1 if narrowing) skip.""" + + def __init__(self, in_channels: int, out_channels: int, mid_channels: int, zq_dim: int) -> None: + super().__init__() + self.norm1 = ModulatedRMSNorm(in_channels, zq_dim) + self.conv1 = ReplicateTimeConv3d(in_channels, mid_channels, kernel_size=3) + self.norm2 = ModulatedRMSNorm(mid_channels, zq_dim) + self.conv2 = ReplicateTimeConv3d(mid_channels, out_channels, kernel_size=3) + self.shortcut: nn.Module = (nn.Identity() if in_channels == out_channels else ReplicateTimeConv3d( + in_channels, out_channels, kernel_size=1)) + + def forward(self, x: Tensor, zq: Tensor) -> Tensor: + h = F.silu(self.norm1(x, zq)) + h = self.conv1(h) + h = F.silu(self.norm2(h, zq)) + h = self.conv2(h) + return self.shortcut(x) + h + + +def nearest_2x(x: Tensor) -> Tensor: + """Nearest-neighbour 2x resize of H and W of ``[B, C, T, H, W]`` (T folded into the batch).""" + b, c, t, h, w = x.shape + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = F.interpolate(x, scale_factor=2, mode="nearest") + return x.reshape(b, t, c, 2 * h, 2 * w).permute(0, 2, 1, 3, 4) + + +class PXSUpsample(nn.Module): + """2x spatial upsample ``linear(up + spatial_conv(up))`` with ``up = nearest_2x(x)``.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.spatial_conv = nn.Conv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + self.linear = nn.Conv3d(channels, channels, kernel_size=1) + + def forward(self, x: Tensor) -> Tensor: + up = nearest_2x(x) + return self.linear(up + self.spatial_conv(up)) + + +class X2Finisher(nn.Module): + """The ``PXSUpsample`` conv pair applied on the unchanged grid (no resize).""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.linear = nn.Conv3d(channels, channels, kernel_size=1) + self.spatial_conv = nn.Conv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + + def forward(self, x: Tensor) -> Tensor: + return self.linear(x + self.spatial_conv(x)) + + +def _stem(in_channels: int, out_channels: int) -> nn.Sequential: + return nn.Sequential(ReplicateTimeConv3d(in_channels, out_channels, kernel_size=3)) + + +class OutputHead(nn.Module): + """Named norm/activation/conv projection used by current Diffusers weights.""" + + def __init__(self, channels: int, out_channels: int, zq_dim: int) -> None: + super().__init__() + self.norm = ModulatedRMSNorm(channels, zq_dim) + self.activation = nn.SiLU() + self.conv = ReplicateTimeConv3d(channels, out_channels, kernel_size=3) + + def forward(self, x: Tensor, zq: Tensor) -> Tensor: + return self.conv(self.activation(self.norm(x, zq))) + + +def _residual_stack(count: int, in_channels: int, channels: int, expand_ratio: int, zq_dim: int) -> nn.Sequential: + """``count`` blocks at width ``channels``; the first one narrows from ``in_channels``.""" + return nn.Sequential(*(ResidualBlock(in_channels if index == 0 else channels, channels, channels * expand_ratio, + zq_dim) for index in range(count))) + + +def _run_blocks(blocks: nn.Sequential, x: Tensor, zq: Tensor) -> Tensor: + for block in blocks: + x = block(x, zq) + return x + + +def _second_stage(upsample: nn.Module, blocks: nn.Sequential, head: OutputHead, x: Tensor, zq: Tensor) -> Tensor: + """``upsample -> blocks -> head``, with the head's first norm conditioned on ``zq``.""" + x = _run_blocks(blocks, upsample(x), zq) + return head(x, zq) + + +class X2Branch(nn.Module): + """Weights exclusive to the x2 path: a residual adapter at the input grid, a finisher, and private copies of the + mid stage and second stage (so the x2 path shares no weights with the x4 path).""" + + def __init__(self, config: Kandinsky6SRLatentUpscalerEntryConfig) -> None: + super().__init__() + c = config.in_channels + w1, w2, w3 = config.stage_channels + er = config.expand_ratio + self.adapter = _residual_stack(config.x2_adapter_blocks, w1, w1, er, c) + self.finisher = X2Finisher(w1) + self.upsample = PXSUpsample(w2) + self.blocks = _residual_stack(config.num_post_blocks, w2, w3, er, c) + self.output_proj = OutputHead(w3, c, zq_dim=c) + self.mid_blocks = _residual_stack(config.num_mid_blocks, w1, w2, er, c) + + def forward(self, x: Tensor, zq: Tensor) -> Tensor: + x = _run_blocks(self.adapter, x, zq) + x = self.finisher(x) + x = _run_blocks(self.mid_blocks, x, zq) + return _second_stage(self.upsample, self.blocks, self.output_proj, x, zq) + + +class Kandinsky6SRLatentUpscaler(nn.Module): + """One bank entry: ``[B, C, T, h, w]`` -> ``[B, C, T, s*h, s*w]`` for ``s`` = 4, or 2 with ``enable_x2_entry``.""" + + def __init__(self, config: Kandinsky6SRLatentUpscalerEntryConfig) -> None: + super().__init__() + c = config.in_channels + w1, w2, w3 = config.stage_channels + er = config.expand_ratio + self.input_proj = _stem(c, w1) + # The 2x deep-supervision head of training: part of the checkpoint, never run at inference. + self.mid_output_head = nn.Sequential(RMSNorm(w2), nn.SiLU(), ReplicateTimeConv3d(w2, c, kernel_size=3)) + self.output_proj = OutputHead(w3, c, zq_dim=c) + self.upsample_1 = PXSUpsample(w1) + self.upsample_2 = PXSUpsample(w2) + self.pre_blocks = _residual_stack(config.num_pre_blocks, w1, w1, er, c) + self.mid_blocks = _residual_stack(config.num_mid_blocks, w1, w2, er, c) + self.post_blocks = _residual_stack(config.num_post_blocks, w2, w3, er, c) + self.mid_input_proj: nn.Sequential | None = None + self.x2_branch: X2Branch | None = None + if config.enable_x2_entry: + self.mid_input_proj = _stem(c, config.hidden_channels) + self.x2_branch = X2Branch(config) + + def forward(self, z: Tensor, scale: int) -> Tensor: + zq = z + if scale == 2: + if self.x2_branch is None or self.mid_input_proj is None: + raise ValueError("this latent-upscaler entry has no x2 path (enable_x2_entry=false)") + return self.x2_branch(self.mid_input_proj(z), zq) + if scale != 4: + raise ValueError(f"latent-upscaler entries upsample by 2 or 4, got scale={scale!r}") + x = _run_blocks(self.pre_blocks, self.input_proj(z), zq) + x = _run_blocks(self.mid_blocks, self.upsample_1(x), zq) + return _second_stage(self.upsample_2, self.post_blocks, self.output_proj, x, zq) + + +class Kandinsky6SRLatentUpscalerBank(nn.Module): + """The x2 / x4 latent-upscaler bank of a Kandinsky6 SR bundle (``latent_upscaler/``). + + ``forward(z, scale)`` upsamples an already scaled (``latent * scaling_factor``) KVAE latent ``[B, C, T, h, w]`` by + ``scale`` in H and W with the entry serving that scale. + """ + + def __init__(self, config: Kandinsky6SRLatentUpscalerConfig) -> None: + super().__init__() + if not config.models: + raise ValueError("latent upscaler config declares no `models` entries") + self.config = config + self.scaling_factor = config.scaling_factor + entries = {entry.target_scale: entry for entry in config.models} + if set(entries) != set(config.scales): + raise ValueError(f"Latent upscaler models {tuple(entries)} must match scales={config.scales}") + self._models = nn.ModuleList([Kandinsky6SRLatentUpscaler(entries[scale]) for scale in config.scales]) + + @property + def scales(self) -> tuple[int, ...]: + """The upscale factors this bank serves, ascending.""" + return tuple(sorted(self.config.scales)) + + def forward(self, z: Tensor, scale: int) -> Tensor: + if not float(scale).is_integer() or int(scale) not in self.config.scales: + raise ValueError(f"the latent-upscaler bank has no x{scale} entry (scales: {self.scales})") + return self._models[self.config.scales.index(int(scale))](z, int(scale)) + + +EntryClass = Kandinsky6SRLatentUpscalerBank diff --git a/fastvideo/models/vaes/kandinsky6_sr.py b/fastvideo/models/vaes/kandinsky6_sr.py new file mode 100644 index 000000000..2134054eb --- /dev/null +++ b/fastvideo/models/vaes/kandinsky6_sr.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR causal video VAE (KVAE): spatial x16, temporal x4 (causal, ``1 + 4k`` frames). + +Clips are encoded / decoded in temporal segments (16 pixel frames, the first one 17; 4 latent frames, the first one 5). +Every causal conv carries its last input frames to the next segment, so memory is bounded by one segment whatever the +clip length, and the segmentation is part of the numerics: it fixes which frames the first-frame special cases see. + +Conventions: ``normalize_data`` maps ``[0, 255]`` pixels to ``x / 128 - 1``; ``encode`` returns ``(latent, +split_list)`` where ``latent`` is the raw (unscaled) posterior mean and ``split_list`` the pixel-frame segment sizes; +``decode`` returns an object whose ``.sample`` is in the normalized pixel range. State-dict keys are the checkpoint's +own ``encoder.*`` / ``decoder.*`` keys, so the loader loads strictly without renaming. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEArchConfig, Kandinsky6SRVAEConfig +from fastvideo.models.vaes.wanvae import WanRMS_norm + +# Pixel frames per encode segment after the first (which has one more: the causal first frame). +_SEGMENT_FRAMES = 16 +# Inputs above this many elements are convolved in temporal chunks: at 4x SR decode tiles (~1536x832) the +# full-resolution 256-channel convs of the last decoder level see ~5.6e9 elements per segment, beyond 32-bit +# indexing, and chunking also keeps the transient conv workspace to a fraction of the activation. +_MAX_CONV_NUMEL = 2 * 10**9 + + +@dataclass +class DecoderOutput: + sample: torch.Tensor + + +class _SegmentCache: + """Causal state carried from one temporal segment to the next during a single encode / decode call.""" + + def __init__(self) -> None: + self.first = True + self.padding: dict[nn.Module, torch.Tensor] = {} + + +def _silu(x: torch.Tensor) -> torch.Tensor: + # The output norms use this form and the resblocks use F.silu; the two are not bitwise equal in bf16. + return x * torch.sigmoid(x) + + +class _K6RMSNorm(WanRMS_norm): + """``WanRMS_norm`` with an explicit float32 upcast around the normalize, matching the diffusers reference's + ``Kandinsky6SRRMSNorm`` (which upcasts unconditionally, not only for fp16/bf16/fp8 inputs): at bf16/fp16 + activations, normalizing without upcasting measurably drifts from the reference over many resnet blocks.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + normed = F.normalize(x.float(), dim=1).to(x.dtype) + return normed * self.scale * self.gamma + self.bias + + +class _ChunkedConv3d(nn.Conv3d): + """``nn.Conv3d`` that convolves inputs larger than ``_MAX_CONV_NUMEL`` in overlapping temporal chunks.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + num_parts = math.ceil(x.numel() / _MAX_CONV_NUMEL) + if num_parts <= 1: + return super().forward(x) + k, frames = self.kernel_size[0], x.size(2) + if self.padding[0] != 0: + raise ValueError(f"temporal chunking needs an unpadded time axis, got kernel {self.kernel_size}, " + f"padding {self.padding}") + if self.stride[0] != 1: + # A strided conv can't reuse the stride-1 path's "prefix each chunk with the previous + # chunk's tail" trick below (a strided conv's output doesn't align with input chunk + # boundaries), so slide one kernel-width window at a time instead -- matches the diffusers + # reference's Kandinsky6SRSafeConv3d, only raising if even a single window would overflow. + stride = self.stride[0] + window_numel = x.shape[0] * x.shape[1] * k * x.shape[3] * x.shape[4] + if window_numel >= _MAX_CONV_NUMEL: + raise ValueError(f"frames are too big for Conv3d even one stride-{stride} window at a time " + f"({window_numel} elements per window)") + # Bare super() doesn't resolve inside a list comprehension's own implicit scope, hence the loop. + window_outputs = [] + for i in range(0, frames - k + 1, stride): + window_outputs.append(super().forward(x[:, :, i:i + k])) + return torch.cat(window_outputs, dim=2) + step = math.ceil(frames / num_parts) + last = frames - step * (math.ceil(frames / step) - 1) + if k > 1 and min(step, last) < k: + # A chunk would be shorter than the kernel: convolve one output frame at a time instead. + windows = [(i, i + k) for i in range(frames - k + 1)] + else: + # Each input chunk [start, end) is prefixed with the previous chunk's last k - 1 frames. + windows = [(max(start - k + 1, 0), min(start + step, frames)) for start in range(0, frames, step)] + out: torch.Tensor | None = None + for lo, hi in windows: + y = super().forward(x[:, :, lo:hi]) + if out is None: + out = y.new_empty((*y.shape[:2], frames - k + 1, *y.shape[3:])) + out[:, :, lo:lo + y.size(2)] = y + assert out is not None + return out + + +class CausalConv3d(nn.Module): + """Conv3d that is causal in time: the first segment is left-padded with copies of its first frame, later + segments with the tail of the previous one; height / width are zero padded.""" + + def __init__(self, chan_in: int, chan_out: int, kernel_size: int | tuple[int, int, int], + stride: tuple[int, int, int] = (1, 1, 1)) -> None: + super().__init__() + if isinstance(kernel_size, int): + kernel_size = (kernel_size, ) * 3 + time_kernel, height_kernel, width_kernel = kernel_size + self.height_pad = height_kernel // 2 + self.width_pad = width_kernel // 2 + self.time_pad = time_kernel - 1 + self.time_kernel = time_kernel + self.time_stride = stride[0] + self.conv = _ChunkedConv3d(chan_in, chan_out, kernel_size, stride=stride) + + def forward(self, x: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + s, k = self.time_stride, self.time_kernel + batch, _, frames, height, width = x.shape + x = F.pad(x, (self.width_pad, self.width_pad, self.height_pad, self.height_pad, 0, 0)) + if cache.first: + first_frame = x[:, :, :1] + padding = first_frame.expand(-1, -1, self.time_pad, -1, -1) + else: + padding = cache.padding[self] + + # conv(cat([padding, x])) without materializing the concatenation of the whole segment: only the outputs + # whose window overlaps ``padding`` are computed on a small concatenated head. + out_frames = frames if s == 1 else (frames + 1) // 2 + output = x.new_empty((batch, self.conv.out_channels, out_frames, height, width)) + head_out = math.ceil(padding.size(2) / s) + head_in = head_out * s - padding.size(2) + if head_out > 0: + output[:, :, :head_out] = self.conv(torch.cat([padding, x[:, :, :head_in + k - s]], dim=2)) + if head_out < out_frames: + output[:, :, head_out:] = self.conv(x[:, :, head_in:]) + + # Keep exactly the input frames the next segment's first output window needs. + pad_offset = head_in + s * math.trunc((frames - head_in - k) / s) + s + if pad_offset < 0: # segment shorter than the kernel + cache.padding[self] = torch.cat([padding[:, :, pad_offset:], x], dim=2) + else: + cache.padding[self] = x[:, :, pad_offset:].clone() + return output + + +class ResnetBlock3D(nn.Module): + """norm -> SiLU -> causal conv, twice, plus a 1x1 shortcut when the width changes. Decoder blocks use the + latent-conditioned ``SpatialNorm3D`` (``zq_channels`` set).""" + + def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None = None) -> None: + super().__init__() + if zq_channels is None: + self.norm1: nn.Module = _K6RMSNorm(in_channels, images=False) + self.norm2: nn.Module = _K6RMSNorm(out_channels, images=False) + else: + self.norm1 = SpatialNorm3D(in_channels, zq_channels) + self.norm2 = SpatialNorm3D(out_channels, zq_channels) + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3) + self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3) + if in_channels != out_channels: + self.nin_shortcut = _ChunkedConv3d(in_channels, out_channels, kernel_size=1) + + def _norm(self, norm: nn.Module, h: torch.Tensor, zq: torch.Tensor | None, cache: _SegmentCache) -> torch.Tensor: + return norm(h) if zq is None else norm(h, zq, cache) + + def forward(self, x: torch.Tensor, cache: _SegmentCache, zq: torch.Tensor | None = None) -> torch.Tensor: + h = F.silu(self._norm(self.norm1, x, zq, cache), inplace=True) + h = self.conv1(h, cache) + h = F.silu(self._norm(self.norm2, h, zq, cache), inplace=True) + h = self.conv2(h, cache) + if hasattr(self, "nin_shortcut"): + x = self.nin_shortcut(x) + return x + h + + +def _chunked_interpolate_nearest(x: torch.Tensor, size: tuple[int, int, int], channels: int = 32) -> torch.Tensor: + """``F.interpolate(x, size, mode="nearest")`` in ``channels``-wide chunks along dim 1, to bound peak memory at + large SR-decode tile sizes (matches the diffusers reference's chunked ``zq`` interpolation).""" + if x.shape[1] <= channels: + return F.interpolate(x, size=size, mode="nearest") + return torch.cat([F.interpolate(chunk, size=size, mode="nearest") for chunk in torch.split(x, channels, dim=1)], + dim=1) + + +class SpatialNorm3D(nn.Module): + """RMS norm modulated by the latent: ``norm(f) * conv_y(zq) + conv_b(zq)`` with ``zq`` nearest-resized to ``f``.""" + + def __init__(self, f_channels: int, zq_channels: int) -> None: + super().__init__() + self.norm_layer = _K6RMSNorm(f_channels, images=False) + self.conv_y = _ChunkedConv3d(zq_channels, f_channels, kernel_size=1) + self.conv_b = _ChunkedConv3d(zq_channels, f_channels, kernel_size=1) + + def forward(self, f: torch.Tensor, zq: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + frames, height, width = f.shape[-3:] + if cache.first: + # The causal first latent frame maps to exactly one pixel frame; the rest cover the remaining frames. + zq_first = _chunked_interpolate_nearest(zq[:, :, :1], (1, height, width)) + if zq.size(2) > 1: + zq_rest = _chunked_interpolate_nearest(zq[:, :, 1:], (frames - 1, height, width)) + zq = torch.cat([zq_first, zq_rest], dim=2) + else: + zq = zq_first + else: + zq = _chunked_interpolate_nearest(zq, (frames, height, width)) + norm_f = self.norm_layer(f) + norm_f.mul_(self.conv_y(zq)) + norm_f.add_(self.conv_b(zq)) + return norm_f + + +class PXSDownsample(nn.Module): + """x2 spatial (strided conv + channel-averaged pixel-unshuffle) and optional x2 causal temporal (two causal convs + + average pooling) downsample; doubles the channels.""" + + def __init__(self, in_channels: int, compress_time: bool) -> None: + super().__init__() + out_channels = 2 * in_channels + self.spatial_conv = _ChunkedConv3d(in_channels, out_channels, kernel_size=(1, 3, 3), stride=(1, 2, 2), + padding=(0, 1, 1)) + self.compress_time = compress_time + if compress_time: + self.temporal_conv = nn.Sequential( + CausalConv3d(out_channels, out_channels, kernel_size=(2, 1, 1)), + CausalConv3d(out_channels, out_channels, kernel_size=(2, 1, 1), stride=(2, 1, 1)), + ) + self.linear = _ChunkedConv3d(out_channels, out_channels, kernel_size=1) + + def _spatial(self, x: torch.Tensor) -> torch.Tensor: + b, c, t, h, w = x.shape + pxs = F.pixel_unshuffle(x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w), 2) + pxs = pxs.view(b * t, 2 * c, 2, h // 2, w // 2).mean(dim=2) + pxs = pxs.reshape(b, t, 2 * c, h // 2, w // 2).permute(0, 2, 1, 3, 4) + return self.spatial_conv(x) + pxs + + def _temporal(self, x: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + b, c, t, h, w = x.shape + pooled = x.permute(0, 3, 4, 1, 2).reshape(b * h * w, c, t) + if cache.first: # the causal first frame is kept as is + first, rest = pooled[..., :1], pooled[..., 1:] + pooled = torch.cat([first, F.avg_pool1d(rest, kernel_size=2, stride=2)], dim=-1) if t > 1 else first + else: + pooled = F.avg_pool1d(pooled, kernel_size=2, stride=2) + pooled = pooled.reshape(b, h, w, c, -1).permute(0, 3, 4, 1, 2) + conv = self.temporal_conv[1](self.temporal_conv[0](x, cache), cache) + return conv + pooled + + def forward(self, x: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + out = self._spatial(x) + if self.compress_time: + out = self._temporal(out, cache) + return self.linear(out) + + +class PXSUpsample(nn.Module): + """Optional x2 causal temporal (frame repeat + causal conv residual) then x2 spatial (nearest + conv residual) + upsample.""" + + def __init__(self, in_channels: int, compress_time: bool) -> None: + super().__init__() + self.spatial_conv = _ChunkedConv3d(in_channels, in_channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + self.compress_time = compress_time + if compress_time: + self.temporal_conv = CausalConv3d(in_channels, in_channels, kernel_size=(3, 1, 1)) + self.linear = _ChunkedConv3d(in_channels, in_channels, kernel_size=1) + + def forward(self, x: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + if self.compress_time: + x = x.repeat_interleave(2, dim=2) + if cache.first: # the causal first frame is not duplicated + x = x[:, :, 1:] + x = self.temporal_conv(x, cache) + x + frames, height, width = x.shape[-3:] + # Plain F.interpolate silently corrupts most elements once the tensor exceeds ~2**31 elements (it appears + # to do 32-bit index arithmetic internally): chunk along channels like SpatialNorm3D's zq resize does. + x = _chunked_interpolate_nearest(x, (frames, height * 2, width * 2)) + x.add_(self.spatial_conv(x)) + return self.linear(x) + + +def _level_module(**children: nn.Module) -> nn.Module: + module = nn.Module() + for name, child in children.items(): + setattr(module, name, child) + return module + + +class Encoder3D(nn.Module): + + def __init__(self, *, ch: int, ch_mult: list[float], num_res_blocks: int, in_channels: int, z_channels: int, + temporal_compress_times: int, temporal_compress_start_level: int) -> None: + super().__init__() + self.num_res_blocks = num_res_blocks + time_levels = range(temporal_compress_start_level, + temporal_compress_start_level + int(math.log2(temporal_compress_times))) + self.conv_in = CausalConv3d(in_channels, round(ch * ch_mult[0]), kernel_size=3) + self.down = nn.ModuleList() + block_in = round(ch * ch_mult[0]) + for level, mult in enumerate(ch_mult): + block_out = round(ch * mult) + blocks = nn.ModuleList() + for _ in range(num_res_blocks): + blocks.append(ResnetBlock3D(block_in, block_out)) + block_in = block_out + down = _level_module(block=blocks) + if level != len(ch_mult) - 1: + down.downsample = PXSDownsample(block_in, compress_time=level in time_levels) + block_in *= 2 + self.down.append(down) + self.mid = _level_module(block_1=ResnetBlock3D(block_in, block_in), block_2=ResnetBlock3D(block_in, block_in)) + self.norm_out = _K6RMSNorm(block_in, images=False) + self.conv_out = CausalConv3d(block_in, 2 * z_channels, kernel_size=3) + + def forward(self, x: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + h = self.conv_in(x, cache) + for down in self.down: + for block in down.block: + h = block(h, cache) + if hasattr(down, "downsample"): + h = down.downsample(h, cache) + h = self.mid.block_2(self.mid.block_1(h, cache), cache) + return self.conv_out(_silu(self.norm_out(h)), cache) + + +class Decoder3D(nn.Module): + + def __init__(self, *, ch: int, ch_mult: list[float], num_res_blocks: int, out_ch: int, z_channels: int, + temporal_compress_times: int, temporal_compress_start_level: int) -> None: + super().__init__() + num_levels = len(ch_mult) + # Mirror of the encoder's temporally compressing levels. + time_levels = range(num_levels - temporal_compress_start_level - int(math.log2(temporal_compress_times)), + num_levels - temporal_compress_start_level) + block_in = round(ch * ch_mult[-1]) + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3) + self.mid = _level_module(block_1=ResnetBlock3D(block_in, block_in, z_channels), + block_2=ResnetBlock3D(block_in, block_in, z_channels)) + up_levels: list[nn.Module] = [] + for level in reversed(range(num_levels)): + block_out = round(ch * ch_mult[level]) + blocks = nn.ModuleList() + for _ in range(num_res_blocks + 1): + blocks.append(ResnetBlock3D(block_in, block_out, z_channels)) + block_in = block_out + up = _level_module(block=blocks) + if level != 0: + up.upsample = PXSUpsample(block_in, compress_time=level in time_levels) + up_levels.insert(0, up) + self.up = nn.ModuleList(up_levels) + self.norm_out = SpatialNorm3D(block_in, z_channels) + self.conv_out = CausalConv3d(block_in, out_ch, kernel_size=3) + + def forward(self, z: torch.Tensor, cache: _SegmentCache) -> torch.Tensor: + h = self.conv_in(z, cache) + h = self.mid.block_2(self.mid.block_1(h, cache, z), cache, z) + for up in reversed(self.up): + for block in up.block: + h = block(h, cache, z) + if hasattr(up, "upsample"): + h = up.upsample(h, cache) + return self.conv_out(_silu(self.norm_out(h, z, cache)), cache) + + +# Code-path-selecting knobs of the released checkpoint; anything else would need layers this module does not have. +_REQUIRED_ENCODER_KNOBS = {"norm_type": "rms_norm", "padding_mode": "zeros", "downsample_version": 2, + "fix_pxs": True, "double_z": True} +_REQUIRED_DECODER_KNOBS = {"norm_type": "rms_norm", "padding_mode": "zeros"} +_ARCH_KEYS = ("ch", "ch_mult", "num_res_blocks", "z_channels", "temporal_compress_times", + "temporal_compress_start_level") + + +def _arch_kwargs(conf: dict[str, Any], required: dict[str, Any], extra: str, which: str) -> dict[str, Any]: + for key, value in required.items(): + if conf.get(key) != value: + raise ValueError(f"Kandinsky6SRVAE supports only {which}_config[{key!r}] == {value!r}, " + f"got {conf.get(key)!r}") + missing = [key for key in (*_ARCH_KEYS, extra) if key not in conf] + if missing: + raise ValueError(f"Kandinsky6SRVAE {which}_config is missing {missing}") + return {key: conf[key] for key in (*_ARCH_KEYS, extra)} + + +class Kandinsky6SRVAE(nn.Module): + + def __init__(self, config: Kandinsky6SRVAEConfig) -> None: + super().__init__() + self.config = config + arch = config.arch_config + assert isinstance(arch, Kandinsky6SRVAEArchConfig) + if not arch.encoder_config or not arch.decoder_config: + raise ValueError("Kandinsky6SRVAE needs `encoder_config` and `decoder_config` (KVAE architecture) in the " + "component config.json.") + enc = _arch_kwargs(dict(arch.encoder_config), _REQUIRED_ENCODER_KNOBS, "in_channels", "encoder") + dec = _arch_kwargs(dict(arch.decoder_config), _REQUIRED_DECODER_KNOBS, "out_ch", "decoder") + self.temporal_compression = enc["temporal_compress_times"] + self.encoder = Encoder3D(**enc) + self.decoder = Decoder3D(**dec) + + @property + def scaling_factor(self) -> float: + return float(self.config.arch_config.scaling_factor) + + @property + def spatial_factor(self) -> int: + return int(self.config.arch_config.spatial_factor) + + @property + def temporal_factor(self) -> int: + return int(self.config.arch_config.temporal_factor) + + @staticmethod + def normalize_data(data: torch.Tensor) -> torch.Tensor: + return data / 128 - 1.0 + + @staticmethod + def denormalize_data(data: torch.Tensor) -> torch.Tensor: + return (data + 1) * 128 + + def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, list[int]]: + """``[B, C, T, H, W]`` normalized pixels -> ``(raw latent mean, pixel-frame segment sizes)``.""" + split_list = [_SEGMENT_FRAMES + 1] + remaining = x.size(2) - split_list[0] + while remaining > 0: + split_list.append(_SEGMENT_FRAMES) + remaining -= _SEGMENT_FRAMES + split_list[-1] += remaining + cache = _SegmentCache() + latents = [] + for segment in torch.split(x, split_list, dim=2): + moments = self.encoder(segment, cache) + cache.first = False + latents.append(moments.chunk(2, dim=1)[0]) # deterministic posterior: the mean + return torch.cat(latents, dim=2), split_list + + def decode(self, z: torch.Tensor) -> DecoderOutput: + """``[B, C, T', h, w]`` raw latent -> ``.sample`` ``[B, 3, T, H, W]`` in the normalized pixel range.""" + segment = _SEGMENT_FRAMES // self.temporal_compression + num_frames = z.size(2) + if num_frames == 1: + split_list = [1] + else: + split_list = [segment] * ((num_frames - 1) // segment) + if (num_frames - 1) % segment: + split_list.append((num_frames - 1) % segment) + split_list[0] += 1 + cache = _SegmentCache() + samples = [] + for chunk in torch.split(z, split_list, dim=2): + samples.append(self.decoder(chunk, cache)) + cache.first = False + return DecoderOutput(sample=torch.cat(samples, dim=2)) + + +EntryClass = Kandinsky6SRVAE diff --git a/fastvideo/pipelines/basic/kandinsky6/__init__.py b/fastvideo/pipelines/basic/kandinsky6/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/kandinsky6/kandinsky6_ti2va_pipeline.py b/fastvideo/pipelines/basic/kandinsky6/kandinsky6_ti2va_pipeline.py new file mode 100644 index 000000000..66c036022 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6/kandinsky6_ti2va_pipeline.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages.input_validation import InputValidationStage +from fastvideo.pipelines.stages.kandinsky6 import ( + Kandinsky6AudioDecodingStage, + Kandinsky6CFGResolutionStage, + Kandinsky6DecodingStage, + Kandinsky6DenoisingStage, + Kandinsky6ImageEncodingStage, + Kandinsky6LatentPreparationStage, + Kandinsky6TextEncodingStage, +) +from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage + + +class Kandinsky6TI2VAPipeline(ComposedPipelineBase): + """Kandinsky6 TI2VA pipeline: text, optionally plus a conditioning image, + to synchronized video+audio, for the base (flow-matching) and the pi-Flow + distilled checkpoint alike (the bundle's scheduler component decides the + sampler). There is deliberately no separate T2VA/I2VA pipeline class -- + one stage chain handles both, like the official Diffusers + Kandinsky6TI2VAPipeline: Kandinsky6ImageEncodingStage is a no-op when the + request carries no conditioning image and injects it as an extra + reference frame when one is supplied. + """ + + # model_index.json declares "audio_vae" (mel-VAE decoder) and "vocoder" (BigVGAN-v2, + # mel->waveform) as two separate components; Kandinsky6AudioDecodingStage calls + # audio_vae.decode() then vocoder(mel), like LTX-2's audio_vae/vocoder pair. + _required_config_modules = [ + "scheduler", + "text_encoder", + "text_encoder_2", + "tokenizer", + "tokenizer_2", + "transformer", + "vae", + "audio_vae", + "vocoder", + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + + # Recomputes do_classifier_free_guidance under Kandinsky6's own CFG rule (abs(g-1)>1e-6, + # forced off for PiFlow) before text encoding decides whether to encode a negative prompt. + self.add_stage( + stage_name="cfg_resolution_stage", + stage=Kandinsky6CFGResolutionStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="text_encoding_stage", + stage=Kandinsky6TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2"), + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2"), + ], + ), + ) + + self.add_stage( + stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + + self.add_stage( + stage_name="latent_preparation_stage", + stage=Kandinsky6LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + ) + + # No-op unless the request supplied a conditioning image; see + # Kandinsky6ImageEncodingStage's docstring. + self.add_stage( + stage_name="image_encoding_stage", + stage=Kandinsky6ImageEncodingStage(vae=self.get_module("vae")), + ) + + self.add_stage( + stage_name="denoising_stage", + stage=Kandinsky6DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + + self.add_stage( + stage_name="audio_decoding_stage", + stage=Kandinsky6AudioDecodingStage( + audio_vae=self.get_module("audio_vae"), + vocoder=self.get_module("vocoder"), + ), + ) + + self.add_stage( + stage_name="decoding_stage", + stage=Kandinsky6DecodingStage(vae=self.get_module("vae"), pipeline=self), + ) + + +EntryClass = Kandinsky6TI2VAPipeline diff --git a/fastvideo/pipelines/basic/kandinsky6/presets.py b/fastvideo/pipelines/basic/kandinsky6/presets.py new file mode 100644 index 000000000..baeb38544 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6/presets.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 model family pipeline presets.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +# Byte-identical to the diffusers Kandinsky6TI2VAPipeline's +# _DEFAULT_NEGATIVE_PROMPT (and to Kandinsky5's own default negative prompt -- +# same lineage, same default). +_NEGATIVE_PROMPT = ("Static, 2D cartoon, cartoon, 2d animation, paintings, images, worst quality, low quality, ugly, " + "deformed, walking backwards") + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Main Kandinsky6 joint video+audio denoising pass", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +# Diffusers Kandinsky6TI2VAPipeline defaults +# (height=512, width=768, num_frames=121, num_inference_steps=50, +# guidance_scale=5.0, sample_fps=24.0). One preset serves the single TI2VA +# pipeline whether or not a conditioning image_path is passed at call time. +KANDINSKY6_TI2VA = InferencePreset( + name="kandinsky6_ti2va", + version=1, + model_family="kandinsky6", + description="Kandinsky6 TI2VA (Pro)", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 512, + "width": 768, + "num_frames": 121, + "fps": 24, + "guidance_scale": 5.0, + "num_inference_steps": 50, + "negative_prompt": _NEGATIVE_PROMPT, + }, +) + +# pi-Flow distilled checkpoint (Kandinsky-6.0-Pro-distill-5s-Diffusers): distilled for 10 steps (model card; its +# scheduler_config.json leaves nfe unset) and the pi-Flow policy has no classifier-free guidance, so the run is 10 +# steps at guidance_scale=1.0 (the diffusers pipeline rejects any other guidance for a PiflowScheduler). +KANDINSKY6_TI2VA_DISTILLED = InferencePreset( + name="kandinsky6_ti2va_distilled", + version=1, + model_family="kandinsky6", + description="Kandinsky6 TI2VA pi-Flow distilled (Pro-distill)", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 512, + "width": 768, + "num_frames": 121, + "fps": 24, + "guidance_scale": 1.0, + "num_inference_steps": 10, + "negative_prompt": _NEGATIVE_PROMPT, + }, +) + +ALL_PRESETS = (KANDINSKY6_TI2VA, KANDINSKY6_TI2VA_DISTILLED) diff --git a/fastvideo/pipelines/basic/kandinsky6_sr/__init__.py b/fastvideo/pipelines/basic/kandinsky6_sr/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6_sr/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/kandinsky6_sr/kandinsky6_sr_pipeline.py b/fastvideo/pipelines/basic/kandinsky6_sr/kandinsky6_sr_pipeline.py new file mode 100644 index 000000000..f670e7bf2 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6_sr/kandinsky6_sr_pipeline.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 video super-resolution (VSR) pipeline: source video -> x2 / x2.25 / x4 video with the source audio. + +The model is text-free. An official ``Kandinsky6SRPipeline`` bundle holds ``transformer`` (SR DiT), ``vae`` (KVAE), +``latent_upscaler`` (x2 / x4 latent upscalers) and ``scheduler``: ``FlowMatchEulerDiscreteScheduler`` for +``Kandinsky-6.0-VSR-5s-Diffusers`` and ``PiflowScheduler`` for the 2-step distilled +``Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers``. See ``docs/inference/kandinsky6_sr.md``. +""" +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages.kandinsky6_sr import (Kandinsky6SRDecodingStage, Kandinsky6SRDenoisingStage, + Kandinsky6SRLatentPreparationStage, + Kandinsky6SRVideoEncodingStage) + + +class Kandinsky6SRPipeline(ComposedPipelineBase): + + _required_config_modules = ["transformer", "vae", "latent_upscaler", "scheduler"] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + transformer = self.get_module("transformer") + vae = self.get_module("vae") + self.add_stage(stage_name="video_encoding_stage", stage=Kandinsky6SRVideoEncodingStage(vae=vae)) + self.add_stage(stage_name="latent_preparation_stage", + stage=Kandinsky6SRLatentPreparationStage(latent_upscaler=self.get_module("latent_upscaler"), + vae=vae, + transformer=transformer)) + self.add_stage(stage_name="denoising_stage", + stage=Kandinsky6SRDenoisingStage(transformer=transformer, + scheduler=self.get_module("scheduler"))) + self.add_stage(stage_name="decoding_stage", stage=Kandinsky6SRDecodingStage(vae=vae, transformer=transformer)) + + +EntryClass = Kandinsky6SRPipeline diff --git a/fastvideo/pipelines/basic/kandinsky6_sr/presets.py b/fastvideo/pipelines/basic/kandinsky6_sr/presets.py new file mode 100644 index 000000000..b4d672346 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6_sr/presets.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 video super-resolution (VSR) presets.""" + +from dataclasses import asdict + +from fastvideo.api.presets import InferencePreset, PresetStageSpec +from fastvideo.configs.pipelines.kandinsky6_sr_options import Kandinsky6SROptions, SR_OPTION_FIELDS + +_SR_STAGE = PresetStageSpec( + name="sr", + kind="super_resolution", + description="Tiled latent-upscaler + SR DiT pass", + allowed_overrides=frozenset({ + "num_inference_steps", + *SR_OPTION_FIELDS, + }), +) + +# Geometry and frame rate come from the source video; height / width / num_frames / fps below only keep the generic +# request validation and logs sensible. The DiT is text-free, so there is no classifier-free guidance. +_COMMON_DEFAULTS = { + "height": 512, + "width": 768, + "num_frames": 121, + "fps": 24, + "seed": 42, + "guidance_scale": 1.0, + "save_video": True, + "return_frames": False, + **asdict(Kandinsky6SROptions()), +} + +KANDINSKY6_SR = InferencePreset( + name="kandinsky6_sr", + version=2, + model_family="kandinsky6_sr", + description="Kandinsky6 video super-resolution, flow matching (4 Euler steps per tile)", + workload_type=None, + stage_schemas=(_SR_STAGE, ), + defaults={ + **_COMMON_DEFAULTS, "num_inference_steps": 4 + }, +) + +KANDINSKY6_SR_DISTILLED = InferencePreset( + name="kandinsky6_sr_distilled", + version=1, + model_family="kandinsky6_sr", + description="Kandinsky6 video super-resolution, 2-step pi-Flow distilled", + workload_type=None, + stage_schemas=(_SR_STAGE, ), + defaults={ + **_COMMON_DEFAULTS, "num_inference_steps": 2 + }, +) + +ALL_PRESETS = (KANDINSKY6_SR, KANDINSKY6_SR_DISTILLED) diff --git a/fastvideo/pipelines/basic/kandinsky6_sr/sr_io.py b/fastvideo/pipelines/basic/kandinsky6_sr/sr_io.py new file mode 100644 index 000000000..de4230079 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6_sr/sr_io.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Video / audio input and delivery resizing for Kandinsky6 video super-resolution. + +The SR model has its own input contract, so it does not go through the generic ``load_video`` (which decodes every +frame and drops the audio) or ``InputValidationStage`` (which resizes to the request geometry): the source is resampled +to 24 fps by a fixed stride, capped at 121 frames aligned to ``1 + 8k`` and kept at its own resolution. Frames are +streamed and only the selected ones are kept; the audio track is decoded separately so it can be muxed back. +""" +from __future__ import annotations + +import math +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch +from torch.nn import functional + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +TARGET_FPS = 24 +MAX_NUM_FRAMES = 121 +# 23.98 and 25 fps sources are treated as already at the training rate. +RESAMPLE_FPS_TOLERANCE = 1.5 +AUDIO_SAMPLE_RATE = 44100 + +#: Delivery tiers ``name -> (H, W)`` candidates (landscape, square, portrait); the closest aspect ratio wins. +TARGET_RESOLUTIONS: dict[str, tuple[tuple[int, int], ...]] = { + "hd": ((720, 1280), (720, 720), (1280, 720)), + "fullhd": ((1080, 1920), (1080, 1080), (1920, 1080)), + "2k": ((1440, 2560), (1440, 1440), (2560, 1440)), +} + + +@dataclass(frozen=True) +class FramePlan: + """Stream indices to keep and the rate they play at. + + When downsampling, the kept count is ``int(total_frames / step)``, which depends on the true stream length: + ``lookahead`` frames must be read before the selection is final. + """ + + indices: tuple[int, ...] + out_fps: int + step: float | None = None + lookahead: int = 0 + + +def plan_frame_selection(src_fps: float) -> FramePlan: + if src_fps <= 0: + raise ValueError(f"Source video reports no usable fps (got {src_fps}).") + if abs(src_fps - TARGET_FPS) < RESAMPLE_FPS_TOLERANCE: + return FramePlan(tuple(range(MAX_NUM_FRAMES)), round(src_fps)) + if src_fps < TARGET_FPS: + logger.warning("Source fps %.2f < %d: keeping the native frame rate (mildly out of distribution).", src_fps, + TARGET_FPS) + return FramePlan(tuple(range(MAX_NUM_FRAMES)), round(src_fps)) + step = src_fps / TARGET_FPS + virtual_total = math.ceil(MAX_NUM_FRAMES * step) + 1 + indices = [round(i * step) for i in range(int(virtual_total / step))] + indices = [i for i in indices if i < virtual_total][:MAX_NUM_FRAMES] + logger.warning("Source fps %.2f > %d: downsampling by a fixed stride.", src_fps, TARGET_FPS) + return FramePlan(tuple(indices), TARGET_FPS, step, max(math.ceil(MAX_NUM_FRAMES * step), indices[-1] + 1)) + + +def _select_frames(frames: Iterator[torch.Tensor], plan: FramePlan) -> list[torch.Tensor]: + wanted = frozenset(plan.indices) + stop_after = max(plan.indices[-1] + 1, plan.lookahead) + kept: list[torch.Tensor] = [] + seen = 0 + for index, frame in enumerate(frames): + if index in wanted: + kept.append(frame) + seen = index + 1 + if seen >= stop_after: + break + if plan.step is not None and seen < plan.lookahead: + kept = kept[:int(seen / plan.step)] + return kept + + +def read_video(path: str | Path) -> tuple[torch.Tensor, int]: + """Decode ``path`` to ``([T, C, H, W] uint8, output_fps)`` following the SR input contract.""" + import av + + path = Path(path) + if not path.is_file(): + raise FileNotFoundError(f"Input video does not exist: {path}") + with av.open(str(path)) as container: + if not container.streams.video: + raise ValueError(f"Video {path.name} has no video stream.") + stream = container.streams.video[0] + plan = plan_frame_selection(float(stream.average_rate or stream.base_rate or 0.0)) + decoded = (torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1) + for frame in container.decode(stream)) + frames = _select_frames(decoded, plan) + if not frames: + raise ValueError(f"Video {path.name} has no readable frames.") + aligned = 1 + 8 * ((min(len(frames), MAX_NUM_FRAMES) - 1) // 8) + return torch.stack(frames[:aligned]).contiguous(), plan.out_fps + + +def read_audio(path: str | Path, num_frames: int, fps: int) -> torch.Tensor | None: + """Mono float audio of ``path`` at :data:`AUDIO_SAMPLE_RATE`, trimmed to the span of ``num_frames`` frames. + + Returns ``None`` when there is no audio track or it cannot be decoded: audio is best effort and never fails the + video. + """ + import av + + keep = int(round(num_frames / fps * AUDIO_SAMPLE_RATE)) + chunks: list[np.ndarray] = [] + produced = 0 + try: + with av.open(str(path)) as container: + if not container.streams.audio: + return None + resampler = av.AudioResampler(format="fltp", layout="mono", rate=AUDIO_SAMPLE_RATE) + for frame in container.decode(audio=0): + for resampled in resampler.resample(frame): + chunks.append(resampled.to_ndarray().reshape(-1)) + produced += chunks[-1].shape[0] + if produced >= keep + AUDIO_SAMPLE_RATE: + break + else: + chunks.extend(resampled.to_ndarray().reshape(-1) for resampled in resampler.resample(None)) + except Exception as exc: # noqa: BLE001 + logger.warning("Could not decode the audio of %s (%s); the output will have no audio.", path, exc) + return None + if not chunks: + return None + audio = np.clip(np.concatenate(chunks), -1.0, 1.0).astype(np.float32, copy=False)[:keep] + return torch.from_numpy(audio.copy()) if audio.size else None + + +def validate_target_spec(spec: str | None, mode: str) -> None: + """Reject a bad delivery spec before any model runs.""" + resolve_target_hw(spec, (1, 1), mode, validate_only=True) + + +def resolve_target_hw(spec: str | None, + sr_hw: tuple[int, int], + mode: str, + validate_only: bool = False) -> tuple[int, int] | None: + """Delivery size for ``spec`` (a tier name or ``WxH``); ``None`` when the SR result is kept as is. + + ``fit`` shrinks isotropically into the bucket (never upscales, even sizes for the codec); ``exact`` returns the + bucket itself. + """ + if mode not in ("fit", "exact"): + raise ValueError(f"sr_target_resize_mode must be 'fit' or 'exact', got {mode!r}") + if spec is None or spec.lower() == "none": + return None + key = spec.lower() + if key in TARGET_RESOLUTIONS: + ratio = sr_hw[1] / sr_hw[0] + bucket = min(TARGET_RESOLUTIONS[key], key=lambda hw: abs(hw[1] / hw[0] - ratio)) + else: + parts = key.split("x") + if len(parts) != 2 or not all(part.strip().isdigit() for part in parts): + raise ValueError(f"sr_target_resolution must be one of {sorted(TARGET_RESOLUTIONS)} or WxH, got {spec!r}") + bucket = (int(parts[1]), int(parts[0])) + if mode == "exact" or validate_only: + return bucket + scale = min(bucket[0] / sr_hw[0], bucket[1] / sr_hw[1]) + if scale >= 1.0: + logger.warning("SR result %dx%d already fits the %dx%d bucket; keeping it as is.", sr_hw[1], sr_hw[0], + bucket[1], bucket[0]) + return None + return tuple(min(b, round(side * scale / 2) * 2) for b, side in zip(bucket, sr_hw, strict=True)) # type: ignore + + +def resize_video(video: torch.Tensor, target_hw: tuple[int, int], frame_chunk: int = 8) -> torch.Tensor: + """Antialiased downscale of a ``[C, T, H, W]`` uint8 video to ``target_hw`` (no crop, never upscales).""" + channels, frames, height, width = video.shape + target_h, target_w = target_hw + if (height, width) == (target_h, target_w): + return video + if target_h > height or target_w > width: + raise ValueError(f"Target {target_w}x{target_h} exceeds the SR result {width}x{height}; pick a lower tier or " + "a higher sr_resolution_scale.") + if abs((width / height) / (target_w / target_h) - 1) > 0.02: + logger.warning("Resizing %dx%d to %dx%d changes the aspect ratio (no crop).", width, height, target_w, target_h) + out = torch.empty((channels, frames, target_h, target_w), dtype=torch.uint8) + for t0 in range(0, frames, frame_chunk): + chunk = video[:, t0:t0 + frame_chunk].permute(1, 0, 2, 3).float() + resized = functional.interpolate(chunk, size=target_hw, mode="bilinear", antialias=True, align_corners=False) + out[:, t0:t0 + frame_chunk] = resized.clamp(0, 255).round().to(torch.uint8).permute(1, 0, 2, 3) + return out diff --git a/fastvideo/pipelines/basic/kandinsky6_sr/tiling.py b/fastvideo/pipelines/basic/kandinsky6_sr/tiling.py new file mode 100644 index 000000000..9e78807c3 --- /dev/null +++ b/fastvideo/pipelines/basic/kandinsky6_sr/tiling.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tile planning and stitching for Kandinsky6 video super-resolution. + +The SR DiT was trained on fixed base resolutions, so a source frame is cut into overlapping tiles whose upscaled size +is one of those bases, each tile is super-resolved independently and the decoded tiles are blended back with a Hanning +window. Tile positions are aligned to the VAE spatial factor so the same grid addresses both the latent and the +decoded pixels. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from torch.nn import functional + +#: Trained base resolutions ``visual_size -> [(H, W), ...]`` of the SR DiT. +RESOLUTIONS: dict[int, tuple[tuple[int, int], ...]] = { + 512: ((512, 512), (512, 768), (768, 512)), +} + + +@dataclass(frozen=True) +class TileGrid: + """Tile size and per-axis tile start positions, in pixels of the (pre-upscaled) source frame.""" + + tile_h: int + tile_w: int + tops: tuple[int, ...] + lefts: tuple[int, ...] + + @property + def num_tiles(self) -> int: + return len(self.tops) * len(self.lefts) + + def to_latent(self, spatial_factor: int) -> TileGrid: + values = (self.tile_h, self.tile_w, *self.tops, *self.lefts) + if any(value % spatial_factor for value in values): + raise ValueError(f"Tile grid {self} is not aligned to the VAE spatial factor {spatial_factor}.") + return TileGrid(self.tile_h // spatial_factor, self.tile_w // spatial_factor, + tuple(top // spatial_factor for top in self.tops), + tuple(left // spatial_factor for left in self.lefts)) + + def extract(self, video: torch.Tensor) -> list[torch.Tensor]: + """``[..., H, W]`` -> row-major list of ``[..., tile_h, tile_w]`` views.""" + return [video[..., top:top + self.tile_h, left:left + self.tile_w] for top in self.tops for left in self.lefts] + + +def _axis_positions(length: int, tile: int, min_overlap: float, snap: int) -> tuple[int, ...]: + """Fewest evenly spaced starts covering ``[0, length - tile]`` with at least ``min_overlap`` overlap. + + Positions are multiples of ``snap`` when both ``length`` and ``tile`` are; otherwise the same layout is computed at + pixel precision and :meth:`TileGrid.to_latent` rejects the unaligned grid. + """ + if tile >= length: + return (0, ) + unit = snap if length % snap == 0 and tile % snap == 0 else 1 + span_units = (length - tile) // unit + max_stride_units = max(1, math.floor(tile * (1.0 - min_overlap) / unit)) + count = math.ceil(span_units / max_stride_units) + 1 + return tuple(round(i * span_units / (count - 1)) * unit for i in range(count)) + + +def plan_tiles(height: int, width: int, visual_size: int, scale: int, min_overlap: float, + spatial_factor: int) -> TileGrid: + """Tile grid of a ``height x width`` source frame for an integer tiling ``scale``. + + The base resolution is the trained one closest in aspect ratio; each tile is ``base / scale`` source pixels, so its + upscaled size is exactly the base resolution the DiT was trained on. + """ + if visual_size not in RESOLUTIONS: + raise ValueError(f"Unsupported SR visual_size={visual_size}; known sizes: {sorted(RESOLUTIONS)}") + if not 0.0 <= min_overlap < 1.0: + raise ValueError(f"sr_tile_min_overlap must be in [0, 1), got {min_overlap}") + ratio = width / height + base_h, base_w = min(RESOLUTIONS[visual_size], key=lambda hw: abs(hw[1] / hw[0] - ratio)) + if base_h % scale or base_w % scale: + raise ValueError(f"Tiling scale {scale} does not divide the base resolution {base_h}x{base_w}.") + tile_h, tile_w = base_h // scale, base_w // scale + return TileGrid(tile_h, tile_w, _axis_positions(height, tile_h, min_overlap, spatial_factor), + _axis_positions(width, tile_w, min_overlap, spatial_factor)) + + +def stitch_tiles(tiles: list[torch.Tensor], + grid: TileGrid, + height: int, + width: int, + scale: int, + frame_chunk: int = 16) -> torch.Tensor: + """Blend row-major float ``[0, 255]`` tiles ``[C, T, tile_h * scale, tile_w * scale]`` into a single uint8 + ``[C, T, H * scale, W * scale]`` video, quantizing once at the very end (not per tile): pre-quantizing each + tile to uint8 before blending would blend already-rounded values instead of the continuous pixel values, + compounding rounding error at every overlap -- matches the diffusers reference's ``decode_latents``/ + ``__call__``, which accumulates ``video_acc += tile * window`` in float and rounds only the final result. + + Each tile is weighted by a 2D Hanning window without zero endpoints (so border pixels covered by a single tile keep + a non-zero weight) and the accumulation is normalised per pixel, which also handles uneven overlaps. Frames are + blended in chunks to bound host memory; frames are independent, so chunking does not change the result. + """ + channels, frames, hr_h, hr_w = tiles[0].shape + window_h = torch.hann_window(hr_h + 2)[1:-1] + window_w = torch.hann_window(hr_w + 2)[1:-1] + window = (window_h[:, None] * window_w[None, :])[None, None] + out = torch.empty((channels, frames, height * scale, width * scale), dtype=torch.uint8) + for t0 in range(0, frames, max(1, frame_chunk)): + t1 = min(frames, t0 + frame_chunk) + acc = torch.zeros((channels, t1 - t0, height * scale, width * scale)) + weight = torch.zeros((1, 1, height * scale, width * scale)) + tile_iter = iter(tiles) + for top in grid.tops: + for left in grid.lefts: + y, x = top * scale, left * scale + acc[:, :, y:y + hr_h, x:x + hr_w] += next(tile_iter)[:, t0:t1] * window + weight[:, :, y:y + hr_h, x:x + hr_w] += window + out[:, t0:t1] = (acc / weight.clamp(min=1e-6)).clamp(0, 255).to(torch.uint8) + return out + + +def resolve_scale(scale: float) -> tuple[int, float]: + """Requested total upscale -> ``(tiling_scale, pixel_pre_upscale)``. + + The latent upscaler only has x2 and x4 entries; the fractional x2.25 is an x1.125 pixel pre-upscale followed by x2 + tiling. + """ + if scale == 2.25: + return 2, 1.125 + if scale in (2.0, 4.0): + return int(scale), 1.0 + raise ValueError(f"sr_resolution_scale must be 2, 4 or 2.25, got {scale}") + + +def pre_upscale(video: torch.Tensor, factor: float, spatial_multiple: int, frame_chunk: int = 16) -> torch.Tensor: + """Bilinear pixel upscale of a ``[T, C, H, W]`` uint8 video, target size rounded to ``spatial_multiple``.""" + height, width = video.shape[-2:] + size = (max(spatial_multiple, + round(height * factor / spatial_multiple) * spatial_multiple), + max(spatial_multiple, + round(width * factor / spatial_multiple) * spatial_multiple)) + chunks = [] + for t0 in range(0, video.shape[0], frame_chunk): + resized = functional.interpolate(video[t0:t0 + frame_chunk].float(), + size=size, + mode="bilinear", + align_corners=False) + chunks.append(resized.round_().clamp_(0, 255).to(torch.uint8)) + return torch.cat(chunks, dim=0) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index d5c0d8230..7f85149df 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -586,6 +586,17 @@ class ComposedPipelineBase(ABC): ) continue transformers_or_diffusers = module_spec[0] + # model_index.json's [library, class] spec is the only place some + # checkpoints declare a component's class name at all, when the + # component's config.json omits `_class_name`. Stash it so component loaders (which + # only see component_model_path, not this spec) can fall back to + # it instead of crashing on a KeyError. + if len(module_spec) > 1 and module_spec[1] is not None: + hints = getattr(fastvideo_args, "_model_index_class_names", None) + if hints is None: + hints = {} + fastvideo_args._model_index_class_names = hints + hints[module_name] = module_spec[1] if transformers_or_diffusers is None: logger.warning("Module %s in model_index.json has null value, removing from required_config_modules", module_name) diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 639b1556e..4667d2216 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -34,6 +34,9 @@ from fastvideo.pipelines.stages.matrixgame3_denoising import MatrixGame3Denoisin from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage from fastvideo.pipelines.stages.kandinsky5 import (Kandinsky5DecodingStage, Kandinsky5DenoisingStage, Kandinsky5LatentPreparationStage) +from fastvideo.pipelines.stages.kandinsky6 import (Kandinsky6AudioDecodingStage, Kandinsky6DecodingStage, + Kandinsky6DenoisingStage, Kandinsky6ImageEncodingStage, + Kandinsky6LatentPreparationStage) from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage from fastvideo.pipelines.stages.gen3c_stages import (Gen3CCFGPolicyStage, Gen3CConditioningStage, Gen3CDenoisingStage, Gen3CLatentPreparationStage) @@ -69,6 +72,11 @@ __all__ = [ "Kandinsky5DecodingStage", "Kandinsky5DenoisingStage", "Kandinsky5LatentPreparationStage", + "Kandinsky6AudioDecodingStage", + "Kandinsky6DecodingStage", + "Kandinsky6DenoisingStage", + "Kandinsky6ImageEncodingStage", + "Kandinsky6LatentPreparationStage", "GameCraftDenoisingStage", "Gen3CCFGPolicyStage", "Gen3CConditioningStage", diff --git a/fastvideo/pipelines/stages/kandinsky6.py b/fastvideo/pipelines/stages/kandinsky6.py new file mode 100644 index 000000000..6ef2c231e --- /dev/null +++ b/fastvideo/pipelines/stages/kandinsky6.py @@ -0,0 +1,707 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +import contextlib +from copy import deepcopy +import math +from typing import Any + +import PIL +import torch +from diffusers.utils.torch_utils import randn_tensor +from tqdm.auto import tqdm + +from fastvideo.attention.backends.nabla import NablaAttentionMetadataBuilder +from fastvideo.configs.pipelines.kandinsky6 import ( + KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX, ) +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader, VAELoader +from fastvideo.models.vaes.common import ParallelTiledVAE +from fastvideo.models.vision_utils import normalize, numpy_to_pt, pil_to_numpy +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.decoding import DecodingStage +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.utils import PRECISION_TO_TYPE, get_mixed_precision_state + +logger = init_logger(__name__) + +# batch.extra keys used to pass state between the Kandinsky6 stages -- the +# same "stash it on batch.extra, downstream stages check for it" pattern +# fastvideo/pipelines/basic/ltx2 uses for its audio latents. +_TAIL_COND_KEY = "kandinsky6_tail_cond_active" +_TOKEN_TYPE_IDS_KEY = "kandinsky6_visual_token_type_ids" +_AUDIO_SAMPLE_RATE_DEFAULT = 44100 + + +def audio_latent_duration(video_latent_frames: int, *, fps: float, audio_sample_rate: int, + audio_downsample_factor: int) -> int: + """Audio latent length matching the diffusers Kandinsky6 T2VA reference: + ``ceil(pixel_frames / fps * audio_sample_rate / audio_downsample_factor)``, + where ``pixel_frames = (video_latent_frames - 1) * 4 + 1`` is the causal + video VAE's temporal-compression convention (matches + HunyuanVAEConfig.temporal_compression_ratio == 4). + """ + pixel_frames = (video_latent_frames - 1) * 4 + 1 + return int(math.ceil(pixel_frames / fps * audio_sample_rate / audio_downsample_factor)) + + +class Kandinsky6CFGResolutionStage(PipelineStage): + """Recomputes ``batch.do_classifier_free_guidance`` before text encoding using Kandinsky6's CFG + contract. The diffusers reference (``pipeline_kandinsky6_ti2va.py``)'s ``do_classifier_free_guidance`` + property gates on the standard ``guidance_scale > 1.0`` (any guidance_scale <= 1.0, including below 1, + runs cond-only with no negative-prompt encoding); this stage mirrors that. A PiFlow (distilled) + scheduler never uses CFG regardless of the requested guidance_scale (its cond-only contract, and the + guidance==1.0 requirement, are enforced by Kandinsky6DenoisingStage instead), so this also + short-circuits to False for it and skips encoding a negative prompt that would otherwise go unused. + """ + + def __init__(self, scheduler) -> None: + super().__init__() + self.scheduler = scheduler + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if bool(getattr(self.scheduler, "is_piflow", False)): + batch.do_classifier_free_guidance = False + else: + batch.do_classifier_free_guidance = batch.guidance_scale > 1.0 + return batch + + +class Kandinsky6TextEncodingStage(TextEncodingStage): + """Maps the request's ``max_sequence_length`` onto the Qwen encoder only. + + Matches the diffusers reference, whose Qwen ``max_length = max_sequence_length + 129`` (the + prompt-template prefix crop, ``ENCODE_START_IDX``). The CLIP pooled encoder always keeps its fixed + 77-token tokenizer config regardless of the request: applying the same override to it (the + generic ``TextEncodingStage``'s default behavior, one ``max_length`` for every encoder) makes a + CLIP tokenizer with only 77 positions try to pad to whatever length was requested for Qwen and + crash. + """ + + def _resolve_max_length(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> int | Sequence[int | None] | None: + if batch.max_sequence_length is None: + return None + return [KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX + batch.max_sequence_length, None] + + +class Kandinsky6LatentPreparationStage(PipelineStage): + """Draw initial video *and* audio noise latents. + + Runs for both T2VA and IT2VA calls -- the conditioning image (if any) is + applied afterward by Kandinsky6ImageEncodingStage, which appends an extra + reference frame rather than overwriting one drawn here, so the RNG order + here does not depend on whether an image was supplied. + """ + + def __init__(self, scheduler, transformer) -> None: + super().__init__() + self.scheduler = scheduler + self.transformer = transformer + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.height is None or batch.width is None: + raise ValueError("height and width must be provided for Kandinsky6.") + height = int(batch.height) + width = int(batch.width) + num_frames = int(batch.num_frames) + + pipeline_config = fastvideo_args.pipeline_config + temporal_ratio = pipeline_config.vae_config.arch_config.temporal_compression_ratio + spatial_ratio = pipeline_config.vae_config.arch_config.spatial_compression_ratio + patch_size = pipeline_config.dit_config.arch_config.patch_size + + # Round to the nearest k*temporal_ratio+1, matching the diffusers reference's + # `check_inputs`/`__call__` num_frames adjustment (`num_frames // temporal_ratio * + # temporal_ratio + 1`, warning first): this floors for a non-exact multiple (119 -> 117) but + # rounds *up* for an exact multiple of temporal_ratio (120 -> 121, 96 -> 97), unlike a pure + # floor formula, which would floor every case (120 -> 117, 96 -> 93); 121 -> 121 unchanged + # either way. + if num_frames % temporal_ratio != 1: + rounded_num_frames = max(num_frames // temporal_ratio * temporal_ratio + 1, 1) + logger.warning("Kandinsky6: num_frames=%d is not of the form k*%d+1; rounding to %d.", num_frames, + temporal_ratio, rounded_num_frames) + num_frames = rounded_num_frames + batch.num_frames = num_frames + + required_divisor_h = spatial_ratio * patch_size[1] + required_divisor_w = spatial_ratio * patch_size[2] + if height % required_divisor_h != 0 or width % required_divisor_w != 0: + raise ValueError(f"Kandinsky6 height must be divisible by {required_divisor_h} and width by " + f"{required_divisor_w}; got height={height}, width={width}.") + + if isinstance(batch.prompt, list): + batch_size = len(batch.prompt) + elif batch.prompt is not None: + batch_size = 1 + else: + batch_size = batch.prompt_embeds[0].shape[0] + batch_size *= batch.num_videos_per_prompt + + dtype = PRECISION_TO_TYPE[pipeline_config.dit_precision] + device = get_local_torch_device() + num_latent_frames = (num_frames - 1) // temporal_ratio + 1 + num_channels = getattr(self.transformer, "in_visual_dim", pipeline_config.dit_config.arch_config.in_visual_dim) + video_shape = (batch_size, num_latent_frames, height // spatial_ratio, width // spatial_ratio, num_channels) + + if batch.latents is None: + video = randn_tensor(video_shape, generator=batch.generator, device=device, dtype=dtype) + if hasattr(self.scheduler, "init_noise_sigma"): + video = video * self.scheduler.init_noise_sigma + else: + video = batch.latents.to(device=device, dtype=dtype) + if tuple(video.shape) != video_shape: + raise ValueError(f"Provided latents shape {list(video.shape)} does not match expected " + f"Kandinsky6 video latent shape {list(video_shape)}.") + + visual_cond = getattr(self.transformer, "visual_cond", False) + if visual_cond: + cond = torch.zeros_like(video) + mask = torch.zeros((*video.shape[:-1], 1), device=video.device, dtype=video.dtype) + video = torch.cat([video, cond, mask], dim=-1) + + num_audio_channels = getattr(self.transformer, "in_audio_dim", + pipeline_config.dit_config.arch_config.in_audio_dim) + # Matches the diffusers reference's `sample_fps` request field, which + # drives the audio latent length directly (not just the saved video's + # playback rate): a request fps of 30 with 121 frames yields 174 + # audio latents there, not the 218 a fixed 24 fps would give. + requested_fps = batch.fps + if isinstance(requested_fps, list): + requested_fps = requested_fps[0] if requested_fps else None + fps = float(requested_fps) if requested_fps else pipeline_config.sample_fps + num_audio_frames = audio_latent_duration( + num_latent_frames, + fps=fps, + audio_sample_rate=pipeline_config.audio_sample_rate, + audio_downsample_factor=pipeline_config.audio_downsample_factor, + ) + audio_shape = (batch_size, num_audio_frames, num_audio_channels) + if batch.audio_latents is None: + # Drawn from the same generator right after the video noise, so it + # is automatically an independent sample (no manual seed offset + # needed -- unlike the diffusers reference, which threads a bare + # int seed through its packed-latent-prep helpers and has to + # offset it by hand). + audio = randn_tensor(audio_shape, generator=batch.generator, device=device, dtype=dtype) + else: + audio = batch.audio_latents.to(device=device, dtype=dtype) + if tuple(audio.shape) != audio_shape: + raise ValueError(f"Provided audio_latents shape {list(audio.shape)} does not match expected " + f"Kandinsky6 audio latent shape {list(audio_shape)}.") + + batch.latents = video + batch.audio_latents = audio + batch.raw_latent_shape = (batch_size, num_channels, num_latent_frames, height // spatial_ratio, + width // spatial_ratio) + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("num_frames", batch.num_frames, V.positive_int) + result.add_check("height", batch.height, V.positive_int) + result.add_check("width", batch.width, V.positive_int) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("audio_latents", batch.audio_latents, [V.is_tensor, V.with_dims(3)]) + return result + + +class Kandinsky6ImageEncodingStage(PipelineStage): + """Optional IT2VA conditioning: no-op for a pure T2VA call. + + When ``batch.pil_image`` is supplied, encodes it through the shared video + VAE and appends it as one extra "clean" reference frame at the end of the + video latent sequence (the diffusers reference's default + ``tail_cond_first_frame`` scheme), tagged via a token-type id so the + transformer's ``visual_token_type_embeddings`` can distinguish it from + generated frames. Kandinsky6DenoisingStage re-pins this frame every step + and strips it back out after the loop. + """ + + def __init__(self, vae: ParallelTiledVAE) -> None: + self.vae = vae + + @staticmethod + def _cover_resize_dims(src_h: int, src_w: int, height: int, width: int) -> tuple[int, int]: + """Resize target for a subsequent centre crop to (height, width), matching the diffusers + reference's ``encode_i2va_first_frame``: ``scale = min(src_h/h, src_w/w)`` so the resized + image covers the target box on both axes (never smaller than it), and the excess is + centre-cropped afterward -- unlike squeezing the whole source image to fit, which distorts + its aspect ratio. + """ + scale = min(src_h / height, src_w / width) + # `int(src / scale)` matches the diffusers reference formula, but on the constraining axis + # (the one that produced `scale`) it can land a hair below the target -- e.g. 831 for 832 -- + # from float error in the scale round-trip, which would give a negative centre-crop offset. + # Clamp each axis to its target as a floor; this is a no-op whenever `int()` already reaches + # the target, so it doesn't change the (non-constraining-axis) diffusers-parity values. + new_h = max(height, int(src_h / scale)) + new_w = max(width, int(src_w / scale)) + return new_h, new_w + + @classmethod + def _preprocess(cls, image, height: int, width: int) -> torch.Tensor: + if isinstance(image, PIL.Image.Image): + src_w, src_h = image.size + new_h, new_w = cls._cover_resize_dims(src_h, src_w, height, width) + image = image.resize((new_w, new_h), resample=PIL.Image.BILINEAR) + top, left = (new_h - height) // 2, (new_w - width) // 2 + image = image.crop((left, top, left + width, top + height)) + image = numpy_to_pt(pil_to_numpy(image)) + return normalize(image) + if image.min() >= 0: + logger.warning("Kandinsky6 conditioning image tensor has no negative values; " + "assuming range [0, 1] and normalizing to [-1, 1]. " + "Pass a [-1, 1] tensor with negative values to skip normalization.") + image = normalize(image) + if image.ndim == 3: + image = image.unsqueeze(0) + src_h, src_w = image.shape[-2], image.shape[-1] + if (src_h, src_w) != (height, width): + new_h, new_w = cls._cover_resize_dims(src_h, src_w, height, width) + image = torch.nn.functional.interpolate(image.float(), size=(new_h, new_w), mode="bilinear", antialias=True) + top, left = (new_h - height) // 2, (new_w - width) // 2 + image = image[..., top:top + height, left:left + width] + return image + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.pil_image is None: + return batch + + if not fastvideo_args.model_loaded["vae"]: + vae = getattr(self, "vae", None) + if vae is None: + loader = VAELoader() + vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args) + self.vae = vae + fastvideo_args.model_loaded["vae"] = True + + device = get_local_torch_device() + vae = self.vae.to(device) + self.vae = vae + vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast + + image = self._preprocess(batch.pil_image, int(batch.height), int(batch.width)) + image = image.to(device=device, dtype=torch.float32).unsqueeze(2) # [B,C,1,H,W] + + prev_use_tiling = vae.use_tiling + vae.use_tiling = False + try: + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): + if not vae_autocast_enabled: + image = image.to(vae_dtype) + generator = batch.generator + if isinstance(generator, list) and len(generator) != image.shape[0]: + generator = generator[0] + image_latent = vae.encode(image).sample(generator=generator) + finally: + vae.use_tiling = prev_use_tiling + + image_latent = image_latent * vae.scaling_factor + # [B,C,1,H,W] -> [B,1,H,W,C] channel-last, matching the video latent. + image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous() + batch.image_latent = image_latent + + latents = batch.latents + image_latent = image_latent.to(device=latents.device, dtype=latents.dtype) + num_channels = image_latent.shape[-1] + + ref_frame = image_latent + if latents.shape[-1] > num_channels: + # visual_cond channel layout: [real, cond, mask]. Per the + # diffusers reference's `_build_video_input` (pipeline_kandinsky6_ + # ti2va.py), the tail_cond_first_frame scheme leaves the cond block + # at zero and only writes the real channel block (done above via + # `ref_frame = image_latent`) and mask=1 -- unlike Kandinsky5's + # I2V "pretrain"/"i2v" schemes (fastvideo/pipelines/stages/ + # kandinsky5.py), which duplicate the image latent into cond. + # Different scheme, don't carry that convention over. + cond_block = torch.zeros_like(image_latent) + mask_block = torch.ones_like(image_latent[..., :1]) + ref_frame = torch.cat([image_latent, cond_block, mask_block], dim=-1) + + if ref_frame.shape[0] != latents.shape[0]: + # One image is encoded per request regardless of + # num_videos_per_prompt (there is only ever one conditioning + # image); broadcast it to every video in the batch instead of + # letting the concatenation below fail with a tensor-size + # mismatch. + if ref_frame.shape[0] != 1: + raise ValueError(f"Kandinsky6 image conditioning batch size {ref_frame.shape[0]} does not match " + f"the video latent batch size {latents.shape[0]} and is not broadcastable.") + ref_frame = ref_frame.expand(latents.shape[0], *ref_frame.shape[1:]) + + batch.latents = torch.cat([latents, ref_frame], dim=1) + batch.extra[_TAIL_COND_KEY] = True + + num_video_frames = latents.shape[1] + token_type_ids = torch.zeros((latents.shape[0], num_video_frames + 1), dtype=torch.long, device=device) + token_type_ids[:, -1] = 1 + batch.extra[_TOKEN_TYPE_IDS_KEY] = token_type_ids + + if fastvideo_args.vae_cpu_offload: + vae.to("cpu") + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + return result + + +class Kandinsky6DenoisingStage(PipelineStage): + """Joint video/audio denoising with independent scheduler state per modality. + + Mirrors the diffusers reference's ``denoise_loop`` (pipeline_kandinsky6_ + ti2va.py): one joint transformer call per step for both modalities, CFG + via ``uncond + w*(cond-uncond)``, video advanced through the shared + scheduler's ``step()`` (which owns the sigma index), audio advanced with + a manual Euler update using the *same* per-step sigma delta so the + scheduler's step index is never double-advanced. + """ + + def __init__(self, transformer, scheduler) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + + @staticmethod + def _text_rope_pos(mask: torch.Tensor, device: torch.device) -> torch.Tensor: + seq_len = int(mask.sum(1).max().item()) + return torch.arange(seq_len, device=device) + + def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -> torch.dtype: + """See Kandinsky5DenoisingStage._resolve_target_dtype: trusts + ``dit_precision`` for a plain load, but FSDP2 always computes in the + policy's ``param_dtype`` regardless of parameter storage dtype.""" + declared_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] + try: + from torch.distributed.fsdp import FSDPModule + except Exception: # pragma: no cover - FSDP not always available + return declared_dtype + if not isinstance(self.transformer, FSDPModule): + return declared_dtype + try: + policy_dtype = get_mixed_precision_state().param_dtype + except ValueError: + policy_dtype = None + return policy_dtype if policy_dtype is not None else torch.bfloat16 + + @staticmethod + def fast_sta_nabla(T: int, H: int, W: int, wT: int, wH: int, wW: int, device: torch.device | str) -> torch.Tensor: + max_extent = int(torch.tensor([T, H, W], device=device).amax().item()) + r = torch.arange(0, max_extent, 1, dtype=torch.int16, device=device) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + sta_t = (mat[:T, :T].flatten() <= wT // 2) + sta_h = (mat[:H, :H].flatten() <= wH // 2) + sta_w = (mat[:W, :W].flatten() <= wW // 2) + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + def get_sparse_params(self, video: torch.Tensor, device: torch.device) -> dict[str, Any] | None: + cfg = self.transformer.config + if cfg.attention_engine != "nabla": + return None + assert cfg.patch_size[0] == 1 + _, T, H, W, _ = video.shape + T, H, W = T // cfg.patch_size[0], H // cfg.patch_size[1], W // cfg.patch_size[2] + sta_mask = self.fast_sta_nabla(T, + H // 8, + W // 8, + cfg.attention_wT, + cfg.attention_wH, + cfg.attention_wW, + device=device) + return { + "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0), + "attention_type": cfg.attention_engine, + "to_fractal": True, + "P": cfg.attention_P, + "wT": cfg.attention_wT, + "wW": cfg.attention_wW, + "wH": cfg.attention_wH, + "add_sta": cfg.attention_add_sta, + "visual_shape": (T, H, W), + "method": cfg.attention_method, + } + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + use_piflow = bool(getattr(self.scheduler, "is_piflow", False)) + if use_piflow and (not math.isfinite(batch.guidance_scale) or abs(batch.guidance_scale - 1.0) > 1e-6): + nfe = getattr(self.scheduler.config, "nfe", None) + raise ValueError( + f"Kandinsky6 PiFlow (distilled) checkpoints need guidance_scale=1.0 (got {batch.guidance_scale}); " + "pi-Flow runs without classifier-free guidance, so the official Diffusers pipeline rejects any " + f"other value too. num_inference_steps is not constrained by the checkpoint -- {nfe} is only the " + "scheduler's nfe, the value this checkpoint was distilled for and the default of the " + "'kandinsky6_ti2va_distilled' preset (itself the default for " + "'kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers' and local copies named like it, a " + "Kandinsky-6 name containing 'distill'). Pass guidance_scale=1.0 explicitly otherwise.") + audio_scheduler = deepcopy(self.scheduler) if use_piflow else None + if batch.timesteps is None: + raise ValueError("timesteps must be prepared before Kandinsky6 denoising.") + if batch.latents is None or batch.audio_latents is None: + raise ValueError("video and audio latents must be prepared before Kandinsky6 denoising.") + if not fastvideo_args.model_loaded["transformer"]: + loader = TransformerLoader() + self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args) + fastvideo_args.model_loaded["transformer"] = True + + device = get_local_torch_device() + target_dtype = self._resolve_target_dtype(fastvideo_args) + autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast + video = batch.latents + audio = batch.audio_latents + num_channels = getattr(self.transformer, "in_visual_dim", + fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim) + + tail_cond = bool(batch.extra.get(_TAIL_COND_KEY, False)) + visual_token_type_ids = batch.extra.get(_TOKEN_TYPE_IDS_KEY) + + prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype) + pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype) + if not batch.prompt_attention_mask: + raise ValueError("Kandinsky6 requires Qwen prompt attention masks.") + text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device) + + neg_prompt_embeds = None + neg_pooled = None + negative_text_rope_pos = None + if batch.do_classifier_free_guidance and batch.negative_prompt_embeds: + neg_prompt_embeds = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype) + neg_pooled = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype) + if not batch.negative_attention_mask: + raise ValueError("Kandinsky6 requires Qwen negative attention masks for CFG.") + negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask[0].to(device), device) + + height = int(batch.height) + width = int(batch.width) + pipeline_config = fastvideo_args.pipeline_config + spatial_ratio = pipeline_config.vae_config.arch_config.spatial_compression_ratio + patch_size = pipeline_config.dit_config.arch_config.patch_size + + num_video_frames = video.shape[1] - (1 if tail_cond else 0) + t_positions = torch.arange(num_video_frames, device=device) + if tail_cond: + # The appended reference frame reuses T-position 0's rope row + # (matches the diffusers reference's `torch.cat([rope, rope[:1]])` + # duplication -- equivalent since RoPE3D is a pure position -> + # table lookup). + t_positions = torch.cat([t_positions, t_positions.new_zeros(1)]) + visual_rope_pos = [ + t_positions, + torch.arange(height // spatial_ratio // patch_size[1], device=device), + torch.arange(width // spatial_ratio // patch_size[2], device=device), + ] + # Read straight off the loaded transformer (falls back to the pipeline's static arch config + # for a not-yet-loaded model, e.g. a verify_input probe): both official checkpoints declare + # (1.0, 2.0, 2.0) in transformer/config.json. Diffusers uses this value unconditionally, not + # a height/width-dependent heuristic. + scale_factor = getattr(self.transformer, "scale_factor", + tuple(float(v) for v in pipeline_config.dit_config.arch_config.scale_factor)) + sparse_params = self.get_sparse_params(video, device) + + # PiFlow's noisy latent state is tracked in fp32 across steps here, separately from the + # `video` buffer's storage dtype -- matches the diffusers reference, whose `PiflowScheduler. + # step` returns fp32 and whose pipeline rebinds the video state to that fp32 result every + # step rather than writing it back into a lower-precision buffer (`scheduling_piflow.py`, + # `pipeline_kandinsky6_ti2va.py`'s `denoise_loop`). flow-Euler does not need this: the + # diffusers reference also downcasts its state to the parameter dtype every step there. + # Always a Tensor (not `Tensor | None`) even on the flow-Euler path, where it is simply + # unused: the initial cast is a single cheap op, and keeping it non-Optional avoids + # re-narrowing `video_state` after every `if use_piflow:` branch re-entry in the loop below. + video_state = video[..., :num_channels].to(torch.float32) + + with tqdm(total=batch.num_inference_steps, desc="Kandinsky6 Denoising") as progress_bar: + for i, timestep in enumerate(batch.timesteps): + if hasattr(self, "interrupt") and self.interrupt: + break + + # The scheduler timestep is fp32 in the diffusers reference (`t.unsqueeze(0). + # expand(bs)`, embedded via `time.float()`); Kandinsky6TimeEmbeddings' sinusoid + # table (`torch.outer(time, freqs)`) needs that precision directly; rounding the + # timestep itself to bf16 first (as opposed to computing the embedding in bf16 from + # an already-fp32 timestep) is the largest single numeric gap against the reference. + t_expand = timestep.unsqueeze(0).repeat(video.shape[0]).to(device=device, dtype=torch.float32) + if use_piflow: + video_input = video_state.to(dtype=target_dtype) + if video.shape[-1] > num_channels: + video_input = torch.cat([video_input, video[..., num_channels:].to(dtype=target_dtype)], dim=-1) + else: + video_input = video.to(dtype=target_dtype) + attn_metadata = None + if sparse_params is not None: + attn_metadata = NablaAttentionMetadataBuilder().build( + current_timestep=i, + sta_mask=sparse_params["sta_mask"], + P=sparse_params["P"], + visual_shape=sparse_params["visual_shape"], + ) + autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled) + if device.type == "cuda" else contextlib.nullcontext()) + with set_forward_context(current_timestep=i, attn_metadata=attn_metadata, + forward_batch=batch), autocast_ctx: + video_vel, audio_vel = self.transformer( + hidden_states=video_input, + hidden_states_audio=audio.to(dtype=target_dtype), + encoder_hidden_states=prompt_embeds, + pooled_projections=pooled, + timestep=t_expand, + visual_rope_pos=visual_rope_pos, + text_rope_pos=text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + visual_token_type_ids=visual_token_type_ids, + return_dict=True, + ).sample + + if neg_prompt_embeds is not None and neg_pooled is not None: + uncond_video_vel, uncond_audio_vel = self.transformer( + hidden_states=video_input, + hidden_states_audio=audio.to(dtype=target_dtype), + encoder_hidden_states=neg_prompt_embeds, + pooled_projections=neg_pooled, + timestep=t_expand, + visual_rope_pos=visual_rope_pos, + text_rope_pos=negative_text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + visual_token_type_ids=visual_token_type_ids, + return_dict=True, + ).sample + video_vel = uncond_video_vel + batch.guidance_scale * (video_vel - uncond_video_vel) + audio_vel = uncond_audio_vel + batch.guidance_scale * (audio_vel - uncond_audio_vel) + + if use_piflow: + video_state = self.scheduler.step(video_vel, timestep, video_state, return_dict=False)[0] + video[..., :num_channels] = video_state.to(video.dtype) + audio = audio_scheduler.step(audio_vel, timestep, audio, return_dict=False)[0] + else: + video[..., :num_channels] = self.scheduler.step(video_vel, + timestep, + video[..., :num_channels], + return_dict=False)[0] + # Manual Euler update for audio: one `scheduler.step()` call + # per iteration total (against video, above) so its internal + # sigma index never double-advances; audio uses the same + # sigma delta directly, matching the diffusers reference. + step_size = self.scheduler.sigmas[i + 1] - self.scheduler.sigmas[i] + audio = audio + step_size.to(device=audio.device, dtype=audio.dtype) * audio_vel + if tail_cond: + ref_frame = batch.image_latent.to(device=video.device, dtype=video.dtype) + video[:, -1:, :, :, :num_channels] = ref_frame + if use_piflow: + video_state[:, -1:] = ref_frame.to(torch.float32) + + if i == len(batch.timesteps) - 1 or (i + 1) % self.scheduler.order == 0: + progress_bar.update() + + video = video[..., :num_channels] + if tail_cond: + video = video[:, :-1] + + batch.latents = video + batch.audio_latents = audio + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("audio_latents", batch.audio_latents, [V.is_tensor, V.with_dims(3)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2)) + return result + + +class Kandinsky6DecodingStage(DecodingStage): + """Channel-last [B,T,H,W,C] -> channel-first [B,C,T,H,W] then the + generic VAE decode, matching Kandinsky5DecodingStage.""" + + def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None: + super().__init__(vae=vae, pipeline=pipeline) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.latents is None: + raise ValueError("latents must be available before Kandinsky6 decoding.") + batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous() + return super().forward(batch, fastvideo_args) + + +class Kandinsky6AudioDecodingStage(PipelineStage): + """Decode Kandinsky6 audio latents into a waveform via two separate + components -- ``audio_vae`` (mel-VAE decoder) and ``vocoder`` + (BigVGAN-v2, mel->waveform) -- like LTX-2's audio_vae/vocoder pair. + Writes ``batch.extra["audio"]`` / ``["audio_sample_rate"]``, which + ``VideoGenerator`` already knows how to mux into the saved video file. + """ + + def __init__(self, audio_vae, vocoder) -> None: + super().__init__() + self.audio_vae = audio_vae + self.vocoder = vocoder + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.audio_latents is None: + return batch + + device = get_local_torch_device() + self.audio_vae = self.audio_vae.to(device) + self.vocoder = self.vocoder.to(device) + + # Cast to the module's own loaded weight dtype directly (no autocast + # / precision-config knob) -- matches the existing MMAudio audio + # decoding stage (fastvideo/pipelines/basic/mmaudio/stages.py). + decoder_dtype = next(self.audio_vae.parameters()).dtype + + # Diffusion-latent-level denorm (matches the diffusers reference's + # postprocess_audio: `audio / scaling_factor`). Read from the loaded + # module itself -- this value lives in the checkpoint's own + # audio_vae/config.json, not a pipeline-level default. + scaling_factor = getattr(self.audio_vae, "scaling_factor", 1.0) + latents = batch.audio_latents / scaling_factor + + # [B, A, D] -> [B, D, A] to match the audio VAE's 1D-conv (channel, + # length) convention. + latents = latents.to(device=device, dtype=decoder_dtype).transpose(1, 2) + mel = self.audio_vae.decode(latents) # [B, num_mels, T] + waveform = self.vocoder(mel.to(next(self.vocoder.parameters()).dtype)) # [B, 1, samples] + + # Batch size is 1 in practice (Kandinsky6, like Kandinsky5, is a + # single-prompt pipeline); VideoGenerator's audio mux path + # (`_audio_to_int16`) expects one plain [-1,1] float array/tensor per + # call, not a per-batch-item list, and does its own int16 conversion. + batch.extra["audio"] = waveform[0, 0].clamp(-1.0, 1.0).float().cpu() + batch.extra["audio_sample_rate"] = _resolve_audio_sample_rate(self.audio_vae, fastvideo_args.pipeline_config) + + if fastvideo_args.vae_cpu_offload: + self.audio_vae.to("cpu") + self.vocoder.to("cpu") + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("audio_latents", batch.audio_latents, V.none_or_tensor) + return result + + +def _resolve_audio_sample_rate(vocoder, pipeline_config) -> int: + model = getattr(vocoder, "module", vocoder) + for attr in ("output_sampling_rate", "output_sample_rate", "sample_rate"): + if hasattr(model, attr): + return int(getattr(model, attr)) + return getattr(pipeline_config, "audio_sample_rate", _AUDIO_SAMPLE_RATE_DEFAULT) diff --git a/fastvideo/pipelines/stages/kandinsky6_sr.py b/fastvideo/pipelines/stages/kandinsky6_sr.py new file mode 100644 index 000000000..f941bce22 --- /dev/null +++ b/fastvideo/pipelines/stages/kandinsky6_sr.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stages of the Kandinsky6 video super-resolution (SR) pipeline. + + video encoding -> latent preparation -> denoising -> decoding + +The source clip is encoded once. Its latent is cut into overlapping tiles whose upscaled size is a resolution the SR +DiT was trained on (``basic/kandinsky6_sr/tiling.py``); every tile is upscaled by the latent upscaler, noised, denoised +with the bundle's scheduler (flow-matching Euler or the distilled pi-Flow) and decoded, and the decoded tiles are +blended into the output video. The denoising loop runs per tile, which is why these stages replace the generic +latent-preparation / denoising / decoding stages. +""" +from __future__ import annotations + +import contextlib +from typing import Any + +import torch +from tqdm.auto import tqdm + +from fastvideo.configs.pipelines.kandinsky6_sr_options import Kandinsky6SROptions +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader, UpsamplerLoader, VAELoader +from fastvideo.pipelines.basic.kandinsky6_sr import sr_io +from fastvideo.pipelines.basic.kandinsky6_sr.tiling import (TileGrid, plan_tiles, pre_upscale, resolve_scale, + stitch_tiles) +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.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + +# Optional request inputs routed into ``batch.extra`` (``REQUEST_BATCH_EXTRA_PASSTHROUGH_FIELDS``): a raw KVAE latent +# of the source instead of ``video_path``, and an audio track to mux instead of the source's own. +_LR_LATENT_KEY = "sr_lr_latent" +_AUDIO_KEY = "sr_audio" +_AUDIO_RATE_KEY = "sr_audio_sample_rate" + + +def _resolve(stage: PipelineStage, name: str, fastvideo_args: FastVideoArgs, loader: Any) -> Any: + """Reload a component that was released after a previous request.""" + if not fastvideo_args.model_loaded.get(name, True): + setattr(stage, name, loader.load(fastvideo_args.model_paths[name], fastvideo_args)) + fastvideo_args.model_loaded[name] = True + return getattr(stage, name) + + +def _module_dtype(module: torch.nn.Module) -> torch.dtype: + return next(module.parameters()).dtype + + +def _tile_grid(batch: ForwardBatch, transformer: torch.nn.Module, vae: torch.nn.Module) -> tuple[TileGrid, int]: + """Pixel tile grid of the encoded source (``batch.lq_latents``) and the integer tiling scale.""" + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + scale, _ = resolve_scale(float(sr_options.sr_resolution_scale)) + spatial_factor = vae.spatial_factor + _, _, height, width = batch.lq_latents.shape + grid = plan_tiles(height * spatial_factor, width * spatial_factor, transformer.config.arch_config.visual_size, + scale, float(sr_options.sr_tile_min_overlap), spatial_factor) + return grid, scale + + +class Kandinsky6SRVideoEncodingStage(PipelineStage): + """Read the source video (and its audio), apply the x2.25 pre-upscale and KVAE-encode it into ``lq_latents``.""" + + def __init__(self, vae: Any) -> None: + super().__init__() + self.vae = vae + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + _, pre_upscale_factor = resolve_scale(float(sr_options.sr_resolution_scale)) + sr_io.validate_target_spec(sr_options.sr_target_resolution, sr_options.sr_target_resize_mode) + if isinstance(batch.video_path, list): + raise ValueError("Kandinsky6 SR processes one video per request; pass a single video_path.") + vae = _resolve(self, "vae", fastvideo_args, VAELoader()) + + lr_latent = batch.extra.pop(_LR_LATENT_KEY, None) + if lr_latent is not None: + if batch.video_path is not None: + raise ValueError("Kandinsky6 SR: pass either video_path or sr_lr_latent, not both.") + if pre_upscale_factor != 1.0: + raise ValueError("Kandinsky6 SR: sr_resolution_scale=2.25 needs a pixel pre-upscale and cannot start " + "from sr_lr_latent; use scale 2 or 4, or pass video_path.") + if not isinstance(lr_latent, torch.Tensor) or lr_latent.ndim != 4: + raise ValueError("sr_lr_latent must be a raw (unscaled) KVAE latent tensor [T, C, H, W].") + batch.lq_latents = lr_latent.float() + fps = sr_io.TARGET_FPS + else: + video, fps = self._read_source(batch, pre_upscale_factor, vae.spatial_factor) + batch.lq_latents = self._encode(video, vae, fastvideo_args) + + audio = batch.extra.pop(_AUDIO_KEY, None) + audio_rate = batch.extra.pop(_AUDIO_RATE_KEY, sr_io.AUDIO_SAMPLE_RATE) + if audio is None and batch.video_path is not None: + audio = sr_io.read_audio(batch.video_path, int(batch.num_frames), fps) + audio_rate = sr_io.AUDIO_SAMPLE_RATE + if audio is not None: + batch.extra["audio"] = audio + batch.extra["audio_sample_rate"] = int(audio_rate) + # ``batch`` is the worker-side copy; the frame rate the mp4 must use travels back through ``extra``. + batch.fps = fps + batch.extra["output_fps"] = fps + return batch + + def _read_source(self, batch: ForwardBatch, factor: float, spatial_factor: int) -> tuple[torch.Tensor, int]: + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + if batch.video_path is None: + if batch.save_video or batch.return_frames: + raise ValueError("Kandinsky6 SR is a video-to-video pipeline: pass the source clip via video_path.") + # Warm-up / health-check request: one tile of deterministic noise. + generator = torch.Generator().manual_seed(0) + video, fps = torch.randint(0, 256, (9, 3, 256, 384), dtype=torch.uint8, generator=generator), 24 + else: + video, fps = sr_io.read_video(batch.video_path) + if factor != 1.0: + video = pre_upscale(video, factor, spatial_factor) + batch.num_frames, batch.height, batch.width = video.shape[0], video.shape[2], video.shape[3] + logger.info("Kandinsky6 SR input: %d frames %dx%d @ %d fps -> x%s", video.shape[0], video.shape[3], + video.shape[2], fps, sr_options.sr_resolution_scale) + return video, fps + + @staticmethod + def _encode(video: torch.Tensor, vae: torch.nn.Module, fastvideo_args: FastVideoArgs) -> torch.Tensor: + """``[T, C, H, W]`` uint8 -> raw ``[T', C', h, w]`` fp32 latent (the latent upscaler applies the scaling).""" + device = get_local_torch_device() + vae.to(device) + pixels = vae.normalize_data(video.permute(1, 0, 2, 3).unsqueeze(0).to(device).float()) + latent = vae.encode(pixels.to(_module_dtype(vae)))[0] + if fastvideo_args.vae_cpu_offload: + vae.to("cpu") + return latent.squeeze(0).permute(1, 0, 2, 3).float() + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("lq_latents", batch.lq_latents, [V.is_tensor, V.with_dims(4)]) + return result + + +class Kandinsky6SRLatentPreparationStage(PipelineStage): + """Upscale every latent tile and build the SR DiT input ``[noised upscaled latent | anchor | anchor mask]``. + + The released checkpoints are trained with an HR anchor channel group; SR of a real low-quality clip has no anchor, + so it is zero with a zero mask. The starting point is the upscaled latent mixed with Gaussian noise + (variance-preserving, ``lq_noise_scale``). Noise is drawn per group of ``sr_tiles_batch_size`` tiles from a + generator seeded with ``seed + first tile index``, so results do not depend on how many tile groups ran before. + Output: ``batch.latents`` ``[num_tiles, T', H, W, 2C + 1]`` fp32. + """ + + def __init__(self, latent_upscaler: Any, vae: Any, transformer: Any) -> None: + super().__init__() + self.latent_upscaler = latent_upscaler + self.vae = vae + self.transformer = transformer + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + upscaler = _resolve(self, "latent_upscaler", fastvideo_args, UpsamplerLoader()) + grid, scale = _tile_grid(batch, self.transformer, self.vae) + if scale not in upscaler.scales: + # The diffusers reference falls back to a pixel-space path here (bilinearly upscale the + # pixel tile to the base resolution, then VAE-encode it directly) when its latent_upscaler + # doesn't cover the requested scale. FastVideo makes latent_upscaler mandatory instead and + # fails fast: by this stage the pipeline only carries the already-KVAE-encoded + # batch.lq_latents (Kandinsky6SRVideoEncodingStage encodes the whole clip upfront, before + # tiling), not the source pixel video, so matching the reference exactly would mean + # threading the pixel video through additional pipeline state all the way to this later + # stage -- a real architectural change for a path no released checkpoint's scales/tile_sizes + # combination is known to hit (every supported SR bundle ships a complete upscaler bank). + raise ValueError(f"The latent upscaler has no x{scale} model (available: {upscaler.scales}).") + arch = self.transformer.config.arch_config + device = get_local_torch_device() + upscaler.to(device) + upscaler_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.upsampler_precision] + autocast = (torch.autocast("cuda", dtype=upscaler_dtype) if device.type == "cuda" else contextlib.nullcontext()) + scaling_factor = float(self.vae.scaling_factor) + + tiles = grid.to_latent(self.vae.spatial_factor).extract(batch.lq_latents) + group = int(sr_options.sr_tiles_batch_size) + if group < 1: + raise ValueError(f"sr_tiles_batch_size must be >= 1, got {group}") + prepared = [] + for start in range(0, len(tiles), group): + upscaled = [] + for tile in tiles[start:start + group]: + z = tile.permute(1, 0, 2, 3).unsqueeze(0).to(device=device, dtype=_module_dtype(upscaler)) + with autocast: + out = upscaler(z * scaling_factor, scale) + upscaled.append(out.squeeze(0).permute(1, 2, 3, 0).float()) # [T', H, W, C] + prepared.append(self._initial_latent(torch.stack(upscaled), int(batch.seed) + start, arch.lq_noise_scale)) + batch.latents = torch.cat(prepared) + if fastvideo_args.vae_cpu_offload: + upscaler.to("cpu") + return batch + + @staticmethod + def _initial_latent(upscaled: torch.Tensor, seed: int, noise_scale: float) -> torch.Tensor: + # One draw for the whole tile group, in the ``[group * T', H, W, C]`` shape the model was trained with. + flat = upscaled.flatten(0, 1) + generator = torch.Generator(device=flat.device).manual_seed(seed) + noise = torch.randn(flat.shape, device=flat.device, dtype=flat.dtype, generator=generator) + start = (1 - noise_scale**2)**0.5 * flat + noise_scale * noise + anchor_and_mask = torch.zeros((*flat.shape[:-1], flat.shape[-1] + 1), device=flat.device, dtype=flat.dtype) + return torch.cat([start, anchor_and_mask], dim=-1).unflatten(0, upscaled.shape[:2]) + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + return result + + +class Kandinsky6SRDenoisingStage(PipelineStage): + """Denoise every tile with the bundle's scheduler; ``batch.latents`` becomes ``[num_tiles, T', H, W, C]``. + + ``FlowMatchEulerDiscreteScheduler`` bundles run ``num_inference_steps`` Euler steps over the shifted + ``linspace(1, 0)`` grid; ``PiflowScheduler`` bundles integrate the distilled policy (the DiT head then holds + ``n_grid`` predictions per channel). The latent state stays fp32 across steps. Only the first ``C`` channels are + denoised; the anchor channels are fixed conditioning. + + ``num_inference_steps`` is the number of DiT calls per tile for both schedulers, as elsewhere in FastVideo. The + upstream Diffusers pipeline counts timestep grid points instead, so its default 5 is 4 steps here. + """ + + def __init__(self, transformer: Any, scheduler: Any) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + transformer = _resolve(self, "transformer", fastvideo_args, TransformerLoader()) + arch = transformer.config.arch_config + device = get_local_torch_device() + if not fastvideo_args.dit_layerwise_offload: + transformer.to(device) + target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] + autocast_enabled = device.type == "cuda" and target_dtype != torch.float32 and not fastvideo_args.disable_autocast + use_piflow = bool(getattr(self.scheduler, "is_piflow", False)) + channels = arch.in_visual_dim + num_steps = int(batch.num_inference_steps) + self._check_scheduler_fits(arch, use_piflow, num_steps) + + latents = batch.latents + num_tiles, duration, height, width, _ = latents.shape + rope_pos = [ + torch.arange(duration, device=device), + torch.arange(height // arch.patch_size[1], device=device), + torch.arange(width // arch.patch_size[2], device=device), + ] + group = int(sr_options.sr_tiles_batch_size) + denoised = [] + with tqdm(total=num_tiles * num_steps, desc="Kandinsky6 SR denoising [tiles x steps]") as progress_bar: + for start in range(0, num_tiles, group): + x = latents[start:start + group].to(device) + state, cond = x[..., :channels], x[..., channels:] + self._set_timesteps(num_steps, device, use_piflow) + for i, t in enumerate(self.scheduler.timesteps): + timestep = t.to(device=device, dtype=torch.float32).expand(x.shape[0]) + autocast = (torch.autocast("cuda", dtype=target_dtype) + if autocast_enabled else contextlib.nullcontext()) + with set_forward_context(current_timestep=i, attn_metadata=None, forward_batch=batch), autocast: + prediction = transformer(hidden_states=torch.cat([state, cond], dim=-1), + timestep=timestep, + visual_rope_pos=rope_pos, + scale_factor=arch.rope_scale_factor) + if not use_piflow: + # The Euler step casts its result to the prediction dtype; keep the state fp32. + prediction = prediction.float() + state = self.scheduler.step(prediction, t, state, return_dict=False)[0] + progress_bar.update(x.shape[0]) + denoised.append(state.float()) + batch.latents = torch.cat(denoised) + self._offload(transformer, fastvideo_args) + return batch + + def _check_scheduler_fits(self, arch: Any, use_piflow: bool, num_steps: int) -> None: + """Reject a transformer / scheduler pair from different bundles before the first DiT call. + + The flow-matching DiT predicts ``in_visual_dim`` channels; the distilled one ``n_grid * in_visual_dim``, which + only ``PiflowScheduler`` can integrate. Mixing them would otherwise fail on a shape error mid-run. + """ + if num_steps < 1: + raise ValueError(f"Kandinsky6 SR: num_inference_steps must be >= 1, got {num_steps}") + scheduler_name = type(self.scheduler).__name__ + expected = arch.in_visual_dim * (int(self.scheduler.n_grid) if use_piflow else 1) + if arch.out_visual_dim != expected: + raise ValueError( + f"Kandinsky6 SR: the transformer head is {arch.out_visual_dim} channels wide, but {scheduler_name} " + f"needs {expected} (in_visual_dim={arch.in_visual_dim}). The transformer and scheduler come from " + "different bundles: the flow-matching transformer needs FlowMatchEulerDiscreteScheduler, the distilled " + "one PiflowScheduler with the n_grid it was trained with.") + nfe = getattr(self.scheduler.config, "nfe", None) if use_piflow else None + if nfe is not None and num_steps != int(nfe): + logger.warning( + "Kandinsky6 SR: running %d pi-Flow steps per tile; this checkpoint was distilled for nfe=%d.", + num_steps, int(nfe)) + + def _set_timesteps(self, num_steps: int, device: torch.device, use_piflow: bool) -> None: + if use_piflow: + self.scheduler.set_timesteps(num_steps, device=device) + else: + # Uniform grid from pure noise; the scheduler applies its shift and appends the terminal 0. + sigmas = torch.linspace(1.0, 0.0, num_steps + 1)[:-1].tolist() + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + + @staticmethod + def _offload(transformer: torch.nn.Module, fastvideo_args: FastVideoArgs) -> None: + if fastvideo_args.dit_layerwise_offload: + manager = getattr(transformer, "_layerwise_offload_manager", None) + if manager is not None and getattr(manager, "enabled", False): + manager.release_all() + elif fastvideo_args.dit_cpu_offload and not fastvideo_args.use_fsdp_inference: + transformer.to("cpu") + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + return result + + +class Kandinsky6SRDecodingStage(PipelineStage): + """KVAE-decode every tile, blend the tiles and apply the optional delivery resize. + + Tiles stay float through blending and are quantized to uint8 once, inside ``stitch_tiles`` -- + matching the diffusers reference's ``decode_latents``/``__call__`` (it accumulates + ``video_acc += tile * window`` in float and only rounds the final ``video_acc / weight_acc``), not + the lossier quantize-per-tile-then-blend-already-rounded-values approach. ``batch.output`` stores + code ``v`` as ``(v + 0.5) / 255`` so the truncating ``* 255 -> uint8`` conversion in + ``VideoGenerator`` gives back exactly ``v``. + """ + + def __init__(self, vae: Any, transformer: Any) -> None: + super().__init__() + self.vae = vae + self.transformer = transformer + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + sr_options = Kandinsky6SROptions.from_extra(batch.extra) + vae = _resolve(self, "vae", fastvideo_args, VAELoader()) + grid, scale = _tile_grid(batch, self.transformer, vae) + device = get_local_torch_device() + vae.to(device) + scaling_factor = float(vae.scaling_factor) + tiles = [] + for tile_latent in batch.latents: + z = (tile_latent.to(device) / scaling_factor).permute(3, 0, 1, 2).unsqueeze(0) + # Upcast before the pixel-range rescale/clamp (matches the diffusers reference's decode_latents, + # which does the same right after vae.decode): decoded stays in the VAE's compute dtype otherwise, + # right before the uint8 quantization below. + decoded = vae.decode(z.to(_module_dtype(vae))).sample.float() + pixels = (vae.denormalize_data(decoded) / 255.0).clamp(0.0, 1.0) + tiles.append((pixels * 255.0)[0].cpu()) + if fastvideo_args.vae_cpu_offload: + vae.to("cpu") + + _, _, height, width = batch.lq_latents.shape + spatial_factor = vae.spatial_factor + video = stitch_tiles(tiles, grid, height * spatial_factor, width * spatial_factor, scale) + target_hw = sr_io.resolve_target_hw(sr_options.sr_target_resolution, tuple(video.shape[-2:]), + sr_options.sr_target_resize_mode) + if target_hw is not None: + video = sr_io.resize_video(video, target_hw) + + batch.output = video.float().unsqueeze(0).add_(0.5).div_(255.0).clamp_(max=1.0) + batch.num_frames, batch.height, batch.width = video.shape[1], video.shape[2], video.shape[3] + batch.latents = None + batch.lq_latents = None + return batch + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)]) + return result diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 322aa967d..addccf4ba 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -6,6 +6,7 @@ This module contains implementations of prompt encoding stages for diffusion pip """ import torch +from collections.abc import Sequence from typing import Any from torch.distributed.tensor import DTensor @@ -68,12 +69,13 @@ class TextEncodingStage(PipelineStage): assert batch.prompt is not None prompt_text: str | list[str] = batch.prompt all_indices: list[int] = list(range(len(self.text_encoders))) + max_length = self._resolve_max_length(batch, fastvideo_args) prompt_embeds_list, prompt_masks_list = self.encode_text( prompt_text, fastvideo_args, encoder_index=all_indices, return_attention_mask=True, - max_length=batch.max_sequence_length, + max_length=max_length, ) if self._last_audio_embeds is not None: batch.extra["ltx2_audio_prompt_embeds"] = self._last_audio_embeds @@ -92,7 +94,7 @@ class TextEncodingStage(PipelineStage): fastvideo_args, encoder_index=all_indices, return_attention_mask=True, - max_length=batch.max_sequence_length, + max_length=max_length, ) if self._last_audio_embeds is not None: batch.extra["ltx2_audio_negative_embeds"] = self._last_audio_embeds @@ -106,6 +108,18 @@ class TextEncodingStage(PipelineStage): return batch + def _resolve_max_length(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> int | Sequence[int | None] | None: + """The per-call ``max_length`` passed to :meth:`encode_text`. + + Default: the request's ``max_sequence_length`` applied uniformly to every text encoder + (unchanged behavior). A subclass whose encoders need different per-request semantics (e.g. + Kandinsky6's Qwen-vs-CLIP split, see ``Kandinsky6TextEncodingStage``) overrides this to + return a sequence instead, one entry per encoder in call order (``None`` for "use this + encoder's own configured default"). + """ + return batch.max_sequence_length + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify text encoding stage inputs.""" result = VerificationResult() @@ -128,7 +142,7 @@ class TextEncodingStage(PipelineStage): return_type: str = "list", # one of: "list", "dict", "stack" device: torch.device | str | None = None, dtype: torch.dtype | None = None, - max_length: int | None = None, + max_length: int | Sequence[int | None] | None = None, truncation: bool | None = None, padding: bool | str | None = None, ): @@ -148,7 +162,10 @@ class TextEncodingStage(PipelineStage): new first dimension (requires matching shapes). device: Optional device override for inputs; defaults to local torch device. dtype: Optional dtype to cast returned embeddings to. - max_length: Optional per-call tokenizer override. + max_length: Optional per-call tokenizer override. Either one value applied to every + selected encoder (existing behavior), or a sequence with one entry per encoder in + ``encoder_index`` order -- ``None`` at a position falls back to that encoder's own + configured default instead of the override. truncation: Optional per-call tokenizer override. padding: Optional per-call tokenizer override. @@ -201,7 +218,7 @@ class TextEncodingStage(PipelineStage): target_device = device if device is not None else get_local_torch_device() - for i in indices: + for pos, i in enumerate(indices): tokenizer = self.tokenizers[i] text_encoder = self.text_encoders[i] encoder_config = encoder_cfgs[i] @@ -237,9 +254,14 @@ class TextEncodingStage(PipelineStage): else: input_device = getattr(text_encoder, "_fastvideo_input_device", encoder_device) + if isinstance(max_length, Sequence) and not isinstance(max_length, str): + per_encoder_max_length = max_length[pos] if pos < len(max_length) else None + else: + per_encoder_max_length = max_length + tok_kwargs = dict(encoder_config.tokenizer_kwargs) - if max_length is not None: - tok_kwargs["max_length"] = max_length + if per_encoder_max_length is not None: + tok_kwargs["max_length"] = per_encoder_max_length elif hasattr(fastvideo_args.pipeline_config, "text_encoder_max_lengths"): tok_kwargs["max_length"] = fastvideo_args.pipeline_config.text_encoder_max_lengths[i] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 6e150a2f3..8a8e2f068 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -29,6 +29,8 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu Hunyuan15SR1080PConfig) from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig +from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig +from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig @@ -174,6 +176,19 @@ def get_model_short_name(model_id: str) -> str: return model_id +def _resolve_class_name(config: dict[str, Any]) -> str | None: + """Read `_class_name` from a model_index.json/component config dict. + + Usually a plain string, but some exports write it in the diffusers + `[library, class]` 2-element form instead -- unwrap that form to the bare + class name. + """ + class_name = config.get("_class_name") + if isinstance(class_name, list | tuple): + return class_name[-1] if class_name else None + return class_name + + def _get_config_info( model_path: str, *, @@ -202,7 +217,7 @@ def _get_config_info( else: config = maybe_download_model_index(model_path, revision=revision) - pipeline_name = config.get("_class_name", "").lower() + pipeline_name = (_resolve_class_name(config) or "").lower() matched_model_names: list[str] = [] for model_id, detector in _MODEL_NAME_DETECTORS: @@ -734,6 +749,95 @@ def _register_configs() -> None: pipeline_cls_name="Kandinsky5I2VPipeline", ) + # Kandinsky6 video super-resolution (video -> video). Registered before the Kandinsky6 TI2VA entries: VSR bundle + # names also contain the Kandinsky6 family markers and the first matching detector wins (the TI2VA detector excludes + # VSR bundles too). There is no V2V WorkloadType, so, like the Hunyuan15 SR entry, it declares no workload. Both + # bundles share the pipeline; the distilled one (PiflowScheduler, 2 steps) gets its own preset. + def _is_kandinsky6_sr(path: str) -> bool: + path_lower = path.lower() + if "kandinsky6" not in path_lower and "kandinsky-6" not in path_lower: + return False + # Only the last path component counts, so a parent directory called "sr" does not make a TI2VA checkpoint SR. + name = path_lower.rstrip("/").rsplit("/", 1)[-1].replace("_", "-").replace(".", "-") + return "6sr" in name or "superres" in name or "super-res" in name or any(token in ("sr", "vsr") + for token in name.split("-")) + + def _is_kandinsky6_sr_distilled(path: str) -> bool: + return _is_kandinsky6_sr(path) and "distill" in path.lower().rstrip("/").rsplit("/", 1)[-1] + + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Kandinsky6SRPipelineConfig, + workload_types=(), + hf_model_paths=["kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers"], + model_detectors=[_is_kandinsky6_sr_distilled], + model_family="kandinsky6_sr", + default_preset="kandinsky6_sr_distilled", + pipeline_cls_name="Kandinsky6SRPipeline", + ) + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Kandinsky6SRPipelineConfig, + workload_types=(), + hf_model_paths=["kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers"], + model_detectors=[lambda path: _is_kandinsky6_sr(path) and not _is_kandinsky6_sr_distilled(path)], + model_family="kandinsky6_sr", + default_preset="kandinsky6_sr", + pipeline_cls_name="Kandinsky6SRPipeline", + ) + + # Kandinsky6 (TI2VA) -- a single pipeline handles text-only and + # image+text calls (see Kandinsky6TI2VAPipeline's docstring), so unlike + # Kandinsky5 there is no t2v/i2v detector split: any path/class-name + # containing "kandinsky6" or "kandinsky-6" resolves here (except VSR + # bundles, see above). The only split is base vs pi-Flow distilled, laid + # out like the LTX-2 distilled/base pair: the distilled entry is registered + # FIRST and the base detector excludes distilled names, so neither can + # swallow the other. Both bundles declare the same model_index + # `_class_name`, so a local copy counts as distilled only when its + # directory name is a Kandinsky-6 name containing "distill" + # ("Kandinsky-6.0-Pro-distill-5s-Diffusers"); any other directory resolves + # to the base entry. + def _is_kandinsky6_ti2va(path: str) -> bool: + path_lower = path.lower() + return ("kandinsky6" in path_lower or "kandinsky-6" in path_lower) and not _is_kandinsky6_sr(path_lower) + + def _is_kandinsky6_distilled(path: str) -> bool: + # Looked up in the last path component only, like the VSR markers above. + name = path.lower().rstrip("/").rsplit("/", 1)[-1] + return _is_kandinsky6_ti2va(path) and "distill" in name + + def _is_kandinsky6(path: str) -> bool: + return _is_kandinsky6_ti2va(path) and not _is_kandinsky6_distilled(path) + + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Kandinsky6TI2VAConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=["kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers"], + model_detectors=[ + _is_kandinsky6_distilled, + ], + model_family="kandinsky6", + default_preset="kandinsky6_ti2va_distilled", + pipeline_cls_name="Kandinsky6TI2VAPipeline", + ) + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Kandinsky6TI2VAConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=[ + "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers", + "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", + ], + model_detectors=[ + _is_kandinsky6, + ], + model_family="kandinsky6", + default_preset="kandinsky6_ti2va", + pipeline_cls_name="Kandinsky6TI2VAPipeline", + ) + # LongCat (T2V, I2V, VC use same config; workload varies by path) register_configs( sampling_param_cls=None, @@ -1090,7 +1194,7 @@ def get_model_info( else: config = maybe_download_model_index(model_path, revision=revision) - pipeline_name = config.get("_class_name") + pipeline_name = _resolve_class_name(config) if config_info.pipeline_cls_name is not None: # The resolved (path/detector-based) config pins the pipeline class, # e.g. an I2V checkpoint whose `_class_name` would otherwise resolve @@ -1151,6 +1255,10 @@ def _register_presets() -> None: ALL_PRESETS as HYWORLD_PRESETS, ) from fastvideo.pipelines.basic.kandinsky5.presets import ( ALL_PRESETS as KANDINSKY5_PRESETS, ) + from fastvideo.pipelines.basic.kandinsky6.presets import ( + ALL_PRESETS as KANDINSKY6_PRESETS, ) + from fastvideo.pipelines.basic.kandinsky6_sr.presets import ( + ALL_PRESETS as KANDINSKY6_SR_PRESETS, ) from fastvideo.pipelines.basic.lingbotworld.presets import ( ALL_PRESETS as LINGBOTWORLD_PRESETS, ) from fastvideo.pipelines.basic.lingbotworld2.presets import ( @@ -1192,6 +1300,8 @@ def _register_presets() -> None: HUNYUAN15_PRESETS, HYWORLD_PRESETS, KANDINSKY5_PRESETS, + KANDINSKY6_PRESETS, + KANDINSKY6_SR_PRESETS, LINGBOT_VIDEO_PRESETS, LINGBOTWORLD_PRESETS, LINGBOTWORLD2_PRESETS, diff --git a/fastvideo/tests/api/test_extra_overrides_routing.py b/fastvideo/tests/api/test_extra_overrides_routing.py index 82da48bff..6e518d351 100644 --- a/fastvideo/tests/api/test_extra_overrides_routing.py +++ b/fastvideo/tests/api/test_extra_overrides_routing.py @@ -40,6 +40,18 @@ def test_passthrough_keys_cover_ltx2_audio_conditioning() -> None: "vsa_mode", "vsa_dense_first_n_steps", "vsa_dense_layers", + # Kandinsky6 SR request options and inputs (Kandinsky6SROptions, read by the stages in + # fastvideo/pipelines/stages/kandinsky6_sr.py). validate_request_batch_extra rejects them for every other + # model. The raw-latent input and audio override are SR-prefixed so a bare "audio"/"audio_sample_rate" + # request kwarg stays a no-op for every other pipeline. + "sr_resolution_scale", + "sr_tiles_batch_size", + "sr_tile_min_overlap", + "sr_target_resolution", + "sr_target_resize_mode", + "sr_lr_latent", + "sr_audio", + "sr_audio_sample_rate", } assert set(_BATCH_EXTRA_PASSTHROUGH_KEYS) == expected diff --git a/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py new file mode 100644 index 000000000..6344fafb8 --- /dev/null +++ b/fastvideo/tests/entrypoints/test_video_frame_quantize_oom_fallback.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``_quantize_video_frames_to_uint8`` and its CUDA-OOM fallback. + +Found on a real GPU run of the K6 T2IVA-distilled example (30B params, ``dit_cpu_offload=False``): the worker process +keeps the transformer resident on the device after denoising and decode, so the client process's on-device quantize +(``fastvideo/entrypoints/video_generator.py``) can be the allocation that finally OOMs, even though the data it needs +(``src``, IPC-shared from the worker) is already there and only needs a device->host copy. These tests force that path +on CPU (no GPU here) by making the on-device arithmetic raise ``torch.cuda.OutOfMemoryError`` once, and check the +fallback produces the identical result instead of propagating the error. +""" +from __future__ import annotations + +import logging + +import pytest +import torch + +from fastvideo.entrypoints.video_generator import _quantize_video_frames_to_uint8 + + +def _reference(src: torch.Tensor) -> torch.Tensor: + """Plain (non-fused) equivalent of the function under test, for comparison.""" + return ((src.float() * 255).clamp(0, 255).to(torch.uint8)).permute(2, 0, 1, 3, 4) + + +@pytest.fixture() +def sample_video() -> torch.Tensor: + # [b, c, t, h, w]; includes values outside [0, 1] to exercise clamp_(). + torch.manual_seed(0) + return torch.rand(2, 3, 4, 5, 6) * 1.3 - 0.15 + + +def test_happy_path_matches_the_plain_computation(sample_video): + out = _quantize_video_frames_to_uint8(sample_video.clone()) + assert out.shape == (sample_video.shape[2], sample_video.shape[0], sample_video.shape[1], sample_video.shape[3], + sample_video.shape[4]) + assert out.dtype == torch.uint8 + assert out.device.type == "cpu" + torch.testing.assert_close(out, _reference(sample_video), rtol=0, atol=0) + + +def test_out_of_range_values_are_clamped_not_wrapped(sample_video): + src = sample_video.clone() + src[0, 0, 0, 0, 0] = 5.0 # would wrap past 255 without clamp_() + src[0, 0, 0, 0, 1] = -5.0 + out = _quantize_video_frames_to_uint8(src) + assert out[0, 0, 0, 0, 0].item() == 255 + assert out[0, 0, 0, 0, 1].item() == 0 + + +def test_cuda_oom_on_the_fast_path_falls_back_to_a_cpu_quantize(sample_video, monkeypatch, caplog): + calls = {"n": 0} + real_clamp_ = torch.Tensor.clamp_ + + def flaky_clamp_(self, *args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise torch.cuda.OutOfMemoryError("simulated: device holds the resident DiT, nothing left to allocate") + return real_clamp_(self, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "clamp_", flaky_clamp_) + with caplog.at_level(logging.WARNING): + out = _quantize_video_frames_to_uint8(sample_video.clone()) + + assert calls["n"] == 2 # the failed fast-path attempt, then the fallback + assert out.dtype == torch.uint8 and out.device.type == "cpu" + torch.testing.assert_close(out, _reference(sample_video), rtol=0, atol=0) + assert any("out of memory" in message.lower() for message in caplog.messages) + + +def test_a_non_oom_runtime_error_is_not_swallowed(sample_video, monkeypatch): + def always_fails(self, *args, **kwargs): + raise RuntimeError("some unrelated failure") + + monkeypatch.setattr(torch.Tensor, "clamp_", always_fails) + with pytest.raises(RuntimeError, match="unrelated failure"): + _quantize_video_frames_to_uint8(sample_video.clone()) diff --git a/fastvideo/tests/ssim/test_kandinsky6_similarity.py b/fastvideo/tests/ssim/test_kandinsky6_similarity.py new file mode 100644 index 000000000..52f1f4919 --- /dev/null +++ b/fastvideo/tests/ssim/test_kandinsky6_similarity.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SSIM regression test for Kandinsky-6.0 Pro text+image-to-video (TI2VA). + +Runs a reduced-size generation (256x384, 21 frames) with a fixed seed and compares against a +device-specific reference video via MS-SSIM. The preset default is 512x768 x 121 frames; the reduced +size keeps CI cheap while the patch (1, 2, 2) and VAE temporal-4 / spatial-8 constraints stay satisfied +(dims divisible by 16, num_frames % 4 == 1) -- same constraints as Kandinsky-5's DiT, which Kandinsky6 +extends (see fastvideo/models/dits/kandinsky6.py's module docstring). + +No reference video exists yet for this test; generating and uploading one needs a GPU run against the +checkpoint plus the fastvideo/tests/ssim/reference_videos_cli.py upload flow documented in +fastvideo/tests/ssim/AGENTS.md. Until then this fails with a clear "run reference_videos_cli.py download" message (or xfails, in +bootstrap mode) the same way every other SSIM test does when its reference is missing -- see +_assert_similarity in inference_similarity_utils.py. +""" +import os + +import pytest + +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.kandinsky6.presets import KANDINSKY6_TI2VA +from fastvideo.tests.ssim.inference_similarity_utils import ( + resolve_inference_device_reference_folder, + run_text_to_video_similarity_test, +) + +logger = init_logger(__name__) + +REQUIRED_GPUS = 1 + +device_reference_folder = resolve_inference_device_reference_folder(logger) + +_PRESET_DEFAULTS = KANDINSKY6_TI2VA.defaults + +KANDINSKY6_TI2VA_PARAMS = { + "num_gpus": 1, + "model_path": "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers", + "height": 256, + "width": 384, + "num_frames": 21, + "num_inference_steps": _PRESET_DEFAULTS["num_inference_steps"], + "guidance_scale": _PRESET_DEFAULTS["guidance_scale"], + "seed": 1024, + "sp_size": 1, + "tp_size": 1, + "fps": _PRESET_DEFAULTS["fps"], + "neg_prompt": _PRESET_DEFAULTS["negative_prompt"], +} + +KANDINSKY6_TI2VA_FULL_QUALITY_PARAMS = { + **KANDINSKY6_TI2VA_PARAMS, + "height": _PRESET_DEFAULTS["height"], + "width": _PRESET_DEFAULTS["width"], + # num_frames stays at the inherited reduced 21 even at full quality (preset default: 121). +} + +KANDINSKY6_TI2VA_MODEL_TO_PARAMS = { + "Kandinsky-6.0-Pro-5s-Diffusers": KANDINSKY6_TI2VA_PARAMS, +} +FULL_QUALITY_KANDINSKY6_TI2VA_MODEL_TO_PARAMS = { + "Kandinsky-6.0-Pro-5s-Diffusers": KANDINSKY6_TI2VA_FULL_QUALITY_PARAMS, +} + +KANDINSKY6_TI2VA_TEST_PROMPTS = [ + "A curious raccoon peers through a vibrant field of yellow sunflowers, its " + "eyes wide with interest. The playful yet serene atmosphere is complemented " + "by soft natural light filtering through the petals. Mid-shot, warm and " + "cheerful tones.", +] + + +@pytest.mark.parametrize("prompt", KANDINSKY6_TI2VA_TEST_PROMPTS) +@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"]) +@pytest.mark.parametrize("model_id", list(KANDINSKY6_TI2VA_MODEL_TO_PARAMS.keys())) +def test_kandinsky6_ti2va_inference_similarity( + prompt: str, + attention_backend_name: str, + model_id: str, +) -> None: + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=KANDINSKY6_TI2VA_MODEL_TO_PARAMS, + full_quality_params_map=FULL_QUALITY_KANDINSKY6_TI2VA_MODEL_TO_PARAMS, + min_acceptable_ssim=0.93, + # Kandinsky6 shares Kandinsky5's inference-path conventions (see + # examples/inference/basic/basic_kandinsky6_ti2va.py): no FSDP inference, text encoder offloaded + # to CPU between encodes. + init_kwargs_override={ + "use_fsdp_inference": False, + "text_encoder_cpu_offload": True, + "pin_cpu_memory": True, + }, + ) diff --git a/fastvideo/tests/stages/kandinsky6_sr/conftest.py b/fastvideo/tests/stages/kandinsky6_sr/conftest.py new file mode 100644 index 000000000..a629b055c --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/conftest.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared fixtures of the Kandinsky6 SR CPU tests (tiny random models, no GPU, no weights, no reference package).""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import k6_sr_tiny # noqa: E402 + +_DIST_ENV_DEFAULTS = { + "MASTER_ADDR": "localhost", + "MASTER_PORT": "29531", + "RANK": "0", + "WORLD_SIZE": "1", + "LOCAL_RANK": "0", +} + + +@pytest.fixture(autouse=True) +def _single_process_dist_env(monkeypatch): + """torch.distributed needs a rendezvous even for one process. Fill only what is missing (never overwrite what a CI + runner leased) and undo it after the test so the settings do not leak into other tests of the same session.""" + for name, value in _DIST_ENV_DEFAULTS.items(): + if name not in os.environ: + monkeypatch.setenv(name, value) + + +@pytest.fixture(scope="session") +def distilled_bundle(tmp_path_factory) -> Path: + """Tiny distilled repo (DX head + ``PiflowScheduler``) in the official Diffusers ``Kandinsky6SRPipeline`` layout. + The directory name carries no SR marker on purpose: routing must work from the model_index ``_class_name`` alone.""" + return k6_sr_tiny.write_official_sr_repo(tmp_path_factory.mktemp("k6sr_official") / "repo") + + +@pytest.fixture(scope="session") +def flow_matching_bundle(tmp_path_factory) -> Path: + """Tiny flow-matching repo: a single-width head, the release ``FlowMatchEulerDiscreteScheduler`` config (shift 5.0), + the root ``sr_config.json`` and the dict-valued ``_kandinsky6_sr`` model_index entry.""" + return k6_sr_tiny.write_official_sr_repo(tmp_path_factory.mktemp("k6sr_flow_matching") / "repo", piflow=None) + + +@pytest.fixture() +def cpu_device(monkeypatch): + """Load and run every SR component on CPU (``get_local_torch_device`` falls back to MPS without CUDA).""" + import fastvideo.models.loader.component_loader as component_loader + import fastvideo.pipelines.stages.kandinsky6_sr as sr_stages + + cpu = torch.device("cpu") + monkeypatch.setattr(component_loader, "get_local_torch_device", lambda: cpu) + monkeypatch.setattr(sr_stages, "get_local_torch_device", lambda: cpu) + return cpu + + +@pytest.fixture() +def tiny_resolutions(monkeypatch): + """Shrink the trained base resolutions so the tiny models get tiny tiles.""" + from fastvideo.pipelines.basic.kandinsky6_sr import tiling + + monkeypatch.setitem(tiling.RESOLUTIONS, 512, list(k6_sr_tiny.TINY_RESOLUTIONS[512])) + return k6_sr_tiny.TINY_RESOLUTIONS + + +@pytest.fixture() +def cpu_args(cpu_device): + """Factory of ``FastVideoArgs`` resolved through the registry for a bundle directory, configured for CPU tests.""" + from fastvideo.fastvideo_args import FastVideoArgs + + def make(bundle: Path) -> FastVideoArgs: + args = FastVideoArgs.from_kwargs(model_path=str(bundle), num_gpus=1, dit_cpu_offload=False, + dit_layerwise_offload=False, vae_cpu_offload=False, use_fsdp_inference=False, + pin_cpu_memory=False) + args.pipeline_config.dit_precision = "fp32" + args.pipeline_config.vae_precision = "fp32" + args.pipeline_config.upsampler_precision = "fp32" + return args + + return make + + +@pytest.fixture() +def fv_args(distilled_bundle, cpu_args): + """``FastVideoArgs`` for the session's tiny bundle (see ``distilled_bundle``).""" + return cpu_args(distilled_bundle) diff --git a/fastvideo/tests/stages/kandinsky6_sr/k6_sr_release.py b/fastvideo/tests/stages/kandinsky6_sr/k6_sr_release.py new file mode 100644 index 000000000..34f27e826 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/k6_sr_release.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The real component configs (no weights) of the two official ``Kandinsky6SRPipeline`` release repos, verbatim. + +* ``kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers``: the flow-matching bundle (``FlowMatchEulerDiscreteScheduler``, a + 64-channel head, 4 Euler steps per tile by default); +* ``kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers``: the distilled bundle (``PiflowScheduler`` with + ``nfe: 2``, a 640-channel head = 64 x ``n_grid`` 10, 2 DiT calls per tile). + +The two repos have byte-identical ``vae/`` and ``latent_upscaler/`` configs, so there is one copy of each. The VAE +config stores the *string* "None" for its unused channel fields.""" +import json + +# --- flow-matching bundle (kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers) +FLOW_MATCHING_TRANSFORMER_CONFIG = json.loads(r'''{"attention_params":{"1024":{"causal":false,"chunk":false,"glob":false,"local":false,"type":"flash","window":3},"512":{"P":0.8,"P_warmup_start":0.95,"P_warmup_steps":200,"add_sta":true,"causal":false,"chunk":false,"glob":false,"local":false,"method":"topcdf","type":"nabla","wH":7,"wT":11,"wW":7,"window":3}},"attribute_overrides":{},"axes_dims":[16,24,24],"ff_dim":7168,"in_text_dim":3584,"in_text_dim2":768,"in_visual_dim":64,"instruct_type":"hybrid_anchor","model_dim":1792,"num_text_blocks":2,"num_visual_blocks":32,"out_visual_dim":64,"patch_size":[1,1,1],"sr_params":{"cap_noise_timestep":false,"fps":24,"lq_channel_noise_scale":0.0,"lq_noise_scale":0.7,"lq_noise_type":"ddpm","scale_factor":{"512":[1.0,2.0,2.0]},"scheduler_scale":5.0,"visual_size":[512]},"time_dim":512,"use_text":false,"visual_cond":true}''') +FLOW_MATCHING_SCHEDULER_CONFIG = json.loads(r'''{"_class_name":"FlowMatchEulerDiscreteScheduler","_diffusers_version":"0.39.0","base_image_seq_len":256,"base_shift":0.5,"invert_sigmas":false,"max_image_seq_len":4096,"max_shift":1.15,"num_train_timesteps":1000,"shift":5.0,"shift_terminal":null,"stochastic_sampling":false,"time_shift_type":"exponential","use_beta_sigmas":false,"use_dynamic_shifting":false,"use_exponential_sigmas":false,"use_karras_sigmas":false}''') +FLOW_MATCHING_MODEL_INDEX = json.loads(r'''{"_class_name":"Kandinsky6SRPipeline","_diffusers_version":"0.39.0","_kandinsky6_sr":{"format_version":1,"latent_upscaler":"latent_upscaler","portable_components":true,"source_checkpoint":"9454578e88f86dd9c40df6a8a79529a7e5a5ecc9","uses_patched_diffusers":true,"vae_name":"video-kvae"},"latent_upscaler":["diffusers","Kandinsky6SRLatentUpscalerBank"],"scheduler":["diffusers","FlowMatchEulerDiscreteScheduler"],"transformer":["diffusers","Kandinsky6SRTransformer3DModel"],"vae":["diffusers","Kandinsky6SRVAE"]}''') +FLOW_MATCHING_SR_CONFIG = json.loads(r'''{"component_loading":"Diffusers ModelMixin wrappers","default_resolution_scale":2.25,"description":"Kandinsky 6 SR Diffusers bundle"}''') # root sr_config.json + +# --- distilled bundle (kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers) +DISTILLED_TRANSFORMER_CONFIG = json.loads(r'''{"attention_params":{"1024":{"causal":false,"chunk":false,"glob":false,"local":false,"type":"flash","window":3},"512":{"P":0.8,"P_warmup_start":0.95,"P_warmup_steps":0,"add_sta":true,"causal":false,"chunk":false,"glob":false,"local":false,"method":"topcdf","type":"nabla","wH":7,"wT":11,"wW":7,"window":3}},"attribute_overrides":{},"axes_dims":[16,24,24],"ff_dim":7168,"in_text_dim":3584,"in_text_dim2":768,"in_visual_dim":64,"instruct_type":"hybrid_anchor","model_dim":1792,"num_text_blocks":2,"num_visual_blocks":32,"out_visual_dim":640,"patch_size":[1,1,1],"sr_params":{"cap_noise_timestep":false,"fps":24,"lq_channel_noise_scale":0.0,"lq_noise_scale":0.7,"lq_noise_type":"ddpm","scale_factor":{"512":[1.0,2.0,2.0]},"scheduler_scale":3.5,"visual_size":[512]},"time_dim":512,"use_text":false,"visual_cond":true}''') +DISTILLED_SCHEDULER_CONFIG = json.loads(r'''{"_class_name":"PiflowScheduler","_diffusers_version":"0.41.0.dev0","eps":1e-06,"final_step_size_scale":0.5,"n_grid":10,"nfe":2,"num_policy_substeps":128,"num_train_timesteps":1000,"shift":3.5}''') +DISTILLED_MODEL_INDEX = json.loads(r'''{"_class_name":"Kandinsky6SRPipeline","_diffusers_version":"0.41.0.dev0","latent_upscaler":["diffusers","Kandinsky6SRLatentUpscalerBank"],"scheduler":["diffusers","PiflowScheduler"],"transformer":["diffusers","Kandinsky6SRTransformer3DModel"],"vae":["diffusers","Kandinsky6SRVAE"]}''') + +# --- shared by both bundles +# Two layouts of the official vae/config.json with the same architecture (ch/ch_mult/z_channels/ +# num_res_blocks are identical): OLD_RELEASE_VAE_CONFIG nests the encoder/decoder knobs under +# encoder_config/decoder_config dicts; RELEASE_VAE_CONFIG, the current one, flattens them to the top +# level instead (decoder_ch / decoder_ch_mult are the only decoder-specific override). +# Kandinsky6SRVAEArchConfig.__post_init__ synthesizes the nested shape from the flat one, so both +# configs must build the identical KVAE. +OLD_RELEASE_VAE_CONFIG = json.loads(r'''{"decoder_config":{"ch":256,"ch_mult":[0.0625,1,2,4,8],"checkpoint_list":[true,true,true,true,true,true],"in_channels":"None","norm_type":"rms_norm","num_res_blocks":2,"out_ch":3,"padding_mode":"zeros","resolution":0,"temporal_compress_start_level":1,"temporal_compress_times":4,"z_channels":64},"encoder_config":{"ch":128,"ch_mult":[0.125,1,2,4,8],"checkpoint_list":[true,true,true,true,true,true],"double_z":true,"downsample_version":2,"fix_pxs":true,"in_channels":3,"norm_type":"rms_norm","num_res_blocks":2,"out_ch":"None","padding_mode":"zeros","resolution":0,"temporal_compress_start_level":1,"temporal_compress_times":4,"z_channels":64},"scaling_factor":0.910344004631042,"spatial_factor":16,"temporal_factor":4,"vae_type":"video-kvae"}''') +RELEASE_VAE_CONFIG = json.loads(r'''{"scaling_factor":0.910344004631042,"spatial_factor":16,"temporal_factor":4,"vae_type":"video-kvae","in_channels":3,"out_channels":3,"z_channels":64,"ch":128,"ch_mult":[0.125,1,2,4,8],"decoder_ch":256,"decoder_ch_mult":[0.0625,1,2,4,8],"num_res_blocks":2,"resolution":0,"padding_mode":"zeros","temporal_compress_times":4,"temporal_compress_start_level":1,"norm_type":"rms_norm","double_z":true,"downsample_version":2,"fix_pxs":true}''') +RELEASE_LATENT_UPSCALER_CONFIG = json.loads(r'''{"models":[{"model":{"architecture":"multi_scale","bare_stem":true,"bottleneck_channels":null,"depthwise":false,"dims":3,"enable_x2_entry":false,"expand_ratio":1,"gradient_checkpointing":false,"grn":false,"hidden_channels":2048,"in_channels":64,"input_skip":false,"kernel_size":3,"layer_scale_init":null,"loss_weight_2x":0.0,"loss_weight_4x":1.0,"modulated_norm":true,"modulated_output_proj":true,"num_mid_blocks":3,"num_post_blocks":3,"num_pre_blocks":5,"stage_channels":[2048,1024,512],"stochastic_depth_rate":0.0,"temporal_mix":false,"temporal_padding":"replicate","upsample_mode":"pxs_v2","upsample_padding_mode":"zeros","upscale_factor":4},"target_scale":"4x"},{"model":{"architecture":"multi_scale","bare_stem":true,"bottleneck_channels":null,"depthwise":false,"dims":3,"enable_x2_entry":true,"expand_ratio":1,"gradient_checkpointing":false,"grn":false,"hidden_channels":2048,"in_channels":64,"input_skip":false,"kernel_size":3,"layer_scale_init":null,"loss_weight_2x":0.0,"loss_weight_4x":1.0,"modulated_norm":true,"modulated_output_proj":true,"num_mid_blocks":3,"num_post_blocks":3,"num_pre_blocks":5,"stage_channels":[2048,1024,512],"stochastic_depth_rate":0.0,"temporal_mix":false,"temporal_padding":"replicate","upsample_mode":"pxs_v2","upsample_padding_mode":"zeros","upscale_factor":4,"x2_adapter_blocks":2,"x2_adapter_sources":[0,4],"x2_finisher":"pxs_residual","x2_tail_mode":"private_full"},"target_scale":"2x"}],"scaling_factor":0.910344004631042}''') diff --git a/fastvideo/tests/stages/kandinsky6_sr/k6_sr_tiny.py b/fastvideo/tests/stages/kandinsky6_sr/k6_sr_tiny.py new file mode 100644 index 000000000..3ac5f9333 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/k6_sr_tiny.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tiny random Kandinsky6 SR components and an official-Diffusers-layout SR repo writer (CPU tests only). + +A 5-level KVAE (spatial x16, temporal x4) with 4 latent channels, a 2-block text-free DiT and a multi-scale +latent-upscaler bank, so tiles of 64x64 / 64x96 / 96x64 pixels are 4x4 / 4x6 / 6x4 latents (2x2 / 2x3 / 3x2 DiT tokens). +""" +from __future__ import annotations + +import copy +import json +from pathlib import Path +from typing import Any + +import torch +from k6_sr_release import (DISTILLED_MODEL_INDEX, FLOW_MATCHING_MODEL_INDEX, FLOW_MATCHING_SCHEDULER_CONFIG, + FLOW_MATCHING_SR_CONFIG) +from safetensors.torch import save_file + +# Same code-path knobs as the released vae/config.json, tiny widths. +TINY_KVAE_ENC = dict(ch=8, ch_mult=[1, 1, 2, 2, 2], num_res_blocks=2, in_channels=3, z_channels=4, double_z=True, + temporal_compress_times=4, temporal_compress_start_level=1, norm_type="rms_norm", + padding_mode="zeros", downsample_version=2, fix_pxs=True) +TINY_KVAE_DEC = dict(ch=8, out_ch=3, ch_mult=[1, 1, 2, 2, 2], num_res_blocks=2, z_channels=4, + temporal_compress_times=4, temporal_compress_start_level=1, norm_type="rms_norm", + padding_mode="zeros") +TINY_SCALING_FACTOR = 0.5 +TINY_SR_PARAMS = dict(scale_factor={"512": [1.0, 1.0, 1.0]}, visual_size=[512], scheduler_scale=5.0, + lq_noise_scale=0.7, lq_noise_type="ddpm", lq_channel_noise_scale=0.0, cap_noise_timestep=False, + fps=24) +TINY_DIT_CFG: dict[str, Any] = dict(in_visual_dim=4, in_text_dim=8, in_text_dim2=8, time_dim=16, out_visual_dim=4, + patch_size=[1, 2, 2], model_dim=32, ff_dim=64, num_text_blocks=0, + num_visual_blocks=2, axes_dims=[8, 4, 4], visual_cond=True, + instruct_type="hybrid_anchor", attention_params={"512": {"type": "flash"}}, + use_text=False, attribute_overrides={}, sr_params=TINY_SR_PARAMS) +TINY_PIFLOW = dict(nfe=2, num_policy_substeps=8, final_step_size_scale=0.5, shift=5.0, n_grid=3, eps=1e-6) +# The released latent-upscaler architecture at toy widths (stage widths narrow 8 -> 8 -> 4 like the release's +# 2048 -> 1024 -> 512, so the narrowing shortcut convs are exercised). +TINY_LU_MODEL = {"architecture": "multi_scale", "in_channels": 4, "hidden_channels": 8, "stage_channels": [8, 8, 4], + "num_pre_blocks": 2, "num_mid_blocks": 1, "num_post_blocks": 2, "expand_ratio": 1, "kernel_size": 3, + "dims": 3, "temporal_padding": "replicate", "upsample_mode": "pxs_v2", "upsample_padding_mode": "zeros", + "modulated_norm": True, "modulated_output_proj": True, "bare_stem": True, "input_skip": False, + "temporal_mix": False, "enable_x2_entry": True, "x2_adapter_blocks": 2, "x2_adapter_sources": [0, 1], + "x2_finisher": "pxs_residual", "x2_tail_mode": "private_full"} +TINY_RESOLUTIONS = {512: [(64, 64), (64, 96), (96, 64)]} + + +def randomize(module: torch.nn.Module, seed: int, std: float = 0.05) -> torch.nn.Module: + """Deterministic non-trivial weights (many reference inits are zero, which would hide layout bugs).""" + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in module.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) * std) + return module + + +def dit_config_dict(piflow: dict | None = TINY_PIFLOW, dit_cfg: dict | None = None) -> dict[str, Any]: + """``transformer/config.json`` of the tiny DiT; a pi-Flow head is ``n_grid`` times ``in_visual_dim`` wide.""" + cfg = copy.deepcopy(dit_cfg or TINY_DIT_CFG) + if piflow is not None: + cfg["out_visual_dim"] = cfg["in_visual_dim"] * int(piflow["n_grid"]) + return cfg + + +def build_dit(piflow: dict | None = TINY_PIFLOW, seed: int = 0, dit_cfg: dict | None = None): + from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRConfig + from fastvideo.models.dits.kandinsky6_sr import Kandinsky6SRTransformer3DModel + + cfg = dit_config_dict(piflow, dit_cfg) + config = Kandinsky6SRConfig() + config.update_model_arch(cfg) + return randomize(Kandinsky6SRTransformer3DModel(config, cfg), seed + 2).eval() + + +def vae_config_dict() -> dict[str, Any]: + return dict(vae_type="video-kvae", encoder_config=copy.deepcopy(TINY_KVAE_ENC), + decoder_config=copy.deepcopy(TINY_KVAE_DEC), scaling_factor=TINY_SCALING_FACTOR, spatial_factor=16, + temporal_factor=4) + + +def build_vae(seed: int = 0): + from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEConfig + from fastvideo.models.vaes.kandinsky6_sr import Kandinsky6SRVAE + + config = Kandinsky6SRVAEConfig() + config.update_model_arch(vae_config_dict()) + return randomize(Kandinsky6SRVAE(config), seed + 1, std=0.1).eval() + + +def lu_config_dict(scales=("2x", "4x")) -> dict[str, Any]: + """``latent_upscaler/config.json``. The release config has no ``scales`` field (the bank defaults to (2, 4), which + fixes the ``_models.`` order); any other bank must declare its scales to match its ``models``.""" + config = {"models": [{"target_scale": s, "model": copy.deepcopy(TINY_LU_MODEL)} for s in scales], + "scaling_factor": TINY_SCALING_FACTOR} + if sorted(int(s.removesuffix("x")) for s in scales) != [2, 4]: + config["scales"] = [int(s.removesuffix("x")) for s in scales] + return config + + +def build_lu(scales=("2x", "4x"), seed: int = 0): + from fastvideo.configs.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerConfig + from fastvideo.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerBank + + config = Kandinsky6SRLatentUpscalerConfig(**lu_config_dict(scales)) + return randomize(Kandinsky6SRLatentUpscalerBank(config), seed + 10).eval() + + +def to_official_dit_keys(dit_state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """FastVideo DiT keys -> official Diffusers keys: the inverse of ``Kandinsky6SRArchConfig.param_names_mapping`` + (``feed_forward.mlp.fc_in/fc_out`` -> ``feed_forward.net.0.proj/net.2``, + ``time_embeddings.in/out_layer`` -> ``time_embeddings.timestep_embedder.linear_1/2``).""" + out = {} + for key, value in dit_state.items(): + key = key.replace("feed_forward.mlp.fc_in.", "feed_forward.net.0.proj.") + key = key.replace("feed_forward.mlp.fc_out.", "feed_forward.net.2.") + key = key.replace("time_embeddings.in_layer.", "time_embeddings.timestep_embedder.linear_1.") + key = key.replace("time_embeddings.out_layer.", "time_embeddings.timestep_embedder.linear_2.") + out[key] = value.detach().clone().contiguous() + return out + + +def scheduler_config_dict(piflow: dict = TINY_PIFLOW) -> dict[str, Any]: + """``scheduler/scheduler_config.json`` of a distilled bundle (``PiflowScheduler``).""" + return {"_class_name": "PiflowScheduler", "_diffusers_version": "0.41.0.dev0", "eps": piflow["eps"], + "final_step_size_scale": piflow["final_step_size_scale"], "n_grid": piflow["n_grid"], "nfe": piflow["nfe"], + "num_policy_substeps": piflow["num_policy_substeps"], "num_train_timesteps": 1000, + "shift": piflow["shift"]} + + +def flow_matching_scheduler_config_dict() -> dict[str, Any]: + """``scheduler/scheduler_config.json`` of the flow-matching bundle (the real one: shift 5.0).""" + return copy.deepcopy(FLOW_MATCHING_SCHEDULER_CONFIG) + + +def write_official_sr_repo(root: Path, *, piflow: dict | None = TINY_PIFLOW, dit_cfg: dict | None = None, + seed: int = 0) -> Path: + """Write a tiny repo in the layout of an official ``Kandinsky6SRPipeline`` bundle: + ``transformer/ vae/ latent_upscaler/ scheduler/`` with unprefixed weight keys. + + ``piflow`` picks the bundle: the distilled one (an ``n_grid``-wide head plus a ``PiflowScheduler``) or, for + ``None``, the flow-matching one (a single-width head plus the released ``FlowMatchEulerDiscreteScheduler`` config + and the release repo's extra metadata files).""" + root.mkdir(parents=True, exist_ok=True) + scheduler = scheduler_config_dict(piflow) if piflow is not None else flow_matching_scheduler_config_dict() + flow_matching = piflow is None + cfg = dit_config_dict(piflow, dit_cfg) + dit = build_dit(piflow, seed, dit_cfg) + (root / "transformer").mkdir(exist_ok=True) + (root / "transformer" / "config.json").write_text(json.dumps(cfg)) + save_file(to_official_dit_keys(dit.state_dict()), str(root / "transformer" / "diffusion_pytorch_model.safetensors")) + + vae = build_vae(seed) + (root / "vae").mkdir(exist_ok=True) + (root / "vae" / "config.json").write_text(json.dumps(vae_config_dict())) + save_file({k: v.detach().clone().contiguous() for k, v in vae.state_dict().items()}, + str(root / "vae" / "diffusion_pytorch_model.safetensors")) + + (root / "scheduler").mkdir(exist_ok=True) + (root / "scheduler" / "scheduler_config.json").write_text(json.dumps(scheduler)) + + index = copy.deepcopy(FLOW_MATCHING_MODEL_INDEX if flow_matching else DISTILLED_MODEL_INDEX) + if flow_matching: + (root / "sr_config.json").write_text(json.dumps(FLOW_MATCHING_SR_CONFIG)) + lu = build_lu(seed=seed) + (root / "latent_upscaler").mkdir(exist_ok=True) + (root / "latent_upscaler" / "config.json").write_text(json.dumps(lu_config_dict())) + save_file({k: v.detach().clone().contiguous() for k, v in lu.state_dict().items()}, + str(root / "latent_upscaler" / "diffusion_pytorch_model.safetensors")) + (root / "model_index.json").write_text(json.dumps(index)) + return root + + +def write_mp4(path: Path, num_frames: int, fps, size=(64, 96), audio_seconds: float | None = None, + seed: int = 0) -> None: + """Write a random-noise mpeg4 clip (+ optional 440 Hz aac track); skips the test if PyAV cannot encode it.""" + from fractions import Fraction + + import numpy as np + import pytest + + av = pytest.importorskip("av") + rng = np.random.default_rng(seed) + rate = Fraction(fps).limit_denominator(1001) + try: + with av.open(str(path), "w") as container: + video = container.add_stream("mpeg4", rate=rate) + video.width, video.height, video.pix_fmt = size[1], size[0], "yuv420p" + audio = None + if audio_seconds is not None: + audio = container.add_stream("aac", rate=44100) + audio.layout = "mono" + for i in range(num_frames): + frame = av.VideoFrame.from_ndarray(rng.integers(0, 255, (*size, 3), dtype=np.uint8), format="rgb24") + frame.pts = i + frame.time_base = Fraction(1) / rate + for packet in video.encode(frame): + container.mux(packet) + for packet in video.encode(): + container.mux(packet) + if audio is not None: + t = np.arange(int(44100 * audio_seconds)) / 44100 + wave = (0.3 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)[None, :] + audio_frame = av.AudioFrame.from_ndarray(wave, format="fltp", layout="mono") + audio_frame.sample_rate = 44100 + audio_frame.pts = 0 + audio_frame.time_base = Fraction(1, 44100) + for packet in audio.encode(audio_frame): + container.mux(packet) + for packet in audio.encode(): + container.mux(packet) + except Exception as exc: # pragma: no cover - PyAV build without the mpeg4 / aac encoders + pytest.skip(f"PyAV cannot encode the test clip: {exc}") diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_arch_config.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_arch_config.py new file mode 100644 index 000000000..2d1ca5472 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_arch_config.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR DiT arch config: release configs parse, unsupported configurations fail by name.""" +from __future__ import annotations + +import copy + +import pytest +from k6_sr_release import DISTILLED_TRANSFORMER_CONFIG, FLOW_MATCHING_TRANSFORMER_CONFIG + +from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRArchConfig, Kandinsky6SRConfig + + +def _arch(**overrides) -> Kandinsky6SRArchConfig: + config = Kandinsky6SRConfig() + config.update_model_arch({**copy.deepcopy(FLOW_MATCHING_TRANSFORMER_CONFIG), **overrides}) + return config.arch_config + + +@pytest.mark.parametrize("release", [FLOW_MATCHING_TRANSFORMER_CONFIG, DISTILLED_TRANSFORMER_CONFIG]) +def test_release_configs_expose_the_sampling_parameters(release): + config = Kandinsky6SRConfig() + config.update_model_arch(copy.deepcopy(release)) + arch = config.arch_config + assert arch.visual_size == 512 + assert arch.rope_scale_factor == (1.0, 2.0, 2.0) + assert arch.lq_noise_scale == 0.7 + assert arch.input_channels == 2 * 64 + 1 + assert arch.num_attention_heads == 1792 // 64 + assert set(release) <= arch.known_config_keys() + + +@pytest.mark.parametrize("overrides, match", [ + ({"use_text": True}, "use_text"), + ({"attribute_overrides": {"visual_cond": False}}, "attribute_overrides"), + ({"instruct_type": "noise"}, "hybrid_anchor"), + ({"visual_cond": False}, "hybrid_anchor"), + ({"sr_params": {"lq_noise_type": "linear"}}, "lq_noise_type"), + ({"sr_params": {"lq_channel_noise_scale": 0.1}}, "lq_channel_noise_scale"), + ({"sr_params": {"cap_noise_timestep": True}}, "cap_noise_timestep"), +]) +def test_unsupported_configurations_raise_and_name_the_field(overrides, match): + with pytest.raises(NotImplementedError, match=match): + _arch(**overrides) + + +def test_temporal_patching_is_rejected(): + with pytest.raises(ValueError, match="patch_size"): + _arch(patch_size=[2, 1, 1]) + + +def test_param_names_mapping_maps_the_official_diffusers_names(): + import re + + mapping = Kandinsky6SRArchConfig().param_names_mapping + + def rename(key: str) -> str: + for pattern, replacement in mapping.items(): + key = re.sub(pattern, replacement, key) + return key + + assert rename("visual_transformer_blocks.0.feed_forward.net.0.proj.weight") == ( + "visual_transformer_blocks.0.feed_forward.mlp.fc_in.weight") + assert rename("visual_transformer_blocks.0.feed_forward.net.2.bias") == ( + "visual_transformer_blocks.0.feed_forward.mlp.fc_out.bias") + assert rename("time_embeddings.timestep_embedder.linear_1.weight") == "time_embeddings.in_layer.weight" + assert rename("time_embeddings.timestep_embedder.linear_2.bias") == "time_embeddings.out_layer.bias" + assert rename("visual_transformer_blocks.0.self_attention.to_query.weight") == ( + "visual_transformer_blocks.0.self_attention.to_query.weight") diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_components.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_components.py new file mode 100644 index 000000000..0f7047587 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_components.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR component loading: KVAE wrapper, latent-upscaler bank config, strict loaders, dtype policy (CPU).""" +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import k6_sr_tiny +import pytest +import torch +from safetensors.torch import load_file, save_file + +from fastvideo.configs.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerConfig +from fastvideo.models.loader.component_loader import PipelineComponentLoader + + +# --------------------------------------------------------------------------- KVAE wrapper + + +def test_vae_wrapper_contract(): + vae = k6_sr_tiny.build_vae() + assert vae.scaling_factor == k6_sr_tiny.TINY_SCALING_FACTOR and vae.config.scaling_factor == 0.5 + assert vae.spatial_factor == 16 and vae.temporal_factor == 4 + assert all(key.startswith(("encoder.", "decoder.")) for key in vae.state_dict()) # bundle keys load unchanged + pixels = torch.arange(0, 256, dtype=torch.float32) + torch.testing.assert_close(vae.denormalize_data(vae.normalize_data(pixels)), pixels) + assert vae.normalize_data(torch.tensor(128.0)).item() == 0.0 # x / 128 - 1, the KVAE convention (not / 127.5) + + +def test_vae_encode_decode_geometry(): + vae = k6_sr_tiny.build_vae() + video = torch.rand(1, 3, 9, 32, 48) * 2 - 1 + with torch.no_grad(): + latent, split_list = vae.encode(video) + decoded = vae.decode(latent).sample + assert split_list == [9] + assert latent.shape == (1, 4, (9 - 1) // 4 + 1, 2, 3) # temporal x4 (causal), spatial x16 + assert decoded.shape == video.shape + + +def test_vae_without_architecture_fails_loudly(): + from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEConfig + from fastvideo.models.vaes.kandinsky6_sr import Kandinsky6SRVAE + + with pytest.raises(ValueError, match="encoder_config"): + Kandinsky6SRVAE(Kandinsky6SRVAEConfig()) + config = Kandinsky6SRVAEConfig() + with pytest.raises(ValueError, match="video-kvae"): + config.update_model_arch({"vae_type": "other-vae"}) + + +# --------------------------------------------------------------------------- LU bank + + +def test_lu_bank_state_keys_and_dispatch(): + bank = k6_sr_tiny.build_lu() + assert bank.scales == (2, 4) + assert all(key.startswith("_models.") for key in bank.state_dict()) + z = torch.randn(1, 4, 5, 4, 6) + with torch.no_grad(): + assert bank(z * 0.5, 2).shape == (1, 4, 5, 8, 12) + assert bank(z * 0.5, 4).shape == (1, 4, 5, 16, 24) + with pytest.raises(ValueError, match="no x3 entry"): + bank(z, 3) + only_x4 = k6_sr_tiny.build_lu(scales=("4x", )) + assert only_x4.scales == (4, ) + with pytest.raises(ValueError, match="no x2 entry"): + only_x4(z, 2) + + +def _lu_config(**model_overrides): + config = k6_sr_tiny.lu_config_dict(("4x", )) + config["models"][0]["model"].update(model_overrides) + return config + + +@pytest.mark.parametrize("config,match", [ + ({"models": [{"target_scale": "2x"}]}, "models\\[0\\]"), + ({"models": [{"target_scale": "8x", "model": k6_sr_tiny.TINY_LU_MODEL}]}, "target_scale"), + (_lu_config(bogus=1), "bogus"), + (_lu_config(architecture="flat"), "architecture"), + (_lu_config(input_skip=True), "input_skip"), + (_lu_config(upsample_mode="pixel_shuffle"), "upsample_mode"), + (_lu_config(motion_attention={"num_heads": 2}), "motion_attention"), + (_lu_config(x2_tail_mode="shared"), "x2_tail_mode"), + (_lu_config(enable_x2_entry=False), "enable_x2_entry"), + (_lu_config(hidden_channels=16), "stage_channels"), +]) +def test_lu_config_errors_name_the_offending_entry_or_key(config, match): + with pytest.raises(ValueError, match=match): + Kandinsky6SRLatentUpscalerConfig(**config) + + +def test_lu_config_omitted_key_means_the_training_default(): + config = _lu_config() + del config["models"][0]["model"]["input_skip"] # training default: input_skip=True, not implemented + with pytest.raises(ValueError, match="input_skip.*omitted"): + Kandinsky6SRLatentUpscalerConfig(**config) + + +def test_lu_bank_duplicate_scales_are_rejected(): + with pytest.raises(ValueError, match="duplicate"): + Kandinsky6SRLatentUpscalerConfig(**k6_sr_tiny.lu_config_dict(("2x", "2x"))) + + +def test_lu_config_survives_the_pipeline_config_json_round_trip(): + import dataclasses + + config = Kandinsky6SRLatentUpscalerConfig(**k6_sr_tiny.lu_config_dict()) + restored = Kandinsky6SRLatentUpscalerConfig() + dumped = json.loads(json.dumps(dataclasses.asdict(config))) + dumped.pop("arch_config") + restored.update_model_config(dumped) + assert restored.models == config.models and restored.scaling_factor == config.scaling_factor + + +# --------------------------------------------------------------------------- loaders + + +def _load(name, root, fv_args, subdir): + # Official component configs (like the Hub repo's) carry no ``_class_name``; ``ComposedPipelineBase.load_modules`` + # hands the model_index classes to the loaders through ``fastvideo_args._model_index_class_names``, mirrored here. + index = json.loads((root / "model_index.json").read_text()) + fv_args._model_index_class_names = {key: spec[1] for key, spec in index.items() if isinstance(spec, list)} + return PipelineComponentLoader.load_module(name, str(root / subdir), "diffusers", fv_args) + + +def test_loaders_build_the_components_with_the_intended_dtypes(distilled_bundle, fv_args): + fv_args.pipeline_config.dit_precision = "bf16" + fv_args.pipeline_config.vae_precision = "bf16" + fv_args.pipeline_config.upsampler_precision = "bf16" + dit = _load("transformer", distilled_bundle, fv_args, "transformer") + vae = _load("vae", distilled_bundle, fv_args, "vae") + lu = _load("latent_upscaler", distilled_bundle, fv_args, "latent_upscaler") + + assert dit.compute_dtype == torch.bfloat16 and next(vae.parameters()).dtype == torch.bfloat16 + assert next(lu.parameters()).dtype == torch.bfloat16 and not lu.training + assert all(not p.requires_grad for p in lu.parameters()) + params = dict(dit.named_parameters()) + # fp32 where the reference computes in fp32 (autocast(float32) regions / .float() norms), bf16 for the linears + assert params["pooled_bias"].dtype == torch.float32 + assert params["time_embeddings.in_layer.weight"].dtype == torch.float32 + assert params["visual_transformer_blocks.0.visual_modulation.out_layer.weight"].dtype == torch.float32 + assert params["visual_transformer_blocks.0.self_attention.query_norm.weight"].dtype == torch.float32 + assert params["visual_transformer_blocks.0.self_attention.to_query.weight"].dtype == torch.bfloat16 + assert dit.out_layer.out_layer.weight.shape[0] == 4 * 3 * 4 # C * n_grid * prod(patch) + assert dit.config.arch_config.sr_params == k6_sr_tiny.TINY_SR_PARAMS + assert all(not p.requires_grad for p in dit.parameters()) + assert vae.scaling_factor == 0.5 and lu.scales == (2, 4) + + +@pytest.mark.parametrize("bundle, class_name", [("distilled_bundle", "PiflowScheduler"), + ("flow_matching_bundle", "FlowMatchEulerDiscreteScheduler")]) +def test_each_bundle_loads_its_own_scheduler(request, cpu_args, bundle, class_name): + root = request.getfixturevalue(bundle) + scheduler = _load("scheduler", root, cpu_args(root), "scheduler") + assert type(scheduler).__name__ == class_name + + +def test_latent_upscaler_may_be_declared_with_the_kandinsky6_library(): + # the official save_pretrained writes ["kandinsky6", ...], the Hub repo ["diffusers", ...] + from fastvideo.models.loader.component_loader import ComponentLoader + + for library in ("diffusers", "kandinsky6"): + assert type(ComponentLoader.for_module_type("latent_upscaler", library)).__name__ == "UpsamplerLoader" + with pytest.raises(AssertionError, match="latent_upscaler must be loaded from"): + ComponentLoader.for_module_type("latent_upscaler", "transformers") + + +def test_bf16_dit_stays_close_to_fp32(distilled_bundle, fv_args): + # dtype boundary: the fp32 residual stream / bf16 linear inputs must not drift (a bf16 residual would). + fp32 = _load("transformer", distilled_bundle, fv_args, "transformer") + fv_args.pipeline_config.dit_precision = "bf16" + bf16 = _load("transformer", distilled_bundle, fv_args, "transformer") + x = torch.randn(2, 3, 4, 6, 2 * 4 + 1, generator=torch.Generator().manual_seed(0)) + t = torch.tensor([300.0, 800.0]) + pos = [torch.arange(3), torch.arange(2), torch.arange(3)] + with torch.no_grad(): + ref = fp32(x, t, pos) + out = bf16(x.to(torch.bfloat16), t, pos) + assert torch.isfinite(out).all() + assert (out.float() - ref).abs().max() < 0.05 * ref.abs().max() + + +def _corrupt(bundle: Path, tmp_path: Path, component: str, mutate) -> Path: + copy_root = tmp_path / f"corrupt_{component}" + shutil.copytree(bundle, copy_root) + path = copy_root / component / "diffusion_pytorch_model.safetensors" + tensors = load_file(str(path)) + mutate(tensors) + save_file(tensors, str(path)) + return copy_root + + +def _drop_first(fragment): + def mutate(tensors): + tensors.pop(next(key for key in tensors if fragment in key)) + + return mutate + + +def _add_unexpected(tensors): + tensors["model.unexpected_extra_weight"] = torch.zeros(3) + + +@pytest.mark.parametrize("component,name,mutate", [ + ("transformer", "transformer", _drop_first("visual_transformer_blocks.0.self_attention.to_query.weight")), + ("transformer", "transformer", _add_unexpected), + ("vae", "vae", _drop_first("decoder")), + ("vae", "vae", lambda t: t.__setitem__("model.unexpected_extra_weight", torch.zeros(3))), + ("latent_upscaler", "latent_upscaler", _drop_first("_models.0.")), + ("latent_upscaler", "latent_upscaler", lambda t: t.__setitem__("_models.0.unexpected", torch.zeros(3))), +]) +def test_partial_or_extra_weights_never_load_silently(distilled_bundle, fv_args, tmp_path, component, name, mutate): + # A silently partial load runs plausible-looking garbage: every SR component must refuse it. + corrupted = _corrupt(distilled_bundle, tmp_path, component, mutate) + with pytest.raises((ValueError, RuntimeError)): + _load(name, corrupted, fv_args, component) + + +def test_unknown_component_class_is_reported(distilled_bundle, fv_args, tmp_path): + import json + + copy_root = tmp_path / "bad_class" + shutil.copytree(distilled_bundle, copy_root) + config_path = copy_root / "latent_upscaler" / "config.json" + config = json.loads(config_path.read_text()) + config["_class_name"] = "NoSuchBank" + config_path.write_text(json.dumps(config)) + with pytest.raises(ValueError, match="not supported"): + _load("latent_upscaler", copy_root, fv_args, "latent_upscaler") + + +# --------------------------------------------------------------------------- _ChunkedConv3d + + +@pytest.mark.parametrize("stride", [1, 2]) +def test_chunked_conv3d_matches_a_plain_conv3d(monkeypatch, stride): + # No existing test exercises the actual chunking branch (real inputs never hit _MAX_CONV_NUMEL on + # CPU-test-sized tensors), so force it by lowering the threshold. Covers both the stride-1 path + # (chunk-with-carried-tail) and the strided path (one kernel-width window at a time). + from fastvideo.models.vaes import kandinsky6_sr as k6_vae + + # Small enough that the full 3072-element input needs chunking, but big enough that one + # kernel-width window (1 x 6 x 3 x 8 x 8 = 1152 elements) still fits in a single window. + monkeypatch.setattr(k6_vae, "_MAX_CONV_NUMEL", 1200) + torch.manual_seed(0) + x = torch.randn(1, 6, 8, 8, 8) # numel = 3072 > 1200 + + plain = torch.nn.Conv3d(6, 4, kernel_size=(3, 1, 1), stride=(stride, 1, 1), padding=0) + chunked = k6_vae._ChunkedConv3d(6, 4, kernel_size=(3, 1, 1), stride=(stride, 1, 1), padding=0) + chunked.load_state_dict(plain.state_dict()) + + with torch.no_grad(): + expected = plain(x) + actual = chunked(x) + torch.testing.assert_close(actual, expected) + + +def test_chunked_conv3d_strided_raises_only_when_a_single_window_still_overflows(monkeypatch): + from fastvideo.models.vaes import kandinsky6_sr as k6_vae + + monkeypatch.setattr(k6_vae, "_MAX_CONV_NUMEL", 10) # even one kernel-width window overflows this + x = torch.randn(1, 6, 8, 8, 8) + chunked = k6_vae._ChunkedConv3d(6, 4, kernel_size=(3, 1, 1), stride=(2, 1, 1), padding=0) + with pytest.raises(ValueError, match="too big"): + chunked(x) + + +def test_chunked_conv3d_rejects_temporal_padding(monkeypatch): + from fastvideo.models.vaes import kandinsky6_sr as k6_vae + + monkeypatch.setattr(k6_vae, "_MAX_CONV_NUMEL", 400) # force the chunking branch to be reached at all + conv = k6_vae._ChunkedConv3d(6, 4, kernel_size=(3, 1, 1), padding=(1, 0, 0)) + x = torch.randn(1, 6, 8, 8, 8) + with pytest.raises(ValueError, match="unpadded"): + conv(x) diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_official_release_configs.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_official_release_configs.py new file mode 100644 index 000000000..e0149a431 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_official_release_configs.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The real component configs (no weights) of both official ``Kandinsky6SRPipeline`` repos against the SR modules. + +Each component is built on the meta device from the release config exactly as the loaders would, and the tensor counts / +key groups are compared with the safetensors headers of the Hub repos (459 DiT, 378 KVAE, 489 latent-upscaler tensors). +The flow-matching repo (``Kandinsky-6.0-VSR-5s-Diffusers``) and the distilled repo +(``Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers``) share the VAE and the latent upscaler and differ in the DiT head, +the scheduler and ``sr_params``. This guards the port against drifting away from the official release without needing +a GPU, the weights or the Hub. +""" +from __future__ import annotations + +import copy +import json +from collections import Counter +from typing import Any, NamedTuple + +import pytest +import torch + +from fastvideo.configs.pipelines.kandinsky6_sr_options import Kandinsky6SROptions +from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRConfig +from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEConfig +from fastvideo.pipelines.basic.kandinsky6_sr.presets import KANDINSKY6_SR +from k6_sr_release import (DISTILLED_MODEL_INDEX, DISTILLED_SCHEDULER_CONFIG, DISTILLED_TRANSFORMER_CONFIG, + FLOW_MATCHING_MODEL_INDEX, FLOW_MATCHING_SCHEDULER_CONFIG, FLOW_MATCHING_SR_CONFIG, + FLOW_MATCHING_TRANSFORMER_CONFIG, OLD_RELEASE_VAE_CONFIG, + RELEASE_LATENT_UPSCALER_CONFIG, RELEASE_VAE_CONFIG) + + +class _Bundle(NamedTuple): + transformer: dict[str, Any] + scheduler: dict[str, Any] + model_index: dict[str, Any] + tiny_bundle: str # name of the conftest fixture with the tiny repo of the same kind + head_width: int # ``out_visual_dim``: the DiT output channels + scheduler_scale: float # ``sr_params.scheduler_scale``: the timestep shift the model was trained with + + +_BUNDLES = { + "flow_matching": _Bundle(FLOW_MATCHING_TRANSFORMER_CONFIG, FLOW_MATCHING_SCHEDULER_CONFIG, + FLOW_MATCHING_MODEL_INDEX, "flow_matching_bundle", 64, 5.0), + "distilled": _Bundle(DISTILLED_TRANSFORMER_CONFIG, DISTILLED_SCHEDULER_CONFIG, DISTILLED_MODEL_INDEX, + "distilled_bundle", 640, 3.5), +} +both_bundles = pytest.mark.parametrize("bundle", list(_BUNDLES.values()), ids=list(_BUNDLES)) + + +def _prefixes(keys, depth=2): + return Counter(".".join(key.split(".")[:depth]) for key in keys) + + +def _meta_dit(transformer_config: dict[str, Any]): + from fastvideo.models.dits.kandinsky6_sr import Kandinsky6SRTransformer3DModel + + config = Kandinsky6SRConfig() + config.update_model_arch(transformer_config) + with torch.device("meta"): + return Kandinsky6SRTransformer3DModel(config, transformer_config) # rejects any undeclared config key + + +def test_the_release_transformer_configs_differ_only_in_head_timestep_warp_and_nabla_warmup(): + flow, distilled = copy.deepcopy(FLOW_MATCHING_TRANSFORMER_CONFIG), copy.deepcopy(DISTILLED_TRANSFORMER_CONFIG) + assert (flow.pop("out_visual_dim"), distilled.pop("out_visual_dim")) == (64, 640) + assert (flow["sr_params"].pop("scheduler_scale"), distilled["sr_params"].pop("scheduler_scale")) == (5.0, 3.5) + assert (flow["attention_params"]["512"].pop("P_warmup_steps"), + distilled["attention_params"]["512"].pop("P_warmup_steps")) == (200, 0) + assert flow == distilled + + +@both_bundles +def test_real_transformer_config_builds_the_release_dit(distributed_setup, bundle): + dit = _meta_dit(bundle.transformer) + state = dit.state_dict() + per_block = Counter(key.split(".")[1] for key in state if key.startswith("visual_transformer_blocks.")) + assert len(state) == 459 and len(per_block) == 32 and set(per_block.values()) == {14} + assert len([key for key in state if not key.startswith("visual_transformer_blocks.")]) == 11 + # ``out_visual_dim`` is the full head width (the distilled head is 64 * n_grid 10); only + # ``out_layer.out_layer.{weight,bias}`` differ between the two repos + assert dit.out_layer.out_layer.weight.shape == (bundle.head_width, 1792) + assert dit.out_layer.out_layer.bias.shape == (bundle.head_width, ) + assert dit.input_channels == 2 * 64 + 1 # noised latent | anchor | anchor mask + arch = dit.config.arch_config + assert arch.rope_scale_factor == (1.0, 2.0, 2.0) and arch.visual_size == 512 + assert arch.sr_params["scheduler_scale"] == bundle.scheduler_scale + + +@both_bundles +def test_real_scheduler_config_loads_natively_with_the_training_shift(distributed_setup, fv_args, tmp_path, bundle): + from fastvideo.models.loader.component_loader import SchedulerLoader + + (tmp_path / "scheduler").mkdir() + (tmp_path / "scheduler" / "scheduler_config.json").write_text(json.dumps(bundle.scheduler)) + scheduler = SchedulerLoader().load(str(tmp_path / "scheduler"), fv_args) + assert type(scheduler).__name__ == bundle.scheduler["_class_name"] + # The scheduler's shift is the timestep warp the SR DiT was trained with. + assert scheduler.config.shift == bundle.scheduler_scale + if getattr(scheduler, "is_piflow", False): + assert scheduler.n_grid * 64 == bundle.head_width # the head holds n_grid predictions per channel + + +@both_bundles +def test_real_model_index_passes_through_load_modules(distributed_setup, request, cpu_args, bundle): + from fastvideo.pipelines.basic.kandinsky6_sr.kandinsky6_sr_pipeline import Kandinsky6SRPipeline + + repo = request.getfixturevalue(bundle.tiny_bundle) + index = json.loads((repo / "model_index.json").read_text()) + assert index == bundle.model_index # the tiny repo carries the release index verbatim + pipeline = Kandinsky6SRPipeline(str(repo), cpu_args(repo)) + pipeline.post_init() + # components are the [library, class] entries; the dict-valued ``_kandinsky6_sr`` metadata is skipped, not loaded + assert set(pipeline.modules) == {"transformer", "vae", "scheduler", "latent_upscaler"} + assert type(pipeline.get_module("scheduler")).__name__ == bundle.scheduler["_class_name"] + + +def test_the_flow_matching_release_carries_kandinsky6_sr_metadata_and_the_preset_default_scale(): + assert isinstance(FLOW_MATCHING_MODEL_INDEX["_kandinsky6_sr"], dict) + assert "_kandinsky6_sr" not in DISTILLED_MODEL_INDEX + # root sr_config.json: the default scale of the release is the preset default (and the request default) + default_scale = FLOW_MATCHING_SR_CONFIG["default_resolution_scale"] + assert default_scale == KANDINSKY6_SR.defaults["sr_resolution_scale"] == Kandinsky6SROptions.sr_resolution_scale == 2.25 + + +@pytest.mark.parametrize("vae_config", [RELEASE_VAE_CONFIG, OLD_RELEASE_VAE_CONFIG], + ids=["flat_2026_09_30", "nested_pre_2026_09_30"]) +def test_real_vae_config_builds_the_release_kvae_and_maps_official_keys(vae_config): + # The Hub re-exported vae/config.json's layout (flat vs. nested encoder_config/decoder_config, + # see k6_sr_release.py) between when these two fixtures were captured; both must build the + # identical KVAE since the underlying architecture never changed. + from fastvideo.models.vaes.kandinsky6_sr import Kandinsky6SRVAE + + config = Kandinsky6SRVAEConfig() + config.update_model_arch(dict(vae_config)) + with torch.device("meta"): + vae = Kandinsky6SRVAE(config) + official = vae.state_dict() + assert len(official) == 378 + assert _prefixes(official) == Counter({ + "decoder.up": 238, "encoder.down": 86, "decoder.mid": 28, "encoder.mid": 12, "decoder.norm_out": 5, + "decoder.conv_in": 2, "decoder.conv_out": 2, "encoder.conv_in": 2, "encoder.conv_out": 2, "encoder.norm_out": 1}) + assert vae.scaling_factor == pytest.approx(0.910344004631042) and vae.spatial_factor == 16 + # the official keys load strictly as they are + vae.load_state_dict(official, strict=True, assign=True) + + +def test_real_latent_upscaler_config_builds_the_release_bank(): + from fastvideo.configs.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerConfig + from fastvideo.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerBank + + with torch.device("meta"): + bank = Kandinsky6SRLatentUpscalerBank(Kandinsky6SRLatentUpscalerConfig(**RELEASE_LATENT_UPSCALER_CONFIG)) + state = bank.state_dict() + # Entries are indexed in ``config.scales`` order (default (2, 4)), as in the release safetensors. + assert len(state) == 489 and _prefixes(state) == Counter({"_models.0": 311, "_models.1": 178}) + assert not any("motion_attention" in entry["model"] for entry in RELEASE_LATENT_UPSCALER_CONFIG["models"]) + x2 = [key for key in state if key.startswith("_models.0.")] + assert _prefixes((key.split(".", 2)[2] for key in x2), 2)["x2_branch.blocks"] == 44 + assert _prefixes((key.split(".", 2)[2] for key in x2), 2)["x2_branch.adapter"] == 28 + assert bank.scales == (2, 4) diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_registry_routing.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_registry_routing.py new file mode 100644 index 000000000..9ac1129be --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_registry_routing.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry routing: SR bundles (both official VSR repo ids and local copies) resolve to the SR config, Kandinsky6 TI2VA +paths keep resolving to Kandinsky6.""" +from __future__ import annotations + +import json +from dataclasses import fields +from pathlib import Path + +import pytest + +from fastvideo import registry +from fastvideo.api.presets import get_preset, validate_preset_selection +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig +from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig +from fastvideo.configs.pipelines.kandinsky6_sr_options import Kandinsky6SROptions, SR_OPTION_FIELDS +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +def _model_dir(root: Path, name: str, class_name: str) -> str: + path = root / name + path.mkdir(parents=True) + (path / "model_index.json").write_text(json.dumps({"_class_name": class_name, "_diffusers_version": "0.35.0"})) + return str(path) + + +def _config_cls(path: str): + return registry._get_config_info(path).pipeline_config_cls + + +SR_FLOW_ID = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" +SR_DISTILLED_ID = "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers" +_PRESET = {SR_FLOW_ID: "kandinsky6_sr", SR_DISTILLED_ID: "kandinsky6_sr_distilled"} + + +@pytest.mark.parametrize("repo_id", [SR_FLOW_ID, SR_DISTILLED_ID]) +def test_official_vsr_repo_ids_resolve_to_the_sr_config(repo_id): + # exact registry match: no model_index fetch + assert repo_id in registry.get_registered_model_paths() + assert registry.get_pipeline_config_cls_from_name(repo_id) is Kandinsky6SRPipelineConfig + assert registry.get_preset_selection(repo_id) == (_PRESET[repo_id], "kandinsky6_sr") + info = registry._get_config_info(repo_id) + assert info.pipeline_cls_name == "Kandinsky6SRPipeline" + assert info.workload_types == () + + +@pytest.mark.parametrize("repo_id", [SR_FLOW_ID, SR_DISTILLED_ID]) +def test_get_model_info_resolves_both_vsr_repo_ids_to_the_sr_pipeline(monkeypatch, repo_id): + model_index = {"_class_name": "Kandinsky6SRPipeline", "_diffusers_version": "0.39.0"} + monkeypatch.setattr(registry, "maybe_download_model_index", lambda *_args, **_kwargs: model_index) + info = registry.get_model_info(repo_id) + assert info.pipeline_cls.__name__ == "Kandinsky6SRPipeline" + assert info.pipeline_config_cls is Kandinsky6SRPipelineConfig + + +@pytest.mark.parametrize("name", ["Kandinsky-6.0-VSR-5s-Diffusers", "Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers"]) +def test_local_copies_of_the_official_vsr_repos_resolve_to_the_sr_config(tmp_path, name): + # the "distill" marker of the distilled VSR repo must not route it to the distilled TI2VA entry + path = _model_dir(tmp_path, name, "Kandinsky6SRPipeline") + assert _config_cls(path) is Kandinsky6SRPipelineConfig + assert registry.get_preset_selection(path) == (_PRESET[f"kandinskylab/{name}"], "kandinsky6_sr") + + +@pytest.mark.parametrize("name,class_name", [ + ("Kandinsky-6.0-Pro-5s-Diffusers", "Kandinsky6TI2VAPipeline"), + ("Kandinsky-6.0-Pro-sft-5s-Diffusers", "Kandinsky6TI2VAPipeline"), + ("Kandinsky-6.0-Pro-distill-5s-Diffusers", "Kandinsky6TI2VAPipeline"), + ("kandinsky6-t2va", "Kandinsky6TI2VAPipeline"), + ("my_kandinsky6_export", "Kandinsky6TI2VAPipeline"), + ("kandinsky-6-export", "SomeOtherPipeline"), +]) +def test_kandinsky6_ti2va_paths_still_resolve_to_kandinsky6(tmp_path, name, class_name): + assert _config_cls(_model_dir(tmp_path, name, class_name)) is Kandinsky6TI2VAConfig + + +def test_a_parent_directory_called_sr_does_not_turn_a_ti2va_checkpoint_into_vsr(tmp_path): + path = _model_dir(tmp_path / "sr" / "vsr_experiments", "kandinsky6-export", "Kandinsky6TI2VAPipeline") + assert _config_cls(path) is Kandinsky6TI2VAConfig + + +@pytest.mark.parametrize("name", [ + "Kandinsky-6.0-VSR-distilled2steps-5s", + "kandinsky6_sr", + "kandinsky6-sr-export", + "Kandinsky6SR", + "kandinsky-6-super-res", +]) +def test_sr_bundle_names_resolve_to_the_sr_config(tmp_path, name): + # the model_index class name of an SR bundle is Kandinsky6SRPipeline; the names are also tried on their own + assert _config_cls(_model_dir(tmp_path, name, "Kandinsky6SRPipeline")) is Kandinsky6SRPipelineConfig + + +def test_sr_bundle_resolves_from_the_model_index_class_name_alone(tmp_path): + assert _config_cls(_model_dir(tmp_path, "vsr_bundle_v1", "Kandinsky6SRPipeline")) is Kandinsky6SRPipelineConfig + + +def test_the_ti2va_detector_never_matches_an_sr_bundle(): + # routing must not depend on registration order alone: exactly one config matches SR inputs + sr_inputs = [ + "kandinsky-6.0-vsr-5s-diffusers", "kandinsky-6.0-vsr-distilled2steps-5s-diffusers", + "kandinsky-6.0-vsr-distilled2steps-5s", "kandinsky6srpipeline", "kandinsky6_sr", "kandinsky-6-sr" + ] + for text in sr_inputs: + matching = {registry._CONFIG_REGISTRY[model_id].pipeline_config_cls + for model_id, detector in registry._MODEL_NAME_DETECTORS if detector(text)} + assert matching == {Kandinsky6SRPipelineConfig}, (text, matching) + ti2va_inputs = [ + "kandinsky-6.0-pro-sft-5s-diffusers", "kandinsky-6.0-pro-distill-5s-diffusers", "kandinsky6-t2va", + "kandinsky6ti2vapipeline", "kandinsky-6-export" + ] + for text in ti2va_inputs: + matching = {registry._CONFIG_REGISTRY[model_id].pipeline_config_cls + for model_id, detector in registry._MODEL_NAME_DETECTORS if detector(text)} + assert matching == {Kandinsky6TI2VAConfig}, (text, matching) + + +def test_sr_config_entry_metadata(): + info = registry._get_config_info(_registered_sr_path()) + assert info.pipeline_cls_name == "Kandinsky6SRPipeline" + assert info.workload_types == () # no V2V WorkloadType: not exposed as a workload option + assert info.default_preset == "kandinsky6_sr" and info.model_family == "kandinsky6_sr" + + +def _registered_sr_path() -> str: + import tempfile + + return _model_dir(Path(tempfile.mkdtemp()), "kandinsky6_sr", "Kandinsky6SRPipeline") + + +def test_pipeline_class_is_discovered_under_the_model_index_name(distilled_bundle): + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.pipeline_registry import PipelineType, get_pipeline_registry + + index_name = json.loads((distilled_bundle / "model_index.json").read_text())["_class_name"] + cls = get_pipeline_registry(PipelineType.BASIC).resolve_pipeline_cls(index_name, PipelineType.BASIC, + WorkloadType.T2V) + assert cls.__name__ == "Kandinsky6SRPipeline" == index_name + assert cls._required_config_modules == ["transformer", "vae", "latent_upscaler", "scheduler"] + + +@pytest.mark.parametrize("name, steps", [("Kandinsky-6.0-VSR-5s-Diffusers", 4), + ("Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", 2)]) +def test_preset_defaults_reach_the_sampling_param(tmp_path, name, steps): + preset = get_preset("kandinsky6_sr", "kandinsky6_sr") + assert preset.stage_schemas[0].kind == "super_resolution" + sp = SamplingParam.from_pretrained(_model_dir(tmp_path, name, "Kandinsky6SRPipeline")) + assert sp.num_inference_steps == steps + assert sp.seed == 42 + assert Kandinsky6SROptions.from_extra(preset.defaults) == Kandinsky6SROptions() + assert not any(hasattr(sp, key) for key in SR_OPTION_FIELDS) + assert sp.guidance_scale == 1.0 and sp.return_frames is False and sp.save_video is True + + +def test_stage_overrides_accept_only_the_sr_knobs(): + validate_preset_selection("kandinsky6_sr", "kandinsky6_sr", stage_overrides={"sr": {"num_inference_steps": 3}}) + from fastvideo.api.errors import ConfigValidationError + + with pytest.raises(ConfigValidationError): + validate_preset_selection("kandinsky6_sr", "kandinsky6_sr", stage_overrides={"sr": {"seed": 1}}) + with pytest.raises(ConfigValidationError): + validate_preset_selection("kandinsky6_sr", "kandinsky6_sr", stage_overrides={"denoise": {}}) + + +def test_sr_options_are_not_global_sampling_or_batch_fields(): + for cls in (SamplingParam, ForwardBatch): + assert not set(SR_OPTION_FIELDS).intersection(f.name for f in fields(cls)) + from fastvideo.utils import shallow_asdict + + ForwardBatch(**shallow_asdict(SamplingParam()), eta=0.0, n_tokens=1, VSA_sparsity=0.0) + + +@pytest.mark.parametrize("section", ["extensions", "stage_overrides"]) +def test_request_extensions_reach_sr_options(tmp_path, section): + from fastvideo.api.compat import normalize_generation_request, request_to_batch_extra, request_to_sampling_param + + model_path = _model_dir(tmp_path, "kandinsky6_sr", "Kandinsky6SRPipeline") + overrides = {"sr_resolution_scale": 4, "sr_target_resolution": "fullhd"} + request = normalize_generation_request({ + "inputs": {"video_path": "clip.mp4"}, + section: overrides if section == "extensions" else {"sr": overrides}, + }) + sp = request_to_sampling_param(request, model_path=model_path) + options = Kandinsky6SROptions.from_extra(request_to_batch_extra(request)) + assert sp.video_path == "clip.mp4" + assert options.sr_resolution_scale == 4 and options.sr_target_resolution == "fullhd" + assert options.sr_tiles_batch_size == 1 and options.sr_tile_min_overlap == 0.20 + assert sp.seed == 42 and sp.num_inference_steps == 4 + + +def test_generator_hooks_are_class_attributes_not_dataclass_fields(): + assert Kandinsky6SRPipelineConfig.prompt_optional is True + assert Kandinsky6SRPipelineConfig.output_shape_from_input is True + assert not getattr(Kandinsky6TI2VAConfig, "prompt_optional", False) + names = {f.name for f in fields(Kandinsky6SRPipelineConfig)} + assert "prompt_optional" not in names and "output_shape_from_input" not in names # no inventory entries needed + + +def test_pipeline_config_needs_no_text_encoder(): + config = Kandinsky6SRPipelineConfig() + config.check_pipeline_config() # equal-length (empty) encoder / precision / preprocess lists, vae_sp off + assert config.text_encoder_configs == () and config.dit_precision == "bf16" + assert config.vae_config.load_encoder and config.vae_config.load_decoder diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_stages_pipeline.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_stages_pipeline.py new file mode 100644 index 000000000..36ac6f7c7 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_stages_pipeline.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR pipeline on tiny random bundles (CPU): stage chain, scheduler-driven sampling, inputs and outputs.""" +from __future__ import annotations + +from pathlib import Path + +import k6_sr_tiny +import pytest +import torch + +from fastvideo.configs.pipelines.kandinsky6_sr_options import SR_OPTION_FIELDS +from fastvideo.pipelines.basic.kandinsky6_sr import sr_io, tiling +from fastvideo.pipelines.basic.kandinsky6_sr.kandinsky6_sr_pipeline import Kandinsky6SRPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +pytestmark = pytest.mark.usefixtures("distributed_setup", "tiny_resolutions") + + +def _batch(**kwargs) -> ForwardBatch: + defaults = dict(data_type="video", prompt="", seed=42, sr_resolution_scale=2.0, save_video=True, + return_frames=False, num_inference_steps=2) + defaults.update(kwargs) + extra = dict(defaults.pop("extra", {})) + extra.update({key: defaults.pop(key) for key in SR_OPTION_FIELDS if key in defaults}) + return ForwardBatch(**defaults, extra=extra) + + +def _load(bundle: Path, args) -> Kandinsky6SRPipeline: + pipeline = Kandinsky6SRPipeline(str(bundle), args) + pipeline.post_init() + return pipeline + + +def _frames(result: ForwardBatch) -> torch.Tensor: + return (result.output[0] * 255).clamp_(0, 255).to(torch.uint8) + + +@pytest.fixture() +def clip(tmp_path) -> Path: + path = tmp_path / "clip.mp4" + k6_sr_tiny.write_mp4(path, 20, 24, size=(64, 96), audio_seconds=2.0) + return path + + +@pytest.fixture(params=["distilled", "flow_matching"]) +def bundle_args(request, cpu_args): + bundle = request.getfixturevalue(f"{request.param}_bundle") + return bundle, cpu_args(bundle) + + +def _count_calls(monkeypatch, obj, name) -> list[int]: + calls = [] + original = getattr(obj, name) + + def wrapped(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(obj, name, wrapped) + return calls + + +# x2.25 pre-upscales 64x96 by 1.125 to the 16 px grid (64x112), then tiles at x2. +@pytest.mark.parametrize("scale, size, expected_hw", [(2.0, (64, 96), (128, 192)), (2.25, (64, 96), (128, 224)), + (4.0, (64, 64), (256, 256))]) +def test_pipeline_upscales_and_keeps_fps_and_audio(bundle_args, tmp_path, scale, size, expected_hw): + bundle, args = bundle_args + clip = tmp_path / "clip.mp4" + k6_sr_tiny.write_mp4(clip, 20, 24, size=size, audio_seconds=2.0) + pipeline = _load(bundle, args) + assert list(pipeline._stage_name_mapping) == [ + "video_encoding_stage", "latent_preparation_stage", "denoising_stage", "decoding_stage" + ] + result = pipeline.forward(_batch(video_path=str(clip), sr_resolution_scale=scale), args) + + assert result.output.dtype == torch.float32 and tuple(result.output.shape) == (1, 3, 17, *expected_hw) + assert (result.height, result.width, result.num_frames) == (*expected_hw, 17) + assert result.fps == result.extra["output_fps"] == 24 + assert result.extra["audio_sample_rate"] == sr_io.AUDIO_SAMPLE_RATE + assert result.extra["audio"].shape[0] == round(17 / 24 * sr_io.AUDIO_SAMPLE_RATE) + assert result.latents is None and result.lq_latents is None # large intermediates are not shipped back + + +def test_each_tile_runs_num_inference_steps_through_the_scheduler(bundle_args, clip, monkeypatch): + bundle, args = bundle_args + pipeline = _load(bundle, args) + transformer = pipeline.get_module("transformer") + scheduler = pipeline.get_module("scheduler") + forwards = _count_calls(monkeypatch, transformer, "forward") + steps = _count_calls(monkeypatch, scheduler, "step") + pipeline.forward(_batch(video_path=str(clip), num_inference_steps=3), args) + num_tiles = tiling.plan_tiles(64, 96, 512, 2, 0.2, 16).num_tiles + assert num_tiles > 1 + assert len(forwards) == len(steps) == num_tiles * 3 + + +def test_same_seed_is_reproducible_and_a_new_seed_changes_the_result(bundle_args, clip): + bundle, args = bundle_args + pipeline = _load(bundle, args) + first = _frames(pipeline.forward(_batch(video_path=str(clip)), args)) + again = _frames(pipeline.forward(_batch(video_path=str(clip)), args)) + other = _frames(pipeline.forward(_batch(video_path=str(clip), seed=7), args)) + assert torch.equal(first, again) + assert not torch.equal(first, other) + + +def test_tile_groups_are_seeded_by_their_first_tile(fv_args, distilled_bundle, clip): + pipeline = _load(distilled_bundle, fv_args) + single = _frames(pipeline.forward(_batch(video_path=str(clip)), fv_args)) + grouped = _frames(pipeline.forward(_batch(video_path=str(clip), sr_tiles_batch_size=2), fv_args)) + assert single.shape == grouped.shape + assert not torch.equal(single, grouped) # tile 1 draws its noise from the group seeded with seed + 0 + + +def test_raw_latent_input_skips_the_video_and_keeps_an_audio_override(fv_args, distilled_bundle, monkeypatch): + monkeypatch.setattr(sr_io, "read_video", lambda *a, **k: pytest.fail("the video must not be decoded")) + latent = torch.randn(3, 4, 4, 6, generator=torch.Generator().manual_seed(5)) + audio = torch.linspace(-0.5, 0.5, 1000) + batch = _batch() + batch.extra.update(sr_lr_latent=latent, sr_audio=audio, sr_audio_sample_rate=16000) + result = _load(distilled_bundle, fv_args).forward(batch, fv_args) + assert tuple(result.output.shape[-2:]) == (4 * 16 * 2, 6 * 16 * 2) + assert torch.equal(result.extra["audio"], audio) and result.extra["audio_sample_rate"] == 16000 + assert "sr_lr_latent" not in result.extra and "sr_audio" not in result.extra + + +@pytest.mark.parametrize("extra, kwargs, error", [ + ({"sr_lr_latent": torch.zeros(3, 4, 4, 6)}, {"sr_resolution_scale": 2.25}, "2.25"), + ({"sr_lr_latent": torch.zeros(4, 4, 6)}, {}, r"\[T, C, H, W\]"), + ({"sr_lr_latent": torch.zeros(3, 4, 4, 6)}, {"video_path": "clip.mp4"}, "not both"), + ({}, {"sr_resolution_scale": 3.0}, "2, 4 or 2.25"), + ({}, {"video_path": "clip.mp4", "sr_target_resolution": "8k"}, "sr_target_resolution"), + ({}, {"video_path": ["a.mp4", "b.mp4"]}, "one video"), + ({}, {}, "video_path"), +]) +def test_bad_requests_fail_before_any_model_runs(fv_args, distilled_bundle, extra, kwargs, error): + batch = _batch(**kwargs) + batch.extra.update(extra) + with pytest.raises(ValueError, match=error): + _load(distilled_bundle, fv_args).forward(batch, fv_args) + + +@pytest.mark.parametrize("bundle_name", ["distilled", "flow_matching"]) +def test_a_transformer_from_the_other_bundle_is_rejected(request, cpu_args, clip, monkeypatch, bundle_name): + bundle = request.getfixturevalue(f"{bundle_name}_bundle") + args = cpu_args(bundle) + pipeline = _load(bundle, args) + arch = pipeline.get_module("transformer").config.arch_config + scheduler = pipeline.get_module("scheduler") + # Give each transformer the other bundle's head width. + monkeypatch.setattr(arch, "out_visual_dim", + arch.in_visual_dim if bundle_name == "distilled" else arch.in_visual_dim * 10) + forwards = _count_calls(monkeypatch, pipeline.get_module("transformer"), "forward") + with pytest.raises(ValueError, match=f"head is {arch.out_visual_dim} channels wide, but {type(scheduler).__name__}"): + pipeline.forward(_batch(video_path=str(clip)), args) + assert not forwards + + +def test_zero_steps_are_rejected(fv_args, distilled_bundle, clip): + with pytest.raises(ValueError, match="num_inference_steps must be >= 1"): + _load(distilled_bundle, fv_args).forward(_batch(video_path=str(clip), num_inference_steps=0), fv_args) + + +def test_a_missing_upscaler_scale_is_reported(fv_args, distilled_bundle, tmp_path): + clip = tmp_path / "square.mp4" + k6_sr_tiny.write_mp4(clip, 20, 24, size=(64, 64)) + pipeline = _load(distilled_bundle, fv_args) + bank = pipeline.get_module("latent_upscaler") + bank.config.scales = (2, ) # a bank that only serves x2 + del bank._models[1] + with pytest.raises(ValueError, match="no x4"): + pipeline.forward(_batch(video_path=str(clip), sr_resolution_scale=4.0), fv_args) + + +def test_warmup_request_without_a_video_completes(fv_args, distilled_bundle): + result = _load(distilled_bundle, fv_args).forward(_batch(save_video=False), fv_args) + assert result.output.ndim == 5 + + +def test_output_codes_survive_video_generators_truncating_conversion(fv_args, distilled_bundle, clip): + result = _load(distilled_bundle, fv_args).forward(_batch(video_path=str(clip)), fv_args) + codes = result.output * 255 + fraction = codes - codes.floor() + assert bool((((fraction - 0.5).abs() < 1e-3) | (codes == 255)).all()) + + +def test_delivery_resize_fits_the_bucket(fv_args, distilled_bundle, clip): + result = _load(distilled_bundle, fv_args).forward(_batch(video_path=str(clip), sr_target_resolution="96x64"), + fv_args) + assert (result.height, result.width) == (64, 96) + + +def test_tiles_reach_stitch_tiles_as_float_not_pre_quantized_uint8(fv_args, distilled_bundle, clip, monkeypatch): + # Tiles must stay float through blending and be quantized to uint8 once, inside + # stitch_tiles -- not per tile beforehand (which would blend already-rounded values). Spy on the + # actual call the decoding stage makes. + import fastvideo.pipelines.stages.kandinsky6_sr as k6_sr_stages + + seen_tiles = [] + original = k6_sr_stages.stitch_tiles + + def spy(tiles, *args, **kwargs): + seen_tiles.extend(tiles) + return original(tiles, *args, **kwargs) + + monkeypatch.setattr(k6_sr_stages, "stitch_tiles", spy) + _load(distilled_bundle, fv_args).forward(_batch(video_path=str(clip)), fv_args) + + assert seen_tiles, "stitch_tiles was never called" + for tile in seen_tiles: + assert tile.dtype.is_floating_point, f"tile reached stitch_tiles pre-quantized as {tile.dtype}" diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_tiling.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_tiling.py new file mode 100644 index 000000000..737ecb3e2 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_tiling.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR tile planning, stitching and scale handling (CPU).""" +from __future__ import annotations + +import pytest +import torch + +from fastvideo.pipelines.basic.kandinsky6_sr.tiling import (TileGrid, plan_tiles, pre_upscale, resolve_scale, + stitch_tiles) + + +@pytest.mark.parametrize("size", [(512, 768), (768, 512), (512, 512), (576, 864), (720, 1280), (1088, 1920), + (208, 384), (240, 432)]) +@pytest.mark.parametrize("scale", [2, 4]) +def test_tile_grids_cover_the_frame_overlap_and_stay_latent_aligned(size, scale): + height, width = size + grid = plan_tiles(height, width, 512, scale, 0.2, 16) + assert grid.tops[0] == 0 and grid.lefts[0] == 0 + assert grid.tops[-1] + grid.tile_h == max(height, grid.tile_h) + assert grid.lefts[-1] + grid.tile_w == max(width, grid.tile_w) + assert (grid.tile_h * scale, grid.tile_w * scale) in ((512, 512), (512, 768), (768, 512)) + latent = grid.to_latent(16) + assert latent.tile_h * 16 == grid.tile_h and latent.num_tiles == grid.num_tiles + for a, b in zip(grid.tops, grid.tops[1:]): + assert grid.tile_h - (b - a) >= 0.2 * grid.tile_h - 16 + + +def test_unaligned_grid_is_rejected_for_the_latent_tiles(): + with pytest.raises(ValueError, match="not aligned"): + TileGrid(256, 384, (0, 100), (0, 128)).to_latent(16) + + +def test_extract_returns_row_major_tiles(): + video = torch.arange(4 * 6).reshape(1, 4, 6) + tiles = TileGrid(2, 3, (0, 2), (0, 3)).extract(video) + assert [tile[0, 0, 0].item() for tile in tiles] == [0, 3, 12, 15] + + +def test_scale_that_does_not_divide_the_base_resolution_is_rejected(): + with pytest.raises(ValueError, match="does not divide"): + plan_tiles(512, 768, 512, 3, 0.2, 16) + + +def test_stitching_constant_tiles_reproduces_the_constant(): + grid = TileGrid(32, 48, (0, 16, 32), (0, 24, 48)) # uneven overlaps, as at the frame edges + tiles = [torch.full((3, 4, 64, 96), 173, dtype=torch.uint8) for _ in range(grid.num_tiles)] + out = stitch_tiles(tiles, grid, 64, 96, scale=2, frame_chunk=3) + assert out.shape == (3, 4, 128, 192) and out.dtype == torch.uint8 + assert bool(((out.int() - 173).abs() <= 1).all()) # normalised blend, truncated to uint8 + + +def test_scale_requests(): + assert resolve_scale(2) == (2, 1.0) and resolve_scale(4.0) == (4, 1.0) and resolve_scale(2.25) == (2, 1.125) + for bad in (1, 3, 8, 2.5, 0): + with pytest.raises(ValueError, match="2, 4 or 2.25"): + resolve_scale(bad) + + +def test_pre_upscale_targets_are_vae_aligned_and_reach_the_2_25_total(): + video = torch.randint(0, 256, (2, 3, 512, 768), dtype=torch.uint8) + up = pre_upscale(video, 1.125, 16, frame_chunk=1) + assert up.shape[-2:] == (576, 864) and up.dtype == torch.uint8 + assert torch.equal(up, pre_upscale(video, 1.125, 16)) # frames are independent: chunking changes nothing + odd = pre_upscale(torch.zeros(1, 3, 100, 130, dtype=torch.uint8), 1.125, 16) + assert odd.shape[-2] % 16 == 0 and odd.shape[-1] % 16 == 0 diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_transformer.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_transformer.py new file mode 100644 index 000000000..7a6f312ef --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_transformer.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR DiT: batched forward semantics on tiny random weights (CPU).""" +from __future__ import annotations + +import k6_sr_tiny +import pytest +import torch + +B, T, H, W, C = 3, 3, 4, 6, 4 +IN = 2 * C + 1 # noised latent | anchor | anchor mask + + +def _rope_pos(): + return [torch.arange(T), torch.arange(H // 2), torch.arange(W // 2)] + + +def _inputs(seed=1, batch=B): + generator = torch.Generator().manual_seed(seed) + return torch.randn(batch, T, H, W, IN, generator=generator), torch.linspace(100.0, 900.0, batch) + + +@pytest.mark.parametrize("piflow", [k6_sr_tiny.TINY_PIFLOW, None]) +def test_forward_shape_and_head_width(piflow): + dit = k6_sr_tiny.build_dit(piflow) + x, t = _inputs() + with torch.no_grad(): + y = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + n_grid = piflow["n_grid"] if piflow else 1 + assert y.shape == (B, T, H, W, C * n_grid) + assert torch.isfinite(y).all() + + +def test_samples_of_a_batch_do_not_interact(): + # The reference packs samples into one ragged sequence; the batched layout must keep them independent + # (attention over the wrong tokens or a shared time embedding would leak between tiles). + dit = k6_sr_tiny.build_dit() + x, t = _inputs() + with torch.no_grad(): + batched = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + single = torch.cat([dit(x[i:i + 1], t[i:i + 1], _rope_pos(), (1.0, 1.0, 1.0)) for i in range(B)]) + torch.testing.assert_close(batched, single, rtol=1e-5, atol=1e-5) + + +def test_output_depends_on_the_time_embedding_and_pooled_bias(): + dit = k6_sr_tiny.build_dit() + x, t = _inputs() + with torch.no_grad(): + base = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + other_t = dit(x, t * 0.5, _rope_pos(), (1.0, 1.0, 1.0)) + dit.pooled_bias.add_(1.0) # dropping the bias corrupts the time conditioning (see DiffusionTransformer3D) + shifted = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + assert not torch.allclose(base, other_t) + assert not torch.allclose(base, shifted) + + +def test_rope_scale_factor_changes_the_output(): + dit = k6_sr_tiny.build_dit() + x, t = _inputs() + with torch.no_grad(): + a = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + b = dit(x, t, _rope_pos(), (1.0, 2.0, 2.0)) + assert not torch.allclose(a, b) + + +def test_wrong_channel_count_raises_a_clear_error(): + dit = k6_sr_tiny.build_dit() + narrow = torch.randn(1, T, H, W, C) + with pytest.raises(ValueError, match=f"expects {IN}"): + dit(narrow, torch.tensor([500.0]), _rope_pos()) + + +def test_unsupported_forward_arguments_are_rejected(): + dit = k6_sr_tiny.build_dit() + x, t = _inputs(batch=1) + with pytest.raises(TypeError, match="sparse_params"): + dit(x, t, _rope_pos(), sparse_params={"to_fractal": True}) + + +def test_state_dict_layout_matches_the_bundle_contract(): + dit = k6_sr_tiny.build_dit() + keys = set(dit.state_dict()) + assert "pooled_bias" in keys + assert "time_embeddings.in_layer.weight" in keys and "time_embeddings.out_layer.bias" in keys + assert "visual_embeddings.in_layer.weight" in keys + for name in ("to_query", "to_key", "to_value", "out_layer"): + assert f"visual_transformer_blocks.0.self_attention.{name}.weight" in keys + assert "visual_transformer_blocks.1.self_attention.query_norm.weight" in keys + assert "visual_transformer_blocks.0.feed_forward.mlp.fc_in.weight" in keys + assert "visual_transformer_blocks.0.visual_modulation.out_layer.weight" in keys + assert "out_layer.modulation.out_layer.weight" in keys and "out_layer.out_layer.weight" in keys + # text-free: no text tower, no cross-attention + assert not any("text" in key or "cross_attention" in key for key in keys) + # 6 modulation params (self-attention + ffn) per block; pi-Flow head width C * n_grid * prod(patch) + assert dit.state_dict()["visual_transformer_blocks.0.visual_modulation.out_layer.weight"].shape[0] == 6 * 32 + assert dit.state_dict()["out_layer.out_layer.weight"].shape[0] == 4 * 3 * 4 + + +def test_official_keys_round_trip_through_the_mapping(): + from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + + dit = k6_sr_tiny.build_dit() + official_keys = k6_sr_tiny.to_official_dit_keys(dit.state_dict()) + assert not any(key.startswith("model.") for key in official_keys) # the official checkpoint has no prefix + custom, _ = hf_to_custom_state_dict(official_keys, get_param_names_mapping(dit.param_names_mapping)) + assert set(custom) == set(dit.state_dict()) # every official key lands on a model key, none left over + dit.load_state_dict(custom, strict=True) + + +def test_forward_survives_fp8_quantization(): + # With transformer_quant=get_quantization_config("FP8")() (per docs/inference/optimizations.md), + # convert_model_to_fp8 pops `.weight` off every FP8-tagged ReplicatedLinear (e.g. + # self_attention.to_query), so the blocks must take their compute dtype from the model's + # never-quantized `compute_dtype`, not from a quantized layer's `.weight.dtype`. + from fastvideo.layers.quantization import get_quantization_config + from fastvideo.layers.quantization.fp8_config import convert_model_to_fp8 + + cfg = k6_sr_tiny.dit_config_dict(k6_sr_tiny.TINY_PIFLOW) + from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRConfig + from fastvideo.models.dits.kandinsky6_sr import Kandinsky6SRTransformer3DModel + + config = Kandinsky6SRConfig(quant_config=get_quantization_config("FP8")()) + config.update_model_arch(cfg) + dit = k6_sr_tiny.randomize(Kandinsky6SRTransformer3DModel(config, cfg), seed=2).eval() + + convert_model_to_fp8(dit) + # Sanity: quantization actually removed `.weight` from at least one attention projection, so this + # test would have caught the original bug (it wasn't a no-op). + assert not hasattr(dit.visual_transformer_blocks[0].self_attention.to_query, "weight") + + x, t = _inputs(batch=1) + with torch.no_grad(): + y = dit(x, t, _rope_pos(), (1.0, 1.0, 1.0)) + assert torch.isfinite(y).all() diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_generator_hooks.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_generator_hooks.py new file mode 100644 index 000000000..b56c32b50 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_generator_hooks.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VideoGenerator hooks used by prompt-less, input-shaped pipelines (Kandinsky6 SR).""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +import fastvideo.entrypoints.video_generator as video_generator +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.entrypoints.video_generator import VideoGenerator +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +def _generator(args: FastVideoArgs, executor=None) -> VideoGenerator: + generator = VideoGenerator.__new__(VideoGenerator) + generator.fastvideo_args = args + generator.executor = executor + return generator + + +def _kandinsky6_ti2va_args(tmp_path: Path) -> FastVideoArgs: + model = tmp_path / "kandinsky6-t2va" + model.mkdir() + (model / "model_index.json").write_text(json.dumps({"_class_name": "Kandinsky6TI2VAPipeline", + "_diffusers_version": "0.35.0"})) + return FastVideoArgs.from_kwargs(model_path=str(model), num_gpus=1) + + +def test_prompt_is_optional_for_the_sr_pipeline_and_the_output_is_named_after_the_video(fv_args, tmp_path, monkeypatch): + seen = {} + generator = _generator(fv_args) + monkeypatch.setattr(generator, "_generate_single_video", + lambda prompt, sampling_param, fastvideo_args, **kw: seen.update(prompt=prompt, **kw) or {}) + sampling_param = SamplingParam(video_path="/data/clips/holiday.mp4", output_path=str(tmp_path / "out")) + generator._generate_video_impl(prompt=None, sampling_param=sampling_param, fastvideo_args=fv_args) + assert seen["prompt"] == "" + assert Path(seen["output_path"]) == tmp_path / "out" / "holiday.mp4" + # ... and a second run does not overwrite the first result + (tmp_path / "out" / "holiday.mp4").write_bytes(b"x") + generator._generate_video_impl(prompt=None, sampling_param=sampling_param, fastvideo_args=fv_args) + assert Path(seen["output_path"]).name == "holiday_1.mp4" + + +def test_other_pipelines_still_require_a_prompt(tmp_path): + args = _kandinsky6_ti2va_args(tmp_path) + with pytest.raises(ValueError, match="Either prompt or prompt_txt"): + _generator(args)._generate_video_impl(prompt=None, sampling_param=SamplingParam(video_path="clip.mp4"), + fastvideo_args=args) + + +class _StubExecutor: + + def __init__(self, shape): + self.shape = shape + + def execute_forward(self, batch, args): + out = ForwardBatch(data_type="video") + out.output = torch.rand(*self.shape) + return out + + +def _run(generator, args, monkeypatch): + big_allocations = [] + real_empty = torch.empty + + def spy(*shape, **kwargs): + size = shape[0] if len(shape) == 1 and isinstance(shape[0], (tuple, list)) else shape + if len(size) == 5: + big_allocations.append(tuple(size)) + return real_empty(0) + return real_empty(*shape, **kwargs) + + monkeypatch.setattr(video_generator.torch, "empty", spy) + sampling_param = SamplingParam(save_video=False, return_frames=True, height=720, width=1280, num_frames=125) + result = generator._generate_single_video(prompt="", sampling_param=sampling_param, fastvideo_args=args, + output_path="unused.mp4") + return result, big_allocations + + +def test_sr_does_not_preallocate_the_return_buffer_from_the_request_geometry(fv_args, monkeypatch): + shape = (1, 3, 9, 64, 96) + result, allocations = _run(_generator(fv_args, _StubExecutor(shape)), fv_args, monkeypatch) + assert allocations == [] # (1, 3, 125, 720, 1280) floats would be ~1.3 GB of pinned memory for nothing + assert tuple(result["samples"].shape) == shape and result["size"] == (64, 96, 9) # geometry read from the output + + +def test_other_pipelines_keep_the_preallocation(tmp_path, monkeypatch): + args = _kandinsky6_ti2va_args(tmp_path) + _, allocations = _run(_generator(args, _StubExecutor((1, 3, 9, 64, 96))), args, monkeypatch) + assert allocations == [(1, 3, 125, 720, 1280)] + + +class _ExtraStubExecutor(_StubExecutor): + """Returns ``extra`` the way a worker does: the worker's own ``batch.fps`` never comes back, only ``extra``.""" + + def __init__(self, shape, extra): + super().__init__(shape) + self.extra = extra + + def execute_forward(self, batch, args): + out = super().execute_forward(batch, args) + out.extra.update(self.extra) + return out + + +def _saved_fps(generator, args, monkeypatch, request_fps=24) -> int: + saved = {} + monkeypatch.setattr(generator, "_save_video_with_audio_ffmpeg_pipe", lambda **kw: saved.update(kw) or True) + monkeypatch.setattr(video_generator.imageio, "mimsave", lambda path, frames, fps, format: saved.update(fps=fps)) + sampling_param = SamplingParam(save_video=True, return_frames=False, fps=request_fps) + generator._generate_single_video(prompt="", sampling_param=sampling_param, fastvideo_args=args, + output_path="unused.mp4") + return saved["fps"] + + +@pytest.mark.parametrize("with_audio", [True, False], ids=["audio_single_pass", "video_only"]) +def test_the_mp4_is_written_at_the_fps_reported_by_the_worker(fv_args, monkeypatch, with_audio): + extra = {"output_fps": 16} + if with_audio: + extra.update(audio=torch.zeros(1600), audio_sample_rate=44100) + generator = _generator(fv_args, _ExtraStubExecutor((1, 3, 9, 64, 96), extra)) + assert _saved_fps(generator, fv_args, monkeypatch, request_fps=24) == 16 + + +def test_other_pipelines_keep_writing_at_the_request_fps(fv_args, monkeypatch): + generator = _generator(fv_args, _ExtraStubExecutor((1, 3, 9, 64, 96), {})) + assert _saved_fps(generator, fv_args, monkeypatch, request_fps=30) == 30 + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_sr_options_reach_batch_extra_for_both_generator_apis(fv_args, tmp_path, monkeypatch, legacy): + generator = _generator(fv_args) + seen = {} + monkeypatch.setattr(generator, "_generate_single_video", + lambda **kwargs: seen.update(kwargs) or {}) + monkeypatch.setattr(generator, "_wrap_legacy_result", lambda result: result) + options = {"sr_resolution_scale": 4, "sr_tiles_batch_size": 2, "sr_tile_min_overlap": 0.3, + "sr_target_resolution": "hd", "sr_target_resize_mode": "exact"} + if legacy: + generator.generate_video(video_path="clip.mp4", output_path=str(tmp_path), **options) + else: + generator._generate_single_request(video_generator.normalize_generation_request({ + "inputs": {"video_path": "clip.mp4"}, "output": {"output_path": str(tmp_path)}, + "extensions": options, + })) + assert seen["_extra_overrides"] == options + assert all(not hasattr(seen["sampling_param"], key) for key in options) + + +def test_legacy_generator_rejects_sr_options_for_ti2va(tmp_path): + args = _kandinsky6_ti2va_args(tmp_path) + with pytest.raises(ValueError, match="only supported by Kandinsky6 SR"): + _generator(args).generate_video(prompt="a cat", sr_resolution_scale=4) diff --git a/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_io.py b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_io.py new file mode 100644 index 000000000..55da4030a --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_sr/test_kandinsky6_sr_video_io.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 SR video / audio input and delivery resizing (CPU).""" +from __future__ import annotations + +from fractions import Fraction +from pathlib import Path + +import k6_sr_tiny +import pytest +import torch + +from fastvideo.pipelines.basic.kandinsky6_sr import sr_io + +av = pytest.importorskip("av") + + +def _resample_then_clip(video: torch.Tensor, src_fps: float) -> tuple[torch.Tensor, int]: + """The SR input contract written the straightforward way: decode everything, resample to 24 fps, keep the first + 121 frames aligned to 1 + 8k.""" + if abs(src_fps - 24) < 1.5: + out_fps = round(src_fps) + elif src_fps > 24: + step = src_fps / 24 + indices = [i for i in (round(i * step) for i in range(int(video.shape[0] / step))) if i < video.shape[0]] + video, out_fps = video[indices], 24 + else: + out_fps = round(src_fps) + video = video[:121] + if video.shape[0] == 0: + raise ValueError("no readable frames") + return video[:1 + 8 * ((video.shape[0] - 1) // 8)], out_fps + + +@pytest.mark.parametrize("src_fps", [10, 12, 23.976, 24, 25, 25.5, 29.97, 30, 48, 60, 120, 239.76]) +@pytest.mark.parametrize("total", [1, 2, 9, 17, 25, 100, 121, 122, 151, 241, 300, 1000]) +def test_streaming_selection_equals_decoding_everything(src_fps, total): + ids = torch.arange(total, dtype=torch.int32).view(-1, 1, 1, 1) + plan = sr_io.plan_frame_selection(src_fps) + kept = sr_io._select_frames(iter(list(ids)), plan) + try: + expected, expected_fps = _resample_then_clip(ids, src_fps) + except ValueError: + assert kept == [] + return + got = torch.stack(kept)[:121] + got = got[:1 + 8 * ((got.shape[0] - 1) // 8)] + assert got.flatten().tolist() == expected.flatten().tolist() + assert plan.out_fps == expected_fps + + +def test_decoding_stops_after_the_last_needed_frame(): + consumed = [] + + def frames(): + for i in range(10**7): + consumed.append(i) + yield torch.zeros(1, 1, 1) + + plan = sr_io.plan_frame_selection(60.0) + assert len(sr_io._select_frames(frames(), plan)) == 121 + assert len(consumed) == plan.lookahead == 303 # ceil(121 * 2.5) + + +def _decode_all(path: Path) -> tuple[torch.Tensor, float]: + with av.open(str(path)) as container: + stream = container.streams.video[0] + frames = [torch.from_numpy(f.to_ndarray(format="rgb24")).permute(2, 0, 1) for f in container.decode(stream)] + return torch.stack(frames), float(stream.average_rate) + + +@pytest.mark.parametrize("fps,num_frames", [(30, 50), (24, 20), (Fraction(24000, 1001), 30), (12, 18), (60, 70)]) +def test_read_video_follows_the_sr_input_contract(tmp_path, fps, num_frames): + path = tmp_path / "clip.mp4" + k6_sr_tiny.write_mp4(path, num_frames, fps) + video, out_fps = sr_io.read_video(path) + expected, expected_fps = _resample_then_clip(*_decode_all(path)) + assert video.dtype == torch.uint8 and torch.equal(video, expected) and out_fps == expected_fps + + +def test_missing_input_is_reported(tmp_path): + with pytest.raises(FileNotFoundError): + sr_io.read_video(tmp_path / "nope.mp4") + + +def test_audio_is_mono_trimmed_to_the_processed_frames_and_optional(tmp_path): + with_audio, without = tmp_path / "a.mp4", tmp_path / "b.mp4" + k6_sr_tiny.write_mp4(with_audio, 24, 24, audio_seconds=3.0) + k6_sr_tiny.write_mp4(without, 24, 24) + audio = sr_io.read_audio(with_audio, 17, 24) + assert audio.dtype == torch.float32 and audio.ndim == 1 + assert audio.shape[0] == round(17 / 24 * sr_io.AUDIO_SAMPLE_RATE) + assert 0.05 < audio.abs().max() <= 1.0 + assert sr_io.read_audio(without, 17, 24) is None + + +@pytest.mark.parametrize("spec, mode, expected", [ + (None, "fit", None), + ("hd", "fit", (694, 1280)), # isotropic fit, even sizes + ("hd", "exact", (720, 1280)), + ("640x360", "fit", (346, 640)), + ("2k", "fit", None), # already inside the bucket: never upscales +]) +def test_delivery_size(spec, mode, expected): + assert sr_io.resolve_target_hw(spec, (832, 1536), mode) == expected + + +@pytest.mark.parametrize("spec, mode", [("8k", "fit"), ("12x", "fit"), ("hd", "stretch")]) +def test_bad_delivery_specs_are_rejected(spec, mode): + with pytest.raises(ValueError): + sr_io.validate_target_spec(spec, mode) + + +def test_resize_never_upscales(): + video = torch.zeros(3, 2, 64, 96, dtype=torch.uint8) + assert sr_io.resize_video(video, (32, 48)).shape == (3, 2, 32, 48) + with pytest.raises(ValueError, match="exceeds"): + sr_io.resize_video(video, (128, 192)) diff --git a/fastvideo/tests/stages/kandinsky6_ti2va_tiny.py b/fastvideo/tests/stages/kandinsky6_ti2va_tiny.py new file mode 100644 index 000000000..af36e3e19 --- /dev/null +++ b/fastvideo/tests/stages/kandinsky6_ti2va_tiny.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tiny random Kandinsky6 TI2VA DiT (video+audio, text+pooled conditioning) for CPU tests. + +Small enough to run forward passes in milliseconds: 2 fused decoder blocks, 1 text block per tower, +axes_dims (8, 4, 4) -> head_dim 16, model_dim 32. +""" +from __future__ import annotations + +import copy +from typing import Any + +import torch + +TINY_DIT_CFG: dict[str, Any] = dict( + in_visual_dim=4, + out_visual_dim=4, + in_text_dim=8, + in_text_dim2=8, + time_dim=16, + patch_size=(1, 2, 2), + model_dim=32, + ff_dim=32, + num_text_blocks=1, + num_visual_blocks=2, + axes_dims=(8, 4, 4), + visual_cond=True, + is_multimodal=True, + in_audio_dim=6, + # Set explicitly (not left to the `x or time_dim` default-resolution in Kandinsky6ArchConfig. + # __post_init__): update_model_arch calls __post_init__ twice -- once via the arch_config field's + # default_factory (which already resolves *_a fields from the *default* time_dim/model_dim/ff_dim/ + # axes_dims, not this dict's), and again after applying this dict -- so an unset *_a field here + # would silently keep the stale default-resolved value instead of picking up this dict's values. + model_dim_a=32, + time_dim_a=16, + ff_dim_a=32, + axes_dims_a=(8, 4, 4), + visual_token_type_num_embeddings=2, + attention_engine="auto", +) + + +def dit_config_dict(dit_cfg: dict | None = None) -> dict[str, Any]: + return copy.deepcopy(dit_cfg or TINY_DIT_CFG) + + +def randomize(module: torch.nn.Module, seed: int, std: float = 0.05) -> torch.nn.Module: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in module.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) * std) + return module + + +def build_dit(seed: int = 0, dit_cfg: dict | None = None): + from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6VideoAudioConfig + from fastvideo.models.dits.kandinsky6 import Kandinsky6Transformer3DModel + + cfg = dit_config_dict(dit_cfg) + config = Kandinsky6VideoAudioConfig() + config.update_model_arch(cfg) + return randomize(Kandinsky6Transformer3DModel(config, cfg), seed + 2).eval() diff --git a/fastvideo/tests/stages/test_kandinsky6_arch_config.py b/fastvideo/tests/stages/test_kandinsky6_arch_config.py new file mode 100644 index 000000000..3db731280 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_arch_config.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for Kandinsky6ArchConfig's __post_init__ validation. + +Pure dataclass/logic tests -- no GPU, no model weights. +""" +from __future__ import annotations + +import pytest + +from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig + + +def test_defaults_match_diffusers_reference_not_kandinsky5(): + # Kandinsky5ArchConfig's defaults are much smaller (in_visual_dim=4, + # model_dim=2048, num_visual_blocks=32, axes_dims=(16,24,24)) -- this + # pins Kandinsky6ArchConfig to the actual Kandinsky6Transformer3DModel + # defaults so a future edit can't silently drift back toward Kandinsky5's. + cfg = Kandinsky6ArchConfig() + assert cfg.in_visual_dim == 16 + assert cfg.out_visual_dim == 16 + assert cfg.in_text_dim == 3584 + assert cfg.in_text_dim2 == 768 + assert cfg.time_dim == 1024 + assert cfg.model_dim == 4096 + assert cfg.ff_dim == 16384 + assert cfg.num_text_blocks == 4 + assert cfg.num_visual_blocks == 60 + assert cfg.axes_dims == (32, 48, 48) + assert cfg.in_audio_dim == 20 + # Both official checkpoints (Pro-sft and Pro-distill) declare this fixed value in + # transformer/config.json; it is not derived from height/width. + assert cfg.scale_factor == (1.0, 2.0, 2.0) + + +def test_audio_dims_default_to_matching_video_dims(): + cfg = Kandinsky6ArchConfig(model_dim=48, time_dim=32, ff_dim=128, axes_dims=(8, 8, 8)) + assert cfg.model_dim_a == 48 + assert cfg.time_dim_a == 32 + assert cfg.ff_dim_a == 128 + assert cfg.axes_dims_a == (8, 8, 8) + + +def test_audio_dims_can_be_overridden_independently(): + cfg = Kandinsky6ArchConfig( + model_dim=48, + time_dim=32, + ff_dim=128, + axes_dims=(8, 8, 8), + model_dim_a=48, + time_dim_a=32, # must match time_dim unless fix_modulation=True + ff_dim_a=96, + axes_dims_a=(4, 4, 4), + ) + assert cfg.model_dim_a == 48 + assert cfg.ff_dim_a == 96 + assert cfg.axes_dims_a == (4, 4, 4) + + +def test_model_dim_a_not_divisible_by_head_dim_a_raises(): + with pytest.raises(ValueError, match="model_dim_a"): + Kandinsky6ArchConfig( + model_dim=48, + time_dim=32, + axes_dims=(8, 8, 8), + model_dim_a=50, # 50 % sum((4,4,4))=12 != 0 + time_dim_a=32, + axes_dims_a=(4, 4, 4), + ) + + +def test_mismatched_time_dim_a_without_fix_modulation_raises(): + # Kandinsky6FusedTransformerDecoderBlock's cross-modal modulation is + # driven by the *other* modality's time embedding by default + # (fix_modulation=False) -- va_modulation is built from time_dim but + # invoked with the audio time embedding, so time_dim_a must equal + # time_dim or the fused block's modulation matmul shapes mismatch. + with pytest.raises(ValueError, match="time_dim_a"): + Kandinsky6ArchConfig(model_dim=48, time_dim=32, axes_dims=(8, 8, 8), time_dim_a=16) + + +def test_mismatched_time_dim_a_allowed_with_fix_modulation(): + cfg = Kandinsky6ArchConfig(model_dim=48, time_dim=32, axes_dims=(8, 8, 8), time_dim_a=16, + fix_modulation=True) + assert cfg.time_dim_a == 16 + + +def test_is_multimodal_false_is_rejected(): + # The DiT is always the joint video+audio model; a config asking for anything else fails loudly. + with pytest.raises(ValueError, match="is_multimodal=True"): + Kandinsky6ArchConfig(model_dim=48, time_dim=32, axes_dims=(8, 8, 8), is_multimodal=False) + + +def test_derived_arch_fields_match_dit_arch_config_contract(): + cfg = Kandinsky6ArchConfig(model_dim=48, time_dim=32, axes_dims=(8, 8, 8), in_visual_dim=6, out_visual_dim=6) + assert cfg.hidden_size == 48 + assert cfg.num_attention_heads == 48 // 24 + assert cfg.in_channels == 6 + assert cfg.out_channels == 6 + assert cfg.num_channels_latents == 6 diff --git a/fastvideo/tests/stages/test_kandinsky6_audio_decoding.py b/fastvideo/tests/stages/test_kandinsky6_audio_decoding.py new file mode 100644 index 000000000..53d5ee935 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_audio_decoding.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for Kandinsky6 audio decoding with separate audio_vae (mel-VAE decoder) and vocoder (BigVGAN-v2) +components: model_index.json declares both, and audio_vae/diffusion_pytorch_model.safetensors also carries a +bundled copy of the vocoder weights under a `vocoder.*` prefix, which the loader drops. + +Pure CPU: Kandinsky6AudioVAE's "44k" mode has fixed (not configurably tiny) dimensions, so those tests +build the real-sized module (random-initialized, no real weights) -- still fast, no GPU needed. The +BigVGAN-v2 side does accept a tiny config, so the vocoder-remap tests use one. +""" +from __future__ import annotations + +import types + +import pytest +import torch + +B, A_LEN = 1, 4 # audio latent sequence length + + +def _audio_vae(): + from fastvideo.models.audio.kandinsky6_audio_vae import Kandinsky6AudioVAE + + vae = Kandinsky6AudioVAE({"mode": "44k", "scaling_factor": 0.5302, "need_vae_encoder": False}) + vae.remove_weight_norm() + return vae.eval() + + +def test_audio_vae_has_no_vocoder_submodule_or_state_dict_keys(): + # The checkpoint's audio_vae/*.safetensors still bundles a redundant `vocoder.*` copy for + # backward compat (see AudioDecoderLoader), but this class itself no longer has anywhere to load + # it -- the vocoder is now a fully separate pipeline component. + vae = _audio_vae() + assert not hasattr(vae, "vocoder") + assert not hasattr(vae, "vocode") + assert not hasattr(vae, "wrapped_decode") + keys = set(vae.state_dict()) + assert not any(k.startswith("vocoder.") for k in keys) + assert any(k.startswith("vae.") for k in keys) + assert any(k.startswith("mel_converter.") for k in keys) + + +def test_audio_vae_decode_returns_a_mel_spectrogram(): + vae = _audio_vae() + latents = torch.randn(B, 40, A_LEN) # [B, embed_dim, A], matching the stage's channel-first convention + with torch.no_grad(): + mel = vae.decode(latents) + assert mel.shape[0] == B + assert mel.shape[1] == 128 # num_mels + assert torch.isfinite(mel).all() + + +def test_mmaudio_vocoder_config_remaps_to_bigvgan_with_the_hardcoded_fixed_fields(): + from fastvideo.models.loader.component_loader import _mmaudio_vocoder_to_bigvgan + + # A flat config shaped like the real checkpoint's vocoder/config.json (diffusers' MMAudioVocoder + # schema): no resblock/activation/snake_logscale/use_bias_at_final/use_tanh_at_final keys. + raw = dict(num_mels=128, upsample_initial_channel=1536, upsample_rates=[8, 4, 2, 2, 2, 2], + upsample_kernel_sizes=[16, 8, 4, 4, 4, 4], resblock_kernel_sizes=[3, 7, 11], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]]) + class_name, merged = _mmaudio_vocoder_to_bigvgan("MMAudioVocoder", raw) + assert class_name == "BigVGANV2" + assert merged["resblock"] == "1" + assert merged["activation"] == "snakebeta" + assert merged["snake_logscale"] is True + assert merged["use_bias_at_final"] is False + assert merged["use_tanh_at_final"] is False + # Without this, BigVGANV2's default weight_norm-parametrized construction leaves the fresh + # module's state_dict keys as conv_pre.parametrizations.weight.original0/1 etc., but the + # checkpoint's vocoder/*.safetensors is saved with weight_norm already removed (plain conv_pre.weight) -- + # VocoderLoader's load_state_dict(strict=True) then fails with "Missing key(s)". + assert merged["weight_norm_removed"] is True + # Original fields pass through untouched. + assert merged["num_mels"] == 128 + assert merged["upsample_rates"] == [8, 4, 2, 2, 2, 2] + + +def test_mmaudio_vocoder_remap_is_a_no_op_for_other_class_names(): + from fastvideo.models.loader.component_loader import _mmaudio_vocoder_to_bigvgan + + class_name, cfg = _mmaudio_vocoder_to_bigvgan("BigVGANV2", {"num_mels": 128}) + assert class_name == "BigVGANV2" + assert cfg == {"num_mels": 128} + + +def test_remapped_bigvgan_builds_with_plain_non_parametrized_keys_matching_the_real_checkpoint(): + # A freshly-built BigVGANV2 defaults to weight_norm-parametrized keys + # (conv_pre.parametrizations.weight.original0/1), but the checkpoint's vocoder/*.safetensors has plain + # keys (conv_pre.weight) -- VocoderLoader.load() calls load_state_dict(strict=True) BEFORE + # remove_weight_norm(), so without weight_norm_removed=True set at construction time (asserted in + # the test above), that strict load fails with "Missing key(s)". Reproduces VocoderLoader's exact + # load order: construct -> load_state_dict(strict=True) -> remove_weight_norm() (idempotent). + from fastvideo.models.loader.component_loader import _mmaudio_vocoder_to_bigvgan + from fastvideo.models.registry import ModelRegistry + + tiny = dict(num_mels=8, upsample_initial_channel=16, upsample_rates=[2, 2], upsample_kernel_sizes=[4, 4], + resblock_kernel_sizes=[3], resblock_dilation_sizes=[[1, 3]]) + class_name, cfg = _mmaudio_vocoder_to_bigvgan("MMAudioVocoder", tiny) + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vocoder = model_cls(cfg) + keys = set(vocoder.state_dict()) + assert not any("parametrizations" in key for key in keys), keys + assert "conv_pre.weight" in keys + + # Simulate loading a "real checkpoint" (here, the module's own freshly-initialized weights) the + # same way VocoderLoader does: strict load first, remove_weight_norm second. + own_state = {k: v.clone() for k, v in vocoder.state_dict().items()} + vocoder.load_state_dict(own_state, strict=True) + vocoder.remove_weight_norm() # idempotent no-op here; must not raise + + +def test_remapped_bigvgan_decodes_without_tanh_saturation_or_a_final_bias(): + # Exercises the two fields that would otherwise default wrong (use_bias_at_final / + # use_tanh_at_final default to True in BigVGANV2 itself). + from fastvideo.models.loader.component_loader import _mmaudio_vocoder_to_bigvgan + from fastvideo.models.registry import ModelRegistry + + tiny = dict(num_mels=8, upsample_initial_channel=16, upsample_rates=[2, 2], upsample_kernel_sizes=[4, 4], + resblock_kernel_sizes=[3], resblock_dilation_sizes=[[1, 3]]) + class_name, cfg = _mmaudio_vocoder_to_bigvgan("MMAudioVocoder", tiny) + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + vocoder = model_cls(cfg) + assert vocoder.conv_post.bias is None + assert vocoder.use_tanh_at_final is False + vocoder.remove_weight_norm() + + mel = torch.randn(1, 8, 5) + with torch.no_grad(): + waveform = vocoder(mel) + assert waveform.shape[0] == 1 and waveform.shape[1] == 1 + assert torch.isfinite(waveform).all() + + +class _StubAudioVAE(torch.nn.Module): + """Records its decode() call's input and returns a deterministic fake mel.""" + + def __init__(self) -> None: + super().__init__() + self.scaling_factor = 2.0 + self.linear = torch.nn.Linear(1, 1) # so next(self.parameters()) works + self.seen_latents: torch.Tensor | None = None + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + self.seen_latents = latents + batch, _, length = latents.shape + return torch.full((batch, 8, length), 0.5) + + +class _StubVocoder(torch.nn.Module): + """Records its forward() call's input (the VAE's mel output) and returns a fake waveform.""" + + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(1, 1) + self.seen_mel: torch.Tensor | None = None + + def forward(self, mel: torch.Tensor) -> torch.Tensor: + self.seen_mel = mel + batch, _, length = mel.shape + return torch.full((batch, 1, length * 4), 2.0) # out of [-1, 1] to exercise the stage's clamp + + +@pytest.fixture() +def cpu_device(monkeypatch): + import fastvideo.pipelines.stages.kandinsky6 as k6_stages + cpu = torch.device("cpu") + monkeypatch.setattr(k6_stages, "get_local_torch_device", lambda: cpu) + return cpu + + +def test_audio_decoding_stage_calls_vae_decode_then_vocoder_in_sequence(cpu_device): + from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6AudioDecodingStage + + audio_vae = _StubAudioVAE() + vocoder = _StubVocoder() + stage = Kandinsky6AudioDecodingStage(audio_vae=audio_vae, vocoder=vocoder) + + audio_latents = torch.ones(1, A_LEN, 40) * 4.0 # [B, A, D]; scaling_factor=2.0 -> latents=2.0 + batch = types.SimpleNamespace(audio_latents=audio_latents, extra={}) + fastvideo_args = types.SimpleNamespace(vae_cpu_offload=False, + pipeline_config=types.SimpleNamespace()) + + out = stage.forward(batch, fastvideo_args) + + # audio_vae.decode saw latents already divided by scaling_factor (no + mean_value term: the new + # diffusers reference's postprocess_audio dropped that term) and transposed to [B, D, A]. + assert audio_vae.seen_latents is not None + assert audio_vae.seen_latents.shape == (1, 40, A_LEN) + torch.testing.assert_close(audio_vae.seen_latents, torch.full((1, 40, A_LEN), 2.0)) + + # vocoder saw exactly the VAE's mel output (the two-step decode -> vocode chain). + assert vocoder.seen_mel is not None + torch.testing.assert_close(vocoder.seen_mel, torch.full((1, 8, A_LEN), 0.5)) + + # Final waveform is clamped to [-1, 1] even though the stub vocoder returned 2.0. + assert out.extra["audio"].max().item() == pytest.approx(1.0) + assert out.extra["audio"].shape == (A_LEN * 4,) + + +def test_audio_decoding_stage_is_a_noop_without_audio_latents(cpu_device): + from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6AudioDecodingStage + + audio_vae, vocoder = _StubAudioVAE(), _StubVocoder() + stage = Kandinsky6AudioDecodingStage(audio_vae=audio_vae, vocoder=vocoder) + batch = types.SimpleNamespace(audio_latents=None, extra={}) + out = stage.forward(batch, types.SimpleNamespace(vae_cpu_offload=False)) + assert out.extra == {} + assert audio_vae.seen_latents is None + assert vocoder.seen_mel is None diff --git a/fastvideo/tests/stages/test_kandinsky6_audio_latent_duration.py b/fastvideo/tests/stages/test_kandinsky6_audio_latent_duration.py new file mode 100644 index 000000000..c9de69d0a --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_audio_latent_duration.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for ``audio_latent_duration``, the formula that keeps +Kandinsky6's audio and video latents time-aligned. + +Byte-for-byte match to the diffusers reference's +``pipeline_kandinsky6_ti2va.audio_latent_duration``: +``ceil(((T_lat-1)*4+1) / fps * audio_sample_rate / audio_downsample_factor)``, +where ``(T_lat-1)*4+1`` is the causal video VAE's pixel-frame count for a +given latent-frame count (temporal_compression_ratio == 4, matching +HunyuanVAEConfig -- the video VAE Kandinsky6 reuses from Kandinsky5). + +Pure Python/math, no GPU or model weights needed. +""" +from __future__ import annotations + +import math + +from fastvideo.pipelines.stages.kandinsky6 import audio_latent_duration + + +def test_audio_latent_duration_matches_diffusers_reference_default_5s_clip(): + # Native T2VA default operating point: num_frames=121 -> T_lat=31 at + # temporal_compression_ratio=4 ((121-1)//4+1 == 31). + result = audio_latent_duration(31, fps=24.0, audio_sample_rate=44100, audio_downsample_factor=1024) + pixel_frames = (31 - 1) * 4 + 1 + expected = math.ceil(pixel_frames / 24.0 * 44100 / 1024) + assert result == expected + assert result == 218 + + +def test_audio_latent_duration_single_latent_frame(): + # T_lat=1 -> pixel_frames=1 (the causal VAE's "+1 first frame" edge case). + result = audio_latent_duration(1, fps=24.0, audio_sample_rate=44100, audio_downsample_factor=1024) + assert result == math.ceil(1 / 24.0 * 44100 / 1024) + + +def test_audio_latent_duration_scales_with_video_length(): + short = audio_latent_duration(16, fps=24.0, audio_sample_rate=44100, audio_downsample_factor=1024) + long = audio_latent_duration(31, fps=24.0, audio_sample_rate=44100, audio_downsample_factor=1024) + assert long > short + + +def test_audio_latent_duration_returns_int(): + result = audio_latent_duration(31, fps=24.0, audio_sample_rate=44100, audio_downsample_factor=1024) + assert isinstance(result, int) diff --git a/fastvideo/tests/stages/test_kandinsky6_cfg_resolution.py b/fastvideo/tests/stages/test_kandinsky6_cfg_resolution.py new file mode 100644 index 000000000..a646a48fd --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_cfg_resolution.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for ``Kandinsky6CFGResolutionStage``, which recomputes +``batch.do_classifier_free_guidance`` under Kandinsky6's own CFG contract, mirroring the diffusers +reference's (``pipeline_kandinsky6_ti2va.py``) ``do_classifier_free_guidance`` property: the uncond pass +runs whenever ``guidance_scale > 1.0``, same as the generic rule (``ForwardBatch.__post_init__``); any +guidance_scale <= 1.0 (including below 1.0) runs cond-only with no negative-prompt encoding. A PiFlow +(distilled) scheduler never runs CFG regardless of the requested guidance_scale. + +Pure logic test -- no GPU, no model weights. +""" +from __future__ import annotations + +import types + +import pytest + +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6CFGResolutionStage + + +def _stage(scheduler) -> Kandinsky6CFGResolutionStage: + stage = Kandinsky6CFGResolutionStage.__new__(Kandinsky6CFGResolutionStage) + stage.scheduler = scheduler + return stage + + +def _flow_scheduler() -> types.SimpleNamespace: + return types.SimpleNamespace(is_piflow=False) + + +def _piflow_scheduler() -> types.SimpleNamespace: + return types.SimpleNamespace(is_piflow=True) + + +@pytest.mark.parametrize("guidance_scale,expected", [ + (5.0, True), + (2.0, True), + (1.0 + 1e-5, True), + (1.0, False), + (1.0 - 1e-9, False), + (0.5, False), + (0.0, False), +]) +def test_cfg_trigger_uses_guidance_scale_greater_than_1(guidance_scale, expected): + stage = _stage(_flow_scheduler()) + batch = types.SimpleNamespace(guidance_scale=guidance_scale, do_classifier_free_guidance=None) + + out = stage.forward(batch, fastvideo_args=None) + + assert out.do_classifier_free_guidance is expected + + +@pytest.mark.parametrize("guidance_scale", [1.0, 5.0, 0.0, float("nan")]) +def test_piflow_scheduler_never_triggers_cfg(guidance_scale): + stage = _stage(_piflow_scheduler()) + batch = types.SimpleNamespace(guidance_scale=guidance_scale, do_classifier_free_guidance=None) + + out = stage.forward(batch, fastvideo_args=None) + + assert out.do_classifier_free_guidance is False + + +def test_nan_guidance_on_a_flow_scheduler_does_not_trigger_cfg(): + stage = _stage(_flow_scheduler()) + batch = types.SimpleNamespace(guidance_scale=float("nan"), do_classifier_free_guidance=None) + + out = stage.forward(batch, fastvideo_args=None) + + assert out.do_classifier_free_guidance is False diff --git a/fastvideo/tests/stages/test_kandinsky6_denoising_precision.py b/fastvideo/tests/stages/test_kandinsky6_denoising_precision.py new file mode 100644 index 000000000..bb352b77b --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_denoising_precision.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for three ``Kandinsky6DenoisingStage.forward`` numerics gaps against the diffusers +reference (``pipeline_kandinsky6_ti2va.py`` / ``transformer_kandinsky6.py``): + +* the scheduler timestep must reach the transformer as fp32, in both the conditional and the + unconditional call -- not rounded to ``dit_precision`` (bf16) first; +* the PiFlow (distilled) noisy latent state must stay fp32 across steps, separately from the bf16 + ``video`` buffer; +* the RoPE ``scale_factor`` must come from the transformer's static config, not a height/width + heuristic. + +Pure CPU harness: a fake transformer records the kwargs of every call instead of computing anything, +and a fake scheduler stands in for ``FlowMatchEulerDiscreteScheduler``/``PiflowScheduler`` so no real +weights or scheduler-config parsing is needed. +""" +from __future__ import annotations + +import types + +import pytest +import torch + +from fastvideo.pipelines.stages import kandinsky6 as k6_stages +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6DenoisingStage + + +class _RecordingTransformer: + """Stands in for Kandinsky6Transformer3DModel: records every call's kwargs and returns zeros of + the right shape instead of computing anything.""" + + in_visual_dim = 2 + in_audio_dim = 2 + + def __init__(self, scale_factor=(1.0, 2.0, 2.0)): + self.scale_factor = scale_factor + # get_sparse_params reads self.transformer.config.attention_engine; "sdpa" (dense) keeps it a + # no-op, matching every real checkpoint observed so far. + self.config = types.SimpleNamespace(attention_engine="sdpa") + self.calls: list[dict] = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + video = kwargs["hidden_states"] + audio = kwargs["hidden_states_audio"] + video_vel = torch.zeros_like(video[..., :self.in_visual_dim]) + audio_vel = torch.zeros_like(audio) + return types.SimpleNamespace(sample=(video_vel, audio_vel)) + + +class _FakeFlowScheduler: + is_piflow = False + order = 1 + + def __init__(self, num_steps: int): + self.sigmas = torch.linspace(1.0, 0.0, num_steps + 1) + + def step(self, vel, timestep, sample, return_dict=False): + return (sample - 0.01 * vel.to(sample.dtype), ) + + +class _FakePiflowScheduler: + is_piflow = True + order = 1 + + def __init__(self, num_steps: int): + self.sigmas = torch.linspace(1.0, 0.0, num_steps + 1) + self.step_sample_dtypes: list[torch.dtype] = [] + + def step(self, vel, timestep, sample, return_dict=False): + self.step_sample_dtypes.append(sample.dtype) + # PiflowScheduler.step returns fp32 regardless of the input sample's dtype (the + # reference's policy math runs in fp32); mirrored here to make the fake honest. + return ((sample.float() - 0.01 * vel.float()), ) + + +def _fastvideo_args(scale_factor=(1.0, 2.0, 2.0)) -> types.SimpleNamespace: + dit_arch = types.SimpleNamespace(in_visual_dim=2, patch_size=(1, 1, 1), scale_factor=scale_factor) + vae_arch = types.SimpleNamespace(spatial_compression_ratio=1) + pipeline_config = types.SimpleNamespace( + dit_precision="bf16", + dit_config=types.SimpleNamespace(arch_config=dit_arch), + vae_config=types.SimpleNamespace(arch_config=vae_arch), + ) + return types.SimpleNamespace(model_loaded={"transformer": True}, pipeline_config=pipeline_config, + disable_autocast=True) + + +def _batch(timesteps: torch.Tensor, *, guidance_scale: float = 1.0, do_cfg: bool = False, + height: int = 2, width: int = 2) -> types.SimpleNamespace: + video = torch.zeros(1, 1, 2, 2, 2, dtype=torch.bfloat16) # [B,T,H,W,C] + audio = torch.zeros(1, 3, 2, dtype=torch.bfloat16) + prompt_embeds = torch.zeros(1, 4, 8, dtype=torch.bfloat16) + pooled = torch.zeros(1, 8, dtype=torch.bfloat16) + mask = torch.ones(1, 4, dtype=torch.long) + return types.SimpleNamespace( + guidance_scale=guidance_scale, + timesteps=timesteps, + latents=video, + audio_latents=audio, + extra={}, + prompt_embeds=[prompt_embeds, pooled], + prompt_attention_mask=[mask], + do_classifier_free_guidance=do_cfg, + negative_prompt_embeds=[prompt_embeds, pooled] if do_cfg else [], + negative_attention_mask=[mask] if do_cfg else None, + height=height, + width=width, + num_inference_steps=len(timesteps), + ) + + +def _stage(transformer, scheduler) -> Kandinsky6DenoisingStage: + stage = Kandinsky6DenoisingStage.__new__(Kandinsky6DenoisingStage) + stage.transformer = transformer + stage.scheduler = scheduler + return stage + + +@pytest.fixture() +def cpu_device(monkeypatch): + cpu = torch.device("cpu") + monkeypatch.setattr(k6_stages, "get_local_torch_device", lambda: cpu) + return cpu + + +def test_timestep_reaches_the_transformer_as_fp32_not_bf16(cpu_device): + transformer = _RecordingTransformer() + scheduler = _FakeFlowScheduler(num_steps=2) + stage = _stage(transformer, scheduler) + # A value that is NOT exactly representable in bf16 (bf16 has an 8-bit mantissa): if the stage + # rounds it to dit_precision before the transformer sees it, this exact value is lost. + timesteps = torch.tensor([517.3, 233.7], dtype=torch.float32) + batch = _batch(timesteps) + + stage.forward(batch, _fastvideo_args()) + + assert len(transformer.calls) == 2 + for call in transformer.calls: + assert call["timestep"].dtype == torch.float32 + # The exact fp32 value survives -- not just the dtype -- confirming the timestep was never + # rounded through bf16 (517.3 and 233.7 are not exactly bf16-representable) before this point. + assert torch.equal(transformer.calls[0]["timestep"], torch.tensor([517.3])) + assert torch.equal(transformer.calls[1]["timestep"], torch.tensor([233.7])) + + +def test_timestep_is_fp32_in_both_the_cond_and_uncond_call(cpu_device): + transformer = _RecordingTransformer() + scheduler = _FakeFlowScheduler(num_steps=1) + stage = _stage(transformer, scheduler) + timesteps = torch.tensor([517.3], dtype=torch.float32) + batch = _batch(timesteps, guidance_scale=5.0, do_cfg=True) + + stage.forward(batch, _fastvideo_args()) + + # One cond + one uncond call for the single step. + assert len(transformer.calls) == 2 + for call in transformer.calls: + assert call["timestep"].dtype == torch.float32 + + +def test_piflow_latent_state_stays_fp32_across_steps(cpu_device): + transformer = _RecordingTransformer() + scheduler = _FakePiflowScheduler(num_steps=3) + stage = _stage(transformer, scheduler) + timesteps = torch.tensor([700.0, 400.0, 100.0]) + batch = _batch(timesteps, guidance_scale=1.0, do_cfg=False) + + stage.forward(batch, _fastvideo_args()) + + assert scheduler.step_sample_dtypes == [torch.float32] * 3 + + +def test_flow_euler_state_is_unaffected_by_the_piflow_fp32_change(cpu_device): + # flow-Euler keeps writing through the bf16 `video` buffer every step (matches the diffusers + # reference, which also downcasts its state to the parameter dtype there) -- only PiFlow gets the + # separate fp32 tracker. + transformer = _RecordingTransformer() + scheduler = _FakeFlowScheduler(num_steps=2) + stage = _stage(transformer, scheduler) + timesteps = torch.tensor([700.0, 100.0]) + batch = _batch(timesteps) + + out = stage.forward(batch, _fastvideo_args()) + + assert out.latents.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("height,width", [(480, 864), (720, 1280)]) +def test_scale_factor_is_read_from_the_transformer_config_not_a_resolution_heuristic(cpu_device, height, width): + # The old heuristic special-cased 480<=h,w<=854 -> (1,2,2), else (1,3.16,3.16). Both resolutions + # here take the "else" branch under that heuristic (864 and 1280 both exceed 854), but both + # official Kandinsky6 checkpoints declare scale_factor=(1,2,2) unconditionally in + # transformer/config.json, including for the 480x864 docstring example. + transformer = _RecordingTransformer(scale_factor=(1.0, 2.0, 2.0)) + scheduler = _FakeFlowScheduler(num_steps=1) + stage = _stage(transformer, scheduler) + timesteps = torch.tensor([500.0]) + batch = _batch(timesteps, height=height, width=width) + + stage.forward(batch, _fastvideo_args()) + + assert transformer.calls[0]["scale_factor"] == (1.0, 2.0, 2.0) + + +def test_scale_factor_falls_back_to_the_static_arch_config_when_the_transformer_lacks_the_attribute(cpu_device): + transformer = _RecordingTransformer() + del transformer.scale_factor # simulate a not-yet-loaded transformer / older checkpoint object + scheduler = _FakeFlowScheduler(num_steps=1) + stage = _stage(transformer, scheduler) + timesteps = torch.tensor([500.0]) + batch = _batch(timesteps, height=720, width=1280) + + stage.forward(batch, _fastvideo_args(scale_factor=(1.0, 2.0, 2.0))) + + assert transformer.calls[0]["scale_factor"] == (1.0, 2.0, 2.0) diff --git a/fastvideo/tests/stages/test_kandinsky6_empty_prompt.py b/fastvideo/tests/stages/test_kandinsky6_empty_prompt.py new file mode 100644 index 000000000..c28fe7669 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_empty_prompt.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression test for ``kandinsky6_qwen_preprocess_text`` against the diffusers reference +(``pipeline_kandinsky6_ti2va.py``): an empty prompt encodes an empty user turn. The old code substituted +``"."`` for an empty prompt, which tokenizes differently from the reference's empty user turn (the +template's system half is never empty, so the Qwen tokenizer never sees zero tokens either way -- the +"." guard was unnecessary here, unlike the generic `treat_empty_as_dot` escape hatch other encoders use +for tokenizers that genuinely produce zero tokens on ""). + +Pure string test -- no GPU, no model weights, no tokenizer. +""" +from __future__ import annotations + +from fastvideo.configs.pipelines.kandinsky6 import ( + KANDINSKY6_PROMPT_TEMPLATE, + kandinsky6_qwen_preprocess_text, +) + + +def test_empty_prompt_encodes_an_empty_user_turn_not_a_dot(): + result = kandinsky6_qwen_preprocess_text("") + + assert result == KANDINSKY6_PROMPT_TEMPLATE.format("") + assert "user\n<|im_end|>" in result + assert "user\n.<|im_end|>" not in result + + +def test_whitespace_only_prompt_is_also_left_as_is(): + result = kandinsky6_qwen_preprocess_text(" ") + + assert result == KANDINSKY6_PROMPT_TEMPLATE.format(" ") + + +def test_non_empty_prompt_is_unaffected(): + result = kandinsky6_qwen_preprocess_text("a chef chops vegetables") + + assert result == KANDINSKY6_PROMPT_TEMPLATE.format("a chef chops vegetables") diff --git a/fastvideo/tests/stages/test_kandinsky6_image_conditioning.py b/fastvideo/tests/stages/test_kandinsky6_image_conditioning.py new file mode 100644 index 000000000..aa73da57f --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_image_conditioning.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for two ``Kandinsky6ImageEncodingStage`` gaps against the diffusers reference's +``encode_i2va_first_frame``: + +* the conditioning image must be cover-resized (aspect-preserving) and centre-cropped to (height, + width), not stretched to fit -- ``_cover_resize_dims`` replicates ``scale=min(src_h/h, src_w/w)``; + ``resize=(int(src_h/scale), int(src_w/scale))``, ``resize -> [h,w] centre crop``; +* a conditioning image must broadcast across ``num_videos_per_prompt>1`` instead of crashing + ``torch.cat`` with a batch-size mismatch. + +Pure CPU harness: a fake VAE returns zeros of the right shape instead of really encoding, so no +weights are needed. +""" +from __future__ import annotations + +import types + +import numpy as np +import PIL.Image +import pytest +import torch + +from fastvideo.pipelines.stages import kandinsky6 as k6_stages +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6ImageEncodingStage + + +@pytest.fixture() +def cpu_device(monkeypatch): + cpu = torch.device("cpu") + monkeypatch.setattr(k6_stages, "get_local_torch_device", lambda: cpu) + return cpu + + +# -------------------------------------------------------------------------------------------------- +# P10: cover-resize + centre crop +# -------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "src_h,src_w,height,width,expected", + [ + # The audit's worked example: 1024x1024 source -> 512x768 target. scale=min(1024/512, + # 1024/768)=1.333(3); resize to (int(1024/1.3333), int(1024/1.3333)) = (768, 768). + (1024, 1024, 512, 768, (768, 768)), + # A portrait target from a landscape source: scale=min(480/512, 864/768)=0.9375; resize to + # (int(480/0.9375), int(864/0.9375)) = (512, 921). + (480, 864, 512, 768, (512, 921)), + # Source already exactly the target size: no-op. + (512, 768, 512, 768, (512, 768)), + ], +) +def test_cover_resize_dims_matches_the_diffusers_formula(src_h, src_w, height, width, expected): + assert Kandinsky6ImageEncodingStage._cover_resize_dims(src_h, src_w, height, width) == expected + + +def test_cover_resize_dims_does_not_underflow_the_target_on_float_error(): + # A 510x510 source at a 480x832 target gives int(510 / (510/832)) == 831 (one pixel short of + # 832) because of float error in the scale round-trip, which would mean a negative centre-crop + # offset. Both axes must be >= the target. + new_h, new_w = Kandinsky6ImageEncodingStage._cover_resize_dims(510, 510, 480, 832) + assert new_h >= 480 + assert new_w >= 832 + + +def test_preprocess_crops_the_top_and_bottom_instead_of_stretching_the_whole_image(): + # 1024x1024 source with a marker confined to the first few rows only. A cover-resize to (768,768) + # + centre crop to (512,768) keeps rows [128:640] of the resized image (the audit: "keeps a + # 1024x683 centre region") -- far from the marker, so it must be cropped away entirely. The old + # code (plain stretch to (512,768), no crop) would keep *some* of every source row, including the + # marker, in the output's top rows. + src = torch.zeros(1, 3, 1024, 1024) + src[:, :, :4, :] = 1.0 + src = src * 2 - 1 # move to [-1,1] with real negative values so _preprocess skips its [0,1] guess + + out = Kandinsky6ImageEncodingStage._preprocess(src, height=512, width=768) + + assert out.shape[-2:] == (512, 768) + assert float(out.max()) < 0.0 # the positive marker did not survive the crop + + +def test_preprocess_pil_path_also_crops_not_stretches(): + arr = np.zeros((1024, 1024, 3), dtype=np.uint8) + arr[:4, :, :] = 255 # marker confined to the first 4 rows, full width + src = PIL.Image.fromarray(arr) + + out = Kandinsky6ImageEncodingStage._preprocess(src, height=512, width=768) + + assert out.shape[-2:] == (512, 768) + assert float(out.max()) < 0.0 + + +def test_preprocess_same_aspect_ratio_needs_no_crop_and_matches_a_plain_resize(): + # A source already at the target aspect ratio: cover-resize lands exactly on (height, width), so + # there is nothing to crop and the result should equal a plain resize. + src = torch.rand(1, 3, 256, 384) * 2 - 1 # 2:3 aspect, same as height=512,width=768 + + out = Kandinsky6ImageEncodingStage._preprocess(src, height=512, width=768) + + expected = torch.nn.functional.interpolate(src, size=(512, 768), mode="bilinear", antialias=True) + assert out.shape == expected.shape + assert torch.allclose(out, expected, atol=1e-5) + + +# -------------------------------------------------------------------------------------------------- +# P11: image conditioning + num_videos_per_prompt > 1 +# -------------------------------------------------------------------------------------------------- + + +class _FakeVAEOutput: + + def __init__(self, latent: torch.Tensor) -> None: + self._latent = latent + + def sample(self, generator=None) -> torch.Tensor: + return self._latent + + +class _FakeVAE: + scaling_factor = 1.0 + use_tiling = False + + def to(self, device): + return self + + def encode(self, image: torch.Tensor) -> _FakeVAEOutput: + b, _c, t, h, w = image.shape + return _FakeVAEOutput(torch.zeros(b, 4, t, h, w)) + + +def _image_encoding_stage() -> Kandinsky6ImageEncodingStage: + stage = Kandinsky6ImageEncodingStage.__new__(Kandinsky6ImageEncodingStage) + stage.vae = _FakeVAE() + return stage + + +def _fastvideo_args() -> types.SimpleNamespace: + return types.SimpleNamespace( + model_loaded={"vae": True}, + pipeline_config=types.SimpleNamespace(vae_precision="fp32"), + disable_autocast=True, + vae_cpu_offload=False, + ) + + +def test_image_plus_multiple_videos_per_prompt_no_longer_crashes(cpu_device): + stage = _image_encoding_stage() + num_videos_per_prompt = 2 + latents = torch.zeros(num_videos_per_prompt, 3, 8, 8, 4) # [B,T,H,W,C], visual_cond off (C==4) + batch = types.SimpleNamespace( + pil_image=PIL.Image.new("RGB", (8, 8)), + height=8, + width=8, + generator=None, + latents=latents, + extra={}, + ) + + out = stage.forward(batch, _fastvideo_args()) + + assert out.latents.shape == (num_videos_per_prompt, 4, 8, 8, 4) + assert out.extra["kandinsky6_tail_cond_active"] is True + assert out.extra["kandinsky6_visual_token_type_ids"].shape == (num_videos_per_prompt, 4) + + +def test_image_conditioning_still_works_for_a_single_video(cpu_device): + stage = _image_encoding_stage() + latents = torch.zeros(1, 3, 8, 8, 4) + batch = types.SimpleNamespace( + pil_image=PIL.Image.new("RGB", (8, 8)), + height=8, + width=8, + generator=None, + latents=latents, + extra={}, + ) + + out = stage.forward(batch, _fastvideo_args()) + + assert out.latents.shape == (1, 4, 8, 8, 4) + + +def test_image_batch_that_is_neither_one_nor_the_video_batch_size_raises_a_clear_error(cpu_device, monkeypatch): + stage = _image_encoding_stage() + + class _MismatchedVAE(_FakeVAE): + + def encode(self, image): + # Pretend the VAE itself produced 3 reference frames -- neither 1 (broadcastable) nor + # matching the video batch size of 2. + b, _c, t, h, w = image.shape + return _FakeVAEOutput(torch.zeros(3, 4, t, h, w)) + + stage.vae = _MismatchedVAE() + latents = torch.zeros(2, 3, 8, 8, 4) + batch = types.SimpleNamespace( + pil_image=PIL.Image.new("RGB", (8, 8)), + height=8, + width=8, + generator=None, + latents=latents, + extra={}, + ) + + with pytest.raises(ValueError, match="not broadcastable"): + stage.forward(batch, _fastvideo_args()) diff --git a/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py b/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py new file mode 100644 index 000000000..49defaf61 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_latent_preparation.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for two ``Kandinsky6LatentPreparationStage.forward`` gaps against the diffusers +reference (``pipeline_kandinsky6_ti2va.py``): + +* ``num_frames`` must round to the nearest ``4k+1`` the same way the diffusers reference's + ``num_frames // temporal_ratio * temporal_ratio + 1`` does: floor for a non-exact multiple of + ``temporal_ratio``, but round *up* for an exact multiple (120 -> 121, not 117); +* the audio latent length must follow the request's ``fps`` (``sample_fps``), not always the + pipeline's static default. + +Pure CPU harness: a fake transformer/scheduler stand in for real ones (only ``in_visual_dim``/ +``in_audio_dim``/``visual_cond`` and an optional ``init_noise_sigma`` are read), so no weights are +needed. +""" +from __future__ import annotations + +import types + +import pytest +import torch + +from fastvideo.pipelines.stages import kandinsky6 as k6_stages +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6LatentPreparationStage + + +@pytest.fixture() +def cpu_device(monkeypatch): + cpu = torch.device("cpu") + monkeypatch.setattr(k6_stages, "get_local_torch_device", lambda: cpu) + return cpu + + +def _fastvideo_args(*, sample_fps: float = 24.0) -> types.SimpleNamespace: + pipeline_config = types.SimpleNamespace( + vae_config=types.SimpleNamespace( + arch_config=types.SimpleNamespace(temporal_compression_ratio=4, spatial_compression_ratio=8)), + dit_config=types.SimpleNamespace( + arch_config=types.SimpleNamespace(patch_size=(1, 2, 2), in_visual_dim=16, in_audio_dim=20)), + dit_precision="fp32", + sample_fps=sample_fps, + audio_sample_rate=44100, + audio_downsample_factor=1024, + ) + return types.SimpleNamespace(pipeline_config=pipeline_config) + + +def _stage() -> Kandinsky6LatentPreparationStage: + stage = Kandinsky6LatentPreparationStage.__new__(Kandinsky6LatentPreparationStage) + stage.scheduler = types.SimpleNamespace() # no init_noise_sigma attribute + stage.transformer = types.SimpleNamespace(in_visual_dim=16, in_audio_dim=20, visual_cond=False) + return stage + + +def _batch(num_frames: int, *, fps=None) -> types.SimpleNamespace: + return types.SimpleNamespace( + height=512, + width=768, + num_frames=num_frames, + prompt="a chef chops vegetables", + num_videos_per_prompt=1, + latents=None, + audio_latents=None, + generator=None, + fps=fps, + ) + + +# -------------------------------------------------------------------------------------------------- +# P5: num_frames rounds to the nearest 4k+1 (floor for a non-exact multiple, round up for an exact one) +# -------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("requested,expected", [ + (120, 121), # exact multiple of 4 -> rounds up, not down + (96, 97), # ditto + (121, 121), # already valid -> unchanged + (60, 61), # exact multiple -> rounds up + (124, 125), # exact multiple -> rounds up + (119, 117), # non-exact multiple -> floors, same as before + (122, 121), # non-exact multiple -> floors, same as before +]) +def test_num_frames_rounds_to_nearest_4k_plus_1(cpu_device, requested, expected): + stage = _stage() + batch = _batch(requested) + + out = stage.forward(batch, _fastvideo_args()) + + assert out.num_frames == expected + assert out.latents.shape[1] == (expected - 1) // 4 + 1 + + +# -------------------------------------------------------------------------------------------------- +# P17: audio latent length follows batch.fps +# -------------------------------------------------------------------------------------------------- + + +def test_audio_latent_length_follows_the_requested_fps_not_the_pipeline_default(cpu_device): + stage = _stage() + batch = _batch(121, fps=30) + + out = stage.forward(batch, _fastvideo_args(sample_fps=24.0)) + + # T_lat=31 ((121-1)//4+1); pixel_frames=(31-1)*4+1=121; ceil(121/30*44100/1024) == 174, matching + # the diffusers reference's sample_fps=30 case exactly -- a fixed 24fps default would give 218. + assert out.audio_latents.shape[1] == 174 + + +def test_audio_latent_length_falls_back_to_the_pipeline_default_fps_when_the_request_has_none(cpu_device): + stage = _stage() + batch = _batch(121, fps=None) + + out = stage.forward(batch, _fastvideo_args(sample_fps=24.0)) + + assert out.audio_latents.shape[1] == 218 + + +def test_audio_latent_length_uses_the_first_entry_of_a_per_prompt_fps_list(cpu_device): + stage = _stage() + batch = _batch(121, fps=[30, 30]) + + out = stage.forward(batch, _fastvideo_args(sample_fps=24.0)) + + assert out.audio_latents.shape[1] == 174 diff --git a/fastvideo/tests/stages/test_kandinsky6_max_sequence_length.py b/fastvideo/tests/stages/test_kandinsky6_max_sequence_length.py new file mode 100644 index 000000000..8e09d8a13 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_max_sequence_length.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for ``max_sequence_length`` semantics against the diffusers reference +(``pipeline_kandinsky6_ti2va.py``): a request's ``max_sequence_length`` maps to the Qwen encoder only +(``max_length = max_sequence_length + 129``, the prompt-template prefix crop) -- not, as the generic +``TextEncodingStage`` does by default, to every encoder uniformly. Applying it to the fixed-77-position +CLIP tokenizer too (the old behavior) makes ``padding="max_length"`` on a length bigger than CLIP's own +77 positions crash. + +Also proves the shared hook this needed in ``TextEncodingStage`` (``_resolve_max_length`` + +``encode_text``'s per-encoder ``max_length`` sequence) leaves every *other* model's behavior identical: +a plain int or ``None`` (everyone else's call shape) still applies uniformly, unchanged. + +Pure CPU harness: fake tokenizers/encoders record their kwargs instead of really tokenizing/encoding, +so no weights are needed. +""" +from __future__ import annotations + +import types + +import pytest +import torch + +from fastvideo.configs.pipelines.kandinsky6 import KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX +from fastvideo.pipelines.stages import text_encoding as text_encoding_module +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6TextEncodingStage +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage + + +@pytest.fixture(autouse=True) +def _cpu_device(monkeypatch): + # get_local_torch_device() falls back to "mps" without CUDA (see + # fastvideo/distributed/parallel_state.py); force plain CPU tensors so this test does not depend + # on MPS being usable in the sandbox that runs it. + monkeypatch.setattr(text_encoding_module, "get_local_torch_device", lambda: torch.device("cpu")) + + +class _TensorDict(dict): + + def to(self, device): + return self + + +class _FakeTokenizer: + + def __init__(self) -> None: + self.calls: list[dict] = [] + + def __call__(self, texts, **kwargs): + self.calls.append(dict(kwargs)) + seq_len = int(kwargs.get("max_length", 8)) + n = len(texts) + return _TensorDict(input_ids=torch.zeros(n, seq_len, dtype=torch.long), + attention_mask=torch.ones(n, seq_len, dtype=torch.long)) + + +class _FakeTextEncoder(torch.nn.Module): + + def __init__(self, hidden_size: int = 4) -> None: + super().__init__() + self.hidden_size = hidden_size + + def forward(self, input_ids, attention_mask, output_hidden_states=False): + b, s = input_ids.shape + return types.SimpleNamespace(last_hidden_state=torch.zeros(b, s, self.hidden_size), hidden_states=None, + pooler_output=torch.zeros(b, self.hidden_size)) + + +def _encoder_config(default_max_length: int) -> types.SimpleNamespace: + return types.SimpleNamespace( + tokenizer_kwargs={"max_length": default_max_length}, + is_chat_model=False, + hidden_size=4, + arch_config=types.SimpleNamespace(output_hidden_states=False), + ) + + +def _fastvideo_args(qwen_default: int, clip_default: int) -> types.SimpleNamespace: + return types.SimpleNamespace( + pipeline_config=types.SimpleNamespace( + text_encoder_configs=[_encoder_config(qwen_default), _encoder_config(clip_default)], + preprocess_text_funcs=[lambda s: s, lambda s: s], + postprocess_text_funcs=[lambda o: o.last_hidden_state, lambda o: o.pooler_output], + dit_config=types.SimpleNamespace(prefix="kandinsky6"), + text_encoder_max_lengths=(qwen_default, clip_default), + ), + text_encoder_cpu_offload=False, + ) + + +def _make_stage(cls) -> tuple[TextEncodingStage, list[_FakeTokenizer], list[_FakeTextEncoder]]: + tokenizers = [_FakeTokenizer(), _FakeTokenizer()] + encoders = [_FakeTextEncoder(), _FakeTextEncoder()] + stage = cls(text_encoders=encoders, tokenizers=tokenizers) + return stage, tokenizers, encoders + + +# -------------------------------------------------------------------------------------------------- +# Kandinsky6TextEncodingStage._resolve_max_length +# -------------------------------------------------------------------------------------------------- + + +def test_resolve_max_length_is_none_when_the_request_does_not_set_one(): + stage, _, _ = _make_stage(Kandinsky6TextEncodingStage) + batch = types.SimpleNamespace(max_sequence_length=None) + + assert stage._resolve_max_length(batch, fastvideo_args=None) is None + + +def test_resolve_max_length_maps_to_qwen_only_leaving_clip_at_its_own_default(): + stage, _, _ = _make_stage(Kandinsky6TextEncodingStage) + batch = types.SimpleNamespace(max_sequence_length=1024) + + resolved = stage._resolve_max_length(batch, fastvideo_args=None) + + assert resolved == [KANDINSKY6_PROMPT_TEMPLATE_ENCODE_START_IDX + 1024, None] + assert resolved[0] == 1153 + + +def test_default_text_encoding_stage_resolve_max_length_is_unchanged_passthrough(): + # Every other model family uses the base TextEncodingStage: the hook's default implementation + # must be exactly the old inlined behavior (`batch.max_sequence_length`, applied uniformly by + # encode_text below) -- proof that adding the hook did not change anyone else's semantics. + stage, _, _ = _make_stage(TextEncodingStage) + for value in (None, 77, 1024): + batch = types.SimpleNamespace(max_sequence_length=value) + assert stage._resolve_max_length(batch, fastvideo_args=None) == value + + +# -------------------------------------------------------------------------------------------------- +# encode_text: per-encoder override vs. the old uniform-int/None behavior +# -------------------------------------------------------------------------------------------------- + + +def test_qwen_gets_the_overridden_length_and_clip_keeps_its_own_fixed_77(): + stage, tokenizers, _ = _make_stage(TextEncodingStage) + fastvideo_args = _fastvideo_args(qwen_default=641, clip_default=77) + + stage.encode_text("a cat", fastvideo_args, encoder_index=[0, 1], max_length=[1153, None], + device=torch.device("cpu")) + + assert tokenizers[0].calls[-1]["max_length"] == 1153 + # GAP: the old code applied the same override to every encoder, which would have set CLIP's + # max_length to 1153 too -- padding="max_length" with only 77 positions crashes for real. + assert tokenizers[1].calls[-1]["max_length"] == 77 + + +def test_no_override_falls_back_to_each_encoders_configured_default(): + stage, tokenizers, _ = _make_stage(TextEncodingStage) + fastvideo_args = _fastvideo_args(qwen_default=641, clip_default=77) + + stage.encode_text("a cat", fastvideo_args, encoder_index=[0, 1], max_length=None, device=torch.device("cpu")) + + assert tokenizers[0].calls[-1]["max_length"] == 641 + assert tokenizers[1].calls[-1]["max_length"] == 77 + + +def test_a_plain_int_max_length_still_applies_uniformly_to_every_encoder(): + # Regression: every non-Kandinsky6 model passes a single int (or None) as `max_length`, never a + # sequence -- that call shape must keep behaving exactly as before. + stage, tokenizers, _ = _make_stage(TextEncodingStage) + fastvideo_args = _fastvideo_args(qwen_default=641, clip_default=77) + + stage.encode_text("a cat", fastvideo_args, encoder_index=[0, 1], max_length=999, device=torch.device("cpu")) + + assert tokenizers[0].calls[-1]["max_length"] == 999 + assert tokenizers[1].calls[-1]["max_length"] == 999 + + +def test_kandinsky6_text_encoding_stage_forward_end_to_end_respects_the_split(): + stage, tokenizers, _ = _make_stage(Kandinsky6TextEncodingStage) + fastvideo_args = _fastvideo_args(qwen_default=641, clip_default=77) + batch = types.SimpleNamespace( + prompt="a cat", + max_sequence_length=1024, + prompt_embeds=[], + prompt_attention_mask=[], + do_classifier_free_guidance=False, + negative_prompt_embeds=None, + negative_attention_mask=None, + extra={}, + ) + + out = stage.forward(batch, fastvideo_args) + + assert len(out.prompt_embeds) == 2 + assert tokenizers[0].calls[-1]["max_length"] == 1153 + assert tokenizers[1].calls[-1]["max_length"] == 77 diff --git a/fastvideo/tests/stages/test_kandinsky6_parameter_dtype.py b/fastvideo/tests/stages/test_kandinsky6_parameter_dtype.py new file mode 100644 index 000000000..7735fd008 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_parameter_dtype.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for ``Kandinsky6Transformer3DModel._get_parameter_dtype`` (A13): the distill +checkpoint's safetensors file stores its modulation and time-embedding tensors as F32, and diffusers' +``from_pretrained(torch_dtype=bf16)`` keeps only an *exact* dotted-path-component match against +``visual_modulation``/``text_modulation``/``modulation`` fp32 -- not a substring match, so the +differently-named ``va_modulation``/``av_modulation`` cross-modal gates and the ``*_time_embeddings`` +towers are rounded to bf16 despite also being fp32 in the raw checkpoint file. + +``_get_parameter_dtype`` is a pure function of its ``name``/``default_dtype`` arguments (it never reads +``self``), so the model is constructed with ``__new__`` and never initialized -- no weights, no +``config`` object, no GPU needed. ``TransformerLoader``/``maybe_load_fsdp_model`` calling this hook is +exercised elsewhere (it is the same generic mechanism ``Kandinsky6SRTransformer3DModel`` already uses); +this file only pins the K6-specific fp32/bf16 split. +""" +from __future__ import annotations + +import torch + +from fastvideo.models.dits.kandinsky6 import Kandinsky6Transformer3DModel +from fastvideo.models.dits.kandinsky6_sr import Kandinsky6SRTransformer3DModel + + +def _dit() -> Kandinsky6Transformer3DModel: + return Kandinsky6Transformer3DModel.__new__(Kandinsky6Transformer3DModel) + + +def test_exact_modulation_components_stay_fp32(): + model = _dit() + fp32_names = [ + "text_transformer_blocks.0.text_modulation.out_layer.weight", + "video_text_transformer_blocks.2.text_modulation.out_layer.bias", + "audio_text_transformer_blocks.1.text_modulation.out_layer.weight", + "visual_transformer_blocks.0.visual_modulation.out_layer.weight", # non-multimodal decoder block + "visual_transformer_blocks.3.videoT.visual_modulation.out_layer.weight", + "visual_transformer_blocks.3.audioT.visual_modulation.out_layer.bias", + "out_layer.modulation.out_layer.weight", + "audio_out_layer.modulation.out_layer.bias", + ] + for name in fp32_names: + assert model._get_parameter_dtype(name, torch.bfloat16) == torch.float32, name + + +def test_differently_named_cross_modal_gates_and_time_towers_round_to_bf16(): + # These are fp32 in the raw checkpoint file too, but diffusers' EXACT component match does not + # keep them fp32 -- "va_modulation" as a whole path component is not "modulation". + model = _dit() + bf16_names = [ + "visual_transformer_blocks.5.va_modulation.out_layer.weight", + "visual_transformer_blocks.5.av_modulation.out_layer.bias", + "video_time_embeddings.in_layer.weight", + "audio_time_embeddings.out_layer.bias", + "time_embeddings.in_layer.weight", + "visual_embeddings.in_layer.weight", + "text_embeddings.in_layer.weight", + "video_text_embeddings.norm.weight", + ] + for name in bf16_names: + assert model._get_parameter_dtype(name, torch.bfloat16) == torch.bfloat16, name + + +def test_default_dtype_is_returned_unchanged_for_non_fp32_parameters(): + model = _dit() + assert model._get_parameter_dtype("visual_embeddings.in_layer.weight", torch.float16) == torch.float16 + + +def test_k6_and_sr_dit_hooks_are_independent_and_do_not_leak_into_each_other(): + # The K6 (T2VA) DiT and the SR DiT are separate classes with separate hooks and separate marker + # sets; adding _get_parameter_dtype to Kandinsky6Transformer3DModel must not change + # Kandinsky6SRTransformer3DModel's behavior (or vice versa) -- e.g. the SR DiT keeps + # "time_embeddings." fp32 (single-tower, no video/audio-prefixed names), the K6 DiT does not. + assert Kandinsky6Transformer3DModel._get_parameter_dtype is not Kandinsky6SRTransformer3DModel._get_parameter_dtype + + k6 = _dit() + sr = Kandinsky6SRTransformer3DModel.__new__(Kandinsky6SRTransformer3DModel) + + assert sr._get_parameter_dtype("time_embeddings.in_layer.weight", torch.bfloat16) == torch.float32 + assert k6._get_parameter_dtype("time_embeddings.in_layer.weight", torch.bfloat16) == torch.bfloat16 + + +def test_other_dit_classes_have_no_parameter_dtype_hook(): + # Sanity/scope check: the loader's mixed-dtype mechanism (component_loader.py / + # fsdp_load.py's `getattr(model, "_get_parameter_dtype", None)`) is opt-in per DiT class; a + # plain DiT untouched by this change must not suddenly grow the hook. + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + + assert getattr(WanTransformer3DModel, "_get_parameter_dtype", None) is None diff --git a/fastvideo/tests/stages/test_kandinsky6_piflow_guidance.py b/fastvideo/tests/stages/test_kandinsky6_piflow_guidance.py new file mode 100644 index 000000000..821b6128a --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_piflow_guidance.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``Kandinsky6DenoisingStage`` guidance guard for pi-Flow (distilled) checkpoints. + +A ``PiflowScheduler`` bundle runs without classifier-free guidance, so the stage refuses any ``guidance_scale`` other +than 1.0; the error must tell the user what to change. The stage is built without ``__init__`` and every call stops at +the first guard or at the missing-timesteps check, so no model is needed. +""" +from __future__ import annotations + +import types + +import pytest + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from fastvideo.models.schedulers.scheduling_piflow import PiflowScheduler +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6DenoisingStage + + +def _stage(scheduler) -> Kandinsky6DenoisingStage: + stage = Kandinsky6DenoisingStage.__new__(Kandinsky6DenoisingStage) + stage.scheduler = scheduler + return stage + + +def _batch(guidance_scale: float) -> types.SimpleNamespace: + return types.SimpleNamespace(guidance_scale=guidance_scale, timesteps=None) + + +@pytest.mark.parametrize("guidance_scale", [5.0, 2.0, 0.0, float("nan"), float("inf")]) +def test_piflow_rejects_any_guidance_but_one_with_an_actionable_message(guidance_scale): + stage = _stage(PiflowScheduler(nfe=16, n_grid=8, shift=5.0)) + + with pytest.raises(ValueError) as excinfo: + stage.forward(_batch(guidance_scale), None) + + message = str(excinfo.value) + assert "guidance_scale=1.0" in message + assert "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" in message + assert "kandinsky6_ti2va_distilled" in message + # num_inference_steps is not enforced -- neither Diffusers nor FastVideo + # requires it to equal the scheduler's nfe (PiflowScheduler.set_timesteps + # accepts any step count and ignores nfe); the message must not claim + # otherwise, only that 16 is the value the checkpoint was distilled for. + assert "num_inference_steps is not constrained" in message + assert "distilled for" in message + + +@pytest.mark.parametrize("guidance_scale", [1.0, 1.0 + 1e-7]) +def test_piflow_accepts_guidance_one(guidance_scale): + stage = _stage(PiflowScheduler(nfe=16, n_grid=8, shift=5.0)) + + with pytest.raises(ValueError, match="timesteps must be prepared"): + stage.forward(_batch(guidance_scale), None) + + +def test_flow_matching_scheduler_keeps_classifier_free_guidance(): + stage = _stage(FlowMatchEulerDiscreteScheduler(shift=5.0)) + + with pytest.raises(ValueError, match="timesteps must be prepared"): + stage.forward(_batch(5.0), None) diff --git a/fastvideo/tests/stages/test_kandinsky6_registry.py b/fastvideo/tests/stages/test_kandinsky6_registry.py new file mode 100644 index 000000000..4dfd6c809 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_registry.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Registry routing and presets of the Kandinsky6 TI2VA checkpoints: base (Pro-sft) and pi-Flow distilled (Pro-distill). + +Both official Diffusers repos declare the same model_index ``_class_name`` (``Kandinsky6TI2VAPipeline``), so they are +told apart by repo id or directory name, laid out like the LTX-2 distilled/base pair: the distilled entry is registered +first and the base detector excludes distilled names. Hub ids resolve by exact match (no network); local checkpoints are +model_index.json-only directories. The VSR ids are covered in ``kandinsky6_sr/test_kandinsky6_sr_registry_routing.py``. +""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from fastvideo import registry +from fastvideo.api.compat import normalize_generation_request, request_to_sampling_param +from fastvideo.api.presets import get_preset, validate_preset_selection +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig +from fastvideo.fastvideo_args import WorkloadType +from fastvideo.pipelines.basic.kandinsky6.presets import ( + ALL_PRESETS, + KANDINSKY6_TI2VA_DISTILLED, + KANDINSKY6_TI2VA, +) + +BASE_ID = "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" +OLD_BASE_ID = "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers" +DISTILLED_ID = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" +BASE_PRESET = "kandinsky6_ti2va" +DISTILLED_PRESET = "kandinsky6_ti2va_distilled" + + +def _model_dir(root: Path, name: str) -> str: + path = root / name + path.mkdir(parents=True) + model_index = {"_class_name": "Kandinsky6TI2VAPipeline", "_diffusers_version": "0.41.0.dev0"} + (path / "model_index.json").write_text(json.dumps(model_index)) + return str(path) + + +def _matching_presets(text: str) -> set[str]: + """Default presets of the Kandinsky6 registry entries whose detector claims ``text``.""" + entries = (registry._CONFIG_REGISTRY[model_id] for model_id, detector in registry._MODEL_NAME_DETECTORS + if detector(text)) + return {entry.default_preset for entry in entries if entry.model_family in ("kandinsky6", "kandinsky6_sr")} + + +@pytest.mark.parametrize("repo_id,preset", [(BASE_ID, BASE_PRESET), (OLD_BASE_ID, BASE_PRESET), + (DISTILLED_ID, DISTILLED_PRESET)]) +def test_official_repo_ids_resolve_to_their_presets(repo_id, preset): + assert repo_id in registry.get_registered_model_paths() + assert registry.get_pipeline_config_cls_from_name(repo_id) is Kandinsky6TI2VAConfig + assert registry.get_default_preset(repo_id) == preset + assert registry.get_preset_selection(repo_id) == (preset, "kandinsky6") + assert registry.get_model_family(repo_id) == "kandinsky6" + info = registry._get_config_info(repo_id) + assert info.pipeline_cls_name == "Kandinsky6TI2VAPipeline" + assert info.workload_types == (WorkloadType.T2V, WorkloadType.I2V) + + +@pytest.mark.parametrize("repo_id", [BASE_ID, DISTILLED_ID]) +def test_get_model_info_resolves_both_repo_ids_to_the_ti2va_pipeline(monkeypatch, repo_id): + model_index = {"_class_name": "Kandinsky6TI2VAPipeline", "_diffusers_version": "0.41.0.dev0"} + monkeypatch.setattr(registry, "maybe_download_model_index", lambda *_args, **_kwargs: model_index) + info = registry.get_model_info(repo_id) + assert info.pipeline_cls.__name__ == "Kandinsky6TI2VAPipeline" + assert info.pipeline_config_cls is Kandinsky6TI2VAConfig + + +def test_registered_model_listing_offers_both_repo_ids_for_text_and_image_workloads(): + listed = {model["id"]: model["workload_types"] for model in registry.get_registered_models_with_workloads()} + assert listed[BASE_ID] == listed[DISTILLED_ID] == ["t2v", "i2v"] + + +@pytest.mark.parametrize("name,preset", [ + ("Kandinsky-6.0-Pro-5s-Diffusers", BASE_PRESET), + ("Kandinsky-6.0-Pro-sft-5s-Diffusers", BASE_PRESET), + ("Kandinsky-6.0-Pro-distill-5s-Diffusers", DISTILLED_PRESET), + ("kandinsky6-distilled-export", DISTILLED_PRESET), + ("kandinsky6-export", BASE_PRESET), + ("my_export", BASE_PRESET), +]) +def test_local_directories_route_by_name(tmp_path, name, preset): + path = _model_dir(tmp_path, name) + assert registry.get_pipeline_config_cls_from_name(path) is Kandinsky6TI2VAConfig + assert registry.get_preset_selection(path) == (preset, "kandinsky6") + + +def test_a_trailing_slash_does_not_hide_the_distilled_name(tmp_path): + assert registry.get_default_preset(_model_dir(tmp_path, "kandinsky6-distilled-export") + "/") == DISTILLED_PRESET + + +def test_distill_in_a_parent_directory_does_not_make_a_base_checkpoint_distilled(tmp_path): + path = _model_dir(tmp_path / "distill_experiments", "kandinsky6-export") + assert registry.get_default_preset(path) == BASE_PRESET + + +@pytest.mark.parametrize("text,expected", [ + ("kandinsky-6.0-pro-distill-5s-diffusers", {DISTILLED_PRESET}), + ("/models/kandinsky6-distilled-export", {DISTILLED_PRESET}), + ("/models/kandinsky6-distilled-export/", {DISTILLED_PRESET}), + ("kandinsky-6.0-pro-sft-5s-diffusers", {BASE_PRESET}), + ("kandinsky6ti2vapipeline", {BASE_PRESET}), + ("/distill_experiments/kandinsky6-export", {BASE_PRESET}), +]) +def test_exactly_one_kandinsky6_detector_claims_a_name(text, expected): + # routing must not depend on registration order alone: the base detector excludes distilled names + assert _matching_presets(text) == expected + + +def test_ti2va_presets_are_registered(): + assert ALL_PRESETS == (KANDINSKY6_TI2VA, KANDINSKY6_TI2VA_DISTILLED) + for preset in ALL_PRESETS: + assert get_preset(preset.name, "kandinsky6") is preset + + +def test_base_preset_keeps_the_diffusers_defaults(): + defaults = get_preset(BASE_PRESET, "kandinsky6").defaults + assert (defaults["num_inference_steps"], defaults["guidance_scale"]) == (50, 5.0) + assert (defaults["height"], defaults["width"], defaults["num_frames"], defaults["fps"]) == (512, 768, 121, 24) + + +def _without_steps_and_guidance(defaults: dict) -> dict: + return {key: value for key, value in defaults.items() if key not in ("num_inference_steps", "guidance_scale")} + + +def test_distilled_preset_differs_from_the_base_preset_only_in_steps_and_guidance(): + base = get_preset(BASE_PRESET, "kandinsky6") + distilled = get_preset(DISTILLED_PRESET, "kandinsky6") + assert (distilled.defaults["num_inference_steps"], distilled.defaults["guidance_scale"]) == (10, 1.0) + assert _without_steps_and_guidance(distilled.defaults) == _without_steps_and_guidance(base.defaults) + assert (distilled.model_family, distilled.workload_type) == (base.model_family, base.workload_type) + assert distilled.stage_schemas == base.stage_schemas + + +@pytest.mark.parametrize("preset", [BASE_PRESET, DISTILLED_PRESET]) +def test_denoise_stage_overrides_accept_steps_and_guidance(preset): + overrides = {"denoise": {"num_inference_steps": 8, "guidance_scale": 1.0}} + validate_preset_selection(preset, "kandinsky6", stage_overrides=overrides) + + +@pytest.mark.parametrize("repo_id,steps,guidance", [(BASE_ID, 50, 5.0), (DISTILLED_ID, 10, 1.0)]) +def test_sampling_param_defaults_follow_the_repo_id(repo_id, steps, guidance): + sampling_param = SamplingParam.from_pretrained(repo_id) + assert (sampling_param.num_inference_steps, sampling_param.guidance_scale) == (steps, guidance) + assert (sampling_param.height, sampling_param.width) == (512, 768) + assert (sampling_param.num_frames, sampling_param.fps) == (121, 24) + + +def _request_sampling_param(model_path: str, **sampling) -> SamplingParam: + request = normalize_generation_request({"prompt": "a chef chops vegetables", "sampling": sampling}) + return request_to_sampling_param(request, model_path=model_path) + + +def test_explicit_request_values_win_over_the_distilled_preset(): + sampling_param = _request_sampling_param(DISTILLED_ID, num_inference_steps=8) + assert (sampling_param.num_inference_steps, sampling_param.guidance_scale) == (8, 1.0) + + +def test_explicit_request_values_win_over_the_base_preset(): + sampling_param = _request_sampling_param(BASE_ID, num_inference_steps=30, guidance_scale=3.0) + assert (sampling_param.num_inference_steps, sampling_param.guidance_scale) == (30, 3.0) + sampling_param = _request_sampling_param(BASE_ID, guidance_scale=3.0) + assert (sampling_param.num_inference_steps, sampling_param.guidance_scale) == (50, 3.0) diff --git a/fastvideo/tests/stages/test_kandinsky6_target_dtype.py b/fastvideo/tests/stages/test_kandinsky6_target_dtype.py new file mode 100644 index 000000000..5c7385637 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_target_dtype.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for ``Kandinsky6DenoisingStage._resolve_target_dtype``. + +Same logic and same rationale as +``fastvideo/tests/stages/test_kandinsky5_target_dtype.py`` (Kandinsky6DenoisingStage +copies Kandinsky5DenoisingStage's dtype-resolution helper verbatim): a plain +(non-FSDP) transformer must honor ``pipeline_config.dit_precision`` exactly, +while an FSDP2-wrapped transformer must read the active +``MixedPrecisionPolicy`` (always recorded by ``maybe_load_fsdp_model`` before +sharding) rather than the parameter storage dtype, since FSDP2 computes in +the policy's ``param_dtype`` regardless of what dtype parameters are stored +in. + +Pure logic tests -- the FSDP "wrap" only swaps ``__class__`` the same way +``fully_shard`` does, without any process group, GPU, or model load. +""" +from __future__ import annotations + +import types + +import pytest +import torch + +from fastvideo import utils as fastvideo_utils +from fastvideo.pipelines.stages.kandinsky6 import Kandinsky6DenoisingStage +from fastvideo.utils import set_mixed_precision_policy + + +def _make_stage(transformer: torch.nn.Module) -> Kandinsky6DenoisingStage: + """Bypass __init__: only set the field _resolve_target_dtype reads.""" + stage = Kandinsky6DenoisingStage.__new__(Kandinsky6DenoisingStage) + stage.transformer = transformer + return stage + + +def _fastvideo_args(dit_precision: str) -> types.SimpleNamespace: + return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(dit_precision=dit_precision)) + + +def _fsdp_wrap(module: torch.nn.Module) -> torch.nn.Module: + from torch.distributed.fsdp import FSDPModule + + orig_cls = module.__class__ + module.__class__ = type(f"FSDP{orig_cls.__name__}", (FSDPModule, orig_cls), {}) + return module + + +@pytest.fixture() +def mixed_precision_state_reset(): + state_holder = fastvideo_utils._mixed_precision_state + had_state = hasattr(state_holder, "state") + prev_state = getattr(state_holder, "state", None) + if had_state: + del state_holder.state + yield + if had_state: + state_holder.state = prev_state + elif hasattr(state_holder, "state"): + del state_holder.state + + +def test_resolve_target_dtype_honors_explicit_fp32_for_plain_transformer(): + transformer = torch.nn.Linear(4, 4).to(torch.float32) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("fp32")) + + assert resolved == torch.float32 + + +def test_resolve_target_dtype_honors_explicit_bf16_for_plain_transformer(): + transformer = torch.nn.Linear(4, 4).to(torch.bfloat16) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("bf16")) + + assert resolved == torch.bfloat16 + + +def test_resolve_target_dtype_fsdp_reads_policy_not_parameter_storage(mixed_precision_state_reset): + set_mixed_precision_policy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) + transformer = _fsdp_wrap(torch.nn.Linear(4, 4).to(torch.float16)) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("fp16")) + + assert resolved == torch.bfloat16 + + +def test_resolve_target_dtype_fsdp_defaults_to_bf16_without_policy_state(mixed_precision_state_reset): + transformer = _fsdp_wrap(torch.nn.Linear(4, 4).to(torch.float16)) + stage = _make_stage(transformer) + + resolved = stage._resolve_target_dtype(_fastvideo_args("fp16")) + + assert resolved == torch.bfloat16 diff --git a/fastvideo/tests/stages/test_kandinsky6_ti2va_transformer.py b/fastvideo/tests/stages/test_kandinsky6_ti2va_transformer.py new file mode 100644 index 000000000..dd04e2542 --- /dev/null +++ b/fastvideo/tests/stages/test_kandinsky6_ti2va_transformer.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Kandinsky6 TI2VA DiT: batched forward semantics on tiny random weights (CPU). + +The model always builds the fused video+audio architecture, and a video-only forward call +(hidden_states_audio=None) is a supported case. +""" +from __future__ import annotations + +import kandinsky6_ti2va_tiny as k6_tiny +import pytest +import torch + +B, T, H, W = 2, 2, 4, 4 +VISUAL_DIM = 2 * k6_tiny.TINY_DIT_CFG["in_visual_dim"] + 1 # visual_cond=True: noised | cond | mask +A = 3 # audio sequence length +L = 5 # text token length + + +def _rope_pos(): + return [torch.arange(T), torch.arange(H // 2), torch.arange(W // 2)] + + +def _inputs(seed=1, batch=B, with_audio=True): + generator = torch.Generator().manual_seed(seed) + hidden_states = torch.randn(batch, T, H, W, VISUAL_DIM, generator=generator) + encoder_hidden_states = torch.randn(batch, L, k6_tiny.TINY_DIT_CFG["in_text_dim"], generator=generator) + pooled = torch.randn(batch, k6_tiny.TINY_DIT_CFG["in_text_dim2"], generator=generator) + timestep = torch.linspace(100.0, 900.0, batch) + kwargs = dict( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + pooled_projections=pooled, + visual_rope_pos=_rope_pos(), + text_rope_pos=torch.arange(L), + ) + if with_audio: + kwargs["hidden_states_audio"] = torch.randn(batch, A, k6_tiny.TINY_DIT_CFG["in_audio_dim"], generator=generator) + return kwargs + + +def test_multimodal_forward_shapes(): + dit = k6_tiny.build_dit() + with torch.no_grad(): + video_out, audio_out = dit(**_inputs()) + assert video_out.shape == (B, T, H, W, k6_tiny.TINY_DIT_CFG["out_visual_dim"]) + assert audio_out.shape == (B, A, k6_tiny.TINY_DIT_CFG["in_audio_dim"]) # out_audio_dim defaults to in_audio_dim + assert torch.isfinite(video_out).all() + assert torch.isfinite(audio_out).all() + + +def test_video_only_forward_does_not_raise_and_returns_a_single_tensor(): + # Pre-collapse this raised NotImplementedError for a multimodal checkpoint; the fused block now + # guards aud=None exactly like the diffusers reference. + dit = k6_tiny.build_dit() + kwargs = _inputs(with_audio=False) + with torch.no_grad(): + result = dit(**kwargs) + assert isinstance(result, torch.Tensor) + assert result.shape == (B, T, H, W, k6_tiny.TINY_DIT_CFG["out_visual_dim"]) + assert torch.isfinite(result).all() + + +def test_video_only_output_matches_the_video_half_of_a_multimodal_call_given_the_same_video_inputs(): + # The video stream's own self-attention / text-cross-attention / feed-forward do not depend on + # audio being present; only the cross-modal va/av mixing (skipped when aud=None) does. So with the + # same seed, video-only output should differ from the multimodal call's video_out (the cross-modal + # terms are real contributions, not zero), but both must be finite and distinctly computed -- + # this test pins that video-only is NOT silently identical to just truncating a multimodal run. + dit = k6_tiny.build_dit() + kwargs = _inputs() + with torch.no_grad(): + video_out_multimodal, _ = dit(**kwargs) + video_only_kwargs = {k: v for k, v in kwargs.items() if k != "hidden_states_audio"} + video_out_alone = dit(**video_only_kwargs) + assert not torch.allclose(video_out_multimodal, video_out_alone) + + +def test_samples_of_a_batch_do_not_interact(): + dit = k6_tiny.build_dit() + kwargs = _inputs() + with torch.no_grad(): + video_batched, audio_batched = dit(**kwargs) + singles = [ + dit( + hidden_states=kwargs["hidden_states"][i:i + 1], + encoder_hidden_states=kwargs["encoder_hidden_states"][i:i + 1], + timestep=kwargs["timestep"][i:i + 1], + pooled_projections=kwargs["pooled_projections"][i:i + 1], + hidden_states_audio=kwargs["hidden_states_audio"][i:i + 1], + visual_rope_pos=_rope_pos(), + text_rope_pos=torch.arange(L), + ) for i in range(B) + ] + video_single = torch.cat([s[0] for s in singles]) + audio_single = torch.cat([s[1] for s in singles]) + torch.testing.assert_close(video_batched, video_single, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(audio_batched, audio_single, rtol=1e-5, atol=1e-5) + + +def test_state_dict_has_prefixed_video_audio_towers_and_no_bare_text_tower(): + # Checkpoints use the video_/audio_-prefixed text towers; there is no unprefixed text tower. + dit = k6_tiny.build_dit() + keys = set(dit.state_dict()) + assert "video_text_embeddings.in_layer.weight" in keys + assert "audio_text_embeddings.in_layer.weight" in keys + assert "video_time_embeddings.in_layer.weight" in keys + assert "audio_time_embeddings.in_layer.weight" in keys + assert "visual_transformer_blocks.0.videoT.self_attention.to_query.weight" in keys + assert "visual_transformer_blocks.0.audioT.self_attention.to_query.weight" in keys + assert "audio_out_layer.out_layer.weight" in keys + assert not any(key.startswith("text_embeddings.") or key.startswith("text_transformer_blocks.") + for key in keys) + + +def test_gradient_checkpointing_matches_no_checkpointing_for_a_multimodal_call(): + dit = k6_tiny.build_dit() + dit.train() + kwargs = _inputs() + dit.gradient_checkpointing = False + video_a, audio_a = dit(**kwargs) + dit.gradient_checkpointing = True + video_b, audio_b = dit(**kwargs) + torch.testing.assert_close(video_a, video_b, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(audio_a, audio_b, rtol=1e-5, atol=1e-5) + + +def test_gradient_checkpointing_matches_no_checkpointing_for_a_video_only_call(): + # aud=None flows through torch.utils.checkpoint.checkpoint as a plain positional arg; this pins + # that None-valued args survive the checkpoint round-trip correctly. + dit = k6_tiny.build_dit() + dit.train() + kwargs = _inputs(with_audio=False) + dit.gradient_checkpointing = False + video_a = dit(**kwargs) + dit.gradient_checkpointing = True + video_b = dit(**kwargs) + torch.testing.assert_close(video_a, video_b, rtol=1e-5, atol=1e-5) + + +def test_visual_transformer_blocks_is_the_first_registered_module_list(): + # fastvideo/hooks/layerwise_offload.py's enable_layerwise_offload hooks only the first top-level + # nn.ModuleList found via named_children() (attribute-registration order) and stops there -- with + # dit_layerwise_offload=True (the default), that must be visual_transformer_blocks, not the 4-entry + # video_text_transformer_blocks. Not itself a test of the (CUDA-only) hook helper; pins the + # registration order the helper relies on. + dit = k6_tiny.build_dit() + first_module_list_name = next(name for name, child in dit.named_children() if isinstance(child, torch.nn.ModuleList)) + assert first_module_list_name == "visual_transformer_blocks" diff --git a/fastvideo/utils.py b/fastvideo/utils.py index 10d142b76..1c7a68a5d 100644 --- a/fastvideo/utils.py +++ b/fastvideo/utils.py @@ -664,8 +664,10 @@ def verify_model_config_and_directory( # Diffusers convention: component entries start with [library, class]. # Modular manifests may append loading metadata, which FastVideo does not # need because published component subfolders match their manifest keys. - # Non-list entries are scalar metadata - # (e.g. boundary_ratio); a None first element marks a disabled + # Non-list entries, and list entries whose first element isn't a string, + # are scalar/tuple pipeline metadata (e.g. boundary_ratio, or a + # tuple-valued field registered via register_to_config, which serializes + # to a JSON list just like a component entry); a None first element marks a disabled # component (matches composed_pipeline_base.py). Pipelines that # lazy-load shared components from upstream HF repos simply omit the # key, so we only enforce "declared, active, but missing on disk". @@ -677,7 +679,7 @@ def verify_model_config_and_directory( for key, value in config.items(): if key.startswith("_") or key == "transformer" or key.startswith("tokenizer"): continue - if not isinstance(value, list) or len(value) < 1 or value[0] is None: + if (not isinstance(value, list) or len(value) < 1 or value[0] is None or not isinstance(value[0], str)): continue subdir = os.path.join(model_path, key) if not os.path.exists(subdir): diff --git a/mkdocs.yml b/mkdocs.yml index 81c084016..a2d92045b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -165,6 +165,7 @@ nav: - LTX: cookbook/ltx.md - Hunyuan: cookbook/hunyuan.md - Kandinsky 5: cookbook/kandinsky5.md + - Kandinsky 6: cookbook/kandinsky6.md - LongCat: cookbook/longcat.md - TurboDiffusion: cookbook/turbodiffusion.md - Image: @@ -185,6 +186,8 @@ nav: - Optimizations: inference/optimizations.md - ComfyUI: inference/comfyui.md - Support Matrix: inference/support_matrix.md + - Kandinsky 6 T2IVA: inference/kandinsky6.md + - Kandinsky 6 Video SR: inference/kandinsky6_sr.md - CLI: inference/cli.md - Add Pipeline: inference/add_pipeline.md - Examples: diff --git a/tests/local_tests/README.md b/tests/local_tests/README.md index f3bd31e1c..0a819a263 100644 --- a/tests/local_tests/README.md +++ b/tests/local_tests/README.md @@ -27,6 +27,8 @@ tests/local_tests// | Hunyuan GameCraft | T2V / I2V | [`gamecraft/`](./gamecraft/) | | GEN3C | T2V | [`gen3c/`](./gen3c/) | | Kandinsky-5 | T2V | [`kandinsky5/`](./kandinsky5/) | +| Kandinsky-6 | TI2VA (T2V/I2V + audio) | [`kandinsky6/`](./kandinsky6/) | +| Kandinsky-6 SR | V2V (upscale) | [`kandinsky6_sr/README.md`](./kandinsky6_sr/README.md) | | LTX-2 | T2V (+ audio) | [`ltx2/`](./ltx2/) | | Stable Diffusion 3.5 | T2I | [`sd35/`](./sd35/) | | Stable Audio Open 1.0 | T2A | [`stable_audio/`](./stable_audio/) | diff --git a/tests/local_tests/kandinsky6/README.md b/tests/local_tests/kandinsky6/README.md new file mode 100644 index 000000000..e21f64f25 --- /dev/null +++ b/tests/local_tests/kandinsky6/README.md @@ -0,0 +1,79 @@ +# Kandinsky-6 Local Tests + +Local-only parity tests for the `kandinsky6` FastVideo port (T2IVA and T2IVA distilled: text, optionally plus an +image, to video with audio). Compares the FastVideo Kandinsky6 transformer against the Diffusers reference +(`diffusers.Kandinsky6Transformer3DModel`); skipped in CI. + +## Reference Assets + +| Field | Value | +|---|---| +| Model family | `kandinsky6` | +| Workload types | `T2V`, `I2V` (one merged TI2VA pipeline serves both) | +| Official reference | `diffusers.Kandinsky6Transformer3DModel` (used by `diffusers.Kandinsky6TI2VAPipeline`) | +| Local reference dir | set `KANDINSKY6_DIFFUSERS_REPO_PATH` to a `diffusers` checkout that provides Kandinsky6 (its `src/` is prepended to `sys.path`); only needed when the installed `diffusers` does not | +| Official commit/version | not pinned; the official repos are exported with `_diffusers_version` `0.41.0.dev0` | +| HF weights | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` (T2IVA), `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` (T2IVA distilled) | +| HF revision | `main` | +| Local weights dir | `official_weights/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` (env: `KANDINSKY6_DIFFUSERS_PATH`, `KANDINSKY6_TRANSFORMER_PATH`) | +| Source layout | `diffusers` | +| Needs conversion | no: the official repos are in Diffusers layout and load directly | + +> Use only the env-var **name** for tokens (e.g., `HF_TOKEN`). Never paste a token value. + +## Why this differs from `tests/local_tests/kandinsky5/` + +Kandinsky5's parity test imports its diffusers reference class directly from the pip-installed `diffusers` package. +`test_kandinsky6_transformer_parity.py` tries the same import first and, when the installed `diffusers` does not provide +`Kandinsky6Transformer3DModel`, falls back to a `diffusers` checkout at `KANDINSKY6_DIFFUSERS_REPO_PATH`. The test +skips when the transformer weights are missing, when CUDA is unavailable, or when the reference cannot be imported; a +skip is not a verified pass (see `.agents/skills/add-model/shared/common_rules.md`). + +The reference's `forward()` also uses a packed/ragged `(sum_T, H, W, C)` layout (`cu_seqlens`-addressed) rather than +FastVideo's ordinary batched `(B, T, H, W, C)` tensor. The test converts between them (batch size 1, so it is a plain +squeeze/unsqueeze). + +## Shared Environment Setup + +Run from the FastVideo repo root in the same env used for FastVideo. + +```bash +export KANDINSKY6_DIFFUSERS_REPO_PATH=/path/to/diffusers-checkout-with-kandinsky6 +export KANDINSKY6_DIFFUSERS_PATH=official_weights/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers +``` + +Do not change core dependency versions (`torch`, `diffusers`, `transformers`, +`flash-attn`, `triton`, CUDA packages) without explicit approval. + +## Weight Setup + +```bash +python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \ + "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" \ + "official_weights/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" +``` + +The base repo is about 90 GB (the transformer alone is 70 GB); set `KANDINSKY6_TRANSFORMER_PATH` to a different +`transformer/` directory to test another checkpoint. + +## Tests in this directory + +```bash +pytest tests/local_tests/kandinsky6/ -v +``` + +| Component | Test | Concerns | Status | +|---|---|---|---| +| `transformer` (Kandinsky6Transformer3DModel) | [`test_kandinsky6_transformer_parity.py`](./test_kandinsky6_transformer_parity.py) | Packed-vs-batched shape conversion; needs a CUDA GPU that holds the 60 GB transformer | `scaffold_skip` | + +Audio codec (`MMAudioVAE` + `BigVGANV2`) has no dedicated parity test here: +Kandinsky6 reuses those classes unmodified from FastVideo's existing MMAudio +port (verified via byte-identical `DATA_MEAN_128D`/`DATA_STD_128D` constants +against the diffusers reference), so their existing parity coverage applies. + +## Review Notes + +- Required before handoff: non-skip PASS for each component parity test, + including reused components that own weights or numerical behavior. +- Pipeline parity may start as a scaffold; final handoff requires non-skip + PASS or an explicit blocker accepted via the escape-hatch process. diff --git a/tests/local_tests/kandinsky6/test_kandinsky6_transformer_parity.py b/tests/local_tests/kandinsky6/test_kandinsky6_transformer_parity.py new file mode 100644 index 000000000..9288782a8 --- /dev/null +++ b/tests/local_tests/kandinsky6/test_kandinsky6_transformer_parity.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numeric parity test: FastVideo's Kandinsky6Transformer3DModel vs. the +diffusers reference (see tests/local_tests/kandinsky6/README.md). + +The reference classes must be importable (from the installed `diffusers` if it +provides Kandinsky6, else from a `diffusers` checkout at +KANDINSKY6_DIFFUSERS_REPO_PATH), and the transformer weights of an official +Diffusers repo (default: kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers) must +exist on disk. Absent either, this skips -- a skip is not a verified pass (see +.agents/skills/add-model/shared/common_rules.md). + +The reference's forward() uses a packed/ragged (sum_T, H, W, C) layout keyed +by cu_seqlens (batch size folded into the leading dim); FastVideo's port uses +an ordinary batched (B, T, H, W, C) tensor. This test's batch size is 1, so +the two are related by a plain squeeze/unsqueeze. +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29516") +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "native") + + +_DEFAULT_WEIGHTS_DIR = "official_weights/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" + + +def _resolve_transformer_path() -> Path: + root = Path(os.getenv("KANDINSKY6_DIFFUSERS_PATH", _DEFAULT_WEIGHTS_DIR)) + return Path(os.getenv("KANDINSKY6_TRANSFORMER_PATH", str(root / "transformer"))) + + +def _import_diffusers_kandinsky6(): + """Try the installed `diffusers` first, then a diffusers checkout. + + KANDINSKY6_DIFFUSERS_REPO_PATH should point at a diffusers checkout that + provides Kandinsky6 -- its src/ directory is prepended to sys.path so + `import diffusers` resolves there instead of the installed copy. + """ + try: + from diffusers import Kandinsky6Transformer3DModel + return Kandinsky6Transformer3DModel + except ImportError: + pass + + repo_path = os.getenv("KANDINSKY6_DIFFUSERS_REPO_PATH") + if not repo_path: + return None + src_path = str(Path(repo_path) / "src") + if src_path not in sys.path: + sys.path.insert(0, src_path) + # Drop any partially-imported `diffusers` from the earlier failed + # attempt so the retry actually re-resolves against the new sys.path. + sys.modules.pop("diffusers", None) + try: + from diffusers import Kandinsky6Transformer3DModel + return Kandinsky6Transformer3DModel + except ImportError: + return None + + +def test_kandinsky6_transformer_parity(): + transformer_path = _resolve_transformer_path() + if not transformer_path.exists(): + pytest.skip(f"Kandinsky6 transformer weights not found at {transformer_path}") + + if not torch.cuda.is_available(): + pytest.skip("Kandinsky6 transformer parity test requires CUDA for practical runtime.") + + DiffusersKandinsky6 = _import_diffusers_kandinsky6() + if DiffusersKandinsky6 is None: + pytest.skip("diffusers.Kandinsky6Transformer3DModel is not importable: the installed diffusers lacks it and " + "KANDINSKY6_DIFFUSERS_REPO_PATH is unset or doesn't resolve. See README.md.") + + try: + from fastvideo.configs.models.dits import Kandinsky6VideoAudioConfig + from fastvideo.configs.pipelines import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import TransformerLoader + except Exception as exc: + pytest.skip(f"FastVideo imports unavailable for parity run: {exc}") + + torch.manual_seed(42) + device = torch.device("cuda:0") + precision = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + precision_str = "bf16" if precision == torch.bfloat16 else "fp16" + + reference_model = DiffusersKandinsky6.from_pretrained(transformer_path).to(device=device, dtype=precision) + reference_model.eval() + + config = Kandinsky6VideoAudioConfig() + args = FastVideoArgs( + model_path=str(transformer_path), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + pipeline_config=PipelineConfig(dit_config=config, dit_precision=precision_str), + ) + args.device = device + # The official transformer/config.json has no _class_name; the pipeline loader takes it from model_index.json + # (composed_pipeline_base.load_modules), so hand the loader the same hint when loading the component directly. + args._model_index_class_names = {"transformer": "Kandinsky6Transformer3DModel"} + fastvideo_model = TransformerLoader().load(str(transformer_path), args).to(device=device, dtype=precision) + fastvideo_model.eval() + + arch = reference_model.config + in_visual_dim = arch.in_visual_dim + visual_cond = bool(getattr(arch, "visual_cond", False)) + in_text_dim = arch.in_text_dim + in_text_dim2 = arch.in_text_dim2 + in_audio_dim = arch.in_audio_dim + patch_size = arch.patch_size + + batch_size = 1 + grid_t, grid_h, grid_w = 2, 4, 4 + latent_t = grid_t * patch_size[0] + latent_h = grid_h * patch_size[1] + latent_w = grid_w * patch_size[2] + audio_len = 12 + text_len = 8 + + base_video = torch.randn(batch_size, latent_t, latent_h, latent_w, in_visual_dim, device=device, dtype=precision) + if visual_cond: + cond = torch.zeros_like(base_video) + mask = torch.zeros(batch_size, latent_t, latent_h, latent_w, 1, device=device, dtype=precision) + video = torch.cat([base_video, cond, mask], dim=-1) + else: + video = base_video + audio = torch.randn(batch_size, audio_len, in_audio_dim, device=device, dtype=precision) + encoder_hidden_states = torch.randn(batch_size, text_len, in_text_dim, device=device, dtype=precision) + pooled_projections = torch.randn(batch_size, in_text_dim2, device=device, dtype=precision) + timestep = torch.tensor([500.0], device=device, dtype=precision) + + visual_rope_pos = [ + torch.arange(grid_t, device=device), + torch.arange(grid_h, device=device), + torch.arange(grid_w, device=device), + ] + text_rope_pos = torch.arange(text_len, device=device) + + sdpa_math_ctx = torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH) + with torch.no_grad(), sdpa_math_ctx: + fv_video_out, fv_audio_out = fastvideo_model( + hidden_states=video, + hidden_states_audio=audio, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + timestep=timestep, + visual_rope_pos=visual_rope_pos, + text_rope_pos=text_rope_pos, + scale_factor=(1.0, 1.0, 1.0), + sparse_params=None, + return_dict=False, + ) + + # Reference forward: packed (sum_T, H, W, C) / (sum_A, D) layout with + # batch_size folded into the leading dim -- batch_size=1 here, so + # this is a pure squeeze of the batch axis, no reshuffling needed. + ref_video_in = video.squeeze(0) + ref_audio_in = audio.squeeze(0) + ref_out = reference_model( + x_video=ref_video_in, + x_audio=ref_audio_in, + text_embed=encoder_hidden_states.squeeze(0), + pooled_text_embed=pooled_projections.squeeze(0), + time=timestep, + visual_rope=reference_model.visual_rope_embeddings((grid_t, grid_h, grid_w), visual_rope_pos, + (1.0, 1.0, 1.0)), + audio_rope=reference_model.audio_rope_embeddings(torch.arange(audio_len, device=device)), + text_rope=reference_model.video_text_rope_embeddings(text_rope_pos), + sparse_params=None, + return_dict=False, + ) + ref_video_out, ref_audio_out = ref_out + + assert fv_video_out.shape == ref_video_out.unsqueeze(0).shape + assert fv_audio_out.shape == ref_audio_out.unsqueeze(0).shape + tol = 1e-4 if precision == torch.bfloat16 else 2e-4 + assert_close(fv_video_out.squeeze(0), ref_video_out, atol=tol, rtol=tol) + assert_close(fv_audio_out.squeeze(0), ref_audio_out, atol=tol, rtol=tol) diff --git a/tests/local_tests/kandinsky6_sr/README.md b/tests/local_tests/kandinsky6_sr/README.md new file mode 100644 index 000000000..8eccfaca7 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/README.md @@ -0,0 +1,41 @@ +# Kandinsky-6 Video Super-Resolution (SR) Local Tests + +Local-only parity checks for the Kandinsky6 SR port. CI runs the CPU unit tests in `fastvideo/tests/stages/kandinsky6_sr/` +(tiny random models, both bundle kinds through the full pipeline); the tests here need a GPU, the released weights, a +checkout of the previous port or the official reference, and skip without them. + +| Field | Value | +|---|---| +| Official reference | `k6_video/src/kandinsky_sr/` (`KANDINSKY_SR_SRC`, default `../k6_video/src`), HEAD `473ea253603c98e23130da8f67c7f00504f33d4e` | +| HF weights | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`, `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | +| Previous port | PR head `12be8fdd` (own sampler loops, vendored upscaler / KVAE packages) | + +## Tests + +| Test | Compares | Needs | +|---|---|---| +| `../pipelines/test_kandinsky6_sr_pipeline_parity.py` | full pipeline vs the previous port, both bundles, x2 / x2.25 / x4 (PSNR >= 48 dB) | CUDA, `KANDINSKY6_SR_MODELS`, `KANDINSKY6_SR_VIDEO`, `KANDINSKY6_SR_PREVIOUS_ROOT` | +| `test_kandinsky6_sr_latent_upscaler_parity.py` | latent upscaler vs the previous port, bitwise, x2 / x4, fp32 / bf16 | CUDA, `KANDINSKY6_SR_BUNDLE`, `KANDINSKY6_SR_PREVIOUS_ROOT` | +| `test_kandinsky6_sr_vae_parity.py` | KVAE encode / decode vs the previous port, bitwise | CUDA, `KANDINSKY6_SR_VAE_DIR`, `KANDINSKY6_SR_VAE_REFERENCE` | +| `test_kandinsky6_sr_dit_parity.py` | SR DiT vs the reference DiT, tiny random weights | `KANDINSKY_SR_SRC` | +| `test_kandinsky6_sr_real_weights.py` | real-weights smoke; optional PSNR vs a `kandy-sr` reference output | CUDA, `KANDINSKY6_SR_BUNDLE`, `KANDINSKY6_SR_TEST_VIDEO` | + +The previous-port comparisons run each implementation in its own subprocess, because both are the package `fastvideo`. + +## Commands + +```bash +# CPU unit tests (the CI set) +pytest fastvideo/tests/stages/kandinsky6_sr -q + +# pipeline parity on a GPU +KANDINSKY6_SR_MODELS=/path/to/dir/with/both/bundles KANDINSKY6_SR_VIDEO=/path/to/clip_17f.mp4 \ +KANDINSKY6_SR_PREVIOUS_ROOT=/path/to/previous/checkout \ +pytest tests/local_tests/pipelines/test_kandinsky6_sr_pipeline_parity.py -v -s + +# component parity on a GPU +KANDINSKY6_SR_BUNDLE=/path/to/Kandinsky-6.0-VSR-5s-Diffusers KANDINSKY6_SR_PREVIOUS_ROOT=/path/to/previous/checkout \ +pytest tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py -v -s +KANDINSKY6_SR_VAE_DIR=/path/to/Kandinsky-6.0-VSR-5s-Diffusers/vae KANDINSKY6_SR_VAE_REFERENCE=/path/to/previous/checkout \ +pytest tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py -v -s +``` diff --git a/tests/local_tests/kandinsky6_sr/__init__.py b/tests/local_tests/kandinsky6_sr/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/kandinsky6_sr/lu_parity_common.py b/tests/local_tests/kandinsky6_sr/lu_parity_common.py new file mode 100644 index 000000000..e7e05c574 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/lu_parity_common.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared driver of the latent-upscaler parity test: runs one implementation on the real weights. + +Used in two processes, because the new and the previous implementation both live in a package named ``fastvideo``: + + python lu_parity_common.py --impl previous --lu-dir /latent_upscaler --out ref.pt # PYTHONPATH= + +and imported by ``test_kandinsky6_sr_latent_upscaler_parity.py`` for the current implementation. Both sides build +the bank on the meta device, load the checkpoint with ``strict=True, assign=True`` and run the same inputs, so any +difference comes from the module code. +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +from collections.abc import Callable + +import torch + +LATENT_SHAPE = (1, 64, 5, 26, 48) # one 5-frame KVAE latent tile of a ~416x768 clip +SEED = 0 +SCALES = (2, 4) +DTYPES = ("fp32", "bf16") + + +def load_checkpoint(lu_dir: str) -> dict[str, torch.Tensor]: + from safetensors.torch import load_file + + state: dict[str, torch.Tensor] = {} + for path in sorted(glob.glob(os.path.join(lu_dir, "*.safetensors"))): + state.update(load_file(path)) + return state + + +def backend_flags() -> dict[str, bool]: + return { + "cudnn.allow_tf32": torch.backends.cudnn.allow_tf32, + "cudnn.benchmark": torch.backends.cudnn.benchmark, + "cudnn.deterministic": torch.backends.cudnn.deterministic, + "matmul.allow_tf32": torch.backends.cuda.matmul.allow_tf32, + } + + +def run_outputs(bank: torch.nn.Module, forward: Callable[[torch.nn.Module, torch.Tensor, int], torch.Tensor], + device: str = "cuda") -> dict[str, torch.Tensor]: + """``{"_x": output}``: fp32 weights without autocast, then bf16 weights under bf16 autocast (the + production LU call: input cast to the module dtype, ``torch.autocast("cuda", bfloat16)``).""" + torch.backends.cudnn.benchmark = False + z = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(SEED)) + outputs = {} + with torch.no_grad(): + bank.to(device=device, dtype=torch.float32) + for scale in SCALES: + outputs[f"fp32_x{scale}"] = forward(bank, z.to(device), scale).float().cpu() + bank.to(dtype=torch.bfloat16) + for scale in SCALES: + with torch.autocast("cuda", dtype=torch.bfloat16): + outputs[f"bf16_x{scale}"] = forward(bank, z.to(device, torch.bfloat16), scale).float().cpu() + return outputs + + +def _previous(lu_dir: str) -> dict: + import fastvideo + from fastvideo.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerBank + from fastvideo.models.upsamplers.kandinsky6_sr_lu import latent_upscaler_for_scale, run_latent_upscaler + + config = json.load(open(os.path.join(lu_dir, "config.json"))) + with torch.device("meta"): + bank = Kandinsky6SRLatentUpscalerBank(config) + bank.load_state_dict(load_checkpoint(lu_dir), strict=True, assign=True) + bank.eval() + outputs = run_outputs(bank, lambda b, z, s: run_latent_upscaler(latent_upscaler_for_scale(b, s), z)) + return { + "fastvideo_file": fastvideo.__file__, + "outputs": outputs, + "num_parameters": sum(p.numel() for p in bank.parameters()), + "state_keys": list(bank.state_dict()), + "backend_flags": backend_flags(), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--impl", choices=["previous"], required=True) + parser.add_argument("--lu-dir", required=True) + parser.add_argument("--out", required=True) + args = parser.parse_args() + torch.save(_previous(args.lu_dir), args.out) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/kandinsky6_sr/reference.py b/tests/local_tests/kandinsky6_sr/reference.py new file mode 100644 index 000000000..55b82a16c --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/reference.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Loads the k6_video reference (``kandinsky_sr``) for the Kandinsky6 SR parity tests. + +Local-only tooling: the reference package is not a FastVideo dependency. It is found via +``$KANDINSKY_SR_SRC`` (the ``k6_video/src`` directory) or a ``k6_video`` checkout next to this repository. On a +machine without CUDA / flash-attn / loguru small stubs are installed (per-sequence SDPA instead of flash-attn varlen, +no-op logger) and the base resolutions are shrunk, so the whole reference pipeline runs on CPU with tiny random +models. On a GPU box with the real packages nothing is stubbed. + +IMPORTANT: import FastVideo (``import fastvideo...``) BEFORE calling :func:`load_reference`; diffusers probes +``flash_attn`` at import time and chokes on a stub module. +""" +from __future__ import annotations + +import copy +import importlib +import os +import sys +import types +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT / "fastvideo" / "tests" / "stages" / "kandinsky6_sr")) + +import k6_sr_tiny # noqa: E402 (tiny configs shared with the CPU unit tests) + +TINY_RESOLUTIONS = k6_sr_tiny.TINY_RESOLUTIONS + + +def find_reference_src() -> Path | None: + """``$KANDINSKY_SR_SRC`` (``none`` disables the lookup, forcing the tests to skip) or ``../k6_video/src``.""" + if os.environ.get("KANDINSKY_SR_SRC", "").lower() == "none": + return None + candidates = [os.environ.get("KANDINSKY_SR_SRC"), str(REPO_ROOT.parent / "k6_video" / "src")] + for candidate in candidates: + if candidate and (Path(candidate) / "kandinsky_sr").is_dir(): + return Path(candidate) + return None + + +def _varlen_sdpa(q, k, v, cu_q, cu_k): + import torch.nn.functional as F + + outs = [] + for i in range(len(cu_q) - 1): + qi = q[int(cu_q[i]):int(cu_q[i + 1])].transpose(0, 1).unsqueeze(0) + ki = k[int(cu_k[i]):int(cu_k[i + 1])].transpose(0, 1).unsqueeze(0) + vi = v[int(cu_k[i]):int(cu_k[i + 1])].transpose(0, 1).unsqueeze(0) + outs.append(F.scaled_dot_product_attention(qi, ki, vi).squeeze(0).transpose(0, 1)) + return torch.cat(outs, dim=0) + + +def _install_stubs() -> None: + os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") + if "flash_attn" not in sys.modules: + try: + importlib.import_module("flash_attn") + except ImportError: + mod = types.ModuleType("flash_attn") + + def varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0, + softmax_scale=None, causal=False, return_attn_probs=False, **_): + out = _varlen_sdpa(q, k, v, cu_seqlens_q, cu_seqlens_k) + return (out, torch.zeros(1), None) if return_attn_probs else out + + def qkvpacked(qkv, cu_seqlens, max_seqlen, dropout_p=0.0, softmax_scale=None, causal=False, + return_attn_probs=False, **_): + q, k, v = qkv.unbind(dim=1) + out = _varlen_sdpa(q, k, v, cu_seqlens, cu_seqlens) + return (out, torch.zeros(1), None) if return_attn_probs else out + + mod.flash_attn_varlen_func = varlen + mod.flash_attn_varlen_qkvpacked_func = qkvpacked + sys.modules["flash_attn"] = mod + if "loguru" not in sys.modules: + try: + importlib.import_module("loguru") + except ImportError: + mod = types.ModuleType("loguru") + + class _Logger: + + def __getattr__(self, name): + return lambda *a, **k: self + + mod.logger = _Logger() + sys.modules["loguru"] = mod + + +@dataclass +class Reference: + """Handles on the reference modules used by the parity tests.""" + + constants: Any + dit_mod: Any + dx_dit_mod: Any + kvae_mod: Any + sr_pipeline: Any + stages: Any + upscale_utils: Any + lu_algo: Any + lu_config: Any + lu_factory: Any + piflow_sampler: Any + algo_utils: Any + tile_grid: Any + tiling_utils: Any + video_io: Any + tiny: bool + + +_reference: Reference | None = None + + +def load_reference(tiny: bool | None = None) -> Reference: + """Import the reference (skips the calling test when it is unavailable). + + ``tiny`` (default: True unless CUDA is available) shrinks ``constants.RESOLUTIONS`` in place to 64/96 px bases; + it only takes effect on the first call (tests needing the trained bases patch the dict, see the tiling tests). + """ + global _reference + src = find_reference_src() + if src is None: + pytest.skip("k6_video reference (kandinsky_sr) not found: set KANDINSKY_SR_SRC to the k6_video/src directory") + if tiny is None: + tiny = not torch.cuda.is_available() + if _reference is None: + _install_stubs() + if str(src) not in sys.path: + sys.path.insert(0, str(src)) + try: + imp = importlib.import_module + _reference = Reference( + constants=imp("kandinsky_sr.constants"), + dit_mod=imp("kandinsky_sr.core.components.model.dit"), + dx_dit_mod=imp("kandinsky_sr.core.components.model.dx_dit"), + kvae_mod=imp("kandinsky_sr.core.components.video_kvae.cached_model"), + sr_pipeline=imp("kandinsky_sr.pipeline.sr_pipeline"), + stages=imp("kandinsky_sr.pipeline.stages"), + upscale_utils=imp("kandinsky_sr.pipeline.upscale_utils"), + lu_algo=imp("kandinsky_sr.core.algo.latent_upscaler"), + lu_config=imp("kandinsky_sr.core.components.latent_upscaler.config"), + lu_factory=imp("kandinsky_sr.core.components.latent_upscaler.model.factory"), + piflow_sampler=imp("kandinsky_sr.core.algo.piflow_sampler"), + algo_utils=imp("kandinsky_sr.core.algo.utils"), + tile_grid=imp("kandinsky_sr.pipeline.tile_grid"), + tiling_utils=imp("kandinsky_sr.core.algo.tiling_utils"), + video_io=imp("kandinsky_sr.pipeline.video_io"), + tiny=tiny, + ) + except Exception as exc: # missing optional dependency of the reference on this machine + pytest.skip(f"k6_video reference cannot be imported here: {exc!r}") + if tiny: + _reference.constants.RESOLUTIONS.clear() + _reference.constants.RESOLUTIONS.update(TINY_RESOLUTIONS) + return _reference + + +def randomize(module: torch.nn.Module, seed: int, std: float = 0.05) -> None: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for parameter in module.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator) * std) + + +def build_reference_dit(ref: Reference, piflow: dict | None = k6_sr_tiny.TINY_PIFLOW, dit_cfg: dict | None = None, + seed: int = 0): + cfg = copy.deepcopy(dit_cfg or k6_sr_tiny.TINY_DIT_CFG) + cfg.pop("sr_params", None) # bundle metadata, not constructor arguments of the reference model + cfg.pop("attribute_overrides", None) + cfg["patch_size"] = tuple(cfg["patch_size"]) + cfg["axes_dims"] = tuple(cfg["axes_dims"]) + cfg["attention_params"] = {512: {"type": "flash"}, "512": {"type": "flash"}} + if piflow is not None: + dit = ref.dx_dit_mod.DXDiTWrapper(cfg, out_visual_dim=cfg["out_visual_dim"], n_grid=piflow["n_grid"]).eval() + dit.piflow_params = dict(piflow) + else: + dit = ref.dit_mod.get_dit(cfg).eval() + randomize(dit, seed + 2) + return dit + + +def load_into_port(reference_module: torch.nn.Module, port_module: torch.nn.Module, prefix: str = "") -> None: + """Load ``reference.state_dict()`` (the unprefixed official Diffusers key layout) into a port module, strictly.""" + from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + + state = {f"{prefix}{key}": value for key, value in reference_module.state_dict().items()} + mapping = get_param_names_mapping(getattr(port_module, "param_names_mapping", {})) + custom, _ = hf_to_custom_state_dict(state, mapping) + port_module.load_state_dict(custom, strict=True) diff --git a/tests/local_tests/kandinsky6_sr/sr_pipeline_capture.py b/tests/local_tests/kandinsky6_sr/sr_pipeline_capture.py new file mode 100644 index 000000000..207bab2f6 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/sr_pipeline_capture.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run one Kandinsky6 SR bundle end to end and save the output frames (driver of the pipeline parity test). + +Run in a subprocess per implementation, because the current and the previous implementation are both the package +``fastvideo``. ``--legacy-steps`` passes the step count the way the previous implementation took it (``sr_num_steps`` +Euler grid points instead of ``num_inference_steps`` DiT calls). +""" +import argparse +import json +from pathlib import Path + +import numpy as np + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--cases", required=True, help="JSON list of {name, video, scale, steps}") + parser.add_argument("--out", required=True) + parser.add_argument("--legacy-steps", action="store_true") + args = parser.parse_args() + + import fastvideo + from fastvideo import VideoGenerator + + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + generator = VideoGenerator.from_pretrained(args.model, num_gpus=1, dit_cpu_offload=False, vae_cpu_offload=False, + dit_layerwise_offload=False) + for case in json.loads(args.cases): + kwargs = {"video_path": case["video"], "sr_resolution_scale": case["scale"]} + if case.get("steps") is not None: + if args.legacy_steps: + kwargs["sr_num_steps"] = case["steps"] + 1 + else: + kwargs["num_inference_steps"] = case["steps"] + result = generator.generate_video(None, seed=42, save_video=False, return_frames=True, **kwargs) + np.save(out / f"{case['name']}.npy", np.stack([np.asarray(frame) for frame in result["frames"]])) + generator.shutdown() + (out / "fastvideo_path.txt").write_text(str(Path(fastvideo.__file__).resolve().parent)) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_dit_parity.py b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_dit_parity.py new file mode 100644 index 000000000..407020ed0 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_dit_parity.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo Kandinsky6 SR DiT vs the k6_video reference DiT (random tiny weights, batched vs packed layout). + +The reference ``state_dict()`` (unprefixed, i.e. the official Diffusers layout) is loaded into the port through the arch config's +``param_names_mapping`` with ``strict=True``; outputs are compared in fp32. On CPU the result is bit-identical; +on CUDA (flash-attn vs SDPA) a small tolerance applies. Skips when the reference package is unavailable. +""" +from __future__ import annotations + +import pytest +import torch + +# FastVideo is imported before the reference is loaded (see reference.py). +from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + +from . import reference as ref_env + +k6_sr_tiny = ref_env.k6_sr_tiny + +TOL = dict(rtol=0.0, atol=0.0) if not torch.cuda.is_available() else dict(rtol=1e-3, atol=1e-3) +B, T, H, W, C = 2, 3, 4, 6, 4 + + +def _packed_call(ref_dit, x, t, extra=None): + """Run the reference DiT on the packed ragged layout that the reference pipeline uses.""" + batch, frames = x.shape[:2] + packed = x.reshape(batch * frames, *x.shape[2:]) + cu = frames * torch.arange(batch + 1, dtype=torch.int32) + pos = [torch.cat([torch.arange(frames) for _ in range(batch)]), torch.arange(x.shape[2] // 2), + torch.arange(x.shape[3] // 2)] + kwargs = dict(scale_factor=(1.0, 1.0, 1.0)) + kwargs.update(extra or {}) + with torch.no_grad(): + return ref_dit(packed, torch.zeros(0, 1), torch.zeros(batch, 1), t, cu, + torch.zeros(batch + 1, dtype=torch.int32), pos, torch.zeros(0, dtype=torch.long), **kwargs) + + +def _port_call(dit, x, t, extra=None): + pos = [torch.arange(x.shape[1]), torch.arange(x.shape[2] // 2), torch.arange(x.shape[3] // 2)] + with torch.no_grad(): + return dit(x, t, pos, (1.0, 1.0, 1.0), **(extra or {})) + + +@pytest.mark.parametrize("piflow", [k6_sr_tiny.TINY_PIFLOW, None]) +def test_dit_forward_matches_the_reference(piflow): + ref = ref_env.load_reference() + ref_dit = ref_env.build_reference_dit(ref, piflow) + dit = k6_sr_tiny.build_dit(piflow) + ref_env.load_into_port(ref_dit, dit) # strict + x = torch.randn(B, T, H, W, 2 * C + 1, generator=torch.Generator().manual_seed(1)) + t = torch.tensor([300.0, 700.0]) + expected = _packed_call(ref_dit, x, t) + got = _port_call(dit, x, t) + if piflow is not None: # DX wrapper already moved the grid axis: [B*T, n_grid, H, W, out] + n_grid = piflow["n_grid"] + got = got.reshape(B * T, H, W, n_grid, C).movedim(-2, 1) + else: + got = got.reshape(B * T, H, W, C) + torch.testing.assert_close(got, expected, **TOL) + + +def test_reference_state_dict_keys_map_onto_the_port_exactly(): + ref = ref_env.load_reference() + ref_dit = ref_env.build_reference_dit(ref) + dit = k6_sr_tiny.build_dit() + mapping = get_param_names_mapping(dit.param_names_mapping) + custom, _ = hf_to_custom_state_dict(dict(ref_dit.state_dict()), mapping) # unprefixed = the official layout + assert set(custom) == set(dit.state_dict()) + assert all(custom[k].shape == v.shape for k, v in dit.state_dict().items()) diff --git a/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py new file mode 100644 index 000000000..301ed2b58 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real-weights parity of the Kandinsky6 SR latent upscaler against the previous FastVideo implementation. + +The previous implementation (a general re-implementation of the training code's upscaler family) was itself checked +against the ``k6_video`` reference; the current module implements only the released architecture. Both run the +released ``latent_upscaler/`` weights on one fixed latent tile for x2 and x4, in fp32 and in bf16 under bf16 autocast +(the production call), and the outputs must be bitwise equal. + +Needs a CUDA device and (skips otherwise): + + KANDINSKY6_SR_BUNDLE local copy of an official Diffusers SR repo (both repos share ``latent_upscaler/``) + KANDINSKY6_SR_PREVIOUS_ROOT a FastVideo checkout that still has ``fastvideo/models/upsamplers/kandinsky6_sr_lu/`` + (e.g. commit 12be8fdd); it runs in a subprocess because both packages are + named ``fastvideo`` + +Run: ``pytest tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_latent_upscaler_parity.py -v -s``. +""" +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from . import lu_parity_common as common + +BUNDLE = os.environ.get("KANDINSKY6_SR_BUNDLE") +PREVIOUS_ROOT = os.environ.get("KANDINSKY6_SR_PREVIOUS_ROOT") +LU_DIR = str(Path(BUNDLE or "") / "latent_upscaler") + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="latent-upscaler parity needs a CUDA device"), + pytest.mark.skipif(not BUNDLE or not list(Path(LU_DIR).glob("*.safetensors")), + reason="set KANDINSKY6_SR_BUNDLE to a local SR repo with latent_upscaler weights"), +] + + +@pytest.fixture(scope="module") +def previous(tmp_path_factory): + if not PREVIOUS_ROOT or not (Path(PREVIOUS_ROOT) / "fastvideo/models/upsamplers/kandinsky6_sr_lu").is_dir(): + pytest.skip("set KANDINSKY6_SR_PREVIOUS_ROOT to a FastVideo checkout with the previous latent upscaler") + out = tmp_path_factory.mktemp("lu_parity") / "previous.pt" + env = dict(os.environ, PYTHONPATH=PREVIOUS_ROOT) + subprocess.run([sys.executable, common.__file__, "--impl", "previous", "--lu-dir", LU_DIR, "--out", + str(out)], + cwd=PREVIOUS_ROOT, + env=env, + check=True) + result = torch.load(out) + assert Path(result["fastvideo_file"]).resolve().is_relative_to(Path(PREVIOUS_ROOT).resolve()) + return result + + +def _build_current(): + from fastvideo.configs.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerConfig + from fastvideo.models.upsamplers.kandinsky6_sr import Kandinsky6SRLatentUpscalerBank + + config = json.loads((Path(LU_DIR) / "config.json").read_text()) + with torch.device("meta"): + return Kandinsky6SRLatentUpscalerBank(Kandinsky6SRLatentUpscalerConfig(**config)) + + +@pytest.fixture(scope="module") +def current(): + bank = _build_current() + bank.load_state_dict(common.load_checkpoint(LU_DIR), strict=True, assign=True) + bank.eval() + outputs = common.run_outputs(bank, lambda b, z, s: b(z, s)) + return { + "outputs": outputs, + "num_parameters": sum(p.numel() for p in bank.parameters()), + "state_keys": list(bank.state_dict()), + "backend_flags": common.backend_flags(), + } + + +def test_state_dict_matches_the_previous_implementation(previous, current): + assert current["state_keys"] == previous["state_keys"] + assert current["num_parameters"] == previous["num_parameters"] + + +def test_partial_checkpoint_fails_the_strict_load(): + state = common.load_checkpoint(LU_DIR) + state.pop(next(key for key in state if key.startswith("_models.2x.x2_branch."))) + with pytest.raises(RuntimeError, match="Missing key"): + _build_current().load_state_dict(state, strict=True, assign=True) + + +@pytest.mark.parametrize("dtype", common.DTYPES) +@pytest.mark.parametrize("scale", common.SCALES) +def test_outputs_are_bitwise_equal(previous, current, dtype, scale): + assert current["backend_flags"] == previous["backend_flags"] + key = f"{dtype}_x{scale}" + got, want = current["outputs"][key], previous["outputs"][key] + n, c, t, h, w = common.LATENT_SHAPE + assert got.shape == want.shape == (n, c, t, h * scale, w * scale) + assert torch.isfinite(got).all() + max_diff = (got - want).abs().max().item() + print(f"{key}: max |diff| = {max_diff}") + assert torch.equal(got, want), f"{key}: max |diff| = {max_diff}" + + +def test_production_loader_matches_the_previous_bf16_outputs(previous): + """``UpsamplerLoader`` (config from the pipeline config, bf16 build, strict load) reproduces the bf16 outputs.""" + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.loader.component_loader import PipelineComponentLoader + + args = FastVideoArgs.from_kwargs(model_path=BUNDLE, num_gpus=1, use_fsdp_inference=False, dit_cpu_offload=False, + vae_cpu_offload=False, pin_cpu_memory=False) + index = json.loads((Path(BUNDLE) / "model_index.json").read_text()) + args._model_index_class_names = {key: spec[1] for key, spec in index.items() if isinstance(spec, list)} + library = index["latent_upscaler"][0] + bank = PipelineComponentLoader.load_module("latent_upscaler", LU_DIR, library, args) + assert type(bank).__name__ == "Kandinsky6SRLatentUpscalerBank" and bank.scales == (2, 4) + assert next(bank.parameters()).dtype == torch.bfloat16 and not bank.training + assert all(not p.requires_grad for p in bank.parameters()) + z = torch.randn(common.LATENT_SHAPE, generator=torch.Generator().manual_seed(common.SEED)) + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + for scale in common.SCALES: + got = bank(z.to("cuda", torch.bfloat16), scale).float().cpu() + assert torch.equal(got, previous["outputs"][f"bf16_x{scale}"]), f"x{scale}" diff --git a/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_real_weights.py b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_real_weights.py new file mode 100644 index 000000000..2e3f5623f --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_real_weights.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real-weights smoke / reference-comparison scaffold for the Kandinsky6 SR pipeline (CUDA + an official SR repo). + +STATUS: scaffold. It needs a CUDA device and a local copy of an official Diffusers SR repo, and skips without them. +Run it once per repo: the flow-matching ``kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`` and the distilled +``kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers`` (the reference comparison needs the reference output of +the same model). Environment variables: + + KANDINSKY6_SR_BUNDLE local directory of the official Diffusers SR repo under test + KANDINSKY6_SR_TEST_VIDEO a short low-resolution clip (e.g. 16-24 frames at 24 fps, a few hundred pixels wide) + KANDINSKY6_SR_REFERENCE_VIDEO optional: the ``kandy-sr from-video`` output for the same clip, scale and seed + KANDINSKY6_SR_REFERENCE_SCALE the ``--resolution-scale`` used for that reference run (default 2) + KANDINSKY6_SR_MIN_PSNR optional dB floor for the comparison (default 30; NOT calibrated, see README) + +Run on a CUDA machine: ``pytest tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_real_weights.py -v -s``. +""" +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +BUNDLE = os.environ.get("KANDINSKY6_SR_BUNDLE") +VIDEO = os.environ.get("KANDINSKY6_SR_TEST_VIDEO") +REFERENCE_VIDEO = os.environ.get("KANDINSKY6_SR_REFERENCE_VIDEO") +MIN_PSNR = float(os.environ.get("KANDINSKY6_SR_MIN_PSNR", "30")) + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="real-weights SR validation needs a CUDA device"), + pytest.mark.skipif(not BUNDLE or not Path(BUNDLE or "").is_dir(), + reason="set KANDINSKY6_SR_BUNDLE to a local copy of the official Diffusers SR repo"), + pytest.mark.skipif(not VIDEO or not Path(VIDEO or "").is_file(), + reason="set KANDINSKY6_SR_TEST_VIDEO to a short low-resolution clip"), +] + + +def decode_rgb(path) -> np.ndarray: + """All frames of a video as a ``[T, H, W, 3]`` uint8 array.""" + av = pytest.importorskip("av") + frames = [] + with av.open(str(path)) as container: + stream = container.streams.video[0] + for frame in container.decode(stream): + frames.append(frame.to_ndarray(format="rgb24")) + return np.stack(frames) + + +def per_frame_psnr(a: np.ndarray, b: np.ndarray) -> np.ndarray: + assert a.shape == b.shape, f"shape mismatch {a.shape} vs {b.shape}" + mse = ((a.astype(np.float64) - b.astype(np.float64))**2).mean(axis=(1, 2, 3)) + return 10.0 * np.log10(255.0**2 / np.maximum(mse, 1e-10)) + + +@pytest.fixture(scope="module") +def generator(): + from fastvideo import VideoGenerator + + # Strict loading of the DiT / KVAE / latent-upscaler weights happens here: a missing or unexpected key fails + # from_pretrained instead of producing garbage later. + gen = VideoGenerator.from_pretrained( + BUNDLE, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + ) + yield gen + gen.shutdown() + + +def _sr(generator, out_dir: Path, **extensions): + return generator.generate({ + "inputs": { + "video_path": VIDEO + }, + "output": { + "output_path": str(out_dir), + "save_video": True, + "return_frames": False, + }, + "extensions": extensions, + }) + + +@pytest.mark.parametrize("scale", [2, 2.25]) +def test_sr_generates_a_decodable_video_of_the_expected_geometry(generator, tmp_path, scale): + from fastvideo.pipelines.basic.kandinsky6_sr import sr_io + + result = _sr(generator, tmp_path, sr_resolution_scale=scale) + assert result.video_path and Path(result.video_path).is_file() + + source = decode_rgb(VIDEO) + out = decode_rgb(result.video_path) + expected_frames = sr_io.read_video(VIDEO)[0].shape[0] # fps rule + 1+8k alignment of the request contract + assert out.shape[0] == expected_frames + assert out.shape[1] >= int(source.shape[1] * scale) - 16 and out.shape[2] >= int(source.shape[2] * scale) - 16 + assert out.dtype == np.uint8 and out.std() > 1.0, "output is constant / empty" + + +def test_sr_is_deterministic_for_a_fixed_seed(generator, tmp_path): + first = decode_rgb(_sr(generator, tmp_path / "a", sr_resolution_scale=2).video_path) + second = decode_rgb(_sr(generator, tmp_path / "b", sr_resolution_scale=2).video_path) + # Same seed, same clip: drift here is sampling nondeterminism (the 40 dB floor is a starting value, uncalibrated). + assert per_frame_psnr(first, second).min() > 40.0 + + +@pytest.mark.skipif(not REFERENCE_VIDEO or not Path(REFERENCE_VIDEO or "").is_file(), + reason="set KANDINSKY6_SR_REFERENCE_VIDEO to the kandy-sr output for the same clip") +def test_sr_is_close_to_the_reference_output(generator, tmp_path): + scale = float(os.environ.get("KANDINSKY6_SR_REFERENCE_SCALE", "2")) + ours = decode_rgb(_sr(generator, tmp_path, sr_resolution_scale=scale).video_path) + reference = decode_rgb(REFERENCE_VIDEO) + psnr = per_frame_psnr(ours, reference) + print(f"PSNR vs reference: mean {psnr.mean():.2f} dB, min {psnr.min():.2f} dB") + assert psnr.mean() >= MIN_PSNR diff --git a/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py new file mode 100644 index 000000000..d82dfdb83 --- /dev/null +++ b/tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise parity of the native ``Kandinsky6SRVAE`` against a reference checkout of the previous port, on the real +checkpoint (CUDA). + +The reference is another FastVideo checkout (its package is also named ``fastvideo``), so it runs in a subprocess with +``PYTHONPATH`` pointing at it and saves its outputs to a ``.pt`` file; this process runs the native VAE and requires +``torch.equal`` on every output. Cases follow the pipeline's use: whole-clip encode in ``encode_lq_video_to_lr_latent`` +style (uint8 -> float -> ``normalize_data`` -> VAE dtype -> ``encode(x)[0]``) and ``decode(z).sample``, in fp32 and bf16 +(the pipeline's ``vae_precision``) at an LQ-sized clip, plus a bf16 4x-SR-sized tile (832x1536) whose decoder +convolutions exceed the chunked-conv threshold. + +Environment variables (the test skips without the weights or the reference checkout): + + KANDINSKY6_SR_VAE_DIR ``vae/`` of an official Diffusers SR repo (config.json + safetensors) + KANDINSKY6_SR_VAE_REFERENCE root of the reference FastVideo checkout + +GPU: needs ~80 GB free for the 832x1536 bf16 decode (run on a GB200 / H200 class device). +Run: ``pytest tests/local_tests/kandinsky6_sr/test_kandinsky6_sr_vae_parity.py -v -s``. +""" +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import torch + +VAE_DIR = Path(os.environ.get("KANDINSKY6_SR_VAE_DIR", "/nonexistent")) +REFERENCE_ROOT = Path(os.environ.get("KANDINSKY6_SR_VAE_REFERENCE", "/nonexistent")) + +# (name, dtype, frames, height, width): 33 frames = two encode segments (17 + 16) and two decode segments (5 + 4). +CASES = [ + ("lq_fp32", torch.float32, 33, 208, 384), + ("lq_bf16", torch.bfloat16, 33, 208, 384), + ("tile4x_bf16", torch.bfloat16, 33, 832, 1536), +] + + +def make_clip(frames: int, height: int, width: int) -> torch.Tensor: + """Deterministic ``[T, 3, H, W]`` uint8 clip with moving structure and texture (not flat, not pure noise).""" + t = torch.arange(frames, dtype=torch.float32)[:, None, None] + y = torch.linspace(0, 1, height)[None, :, None] + x = torch.linspace(0, 1, width)[None, None, :] + channels = [ + 128 + 100 * torch.sin(12 * x + 0.3 * t) * torch.cos(7 * y - 0.2 * t), + 128 + 90 * torch.cos(20 * (x + y) - 0.4 * t), + 255 * ((x - 0.5 - 0.01 * t)**2 + (y - 0.5)**2 < 0.05).float() + 20 * torch.sin(60 * x * y + t), + ] + noise = torch.randn((frames, 3, height, width), generator=torch.Generator().manual_seed(0)) * 8 + return (torch.stack(channels, dim=1) + noise).clamp(0, 255).to(torch.uint8) + + +def load_vae(vae_dir: Path, device: str): + """Build ``Kandinsky6SRVAE`` from ``config.json`` like ``VAELoader`` and load the checkpoint strictly (fp32).""" + from safetensors.torch import load_file + + from fastvideo.configs.models.vaes.kandinsky6_sr import Kandinsky6SRVAEConfig + from fastvideo.models.vaes.kandinsky6_sr import Kandinsky6SRVAE + + config_dict = json.loads((vae_dir / "config.json").read_text()) + config_dict.pop("_class_name", None) + config = Kandinsky6SRVAEConfig() + config.update_model_arch(config_dict) + with torch.device(device): + vae = Kandinsky6SRVAE(config) + state: dict[str, torch.Tensor] = {} + for path in sorted(vae_dir.glob("*.safetensors")): + state.update(load_file(str(path), device=device)) + vae.load_state_dict(state, strict=True) + return vae.eval(), sorted(state) + + +@torch.no_grad() +def run_cases(vae_dir: Path, device: str = "cuda") -> dict[str, object]: + import fastvideo + + vae, checkpoint_keys = load_vae(vae_dir, device) + outputs: dict[str, object] = { + "fastvideo_file": fastvideo.__file__, + "checkpoint_keys": checkpoint_keys, + "state_keys": sorted(vae.state_dict()), + "num_params": sum(p.numel() for p in vae.parameters()), + } + for name, dtype, frames, height, width in CASES: + vae.to(dtype) + clip = make_clip(frames, height, width) + pixel = vae.normalize_data(clip.permute(1, 0, 2, 3).unsqueeze(0).to(device).float()).to(dtype) + latent = vae.encode(pixel)[0] + decoded = vae.decode(latent).sample + outputs[f"{name}/latent"] = latent.cpu() + outputs[f"{name}/decoded"] = decoded.cpu() + del pixel, latent, decoded + torch.cuda.empty_cache() + return outputs + + +def _run_reference(out_path: Path) -> dict[str, object]: + env = dict(os.environ, PYTHONPATH=str(REFERENCE_ROOT)) + subprocess.run([sys.executable, __file__, str(VAE_DIR), str(out_path)], cwd=REFERENCE_ROOT, env=env, check=True) + return torch.load(out_path, weights_only=False) + + +def test_native_vae_is_bitwise_equal_to_reference_on_real_weights(tmp_path): + import pytest + + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + if not (VAE_DIR / "config.json").is_file() or not any(VAE_DIR.glob("*.safetensors")): + pytest.skip(f"no Kandinsky6 SR VAE weights at {VAE_DIR} (set KANDINSKY6_SR_VAE_DIR)") + if not (REFERENCE_ROOT / "fastvideo" / "models" / "vaes" / "kandinsky6_sr_kvae").is_dir(): + pytest.skip(f"no reference checkout at {REFERENCE_ROOT} (set KANDINSKY6_SR_VAE_REFERENCE)") + + reference = _run_reference(tmp_path / "reference.pt") + assert str(reference["fastvideo_file"]).startswith(str(REFERENCE_ROOT)) + native = run_cases(VAE_DIR) + assert not str(native["fastvideo_file"]).startswith(str(REFERENCE_ROOT)) + + # Strict load already passed in both; the native keys are the checkpoint's own, the reference's carry "model.". + assert native["state_keys"] == native["checkpoint_keys"] + assert native["state_keys"] == sorted(key.removeprefix("model.") for key in reference["state_keys"]) + assert native["num_params"] == reference["num_params"] + + failures = [] + for name, *_ in CASES: + for kind in ("latent", "decoded"): + key = f"{name}/{kind}" + ours, theirs = native[key], reference[key] + assert ours.shape == theirs.shape and ours.dtype == theirs.dtype, key + equal = torch.equal(ours, theirs) + diff = (ours.float() - theirs.float()).abs().max().item() + print(f"{key}: shape={tuple(ours.shape)} dtype={ours.dtype} bitwise={equal} max_abs_diff={diff:.3e}") + if not equal: + failures.append(f"{key} max_abs_diff={diff:.3e}") + assert not failures, failures + + +if __name__ == "__main__": + # Reference worker: python test_kandinsky6_sr_vae_parity.py + torch.save(run_cases(Path(sys.argv[1])), sys.argv[2]) diff --git a/tests/local_tests/pipelines/test_kandinsky6_sr_pipeline_parity.py b/tests/local_tests/pipelines/test_kandinsky6_sr_pipeline_parity.py new file mode 100644 index 000000000..6b858615b --- /dev/null +++ b/tests/local_tests/pipelines/test_kandinsky6_sr_pipeline_parity.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real-weights end-to-end parity of the Kandinsky6 SR pipeline against a previous FastVideo checkout. + +The previous implementation ran its own Euler / pi-Flow loops; the current pipeline steps the bundle's scheduler. The +two differ only in scheduler arithmetic: the Euler update ``dt * velocity`` is fp32 (the previous loop multiplied in +bf16) and ``PiflowScheduler`` computes its schedule in float32 with the final segment ending at ``eps``. With those two +conventions patched back, the outputs were bitwise identical on GB200; as shipped, the observed PSNR was 50.9-56.5 dB. + +Environment (the test skips without them): + KANDINSKY6_SR_MODELS directory holding ``Kandinsky-6.0-VSR-5s-Diffusers`` and + ``Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers`` + KANDINSKY6_SR_VIDEO a short low-resolution clip (e.g. 17 frames at 384x208) + KANDINSKY6_SR_PREVIOUS_ROOT checkout of the previous implementation (PR head 12be8fdd) +""" +import json +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +_HELPER = Path(__file__).resolve().parents[1] / "kandinsky6_sr" / "sr_pipeline_capture.py" +_MIN_PSNR_DB = 48.0 +_BUNDLES = { + "flow_matching": ("Kandinsky-6.0-VSR-5s-Diffusers", [("x2", 2.0, None), ("x2_25", 2.25, None), ("x2_steps2", 2.0, 2), + ("x4", 4.0, None)]), + "distilled": ("Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", [("x2", 2.0, None), ("x4", 4.0, None)]), +} + + +def _env(name: str) -> str: + value = os.environ.get(name) + if not value: + pytest.skip(f"{name} is not set") + return value + + +def _capture(root: Path, model: Path, cases: list[dict], out: Path, legacy: bool) -> None: + env = dict(os.environ, PYTHONPATH=os.pathsep.join([str(root), os.environ.get("PYTHONPATH", "")])) + command = [sys.executable, str(_HELPER), "--model", str(model), "--cases", json.dumps(cases), "--out", str(out)] + if legacy: + command.append("--legacy-steps") + subprocess.run(command, cwd=root, env=env, check=True) + assert Path((out / "fastvideo_path.txt").read_text()) == (root / "fastvideo").resolve() + + +def _psnr(a: np.ndarray, b: np.ndarray) -> float: + mse = float(((a.astype(np.float64) - b.astype(np.float64))**2).mean()) + return float("inf") if mse == 0 else 10 * np.log10(255**2 / mse) + + +@pytest.mark.parametrize("bundle", list(_BUNDLES)) +def test_pipeline_matches_the_previous_implementation(bundle, tmp_path): + import torch + + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + models, video = Path(_env("KANDINSKY6_SR_MODELS")), _env("KANDINSKY6_SR_VIDEO") + previous = Path(_env("KANDINSKY6_SR_PREVIOUS_ROOT")) + name, case_specs = _BUNDLES[bundle] + cases = [{"name": case, "video": video, "scale": scale, "steps": steps} for case, scale, steps in case_specs] + current = Path(__file__).resolve().parents[3] + _capture(current, models / name, cases, tmp_path / "current", legacy=False) + _capture(previous, models / name, cases, tmp_path / "previous", legacy=True) + for case in cases: + got = np.load(tmp_path / "current" / f"{case['name']}.npy") + expected = np.load(tmp_path / "previous" / f"{case['name']}.npy") + assert got.shape == expected.shape, case + psnr = _psnr(got, expected) + print(f"{bundle}/{case['name']}: PSNR {psnr:.2f} dB, max diff {np.abs(got.astype(int) - expected).max()}") + assert psnr >= _MIN_PSNR_DB, (case, psnr) diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index 2903f6b9c..f90b17c9d 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -257,3 +257,30 @@ def test_h3_command_blocks_have_unique_copy_targets(): "cookbook-client-code", ] assert len(set(parser.ids)) == len(parser.ids) + + +def test_kandinsky6_cookbook_recipes_match_maintained_examples(): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + family = [recipe for recipe in recipes if recipe["family"] == "kandinsky6"] + assert [recipe["id"] for recipe in family] == [ + "kandinsky6-ti2va-base", + "kandinsky6-ti2va-piflow", + "kandinsky6-vsr", + "kandinsky6-vsr-distilled", + ] + assert all((ROOT / recipe["source"]).is_file() for recipe in family) + assert all(recipe["evidence"] == "Source-backed" for recipe in family) + assert all(recipe["hardware"] == { + "platform": "cuda", + "gpu_count": 1, + "evidence": "source-configured", + } for recipe in family) + assert all("accelerator" not in recipe["hardware"] and "peak_memory" not in recipe["hardware"] + for recipe in family) + assert family[1]["model"].endswith("Pro-distill-5s-Diffusers") + # The distilled recipe reruns the shared example with KANDINSKY6_MODEL_PATH; the repo id selects the 10-step preset. + assert family[1]["command"].startswith("KANDINSKY6_MODEL_PATH=" + family[1]["model"] + " ") + assert "KANDINSKY6_MODEL_PATH" in (ROOT / family[1]["source"]).read_text() + assert all("INPUT_VIDEO" in recipe["command"] for recipe in family[2:]) + assert "VSR-distilled2steps" in family[3]["command"] + validate_cookbook() From 527a3165b9af3ab2a48a622c3de8e0b5e5828121 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 6 Oct 2026 11:05:40 -0700 Subject: [PATCH 051/151] [docs]: add Kandinsky 6 to the README news (#1925) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index dbca636b6..69c94a5e4 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.** ## NEWS +- `2026/10/06`: FastVideo now supports [Kandinsky 6](https://x.com/kandinskylab_ai/status/2107374635218055345) from Kandinsky Lab: text- and image-to-video with synchronized audio (base and 10-step distilled pi-Flow checkpoints) plus video super-resolution up to 4x. See the [Kandinsky 6 recipes](https://haoailab.com/FastVideo/cookbook/kandinsky6/). - `2026/09/15`: Release [FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2), an eight-forward data-free DMD2 checkpoint distilled from MiniMax-H3 with 80% Video Sparse Attention. Run it with `examples/inference/basic/basic_fasth3_8step.py` or the [FastH3 8-Step V2 recipe](https://haoailab.com/FastVideo/cookbook/minimax-h3/). - `2026/09/01`: FastH3 now runs locally on Apple Silicon through MLX and on NVIDIA DGX Spark through CUDA 13, including two-Spark inference. Follow the [FastH3 recipes](https://haoailab.com/FastVideo/cookbook/minimax-h3/) and read the [Blog](https://haoailab.com/blogs/fasth3-local/). - `2026/08/27`: [FastH3 Preview v1](https://haoailab.com/blogs/fasth3-preview/) is an open-weight 4-step sparse-distilled MiniMax-H3 model for synchronized video-and-audio generation, developed in collaboration with [Nuva Lab](https://nuvalab.ai/) and the [NVIDIA FastGen team](https://github.com/NVlabs/FastGen). Download the recommended [VSA / Data-Free weights](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree), or see the [full FastH3 collection](https://huggingface.co/collections/FastVideo/fastvideo-fasth3). From 8b51380466a963692a8c0afa254417bfcc2ba8bc Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:37:01 -0700 Subject: [PATCH 052/151] [feat]: FastH3 on single RTX GPUs (#1919) Co-authored-by: Aryan Kumar Co-authored-by: Aryan Kumar Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 --- .gitignore | 2 + docs/assets/cookbook-recipes.json | 73 ++- docs/contributing/env_vars.md | 29 + docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/kandinsky6.md | 2 +- docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 28 +- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 2 +- docs/cookbook/z-image.md | 2 +- docs/inference/fasth3_rtx_pro_6000.md | 260 ++++++++ docs/inference/optimizations.md | 5 + .../basic/basic_compacth3_rtx5090.yaml | 49 ++ .../basic/basic_compacth3_rtx_pro6000.yaml | 49 ++ .../serving/openai_compacth3_rtx5090.yaml | 53 ++ .../serving/openai_compacth3_rtx_pro6000.yaml | 53 ++ fastvideo-kernel/attn_qat_infer/api.py | 210 +++++- .../attn_qat_infer/blackwell/api.cu | 87 ++- .../attn_qat_infer/blackwell/kernel_ws.h | 6 +- .../attn_qat_infer/blackwell/launch.h | 4 +- .../blackwell/mainloop_tma_ws.h | 149 ++++- .../attn_qat_infer/blackwell/params.h | 17 + .../tests/test_attn_qat_infer_sparse.py | 155 +++++ fastvideo/api/compat.py | 14 +- fastvideo/api/schema.py | 5 + .../backends/minimax_h3_sparse_int8.py | 135 ++++ .../backends/video_sparse_attn_h3.py | 54 +- fastvideo/envs.py | 115 ++++ fastvideo/hooks/layerwise_offload.py | 95 ++- fastvideo/hooks/pinned_memory.py | 86 +++ fastvideo/layers/quantization/fp8_config.py | 21 + fastvideo/layers/quantization/fp8_kernels.py | 84 +++ fastvideo/layers/quantization/nvfp4_config.py | 261 +++++++- .../layers/quantization/nvfp4_dequant.py | 54 ++ fastvideo/models/dits/minimax_h3.py | 222 ++++++- fastvideo/models/dits/minimax_h3_vsa_fp4.py | 440 +++++++++++++ .../encoders/minimax_h3_checkpoint_nvfp4.py | 78 ++- .../models/encoders/minimax_h3_qwen3_vl.py | 29 +- fastvideo/models/loader/component_loader.py | 69 +- fastvideo/models/loader/fsdp_load.py | 150 ++++- .../models/vaes/minimax_h3_int8_convrot.py | 353 ++++++++++ .../models/vaes/minimax_h3_int8_kernels.py | 46 ++ fastvideo/models/vaes/minimax_h3_video.py | 41 +- .../basic/minimax_h3/minimax_h3_pipeline.py | 173 ++++- .../stages/minimax_h3_conditioning.py | 3 +- .../minimax_h3/stages/minimax_h3_denoising.py | 6 + .../pipelines/basic/minimax_h3/vsa_guard.py | 42 ++ fastvideo/pipelines/stages/base.py | 8 + fastvideo/tests/api/test_typed_quant_flow.py | 15 + .../attention/test_minimax_h3_sparse_int8.py | 89 +++ .../test_minimax_h3_encoder_layerwise.py | 102 +++ ...st_minimax_h3_qwen3_vl_checkpoint_nvfp4.py | 26 +- fastvideo/tests/hooks/test_pinned_memory.py | 109 ++++ ...st_minimax_h3_modelopt_activation_scale.py | 48 ++ ...est_minimax_h3_modelopt_nvfp4_converter.py | 120 ++++ .../ops/quantization/test_nvfp4_config.py | 8 + .../ops/quantization/test_nvfp4_dequant.py | 20 + .../quantization/test_nvfp4_h3_dit_export.py | 87 +++ .../test_nvfp4_minimax_h3_wiring.py | 29 + .../ops/quantization/test_nvfp4_purge.py | 21 + .../test_minimax_h3_sequential_start.py | 244 ++++++- .../tests/stages/test_minimax_h3_vsa_guard.py | 88 +++ .../test_minimax_h3_tile_first.py | 108 +++ .../vaes/test_minimax_h3_int8_convrot.py | 187 ++++++ .../tests/vaes/test_minimax_h3_int8_shared.py | 95 +++ fastvideo/worker/gpu_worker.py | 40 ++ mkdocs.yml | 1 + scripts/benchmarks/minimax_h3_4090/README.md | 422 ++++++++++++ .../minimax_h3_4090/bench_encoder_dequant.py | 59 ++ .../benchmarks/minimax_h3_4090/bench_pod.py | 187 ++++++ .../minimax_h3_4090/bench_sparse_qkv.py | 57 ++ .../minimax_h3_4090/prompts_1k.json | 4 + .../benchmarks/minimax_h3_4090/summarize.py | 63 ++ scripts/benchmarks/minimax_h3_pro6000/a2a8.py | 72 ++ scripts/benchmarks/minimax_h3_pro6000/app.py | 613 ++++++++++++++++++ .../minimax_h3_pro6000/bench_code.py | 338 ++++++++++ .../minimax_h3_pro6000/bench_headline.py | 141 ++++ .../minimax_h3_pro6000/headline_prompts.json | 4 + .../minimax_h3_pro6000/showcase_prompts.json | 6 + .../convert_minimax_h3_modelopt_nvfp4_dit.py | 261 ++++++++ 89 files changed, 7352 insertions(+), 155 deletions(-) create mode 100644 docs/inference/fasth3_rtx_pro_6000.md create mode 100644 examples/inference/basic/basic_compacth3_rtx5090.yaml create mode 100644 examples/inference/basic/basic_compacth3_rtx_pro6000.yaml create mode 100644 examples/serving/openai_compacth3_rtx5090.yaml create mode 100644 examples/serving/openai_compacth3_rtx_pro6000.yaml create mode 100644 fastvideo-kernel/tests/test_attn_qat_infer_sparse.py create mode 100644 fastvideo/attention/backends/minimax_h3_sparse_int8.py create mode 100644 fastvideo/hooks/pinned_memory.py create mode 100644 fastvideo/layers/quantization/fp8_kernels.py create mode 100644 fastvideo/layers/quantization/nvfp4_dequant.py create mode 100644 fastvideo/models/dits/minimax_h3_vsa_fp4.py create mode 100644 fastvideo/models/vaes/minimax_h3_int8_convrot.py create mode 100644 fastvideo/models/vaes/minimax_h3_int8_kernels.py create mode 100644 fastvideo/pipelines/basic/minimax_h3/vsa_guard.py create mode 100644 fastvideo/tests/attention/test_minimax_h3_sparse_int8.py create mode 100644 fastvideo/tests/encoders/test_minimax_h3_encoder_layerwise.py create mode 100644 fastvideo/tests/hooks/test_pinned_memory.py create mode 100644 fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_activation_scale.py create mode 100644 fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_nvfp4_converter.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_dequant.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py create mode 100644 fastvideo/tests/stages/test_minimax_h3_vsa_guard.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_tile_first.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py create mode 100644 fastvideo/tests/vaes/test_minimax_h3_int8_shared.py create mode 100644 scripts/benchmarks/minimax_h3_4090/README.md create mode 100644 scripts/benchmarks/minimax_h3_4090/bench_encoder_dequant.py create mode 100644 scripts/benchmarks/minimax_h3_4090/bench_pod.py create mode 100644 scripts/benchmarks/minimax_h3_4090/bench_sparse_qkv.py create mode 100644 scripts/benchmarks/minimax_h3_4090/prompts_1k.json create mode 100644 scripts/benchmarks/minimax_h3_4090/summarize.py create mode 100644 scripts/benchmarks/minimax_h3_pro6000/a2a8.py create mode 100644 scripts/benchmarks/minimax_h3_pro6000/app.py create mode 100644 scripts/benchmarks/minimax_h3_pro6000/bench_code.py create mode 100644 scripts/benchmarks/minimax_h3_pro6000/bench_headline.py create mode 100644 scripts/benchmarks/minimax_h3_pro6000/headline_prompts.json create mode 100644 scripts/benchmarks/minimax_h3_pro6000/showcase_prompts.json create mode 100644 scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py diff --git a/.gitignore b/.gitignore index 13c5fbc87..7764e2ea7 100644 --- a/.gitignore +++ b/.gitignore @@ -142,3 +142,5 @@ fastvideo/tests/ssim/reference_videos/** *.nvimlog .nvimlog .python-version +scripts/benchmarks/minimax_h3_pro6000/headline_results/ +fastvideo/tests/ssim/.reference_videos_download.lock diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index c92c95e22..63a0cdbf5 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 13, + "version": 14, "recipes": [ { "id": "fastwan21-t2v", @@ -665,6 +665,77 @@ "Height, width, frames, and steps in the YAML are examples. Edit them or pass CLI flags. See docs/getting_started/installation/spark_pair.md." ] }, + { + "id": "compacth3-rtx5090", + "group": "compacth3-rtx5090", + "group_label": "CompactH3 on RTX 5090", + "group_task": "4-step 42-block text to video + audio", + "family": "minimax_h3", + "stage": "inference", + "task": "Few-step text to video (with audio)", + "label": "CompactH3 NVFP4 on RTX 5090", + "summary": "Run the 42-block CompactH3 NVFP4 DiT with the NVFP4 Qwen3-VL encoder, Comfy int8-convrot VAE, and SageAttention3 FP4 on one 32 GB RTX 5090. Sequential load parks the encoder in pinned host RAM.", + "model": "./CompactH3", + "source": "examples/inference/basic/basic_compacth3_rtx5090.yaml", + "serving": { + "source": "examples/serving/openai_compacth3_rtx5090.yaml", + "install": "UV_TORCH_BACKEND=cu128 uv pip install -e \".[fasth3]\"", + "prepare": "hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree --local-dir ./CompactH3 --include model_index.json --include \"tokenizer/**\" --include \"processor/**\" --include \"scheduler/**\" --include \"audio_scheduler/**\" --include \"audio_vae/**\" --include \"vae/**\"\nhf download aryan5v/FastH3-20B-42block-dmd2-ckpt1400-bf16 --local-dir ./CompactH3/transformer\nhf download aryan5v/FastH3-20B-42block-dmd2-ckpt1400-nvfp4 --local-dir ./CompactH3/transformer --include nvfp4_weights.safetensors\nhf download KyleNeverGivesUp/FastH3-text-encoder-nvfp4 --local-dir ./CompactH3/text_encoder\nhf download Comfy-Org/MiniMax-H3 --local-dir ./Comfy-MiniMax-H3 --include vae/minimax_h3_video_vae_int8_convrot.safetensors\ncp ./Comfy-MiniMax-H3/vae/minimax_h3_video_vae_int8_convrot.safetensors ./CompactH3/vae/", + "env": "FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FLASHINFER_CUDA_ARCH_LIST=12.0a" + }, + "command": "FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_compacth3_rtx5090.yaml", + "gpu_types": ["NVIDIA"], + "hardware": { + "platform": "cuda", + "gpu_count": 1, + "evidence": "source-configured" + }, + "evidence": "Source-backed", + "expected_artifact": "MP4 under outputs/compacth3_rtx5090/", + "modes": ["T2VA", "CompactH3 NVFP4", "RTX 5090"], + "limitations": [ + "Assemble ./CompactH3 before running. The DiT NVFP4 export and encoder snapshot are gated; run huggingface-cli login and accept each repo license.", + "32 GB cannot keep the NVFP4 encoder and DiT on the GPU together. Keep h3_sequential_load on and lazy_module_load off.", + "Blackwell sm_120 needs ATTN_QAT_INFER, FLASHINFER_CUDA_ARCH_LIST=12.0a, and a CUDA 12.8 PyTorch wheel.", + "Legal num_frames values are 17n+5, capped at 362 (15.08 s). Native 16:9 sizes include 832x480 and 1344x768; 1344x768 on 32 GB is unmeasured." + ] + }, + { + "id": "compacth3-rtx-pro6000", + "group": "compacth3-rtx-pro6000", + "group_label": "CompactH3 on RTX PRO 6000", + "group_task": "4-step 42-block text to video + audio", + "family": "minimax_h3", + "stage": "inference", + "task": "Few-step text to video (with audio)", + "label": "CompactH3 NVFP4 on RTX PRO 6000 Blackwell", + "summary": "Run CompactH3 NVFP4 with the encoder, DiT, and int8-convrot VAE resident on one 96 GB RTX PRO 6000 Blackwell. The checked-in example is 1344x768 and 124 frames (5.17 s) with VAE torch.compile. Use 832x480 for clip-queue playground traffic.", + "model": "./CompactH3", + "source": "examples/inference/basic/basic_compacth3_rtx_pro6000.yaml", + "serving": { + "source": "examples/serving/openai_compacth3_rtx_pro6000.yaml", + "install": "UV_TORCH_BACKEND=cu128 uv pip install -e \".[fasth3]\"", + "prepare": "hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree --local-dir ./CompactH3 --include model_index.json --include \"tokenizer/**\" --include \"processor/**\" --include \"scheduler/**\" --include \"audio_scheduler/**\" --include \"audio_vae/**\" --include \"vae/**\"\nhf download aryan5v/FastH3-20B-42block-dmd2-ckpt1400-bf16 --local-dir ./CompactH3/transformer\nhf download aryan5v/FastH3-20B-42block-dmd2-ckpt1400-nvfp4 --local-dir ./CompactH3/transformer --include nvfp4_weights.safetensors\nhf download KyleNeverGivesUp/FastH3-text-encoder-nvfp4 --local-dir ./CompactH3/text_encoder\nhf download Comfy-Org/MiniMax-H3 --local-dir ./Comfy-MiniMax-H3 --include vae/minimax_h3_video_vae_int8_convrot.safetensors\ncp ./Comfy-MiniMax-H3/vae/minimax_h3_video_vae_int8_convrot.safetensors ./CompactH3/vae/", + "env": "FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FLASHINFER_CUDA_ARCH_LIST=12.0a" + }, + "command": "FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 fastvideo generate --config examples/inference/basic/basic_compacth3_rtx_pro6000.yaml", + "gpu_types": ["NVIDIA"], + "hardware": { + "platform": "cuda", + "gpu_count": 1, + "evidence": "source-configured" + }, + "evidence": "Source-backed", + "expected_artifact": "MP4 under outputs/compacth3_rtx_pro6000/", + "modes": ["T2VA", "CompactH3 NVFP4", "RTX PRO 6000"], + "limitations": [ + "Assemble ./CompactH3 before running. The DiT NVFP4 export and encoder snapshot are gated; run huggingface-cli login and accept each repo license.", + "96 GB keeps the encoder, DiT, and VAE on GPU. Do not enable h3_sequential_load or lazy_module_load on this box.", + "Enable compile.vae_enabled. Leave inference_torch_compile off: FlashInfer and Sage3 custom ops cannot be compiled.", + "Blackwell sm_120 needs ATTN_QAT_INFER, FLASHINFER_CUDA_ARCH_LIST=12.0a, and a CUDA 12.8 PyTorch wheel.", + "Legal num_frames values are 17n+5, capped at 362 (15.08 s). Native 16:9 sizes include 832x480 and 1344x768. Dense CompactH3 has zero VSA gates; the pipeline raises if VIDEO_SPARSE_ATTN_H3 is loaded with all-zero to_gate_compress weights." + ] + }, { "id": "fasth3-8step-v2-cuda", "group": "fasth3-8step-v2", diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 4f47ddcb1..883e2d8ac 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -240,6 +240,35 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_VBENCH_FULL_INFO_JSON` | str | unset | eval | Path to VBench_full_info.json, used instead of the vendored copy. Deprecated names: `VBENCH_FULL_INFO_JSON`. | | `FASTVIDEO_FVD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the FVD metric. | | `FASTVIDEO_FAD_REF_FEATURES` | str | unset | eval | Cached reference-feature file for the audio Frechet distance metric. | +| `FASTVIDEO_H3_VSA_FP4` | bool | `0` | attention | Run MiniMax-H3 VSA attention on the block-sparse SageAttention3 FP4 kernel (sm_120, no-grad, single sequence-parallel rank). | +| `FASTVIDEO_H3_VSA_TILE_FIRST` | bool | `0` | attention | Single-rank MiniMax-H3 VSA with one tile gather of the block input instead of separate Q/K/V/gate scatters. | +| `FASTVIDEO_H3_VSA_SM89_KERNEL` | one of original, bf16, int8 | `original` | attention | Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, or int8 (INT8 QK, BF16 PV). | +| `FASTVIDEO_H3_SIM_SP_FP8` | bool | `0` | debug | Simulate the FP8 sequence-parallel exchange of the MiniMax-H3 FP4 VSA path on one rank. | +| `FASTVIDEO_H3_FFN_CHUNK_TOKENS` | int | `0` | performance | Inference-only MiniMax-H3 FFN token chunk size; 0 runs the FFN unchunked. | +| `FASTVIDEO_H3_FP8_ATTENTION` | bool | `0` | performance | With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention projections in FP8. | +| `FASTVIDEO_H3_FP8_GRANULARITY` | one of tensor, channel | `tensor` | performance | FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION. | +| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. | +| `FASTVIDEO_NVFP4_ACT_AMAX` | path | unset | performance | JSON of calibrated NVFP4 input amax per linear, keyed b<block>.<sub> or full prefix; sets a static activation scale. | +| `FASTVIDEO_NVFP4_DYNAMIC_ACT` | str | `""` | performance | NVFP4 linears that derive the activation scale per call: all, or comma-separated layer-name suffixes such as ff.fc_out. | +| `FASTVIDEO_H3_ADALN_CACHE` | bool | `0` | performance | Cache MiniMax-H3 AdaLN modulation per timestep instead of keeping the projection weights resident. | +| `FASTVIDEO_H3_ADALN_TABLE` | path | unset | performance | Precomputed MiniMax-H3 AdaLN modulation table; enables the cache and skips loading the AdaLN projection weights. | +| `FASTVIDEO_H3_ADALN_DUMP` | path | unset | debug | Write the MiniMax-H3 AdaLN modulation table to this path while sampling. | +| `FASTVIDEO_H3_SPLICE_TRANSFORMER` | path | unset | eval | Second MiniMax-H3 transformer that runs the late DMD steps (checkpoint step-splice evaluation). | +| `FASTVIDEO_H3_SPLICE_FROM_STEP` | int | `4` | eval | First denoising step run by FASTVIDEO_H3_SPLICE_TRANSFORMER. | +| `FASTVIDEO_H3_ENCODER_LAYERWISE` | bool | `0` | performance | Stream MiniMax-H3 text-encoder language layers through exact-size pinned host memory (text-only prompts). | +| `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT` | bool | `0` | performance | Expand the serialized NVFP4 MiniMax-H3 text encoder with one fused Triton pass on GPUs without FP4 GEMM. | +| `FASTVIDEO_H3_VAE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per tile. | +| `FASTVIDEO_H3_VAE_INT8_SHARED_QKV` | bool | `0` | performance | Share the INT8 activation rotation and quantization across the MiniMax-H3 VAE Q/K/V projections. | +| `FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW` | bool | `0` | performance | Use transposed weight views in the MiniMax-H3 VAE INT8 projections. | +| `FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT` | bool | `0` | performance | Fused dequantization epilogue for the MiniMax-H3 VAE INT8 projections. | +| `FASTVIDEO_H3_PINNED_SWAP` | bool | `1` | performance | Swap offloaded MiniMax-H3 modules through exact-size pinned host arenas. | +| `FASTVIDEO_H3_PARK_MODULES` | str | unset | performance | Comma-separated MiniMax-H3 denoise modules parked on the host while the text encoder runs, e.g. vae,audio_vae. | +| `FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS` | bool | `0` | performance | Layerwise offload also streams large buffers such as packed FP4/FP8 weights. | +| `FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS` | int | `0` | performance | Keep the first N layerwise-offloaded blocks resident on the GPU. | +| `FASTVIDEO_H3_SP_PROFILE` | bool | `0` | profiling | CUDA-event spans per stage over one MiniMax-H3 FP4 VSA DiT forward. | +| `FASTVIDEO_H3_CAPTURE_QKV` | path | unset | debug | Directory for captured real MiniMax-H3 Q/K/V attention inputs. | +| `FASTVIDEO_CUDA_MEMORY_CAP_GIB` | float | `0.0` | debug | Cap this process's CUDA allocator at this many GiB to emulate a smaller GPU; 0 leaves it uncapped. | +| `FASTVIDEO_MEMORY_REPORT` | bool | `0` | debug | Log bytes held per pipeline component by device and dtype after loading. | | `FASTVIDEO_TEST_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`. | diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 89e39e845..4d9c8c0fb 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
      +
      All model families
      diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index 82508bd81..160c56ca5 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
      +
      All model families
      diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index 51f4a7f28..93f8b5cf2 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
      +
      All model families
      diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index 75caf3144..dede8af94 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
      +
      All model families
      diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index bbd96547b..d899dd977 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
      +
      All model families
      diff --git a/docs/cookbook/kandinsky6.md b/docs/cookbook/kandinsky6.md index 3aac661fa..fc9f32dff 100644 --- a/docs/cookbook/kandinsky6.md +++ b/docs/cookbook/kandinsky6.md @@ -5,7 +5,7 @@ hide: # Kandinsky 6 recipes -
      +
      All model families
      diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index 3b43788d4..28f303deb 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
      +
      All model families
      diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index 53bf0fa14..f751661c9 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
      +
      All model families
      diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index 15b62ae34..b6d4e16cd 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
      +
      All model families
      diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 0aabcb034..e01ee210a 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -5,7 +5,15 @@ hide: # MiniMax H3 recipes -
      +FastH3 is two distilled MiniMax-H3 checkpoints. **V1** is the four-step +launch. Some Hub repo names still say Preview. That name is historical. V1 is +a full model, not a demo. **V2** is the eight-step checkpoint. More forwards +is why V2 is the higher-quality FastH3. The V2 schedule contract is in +[FastH3 distilled checkpoint schedules](../inference/fasth3-distilled.md). +**CompactH3** is the 42-block 20B NVFP4 H3 checkpoint for one Blackwell GPU +(RTX 5090 or RTX PRO 6000). The FastH3 V1 and V2 recipes are unchanged. + +
      All model families
      @@ -15,8 +23,8 @@ hide:

      Primary focus · Inference

      MiniMax H3 recipes

      -

      Generate video and audio with H3. Run a server on CUDA, one DGX Spark, or Apple Silicon MLX to iterate on prompts, or call the pipeline directly from Python.

      - 9 maintained recipes +

      Generate video and audio with H3. Run a server on CUDA, one Blackwell GPU, one DGX Spark, or Apple Silicon MLX to iterate on prompts, or call the pipeline directly from Python.

      + 11 maintained recipes
      @@ -41,7 +49,7 @@ hide:

      Supported modes

      CUDA covers T2VA, FL2VA, and Ref2VA on the full checkpoint, plus FastH3 - V1 and FastH3 V2. FastH3 V1 also has a DGX Spark runtime with + V1 and FastH3 V2, plus CompactH3 NVFP4 on one Blackwell GPU. FastH3 V1 also has a DGX Spark runtime with a 1-Spark or 2-Spark device row. MLX is T2VA only: V1 and V2. Temporal --fast, spatial --fast-spatial, and opt-in VSA are flags on the same MLX script, not extra recipes. @@ -58,7 +66,7 @@ hide: T2VA - Full H3, FastH3 V1, FastH3 LoRA, FastH3 V2 + Full H3, FastH3 V1, FastH3 LoRA, FastH3 V2, CompactH3 NVFP4 FastH3 V1 or FastH3 V2 after a local DiT conversion @@ -96,6 +104,11 @@ hide: FastH3 V1 on one GB10, or two Sparks with Ray sequence parallel (sp_size=2) over QSFP RoCE. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. Not wired + + CompactH3 NVFP4 + 42-block 20B checkpoint on one RTX 5090 (32 GB, sequential encoder offload) or one RTX PRO 6000 Blackwell (96 GB, encoder+DiT+VAE resident). SageAttention3 FP4, packed NVFP4 DiT, Comfy int8-convrot VAE. + Not wired +

      @@ -104,7 +117,7 @@ hide:

      Pick an H3 recipe and runtime

      -

      Choose the result you want, then use a maintained CUDA, DGX Spark, or MLX path. +

      Choose the result you want, then use a maintained CUDA, Blackwell, DGX Spark, or MLX path. Device claims stay tied to checked-in sources and recorded runs.

      @@ -269,7 +282,8 @@ cd FastVideo
    • The MLX source runtime supports T2VA, optional temporal --fast, optional spatial --fast-spatial, and opt-in VSA on --include-vsa checkpoints. FastH3 V2 MLX converts with --include-vsa and runs eight forwards. FL2VA, Ref2VA, and two-pass refinement are not wired.
    • GPU count and VAE decode backend are configurable in the builder above for FastH3 CUDA recipes. Only the value shown by default has a recorded run; other supported values are unmeasured here.
    • DGX Spark is a runtime on FastH3 V1, not a separate family card. Select NVIDIA DGX Spark, then 1 Spark or 2 Sparks. The CUDA GPU-count knob does not apply to Spark.
      • -
    • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM.
      • +
    • CompactH3 NVFP4 is one Blackwell GPU. RTX 5090 (32 GB) parks the encoder in pinned host RAM. RTX PRO 6000 Blackwell (96 GB) keeps encoder, DiT, and VAE resident. Keep FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER, FASTVIDEO_FA4=0, FASTVIDEO_VSA_SM100A=0, and FLASHINFER_CUDA_ARCH_LIST=12.0a. On PRO 6000 enable VAE compile and leave DiT inference_torch_compile off. CompactH3 is a dense prune; do not enable VIDEO_SPARSE_ATTN_H3 until a VSA-trained student exists.
      • +
    • GB10 has no FA4 / sm_100a VSA kernel. Keep FASTVIDEO_FA4=0 and FASTVIDEO_VSA_SM100A=0. Legal num_frames values are 17n+5, capped at 362 (15.08 s). A 345-frame request on one Spark can OOM. Native 16:9 sizes include 832×480 and 1344×768.
    • Gated or missing checkpoints: run huggingface-cli login and confirm you accepted the model's license on Hugging Face.
    diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index a6e5654e7..544241d44 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index b2ccbc533..5fa15b299 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index 653b63677..845a30594 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index 860abee1c..df3dac64f 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index 9e5eb3e2e..7139d02fb 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
    +
    All model families
    diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index 7f39034dc..4db21f991 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
    +
    All model families
    diff --git a/docs/inference/fasth3_rtx_pro_6000.md b/docs/inference/fasth3_rtx_pro_6000.md new file mode 100644 index 000000000..9debdf043 --- /dev/null +++ b/docs/inference/fasth3_rtx_pro_6000.md @@ -0,0 +1,260 @@ +# FastH3 NVFP4 on RTX PRO 6000 (sm_120) + +This page covers serving the FastH3 NVFP4 checkpoints on one RTX PRO 6000 +Blackwell (sm_120, 96 GB) with every component resident: the NVFP4 text +encoder, the NVFP4 denoiser and an INT8 light VAE. It also records what was +measured and tried along the way. Every switch is opt-in and defaults to the +existing behavior. + +## Results + +All results are for one RTX PRO 6000 (Modal), a 10.1 s clip at 1344x768 +(243 frames, 73.6k packed tokens), the light INT8 VAE and warm runs. + +| Checkpoint | Denoise | Video decode | End to end | Peak memory | +| --- | ---: | ---: | ---: | ---: | +| [`FastH3-4-step-Preview-v1-VSA-DataFree-NVFP4`](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree-NVFP4) (4 forwards, VSA 0.9) | 30.1–31.0 s | 9.0 s | **41.2 / 42.4 s** | 79.6 GB | +| [`FastH3-8-Step-V2-NVFP4`](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4) (8 forwards, VSA 0.8) | 75.6 s | 9.0 s | **86.5 s** | — | + +These end-to-end runs predate the warp-skip kernel change below, which cuts +sparse attention by a further 1.5x, so they are upper bounds. Conditioning +takes about 0.12 s, frame post-processing plus MP4 writing about 0.9 s, and +the first clip at a new shape about 200 s (VAE compile). + +At 480p (124 frames, 15.1k tokens) V2 8-step denoises in 12.5 s on a cold +run, down from 16.8 s warm before this work. + +## Usage + +### 1. Convert the checkpoint + +The published NVFP4 checkpoints use ModelOpt's unified Hugging Face layout. +`convert_minimax_h3_modelopt_nvfp4_dit.py` repacks it into FastVideo's packed +export, `transformer/nvfp4_weights.safetensors`. The calibrated FFN weights +and scales are carried over bit for bit and only the scale bytes are +swizzled. Optionally it also quantizes the BF16 attention projections and the +VSA compression gates: + +```bash +python scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py \ + --src /models/FastH3-8-Step-V2-NVFP4/transformer \ + --dst /models/fasth3-v2-fv/transformer \ + --quantize-attention --quantize-gate +``` + +Every exported linear is probed through the loader's `mm_fp4` path against a +BF16 matmul with the dequantized weight. Conversion refuses to write when any +error exceeds `--max-probe-error` (default 0.3; the converted V2 and 4-step +checkpoints probe at 0.14). The rest of the model folder (text encoder, +VAEs, schedulers, `fastvideo_inference.json`) is used as-is; the +NVFP4 text encoder comes from +`convert_minimax_h3_text_encoder_nvfp4.py`. + +The flags select the `layer_profile` to load the result with: + +| Flags | Linears in NVFP4 | `layer_profile` | +| --- | --- | --- | +| *(none)* | FFN `fc_in`/`fc_out` | `h3_dit_ffn` | +| `--quantize-attention` | + attention `to_{q,k,v,out}` | `h3_dit` | +| `--quantize-attention --quantize-gate` | + VSA `to_gate_compress` | `h3_dit_vsa` | + +### 2. Generate + +```python +import os +os.environ.update({ + "FASTVIDEO_H3_VSA_FP4": "1", # sparse FP4 attention + "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", # Triton norm/modulate/RoPE/SwiGLU fusions + "FASTVIDEO_NVFP4_MM_BACKEND": "cutlass", # see "FP4 GEMM backend" below + "FASTVIDEO_H3_VAE_TILE_BATCH": "28", # one decoder call per 1344x768 tile grid +}) +from fastvideo import VideoGenerator + +generator = VideoGenerator.from_config({ + "model_path": "/models/fasth3-v2-fv", + "engine": { + "num_gpus": 1, + "quantization": {"transformer_quant": "NVFP4", "layer_profile": "h3_dit_vsa"}, + "compile": {"enabled": False, "vae_enabled": True}, + }, + "pipeline": {"experimental": {"attention_backend": "VIDEO_SPARSE_ATTN_H3", + "VSA_sparsity": 0.8, "VSA_tile_size": 64}}, +}) +generator.generate({"prompt": "...", "sampling": {"height": 768, "width": 1344, "num_frames": 243, + "num_inference_steps": 9}}) +``` + +Use `VSA_sparsity` 0.9 and `num_inference_steps` 5 for the 4-step checkpoint +(see its `fastvideo_inference.json`). Frame counts must be `17n + 5`: 243 +frames is the closest to 10 s. + +## What changed + +### Block-sparse FP4 attention for VSA tiles (`fastvideo-kernel`) + +SageAttention3's sm_120 FP4 kernel (`attn_qat_infer`) gains a block-sparse +forward, `fwd_sparse`, exposed as `sageattn_blackwell_sparse` (head-major +inputs) and `sageattn_blackwell_sparse_bshd` (sequence-major inputs, quantized +in place without a transpose). `vsa_tile_mask_to_fp4_blocks` turns a VSA tile +mask into the kernel's lists: + +- **Block lists.** Query block `m` visits only the 128-token KV blocks in + `q2k_idx[b, h, m, :q2k_num[b, h, m]]`. +- **Quadrant masks for 64-token tiles.** The kernel computes on 128x128 + blocks, but V2 and the 4-step preview use 64-token VSA tiles. Each listed + block carries a 4-bit `q2k_quad` (one bit per 64x64 quadrant). The kernel + masks unselected quadrants to `-inf`, so the result is exactly VSA's tile-64 + semantics. +- **Valid counts per 64-column half.** `kv_valid` gives the valid tokens in + each 64-column half, so partially filled tiles can pad mid-block. +- **Warp-level skipping.** Each MMA warp owns 16 query rows and so sits + inside one 64-row half. A warp skips a listed block that its half did not + select, and the P·V chunk of a key half it did not select. Masked scores + contribute exactly zero, and the warp sharing its tensor-core partition runs + faster meanwhile. This recovers most of the work that pairing 64-token tiles + into 128-token blocks adds. +- **First-visited block.** Lists run in descending block order because the + kernel visits them last entry first. Block 0 (the first prefix tile, which + VSA-H3's exempt mode gives every query) is therefore visited first, so every + row starts from a finite running max. `validate=True` checks this. The model + integration uses only exempt mode. + +The dense and sparse entry points also stop allocating `delta_s`. With Q +smoothing off (the default), each call used to allocate and zero a +`[B, H, L/128, L]` fp32 tensor: 9.5 GB at 73k tokens. Its int32 batch stride +also overflowed the TMA descriptor ("Failed to initialize the TMA descriptor +1", then an illegal instruction), so FP4 attention could not run 10 s 1344x768 +clips at all. A cached `[B, H, 1, L]` zero row read with `per_block_mean=False` +replaces it, and the outputs are bit-identical. + +Correctness (`fastvideo-kernel/tests/test_attn_qat_infer_sparse.py`, RTX PRO +6000): + +| Layout | Error vs token-masked fp32 | Dense FP4 floor | +| --- | ---: | ---: | +| 64-token tiles, odd count, partial tiles | 0.190 | 0.191 | +| 64-token tiles, even count | 0.191 | 0.193 | +| 256-token tiles, partial tile | 0.195 | 0.196 | + +Errors are relative L2. The sparse kernel sits exactly at the FP4 noise +floor; random Gaussian inputs make that floor large. Full block lists +reproduce the dense kernel bit for bit. + +### Model integration (`FASTVIDEO_H3_VSA_FP4=1`) + +`fastvideo/models/dits/minimax_h3_vsa_fp4.py` replaces only the attention core +of `MiniMaxH3Attention`. VSA-H3's tile pooling, top-k mask, exempt prefix and +gated compression branch are unchanged. Per block, it: + +1. Gathers the attention input into tile order once (one `hidden_size`-wide + pass). Pad rows stay zero, so the q/k/v pad rows are exactly zero through + the bias-free projections, RMSNorm and RoPE. +2. Quantizes that input once and shares it between `to_q`, `to_k` and `to_v`. + NVFP4 activations use a unit global scale, so this is exact. +3. Applies QK-norm and RoPE with tile-ordered `cos`/`sin`, computed once per + step. +4. Runs the sparse FP4 kernel on sequence-major tensors and gathers the output + back to packed order before `to_out`. + +This replaces the generic path's concat, four tile scatters and three +transposes. The route applies only to no-grad, non-compiled, single +sequence-parallel-rank calls in exempt mode; everything else keeps the +existing path. + +Two smaller pieces ship alongside it: + +- **Packed gate check.** With `--quantize-gate`, `to_gate_compress` loses its + BF16 weight, so the gate-activity check reads the packed E2M1 bytes instead. +- **Compiled residual.** With `FASTVIDEO_MINIMAX_H3_FUSIONS` enabled, each + block's final `hidden + gate[indices] * ffn_out` runs as one compiled op + instead of materializing the gathered gate. + +### FP4 GEMM backend (`FASTVIDEO_NVFP4_MM_BACKEND`) + +FlashInfer's `mm_fp4(backend="auto")` picks a kernel about 2x slower than +`cutlass` or `cudnn` on sm_120 once activations reach tens of thousands of +rows. At 15k rows all three match. + +| Linear | 73.6k rows: `auto` | 73.6k rows: `cutlass` | 73.6k rows: `cudnn` | 15.1k rows: `auto` | +| --- | ---: | ---: | ---: | ---: | +| `to_q` (5376→7168) | 7.96 ms | 3.96 ms | 4.24 ms | 0.88 ms | +| `to_out` (7168→5376) | 9.07 ms | 4.09 ms | 4.27 ms | 0.92 ms | +| `fc_in` (5376→28672) | 21.54 ms | 16.59 ms | 16.17 ms | 3.01 ms | +| `fc_out` (14336→5376) | 18.06 ms | 8.09 ms | 8.43 ms | 1.65 ms | + +### Batched VAE tile decode (`FASTVIDEO_H3_VAE_TILE_BATCH`) + +The H3 video VAE decodes 256-pixel spatial tiles one at a time. A 1344x768 +clip is a 4x7 grid per temporal chunk, so a 10 s clip is roughly 400 small +decoder calls. The ViT decoder treats batch entries independently, so +`FASTVIDEO_H3_VAE_TILE_BATCH=N` decodes up to `N` equal-shaped tiles per call; +28 covers a full 1344x768 grid. The decoded tiles are the same as per-tile +decoding. + +## Per-block measurements + +One H3 transformer block (hidden 5376, 56 heads, FFN 14336), RTX PRO 6000: + +| Component | 480p, 124 f (15.1k tokens) | 768p, 243 f (73.6k tokens) | +| --- | ---: | ---: | +| VSA Triton BF16 attention (kernel + pooling/mask) | 14.8 ms | 180.1 ms | +| Tile scatter of q/k/v/gate + gather (generic path) | 2.7 ms | 12.7 ms | +| Dense FP4 attention (SageAttention3) | 10.5 ms | 211.9 ms | +| Sparse FP4, quadrant masks | 7.9 ms | 123.4 ms | +| Sparse FP4, quadrant masks + warp skip | — | **83.0 ms** (VSA 0.8) / **49.3 ms** (VSA 0.9) | +| Dense BF16 SDPA | 18.1 ms | — | +| Modulation: eager / fused / compiled | 4.07 / 1.42 / 0.58 ms | 20.3 / 7.0 / 2.8 ms | +| SwiGLU: eager / fused | 1.53 / 0.88 ms | 7.37 / 4.24 ms | +| QK-norm + RoPE: eager / fused | 4.64 / 1.12 ms | 22.4 / 5.2 ms | + +Before this work, a 480p block cost about 40 ms: 17.7 ms of attention, 12 ms +of linears and 10 ms of eager elementwise ops. Over 50 blocks that is 2.0 s +per step, which matches the measured 2.1 s. + +Block density each kernel granularity computes at 768p (fraction of the +dense attention). "Selected" is what VSA needs; the other columns are what +each block shape computes: + +| VSA sparsity | Selected (64x64) | 128x128 blocks | 64-row x 128 | 128 x 64-col | +| --- | ---: | ---: | ---: | ---: | +| 0.8 | 0.222 | 0.434 | 0.317 | 0.313 | +| 0.9 | 0.125 | 0.254 | 0.180 | 0.177 | + +## What was tried and not shipped + +- **Dense FP4 attention for VSA-trained students.** `ATTN_QAT_INFER` does not + build `to_gate_compress`. A VSA-distilled checkpoint such as V2 carries + trained gates, so the strict loader refuses it ("Parameter + ...to_gate_compress.weight not found"). Dense FP4 attention is also slower + than sparse FP4 at 768p (212 vs 83 ms per block). +- **Multi-GPU (Ulysses) FP8 exchange.** On 8x RTX PRO 6000 (PCIe only), NCCL + all-to-all moves about 21 GB/s per GPU, with NCCL P2P on or off. A BF16 + q/k/v/gate exchange at 73.6k tokens therefore costs 24 ms per block, and the + attention output another 6.6 ms; with FP4 payloads q/k/v/gate drop to 7.2 ms. + The branch `h3-sm120-experimental` keeps a sequence-parallel path that: + - sends q/k/v as FP8 with one scale per token and head; + - never sends the VSA gate, applying it on each rank after a small + all-gather of the per-tile compression output; + - returns the attention output as FP8. + + It is estimated at about 18–20 s per 10 s clip for V2 8-step on 8 GPUs. It + has not executed yet (8-GPU capacity was unavailable), so it is not part of + this change. The same branch holds the Modal drivers behind every number on + this page. +- **64-row query blocks.** The kernel traits allow `kBlockM = 64`, which would + remove the query-side pairing waste. Warp-level skipping recovers most of + that waste without a second kernel instantiation, so it was not built. + +## Known limitations + +- **End-to-end quality.** The kernel matches a masked reference at the FP4 + noise floor. Generated videos have not yet been A/B-compared against the + BF16 Triton VSA path on the H3 audio/video metrics. +- **Activation scales.** The packed export drops ModelOpt's calibrated static + `input_scale` and quantizes activations with a unit global scale and dynamic + per-16 block scales, as FastVideo's NVFP4 linears do elsewhere. +- **Decode cost.** Video decode (9 s at 10 s/768p with the light INT8 VAE) is + the next largest cost after denoising. +- **Hardware.** Everything here targets sm_120. The GeForce RTX 5090 shares + the architecture but has 32 GB, which needs a reduced-AdaLN checkpoint and a + non-resident text encoder at this resolution. diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index da98a0313..adcfc730a 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -198,6 +198,11 @@ The `attn_qat_infer` kernel hard-gates on **sm_120 (consumer Blackwell / RTX 5090)**; on other GPUs the backend logs a notice and falls back to Flash Attention. See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040). +For VSA-distilled MiniMax-H3 students, the same kernel has a block-sparse +forward that runs VSA's 64-token tile selection in FP4 +(`FASTVIDEO_H3_VSA_FP4=1`). See +[FastH3 NVFP4 on RTX PRO 6000](fasth3_rtx_pro_6000.md). + Enable both halves — attention via the env var, linear via `transformer_quant`: ```python diff --git a/examples/inference/basic/basic_compacth3_rtx5090.yaml b/examples/inference/basic/basic_compacth3_rtx5090.yaml new file mode 100644 index 000000000..46ce251b8 --- /dev/null +++ b/examples/inference/basic/basic_compacth3_rtx5090.yaml @@ -0,0 +1,49 @@ +# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +# FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 \ +# fastvideo generate --config examples/inference/basic/basic_compacth3_rtx5090.yaml +generator: + model_path: ./CompactH3 + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit + 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 + compile: + enabled: false + vae_enabled: false + pipeline: + experimental: + attention_backend: ATTN_QAT_INFER + h3_sequential_load: true + inference_torch_compile: false + vae_parallel_decode: false + video_decode_backend: h3-vae +request: + prompt: >- + A wide cinematic shot of an alpine meadow at sunrise, pale pink mountain + peaks above a blue valley filled with thin morning mist. + negative_prompt: "" + sampling: + seed: 2026 + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 5 + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/compacth3_rtx5090/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml b/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml new file mode 100644 index 000000000..6c92ae7e8 --- /dev/null +++ b/examples/inference/basic/basic_compacth3_rtx_pro6000.yaml @@ -0,0 +1,49 @@ +# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +# FLASHINFER_CUDA_ARCH_LIST=12.0a FASTVIDEO_STAGE_LOGGING=1 \ +# fastvideo generate --config examples/inference/basic/basic_compacth3_rtx_pro6000.yaml +generator: + model_path: ./CompactH3 + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit + 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 + compile: + enabled: false + vae_enabled: true + pipeline: + experimental: + attention_backend: ATTN_QAT_INFER + h3_sequential_load: false + inference_torch_compile: false + vae_parallel_decode: false + video_decode_backend: h3-vae +request: + prompt: >- + A wide cinematic shot of an alpine meadow at sunrise, pale pink mountain + peaks above a blue valley filled with thin morning mist. + negative_prompt: "" + sampling: + seed: 2026 + height: 768 + width: 1344 + num_frames: 124 + fps: 24 + num_inference_steps: 5 + guidance_scale: 1.0 + batch_cfg: false + output: + output_path: outputs/compacth3_rtx_pro6000/ + save_video: true + return_frames: false diff --git a/examples/serving/openai_compacth3_rtx5090.yaml b/examples/serving/openai_compacth3_rtx5090.yaml new file mode 100644 index 000000000..136fc0af9 --- /dev/null +++ b/examples/serving/openai_compacth3_rtx5090.yaml @@ -0,0 +1,53 @@ +# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +# FLASHINFER_CUDA_ARCH_LIST=12.0a \ +# fastvideo serve --config examples/serving/openai_compacth3_rtx5090.yaml +generator: + model_path: ./CompactH3 + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: true + image_encoder: true + vae: false + pin_cpu_memory: true + lazy_module_load: false + compile: + enabled: false + vae_enabled: false + pipeline: + workload_type: t2v + experimental: + attention_backend: ATTN_QAT_INFER + h3_sequential_load: true + inference_torch_compile: false + vae_parallel_decode: false + video_decode_backend: h3-vae + +server: + host: 127.0.0.1 + port: 8000 + output_dir: outputs/openai_compacth3_rtx5090 + served_model_name: compacth3 + +default_request: + negative_prompt: "" + sampling: + height: 480 + width: 832 + num_frames: 124 + fps: 24 + num_inference_steps: 5 + guidance_scale: 1.0 + batch_cfg: false + seed: 2026 + output: + return_frames: false diff --git a/examples/serving/openai_compacth3_rtx_pro6000.yaml b/examples/serving/openai_compacth3_rtx_pro6000.yaml new file mode 100644 index 000000000..0f947feec --- /dev/null +++ b/examples/serving/openai_compacth3_rtx_pro6000.yaml @@ -0,0 +1,53 @@ +# FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +# FLASHINFER_CUDA_ARCH_LIST=12.0a \ +# fastvideo serve --config examples/serving/openai_compacth3_rtx_pro6000.yaml +generator: + model_path: ./CompactH3 + engine: + num_gpus: 1 + use_fsdp_inference: false + quantization: + transformer_quant: NVFP4 + layer_profile: h3_dit + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: false + image_encoder: true + vae: false + pin_cpu_memory: true + lazy_module_load: false + compile: + enabled: false + vae_enabled: true + pipeline: + workload_type: t2v + experimental: + attention_backend: ATTN_QAT_INFER + h3_sequential_load: false + inference_torch_compile: false + vae_parallel_decode: false + video_decode_backend: h3-vae + +server: + host: 127.0.0.1 + port: 8000 + output_dir: outputs/openai_compacth3_rtx_pro6000 + served_model_name: compacth3 + +default_request: + negative_prompt: "" + sampling: + height: 768 + width: 1344 + num_frames: 124 + fps: 24 + num_inference_steps: 5 + guidance_scale: 1.0 + batch_cfg: false + seed: 2026 + output: + return_frames: false diff --git a/fastvideo-kernel/attn_qat_infer/api.py b/fastvideo-kernel/attn_qat_infer/api.py index f506d692d..e91196a82 100644 --- a/fastvideo-kernel/attn_qat_infer/api.py +++ b/fastvideo-kernel/attn_qat_infer/api.py @@ -106,12 +106,32 @@ def preprocess_qkv(q: torch.Tensor, if enable_smoothing_q: delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous() else: # used to disable q smoothing - B, H, L, D = q.shape - delta_s = torch.zeros((B, H, L // BLOCK_M, k.shape[2]), device=q.device, dtype=torch.float32) + delta_s = _zero_delta_s(q.shape[0], q.shape[1], k.shape[2], q.device) return q, k, v, delta_s +_ZERO_DELTA_S: dict = {} + + +def _zero_delta_s(batch: int, heads: int, kv_len: int, device: torch.device) -> torch.Tensor: + """Cached all-zero delta_s for unsmoothed Q, read with per_block_mean=False. + + The kernel reads delta_s as a contiguous [B, H, rows, KL] tensor with one + row per query block when per_block_mean is set and a single row otherwise. + With Q smoothing off every row is zero, so one shared row replaces the + [B, H, L/128, KL] tensor that was allocated and zero-filled on every call + (9.5 GB at 73k tokens, whose int32 batch stride also broke the TMA + descriptor). The kernel only reads it. + """ + key = (batch, heads, kv_len, device) + zeros = _ZERO_DELTA_S.get(key) + if zeros is None: + zeros = torch.zeros((batch, heads, 1, kv_len), device=device, dtype=torch.float32) + _ZERO_DELTA_S[key] = zeros + return zeros + + def scale_and_quant_fp4(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert x.ndim == 4 B, H, N, D = x.shape @@ -154,6 +174,191 @@ def blockscaled_fp4_attn(qlist: Tuple, softmax_scale, is_causal, per_block_mean, is_bf16, single_level_p_quant) +def blockscaled_fp4_attn_sparse(qlist: Tuple, + klist: Tuple, + vlist: Tuple, + delta_s: torch.Tensor, + KL: int, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + q2k_quad: torch.Tensor | None = None, + per_block_mean: bool = True, + is_bf16: bool = True, + single_level_p_quant: bool = False, + sm_scale: float | None = None): + softmax_scale = sm_scale if sm_scale is not None else (qlist[0].shape[-1] * 2)**(-0.5) + return fp4attn_cuda.fwd_sparse(qlist[0], klist[0], vlist[0], qlist[1], klist[1], vlist[1], delta_s, KL, None, + softmax_scale, per_block_mean, is_bf16, single_level_p_quant, q2k_idx, q2k_num, + kv_valid, q2k_quad) + + +HALF_N = BLOCK_N // 2 + + +def vsa_tile_mask_to_fp4_blocks( + tile_mask: torch.Tensor, + tile_tokens: int, + tile_valid: torch.Tensor | None = None, + validate: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Convert a VSA tile mask into the FP4 kernel's block lists. + + ``tile_mask`` is a ``[B, H, T, T]`` bool mask over VSA tiles of + ``tile_tokens`` tokens (query tile, key tile); ``tile_tokens`` is 64 or a + multiple of 128. ``tile_valid`` holds each tile's valid token count (valid + tokens first, padding last, as VSA tiling lays them out). The kernel works + on 128x128 blocks, so 64-token tiles are paired: a block is listed when any + of its four 64x64 quadrants is selected, and ``q2k_quad`` carries which. + + Lists run in descending block order because the kernel visits them last + entry first: block 0 (the first prefix tile, which VSA-H3's exempt mode + gives every query) is then visited first, so every query row starts from a + finite running max. ``validate=True`` checks that (host sync). + + Returns ``(q2k_idx [B, H, M, N], q2k_num [B, H, M], kv_valid [2N], + q2k_quad [B, H, M, N] | None)``; the queries cover ``M * 128`` padded rows. + """ + if tile_tokens % HALF_N or (tile_tokens > HALF_N and tile_tokens % BLOCK_N): + raise ValueError(f"tile_tokens={tile_tokens} must be {HALF_N} or a multiple of {BLOCK_N}") + batch, heads, n_tiles, _ = tile_mask.shape + device = tile_mask.device + halves_per_tile = tile_tokens // HALF_N + half_mask = tile_mask + if halves_per_tile > 1: + half_mask = half_mask.repeat_interleave(halves_per_tile, dim=2).repeat_interleave(halves_per_tile, dim=3) + n_halves = n_tiles * halves_per_tile + if tile_valid is None: + half_valid = torch.full((n_halves, ), HALF_N, device=device, dtype=torch.int32) + else: + offsets = torch.arange(halves_per_tile, device=device, dtype=torch.int32) * HALF_N + half_valid = (tile_valid.to(torch.int32)[:, None] - offsets[None, :]).clamp(0, HALF_N).reshape(-1) + if n_halves % 2: + # Pad to whole 128-blocks: the extra key half is empty, the extra query + # half (discarded output) attends block 0 only. + half_mask = F.pad(half_mask, (0, 1, 0, 1), value=False) + half_mask[:, :, -1, 0] = True + half_valid = F.pad(half_valid, (0, 1), value=0) + n_halves += 1 + half_mask = half_mask & (half_valid > 0)[None, None, None, :] + n_blocks = n_halves // 2 + quads = half_mask.view(batch, heads, n_blocks, 2, n_blocks, 2) + weights = torch.tensor([[1, 2], [4, 8]], device=device, dtype=torch.uint8) # [row_half, col_half] + quad = (quads.to(torch.uint8) * weights[None, None, None, :, None, :]).sum(dim=(3, 5), dtype=torch.uint8) + block_mask = quad != 0 + # Descending compaction without a sort: position = running count from the right. + 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=device, dtype=torch.int32).expand_as(pos) + slot = torch.where(rev, pos, torch.full_like(pos, n_blocks)).long() + q2k_idx = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.int32) + q2k_idx.scatter_(-1, slot, cols) + q2k_idx = q2k_idx[..., :n_blocks].contiguous() + kv_valid = half_valid.contiguous() + q2k_quad = None + if halves_per_tile == 1: + q2k_quad = torch.zeros((batch, heads, n_blocks, n_blocks + 1), device=device, dtype=torch.uint8) + q2k_quad.scatter_(-1, slot, quad.flip(-1)) + q2k_quad = q2k_quad[..., :n_blocks].contiguous() + if validate: + if int(q2k_num.min()) < 1: + raise ValueError("every query block must attend to at least one non-empty KV block") + last = (q2k_num - 1).long().unsqueeze(-1) + first_block = q2k_idx.gather(-1, last.int().long()).squeeze(-1).long() + first_quad = (q2k_quad.gather(-1, last).squeeze(-1).int() if q2k_quad is not None else + torch.full_like(first_block, 15, dtype=torch.int32)) + v0 = (kv_valid[2 * first_block] > 0).int() + v1 = (kv_valid[2 * first_block + 1] > 0).int() + row0 = ((first_quad & 1).bool() & v0.bool()) | ((first_quad & 2).bool() & v1.bool()) + row1 = ((first_quad & 4).bool() & v0.bool()) | ((first_quad & 8).bool() & v1.bool()) + if not bool((row0 & row1).all()): + raise ValueError("the first block each query block visits must give both 64-row halves a valid key") + return q2k_idx, q2k_num, kv_valid, q2k_quad + + +def check_sparse_block_lists(q2k_idx: torch.Tensor, q2k_num: torch.Tensor, kv_len: int) -> None: + """Reject block lists the sparse kernel would read out of bounds (host sync). + + Each row needs ``1 <= q2k_num <= q2k_idx.size(-1)`` and every listed index + in ``[0, ceil(kv_len / BLOCK_N))``; a zero count or an out-of-range index + makes the kernel load outside its index row or the KV tensors. + """ + num_kv_blocks = -(-kv_len // BLOCK_N) + if int(q2k_num.min()) < 1 or int(q2k_num.max()) > q2k_idx.size(-1): + raise ValueError(f"q2k_num must be in [1, {q2k_idx.size(-1)}]") + listed = torch.arange(q2k_idx.size(-1), device=q2k_idx.device) < q2k_num.unsqueeze(-1) + idx = q2k_idx[listed] + if idx.numel() and (int(idx.min()) < 0 or int(idx.max()) >= num_kv_blocks): + raise ValueError(f"q2k_idx entries must be in [0, {num_kv_blocks})") + + +def sageattn_blackwell_sparse(q, + k, + v, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + q2k_quad: torch.Tensor | None = None, + per_block_mean=True, + single_level_p_quant=True, + sm_scale: float | None = None, + validate: bool = True): + """Block-sparse SageAttention3 FP4 forward (non-causal). + + Query block ``m`` (``BLOCK_M`` rows) of each (batch, head) attends only to + the ``BLOCK_N``-token KV blocks in ``q2k_idx[b, h, m, :q2k_num[b, h, m]]``, + restricted to the quadrants in ``q2k_quad`` when given; see + :func:`vsa_tile_mask_to_fp4_blocks`. Q/K/V are ``[B, H, L, D]``. + Block lists are checked with :func:`check_sparse_block_lists` (a host + sync) unless ``validate=False``; pass that only for lists built by + :func:`vsa_tile_mask_to_fp4_blocks`, which are in range by construction. + """ + QL = q.size(2) + KL = k.size(2) + if validate: + check_sparse_block_lists(q2k_idx, q2k_num, KL) + is_bf16 = q.dtype == torch.bfloat16 + q, k, v, delta_s = preprocess_qkv(q, k, v, per_block_mean) + per_block_mean = delta_s.shape[2] > 1 + qlist_from_cuda = scale_and_quant_fp4(q) + klist_from_cuda = scale_and_quant_fp4_permute(k) + vlist_from_cuda = scale_and_quant_fp4_transpose(v) + o_fp4 = blockscaled_fp4_attn_sparse(qlist_from_cuda, klist_from_cuda, vlist_from_cuda, delta_s, KL, q2k_idx, + q2k_num, kv_valid, q2k_quad, per_block_mean, is_bf16, single_level_p_quant, + sm_scale)[0][:, :, :QL, :].contiguous() + return o_fp4 + + +def sageattn_blackwell_sparse_bshd(q, + k, + v, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + kv_valid: torch.Tensor | None = None, + q2k_quad: torch.Tensor | None = None, + single_level_p_quant=True, + sm_scale: float | None = None, + validate: bool = True) -> torch.Tensor: + """:func:`sageattn_blackwell_sparse` for ``[B, L, H, D]`` inputs, without copies. + + The FP4 quantizers read strided input, so the sequence-major tensors a + linear produces are quantized in place of a transpose + pad. ``L`` must be + a multiple of ``BLOCK_M`` (callers allocate the padding) and Q is + 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}") + if validate: + check_sparse_block_lists(q2k_idx, q2k_num, seq_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) + 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, + q2k_quad, False, q.dtype == torch.bfloat16, single_level_p_quant, sm_scale)[0] + + def sageattn_blackwell(q, k, v, @@ -191,6 +396,7 @@ def sageattn_blackwell(q, KL = k.size(2) is_bf16 = q.dtype == torch.bfloat16 q, k, v, delta_s = preprocess_qkv(q, k, v, per_block_mean) + per_block_mean = delta_s.shape[2] > 1 qlist_from_cuda = scale_and_quant_fp4(q) klist_from_cuda = scale_and_quant_fp4_permute(k) vlist_from_cuda = scale_and_quant_fp4_transpose(v) diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/api.cu b/fastvideo-kernel/attn_qat_infer/blackwell/api.cu index 243d876db..c7a3624f0 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/api.cu +++ b/fastvideo-kernel/attn_qat_infer/blackwell/api.cu @@ -204,8 +204,17 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split })); } -std::vector -mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head_size // 2) +struct SparseKvLists { + int const *q2k_idx = nullptr; + int const *q2k_num = nullptr; + int q2k_max = 0; + int num_m_blocks = 0; + int const *kv_valid = nullptr; + uint8_t const *q2k_quad = nullptr; +}; + +static std::vector +mha_fwd_impl(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head_size // 2) const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x (head_size // 2) const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x (head_size // 2) const at::Tensor &sfq, @@ -218,7 +227,8 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head_size bool is_causal, bool per_block_mean, bool is_bf16, - bool single_level_p_quant=false // If true, use only per-row scale s_P2 (no per-block s_P1) + bool single_level_p_quant, // If true, use only per-row scale s_P2 (no per-block s_P1) + SparseKvLists const &sparse ) { auto dprops = at::cuda::getCurrentDeviceProperties(); @@ -316,6 +326,12 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head_size // stack-local tensor whose data pointer would dangle after mha_fwd returns // while the async kernel may still be running. params.tile_count_semaphore = nullptr; + params.q2k_idx = sparse.q2k_idx; + params.q2k_num = sparse.q2k_num; + params.q2k_max = sparse.q2k_max; + params.num_m_blocks = sparse.num_m_blocks; + params.kv_valid = sparse.kv_valid; + params.q2k_quad = sparse.q2k_quad; if (seqlen_k > 0) { auto stream = at::cuda::getCurrentCUDAStream().stream(); @@ -341,7 +357,72 @@ mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head_size +std::vector +mha_fwd(at::Tensor &q, const at::Tensor &k, const at::Tensor &v, + const at::Tensor &sfq, const at::Tensor &sfk, const at::Tensor &sfv, + const at::Tensor &delta_s, int unpadded_k, c10::optional &out_, + const float softmax_scale, bool is_causal, bool per_block_mean, bool is_bf16, + bool single_level_p_quant=false) { + return mha_fwd_impl(q, k, v, sfq, sfk, sfv, delta_s, unpadded_k, out_, softmax_scale, + is_causal, per_block_mean, is_bf16, single_level_p_quant, SparseKvLists{}); +} + +// Block-sparse forward: query block m of (b, h) attends only to the KV blocks +// listed in q2k_idx[b, h, m, :q2k_num[b, h, m]] (BLOCK_M x BLOCK_N granularity). +// kv_valid, if given, holds the valid token count of each 64-column half of +// every KV block ([2 * num_kv_blocks], valid tokens first within a half). +// q2k_quad, if given (uint8, same shape as q2k_idx), restricts each listed +// block to the 64x64 quadrants whose bit (2 * row_half + col_half) is set, so +// 64-token VSA tiles run on 128x128 blocks. The block a list visits first +// (its last entry) must leave every query row at least one valid key. +// Non-causal only. +std::vector +mha_fwd_sparse(at::Tensor &q, const at::Tensor &k, const at::Tensor &v, + const at::Tensor &sfq, const at::Tensor &sfk, const at::Tensor &sfv, + const at::Tensor &delta_s, int unpadded_k, c10::optional &out_, + const float softmax_scale, bool per_block_mean, bool is_bf16, + bool single_level_p_quant, + const at::Tensor &q2k_idx, const at::Tensor &q2k_num, + c10::optional &kv_valid_, + c10::optional &q2k_quad_) { + const int batch_size = q.size(0); + const int num_heads = q.size(1); + const int num_m_blocks = (q.size(2) + flash::BLOCK_M - 1) / flash::BLOCK_M; + const int num_n_blocks = (k.size(2) + flash::BLOCK_N - 1) / flash::BLOCK_N; + for (auto const *t : {&q2k_idx, &q2k_num}) { + TORCH_CHECK(t->scalar_type() == torch::kInt32, "q2k_idx / q2k_num must be int32"); + CHECK_DEVICE((*t)); CHECK_CONTIGUOUS((*t)); + } + TORCH_CHECK(q2k_idx.dim() == 4, "q2k_idx must be [batch, heads, num_m_blocks, max_kv_blocks]"); + TORCH_CHECK(q2k_idx.size(0) == batch_size && q2k_idx.size(1) == num_heads && q2k_idx.size(2) == num_m_blocks, + "q2k_idx leading dims must be [batch, heads, ceil(seqlen_q / BLOCK_M)]"); + TORCH_CHECK(q2k_idx.size(3) >= 1 && q2k_idx.size(3) <= num_n_blocks, "q2k_idx last dim must be in [1, num_kv_blocks]"); + CHECK_SHAPE(q2k_num, batch_size, num_heads, num_m_blocks); + SparseKvLists sparse; + sparse.q2k_idx = q2k_idx.data_ptr(); + sparse.q2k_num = q2k_num.data_ptr(); + sparse.q2k_max = q2k_idx.size(3); + sparse.num_m_blocks = num_m_blocks; + if (kv_valid_.has_value()) { + auto const &kv_valid = kv_valid_.value(); + TORCH_CHECK(kv_valid.scalar_type() == torch::kInt32, "kv_valid must be int32"); + CHECK_DEVICE(kv_valid); CHECK_CONTIGUOUS(kv_valid); + CHECK_SHAPE(kv_valid, 2 * num_n_blocks); + sparse.kv_valid = kv_valid.data_ptr(); + } + if (q2k_quad_.has_value()) { + auto const &q2k_quad = q2k_quad_.value(); + TORCH_CHECK(q2k_quad.scalar_type() == torch::kUInt8, "q2k_quad must be uint8"); + CHECK_DEVICE(q2k_quad); CHECK_CONTIGUOUS(q2k_quad); + TORCH_CHECK(q2k_quad.sizes() == q2k_idx.sizes(), "q2k_quad must match q2k_idx's shape"); + sparse.q2k_quad = q2k_quad.data_ptr(); + } + return mha_fwd_impl(q, k, v, sfq, sfk, sfv, delta_s, unpadded_k, out_, softmax_scale, + /*is_causal=*/false, per_block_mean, is_bf16, single_level_p_quant, sparse); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "FlashAttention"; m.def("fwd", &mha_fwd, "Forward pass"); + m.def("fwd_sparse", &mha_fwd_sparse, "Block-sparse forward pass (non-causal)"); } \ No newline at end of file diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h index 675608726..d237499c3 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h @@ -185,14 +185,14 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, auto block_coord = work_tile_info.get_block_coord(scheduler_params); auto [m_block, bidh, bidb] = block_coord; - int n_block_max = collective_mainloop.get_n_block_max(mainloop_params, m_block); - if (Is_causal && n_block_max <= 0) { // We exit early and write 0 to gO and -inf to gLSE. + int n_block_count = collective_mainloop.get_n_block_count(mainloop_params, m_block, bidh, bidb); + if (Is_causal && n_block_count <= 0) { // We exit early and write 0 to gO and -inf to gLSE. collective_epilogue.store_zero(epilogue_params, threadIdx.x - NumCopyThreads, block_coord); continue; } collective_mainloop.mma(mainloop_params, pipeline_q, pipeline_k, pipeline_v, smem_pipe_read_q, smem_pipe_read_k, smem_pipe_read_v, - tOrO, softmax_fused, n_block_max, threadIdx.x - NumCopyThreads, work_idx, m_block, shared_storage); + tOrO, softmax_fused, n_block_count, threadIdx.x - NumCopyThreads, work_idx, m_block, bidh, bidb, shared_storage); barrier_o.wait(); collective_epilogue.mma_store(shared_storage, tiled_mma_pv, tOrO, threadIdx.x - NumCopyThreads); barrier_o.arrive(); diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/launch.h b/fastvideo-kernel/attn_qat_infer/blackwell/launch.h index 91398c5f4..edc167afc 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/launch.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/launch.h @@ -62,7 +62,9 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { static_cast(params.delta_s_ptr), {params.seqlen_s, params.seqlen_k, params.h_k, params.b}, {params.ds_row_stride, _1{}, params.ds_head_stride, params.ds_batch_stride}, - params.scale_softmax_log2 + params.scale_softmax_log2, + params.q2k_idx, params.q2k_num, params.q2k_max, + params.num_m_blocks, params.h, params.kv_valid, params.q2k_quad }); typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments({ diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h index fa2240422..51d9ff23a 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h @@ -174,6 +174,13 @@ struct CollectiveMainloopFwd { ShapeQKV const shape_ds; StrideQKV const stride_ds; float const softmax_scale_log2; + int const* ptr_q2k_idx{nullptr}; + int const* ptr_q2k_num{nullptr}; + int q2k_max{0}; + int num_m_blocks{0}; + int num_heads{0}; + int const* ptr_kv_valid{nullptr}; + uint8_t const* ptr_q2k_quad{nullptr}; }; // Device side kernel params @@ -194,6 +201,13 @@ struct CollectiveMainloopFwd { TMA_SFVt tma_load_SFVt; TMA_DS tma_load_DS; float const softmax_scale_log2; + int const* ptr_q2k_idx; + int const* ptr_q2k_num; + int q2k_max; + int num_m_blocks; + int num_heads; + int const* ptr_kv_valid; + uint8_t const* ptr_q2k_quad; }; @@ -261,7 +275,9 @@ struct CollectiveMainloopFwd { tma_load_K, tma_load_sfk, tma_load_Vt, tma_load_sfvt, tma_load_ds, - args.softmax_scale_log2}; + args.softmax_scale_log2, + args.ptr_q2k_idx, args.ptr_q2k_num, args.q2k_max, + args.num_m_blocks, args.num_heads, args.ptr_kv_valid, args.ptr_q2k_quad}; } /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance @@ -290,6 +306,51 @@ struct CollectiveMainloopFwd { return n_block_max; } + // Number of KV blocks query block m_block visits: all of them when dense, + // its index-list length when block-sparse. + CUTLASS_DEVICE + int get_n_block_count(Params const& mainloop_params, int m_block, int bidh, int bidb) { + if (mainloop_params.ptr_q2k_idx == nullptr) { + return get_n_block_max(mainloop_params, m_block); + } + return mainloop_params.ptr_q2k_num[(bidb * mainloop_params.num_heads + bidh) * mainloop_params.num_m_blocks + m_block]; + } + + // KV block visited at iteration i (iterations run from count-1 down to 0). + CUTLASS_DEVICE + int get_kv_block(Params const& mainloop_params, int m_block, int bidh, int bidb, int i) { + if (mainloop_params.ptr_q2k_idx == nullptr) { + return i; + } + int64_t const row = (int64_t(bidb) * mainloop_params.num_heads + bidh) * mainloop_params.num_m_blocks + m_block; + return mainloop_params.ptr_q2k_idx[row * mainloop_params.q2k_max + i]; + } + + // Valid key columns of each 64-column half of KV block n_block (valid tokens + // first within a half). kv_valid, when given, stores two counts per block so + // 64-token VSA tiles can pad mid-block; otherwise only the sequence tail is + // masked. + CUTLASS_DEVICE + int get_kv_valid_half(Params const& mainloop_params, int n_block, int half, int unpadded_seqlen_k) { + static constexpr int kBlockN = get<1>(TileShape_MNK{}); + static constexpr int kHalfN = kBlockN / 2; + if (mainloop_params.ptr_kv_valid != nullptr) { + return mainloop_params.ptr_kv_valid[2 * n_block + half]; + } + return max(0, min(kHalfN, unpadded_seqlen_k - n_block * kBlockN - half * kHalfN)); + } + + // Quadrant mask of list entry i: bit (2 * row_half + col_half) is set when the + // 64-row query half attends the 64-column key half. 0xF when absent. + CUTLASS_DEVICE + int get_quad(Params const& mainloop_params, int m_block, int bidh, int bidb, int i) { + if (mainloop_params.ptr_q2k_quad == nullptr) { + return 0xF; + } + int64_t const row = (int64_t(bidb) * mainloop_params.num_heads + bidh) * mainloop_params.num_m_blocks + m_block; + return mainloop_params.ptr_q2k_quad[row * mainloop_params.q2k_max + i]; + } + template CUTE_HOST_DEVICE constexpr auto @@ -450,7 +511,7 @@ struct CollectiveMainloopFwd { auto [m_block, bidh, bidb] = work_tile_info.get_block_coord(scheduler_params); - int n_block_max = get_n_block_max(mainloop_params, m_block); + int n_block_count = get_n_block_count(mainloop_params, m_block, bidh, bidb); Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.begin()), SmemLayoutQ{}); Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.begin()), SmemLayoutK{}); @@ -506,7 +567,8 @@ struct CollectiveMainloopFwd { Tensor tDSsDS = group_modes<0, 3>(block_tma_ds.partition_D(sDS)); uint16_t mcast_mask_kv = 0; - int n_block = n_block_max - 1; + int n_iter = n_block_count - 1; + int n_block = get_kv_block(mainloop_params, m_block, bidh, bidb, n_iter); int lane_predicate = cute::elect_one_sync(); if (lane_predicate) { pipeline_q.producer_acquire(smem_pipe_write_q); @@ -529,11 +591,12 @@ struct CollectiveMainloopFwd { ++smem_pipe_write_v; } - n_block--; + --n_iter; if (lane_predicate) { // CUTLASS_PRAGMA_NO_UNROLL #pragma unroll 2 - for (; n_block >= 0; --n_block) { + for (; n_iter >= 0; --n_iter) { + n_block = get_kv_block(mainloop_params, m_block, bidh, bidb, n_iter); pipeline_k.producer_acquire(smem_pipe_write_k); copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), mcast_mask_kv), tKgK(_, n_block), tKsK(_, smem_pipe_write_k.index())); @@ -585,6 +648,8 @@ struct CollectiveMainloopFwd { int thread_idx, int work_idx, int m_block, + int bidh, + int bidb, SharedStorage& shared_storage ) { @@ -667,7 +732,21 @@ struct CollectiveMainloopFwd { int const seqlen_q = get<0>(mainloop_params.shape_Q); int const seqlen_k = get<0>(mainloop_params.shape_K); int const unpadded_seqlen_k = get<0>(mainloop_params.unpadded_shape_K); - int n_block = n_block_count - 1; + int n_iter = n_block_count - 1; + int n_block = get_kv_block(mainloop_params, m_block, bidh, bidb, n_iter); + bool const per_block_masking = mainloop_params.ptr_q2k_idx != nullptr || mainloop_params.ptr_kv_valid != nullptr; + static_assert(kBlockM == 128 && kBlockN == 128, "quadrant masking assumes 128x128 blocks"); + // Each MMA warp owns 16 consecutive query rows, so a warp lies in one + // 64-row half. With quadrant lists a warp skips blocks its half did not + // select and the P.V chunk of a key half it did not select: masked + // scores contribute exactly zero, and the warp sharing its tensor-core + // partition runs faster meanwhile. + int const my_row_half = __shfl_sync(0xffffffff, [&] { + Tensor cS0 = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{})); + Tensor tScS0 = thread_mma_qk.partition_C(cS0); + return int(get<0>(tScS0(0))) >= kBlockM / 2 ? 1 : 0; + }(), 0); + bool const quad_skip = mainloop_params.ptr_q2k_quad != nullptr; auto copy_k_block = [&](auto block_id) { auto tSsK_stage = tSsK(_, _, _, smem_pipe_read_k.index()); @@ -744,6 +823,26 @@ struct CollectiveMainloopFwd { int local = c & 31; return (c & ~31) | (local & 1) | ((local & 24) >> 2) | ((local & 6) << 2); }; + // Sparse lists may visit a block whose 64-token halves are partially valid + // (VSA tile tails, short prefix chunks) or that only one query half + // selected (64-token tiles), at any iteration. + auto apply_sparse_mask = [&](auto& acc, int n_blk, int it) { + int const quad = get_quad(mainloop_params, m_block, bidh, bidb, it); + int const valid0 = get_kv_valid_half(mainloop_params, n_blk, 0, unpadded_seqlen_k); + int const valid1 = get_kv_valid_half(mainloop_params, n_blk, 1, unpadded_seqlen_k); + if (quad == 0xF && valid0 == kBlockN / 2 && valid1 == kBlockN / 2) { return; } + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{})); + Tensor tScS = thread_mma_qk.partition_C(cS); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc); ++i) { + int const col = actual_col(int(get<1>(tScS(i)))); + int const col_half = col >= kBlockN / 2; + int const row_half = int(get<0>(tScS(i))) >= kBlockM / 2; + bool const keep = ((quad >> (2 * row_half + col_half)) & 1) + && (col & (kBlockN / 2 - 1)) < (col_half ? valid1 : valid0); + if (!keep) { acc(i) = -INFINITY; } + } + }; { Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_MNK{})); Tensor tScS = thread_mma_qk.partition_C(cS); @@ -751,7 +850,9 @@ struct CollectiveMainloopFwd { for (int i = 0; i < size(tSrS); ++i) { int col = actual_col(int(get<1>(tScS(i)))); if constexpr (!Is_causal) { // Just masking based on col - if (col >= int(unpadded_seqlen_k - n_block * kBlockN)) { tSrS(i) = -INFINITY; } + if (!per_block_masking) { + if (col >= int(unpadded_seqlen_k - n_block * kBlockN)) { tSrS(i) = -INFINITY; } + } } else { if (col >= std::min(seqlen_k - n_block * kBlockN, col_limit_causal(int(get<0>(tScS(i))), n_block))) { @@ -760,6 +861,9 @@ struct CollectiveMainloopFwd { } } } + if constexpr (!Is_causal) { + if (per_block_masking) { apply_sparse_mask(tSrS, n_block, n_iter); } + } auto quantize = [&](auto mma_k, auto acc_conversion_view) { Tensor AbsMaxP_stagek = AbsMaxP(_, make_coord(_, _, mma_k)); Tensor acc_conversion_stagek = acc_conversion_view(_, _, mma_k); @@ -827,11 +931,12 @@ struct CollectiveMainloopFwd { } } - n_block--; + --n_iter; constexpr int n_masking_steps = !Is_causal ? 1 : cute::ceil_div(kBlockM, kBlockN) + 1; // // Only go through these if Is_causal, since n_masking_steps = 1 when !Is_causal CUTLASS_PRAGMA_UNROLL - for (int masking_step = 0; masking_step < n_masking_steps - 1 && n_block >= 0; ++masking_step, --n_block) { + for (int masking_step = 0; masking_step < n_masking_steps - 1 && n_iter >= 0; ++masking_step, --n_iter) { + n_block = get_kv_block(mainloop_params, m_block, bidh, bidb, n_iter); Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_MNK{})); Tensor tSrS_converion_view = make_tensor(tSrS.data(), flash::convert_to_conversion_layout(tSrS.layout())); consumer_wait(pipeline_k, smem_pipe_read_k); @@ -878,7 +983,22 @@ struct CollectiveMainloopFwd { } #pragma unroll 1 - for (; n_block >= 0; --n_block) { + for (; n_iter >= 0; --n_iter) { + n_block = get_kv_block(mainloop_params, m_block, bidh, bidb, n_iter); + int const row_bits = quad_skip + ? (get_quad(mainloop_params, m_block, bidh, bidb, n_iter) >> (2 * my_row_half)) & 3 + : 3; + if (row_bits == 0) { + // Neither key half is selected for this warp's rows: its softmax + // state and output are unchanged. Keep the pipeline in step. + consumer_wait(pipeline_k, smem_pipe_read_k); + pipeline_k.consumer_release(smem_pipe_read_k); + ++smem_pipe_read_k; + consumer_wait(pipeline_v, smem_pipe_read_v); + pipeline_v.consumer_release(smem_pipe_read_v); + ++smem_pipe_read_v; + continue; + } Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_MNK{})); Tensor tSrS_converion_view = make_tensor(tSrS.data(), flash::convert_to_conversion_layout(tSrS.layout())); consumer_wait(pipeline_k, smem_pipe_read_k); @@ -896,16 +1016,21 @@ struct CollectiveMainloopFwd { } } + if (per_block_masking) { apply_sparse_mask(tSrS, n_block, n_iter); } softmax_fused.template online_softmax_with_quant(tSrS, AbsMaxP, mainloop_params.softmax_scale_log2); Tensor tOrO = make_fragment_like(tOrO_store); + clear(tOrO); consumer_wait(pipeline_v, smem_pipe_read_v); copy_v_block(_0{}); quantize(_0{}, tSrS_converion_view); CUTLASS_PRAGMA_UNROLL for (int v_block = 0; v_block < size<2>(tOrP); ++v_block) { - cute::gemm(tiled_mma_pv, make_zip_tensor(tOrP(_, _, v_block), tOrSFP(_, _, v_block)), - make_zip_tensor(tOrVt(_, _, v_block), tOrSFVt(_, _, v_block)), tOrO); + // v_block spans one 64-column key half (P's K mode is 2 x 64). + if ((row_bits >> v_block) & 1) { + cute::gemm(tiled_mma_pv, make_zip_tensor(tOrP(_, _, v_block), tOrSFP(_, _, v_block)), + make_zip_tensor(tOrVt(_, _, v_block), tOrSFVt(_, _, v_block)), tOrO); + } if (v_block < size<2>(tOrP) - 1) { copy_v_block(v_block + 1); quantize(v_block + 1, tSrS_converion_view); diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/params.h b/fastvideo-kernel/attn_qat_infer/blackwell/params.h index fcf9bf898..867aa7b25 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/params.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/params.h @@ -114,6 +114,23 @@ int * __restrict__ seqused_k; int *__restrict__ blockmask; + + // Block-sparse KV iteration (non-causal only). When q2k_idx is null the + // kernel is dense. Otherwise query block m of (batch b, head h) visits + // the q2k_num[(b*h + h)*num_m_blocks + m] KV blocks listed at + // q2k_idx[((b*h + h)*num_m_blocks + m)*q2k_max + i], in any order. + int const *__restrict__ q2k_idx; + int const *__restrict__ q2k_num; + int q2k_max; + int num_m_blocks; + // Optional valid token count of each 64-column half of every KV block, + // kv_valid[2*n + half] (valid tokens first within a half); null means only + // the sequence tail beyond the unpadded key length is masked. + int const *__restrict__ kv_valid; + // Optional quadrant mask per list entry (same layout as q2k_idx): bit + // (2*row_half + col_half) set when that 64-row query half attends that + // 64-column key half. Lets 64-token VSA tiles run on 128x128 blocks. + uint8_t const *__restrict__ q2k_quad; // The K_new and V_new matrices. void * __restrict__ knew_ptr; diff --git a/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py b/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py new file mode 100644 index 000000000..3684b2156 --- /dev/null +++ b/fastvideo-kernel/tests/test_attn_qat_infer_sparse.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Correctness tests for the block-sparse SageAttention3 FP4 forward (``fwd_sparse``). + +Covers the VSA tile layouts MiniMax-H3 uses: 64-token tiles carried by 64x64 +quadrant masks on the kernel's 128x128 blocks (odd and even tile counts), +256-token tiles, partially valid tiles, and an exempt prefix that every query +attends. The sparse kernel is compared with a token-masked fp32 reference and +must stay at the dense FP4 kernel's own error on the same data. Also checks +that full block lists reproduce the dense kernel bit for bit and that the +sequence-major entry point matches the head-major one bit for bit. + +Requires a Blackwell GPU (sm_120a) and the fp4attn_cuda / fp4quant_cuda +extensions built via ``cd fastvideo-kernel && ./build.sh``. + + pytest tests/test_attn_qat_infer_sparse.py -v +""" + +import math +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +import pytest +import torch +import torch.nn.functional as F + +pytest.importorskip("fp4attn_cuda", reason="ATTN_QAT_INFER FP4 kernels require a sm_120a build") + +from attn_qat_infer.api import (BLOCK_M, BLOCK_N, check_sparse_block_lists, sageattn_blackwell, + sageattn_blackwell_sparse, sageattn_blackwell_sparse_bshd, + vsa_tile_mask_to_fp4_blocks) + +DEVICE = torch.device("cuda") +HEAD_DIM = 128 + + +def _tile_sizes(prefix: tuple[int, ...], video_tiles: int, tile: int, partial: dict[int, int]) -> torch.Tensor: + """Valid tokens per tile: segment-pure prefix chunks, then video tiles (some partial).""" + sizes = [] + for segment in prefix: + full, rem = divmod(segment, tile) + sizes += [tile] * full + ([rem] if rem else []) + sizes += [partial.get(index, tile) for index in range(video_tiles)] + return torch.tensor(sizes, dtype=torch.int32, device=DEVICE) + + +def _exempt_mask(n_prefix: int, n_tiles: int, sparsity: float, heads: int, gen: torch.Generator) -> torch.Tensor: + """VSA-H3 exempt selection: prefix rows dense, prefix keys always, top-k video keys.""" + n_video = n_tiles - n_prefix + k_video = max(1, math.ceil((1 - sparsity) * n_video)) + scores = torch.rand(1, heads, n_tiles, n_video, generator=gen, device=DEVICE) + mask = torch.zeros(1, heads, n_tiles, n_tiles, dtype=torch.bool, device=DEVICE) + mask.scatter_(-1, scores.topk(k_video, dim=-1).indices + n_prefix, True) + mask[..., :n_prefix] = True + mask[:, :, :n_prefix, :] = True + return mask + + +def _case(tile: int, prefix: tuple[int, ...], video_tiles: int, partial: dict[int, int], heads: int = 4): + gen = torch.Generator(device=DEVICE).manual_seed(tile * 1000 + video_tiles) + sizes = _tile_sizes(prefix, video_tiles, tile, partial) + n_tiles = sizes.numel() + n_prefix = n_tiles - video_tiles + rows = math.ceil(n_tiles * tile / BLOCK_M) * BLOCK_M + valid = torch.zeros(rows, dtype=torch.bool, device=DEVICE) + valid[:n_tiles * tile] = (torch.arange(tile, device=DEVICE)[None, :] < sizes[:, None]).reshape(-1) + q, k, v = (torch.randn(1, heads, rows, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16, generator=gen) + for _ in range(3)) + for x in (q, k, v): + x[:, :, ~valid] = 0 + mask = _exempt_mask(n_prefix, n_tiles, 0.8, heads, gen) + return q, k, v, sizes, mask, valid + + +def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a.float() - b.float()).norm() / b.float().norm()).item() + + +@pytest.mark.parametrize( + ("tile", "prefix", "video_tiles", "partial"), + [ + (64, (77, 46), 19, { + 3: 40, + 11: 1 + }), # odd tile count, short prefix chunks + (64, (64, 40), 9, {}), # even tile count + (256, (100, 70), 4, { + 1: 200 + }), # 256-token tiles, partial tile + ], +) +def test_sparse_matches_masked_reference_at_fp4_floor(tile, prefix, video_tiles, partial) -> None: + q, k, v, sizes, mask, valid = _case(tile, prefix, video_tiles, partial) + q2k_idx, q2k_num, kv_valid, q2k_quad = vsa_tile_mask_to_fp4_blocks(mask, tile, sizes, validate=True) + assert (q2k_quad is not None) == (tile < BLOCK_N) + out = sageattn_blackwell_sparse(q.clone(), k.clone(), v.clone(), q2k_idx, q2k_num, kv_valid, q2k_quad) + + n_tok = sizes.numel() * tile + token_tile = torch.arange(sizes.numel(), device=DEVICE).repeat_interleave(tile) + token_mask = mask[:, :, token_tile][:, :, :, token_tile] & valid[None, None, None, :n_tok] + ref = F.scaled_dot_product_attention(q[:, :, :n_tok].float(), + k[:, :, :n_tok].float(), + v[:, :, :n_tok].float(), + attn_mask=token_mask) + rows = valid[:n_tok] + assert torch.isfinite(out[:, :, :n_tok][:, :, rows]).all() + sparse_err = _rel_l2(out[:, :, :n_tok][:, :, rows], ref[:, :, rows]) + + # FP4 floor: the dense kernel against dense fp32 attention on the valid tokens. + qv, kv, vv = (x[:, :, valid].contiguous() for x in (q, k, v)) + floor = _rel_l2(sageattn_blackwell(qv, kv, vv), F.scaled_dot_product_attention(qv.float(), kv.float(), vv.float())) + assert sparse_err <= 1.05 * floor + 1e-3, (sparse_err, floor) + + +def test_full_lists_reproduce_dense_bitwise() -> None: + gen = torch.Generator(device=DEVICE).manual_seed(0) + q, k, v = (torch.randn(1, 4, 2048, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16, generator=gen) for _ in range(3)) + n_blocks = 2048 // BLOCK_N + q2k_idx = torch.arange(n_blocks, device=DEVICE, dtype=torch.int32).expand(1, 4, 2048 // BLOCK_M, + n_blocks).contiguous() + q2k_num = torch.full((1, 4, 2048 // BLOCK_M), n_blocks, device=DEVICE, dtype=torch.int32) + dense = sageattn_blackwell(q.clone(), k.clone(), v.clone()) + sparse = sageattn_blackwell_sparse(q.clone(), k.clone(), v.clone(), q2k_idx, q2k_num) + assert torch.equal(dense, sparse) + + +def test_bshd_entry_matches_bhsd() -> None: + q, k, v, sizes, mask, _ = _case(64, (77, 46), 19, {3: 40}) + lists = vsa_tile_mask_to_fp4_blocks(mask, 64, sizes) + bhsd = sageattn_blackwell_sparse(q.clone(), k.clone(), v.clone(), *lists) + bshd = sageattn_blackwell_sparse_bshd(*(x.transpose(1, 2).contiguous() for x in (q, k, v)), *lists) + assert torch.equal(bhsd, bshd) + + +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. + mask = torch.zeros(1, 1, 4, 4, dtype=torch.bool, device=DEVICE) + mask[0, 0, 0, 3] = True # query tile 0 -> key tile 3 only + mask[0, 0, 1, 1] = True # query tile 1 -> key tile 1 only (block 0, other quadrant) + mask[0, 0, 2:, 2:] = True + with pytest.raises(ValueError, match="first block"): + vsa_tile_mask_to_fp4_blocks(mask, 64, torch.full((4, ), 64, dtype=torch.int32, device=DEVICE), validate=True) + + +def test_block_list_check_rejects_out_of_bounds_lists() -> None: + idx = torch.zeros((1, 1, 2, 2), dtype=torch.int32, device=DEVICE) + num = torch.ones((1, 1, 2), dtype=torch.int32, device=DEVICE) + idx[..., 1] = 99 # beyond q2k_num: never read, so not checked + check_sparse_block_lists(idx, num, 2 * BLOCK_N) + with pytest.raises(ValueError, match="q2k_num"): + check_sparse_block_lists(idx, torch.zeros_like(num), 2 * BLOCK_N) + with pytest.raises(ValueError, match="q2k_idx"): + check_sparse_block_lists(idx, num + 1, 2 * BLOCK_N) diff --git a/fastvideo/api/compat.py b/fastvideo/api/compat.py index 56a54e282..720afc6a9 100644 --- a/fastvideo/api/compat.py +++ b/fastvideo/api/compat.py @@ -287,6 +287,8 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An quantization = engine.quantization if quantization is not None and quantization.text_encoder_quant is not None: kwargs["override_text_encoder_quant"] = quantization.text_encoder_quant + if quantization is not None and quantization.layer_profile is not None and quantization.transformer_quant is None: + raise ValueError("engine.quantization.layer_profile requires transformer_quant (for example NVFP4)") if quantization is not None and quantization.transformer_quant is not None: # Resolve the typed quant name to a concrete ``QuantizationConfig`` # instance and pin it on ``dit_config.quant_config``. The legacy @@ -295,8 +297,18 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An # typed surface accepts a string and does the wiring here so # downstream code can rely on a single source of truth. from fastvideo.layers.quantization import get_quantization_config + from fastvideo.layers.quantization.nvfp4_config import NVFP4Config _resolved_quant_cls = get_quantization_config(quantization.transformer_quant) - kwargs["transformer_quant"] = _resolved_quant_cls() + quant_instance = _resolved_quant_cls() + if quantization.layer_profile is not None: + if not isinstance(quant_instance, NVFP4Config): + raise ValueError("engine.quantization.layer_profile is only valid with transformer_quant NVFP4, " + f"got {type(quant_instance).__name__}") + quant_instance = NVFP4Config( + layer_profile=quantization.layer_profile, + retain_original_weights=quant_instance.retain_original_weights, + ) + kwargs["transformer_quant"] = quant_instance components = normalized.pipeline.components if components.pipeline_config_path is not None: diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 77224d036..521039785 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -76,6 +76,11 @@ class CompileConfig: class QuantizationConfig: text_encoder_quant: str | None = None transformer_quant: str | None = None + # Forwarded onto ``NVFP4Config`` when ``transformer_quant`` is ``NVFP4``. + # ``h3_dit`` selects the packed MiniMax-H3 attention+FFN export. + # ``h3_dit_ffn`` selects a packed FFN-only export (attention stays dense). + # ``h3_dit_vsa`` is ``h3_dit`` plus the VSA compression gates. + layer_profile: str | None = None @dataclass diff --git a/fastvideo/attention/backends/minimax_h3_sparse_int8.py b/fastvideo/attention/backends/minimax_h3_sparse_int8.py new file mode 100644 index 000000000..bab336c30 --- /dev/null +++ b/fastvideo/attention/backends/minimax_h3_sparse_int8.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental sm89 tile-64 VSA with BF16 or INT8 QK and FP32 accumulation. + +Retains each query tile's original key selection and masks partial key tiles. +Unlike a 128-query adapter, it adds no attention blocks. Q/K use per-token +scales; K centering is a softmax-invariant shift. V uses one scale per head +and channel, so its dequantization can be applied once in the epilogue. +PV stays in BF16: FP8 PV had excessive error on real H3 inputs. +Numerical validation and same-seed clip review are required before enabling. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _quantize_qk(X, Mean, VBS, Y, Scale, L: tl.constexpr, D: tl.constexpr, H: tl.constexpr, XB: tl.constexpr, + XH: tl.constexpr, XS: tl.constexpr, XD: tl.constexpr, CENTER: tl.constexpr, ROWS: tl.constexpr): + hz = tl.program_id(1) + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + cols = tl.arange(0, D) + offset = (hz // H) * XB + (hz % H) * XH + rows[:, None] * XS + cols[None, :] * XD + x = tl.load(X + offset, rows[:, None] < L, 0).to(tl.float32) + if CENTER: + mean = tl.load(Mean + hz * D + cols) + valid_size = tl.load(VBS + rows // 64, rows < L, 0) + x = tl.where((rows % 64 < valid_size)[:, None], x - mean[None, :], 0.0) + scale = tl.maximum(tl.max(tl.abs(x), 1) / 127.0, 1e-8) + y = tl.floor(x / scale[:, None] + 0.5).to(tl.int8) + tl.store(Y + (hz * L + rows[:, None]) * D + cols[None, :], y, rows[:, None] < L) + tl.store(Scale + hz * L + rows, scale, rows < L) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=w, num_stages=s) for w, s in ((4, 2), (4, 3), (4, 4), (8, 2), (8, 3))], + key=["L", "D"]) +@triton.jit +def _sparse_int8_fp8(Q, K, V, QS, KS, Index, Count, VBS, Out, L: tl.constexpr, D: tl.constexpr, H: tl.constexpr, + VB: tl.constexpr, VH: tl.constexpr, VS_ROW: tl.constexpr, VD: tl.constexpr, INT8_QK: tl.constexpr): + tile, hz = tl.program_id(0), tl.program_id(1) + nt: tl.constexpr = L // 64 + rows = tile * 64 + tl.arange(0, 64) + cols = tl.arange(0, D) + q = tl.load(Q + (hz * L + rows[:, None]) * D + cols[None, :]) + if INT8_QK: + qs = tl.load(QS + hz * L + rows) + nblocks = tl.load(Count + hz * nt + tile) + m = tl.full((64, ), -float("inf"), tl.float32) + den = tl.zeros((64, ), tl.float32) + acc = tl.zeros((64, D), tl.float32) + for block in range(nblocks): + kv = tl.load(Index + (hz * nt + tile) * nt + block) + key_rows = kv * 64 + tl.arange(0, 64) + k = tl.load(K + (hz * L + key_rows[None, :]) * D + cols[:, None]) + if INT8_QK: + ks = tl.load(KS + hz * L + key_rows) + valid = tl.load(VBS + kv) + if valid > 0: + logits = tl.dot(q, k).to(tl.float32) + if INT8_QK: + logits = logits * qs[:, None] * ks[None, :] + logits = logits * (1.4426950408889634 / D**0.5) + logits = tl.where((tl.arange(0, 64) < valid)[None, :], logits, -float("inf")) + block_max = tl.max(logits, 1) + new_m = tl.maximum(m, block_max) + p = tl.exp2(logits - new_m[:, None]) + alpha = tl.exp2(m - new_m) + den = den * alpha + tl.sum(p, 1) + acc = acc * alpha[:, None] + v = tl.load(V + (hz // H) * VB + (hz % H) * VH + key_rows[:, None] * VS_ROW + cols[None, :] * VD) + acc += tl.dot(p.to(tl.bfloat16), v, out_dtype=tl.float32) + m = new_m + result = acc / den[:, None] + result = tl.where(den[:, None] > 0, result, 0.0) + tl.store(Out + (hz * L + rows[:, None]) * D + cols[None, :], result.to(Out.dtype.element_ty)) + + +def sparse_sm89_attention(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor, + vbs: torch.Tensor, + *, + int8_qk: bool = True) -> torch.Tensor: + """Forward-only ``[B,H,S,128]`` BF16 attention on sm89, with 64-token tiles.""" + if torch.is_grad_enabled(): + raise ValueError("Sparse INT8/FP8 attention is inference-only") + if not q.is_cuda or torch.cuda.get_device_capability(q.device) != (8, 9): + raise ValueError("Sparse INT8/FP8 attention requires sm89 CUDA") + if q.dtype != torch.bfloat16 or q.shape[-1] != 128 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("Sparse INT8/FP8 attention requires matching BF16 Q/K/V with head dimension 128") + b, h, length, dim = q.shape + if length != vbs.numel() * 64 or mask.shape != (b, h, length // 64, length // 64): + raise ValueError("Sparse INT8/FP8 attention requires a tile-64 mask and validity vector") + from fastvideo_kernel.triton_kernels.index import map_to_index + + # The production INT8-QK/BF16-PV route reads BSHD-backed views directly. + # Quantized Q/K and the output remain contiguous BHSD. Other ablations + # retain their established layout and arithmetic. + if not int8_qk: + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + vbs = vbs.to(device=q.device, dtype=torch.int32).contiguous() + grid = (triton.cdiv(length, 16), b * h) + qi, ki = q, k + qs, ks = q, k # unused pointers in the BF16-QK ablation + if int8_qk: + # Tile pads are zero by contract; avoid a full FP32 copy for the reduction. + # Preserve the exact reduction used by the old contiguous adapter; + # its temporary copy dies before Q/K quantization and fine attention. + mean = k.contiguous().sum(dim=2, dtype=torch.float32) / vbs.sum().clamp_min(1) + qi = torch.empty(q.shape, device=q.device, dtype=torch.int8) + ki = torch.empty(k.shape, device=k.device, dtype=torch.int8) + qs = torch.empty((b, h, length), device=q.device, dtype=torch.float32) + ks = torch.empty_like(qs) + _quantize_qk[grid](q, mean, vbs, qi, qs, length, dim, h, *q.stride(), CENTER=False, ROWS=16, num_warps=4) + _quantize_qk[grid](k, mean, vbs, ki, ks, length, dim, h, *k.stride(), CENTER=True, ROWS=16, num_warps=4) + index, count = map_to_index(mask.contiguous()) + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + _sparse_int8_fp8[(length // 64, b * h)](qi, + ki, + v, + qs, + ks, + index, + count, + vbs, + out, + length, + dim, + h, + *v.stride(), + INT8_QK=int8_qk) + return out diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 65401519b..2ab65a6d9 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -547,6 +547,9 @@ class MiniMaxH3VSAImpl(AttentionImpl): self.prefix = prefix self.layer_idx = layer_idx_from_prefix(prefix, default=-1) self.head_size = head_size + self._sm89_kernel = envs.FASTVIDEO_H3_VSA_SM89_KERNEL.get() + if self._sm89_kernel not in {"original", "bf16", "int8"}: + raise ValueError("FASTVIDEO_H3_VSA_SM89_KERNEL must be original, bf16, or int8") # Generic torch.compile must not specialize the shared VSA forward on # the Python ``layer_idx`` value of each of H3's 50 blocks. This # tensor is prepared after weights load and drives only the compiled @@ -778,9 +781,19 @@ class MiniMaxH3VSAImpl(AttentionImpl): # kernels' granularity. These entries take BHSD ([B, H, S_pad, D]); # mirror block_sparse_attn_256_bshd's Triton branch and transpose # around the call. - q_bhsd = query.transpose(1, 2).contiguous() - k_bhsd = key.transpose(1, 2).contiguous() - v_bhsd = value.transpose(1, 2).contiguous() + # One gate for the sm89 QK/PV route below: the kernels are + # inference-only, BF16, head-128, and validated for Ada (8, 9). + # ``sm89_strided`` marks the INT8-QK entry, the only one that reads + # the BSHD strided views directly. + sm89_route = (self._sm89_kernel != "original" and not torch.is_grad_enabled() and not compiling + and query.dtype == torch.bfloat16 and query.shape[-1] == 128 and query.is_cuda + and torch.cuda.get_device_capability(query.device) == (8, 9)) + sm89_strided = sm89_route and self._sm89_kernel == "int8" + q_bhsd = query.transpose(1, 2) + k_bhsd = key.transpose(1, 2) + v_bhsd = value.transpose(1, 2) + if not sm89_strided: + q_bhsd, k_bhsd, v_bhsd = (t.contiguous() for t in (q_bhsd, k_bhsd, v_bhsd)) sm100a_mask = mask sm100a_variable_block_sizes = attn_metadata.variable_block_sizes @@ -877,19 +890,34 @@ class MiniMaxH3VSAImpl(AttentionImpl): ) else: if has_sm100a_pair: - q_bhsd = q_bhsd[:, :, :logical_seq_len].contiguous() - k_bhsd = k_bhsd[:, :, :logical_seq_len].contiguous() - v_bhsd = v_bhsd[:, :, :logical_seq_len].contiguous() - out_bhsd, _ = block_sparse_attn_64_bhsd( - q_bhsd, - k_bhsd, - v_bhsd, - mask, - attn_metadata.variable_block_sizes, - ) + q_bhsd = q_bhsd[:, :, :logical_seq_len] + k_bhsd = k_bhsd[:, :, :logical_seq_len] + 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: + 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, + k_bhsd, + v_bhsd, + mask, + attn_metadata.variable_block_sizes, + int8_qk=self._sm89_kernel == "int8") + else: + out_bhsd, _ = block_sparse_attn_64_bhsd( + q_bhsd, + k_bhsd, + v_bhsd, + mask, + attn_metadata.variable_block_sizes, + ) if has_sm100a_pair and use_sm100a: out_bhsd = out_bhsd[:, :, :logical_seq_len] out = out_bhsd.transpose(1, 2).contiguous() + # Fine attention is complete. Release its layout copies before the + # gated compression merge creates full-sequence temporaries. + del q_bhsd, k_bhsd, v_bhsd, out_bhsd else: out, _ = block_sparse_attn_256_bshd( logical_query, diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 8962a8e70..77ba871a6 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -516,6 +516,121 @@ FASTVIDEO_FAD_REF_FEATURES = EnvStr(None, category="eval", doc="Cached reference-feature file for the audio Frechet distance metric.") +# ================== MiniMax-H3 single-GPU switches ================== + +FASTVIDEO_H3_VSA_FP4 = EnvBool(False, + category="attention", + doc="Run MiniMax-H3 VSA attention on the block-sparse SageAttention3 FP4 kernel " + "(sm_120, no-grad, single sequence-parallel rank).") +FASTVIDEO_H3_VSA_TILE_FIRST = EnvBool(False, + category="attention", + doc="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 = EnvChoice("original", + choices=("original", "bf16", "int8"), + category="attention", + doc="Fine-attention kernel for MiniMax-H3 VSA on sm_89: original, bf16, " + "or int8 (INT8 QK, BF16 PV).") +FASTVIDEO_H3_SIM_SP_FP8 = EnvBool(False, + category="debug", + doc="Simulate the FP8 sequence-parallel exchange of the MiniMax-H3 FP4 VSA " + "path on one rank.") +FASTVIDEO_H3_FFN_CHUNK_TOKENS = EnvInt(0, + category="performance", + doc="Inference-only MiniMax-H3 FFN token chunk size; 0 runs the FFN " + "unchunked.") +FASTVIDEO_H3_FP8_ATTENTION = EnvBool(False, + category="performance", + doc="With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention " + "projections in FP8.") +FASTVIDEO_H3_FP8_GRANULARITY = EnvChoice("tensor", + choices=("tensor", "channel"), + category="performance", + doc="FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION.") +FASTVIDEO_NVFP4_MM_BACKEND = EnvStr("auto", + category="performance", + doc="FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass.") +FASTVIDEO_NVFP4_ACT_AMAX = EnvPath(None, + category="performance", + doc="JSON of calibrated NVFP4 input amax per linear, keyed b. or " + "full prefix; sets a static activation scale.") +FASTVIDEO_NVFP4_DYNAMIC_ACT = EnvStr("", + category="performance", + doc="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 = EnvBool(False, + category="performance", + doc="Cache MiniMax-H3 AdaLN modulation per timestep instead of keeping the " + "projection weights resident.") +FASTVIDEO_H3_ADALN_TABLE = EnvPath(None, + category="performance", + doc="Precomputed MiniMax-H3 AdaLN modulation table; enables the cache and " + "skips loading the AdaLN projection weights.") +FASTVIDEO_H3_ADALN_DUMP = EnvPath(None, + category="debug", + doc="Write the MiniMax-H3 AdaLN modulation table to this path while sampling.") +FASTVIDEO_H3_SPLICE_TRANSFORMER = EnvPath(None, + category="eval", + doc="Second MiniMax-H3 transformer that runs the late DMD steps " + "(checkpoint step-splice evaluation).") +FASTVIDEO_H3_SPLICE_FROM_STEP = EnvInt(4, + category="eval", + doc="First denoising step run by FASTVIDEO_H3_SPLICE_TRANSFORMER.") +FASTVIDEO_H3_ENCODER_LAYERWISE = EnvBool(False, + category="performance", + doc="Stream MiniMax-H3 text-encoder language layers through exact-size " + "pinned host memory (text-only prompts).") +FASTVIDEO_H3_ENCODER_FUSED_DEQUANT = EnvBool(False, + category="performance", + doc="Expand the serialized NVFP4 MiniMax-H3 text encoder with one " + "fused Triton pass on GPUs without FP4 GEMM.") +FASTVIDEO_H3_VAE_TILE_BATCH = EnvInt(1, + category="performance", + doc="Spatial tiles per MiniMax-H3 video VAE decoder call; 1 decodes per " + "tile.") +FASTVIDEO_H3_VAE_INT8_SHARED_QKV = EnvBool(False, + category="performance", + doc="Share the INT8 activation rotation and quantization across the " + "MiniMax-H3 VAE Q/K/V projections.") +FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW = EnvBool(False, + category="performance", + doc="Use transposed weight views in the MiniMax-H3 VAE INT8 " + "projections.") +FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT = EnvBool(False, + category="performance", + doc="Fused dequantization epilogue for the MiniMax-H3 VAE INT8 " + "projections.") +FASTVIDEO_H3_PINNED_SWAP = EnvBool(True, + category="performance", + doc="Swap offloaded MiniMax-H3 modules through exact-size pinned host " + "arenas.") +FASTVIDEO_H3_PARK_MODULES = EnvStr(None, + category="performance", + doc="Comma-separated MiniMax-H3 denoise modules parked on the host while the " + "text encoder runs, e.g. vae,audio_vae.") +FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS = EnvBool(False, + category="performance", + doc="Layerwise offload also streams large buffers such as packed " + "FP4/FP8 weights.") +FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS = EnvInt(0, + category="performance", + doc="Keep the first N layerwise-offloaded blocks resident on the " + "GPU.") +FASTVIDEO_H3_SP_PROFILE = EnvBool(False, + category="profiling", + doc="CUDA-event spans per stage over one MiniMax-H3 FP4 VSA DiT forward.") +FASTVIDEO_H3_CAPTURE_QKV = EnvPath(None, + category="debug", + doc="Directory for captured real MiniMax-H3 Q/K/V attention inputs.") +FASTVIDEO_CUDA_MEMORY_CAP_GIB = EnvFloat(0.0, + category="debug", + doc="Cap this process's CUDA allocator at this many GiB to emulate a " + "smaller GPU; 0 leaves it uncapped.") +FASTVIDEO_MEMORY_REPORT = EnvBool(False, + category="debug", + doc="Log bytes held per pipeline component by device and dtype after " + "loading.") + # ================== Tests ================== FASTVIDEO_TEST_LTX2_OVERFIT_DATA_DIR = EnvStr("data/cats", diff --git a/fastvideo/hooks/layerwise_offload.py b/fastvideo/hooks/layerwise_offload.py index 05a9aacf5..211ea0f3e 100644 --- a/fastvideo/hooks/layerwise_offload.py +++ b/fastvideo/hooks/layerwise_offload.py @@ -1,8 +1,11 @@ from contextlib import contextmanager +from itertools import chain from typing import Any import torch from torch import nn from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager +import fastvideo.envs as envs +from fastvideo.hooks.pinned_memory import PinnedTensorArena from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -14,6 +17,29 @@ def _tensor_placeholder(tensor: torch.Tensor, device: torch.device) -> torch.Ten return torch.empty(shape, device=device, dtype=tensor.dtype) +# Buffers at least this large also stream (e.g. packed NVFP4 weights registered as buffers); +# small ones (scales, caches) stay resident. Opt in with FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS=1. +_BUFFER_OFFLOAD_MIN_BYTES = 1 << 20 + + +def _offload_tensors(module: nn.Module, names: dict[str, torch.Tensor] | None = None): + """``(name, tensor)`` for every parameter and, when enabled, every large buffer. + + ``names`` restricts the walk to the tensors chosen at init: an offloaded buffer is a + zero-element placeholder afterwards and would fail the size test. + """ + if names is not None: + for name, tensor in chain(module.named_parameters(), module.named_buffers()): + if name in names: + yield name, tensor + return + yield from module.named_parameters() + if envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.get(): + for name, buf in module.named_buffers(): + if buf is not None and buf.numel() * buf.element_size() >= _BUFFER_OFFLOAD_MIN_BYTES: + yield name, buf + + class LayerwiseOffloadState: def __init__( @@ -28,6 +54,7 @@ class LayerwiseOffloadState: self.cpu_named_parameters: dict[str, torch.Tensor] = {} self.module_ref: nn.Module = None # type: ignore self.device: torch.device = device + self.cpu_arena: PinnedTensorArena | None = None def _will_offload(self, name: str) -> bool: return True @@ -35,16 +62,27 @@ class LayerwiseOffloadState: @torch.compiler.disable def on_init(self, module: nn.Module): self.module_ref = module - for name, param in self.module_ref.named_parameters(): + self.clear_cpu_storage() + self.cpu_arena = PinnedTensorArena( + (name, param) for name, param in _offload_tensors(module) if self._will_offload(name)) + for name, param in _offload_tensors(self.module_ref): if self._will_offload(name): - self.cpu_named_parameters[name] = (param.data.detach().to("cpu").pin_memory()) + host = self.cpu_arena.empty_like(name, param) + host.copy_(param.data.detach()) + self.cpu_named_parameters[name] = host param.data = _tensor_placeholder(param.data, self.device) + def clear_cpu_storage(self) -> None: + self.cpu_named_parameters.clear() + if self.cpu_arena is not None: + self.cpu_arena.close() + self.cpu_arena = None + @torch.compiler.disable def wait_and_replace_params(self): torch.cuda.current_stream().wait_stream(self.async_copy_stream) # now gpu_named_parameters are ready - for name, param in self.module_ref.named_parameters(): + for name, param in _offload_tensors(self.module_ref, self.cpu_named_parameters): if not self._will_offload(name): continue if name not in self.gpu_named_parameters: @@ -56,7 +94,7 @@ class LayerwiseOffloadState: def prefetch_params(self): compute_stream = torch.cuda.current_stream() with torch.cuda.stream(self.async_copy_stream): - for name, param in self.module_ref.named_parameters(): + for name, param in _offload_tensors(self.module_ref, self.cpu_named_parameters): if not self._will_offload(name): continue assert name not in self.gpu_named_parameters @@ -66,7 +104,7 @@ class LayerwiseOffloadState: @torch.compiler.disable def release_gpu_params(self): - for name, param in self.module_ref.named_parameters(): + for name, param in _offload_tensors(self.module_ref, self.cpu_named_parameters): if self._will_offload(name): param.data = _tensor_placeholder(param.data, self.device) del self.gpu_named_parameters[name] @@ -83,16 +121,17 @@ class LayerwiseOffloadHook(ForwardHook): self.state.on_init(module) # pyright: ignore def on_detach(self, module: nn.Module): - named_parameters = dict(module.named_parameters()) + self.state.async_copy_stream.synchronize() + named_parameters = dict(_offload_tensors(module, self.state.cpu_named_parameters)) for name, cpu_tensor in self.state.cpu_named_parameters.items(): - if name not in self.state.gpu_named_parameters: - if name in named_parameters: - named_parameters[name].data = cpu_tensor.to(device=self.state.device) - else: - logger.warning( - "Parameter {} not found in module during detachment.", - name, - ) + if name in named_parameters: + gpu_tensor = self.state.gpu_named_parameters.get(name) + named_parameters[name].data = gpu_tensor if gpu_tensor is not None else cpu_tensor.to(self.state.device) + else: + logger.warning("Parameter %s not found in module during detachment.", name) + self.state.gpu_named_parameters.clear() + self.state.clear_cpu_storage() + self.state.next_state = None @classmethod def name(cls) -> str: @@ -129,12 +168,15 @@ class LayerwiseOffloadHook(ForwardHook): yield finally: # instead of releasing, we should overwrite the original params since they have been modified - self.state.cpu_named_parameters.clear() self.state.gpu_named_parameters.clear() self.state.on_init(self.state.module_ref) # pyright: ignore -def enable_layerwise_offload(model: nn.Module, is_replace: bool = False): +def enable_layerwise_offload(model: nn.Module, + is_replace: bool = False, + *, + resident_blocks: int | None = None, + cyclic: bool = True): if torch.cuda.is_available(): device = torch.device("cuda", torch.cuda.current_device()) else: @@ -142,9 +184,21 @@ def enable_layerwise_offload(model: nn.Module, is_replace: bool = False): return state_list = [] async_stream = torch.cuda.Stream() + # The first N entries skip offloading and stay wherever the model is placed (normally the + # GPU), so a GPU with spare memory streams only the remainder over PCIe. + if resident_blocks is not None: + resident = max(0, resident_blocks) + else: + try: + resident = max(0, envs.FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS.get()) + except ValueError as error: + logger.warning("Ignoring malformed FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS: %s", error) + resident = 0 for name, submodule in model.named_children(): if isinstance(submodule, nn.ModuleList): for idx, module_entry in enumerate(submodule): + if idx < resident: + continue state = LayerwiseOffloadState(async_copy_stream=async_stream, device=device) state_list.append(state) hook_mgr = ModuleHookManager.get_from_or_default(module_entry) @@ -159,8 +213,13 @@ def enable_layerwise_offload(model: nn.Module, is_replace: bool = False): hook_mgr.append_forward_hook(hook) break if len(state_list) == 0: + if resident > 0: + logger.info("FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS=%d keeps every block resident; nothing to offload", + resident) + return raise ValueError("No nn.ModuleList found in the model for layerwise offloading.") - # circular linking of states + # Repeated DiT steps prefetch the first block after the last. A once-per-request + # encoder can skip that unused copy and release every layer after its forward. for i in range(len(state_list)): - state_list[i].next_state = state_list[(i + 1) % len(state_list)] + state_list[i].next_state = state_list[(i + 1) % len(state_list)] if cyclic or i + 1 < len(state_list) else None diff --git a/fastvideo/hooks/pinned_memory.py b/fastvideo/hooks/pinned_memory.py new file mode 100644 index 000000000..81e7ff8af --- /dev/null +++ b/fastvideo/hooks/pinned_memory.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact-size pinned host storage for inference offload.""" + +import weakref +from collections.abc import Iterable + +import torch + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) +_ALIGNMENT = 256 +_PAGE_ALIGNMENT = 4096 + + +def _unregister(buffer: torch.Tensor, device: int) -> None: + # The allocation must outlive any outstanding nonblocking H2D copies. + try: + with torch.cuda.device(device): + torch.cuda.synchronize() + error = torch.cuda.cudart().cudaHostUnregister(buffer.data_ptr()) + if error != 0: + logger.warning("cudaHostUnregister failed: %s", error) + except Exception as exc: + # CUDA may already be unavailable during interpreter shutdown. + logger.warning("Could not unregister pinned host arena: %s", exc) + + +class PinnedTensorArena: + """Pack tensors into one CUDA-registered CPU allocation, aligned to 256 bytes. + + PyTorch's pinned allocator rounds large allocations to powers of two. Registering + ordinary host storage avoids that overhead. Typed views retain this owner, so + registration survives even if the module or offload state is dropped first. + If registration is unavailable, allocate conventional pinned tensors instead. + Call ``close`` only after all views and pending copies have been released. + """ + + def __init__(self, tensors: Iterable[tuple[str, torch.Tensor]]) -> None: + self.offsets: dict[str, tuple[int, int]] = {} + size = 0 + for name, tensor in tensors: + size = (size + _ALIGNMENT - 1) // _ALIGNMENT * _ALIGNMENT + length = tensor.numel() * tensor.element_size() + self.offsets[name] = (size, length) + size += length + self.nbytes = size + self.buffer: torch.Tensor | None = None + self._finalizer: weakref.finalize | None = None + if not size: + return + # Register dedicated pages: small malloc allocations can otherwise share a + # registered page with another arena. The extra space is bounded by 8 KiB. + span = (size + _PAGE_ALIGNMENT - 1) // _PAGE_ALIGNMENT * _PAGE_ALIGNMENT + allocation = torch.empty(span + _PAGE_ALIGNMENT - 1, dtype=torch.uint8, device="cpu") + start = (-allocation.data_ptr()) % _PAGE_ALIGNMENT + # Give the aligned region its own storage base. Tensor.is_pinned() queries + # the storage pointer, which would precede the registered region for a + # plain narrow() view. The memoryview retains the original allocation. + buffer = torch.frombuffer(memoryview(allocation.numpy())[start:start + span], dtype=torch.uint8) + device = torch.cuda.current_device() + try: + error = torch.cuda.cudart().cudaHostRegister(buffer.data_ptr(), span, 0) + if error != 0: + raise RuntimeError(f"cudaHostRegister returned {error}") + except Exception as exc: + logger.warning("Exact-size host registration failed; using the pinned allocator: %s", exc) + return + self.buffer = buffer + self._finalizer = weakref.finalize(self, _unregister, buffer, device) + + def empty_like(self, name: str, tensor: torch.Tensor) -> torch.Tensor: + """Return a contiguous host view with the source's dtype and shape.""" + if self.buffer is None: + return torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True) + offset, length = self.offsets[name] + host = self.buffer.narrow(0, offset, length).view(tensor.dtype).reshape(tensor.shape) + host._pinned_arena = self + return host + + def close(self) -> None: + """Unregister before releasing storage; safe to call more than once.""" + if self._finalizer is not None: + self._finalizer() + self._finalizer = None + self.buffer = None diff --git a/fastvideo/layers/quantization/fp8_config.py b/fastvideo/layers/quantization/fp8_config.py index e73c372d6..5b7e3f3cb 100644 --- a/fastvideo/layers/quantization/fp8_config.py +++ b/fastvideo/layers/quantization/fp8_config.py @@ -3,6 +3,8 @@ Matches linear layers by suffix (``to_q/k/v/to_out``, ``ffn.fc_in/fc_out``). Supports per-tensor (default, fast) and per-channel (higher accuracy) granularity. +On sm89 the per-channel path runs the per-tensor GEMM plus a Triton scale epilogue +(``fp8_kernels``), since torch's rowwise-scaled kernel there is slower than bf16. Falls back to bf16 dequant on GPUs older than sm89. """ from __future__ import annotations @@ -32,6 +34,9 @@ FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) _FP8_SUFFIXES = ( "ffn.fc_in", "ffn.fc_out", + # MiniMax-H3 blocks name their MLP ``ff``. + "ff.fc_in", + "ff.fc_out", "to_q", "to_k", "to_v", @@ -74,8 +79,18 @@ def _quantize_tensorwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor return x_fp8, x_scale.view(1) +def _rowwise_scaled_mm_is_slow() -> bool: + if getattr(torch.version, "hip", None) or not torch.cuda.is_available(): + return False + from fastvideo.layers.quantization.fp8_kernels import rowwise_scaled_mm_is_slow + return rowwise_scaled_mm_is_slow() + + def _quantize_rowwise(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``.""" + if x_2d.is_cuda and _rowwise_scaled_mm_is_slow(): + from fastvideo.layers.quantization.fp8_kernels import quantize_rowwise_fp8 + return quantize_rowwise_fp8(x_2d) x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float() x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE) x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) @@ -156,6 +171,12 @@ class FP8QuantizeMethod(QuantizeMethodBase): w_fp8 = layer._fp8_weight w_scale = layer._fp8_weight_scale + if self.granularity == "channel" and _rowwise_scaled_mm_is_slow(): + from fastvideo.layers.quantization.fp8_kernels import scaled_mm_token_channel + out = scaled_mm_token_channel(x_fp8, x_scale, w_fp8.t(), w_scale) + if bias is not None: + out = out + bias + return out.view(*original_shape[:-1], out_dim) scale_b = w_scale.view(1, -1) if self.granularity == "channel" else w_scale out = torch._scaled_mm( diff --git a/fastvideo/layers/quantization/fp8_kernels.py b/fastvideo/layers/quantization/fp8_kernels.py new file mode 100644 index 000000000..a589262c3 --- /dev/null +++ b/fastvideo/layers/quantization/fp8_kernels.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Triton helpers for per-token x per-channel FP8 linears on GPUs whose rowwise ``_scaled_mm`` is slow. + +On sm89 (RTX 4090 / L40S / RTX 6000 Ada) ``torch._scaled_mm`` with rowwise scales runs at +~70 TFLOPS, below bf16 (~160), while the per-tensor kernel reaches 220-305 TFLOPS. The +per-token x per-channel result is recovered exactly by running the per-tensor kernel with unit +scales and applying ``out[i, j] *= sx[i] * sw[j]`` in one pass over the output, which costs +5-10% of the GEMM instead of 2-4x. +""" +from __future__ import annotations + +import functools + +import torch +import triton +import triton.language as tl + +FP8_MAX = 448.0 +FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0) + + +@functools.cache +def rowwise_scaled_mm_is_slow() -> bool: + """Ada (sm89) has no fast rowwise-scaled FP8 GEMM in torch; Hopper and Blackwell do.""" + return torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9) + + +@triton.jit +def _quantize_rowwise_kernel(x_ptr, q_ptr, s_ptr, K, stride_x, stride_q, BLOCK_K: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + x_row = x_ptr + row * stride_x + amax = tl.zeros((BLOCK_K, ), dtype=tl.float32) + for k in range(0, K, BLOCK_K): + cols = k + tl.arange(0, BLOCK_K) + amax = tl.maximum(amax, tl.abs(tl.load(x_row + cols, mask=cols < K, other=0.0).to(tl.float32))) + scale = tl.maximum(tl.max(amax, axis=0) / 448.0, 1.0 / (448.0 * 512.0)) + tl.store(s_ptr + row, scale) + inv = 1.0 / scale + for k in range(0, K, BLOCK_K): + cols = k + tl.arange(0, BLOCK_K) + v = tl.load(x_row + cols, mask=cols < K, other=0.0).to(tl.float32) * inv + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(q_ptr + row * stride_q + cols, v.to(tl.float8e4nv), mask=cols < K) + + +def quantize_rowwise_fp8(x_2d: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token FP8 quantization in one launch. Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``.""" + x_2d = x_2d.contiguous() + M, K = x_2d.shape + q = torch.empty((M, K), device=x_2d.device, dtype=torch.float8_e4m3fn) + s = torch.empty((M, 1), device=x_2d.device, dtype=torch.float32) + if M: + _quantize_rowwise_kernel[(M, )](x_2d, q, s, K, x_2d.stride(0), q.stride(0), BLOCK_K=1024, num_warps=4) + return q, s + + +@triton.jit +def _scale_rows_cols_kernel(o_ptr, sx_ptr, sw_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): + rows = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (rows[:, None] < M) & (cols[None, :] < N) + # int64 offsets: a 78k-token fc_in output has 2.2e9 elements, past int32. + ptrs = o_ptr + rows[:, None] * N + cols[None, :] + v = tl.load(ptrs, mask=mask, other=0.0).to(tl.float32) + sx = tl.load(sx_ptr + rows, mask=rows < M, other=0.0) + sw = tl.load(sw_ptr + cols, mask=cols < N, other=0.0) + tl.store(ptrs, (v * sx[:, None] * sw[None, :]).to(o_ptr.dtype.element_ty), mask=mask) + + +def scaled_mm_token_channel(x_fp8: torch.Tensor, x_scale: torch.Tensor, w_fp8_t: torch.Tensor, + w_scale: torch.Tensor) -> torch.Tensor: + """``(x_fp8 * x_scale) @ (w_fp8_t * w_scale)`` in bf16 via the fast per-tensor GEMM plus a scale epilogue. + + The unit-scale GEMM output is at most 448^2 * K, far inside bf16 range, and its relative + precision is that of any bf16 output, so the epilogue loses nothing against rowwise scaling. + """ + one = torch.ones((), device=x_fp8.device, dtype=torch.float32) + out = torch._scaled_mm(x_fp8, w_fp8_t, scale_a=one, scale_b=one, out_dtype=torch.bfloat16) + if isinstance(out, tuple): + out = out[0] + M, N = out.shape + grid = (triton.cdiv(M, 64), triton.cdiv(N, 128)) + _scale_rows_cols_kernel[grid](out, x_scale.reshape(-1), w_scale.reshape(-1), M, N, BLOCK_M=64, BLOCK_N=128) + return out diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index 284622d59..3e683f749 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/fastvideo/layers/quantization/nvfp4_config.py @@ -8,8 +8,9 @@ explicitly so downstream callers don't conflate it with other FP4 variants that may land later (e.g. AMD's MX-FP4 or vendor-neutral e3m0). -The registered config targets the curated LTX-2 deployment set and the -main MiniMax-H3 transformer-block FFN linears. +The registered config targets the curated LTX-2 deployment set, the +main MiniMax-H3 transformer-block FFN linears, and the packed MiniMax-H3 +DiT export (``layer_profile="h3_dit"``) covering attention plus FFN. `flashinfer` is imported lazily inside the call paths that need it. This keeps ``import fastvideo`` cheap on hosts where flashinfer is @@ -19,6 +20,7 @@ use time, with a clear error. from __future__ import annotations import logging +import os import re from typing import Any @@ -26,6 +28,8 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter +import fastvideo.envs as envs + from fastvideo.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -76,6 +80,21 @@ _LTX2_NVFP4_LINEAR_PREFIXES = frozenset(f"ltx2.blocks.{block_idx}.{suffix}" for for suffix in _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES) | frozenset( ("ltx2.adaln_single.linear", )) _MINIMAX_H3_NVFP4_FF_PREFIX = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.ff\.(?:fc_in|fc_out)$") +_MINIMAX_H3_NVFP4_DIT_PREFIX = re.compile( + r"(?:^|\.)transformer_blocks\.\d+\.(?:attn\.to_(?:q|k|v|out)|ff\.(?:fc_in|fc_out))$") +_H3_BLOCK_ATTN_PROJ = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.attn\.(?:to_q|to_k|to_v|to_out)$") +_MINIMAX_H3_NVFP4_VSA_GATE_PREFIX = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.attn\.to_gate_compress$") +H3_NVFP4_DIT_EXPORT_FILENAME = "nvfp4_weights.safetensors" +H3_NVFP4_DIT_KEY_SEP = "::" +H3_NVFP4_DIT_BUFFER_NAMES = ( + "_nvfp4_weight", + "_nvfp4_weight_scale", + "_nvfp4_alpha", + "_weight_global_sf", +) +# 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" def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool: @@ -84,10 +103,43 @@ def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool: def is_minimax_h3_nvfp4_linear_prefix(prefix: str) -> bool: - """Return whether *prefix* is a main MiniMax-H3 transformer-block FFN linear.""" return _MINIMAX_H3_NVFP4_FF_PREFIX.search(prefix) is not None +def is_minimax_h3_nvfp4_dit_linear_prefix(prefix: str) -> bool: + """Return whether *prefix* is a MiniMax-H3 DiT attention or FFN linear. + + This is the packed NVFP4H3 export set: ``attn.to_{q,k,v,out}`` and + ``ff.{fc_in,fc_out}`` in each main transformer block. Token-refiner, + AdaLN, and embedding linears stay dense. + """ + return _MINIMAX_H3_NVFP4_DIT_PREFIX.search(prefix) is not None + + +def is_minimax_h3_nvfp4_dit_export_path(path: str) -> bool: + return os.path.basename(path) == H3_NVFP4_DIT_EXPORT_FILENAME + + +def find_minimax_h3_nvfp4_dit_export(weight_paths: list[str]) -> str | None: + seen: list[str] = [] + for path in weight_paths: + if is_minimax_h3_nvfp4_dit_export_path(path) and os.path.isfile(path): + return path + directory = path if os.path.isdir(path) else os.path.dirname(path) + if directory and directory not in seen: + seen.append(directory) + for directory in seen: + candidate = os.path.join(directory, H3_NVFP4_DIT_EXPORT_FILENAME) + if os.path.isfile(candidate): + return candidate + return None + + +def dense_transformer_safetensors(weight_paths: list[str]) -> list[str]: + """Drop the packed NVFP4 DiT export so it is not loaded as bf16 weights.""" + return [path for path in weight_paths if not is_minimax_h3_nvfp4_dit_export_path(path)] + + def _is_ltx2_refine_only_prefix(prefix: str) -> bool: return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES) @@ -290,6 +342,16 @@ def _mm_fp4( ) +def _mm_fp4_backend() -> str: + """FlashInfer ``mm_fp4`` backend (``FASTVIDEO_NVFP4_MM_BACKEND``, default ``auto``). + + On sm_120 ``auto`` picks a kernel about 2x slower than ``cutlass`` or + ``cudnn`` once activations reach tens of thousands of rows (measured at + 73k rows on an RTX PRO 6000); short sequences are unaffected. + """ + return envs.FASTVIDEO_NVFP4_MM_BACKEND.get() + + def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor: """Coerce an activation to a dtype the FP4 linear accepts. @@ -306,8 +368,24 @@ def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor: return x +_AMAX_TABLES: dict[str, dict[str, float]] = {} + + +def _load_amax_table(path: str) -> dict[str, float]: + if path not in _AMAX_TABLES: + import json + with open(path) as f: + raw = json.load(f) + _AMAX_TABLES[path] = {k: float(v["all"] if isinstance(v, dict) else v) for k, v in raw.items()} + return _AMAX_TABLES[path] + + class NVFP4QuantizeMethod(QuantizeMethodBase): + # Lazily resolved by _static_activation_global_sf; class defaults also cover object.__new__ test doubles. + _static_sf_checked: bool = False + _static_sf: torch.Tensor | None = None + def __init__(self, layer_prefix: str = ""): super().__init__() self.weight_fp4 = None @@ -333,6 +411,42 @@ class NVFP4QuantizeMethod(QuantizeMethodBase): layer.register_parameter("weight", weight) set_weight_attrs(weight, extra_weight_attrs) + def _static_activation_global_sf(self) -> torch.Tensor | None: + """FASTVIDEO_NVFP4_ACT_AMAX: JSON of calibrated input amax per layer ("b." or full prefix).""" + if getattr(self, "_static_sf_checked", False): + return self._static_sf + self._static_sf_checked, self._static_sf = True, None + path = envs.FASTVIDEO_NVFP4_ACT_AMAX.get() + if path: + table = _load_amax_table(path) + prefix = self.layer_prefix or "" + match = re.search(r"transformer_blocks\.(\d+)\.(.+)$", prefix) + keys = [prefix] + ([f"b{match.group(1)}.{match.group(2)}"] if match else []) + amax = next((table[k] for k in keys if k in table), None) + if amax is not None: + self._static_sf = torch.tensor((448.0 * 6.0) / max(amax, 1e-12), dtype=torch.float32, device="cuda") + return self._static_sf + + def _dynamic_activation_scale(self) -> bool: + """FASTVIDEO_NVFP4_DYNAMIC_ACT: "all", or comma-separated layer-name suffixes (e.g. "ff.fc_out").""" + cached = getattr(self, "_dynamic_act_cached", None) + if cached is None: + selected = envs.FASTVIDEO_NVFP4_DYNAMIC_ACT.get() + suffixes = [part.strip() for part in selected.split(",") if part.strip()] + prefix = self.layer_prefix or "" + cached = "all" in suffixes or any(prefix.endswith(suffix) for suffix in suffixes) + self._dynamic_act_cached = cached + return cached + + def uses_unit_activation_scale(self, layer: torch.nn.Module) -> bool: + """Whether ``apply`` quantizes this layer's input with the unit global scale. + + False when a calibrated scale (env table or the export's ``_nvfp4_input_global_sf``) or a dynamic + per-call scale applies; such inputs cannot share one pre-quantized copy across layers. + """ + return (self._static_activation_global_sf() is None and getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None) is None + and not self._dynamic_activation_scale()) + def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: SfLayout, _, _ = _require_flashinfer() x = _coerce_fp4_input_dtype(x) @@ -392,7 +506,17 @@ class NVFP4QuantizeMethod(QuantizeMethodBase): else: x = _coerce_fp4_input_dtype(x) x = x.view(-1, x.shape[-1]) - x_global_sf = self.x_global_sf + static_sf = self._static_activation_global_sf() + if static_sf is None: + static_sf = getattr(layer, H3_NVFP4_DIT_INPUT_SF_NAME, None) + if static_sf is not None: + x_global_sf = static_sf + elif self._dynamic_activation_scale(): + # A unit global scale caps FP8 block scales at |x| = 6 * 448; inputs such as H3's ff.fc_out + # (post-SwiGLU) exceed that, so derive the global scale from this call's amax. + x_global_sf = (448.0 * 6.0) / x.abs().amax().float().clamp(min=1e-12) + else: + x_global_sf = self.x_global_sf x_fp4, x_scale = _nvfp4_quantize( x, x_global_sf, @@ -417,7 +541,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase): alpha, torch.bfloat16, None, - backend='auto', + backend=_mm_fp4_backend(), ) if bias is not None: @@ -431,14 +555,22 @@ class NVFP4Config(QuantizationConfig): NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha, ``layout_128x4`` scale layout, group size 16). LTX-2 uses its curated - attention and FFN deployment set. MiniMax-H3 uses only ``fc_in`` and - ``fc_out`` in each main transformer-block FFN. + attention and FFN deployment set. MiniMax-H3's default profile uses only + ``fc_in`` and ``fc_out`` in each main transformer-block FFN. + ``layer_profile="h3_dit"`` expands that to the packed NVFP4H3 DiT set + (attention ``to_{q,k,v,out}`` plus those FFN linears). + ``layer_profile="h3_dit_ffn"`` loads a packed export of the FFN linears + only, keeping attention projections dense (e.g. calibrated FFN-only + 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). """ def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None): super().__init__() - # ``base``: stage-1 set (no attn2.to_out, no cross-modal AV - # projections). ``refine``: full stage-2 set. + 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}") self.layer_profile = layer_profile # Original bf16 ``layer.weight`` retention after FP4 conversion. # Default (None/False): purge the purgeable originals -- every @@ -473,13 +605,26 @@ class NVFP4Config(QuantizationConfig): def get_quant_method(self, layer: torch.nn.Module, prefix: str): from fastvideo.layers.linear import LinearBase - # LTX-2 switches its active subset by stage at runtime. MiniMax-H3 - # uses the fixed main-transformer FFN set selected by its prefix. - if isinstance(layer, LinearBase) and (is_ltx2_nvfp4_linear_prefix(prefix) - or is_minimax_h3_nvfp4_linear_prefix(prefix)): + if not isinstance(layer, LinearBase): + return None + if self.layer_profile == "h3_dit": + tagged = is_minimax_h3_nvfp4_dit_linear_prefix(prefix) + elif self.layer_profile == "h3_dit_vsa": + tagged = (is_minimax_h3_nvfp4_dit_linear_prefix(prefix) + or _MINIMAX_H3_NVFP4_VSA_GATE_PREFIX.search(prefix) is not None) + elif self.layer_profile == "h3_dit_ffn": + tagged = is_minimax_h3_nvfp4_linear_prefix(prefix) + else: + tagged = is_ltx2_nvfp4_linear_prefix(prefix) or is_minimax_h3_nvfp4_linear_prefix(prefix) + if tagged: method = NVFP4QuantizeMethod(layer_prefix=prefix) method._retain_original_weights = self.retain_original_weights return method + if (self.layer_profile == "h3_dit_ffn" and envs.FASTVIDEO_H3_FP8_ATTENTION.get() + and _H3_BLOCK_ATTN_PROJ.search(prefix) is not None): + # Mixed precision: NVFP4 MLPs, FP8 (per-tensor weight, dynamic per-tensor activation) attention. + from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod + return FP8QuantizeMethod(granularity=envs.FASTVIDEO_H3_FP8_GRANULARITY.get()) return None @@ -534,9 +679,8 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None: if retain: retained += 1 elif isinstance(weight, DTensor): - # ponytail: purging FSDP-sharded originals needs per-shard - # resharding bookkeeping; skip until a sharded deploy needs it. - retained += 1 + raise RuntimeError("NVFP4 cannot purge FSDP-sharded bf16 weights. Use a packed NVFP4 " + "export, or convert without FSDP sharding.") else: purged_bytes += weight.numel() * weight.element_size() purged += 1 @@ -553,10 +697,95 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None: ) +def nvfp4_linear_weight_param_names(model: torch.nn.Module) -> set[str]: + """State-dict names of ``weight`` on layers tagged with ``NVFP4QuantizeMethod``.""" + names: set[str] = set() + for module_name, module in model.named_modules(): + if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod): + names.add(f"{module_name}.weight" if module_name else "weight") + return names + + +def _module_by_nvfp4_export_prefix(modules: dict[str, torch.nn.Module], prefix: str) -> torch.nn.Module | None: + module = modules.get(prefix) + if module is not None: + return module + if prefix.startswith("minimax_h3."): + return modules.get(prefix[len("minimax_h3."):]) + return modules.get(f"minimax_h3.{prefix}") + + +def load_minimax_h3_nvfp4_dit_export( + model: torch.nn.Module, + path: str, + device: torch.device | str, +) -> int: + """Load a packed NVFP4H3 DiT export onto already-tagged NVFP4 linears. + + Keys are ``::`` with the four buffers + ``convert_model_to_nvfp4`` registers, plus an optional calibrated + ``_nvfp4_input_global_sf``. Every export prefix must match an + NVFP4 linear, and every NVFP4 linear must appear in the export. + """ + from safetensors import safe_open + + modules = dict(model.named_modules()) + tagged = { + name + for name, module in modules.items() if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod) + } + groups: dict[str, dict[str, str]] = {} + with safe_open(path, framework="pt", device="cpu") as reader: + for key in reader.keys(): # noqa: SIM118 + if H3_NVFP4_DIT_KEY_SEP not in key: + raise ValueError(f"MiniMax-H3 NVFP4 DiT export key {key!r} is missing {H3_NVFP4_DIT_KEY_SEP!r}") + prefix, buffer_name = key.split(H3_NVFP4_DIT_KEY_SEP, 1) + groups.setdefault(prefix, {})[buffer_name] = key + + loaded_names: set[str] = set() + for prefix, buffers in groups.items(): + missing_buffers = [name for name in H3_NVFP4_DIT_BUFFER_NAMES if name not in buffers] + if missing_buffers: + raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} is missing {missing_buffers}") + extra_buffers = sorted(set(buffers) - set(H3_NVFP4_DIT_BUFFER_NAMES) - {H3_NVFP4_DIT_INPUT_SF_NAME}) + if extra_buffers: + raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} has unknown buffers {extra_buffers}") + module = _module_by_nvfp4_export_prefix(modules, prefix) + if module is None: + raise ValueError(f"MiniMax-H3 NVFP4 DiT export layer {prefix!r} is not in the model") + if not isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod): + raise RuntimeError("MiniMax-H3 NVFP4 DiT export layer " + f"{prefix!r} is not an NVFP4 linear; set NVFP4Config(layer_profile='h3_dit')") + for buffer_name in H3_NVFP4_DIT_BUFFER_NAMES + (H3_NVFP4_DIT_INPUT_SF_NAME, ): + if buffer_name not in buffers: + continue + tensor = reader.get_tensor(buffers[buffer_name]).to(device=device) + module.register_buffer(buffer_name, tensor, persistent=False) + module.register_parameter("weight", None) + loaded_names.add(next(name for name, candidate in modules.items() if candidate is module)) + + missing_layers = tagged - loaded_names + extra_layers = loaded_names - tagged + if missing_layers or extra_layers: + raise RuntimeError("MiniMax-H3 NVFP4 DiT export does not cover the tagged linear set; " + f"missing={sorted(missing_layers)[:8]} extra={sorted(extra_layers)[:8]}") + logger.info("Loaded MiniMax-H3 NVFP4 DiT export: %d linears from %s", len(loaded_names), path) + return len(loaded_names) + + __all__ = [ + "H3_NVFP4_DIT_BUFFER_NAMES", + "H3_NVFP4_DIT_EXPORT_FILENAME", + "H3_NVFP4_DIT_INPUT_SF_NAME", "NVFP4Config", "NVFP4QuantizeMethod", "convert_model_to_nvfp4", + "dense_transformer_safetensors", + "find_minimax_h3_nvfp4_dit_export", "is_ltx2_nvfp4_linear_prefix", + "is_minimax_h3_nvfp4_dit_export_path", + "is_minimax_h3_nvfp4_dit_linear_prefix", "is_minimax_h3_nvfp4_linear_prefix", + "load_minimax_h3_nvfp4_dit_export", + "nvfp4_linear_weight_param_names", ] diff --git a/fastvideo/layers/quantization/nvfp4_dequant.py b/fastvideo/layers/quantization/nvfp4_dequant.py new file mode 100644 index 000000000..5bc3c2b34 --- /dev/null +++ b/fastvideo/layers/quantization/nvfp4_dequant.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One-pass serialized NVFP4 weight expansion for BF16 consumer-GPU compute.""" +import torch +import triton +import triton.language as tl + + +@triton.jit +def _dequantize_nvfp4(P, S, OUT, N: tl.constexpr, K: tl.constexpr, INVERSE_SCALE, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = offsets < N * K + row, col = offsets // K, offsets % K + packed = tl.load(P + row * (K // 2) + col // 2, valid, other=0).to(tl.uint32) + code = (packed >> ((col % 2) * 4)) & 15 + magnitude = (code & 7).to(tl.float32) + value = tl.where(magnitude < 4, magnitude * 0.5, tl.where(magnitude < 6, magnitude - 2, magnitude * 2 - 8)) + value = value * tl.where((code & 8) != 0, -1.0, 1.0) + group = col // 16 + # FlashInfer layout_128x4: [row_tile, col_tile, row%32, row//32%4, col%4]. + scale_index = ((((row // 128) * (K // 64) + group // 4) * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4 + scale = tl.load(S + scale_index, valid, other=0.0).to(tl.float32) + output = (value * scale) * INVERSE_SCALE + tl.store(OUT + offsets, output, valid) + + +def dequantize_nvfp4_cuda(packed: torch.Tensor, + scales: torch.Tensor, + global_scale: float, + dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Expand E2M1 nibbles and swizzled E4M3 scales without full FP32 intermediates.""" + if not packed.is_cuda or scales.device != packed.device: + raise ValueError("NVFP4 fused dequantization requires tensors on the same CUDA device") + if packed.ndim != 2 or packed.dtype != torch.uint8 or scales.dtype != torch.uint8: + raise ValueError("NVFP4 fused dequantization requires packed uint8 weights and scales") + if not packed.is_contiguous() or not scales.is_contiguous(): + raise ValueError("NVFP4 fused dequantization requires contiguous tensors") + rows, cols = packed.shape[0], packed.shape[1] * 2 + if rows % 128 or cols % 64 or scales.numel() != rows * cols // 16: + raise ValueError("NVFP4 fused dequantization requires exact 128x4 scale geometry") + if dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError("NVFP4 fused dequantization requires a floating output dtype") + output = torch.empty((rows, cols), dtype=dtype, device=packed.device) + _dequantize_nvfp4[(triton.cdiv(rows * cols, 1024), )]( + packed, + scales.view(torch.float8_e4m3fn), + output, + rows, + cols, + # Match Torch's CPU-scalar division: form the + # reciprocal in double, then cast to FP32. + 1.0 / global_scale, + BLOCK=1024, + num_warps=4) + return output diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index e25650466..dd9a645a1 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -4,6 +4,7 @@ from __future__ import annotations import contextlib +import json import math from collections.abc import Iterator, Mapping from typing import Any @@ -14,6 +15,8 @@ import torch.nn.functional as F from fastvideo import envs from fastvideo.attention import DistributedAttention +from fastvideo.attention.backends.abstract import layer_idx_from_prefix +from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadata from fastvideo.attention.layer import DistributedAttention_VSA from fastvideo.attention.selector import get_attn_backend from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config @@ -21,7 +24,9 @@ from fastvideo.distributed.communication_op import ( sequence_model_parallel_all_gather_with_unpad, sequence_model_parallel_shard, ) -from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized +from fastvideo.distributed.parallel_state import (get_sp_group, get_sp_world_size, + model_parallel_is_initialized) +from fastvideo.forward_context import get_forward_context from fastvideo.layers.linear import ReplicatedLinear from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod from fastvideo.layers.mlp import MLP @@ -30,6 +35,8 @@ from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp, + vsa_fp4_requested, vsa_tile_first_attention) from fastvideo.models.dits.minimax_h3_fusions import ( HAVE_TRITON, fused_qknorm_rope, @@ -74,6 +81,13 @@ def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool: return HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() +@torch.compile(dynamic=True, fullgraph=True) +def _gated_residual(hidden_states: torch.Tensor, gate_table: torch.Tensor, indices: torch.Tensor, + branch: torch.Tensor) -> torch.Tensor: + """``hidden + gate[indices] * branch`` in one pass (eager materializes the gathered gate).""" + return hidden_states + gate_table.index_select(0, indices) * branch + + class MiniMaxH3RotaryPosEmbed(nn.Module): """Three-axis rotary frequencies over packed `(t, h, w)` coordinates.""" @@ -125,8 +139,22 @@ class MiniMaxH3FeedForward(nn.Module): self.fuse_swiglu = fuse_swiglu self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance( self.fc_out.quant_method, MXFP8QuantizeMethod) + # Inference-only token chunking: the 2 * ffn_dim intermediate is ~5.3x the block input + # (4.5 GiB at 78k tokens), so chunks bound the activation peak on 24-32 GB GPUs. + self.chunk_tokens = envs.FASTVIDEO_H3_FFN_CHUNK_TOKENS.get() def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + tokens = hidden_states.shape[-2] if hidden_states.dim() > 1 else 0 + if (self.chunk_tokens and tokens > self.chunk_tokens and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() and hidden_states.numel() == tokens * hidden_states.shape[-1]): + out = torch.empty_like(hidden_states) + for start in range(0, tokens, self.chunk_tokens): + rows = slice(start, start + self.chunk_tokens) + out[..., rows, :] = self._forward(hidden_states[..., rows, :].contiguous()) + return out + return self._forward(hidden_states) + + def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.use_mxfp8: from fastvideo.layers.mxfp8linear import mxfp8_swiglu_feed_forward @@ -209,6 +237,11 @@ class MiniMaxH3Attention(nn.Module): prefix=prefix, fa4_packed_varlen=fa4_packed_varlen, ) + # 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_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 # skips the gate branch entirely while it is structurally zero. @@ -255,6 +288,12 @@ class MiniMaxH3Attention(nn.Module): return if self._gate_compress_active is None: weight = self.to_gate_compress.weight + if weight is None: + # Packed NVFP4 gate: any nonzero E2M1 magnitude (bits 0x7 of + # either nibble) makes the branch live. + packed = self.to_gate_compress._nvfp4_weight + self._gate_compress_active = bool((packed & 0x77).any()) + return # bool() on a DTensor reduction resolves collectively, so every # rank caches the same answer. self._gate_compress_active = bool((weight != 0).any()) @@ -282,9 +321,41 @@ class MiniMaxH3Attention(nn.Module): rotary_emb: tuple[torch.Tensor, torch.Tensor] | None, original_seq_len: int, ) -> torch.Tensor: - query, _ = self.to_q(hidden_states) - key, _ = self.to_k(hidden_states) - value, _ = self.to_v(hidden_states) + if (self._vsa_fp4 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 + # FP4 kernel relies on to start each row from a finite running max. + if isinstance(meta, MiniMaxH3VSAMetadata) and meta.exempt: + use_fused_rope = self.fuse_qknorm_rope and _can_run_minimax_h3_fusion(hidden_states) + if not model_parallel_is_initialized() or get_sp_world_size() == 1: + hidden_states = vsa_fp4_attention(self, hidden_states, rotary_emb, meta, use_fused_rope) + else: + hidden_states = vsa_fp4_attention_sp(self, hidden_states, rotary_emb, meta, use_fused_rope, + get_sp_group()) + with STAGES.span("out_proj"): + hidden_states, _ = self.to_out(hidden_states) + return hidden_states + if (self._vsa_tile_first and hidden_states.is_cuda and rotary_emb is not None + and not torch.is_grad_enabled() and not torch.compiler.is_compiling() + and (not model_parallel_is_initialized() or get_sp_world_size() == 1)): + meta = get_forward_context().attn_metadata + if isinstance(meta, MiniMaxH3VSAMetadata) and meta.tile_elems == 64: + use_fused_rope = self.fuse_qknorm_rope and _can_run_minimax_h3_fusion(hidden_states) + hidden_states = vsa_tile_first_attention(self, hidden_states, rotary_emb, meta, use_fused_rope) + with STAGES.span("out_proj"): + hidden_states, _ = self.to_out(hidden_states) + return hidden_states + with STAGES.span("qkv_proj"): + # All three projections see the same activations. Reuse their FP8 + # quantization when the loaded methods have identical granularity. + from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections + if not torch.is_grad_enabled() and not torch.compiler.is_compiling(): + query, key, value = _shared_input_projections((self.to_q, self.to_k, self.to_v), hidden_states) + else: + query, _ = self.to_q(hidden_states) + key, _ = self.to_k(hidden_states) + value, _ = self.to_v(hidden_states) query = query.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) key = key.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) value = value.unflatten(-1, (self.num_attention_heads, self.attention_head_dim)) @@ -308,16 +379,18 @@ class MiniMaxH3Attention(nn.Module): gate_compress, _ = self.to_gate_compress(hidden_states) extra_attention_kwargs["gate_compress"] = gate_compress.unflatten( -1, (self.num_attention_heads, self.attention_head_dim)) - hidden_states, _ = self.distributed_attention( - query, - key, - value, - original_seq_len=original_seq_len, - freqs_cis=None, - **extra_attention_kwargs, - ) + with STAGES.span("attention"): + hidden_states, _ = self.distributed_attention( + query, + key, + value, + original_seq_len=original_seq_len, + freqs_cis=None, + **extra_attention_kwargs, + ) hidden_states = hidden_states.flatten(2, 3).type_as(query) - hidden_states, _ = self.to_out(hidden_states) + with STAGES.span("out_proj"): + hidden_states, _ = self.to_out(hidden_states) return hidden_states @@ -421,7 +494,51 @@ class MiniMaxH3AdaLayerNormModulation(nn.Module): prefix=f"{prefix}.linear", ) + def enable_host_cache(self, table: dict | None = None) -> None: + """Keep the projection in pinned host memory and cache its output per timestep set. + + The modulation is a pure function of the timestep embedding, and few-step checkpoints sample a fixed + timestep ladder, so each block's output is a small constant table. The weights (the largest bf16 tensors + in the DiT) then never occupy device memory: a cache miss copies them in for one matmul. + """ + weight, bias = self.linear.weight, self.linear.bias + if table is not None and (weight is None or bias is None): + # Table-only load skipped these tensors entirely. + self._host_weight = self._host_bias = None + self._modulation_cache = dict(table) + self._cache_key = None + return + weight, bias = weight.data, bias.data + if table is None: + self._host_weight = weight.to("cpu").pin_memory() + self._host_bias = bias.to("cpu").pin_memory() + else: + # Precomputed modulation for a fixed timestep ladder: the projection weights are not needed at all. + self._host_weight = self._host_bias = None + self.linear.weight.data = torch.empty(0, dtype=weight.dtype) + self.linear.bias.data = torch.empty(0, dtype=bias.dtype) + self._modulation_cache: dict[Any, torch.Tensor] = dict(table or {}) + self._cache_key: Any = None + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: + cache = getattr(self, "_modulation_cache", None) + if cache is not None: + out = cache.get(self._cache_key) + if out is None: + if self._host_weight is None: + raise RuntimeError(f"No precomputed AdaLN modulation for timestep key {self._cache_key}; the " + "table only covers the checkpoint's fixed ladder. Load without " + "FASTVIDEO_H3_ADALN_TABLE to use other timesteps.") + x = F.silu(temb) if self.apply_silu else temb + weight = self._host_weight.to(temb.device, non_blocking=True) + bias = self._host_bias.to(temb.device, non_blocking=True) + out = F.linear(x.to(weight.dtype), weight, bias) + if self._cache_key is not None: + cache[self._cache_key] = out + if envs.FASTVIDEO_H3_ADALN_DUMP.get(): + # Projection inputs, kept only when dumping, for offline low-rank fits. + self.__dict__.setdefault("_modulation_inputs", {})[self._cache_key] = x.detach() + return out.view(-1, 6 * self.hidden_size).chunk(6, dim=-1) if self.apply_silu: temb = F.silu(temb) temb, _ = self.linear(temb.to(self.linear.weight.dtype)) @@ -546,6 +663,9 @@ class MiniMaxH3TransformerBlock(nn.Module): 1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices) with nvtx_range("minimax_h3.transformer_block.self_attention"): attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len) + # The attention input is dead now. Keeping it until assignment below + # overlaps three full-width activations during the residual fusion. + del norm_hidden_states if use_modulate_fusion: with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( @@ -564,8 +684,11 @@ class MiniMaxH3TransformerBlock(nn.Module): norm_hidden_states = self.norm2(hidden_states) norm_hidden_states = norm_hidden_states * ( 1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices) - with nvtx_range("minimax_h3.transformer_block.feed_forward"): + del attention_output + with nvtx_range("minimax_h3.transformer_block.feed_forward"), STAGES.span("feed_forward"): feed_forward_output = self.ff(norm_hidden_states) + if use_modulate_fusion and not torch.compiler.is_compiling(): + return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output) return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output @@ -891,6 +1014,55 @@ class MiniMaxH3Transformer3DModel(BaseDiT): self._rope_cache = (position_ids, dtype, value) return value + def attach_step_splice(self, late: nn.Module, from_step: int) -> None: + """Hand denoising steps ``from_step`` onward to ``late`` (same architecture, other weights). + + Early DMD steps fix layout and object count, late ones texture and detail, so two checkpoints + can split the trajectory. ``late`` is kept out of this module's children: it is placed, offloaded + and checkpointed on its own. + """ + object.__setattr__(self, "_splice_late", late) + self._splice_from_step = int(from_step) + + def enable_adaln_host_cache(self, table_path: str | None = None) -> None: + """Move every block's AdaLN projection to pinned host memory behind a per-timestep cache. + + With ``table_path`` (written by FASTVIDEO_H3_ADALN_DUMP), the cache is prefilled from precomputed + modulation tables and the projection weights are dropped entirely. + """ + tables = None + if table_path: + import ast + raw = torch.load(table_path, map_location="cpu") + tables = {int(i): {ast.literal_eval(k): v for k, v in blk.items()} for i, blk in raw.items()} + for index, block in enumerate(self.transformer_blocks): + block.adaln_proj.enable_host_cache(None if tables is None else tables[index]) + self._adaln_host_cache = True + self._adaln_dumped_entries = -1 + + def _move_adaln_tables(self, device: torch.device) -> None: + for block in self.transformer_blocks: + cache = block.adaln_proj._modulation_cache + for key, value in cache.items(): + if value.device != device: + cache[key] = value.to(device) + + def _maybe_dump_adaln_tables(self) -> None: + path = envs.FASTVIDEO_H3_ADALN_DUMP.get() + if not path: + return + entries = sum(len(b.adaln_proj._modulation_cache) for b in self.transformer_blocks) + if entries == self._adaln_dumped_entries: + return + self._adaln_dumped_entries = entries + if model_parallel_is_initialized() and get_sp_group().rank_in_group != 0: + return + torch.save({i: {repr(k): v.detach().cpu() for k, v in b.adaln_proj._modulation_cache.items()} + for i, b in enumerate(self.transformer_blocks)}, path) + first = self.transformer_blocks[0].adaln_proj.__dict__.get("_modulation_inputs", {}) + torch.save({repr(k): v.detach().cpu() for k, v in first.items()}, path + ".inputs") + logger.info("Dumped AdaLN modulation tables (%d entries) to %s", entries, path) + def _refined_text(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: """The prompt embedding is constant across the denoising loop and the refiner blocks are timestep-free, so refine once per generation @@ -930,6 +1102,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT): text_indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Predict video and audio velocities from one caller-defined packed layout.""" + late = self.__dict__.get("_splice_late") + if late is not None and get_forward_context().current_timestep >= self._splice_from_step: + return late(hidden_states, audio_hidden_states, encoder_hidden_states, timestep, timestep_indices, + token_tags, position_ids, video_indices, audio_indices, text_indices) if position_ids.ndim != 2 or position_ids.shape[-1] != 3: raise ValueError(f"position_ids must have shape (seq_len, 3), got {tuple(position_ids.shape)}.") sequence_length = position_ids.shape[0] @@ -972,10 +1148,20 @@ class MiniMaxH3Transformer3DModel(BaseDiT): local_timestep_indices, _ = sequence_model_parallel_shard(local_timestep_indices, dim=0) rotary_emb = (rotary_cos, rotary_sin) + if getattr(self, "_adaln_host_cache", False): + # One host read per forward keys every block's modulation cache by the timestep values. + key = (tuple(timestep.reshape(-1).tolist()), tuple(temb.shape), str(temb.dtype)) + for block in self.transformer_blocks: + block.adaln_proj._cache_key = key + self._move_adaln_tables(temb.device) + # The eager driver owns profiling markers while each block's compiled # forward owns the graph that the marker surrounds. for block_index, block in enumerate(self.transformer_blocks): - with nvtx_range(f"minimax_h3.transformer_block.{block_index}"): + if STAGES.enabled: + logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index, + torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30) + with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"): packed_hidden_states = block( packed_hidden_states, temb, @@ -997,6 +1183,12 @@ class MiniMaxH3Transformer3DModel(BaseDiT): video_output = video_output.index_select(1, video_indices) audio_output = audio_output.index_select(1, audio_indices) + if getattr(self, "_adaln_host_cache", False): + self._maybe_dump_adaln_tables() + if STAGES.enabled: + stages = STAGES.flush() + if not model_parallel_is_initialized() or get_sp_group().rank_in_group == 0: + logger.info("H3_STAGE_MS %s", json.dumps(stages)) return video_output, audio_output diff --git a/fastvideo/models/dits/minimax_h3_vsa_fp4.py b/fastvideo/models/dits/minimax_h3_vsa_fp4.py new file mode 100644 index 000000000..6c18e150a --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_vsa_fp4.py @@ -0,0 +1,440 @@ +# 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 +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 +masks on the kernel's 128x128 blocks. + +The attention input is gathered into tile order once per block (one +``hidden_size``-wide pass; pad rows stay zero, so q/k/v pad rows are exactly +zero through the bias-free projections, RMSNorm and RoPE). That replaces the +generic path's concat, four tile scatters and three transposes, and lets q, k +and v share one activation quantization. The output returns to packed order +with one gather before ``to_out``. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch + +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: + """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. + """ + try: + return api.sageattn_blackwell_sparse_bshd(*args, validate=False) + except TypeError as error: + if "validate" not in str(error): + raise + return api.sageattn_blackwell_sparse_bshd(*args) + +_BLOCK = 128 + + +class _StageTimer: + """Opt-in (``FASTVIDEO_H3_SP_PROFILE=1``) CUDA-event spans summed per stage over one DiT forward. + + Events are only recorded during the forward; ``flush`` synchronizes once + and returns milliseconds per stage, so enabling it costs one sync per step. + """ + + def __init__(self) -> None: + self.enabled = envs.FASTVIDEO_H3_SP_PROFILE.get() + 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) + + def _record(self, name: str): + import contextlib + + @contextlib.contextmanager + def ctx(): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + yield + end.record() + self._spans.append((name, start, end)) + + return ctx() + + def flush(self) -> dict[str, float]: + torch.cuda.synchronize() + totals: dict[str, float] = {} + for name, start, end in self._spans: + totals[name] = totals.get(name, 0.0) + start.elapsed_time(end) + self._spans.clear() + return {k: round(v, 1) for k, v in totals.items()} + + +STAGES = _StageTimer() + +_fp4_api: Any = None + + +def vsa_fp4_requested() -> bool: + return envs.FASTVIDEO_H3_VSA_FP4.get() + + +def _api() -> Any: + global _fp4_api + if _fp4_api is None: + import attn_qat_infer.api as api + _fp4_api = api + return _fp4_api + + +class _TileLayout: + """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: + self.tile = int(meta.tile_elems) + self.n_tiles = int(meta.variable_block_sizes.numel()) + self.rows = math.ceil(self.n_tiles * self.tile / _BLOCK) * _BLOCK + self.untile = meta.untile_combined_index + self.row_tile = self.untile // self.tile + self.rotary_src = rotary_emb + cos, sin = rotary_emb + self.cos = cos.new_zeros((self.rows, cos.shape[-1])).index_copy_(0, self.untile, cos) + self.sin = sin.new_zeros((self.rows, sin.shape[-1])).index_copy_(0, self.untile, sin) + self._buf: torch.Tensor | None = None + + def gather_in(self, x: torch.Tensor) -> torch.Tensor: + """Packed ``[B, L, C]`` -> tile-ordered ``[B, rows, C]``; pad rows stay zero. + + The buffer is reused across blocks: pad rows are never written and + every valid row is overwritten, and each block consumes it (q/k/v + projections) before the next block refills it. + """ + shape = (x.shape[0], self.rows, x.shape[-1]) + if self._buf is None or self._buf.shape != shape or self._buf.dtype != x.dtype: + self._buf = x.new_zeros(shape) + return self._buf.index_copy_(1, self.untile, x) + + +def _layout_for(meta: MiniMaxH3VSAMetadata, rotary_emb: tuple[torch.Tensor, torch.Tensor]) -> _TileLayout: + layout = getattr(meta, "_h3_fp4_layout", None) + if layout is None or layout.rotary_src[0] is not rotary_emb[0]: + layout = _TileLayout(meta, rotary_emb) + meta._h3_fp4_layout = layout # type: ignore[attr-defined] + return layout + + +def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list[torch.Tensor]: + """Share compatible FP8 preparation or unit-scale NVFP4 activation quantization. + + Calibrated or dynamic NVFP4 activation scales retain independent preparation. + """ + from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod + from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod + + methods = [getattr(linear, "quant_method", None) for linear in linears] + if any(m is None for m in methods): # LoRA wrappers keep quant_method on base_layer. + return [linear(x)[0] for linear in linears] + same_nvfp4 = all(type(m) is NVFP4QuantizeMethod and m.wants_prequantized_input() + and m.uses_unit_activation_scale(linear) + for m, linear in zip(methods, linears, strict=True)) + same_fp8 = all(type(m) is FP8QuantizeMethod and m.granularity == methods[0].granularity for m in methods) + if not (same_nvfp4 or same_fp8): + return [linear(x)[0] for linear in linears] + pre = methods[0].quantize_input(x) + return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)] + + +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: + """Single-rank BF16 VSA with one input scatter instead of a Q/K/V/gate stack. + + The existing backend computes the same tile-64 mask, valid-key handling, + fine attention and compression branch. Bias-free projections keep pad + rows zero. This path is inference-only and keeps the checkpoint layout. + """ + layout = _layout_for(meta, rotary_emb) + logical = layout.n_tiles * layout.tile + heads, dim = attn.num_attention_heads, attn.attention_head_dim + with STAGES.span("tile_input"): + x_tiles = layout.gather_in(hidden_states)[:, :logical] + with STAGES.span("qkv_proj"): + query, key, value = (t.unflatten(-1, (heads, dim)) + for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles)) + with STAGES.span("qknorm_rope"): + cos, sin = layout.cos[:logical], layout.sin[:logical] + if use_fused_rope: + from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope + query = fused_qknorm_rope(query, attn.norm_q.weight, cos.to(query.dtype), sin.to(query.dtype), attn.norm_q.eps) + key = fused_qknorm_rope(key, attn.norm_k.weight, cos.to(key.dtype), sin.to(key.dtype), attn.norm_k.eps) + else: + query = attn._apply_rotary_emb(attn.norm_q(query), (cos, sin)) + key = attn._apply_rotary_emb(attn.norm_k(key), (cos, sin)) + gate = None + if attn.to_gate_compress is not None and attn._gate_active(): + with STAGES.span("gate_proj"): + gate, _ = attn.to_gate_compress(x_tiles) + gate = gate.unflatten(-1, (heads, dim)) + capture_root = envs.FASTVIDEO_H3_CAPTURE_QKV.get() + if capture_root and attn._layer_idx in (0, 20, 41): + from pathlib import Path + root = Path(capture_root) + root.mkdir(parents=True, exist_ok=True) + capture = root / f"layer-{attn._layer_idx}.pt" + if not capture.exists(): + q_pooled = _pool_tiles(query, meta.variable_block_sizes, meta.tile_elems) + k_pooled = _pool_tiles(key, meta.variable_block_sizes, meta.tile_elems) + scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / dim**0.5 + sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity + mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, + meta.video_tile_spans, meta.span_sparsities) + # Two heads keep the artifact small while retaining all real keys, + # query rows, per-tile selections and partial-tile validity. + torch.save({"q": query[:, :, :2].transpose(1, 2).contiguous().cpu(), + "k": key[:, :, :2].transpose(1, 2).contiguous().cpu(), + "v": value[:, :, :2].transpose(1, 2).contiguous().cpu(), + "mask": mask[:, :2].cpu(), "vbs": meta.variable_block_sizes.cpu(), + "untile": meta.untile_combined_index.cpu()}, capture) + del q_pooled, k_pooled, scores, mask + with STAGES.span("attention"): + out = attn.distributed_attention.attn_impl.forward(query, key, value, gate, meta) + with STAGES.span("untile_output"): + return out.index_select(1, layout.untile).flatten(2, 3) + + +def vsa_fp4_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], + meta: MiniMaxH3VSAMetadata, use_fused_rope: bool) -> torch.Tensor: + """Attention core for ``MiniMaxH3Attention``; returns the pre-``to_out`` ``[B, L, H*D]``.""" + api = _api() + layout = _layout_for(meta, rotary_emb) + heads, dim = attn.num_attention_heads, attn.attention_head_dim + with STAGES.span("qkv_proj_rope"): + x_tiles = layout.gather_in(hidden_states) + query, key, value = (t.unflatten(-1, (heads, dim)) + for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles)) + if use_fused_rope: + from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope + cos, sin = layout.cos.to(query.dtype), layout.sin.to(query.dtype) + query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps) + key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps) + else: + rope = (layout.cos, layout.sin) + query = attn._apply_rotary_emb(attn.norm_q(query), rope) + key = attn._apply_rotary_emb(attn.norm_k(key), rope) + + sim_fp8 = envs.FASTVIDEO_H3_SIM_SP_FP8.get() + if sim_fp8: + query, key, value = (_fp8_roundtrip(t) for t in (query, key, value)) + + vbs = meta.variable_block_sizes + logical = layout.n_tiles * layout.tile + with STAGES.span("select_mask"): + q_pooled = _pool_tiles(query[:, :logical], vbs, layout.tile) + k_pooled = _pool_tiles(key[:, :logical], vbs, layout.tile) + scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / (dim**0.5) + sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity + 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) + 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) + with STAGES.span("out_untile"): + out = out.transpose(1, 2).index_select(1, layout.untile) # [B, L, H, D], packed order + if sim_fp8: + out = _fp8_roundtrip(out) + + if attn.to_gate_compress is not None and attn._gate_active(): + with STAGES.span("gate_compress"): + 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))) + return out.flatten(2, 3) + + +# --------------------------------------------------------------------------- +# Ulysses sequence parallelism: FP8 head/sequence exchange around the FP4 core +# --------------------------------------------------------------------------- +# +# PCIe-only boxes (e.g. 8x RTX PRO 6000) move ~21 GB/s per GPU in an +# all-to-all, so the Ulysses exchange, not compute, bounds multi-GPU latency. +# q/k/v travel as FP8 (one scale per token and head): the attention kernel +# re-quantizes them to FP4 on arrival, so FP8 transport adds error far below +# that floor while halving the bytes. The VSA gate never travels: each rank +# applies it to its own sequence rows after an all-gather of the tiny per-tile +# compression output. The attention output returns as FP8 too; ``to_out`` +# re-quantizes it to FP4. + +_FP8 = torch.float8_e4m3fn +_FP8_MAX = 448.0 + +# Debug: apply the SP path's FP8 rounding on one GPU (q/k/v after RoPE and the +# attention output, one scale per token and head), to separate exchange +# rounding from sharding errors when comparing SP>1 against SP=1. + + +@torch.compile(dynamic=True, fullgraph=True) +def _fp8_roundtrip(x: torch.Tensor) -> torch.Tensor: + """``[..., H, D]`` BF16 -> FP8 with a per-(token, head) scale -> BF16, as the SP exchange does.""" + xf = x.float() + scale = (xf.abs().amax(dim=-1) / _FP8_MAX).clamp_min(1e-12) + return ((xf / scale[..., None]).to(_FP8).float() * scale[..., None]).to(torch.bfloat16) + + +@torch.compile(dynamic=False, fullgraph=True) +def _pack_heads_fp8(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + world: int) -> tuple[torch.Tensor, torch.Tensor]: + """``[rows, H, D]`` x3 -> per-destination ``[W, 3, rows, H/W, D]`` FP8 payload and scales.""" + x = torch.stack((query, key, value)).float() + scale = (x.abs().amax(dim=-1) / _FP8_MAX).clamp_min(1e-12) + payload = (x / scale[..., None]).to(_FP8) + _, rows, heads, dim = x.shape + payload = payload.view(3, rows, world, heads // world, dim).permute(2, 0, 1, 3, 4).contiguous() + scale = scale.view(3, rows, world, heads // world).permute(2, 0, 1, 3).contiguous() + return payload, scale + + +@torch.compile(dynamic=False, fullgraph=True) +def _unpack_seq_fp8(payload: torch.Tensor, scale: torch.Tensor, seq_len: int) -> torch.Tensor: + """``[W, 3, rows, Hs, D]`` from every source rank -> ``[3, seq_len, Hs, D]`` BF16 in packed order.""" + x = payload.float() * scale[..., None] + world, _, rows, heads, dim = x.shape + return x.permute(1, 0, 2, 3, 4).reshape(3, world * rows, heads, dim)[:, :seq_len].to(torch.bfloat16) + + +@torch.compile(dynamic=False, fullgraph=True) +def _pack_seq_fp8(out_bhsd: torch.Tensor, untile: torch.Tensor, world: int, + rows: int) -> tuple[torch.Tensor, torch.Tensor]: + """Tile-ordered ``[1, Hs, R, D]`` -> packed, padded, per-destination ``[W, rows, Hs, D]`` FP8 + scales.""" + x = out_bhsd[0].transpose(0, 1).index_select(0, untile).float() # [L, Hs, D] + x = torch.nn.functional.pad(x, (0, 0, 0, 0, 0, world * rows - x.shape[0])) + scale = (x.abs().amax(dim=-1) / _FP8_MAX).clamp_min(1e-12) + payload = (x / scale[..., None]).to(_FP8) + return payload.view(world, rows, *payload.shape[1:]), scale.view(world, rows, scale.shape[-1]) + + +@torch.compile(dynamic=False, fullgraph=True) +def _unpack_heads_fp8(payload: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """``[W, rows, Hs, D]`` from every head group -> ``[rows, W*Hs, D]`` BF16.""" + x = payload.float() * scale[..., None] + world, rows, heads, dim = x.shape + return x.permute(1, 0, 2, 3).reshape(rows, world * heads, dim).to(torch.bfloat16) + + +@torch.compile(dynamic=False, fullgraph=True) +def _apply_gate(out: torch.Tensor, out_c: torch.Tensor, row_tile: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """``out + out_c[row_tile] * gate`` for one rank's rows; out/gate ``[rows, H, D]``, out_c ``[n_tiles, H, D]``.""" + return out + out_c.index_select(0, row_tile) * gate + + +def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any) -> tuple[torch.Tensor, torch.Tensor]: + import torch.distributed as dist + recv = torch.empty_like(payload) + recv_scale = torch.empty_like(scale) + dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8), group=group) + dist.all_to_all_single(recv_scale, scale, group=group) + return recv, recv_scale + + +class _SPTileLayout: + """Per-step tile-order state for one rank's head subset.""" + + def __init__(self, meta: MiniMaxH3VSAMetadata, rank: int, local_rows: int) -> None: + self.tile = int(meta.tile_elems) + self.n_tiles = int(meta.variable_block_sizes.numel()) + self.rows = math.ceil(self.n_tiles * self.tile / _BLOCK) * _BLOCK + self.seq_len = int(meta.total_seq_length) + self.untile = meta.untile_combined_index + row_tile = self.untile // self.tile + local = torch.arange(rank * local_rows, (rank + 1) * local_rows, device=row_tile.device) + # Rows past the sequence are SP padding; their outputs are discarded. + self.local_row_tile = row_tile[local.clamp_max(self.seq_len - 1)] + self._buf: torch.Tensor | None = None + + def tiles_from(self, qkv: torch.Tensor) -> torch.Tensor: + """``[3, L, Hs, D]`` packed -> ``[3, rows, Hs, D]`` tile order with zero padding (reused buffer).""" + shape = (3, self.rows, *qkv.shape[2:]) + if self._buf is None or self._buf.shape != shape: + self._buf = qkv.new_zeros(shape) + return self._buf.index_copy_(1, self.untile, qkv) + + +def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], + meta: MiniMaxH3VSAMetadata, use_fused_rope: bool, sp_group: Any) -> torch.Tensor: + """Ulysses-SP attention core on local sequence rows ``[1, rows, C]``; returns pre-``to_out`` ``[1, rows, H*D]``.""" + import torch.distributed as dist + + api = _api() + world, rank = sp_group.world_size, sp_group.rank_in_group + heads, dim = attn.num_attention_heads, attn.attention_head_dim + local_rows = hidden_states.shape[1] + layout = getattr(meta, "_h3_fp4_sp_layout", None) + if layout is None: + layout = _SPTileLayout(meta, rank, local_rows) + meta._h3_fp4_sp_layout = layout # type: ignore[attr-defined] + + with STAGES.span("qkv_proj_rope"): + query, key, value = (t.unflatten(-1, (heads, dim)) + for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)) + if use_fused_rope: + from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope + cos, sin = rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype) + query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps) + key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps) + else: + query = attn._apply_rotary_emb(attn.norm_q(query), rotary_emb) + key = attn._apply_rotary_emb(attn.norm_k(key), rotary_emb) + + with STAGES.span("qkv_pack"): + payload, scale = _pack_heads_fp8(query[0], key[0], value[0], world) + with STAGES.span("qkv_all_to_all"): + payload, scale = _all_to_all(payload, scale, sp_group.device_group) + with STAGES.span("qkv_unpack_tile"): + qkv = layout.tiles_from(_unpack_seq_fp8(payload, scale, layout.seq_len)) # [3, R, Hs, D] + q_t, k_t, v_t = qkv[0:1], qkv[1:2], qkv[2:3] + + vbs = meta.variable_block_sizes + logical = layout.n_tiles * layout.tile + 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) + sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity + 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) + 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) + + with STAGES.span("out_pack"): + payload, scale = _pack_seq_fp8(out_bhsd, layout.untile, world, local_rows) + with STAGES.span("out_all_to_all"): + payload, scale = _all_to_all(payload, scale, sp_group.device_group) + with STAGES.span("out_unpack"): + out = _unpack_heads_fp8(payload, scale) # [rows, H, D] + + if attn.to_gate_compress is not None and attn._gate_active(): + with STAGES.span("gate_compress"): + v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile) + out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(out.dtype) # [Hs, n_tiles, D] + gathered = torch.empty((world, *out_c.shape), dtype=out_c.dtype, device=out_c.device) + with STAGES.span("gate_all_gather"): + dist.all_gather_into_tensor(gathered, out_c.contiguous(), group=sp_group.device_group) + with STAGES.span("gate_apply"): + out_c_all = gathered.flatten(0, 1).transpose(0, 1) # [n_tiles, H, D] + gate, _ = attn.to_gate_compress(hidden_states) + out = _apply_gate(out, out_c_all, layout.local_row_tile, gate[0].unflatten(-1, (heads, dim))) + return out.flatten(1, 2).unsqueeze(0) diff --git a/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py index 0c96e636c..933041e88 100644 --- a/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py +++ b/fastvideo/models/encoders/minimax_h3_checkpoint_nvfp4.py @@ -43,6 +43,8 @@ import torch from torch import nn from torch.nn.parameter import Parameter +import fastvideo.envs as envs + from fastvideo.distributed import get_tp_world_size from fastvideo.layers.linear import LinearBase, LinearMethodBase from fastvideo.layers.quantization.base_config import QuantizationConfig @@ -54,6 +56,9 @@ from fastvideo.layers.quantization.nvfp4_config import ( _require_flashinfer, ) from fastvideo.models.utils import set_weight_attrs +from fastvideo.logger import init_logger + +logger = init_logger(__name__) NVFP4_GROUP_SIZE = 16 NVFP4_SCALE_LAYOUT = "128x4" @@ -182,6 +187,38 @@ def _nvfp4_linear( ) +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def _fp4_gemm_supported(device: torch.device) -> bool: + """FP4 tensor-core GEMMs exist on Blackwell (sm_100 / sm_120) and newer.""" + return torch.cuda.is_available() and torch.cuda.get_device_capability(device)[0] >= 10 + + +def unswizzle_128x4_scales(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """FlashInfer ``layout_128x4`` block scales -> row-major ``[rows, cols]`` (E4M3 bytes). + + The swizzled buffer holds the padded matrix ``[ceil(rows/128)*128, ceil(cols/4)*4]`` as + ``(row_tile, col_tile, row % 32, (row // 32) % 4, col % 4)``. + """ + pad_rows, pad_cols = -(-rows // 128) * 128, -(-cols // 4) * 4 + tiles = scale.reshape(-1)[:pad_rows * pad_cols].view(pad_rows // 128, pad_cols // 4, 32, 4, 4) + return tiles.permute(0, 3, 2, 1, 4).reshape(pad_rows, pad_cols)[:rows, :cols] + + +def dequantize_serialized_nvfp4(weight_packed: torch.Tensor, weight_scale: torch.Tensor, global_scale: float, + dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Packed E2M1 ``[out, in // 2]`` + swizzled E4M3 scales -> dense ``[out, in]`` weight.""" + out_features, in_features = weight_packed.shape[0], weight_packed.shape[1] * 2 + lut = torch.tensor(_E2M1_VALUES, device=weight_packed.device, dtype=torch.float32) + packed = weight_packed.view(torch.uint8) + values = torch.stack((lut[(packed & 0x0F).long()], lut[(packed >> 4).long()]), dim=-1).reshape( + out_features, in_features) + scales = unswizzle_128x4_scales(weight_scale.view(torch.uint8), out_features, in_features // NVFP4_GROUP_SIZE) + scales = scales.view(torch.float8_e4m3fn).float().repeat_interleave(NVFP4_GROUP_SIZE, dim=1) + return (values * scales / global_scale).to(dtype) + + class MiniMaxH3SerializedNVFP4Config(QuantizationConfig): """Serialized 16-group NVFP4 contract for the H3 text encoder. @@ -267,15 +304,20 @@ class MiniMaxH3SerializedNVFP4Config(QuantizationConfig): raise RuntimeError(f"MiniMax-H3 serialized NVFP4 requires a CUDA device; got {device.type!r}") capability = torch.cuda.get_device_capability(device) capability_number = capability[0] * 10 + capability[1] - if capability_number < self.get_min_capability(): - raise RuntimeError("MiniMax-H3 serialized NVFP4 requires GPU capability " - f"sm{self.get_min_capability()} or newer, got sm{capability_number}") - if capability[0] not in (10, 12): - raise RuntimeError("MiniMax-H3 serialized NVFP4 runs FlashInfer's Blackwell FP4 GEMM; " - f"got unsupported sm{capability_number}") if get_tp_world_size() > 1: raise NotImplementedError("MiniMax-H3 serialized NVFP4 supports a single GPU: packed FP4 columns and " "128x4 swizzled scale rows cannot be narrowed per tensor-parallel rank") + if capability_number < self.get_min_capability(): + if capability_number < 80: + raise RuntimeError("MiniMax-H3 serialized NVFP4 needs bf16 compute (sm80+) for its de-quantized " + f"fallback, got sm{capability_number}") + logger.warning( + "MiniMax-H3 serialized NVFP4 on sm%d: no FP4 GEMM, each linear de-quantizes its weight to bf16 " + "per call", capability_number) + return + if capability[0] not in (10, 12): + raise RuntimeError("MiniMax-H3 serialized NVFP4 runs FlashInfer's Blackwell FP4 GEMM; " + f"got unsupported sm{capability_number}") sf_layout, _, _ = _require_flashinfer() if not hasattr(sf_layout, "layout_128x4"): raise RuntimeError("The installed flashinfer has no SfLayout.layout_128x4; MiniMax-H3 serialized NVFP4 " @@ -377,6 +419,11 @@ class MiniMaxH3SerializedNVFP4LinearMethod(LinearMethodBase): # ``mm_fp4`` folds both global scales into one multiplier. Activations use a # unit global scale, so the multiplier is the inverse weight global scale. device = weight_scale.device + # Serialized weights are immutable between post-load hooks. Keeping the + # validated scalar on the host avoids a CUDA synchronization per linear + # on the BF16 fallback used by consumer GPUs. + layer._nvfp4_dequant_global_scale = global_scale + layer._nvfp4_fused_dequant = envs.FASTVIDEO_H3_ENCODER_FUSED_DEQUANT.get() layer.register_buffer("_nvfp4_alpha", torch.tensor(1.0 / global_scale, dtype=torch.float32, device=device), persistent=False) layer.register_buffer("_nvfp4_x_global_scale", torch.ones((), dtype=torch.float32, device=device), @@ -385,12 +432,27 @@ class MiniMaxH3SerializedNVFP4LinearMethod(LinearMethodBase): @staticmethod def _apply_finalized(layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: x = _coerce_fp4_input_dtype(x) + if not _fp4_gemm_supported(layer.weight_packed.device): + # Pre-Blackwell GPUs have no FP4 GEMM: expand this layer's weight to bf16 for the one call. + # The encoder runs once per request, so the transient weight is cheaper than keeping a bf16 copy. + dequantize = dequantize_serialized_nvfp4 + if layer._nvfp4_fused_dequant: + from fastvideo.layers.quantization.nvfp4_dequant import dequantize_nvfp4_cuda + dequantize = dequantize_nvfp4_cuda + weight = dequantize(layer.weight_packed, layer.weight_scale, layer._nvfp4_dequant_global_scale, x.dtype) + return torch.nn.functional.linear(x, weight, None if bias is None else bias.to(x.dtype)) original_shape = x.shape if x.numel() == 0: # An empty prompt has nothing to quantize; the FP4 kernels are not defined for zero rows. return x.new_zeros(*original_shape[:-1], layer.output_size_per_partition, dtype=torch.bfloat16) - x_fp4, x_scale = _quantize_activation_nvfp4(x.reshape(-1, original_shape[-1]), layer._nvfp4_x_global_scale) - output = _nvfp4_linear(x_fp4, x_scale, layer.weight_packed, layer.weight_scale, layer._nvfp4_alpha) + x_global_scale, alpha = layer._nvfp4_x_global_scale, layer._nvfp4_alpha + if alpha.device != layer.weight_packed.device: + # A streamed encoder layer is finalized on the host and only its packed weights visit the GPU; + # FlashInfer requires the scalars on the GEMM's device. + x_global_scale = x_global_scale.to(layer.weight_packed.device, non_blocking=True) + alpha = alpha.to(layer.weight_packed.device, non_blocking=True) + x_fp4, x_scale = _quantize_activation_nvfp4(x.reshape(-1, original_shape[-1]), x_global_scale) + output = _nvfp4_linear(x_fp4, x_scale, layer.weight_packed, layer.weight_scale, alpha) if bias is not None: # The GEMM emits bf16; keep it that way whatever dtype the bias was built in. output = output + bias.to(output.dtype) diff --git a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py index c348140d1..4886ae1db 100644 --- a/fastvideo/models/encoders/minimax_h3_qwen3_vl.py +++ b/fastvideo/models/encoders/minimax_h3_qwen3_vl.py @@ -690,8 +690,17 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]): if (pixel_values_videos is None) != (video_grid_thw is None): raise ValueError("pixel_values_videos and video_grid_thw must be provided together") + stream_device = getattr(self, "_h3_encoder_layerwise_device", None) + if stream_device is not None and (pixel_values is not None or pixel_values_videos is not None): + raise ValueError("Layerwise H3 encoder currently supports text-only conditioning; " + "disable FASTVIDEO_H3_ENCODER_LAYERWISE for visual references") + input_ids = input_ids.unsqueeze(0) - inputs_embeds = self.language_model.embed_tokens(input_ids) + embedding_ids = input_ids.to("cpu") if stream_device is not None else input_ids + inputs_embeds = self.language_model.embed_tokens(embedding_ids) + if stream_device is not None: + inputs_embeds = inputs_embeds.to(stream_device) + input_ids = input_ids.to(stream_device) image_mask = None video_mask = None @@ -746,6 +755,24 @@ class MiniMaxH3Qwen3VLConditioner(TextEncoder[torch.Tensor]): raise RuntimeError(f"MiniMax-H3 language model returned unexpected shape={tuple(hidden_states.shape)}") return hidden_states[0] + def prepare_layerwise_offload(self, device: torch.device) -> None: + """Stream language layers for text-only CUDA inference, retaining embeddings on CPU.""" + if getattr(self, "_h3_encoder_layerwise_device", None) is not None: + return + if device.type != "cuda": + raise ValueError("Layerwise H3 encoder requires CUDA") + from fastvideo.distributed import get_tp_world_size + from fastvideo.hooks.layerwise_offload import enable_layerwise_offload + + if get_tp_world_size() != 1: + raise ValueError("Layerwise H3 encoder requires tensor parallel size 1") + self.to("cpu") + self.language_model.rotary_emb.to(device) + if self.language_model.norm is not None: + self.language_model.norm.to(device) + enable_layerwise_offload(self.language_model, resident_blocks=0, cyclic=False) + self._h3_encoder_layerwise_device = device + def forward( self, input_ids: torch.Tensor, diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 7e969d2a4..1647c6a48 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -13,6 +13,8 @@ from typing import Any, cast import torch import torch.distributed as dist + +import fastvideo.envs as envs import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file, safe_open from torch.distributed import init_device_mesh @@ -471,9 +473,23 @@ class TextEncoderLoader(ComponentLoader): # Explicitly move model to target device after loading weights model = model.to(target_device) + prepare_layerwise = getattr(model, "prepare_layerwise_offload", None) + if envs.FASTVIDEO_H3_ENCODER_LAYERWISE.get() and callable(prepare_layerwise): + if target_device.type != "cpu": + raise ValueError("Layerwise H3 encoder requires text_encoder_cpu_offload=True") + prepare_layerwise(runtime_device) + use_cpu_offload = False + logger.info("Enabled text-only layerwise H3 encoder with CPU token embeddings") + from fastvideo.platforms import current_platform - if use_cpu_offload: + if use_cpu_offload and checkpoint_quant_config is not None: + logger.info( + "Skipping FSDP CPU offload for serialized %s text encoder; " + "packed uint8 weights are not FSDP-shardable", + checkpoint_quant_config.get_name(), + ) + elif use_cpu_offload: pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): @@ -882,6 +898,14 @@ class VAELoader(ComponentLoader): # Find all safetensors files safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) + int8_convrot_path = None + if class_name == "AutoencoderKLMiniMaxH3": + from fastvideo.models.vaes.minimax_h3_int8_convrot import ( + dense_vae_safetensors, + find_int8_convrot_vae_path, + ) + int8_convrot_path = find_int8_convrot_vae_path(model_path) + safetensors_list = dense_vae_safetensors(safetensors_list) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") # Common case: a single `.safetensors` checkpoint file. @@ -918,6 +942,9 @@ 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) + 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 == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False) and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")): vae.optimize_memory_format() @@ -1124,6 +1151,8 @@ class TransformerLoader(ComponentLoader): safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) if not safetensors_list: raise ValueError(f"No safetensors files found in {model_path}") + from fastvideo.layers.quantization.nvfp4_config import dense_transformer_safetensors + safetensors_list = dense_transformer_safetensors(safetensors_list) # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts) update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None) @@ -1176,6 +1205,17 @@ class TransformerLoader(ComponentLoader): os.environ.get("RANK", "0"), resolved.name if resolved else "automatic selection", local_main_process_only=False) + # Layerwise offload keeps every block's weights in pinned host memory, so load them on the CPU and + # attach the hooks before anything moves to the GPU; loading on the GPU first would need the whole + # DiT resident once, which is exactly what offload exists to avoid on small cards. + # Scoped to H3: other models keep main's GPU load, which their quantization paths expect. + layerwise_load = (cls_name.startswith("MiniMaxH3") and fastvideo_args.inference_mode + and fastvideo_args.dit_layerwise_offload and not fastvideo_args.use_fsdp_inference) + # The AdaLN host cache also needs the projection weights to stay off the device from the start. + adaln_table = envs.FASTVIDEO_H3_ADALN_TABLE.get() or None + adaln_host_cache = (fastvideo_args.inference_mode and not fastvideo_args.use_fsdp_inference + and (envs.FASTVIDEO_H3_ADALN_CACHE.get() or adaln_table is not None)) + layerwise_load = layerwise_load or adaln_host_cache model = maybe_load_fsdp_model( model_cls=model_cls, init_params={ @@ -1183,7 +1223,7 @@ class TransformerLoader(ComponentLoader): "hf_config": hf_config }, weight_dir_list=safetensors_list, - device=get_local_torch_device(), + device=torch.device("cpu") if layerwise_load else get_local_torch_device(), hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim, hsdp_shard_dim=fastvideo_args.hsdp_shard_dim, strict=strict_load, @@ -1217,15 +1257,40 @@ class TransformerLoader(ComponentLoader): model = model.eval() + if adaln_host_cache and hasattr(model, "enable_adaln_host_cache"): + model.enable_adaln_host_cache(adaln_table) + logger.info("AdaLN modulation: %s", "precomputed tables from " + adaln_table if adaln_table + else "projections in pinned host memory behind a per-timestep cache") + if layerwise_load and not fastvideo_args.dit_layerwise_offload: + model = model.to(get_local_torch_device()) + if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload: # Check if model has nn.ModuleList for layerwise offload compatibility has_module_list = any(isinstance(m, nn.ModuleList) for m in model.children()) if has_module_list: enable_layerwise_offload(model) + # Blocks now hold placeholders; the remaining (non-block) weights and buffers belong on the GPU. + model = model.to(get_local_torch_device()) else: logger.warning( "Layerwise offload requested but model %s does not have " "nn.ModuleList structure. Skipping layerwise offload.", cls_name) + if layerwise_load: + model = model.to(get_local_torch_device()) + # FASTVIDEO_H3_SPLICE_TRANSFORMER=: a second checkpoint of the same architecture + # runs denoising steps FASTVIDEO_H3_SPLICE_FROM_STEP (default 4) onward. + # Only the primary ``transformer`` component splices; the spliced load itself never does. + splice_path = envs.FASTVIDEO_H3_SPLICE_TRANSFORMER.get() + if (splice_path and not getattr(self, "_loading_splice", False) and hasattr(model, "attach_step_splice") + and os.path.basename(os.path.normpath(model_path)) == "transformer"): + self._loading_splice = True + try: + late = self.load(splice_path, fastvideo_args) + finally: + self._loading_splice = False + from_step = envs.FASTVIDEO_H3_SPLICE_FROM_STEP.get() + model.attach_step_splice(late, from_step) + logger.info("Step splice: steps >= %d run the transformer from %s", from_step, splice_path) return model diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 956850a50..e56851026 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -36,6 +36,54 @@ def _summarize_param_names(names: set[str]) -> str: return ", ".join(f"{family} x{count}" if count > 1 else family for family, count in sorted(families.items())) +def _prequantized_fp8_prefixes(weight_files: list[str]) -> list[str]: + """Checkpoint prefixes stored as FP8 W8A8: float8_e4m3fn ``weight`` plus per-channel float32 ``weight_scale``.""" + from safetensors import safe_open + + fp8_weights: set[str] = set() + scales: set[str] = set() + for path in weight_files: + with safe_open(path, framework="pt", device="cpu") as reader: + for key in reader.keys(): # noqa: SIM118 + if key.endswith(".weight") and reader.get_slice(key).get_dtype() == "F8_E4M3": + fp8_weights.add(key[:-len(".weight")]) + elif key.endswith(".weight_scale") and reader.get_slice(key).get_dtype() == "F32": + scales.add(key[:-len(".weight_scale")]) + if fp8_weights - scales: + raise ValueError(f"FP8 checkpoint weights without a weight_scale: {sorted(fp8_weights - scales)[:4]}") + return sorted(fp8_weights) + + +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).""" + from safetensors import safe_open + + from fastvideo.layers.quantization.fp8_config import FP8QuantizeMethod + + modules = dict(model.named_modules()) + wanted = {f"{prefix}.{suffix}" for prefix in targets for suffix in ("weight", "weight_scale")} + tensors: dict[str, torch.Tensor] = {} + for path in weight_files: + with safe_open(path, framework="pt", device="cpu") as reader: + for key in wanted.intersection(reader.keys()): + tensors[key] = reader.get_tensor(key) + for prefix, target in targets.items(): + module = modules.get(target[:-len(".weight")]) + quant_method = getattr(module, "quant_method", None) + if not isinstance(quant_method, FP8QuantizeMethod): + raise RuntimeError(f"Pre-quantized FP8 weight {prefix!r} maps to {target!r}, which is not an FP8 linear; " + "load the checkpoint with an FP8 quant_config covering it.") + # The checkpoint scales are per output channel; activations then use per-token scales. + quant_method.granularity = "channel" + module.register_buffer("_fp8_weight", tensors[f"{prefix}.weight"].to(device).contiguous(), persistent=False) + module.register_buffer("_fp8_weight_scale", + 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)) + + def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lora_merge: bool = False) -> None: """Quantize inference linear weights after checkpoint loading. @@ -75,17 +123,28 @@ def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lor convert_model_to_mxfp8, ) + # NVFP4 may share a model with FP8 linears (e.g. NVFP4 FFN + FP8 attention); handle that pair before the + # per-module walk so the result does not depend on which quantized module comes first. + nvfp4_modules = [m for m in model.modules() if isinstance(getattr(m, "quant_method", None), NVFP4QuantizeMethod)] + if nvfp4_modules: + mixed_fp8 = any(isinstance(getattr(m, "quant_method", None), FP8QuantizeMethod) for m in model.modules()) + if any(getattr(module, "_nvfp4_weight", None) is not None for module in nvfp4_modules): + logger.info("NVFP4 packed export already populated; skipping runtime weight conversion") + elif defer_weight_conversion_until_lora_merge: + logger.info("Deferring NVFP4 weight conversion until the inference LoRA merge completes") + return + else: + logger.info("Converting loaded model weights for NVFP4 linear layers") + convert_model_to_nvfp4(model) + if mixed_fp8: + logger.info("Converting the FP8 linears of a mixed NVFP4/FP8 model") + convert_model_to_fp8(model) + return + qat_train_attached = 0 qat_train_skipped = 0 for mod in model.modules(): qm = getattr(mod, "quant_method", None) - if isinstance(qm, NVFP4QuantizeMethod): - if defer_weight_conversion_until_lora_merge: - logger.info("Deferring NVFP4 weight conversion until the inference LoRA merge completes") - return - logger.info("Converting loaded model weights for NVFP4 linear layers") - convert_model_to_nvfp4(model) - return if isinstance(qm, NVFP4QATQuantizeMethod): logger.info("Converting loaded model weights for NVFP4-QAT linear layers") convert_model_to_fp4(model) @@ -214,6 +273,22 @@ def maybe_load_fsdp_model( """ _validate_fsdp_inference_quantization(init_params, fsdp_inference) + from fastvideo.layers.quantization.nvfp4_config import ( + dense_transformer_safetensors, + find_minimax_h3_nvfp4_dit_export, + load_minimax_h3_nvfp4_dit_export, + nvfp4_linear_weight_param_names, + ) + packed_candidate = find_minimax_h3_nvfp4_dit_export(weight_dir_list) + 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) + if packed_nvfp4_export is not None and lora_path is not None: + raise ValueError("Packed MiniMax-H3 NVFP4 DiT export cannot be combined with lora_path; " + "merge the adapter before exporting, or load without the packed file.") + # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False) @@ -275,11 +350,19 @@ def maybe_load_fsdp_model( fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory) - # Host offload is already disabled on unified memory (GB10). Staging the - # 35B FastH3 DiT on CPU and then copying to CUDA doubled that working set - # and took minutes. Follow cpu_offload: read onto the accelerator. - weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=cpu_offload) - logger.info("Loading transformer weights with to_cpu=%s", cpu_offload) + nvfp4_skip_param_names: set[str] = set() + if packed_nvfp4_export is not None: + nvfp4_skip_param_names = nvfp4_linear_weight_param_names(model) + if not nvfp4_skip_param_names: + logger.warning( + "Found %s next to the transformer shards but no NVFP4 linears; " + "ignoring the packed export. Set NVFP4Config(layer_profile='h3_dit').", + packed_nvfp4_export, + ) + packed_nvfp4_export = None + load_weights_to_cpu = cpu_offload 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) dense_lora_patch = DenseLoRAPatch.from_adapter( lora_path, @@ -287,9 +370,6 @@ def maybe_load_fsdp_model( strength=lora_strength, ) if dense_lora_patch is not None: - # H3's compression gate is created only by the VSA attention backend. Loading a - # VSA student under dense attention would otherwise warn about 50 unmatched - # replacements and continue with a silently incomplete model. model_parameter_names = {name for name, _ in model.named_parameters()} missing_vsa_gates = sorted(name for name in dense_lora_patch.replacement_parameters if "gate_compress" in name and name not in model_parameter_names) @@ -298,6 +378,22 @@ def maybe_load_fsdp_model( "This LoRA adapter provides MiniMax H3 VSA compression gates, but the selected attention backend " "did not construct them. Use attention_backend='VIDEO_SPARSE_ATTN_H3'. Missing parameters: " + ", ".join(missing_vsa_gates[:3]) + (" ..." if len(missing_vsa_gates) > 3 else "")) + skip_param_names = set(nvfp4_skip_param_names) + fp8_targets: dict[str, str] = {} + # Pre-quantized FP8 checkpoints are an H3 release format; other models load as on main. + fp8_prefixes = (_prequantized_fp8_prefixes(weight_dir_list) + if type(model).__name__.startswith("MiniMaxH3") else []) + if fp8_prefixes: + fp8_targets = {prefix: param_names_mapping_fn(f"{prefix}.weight")[0] for prefix in fp8_prefixes} + skip_param_names |= set(fp8_targets.values()) + fp8_keys = {f"{prefix}.{suffix}" for prefix in fp8_prefixes for suffix in ("weight", "weight_scale")} + weight_iterator = ((name, tensor) for name, tensor in weight_iterator if name not in fp8_keys) + if envs.FASTVIDEO_H3_ADALN_TABLE.get(): + # Precomputed AdaLN modulation replaces the per-block projections; never read their weights. + skip_param_names |= {name for name, _ in model.named_parameters() + if re.fullmatch(r"transformer_blocks\.\d+\.adaln_proj\.linear\.(weight|bias)", name)} + logger.info("Skipping %d AdaLN projection tensors (precomputed modulation tables)", + sum("adaln_proj" in n for n in skip_param_names)) load_model_from_full_model_state_dict( model, weight_iterator, @@ -307,7 +403,12 @@ def maybe_load_fsdp_model( cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn, dense_lora_patch=dense_lora_patch, + skip_param_names=skip_param_names or None, ) + if packed_nvfp4_export is not None: + load_minimax_h3_nvfp4_dit_export(model, packed_nvfp4_export, device=device) + if fp8_targets: + _load_prequantized_fp8(model, weight_dir_list, fp8_targets, device=device) if hasattr(model, "materialize_non_persistent_buffers"): model.materialize_non_persistent_buffers(device=device, dtype=default_dtype) for n, p in chain(model.named_parameters(), model.named_buffers()): @@ -585,6 +686,12 @@ def shard_model( fully_shard(model, **root_kwargs) +def _drop_state_dict_parameter(model: nn.Module, param_name: str) -> None: + module_name, _, attr_name = param_name.rpartition(".") + module = model.get_submodule(module_name) if module_name else model + module.register_parameter(attr_name, None) + + # TODO(PY): device mesh for cfg parallel def load_model_from_full_model_state_dict( model: FSDPModule | torch.nn.Module, @@ -596,6 +703,7 @@ def load_model_from_full_model_state_dict( param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, training_mode: bool = True, dense_lora_patch: DenseLoRAPatch | None = None, + skip_param_names: set[str] | None = None, ) -> _IncompatibleKeys: """ Converting full state dict into a sharded state dict @@ -634,6 +742,11 @@ def load_model_from_full_model_state_dict( # set. for target_param_name in list(custom_param_sd): full_tensor = custom_param_sd.pop(target_param_name) + if skip_param_names and target_param_name in skip_param_names: + continue + if "::" in target_param_name: + logger.warning("Skipping packed NVFP4 export key mixed into dense shards: %s", target_param_name) + continue meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: # Some checkpoints include extra entries that are not part of the @@ -701,6 +814,13 @@ def load_model_from_full_model_state_dict( model.reverse_param_names_mapping = reverse_param_names_mapping unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) + skipped_unused = unused_keys & skip_param_names if skip_param_names else set() + for skipped_name in skipped_unused: + _drop_state_dict_parameter(model, skipped_name) + unused_keys -= skipped_unused + if skipped_unused: + logger.info("Deferred %d NVFP4 linear weights to the packed DiT export (%s)", + len(skipped_unused), _summarize_param_names(skipped_unused)) if unused_keys: # Say which of these the adapter is about to fill in. Reporting all of them as # "unloaded" was accurate when zero-init was the only outcome; with an adapter diff --git a/fastvideo/models/vaes/minimax_h3_int8_convrot.py b/fastvideo/models/vaes/minimax_h3_int8_convrot.py new file mode 100644 index 000000000..d9a12a01f --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_int8_convrot.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Comfy ``int8_tensorwise`` + ConvRot overlay for the MiniMax-H3 video VAE decoder. + +The export stores decoder transformer linears as signed int8 with per-output +channel scales and a JSON ``comfy_quant`` marker: + + {"format": "int8_tensorwise", "convrot": true, "convrot_groupsize": 256} + +Weights were rotated offline by a normalized regular Hadamard (group 256). +Inference rotates activations with the same matrix, row-quantizes them, then +runs int8 GEMM. Encoder convolutions stay dense; only the ViT decoder blocks +are quantized. + +Comfy names (``to_qkv``, ``ff.w1`` / ``ff.w2``, ``x_embedder``) are remapped +onto FastVideo's split Q/K/V and ``ff.net`` surface. +""" + +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors.torch import load_file as safetensors_load_file + +import fastvideo.envs as envs + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +INT8_CONVROT_FILENAME = "minimax_h3_video_vae_int8_convrot.safetensors" +_COM_FY_QUANT_SUFFIX = ".comfy_quant" +_DEFAULT_GROUP_SIZE = 256 + +_HADAMARD_CACHE: dict[tuple[int, str, torch.dtype], torch.Tensor] = {} + + +def regular_hadamard(size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Normalized regular Hadamard of order ``4**k`` (ConvRot Theorem 3.3).""" + cache_key = (size, str(device), dtype) + cached = _HADAMARD_CACHE.get(cache_key) + if cached is not None: + return cached + if size < 4 or (size & (size - 1)) != 0 or math.log(size, 4) % 1: + raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") + h4 = torch.tensor( + [[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], + dtype=dtype, + device=device, + ) + hadamard = h4 + current = 4 + while current < size: + hadamard = torch.kron(hadamard, h4) + current *= 4 + hadamard = hadamard / math.sqrt(size) + _HADAMARD_CACHE[cache_key] = hadamard + return hadamard + + +def rotate_activation(x: torch.Tensor, group_size: int) -> torch.Tensor: + features = x.shape[-1] + if features % group_size: + raise ValueError(f"features {features} are not divisible by convrot group_size {group_size}") + groups = features // group_size + hadamard = regular_hadamard(group_size, device=x.device, dtype=x.dtype) + grouped = x.reshape(*x.shape[:-1], groups, group_size) + return torch.matmul(grouped, hadamard).reshape(x.shape) + + +def parse_comfy_quant_marker(blob: torch.Tensor) -> dict[str, Any]: + """Decode the uint8 JSON marker Comfy stores next to each quantized linear.""" + raw = bytes(blob.detach().cpu().contiguous().view(torch.uint8).numpy()) + raw = raw.split(b"\x00", 1)[0] + marker = json.loads(raw.decode("utf-8")) + if not isinstance(marker, dict): + raise ValueError("comfy_quant marker must be a JSON object") + return marker + + +class Int8ConvRotLinear(nn.Module): + """W8A8 linear matching Comfy ``int8_tensorwise`` (+ optional ConvRot).""" + + def __init__( + self, + in_features: int, + out_features: int, + *, + bias: bool, + convrot: bool, + group_size: int, + ) -> None: + super().__init__() + if convrot and in_features % group_size != 0: + raise ValueError( + f"ConvRot weights were rotated with group_size {group_size}, " + f"but in_features {in_features} is not divisible by that group" + ) + self.in_features = in_features + self.out_features = out_features + self.convrot = convrot + self.group_size = group_size + self._transpose_view = envs.FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW.get() + self._fused_dequant = envs.FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT.get() + self.register_buffer("weight", torch.empty(out_features, in_features, dtype=torch.int8)) + self.register_buffer("weight_scale", torch.empty(out_features, 1, dtype=torch.float32)) + if bias: + self.bias = nn.Parameter(torch.zeros(out_features), requires_grad=False) + else: + self.register_parameter("bias", None) + + def _dequant_weight(self, dtype: torch.dtype) -> torch.Tensor: + return self.weight.to(dtype) * self.weight_scale.to(dtype) + + @staticmethod + def _dequant_int8_gemm( + acc: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + ) -> torch.Tensor: + # int32 acc is ~K·127² and overflows fp16 before 1/127 scales land. + return acc.float() * x_scale.float() * weight_scale.t().float() + + def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + if self.convrot: + x_2d = rotate_activation(x_2d, self.group_size) + row_max = x_2d.abs().amax(dim=-1, keepdim=True).clamp_min(1e-30) + x_scale = row_max / 127.0 + x_q = (x_2d / x_scale).round().clamp(-128, 127).to(torch.int8) + rows = x_q.shape[0] + if rows <= 16: + pad = 17 - rows + x_q = F.pad(x_q, (0, 0, 0, pad)) + x_scale = F.pad(x_scale, (0, 0, 0, pad)) + return x_q, x_scale + + def forward_quantized(self, x_q: torch.Tensor, x_scale: torch.Tensor, + original_shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + rows = math.prod(original_shape[:-1]) + weight = self.weight.t() + if not self._transpose_view: + weight = weight.contiguous() + acc = torch._int_mm(x_q, weight)[:rows] + if self._fused_dequant and not torch.is_grad_enabled() and not torch.compiler.is_compiling(): + from fastvideo.models.vaes.minimax_h3_int8_kernels import fused_int8_dequant_bias + return fused_int8_dequant_bias(acc, x_scale[:rows], self.weight_scale, self.bias, dtype).view( + *original_shape[:-1], self.out_features) + out = self._dequant_int8_gemm(acc, x_scale[:rows], self.weight_scale) + if self.bias is not None: + out = out + self.bias.float() + return out.to(dtype=dtype).view(*original_shape[:-1], self.out_features) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + original_shape = x.shape + if x.device.type == "cuda" and x.shape[-1] % 8 == 0: + return self.forward_quantized(*self.quantize_input(x), original_shape, x.dtype) + else: + x_2d = x.reshape(-1, original_shape[-1]).contiguous() + if self.convrot: + x_2d = rotate_activation(x_2d, self.group_size) + out = F.linear(x_2d.float(), self._dequant_weight(torch.float32)) + if self.bias is not None: + out = out + self.bias.float() + return out.to(dtype=x.dtype).view(*original_shape[:-1], self.out_features) + + +def shared_int8_projections(layers: tuple[nn.Module, ...], x: torch.Tensor) -> tuple[torch.Tensor, ...]: + """Reuse identical ConvRot/row quantization while retaining each projection's INT8 GEMM.""" + first = layers[0] + compatible = (x.is_cuda and x.shape[-1] % 8 == 0 + and all(isinstance(layer, Int8ConvRotLinear) for layer in layers)) + if compatible: + compatible = all((layer.in_features, layer.convrot, layer.group_size) + == (first.in_features, first.convrot, first.group_size) for layer in layers) + if not compatible: + return tuple(layer(x) for layer in layers) + x_q, x_scale = first.quantize_input(x) + return tuple(layer.forward_quantized(x_q, x_scale, x.shape, x.dtype) for layer in layers) + + +def _int8_linear_from_tensors( + weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None, + marker: dict[str, Any], +) -> Int8ConvRotLinear: + if weight.dtype != torch.int8: + raise ValueError(f"expected int8 weight, got {weight.dtype}") + if scale.ndim == 1: + scale = scale.unsqueeze(1) + convrot = bool(marker.get("convrot", False)) + group_size = int(marker.get("convrot_groupsize", _DEFAULT_GROUP_SIZE)) if convrot else _DEFAULT_GROUP_SIZE + if marker.get("format") not in (None, "int8_tensorwise"): + raise ValueError(f"unsupported comfy_quant format {marker.get('format')!r}") + layer = Int8ConvRotLinear( + weight.shape[1], + weight.shape[0], + bias=bias is not None, + convrot=convrot, + group_size=group_size, + ) + layer.weight.copy_(weight) + layer.weight_scale.copy_(scale.to(torch.float32)) + if bias is not None and layer.bias is not None: + layer.bias.data.copy_(bias) + return layer + + +def _split_qkv( + weight: torch.Tensor, + scale: torch.Tensor, + bias: torch.Tensor | None, + marker: dict[str, Any], + *, + heads: int, + dim_head: int, +) -> tuple[Int8ConvRotLinear, Int8ConvRotLinear, Int8ConvRotLinear]: + """Undo Comfy fused QKV: per-head ``[q, k, v]`` then stacked heads.""" + expected = heads * 3 * dim_head + if weight.shape[0] != expected: + raise ValueError(f"fused to_qkv out_features {weight.shape[0]} != heads*3*dim_head {expected}") + if scale.ndim == 1: + scale = scale.unsqueeze(1) + in_features = weight.shape[1] + weight = weight.view(heads, 3, dim_head, in_features) + scale = scale.view(heads, 3, dim_head, 1) + bias_view = None if bias is None else bias.view(heads, 3, dim_head) + pieces = [] + for index in range(3): + w = weight[:, index].reshape(heads * dim_head, in_features) + s = scale[:, index].reshape(heads * dim_head, 1) + b = None if bias_view is None else bias_view[:, index].reshape(heads * dim_head) + pieces.append(_int8_linear_from_tensors(w.contiguous(), s.contiguous(), b, marker)) + return pieces[0], pieces[1], pieces[2] + + +def _swap_swiglu_halves(tensor: torch.Tensor) -> torch.Tensor: + """Comfy ``[gate; value]`` → FastVideo ``[value; gate]`` along out_features.""" + if tensor.shape[0] % 2: + raise ValueError(f"SwiGLU packed rows {tensor.shape[0]} are not even") + half = tensor.shape[0] // 2 + return torch.cat([tensor[half:], tensor[:half]], dim=0) + + +def overlay_minimax_h3_int8_convrot_decoder(vae: nn.Module, checkpoint_path: str | Path) -> int: + """Swap H3 VAE decoder transformer linears for Comfy int8-convrot weights. + + Dense decoder tensors that the export still stores in float (embed, norms, + scales, ``proj_out``) are copied onto the matching FastVideo modules. + Returns the number of quantized linears installed. + """ + path = Path(checkpoint_path) + tensors = safetensors_load_file(str(path)) + decoder = getattr(vae, "decoder", None) + if decoder is None: + raise ValueError("MiniMax-H3 VAE overlay expected a `.decoder` module") + + if "decoder.x_embedder.weight" in tensors and hasattr(decoder, "proj_in"): + decoder.proj_in.weight.data.copy_(tensors["decoder.x_embedder.weight"]) + if decoder.proj_in.bias is not None and "decoder.x_embedder.bias" in tensors: + decoder.proj_in.bias.data.copy_(tensors["decoder.x_embedder.bias"]) + for name in ( + "decoder.register_tokens", + "decoder.norm_out.weight", + "decoder.norm_out.bias", + "decoder.proj_out.weight", + "decoder.proj_out.bias", + ): + if name not in tensors: + continue + module_name, _, param_name = name.removeprefix("decoder.").rpartition(".") + target = decoder if not module_name else decoder.get_submodule(module_name) + getattr(target, param_name).data.copy_(tensors[name]) + + quantized = 0 + blocks = decoder.transformer_blocks + for index, block in enumerate(blocks): + prefix = f"decoder.transformer_blocks.{index}" + marker = parse_comfy_quant_marker(tensors[f"{prefix}.attn.to_qkv{_COM_FY_QUANT_SUFFIX}"]) + to_q, to_k, to_v = _split_qkv( + tensors[f"{prefix}.attn.to_qkv.weight"], + tensors[f"{prefix}.attn.to_qkv.weight_scale"], + tensors.get(f"{prefix}.attn.to_qkv.bias"), + marker, + heads=int(block.attn.heads), + dim_head=int(block.attn.dim_head), + ) + block.attn.to_q = to_q + block.attn.to_k = to_k + block.attn.to_v = to_v + quantized += 3 + out_marker = parse_comfy_quant_marker(tensors[f"{prefix}.attn.to_out{_COM_FY_QUANT_SUFFIX}"]) + to_out = _int8_linear_from_tensors( + tensors[f"{prefix}.attn.to_out.weight"], + tensors[f"{prefix}.attn.to_out.weight_scale"], + tensors.get(f"{prefix}.attn.to_out.bias"), + out_marker, + ) + block.attn.to_out[0] = to_out + quantized += 1 + w1_marker = parse_comfy_quant_marker(tensors[f"{prefix}.ff.w1{_COM_FY_QUANT_SUFFIX}"]) + w1_weight = _swap_swiglu_halves(tensors[f"{prefix}.ff.w1.weight"]) + w1_scale = _swap_swiglu_halves(tensors[f"{prefix}.ff.w1.weight_scale"]) + w1_bias = tensors.get(f"{prefix}.ff.w1.bias") + if w1_bias is not None: + w1_bias = _swap_swiglu_halves(w1_bias) + block.ff.net[0].proj = _int8_linear_from_tensors(w1_weight, w1_scale, w1_bias, w1_marker) + quantized += 1 + w2_marker = parse_comfy_quant_marker(tensors[f"{prefix}.ff.w2{_COM_FY_QUANT_SUFFIX}"]) + block.ff.net[2] = _int8_linear_from_tensors( + tensors[f"{prefix}.ff.w2.weight"], + tensors[f"{prefix}.ff.w2.weight_scale"], + tensors.get(f"{prefix}.ff.w2.bias"), + w2_marker, + ) + quantized += 1 + for dense_name in ("norm1.weight", "norm2.weight", "scale1", "scale2"): + key = f"{prefix}.{dense_name}" + if key not in tensors: + continue + if dense_name.endswith(".weight"): + module = block.get_submodule(dense_name.rsplit(".", 1)[0]) + module.weight.data.copy_(tensors[key]) + else: + getattr(block, dense_name).data.copy_(tensors[key]) + + device = next(vae.parameters()).device + vae.decoder.to(device) + logger.info("Overlaid MiniMax-H3 int8-convrot VAE decoder: %s linears from %s", quantized, path) + return quantized + + +def is_int8_convrot_vae_path(path: str | Path) -> bool: + raw = os.path.basename(str(path)) + real = os.path.basename(os.path.realpath(str(path))) + return INT8_CONVROT_FILENAME in (raw, real) or "int8_convrot" in raw or "int8_convrot" in real + + +def dense_vae_safetensors(paths: list[str]) -> list[str]: + """Drop the ConvRot overlay so it is not loaded as a dense VAE shard.""" + return [path for path in paths if not is_int8_convrot_vae_path(path)] + + +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 diff --git a/fastvideo/models/vaes/minimax_h3_int8_kernels.py b/fastvideo/models/vaes/minimax_h3_int8_kernels.py new file mode 100644 index 000000000..e1e5a9fad --- /dev/null +++ b/fastvideo/models/vaes/minimax_h3_int8_kernels.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Eager INT8 VAE epilogue with the reference's separate FP32 operations.""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _dequant_bias(Acc, XScale, WScale, Bias, Out, M: tl.constexpr, N: tl.constexpr, + XS: tl.constexpr, WS: tl.constexpr, HAS_BIAS: tl.constexpr, BLOCK: tl.constexpr): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + rows, cols = offsets // N, offsets % N + valid = rows < M + 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 + if HAS_BIAS: + out = out + tl.load(Bias + cols).to(tl.float32) + tl.store(Out + offsets, out, valid) + + +def fused_int8_dequant_bias(acc: torch.Tensor, x_scale: torch.Tensor, + weight_scale: torch.Tensor, bias: torch.Tensor | None, + dtype: torch.dtype) -> torch.Tensor: + """Avoid full-size FP32 scaling intermediates; retain both rounding steps.""" + rows, cols = acc.shape + if not acc.is_cuda or acc.dtype != torch.int32 or not acc.is_contiguous(): + raise ValueError("INT8 VAE epilogue requires a contiguous CUDA INT32 matrix") + if x_scale.shape != (rows, 1) or weight_scale.shape != (cols, 1): + raise ValueError("INT8 VAE epilogue requires per-row and per-output-channel scales") + if any(t.device != acc.device for t in (x_scale, weight_scale)): + raise ValueError("INT8 VAE epilogue scales must be on the accumulator device") + if bias is not None and (bias.device != acc.device or bias.shape != (cols,) or not bias.is_contiguous()): + raise ValueError("INT8 VAE epilogue bias must be contiguous on the accumulator device") + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise ValueError("INT8 VAE epilogue supports FP32, FP16 and BF16 outputs") + out = torch.empty((rows, cols), device=acc.device, dtype=dtype) + _dequant_bias[(triton.cdiv(rows * cols, 1024),)]( + acc, x_scale, weight_scale, bias if bias is not None else acc, out, + rows, cols, x_scale.stride(0), weight_scale.stride(0), bias is not None, + BLOCK=1024, num_warps=4, enable_fp_fusion=False) + return out diff --git a/fastvideo/models/vaes/minimax_h3_video.py b/fastvideo/models/vaes/minimax_h3_video.py index e116defb7..fca78c244 100644 --- a/fastvideo/models/vaes/minimax_h3_video.py +++ b/fastvideo/models/vaes/minimax_h3_video.py @@ -15,6 +15,8 @@ import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint +import fastvideo.envs as envs + from fastvideo.attention import get_attn_backend from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig from fastvideo.platforms import AttentionBackendEnum @@ -300,6 +302,7 @@ class MiniMaxH3VideoAttention(nn.Module): self.heads = heads self.dim_head = dim_head self.use_bias = bias + self._share_int8_qkv = envs.FASTVIDEO_H3_VAE_INT8_SHARED_QKV.get() inner_dim = heads * dim_head self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False) self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False) @@ -339,9 +342,13 @@ class MiniMaxH3VideoAttention(nn.Module): rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: """Apply dense self-attention to one spatial VAE token sequence.""" - query = self.to_q(hidden_states).unflatten(2, (self.heads, -1)) - key = self.to_k(hidden_states).unflatten(2, (self.heads, -1)) - value = self.to_v(hidden_states).unflatten(2, (self.heads, -1)) + if self._share_int8_qkv and not torch.is_grad_enabled() and not torch.compiler.is_compiling(): + from fastvideo.models.vaes.minimax_h3_int8_convrot import shared_int8_projections + projections = shared_int8_projections((self.to_q, self.to_k, self.to_v), hidden_states) + else: + projections = tuple(layer(hidden_states) for layer in (self.to_q, self.to_k, self.to_v)) + query, key, value = (projection.unflatten(2, (self.heads, -1)) for projection in projections) + del projections query = self.norm_q(query.float()).to(query.dtype) key = self.norm_k(key.float()).to(key.dtype) @@ -522,6 +529,11 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module): ) +def _tile_batch_size() -> int: + """Spatial tiles decoded per decoder call (``FASTVIDEO_H3_VAE_TILE_BATCH``, default 1 = per tile).""" + return max(1, envs.FASTVIDEO_H3_VAE_TILE_BATCH.get()) + + def _is_minimax_h3_video_vae_decoder(name: str, submodule: nn.Module) -> bool: """Select the video decoder that serves the H3 VAE ``decode`` path.""" return name == "decoder" and isinstance(submodule, MiniMaxH3VideoViTDecoder3d) @@ -817,10 +829,31 @@ class AutoencoderKLMiniMaxH3(nn.Module): ratio = self.spatial_compression_ratio rows = [] + if _tile_batch_size() > 1 and len(set(y_lengths)) == 1 and len(set(x_lengths)) == 1: + # Every tile of the grid has one shape, and the ViT decoder + # treats batch entries independently: decode the grid in a few + # large calls instead of one small call per tile. + with nvtx_range("minimax_h3.vae.decode_clip.decode_tile_batches"): + latent_tiles = [ + z[..., y_position // ratio:(y_position + y_length) // ratio, + x_position // ratio:(x_position + x_length) // ratio] + for y_position, y_length in zip(y_indices, y_lengths) + for x_position, x_length in zip(x_indices, x_lengths) + ] + decoded: list[torch.Tensor] = [] + per_call = _tile_batch_size() + for start in range(0, len(latent_tiles), per_call): + batch = torch.cat(latent_tiles[start:start + per_call], dim=0) + # Clone before retaining: under the reduce-overhead compile the decoder output lives in + # a CUDA-graph pool that the next batch's replay overwrites. + out = self.decoder(self._project_decoder_tile(batch)).clone() + decoded.extend(out.split(z.shape[0], dim=0)) + columns = len(x_indices) + rows = [decoded[index:index + columns] for index in range(0, len(decoded), columns)] # The eager tile driver owns NVTX so each marker remains outside # the compiled decoder graph. with nvtx_range("minimax_h3.vae.decode_clip.decode_tiles"): - for row_index, (y_position, y_length) in enumerate(zip(y_indices, y_lengths)): + for row_index, (y_position, y_length) in enumerate(zip(y_indices, y_lengths) if not rows else ()): row = [] for column_index, (x_position, x_length) in enumerate(zip(x_indices, x_lengths)): with nvtx_range(f"minimax_h3.vae.decode_clip.tile.{row_index}.{column_index}"): diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index b39f5c443..a524a7ec9 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -11,6 +11,9 @@ from pathlib import Path from typing import Any import torch +from torch.distributed.tensor import DTensor + +import fastvideo.envs as envs from fastvideo.attention.selector import (_active_component_attention_backend_scope, coerce_attn_backend, get_env_variable_attn_backend) @@ -19,6 +22,7 @@ from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArch from fastvideo.configs.pipelines.minimax_h3 import (FASTH3_INFERENCE_FILE, FASTH3_INFERENCE_SCHEMA, MiniMaxH3PipelineConfig) from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.hooks.pinned_memory import PinnedTensorArena from fastvideo.logger import init_logger from fastvideo.models.hf_transformer_utils import get_diffusers_config from fastvideo.pipelines.basic.minimax_h3.stages import ( @@ -29,6 +33,7 @@ from fastvideo.pipelines.basic.minimax_h3.stages import ( MiniMaxH3LatentPreparationStage, MiniMaxH3VideoDecodingStage, ) +from fastvideo.pipelines.basic.minimax_h3.vsa_guard import refuse_zero_initialized_h3_vsa from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline from fastvideo.pipelines.pipeline_batch_info import ForwardBatch @@ -77,6 +82,50 @@ def _checkpoint_has_vsa_gates(transformer_dir: Path) -> bool: return False +def _pinned_swap(module: Any, device: torch.device) -> None: + """Move a module's tensors between the GPU and a persistent pinned host copy. + + Inference weights never change, so a parameter's pinned copy is made once and parking just repoints the + parameter at it (no transfer); restoring is one pinned host-to-device copy. Buffers are copied every time. + A LoRA swap/merge/unmerge drops the cached copies (``_drop_pinned_host_copies``), so the next park + re-copies whichever parameters the adapter changed. + """ + store = module.__dict__.setdefault("_pinned_host_tensors", {}) + params = dict(module.named_parameters()) + tensors = list(params.items()) + list(module.named_buffers()) + if device.type == "cpu": + missing = [(name, tensor) for name, tensor in tensors + if tensor is not None and tensor.device.type != "cpu" and ( + name not in store or store[name].shape != tensor.shape or store[name].dtype != tensor.dtype)] + arena = PinnedTensorArena(missing) if missing else None + for name, tensor in tensors: + if tensor is None: + continue + if device.type == "cpu": + if tensor.device.type == "cpu": + continue + host = store.get(name) + if host is None or host.shape != tensor.shape or host.dtype != tensor.dtype: + assert arena is not None + host = arena.empty_like(name, tensor) + host.copy_(tensor) + store[name] = host + elif name not in params: + host.copy_(tensor) + tensor.data = host + elif tensor.device != device: + tensor.data = tensor.data.to(device, non_blocking=True) + if device.type != "cpu" and torch.cuda.is_available(): + torch.cuda.current_stream(device).synchronize() + + +def _module_has_dtensor_params(module: Any) -> bool: + parameters = getattr(module, "parameters", None) + if not callable(parameters): + return False + return any(isinstance(parameter, DTensor) for parameter in parameters()) + + @dataclass(frozen=True) class _H3VideoGeometry: spatial_compression_ratio: int @@ -139,8 +188,11 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): # The linears every published H3 adapter targets. Left unset, ``LoRAPipeline`` # wraps *every* linear in the DiT -- including ``proj_in``, whose ``.weight`` the - # forward pass reads directly. ``BaseLayerWithLoRA`` exposes no ``.weight``, so - # that wrapping turns generation into an AttributeError before the first step. + # forward pass reads directly, and the attention projections whose + # ``quant_method`` the packed NVFP4 forward reads. The wrappers resolve those + # reads through ``base_layer`` (with a non-quantized fallback for + # ``quant_method``), but an explicit list avoids wrapping linears no published + # adapter trains. lora_target_modules = [ "attn.to_q", "attn.to_k", @@ -172,6 +224,28 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): # inspects the schedulers, which are never deferred. _lazy_module_names = ("text_encoder", "transformer", "vae", "audio_vae") + def _drop_pinned_host_copies(self) -> None: + """Drop every module's pinned host copies so the next park re-copies the weights. + + ``_pinned_swap`` reuses a parameter's stored host copy across parks; a LoRA + swap/merge/unmerge changes the weights, so the stale copy must not win. + """ + for module in self.modules.values(): + getattr(module, "__dict__", {}).pop("_pinned_host_tensors", None) + + def set_lora_adapter(self, *args, **kwargs): # type: ignore[override] + result = super().set_lora_adapter(*args, **kwargs) + self._drop_pinned_host_copies() + return result + + def merge_lora_weights(self) -> None: + super().merge_lora_weights() + self._drop_pinned_host_copies() + + def unmerge_lora_weights(self) -> None: + super().unmerge_lora_weights() + self._drop_pinned_host_copies() + def __init__(self, *args: Any, **kwargs: Any) -> None: self._ref2va = getattr(self, "_ref2va_default", False) self._denoise_stages_ready = False @@ -199,6 +273,9 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): if shift is None or not math.isfinite(float(shift)) or float(shift) <= 0: raise ValueError(f"MiniMax-H3 {modality} scheduler must expose a positive finite shift, got {shift}.") self._load_checkpoint_schedule(fastvideo_args) + transformer = self.get_module("transformer") + if transformer is not None: + refuse_zero_initialized_h3_vsa(transformer) def _checkpoint_facts(self) -> tuple[Any, bool]: """The checkpoint's parsed ``fastvideo_inference.json`` (None without one) and whether its transformer @@ -390,13 +467,32 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): finally: self._required_config_modules = saved + def _unified_memory_host(self) -> bool: + from fastvideo.pipelines import composed_pipeline_base + from fastvideo.platforms import current_platform + + device = composed_pipeline_base.get_local_torch_device() + device_id = 0 if device.index is None else int(device.index) + return bool(current_platform.has_unified_memory(device_id)) + def _release_text_encoder(self) -> None: + encoder = self.get_module("text_encoder") + if encoder is None: + return + # Discrete GPUs can keep Qwen in host RAM and borrow the GPU only for + # encode. Unified-memory boxes cannot keep it around even on "CPU", and + # an encoder that cannot move (sharded DTensor params) is released, as + # on main. + if not self._unified_memory_host() and self._move_module(encoder, "cpu"): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + logger.info("Moved MiniMax-H3 text encoder to CPU after conditioning") + return stage = self._stage_name_mapping.get("conditioning_stage") if stage is not None: stage.conditioner = None - encoder = self.modules.pop("text_encoder", None) - if encoder is None: - return + self.modules.pop("text_encoder", None) logger.info("Released MiniMax-H3 text encoder after conditioning") del encoder gc.collect() @@ -408,6 +504,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): encoder = self.get_module("text_encoder") stage = self._stage_name_mapping.get("conditioning_stage") if encoder is not None: + if not self._unified_memory_host(): + from fastvideo.pipelines import composed_pipeline_base + + self._move_module(encoder, composed_pipeline_base.get_local_torch_device()) if stage is not None and getattr(stage, "conditioner", None) is None: stage.conditioner = encoder return @@ -424,11 +524,70 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): if stage is not None: stage.conditioner = self.get_module("text_encoder") + def _move_module(self, module: Any, device: str | torch.device) -> bool: + if getattr(module, "_h3_encoder_layerwise_device", None) is not None: + # Layer hooks own placement; moving the whole encoder would restore + # every weight at once and defeat its VRAM bound. + return True + if _module_has_dtensor_params(module): + return False + if envs.FASTVIDEO_H3_PINNED_SWAP.get(): + _pinned_swap(module, torch.device(device)) + else: + module.to(device) + return True + + @staticmethod + def _parked_module_names() -> tuple[str, ...]: + """Denoise modules parked on the host while the text encoder runs (FASTVIDEO_H3_PARK_MODULES). + + Cards with room for the DiT next to the encoder park only the VAEs and keep the DiT resident. + """ + requested = envs.FASTVIDEO_H3_PARK_MODULES.get() + if not requested: + return _DENOISE_MODULE_NAMES + return tuple(name for name in requested.split(",") if name in _DENOISE_MODULE_NAMES) + + def _park_denoise_modules(self) -> None: + parked = False + for name in self._parked_module_names(): + module = self.get_module(name) + if module is None: + continue + if self._move_module(module, "cpu"): + parked = True + if parked: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + logger.info("Parked MiniMax-H3 denoise modules on CPU for text encode") + + def _restore_denoise_modules(self, fastvideo_args: FastVideoArgs) -> None: + from fastvideo.pipelines import composed_pipeline_base + + device = composed_pipeline_base.get_local_torch_device() + restored = False + for name in _DENOISE_MODULE_NAMES: + # Encode/decode stages move each VAE to the device when consumed. + # Keeping offloaded VAEs on the host leaves room for DiT activations + # and resident blocks throughout the denoising loop. + if name in {"vae", "audio_vae"} and fastvideo_args.vae_cpu_offload: + continue + module = self.get_module(name) + if module is None: + continue + if self._move_module(module, device): + restored = True + if restored: + 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() self._load_denoise_modules(fastvideo_args) + if not self._unified_memory_host(): + self._restore_denoise_modules(fastvideo_args) if not self._denoise_stages_ready: self._add_denoise_stages(ref2va=self._ref2va) for name in ( @@ -482,6 +641,8 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): def _add_denoise_stages(self, *, ref2va: bool) -> None: transformer = self.get_module("transformer") + if transformer is not None: + refuse_zero_initialized_h3_vsa(transformer) vae = self.get_module("vae") audio_vae = self.get_module("audio_vae") scheduler = self.get_module("scheduler") @@ -530,6 +691,8 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): # when neither deferral flag was on. if self._defer_denoise_modules(fastvideo_args): try: + if not self._unified_memory_host(): + self._park_denoise_modules() self._ensure_text_encoder(fastvideo_args) if self._denoise_stages_ready: logger.info("Running MiniMax-H3 condition stages before denoise (subsequent request)") 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 d39128e91..1e26ee540 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_conditioning.py @@ -261,7 +261,8 @@ class MiniMaxH3ConditioningStage(PipelineStage): 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)) + 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) try: 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 c59ac6f11..83d9a004a 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -22,6 +22,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import ( build_row_timesteps, ) from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.basic.minimax_h3.vsa_guard import refuse_zero_initialized_h3_vsa +from fastvideo.pipelines.lazy_module import is_lazy_module from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage from fastvideo.pipelines.stages.validators import StageValidators as V @@ -191,6 +193,10 @@ class MiniMaxH3DenoisingStage(PipelineStage): @torch.no_grad() def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: """Denoise the packed H3 video and audio streams over one shared schedule.""" + transformer = self.transformer + if is_lazy_module(transformer): + transformer.materialize() + refuse_zero_initialized_h3_vsa(transformer) layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY) if not isinstance(layout, MiniMaxH3PackedLayout): raise ValueError("MiniMax-H3 packed layout is missing before denoising.") diff --git a/fastvideo/pipelines/basic/minimax_h3/vsa_guard.py b/fastvideo/pipelines/basic/minimax_h3/vsa_guard.py new file mode 100644 index 000000000..5820252ac --- /dev/null +++ b/fastvideo/pipelines/basic/minimax_h3/vsa_guard.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from fastvideo.pipelines.lazy_module import is_lazy_module + + +def refuse_zero_initialized_h3_vsa(transformer: Any) -> None: + if is_lazy_module(transformer) and not transformer.is_materialized: + return + blocks = getattr(transformer, "transformer_blocks", None) + if blocks is None: + return + saw_weight = False + any_trained = False + for block in blocks: + attention = getattr(block, "attn", None) + gate = getattr(attention, "to_gate_compress", None) if attention is not None else None + if gate is None: + continue + weight = getattr(gate, "weight", None) + packed = getattr(gate, "_nvfp4_weight", None) + if (weight is None or weight.numel() == 0) and packed is not None and packed.numel() > 0: + # h3_dit_vsa: the gate is packed NVFP4 (weight is None). Two E2M1 codes per byte; the low three + # bits of each nibble are the magnitude, so a gate whose codes are all +-0 is a zero gate. + saw_weight = True + if bool(((packed & 0x77) != 0).any()): + any_trained = True + break + continue + # Layerwise offload leaves a zero-element placeholder on the module; it says nothing about the gate. + if weight is None or weight.numel() == 0: + continue + saw_weight = True + if bool((weight != 0).any()): + any_trained = True + break + if saw_weight and not any_trained: + raise RuntimeError("VIDEO_SPARSE_ATTN_H3 is loaded but every to_gate_compress weight is zero. " + "Load a VSA-trained student or a LoRA that carries gates; dense CompactH3 cannot run VSA.") diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 838ce677a..aa086a7d8 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -188,6 +188,14 @@ class PipelineStage(ABC): torch.cuda.synchronize() execution_time = time.perf_counter() - start_time logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000) + if torch.cuda.is_available(): + gib = 1024**3 + logger.info("[%s] Memory peak_allocated=%.2f GiB reserved=%.2f GiB resident_after=%.2f GiB", + stage_name, + torch.cuda.max_memory_allocated() / gib, + torch.cuda.memory_reserved() / gib, + torch.cuda.memory_allocated() / gib) + torch.cuda.reset_peak_memory_stats() batch.logging_info.add_stage_execution_time(stage_key, execution_time) batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name) component_metric = self.performance_component_metric diff --git a/fastvideo/tests/api/test_typed_quant_flow.py b/fastvideo/tests/api/test_typed_quant_flow.py index b7820d987..1a3c6cf9a 100644 --- a/fastvideo/tests/api/test_typed_quant_flow.py +++ b/fastvideo/tests/api/test_typed_quant_flow.py @@ -56,6 +56,21 @@ def test_typed_transformer_quant_resolves_to_nvfp4_instance(captured_kwargs) -> assert isinstance(captured_kwargs["transformer_quant"], NVFP4Config), (f"Expected NVFP4Config instance, got " f"{type(captured_kwargs['transformer_quant']).__name__}") + assert captured_kwargs["transformer_quant"].layer_profile == "refine" + + +def test_typed_nvfp4_layer_profile_pins_h3_dit(captured_kwargs) -> None: + from fastvideo.layers.quantization.nvfp4_config import NVFP4Config as CurrentNVFP4Config + + cfg = GeneratorConfig( + model_path="FastVideo/LTX2-Distilled-Diffusers", + engine=EngineConfig(quantization=QuantizationConfig(transformer_quant="NVFP4", layer_profile="h3_dit"), ), + ) + generator_config_to_fastvideo_args(cfg) + quant = captured_kwargs["transformer_quant"] + assert type(quant).__name__ == "NVFP4Config" + assert getattr(quant, "layer_profile") == "h3_dit" + assert isinstance(quant, CurrentNVFP4Config) def test_no_typed_quant_omits_transformer_quant_kwarg(captured_kwargs) -> None: diff --git a/fastvideo/tests/attention/test_minimax_h3_sparse_int8.py b/fastvideo/tests/attention/test_minimax_h3_sparse_int8.py new file mode 100644 index 000000000..0b92f1d73 --- /dev/null +++ b/fastvideo/tests/attention/test_minimax_h3_sparse_int8.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +"""sm89 INT8-QK/FP8-PV regression against dense masked BF16 attention.""" +from __future__ import annotations + +import pytest +import torch + + +def _cuda_sm89(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (8, 9): + pytest.skip("RTX 4090 / sm89 CUDA is required") + + +@pytest.mark.parametrize("partial", [False, True]) +def test_sparse_int8_preserves_tile_selection_and_valid_keys(partial): + _cuda_sm89() + from fastvideo.attention.backends.minimax_h3_sparse_int8 import sparse_sm89_attention + + torch.manual_seed(42) + q, k, v = (torch.randn(1, 2, 256, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + vbs = torch.tensor([64, 7 if partial else 64, 31 if partial else 64, 64], device="cuda", dtype=torch.int32) + valid = torch.arange(256, device="cuda") % 64 < vbs.repeat_interleave(64) + k[..., ~valid, :] = 0 + v[..., ~valid, :] = 0 + # Adjacent query tiles deliberately select different keys. A paired-query + # OR adapter would fail this regression even with perfect quantization. + mask = torch.tensor([[1, 0, 0, 1], [0, 1, 0, 0], [1, 0, 1, 0], [0, 0, 1, 1]], + device="cuda", dtype=torch.bool)[None, None].expand(1, 2, -1, -1).contiguous() + dense_mask = mask.repeat_interleave(64, -2).repeat_interleave(64, -1) & valid[None, None, None, :] + with torch.inference_mode(): + expected = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), + attn_mask=dense_mask) + output = sparse_sm89_attention(q, k, v, mask, vbs) + assert torch.isfinite(output).all() + relative_error = (output.float() - expected).norm() / expected.norm() + assert relative_error < 0.055, float(relative_error) + torch.testing.assert_close(output.float(), expected, atol=0.05, rtol=0.15) + + +def test_sparse_int8_handles_empty_selection(): + _cuda_sm89() + from fastvideo.attention.backends.minimax_h3_sparse_int8 import sparse_sm89_attention + + q = torch.zeros(1, 1, 128, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + out = sparse_sm89_attention(q, q, q, torch.zeros(1, 1, 2, 2, device="cuda", dtype=torch.bool), + torch.tensor([64, 64], device="cuda", dtype=torch.int32)) + assert torch.count_nonzero(out) == 0 + + +@pytest.mark.parametrize("batch,heads", [(1, 2), (2, 3)]) +@pytest.mark.parametrize("partner_pad", [False, True]) +def test_int8_bshd_views_match_contiguous_and_reduce_peak(batch, heads, partner_pad): + """Read production BSHD views without retaining three BHSD copies.""" + _cuda_sm89() + from fastvideo.attention.backends.minimax_h3_sparse_int8 import sparse_sm89_attention + + torch.manual_seed(113) + length, dim = 1024, 128 + storage_length = length + (64 if partner_pad else 0) + tensors = [torch.randn(batch, storage_length, heads, dim, device="cuda", dtype=torch.bfloat16) + for _ in range(3)] + q, k, v = [tensor[:, :length].transpose(1, 2) for tensor in tensors] + vbs = torch.full((length // 64,), 64, device="cuda", dtype=torch.int32) + vbs[1], vbs[4] = 7, 31 + valid = torch.arange(length, device="cuda") % 64 < vbs.repeat_interleave(64) + k[:, :, ~valid] = 0 + v[:, :, ~valid] = 0 + mask = torch.rand(batch, heads, length // 64, length // 64, device="cuda") > 0.8 + mask[:, :, 0] = False + with torch.inference_mode(): + # Populate autotuning/compilation caches before measuring allocations. + warm = sparse_sm89_attention(q, k, v, mask, vbs) + del warm + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + reference = sparse_sm89_attention(q.contiguous(), k.contiguous(), v.contiguous(), mask, vbs) + torch.cuda.synchronize() + copy_peak = torch.cuda.max_memory_allocated() - baseline + expected = reference.cpu() + del reference + torch.cuda.reset_peak_memory_stats() + actual = sparse_sm89_attention(q, k, v, mask, vbs) + torch.cuda.synchronize() + view_peak = torch.cuda.max_memory_allocated() - baseline + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + tensor_bytes = batch * heads * length * dim * 2 + assert copy_peak - view_peak >= tensor_bytes, (copy_peak, view_peak) diff --git a/fastvideo/tests/encoders/test_minimax_h3_encoder_layerwise.py b/fastvideo/tests/encoders/test_minimax_h3_encoder_layerwise.py new file mode 100644 index 000000000..b96d0042f --- /dev/null +++ b/fastvideo/tests/encoders/test_minimax_h3_encoder_layerwise.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Text-only H3 encoder streaming parity and placement contracts.""" +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +import fastvideo.envs as envs + +from fastvideo.hooks.hooks import ModuleHookManager +from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import MiniMaxH3SerializedNVFP4Config +from fastvideo.models.encoders.minimax_h3_qwen3_vl import MiniMaxH3Qwen3VLConditioner +from fastvideo.models.loader.text_encoder_quantization import _process_quantized_text_encoder_weights +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MiniMaxH3ConditioningStage +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for encoder streaming") +@pytest.mark.parametrize("quantized,fused", [(False, False), (True, False), (True, True)]) +def test_streamed_encoder_matches_resident_and_releases_layers(distributed_setup, monkeypatch, env_overrides, quantized, + fused): + # DiT residency must not accidentally keep encoder layers resident too. + env_overrides.enter_context(envs.FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS.override(6)) + env_overrides.enter_context(envs.FASTVIDEO_H3_ENCODER_FUSED_DEQUANT.override(False)) + config = MiniMaxH3Qwen3VLConfig() + config.arch_config = MiniMaxH3Qwen3VLArchConfig( + vocab_size=64, hidden_size=128, intermediate_size=256, + num_hidden_layers=3, num_hidden_layers_override=2, output_hidden_state_index=2, + num_attention_heads=1, num_key_value_heads=1, head_dim=128, + rope_scaling={"mrope_interleaved": True, "mrope_section": [32, 16, 16], "rope_type": "default"}, + vision_depth=1, vision_hidden_size=64, vision_intermediate_size=128, + vision_num_heads=1, vision_deepstack_visual_indexes=(), vision_out_hidden_size=128, + ) + config.quant_config = MiniMaxH3SerializedNVFP4Config() if quantized else None + torch.manual_seed(81) + model = MiniMaxH3Qwen3VLConditioner(config).to(dtype=torch.bfloat16).eval() + for name, parameter in model.named_parameters(): + if name.endswith("weight_packed"): + parameter.data.random_(0, 256) + elif name.endswith("weight_scale"): + parameter.data.fill_(0x38) + elif name.endswith("weight_global_scale"): + parameter.data.fill_(2.7) + else: + parameter.data.normal_(std=0.02) + if quantized: + _process_quantized_text_encoder_weights(model, torch.device("cuda")) + linear = model.language_model.layers[0].self_attn.q_proj.to("cuda") + x = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16) + expected_linear = linear(x)[0] + with patch.object(torch.Tensor, "item", side_effect=AssertionError("Unexpected device scalar read")): + actual_linear = linear(x)[0] + torch.testing.assert_close(actual_linear, expected_linear, rtol=0, atol=0) + ids = torch.tensor([1, 7, 4, 21, 5, 31, 18], device="cuda") + model.to("cuda") + expected = model.encode_ids(ids) + assert torch.isfinite(expected).all() + if fused: + for layer in model.modules(): + if hasattr(layer, "_nvfp4_fused_dequant"): + layer._nvfp4_fused_dequant = True + model.to("cpu") + model.prepare_layerwise_offload(torch.device("cuda")) + model.prepare_layerwise_offload(torch.device("cuda")) # repeated setup is harmless + assert model.language_model.embed_tokens.weight.device.type == "cpu" + assert next(model.visual.parameters()).device.type == "cpu" + for _ in range(2): + actual = model.encode_ids(ids) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for layer in model.language_model.layers: + assert all(parameter.numel() == 0 for parameter in layer.parameters()) + manager = ModuleHookManager.get_from(layer) + assert manager is not None + assert not manager.forward_hooks["LayerwiseOffloadHook"].state.gpu_named_parameters + with pytest.raises(ValueError, match="text-only"): + model.encode_ids(ids, pixel_values=torch.zeros(1, device="cuda"), + image_grid_thw=torch.ones(1, 3, device="cuda", dtype=torch.int64)) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_pipeline_does_not_move_streamed_encoder_whole(device): + module = SimpleNamespace(_h3_encoder_layerwise_device=torch.device("cuda")) + module.to = lambda *_: pytest.fail("Whole encoder move defeats streaming") + assert MiniMaxH3Pipeline._move_module(None, module, device) + + +def test_conditioning_stage_keeps_streamed_encoder_placement(monkeypatch): + import fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning as conditioning + + module = SimpleNamespace(_h3_encoder_layerwise_device=torch.device("cuda")) + module.parameters = lambda: iter([torch.empty(1)]) + module.to = lambda *_: pytest.fail("Conditioning must retain layerwise placement") + stage = MiniMaxH3ConditioningStage.__new__(MiniMaxH3ConditioningStage) + stage.conditioner, stage.ref2va = module, False + stage._encode_fl2va = lambda *_: (torch.zeros(1, 2, 128), torch.zeros(2, dtype=torch.int32)) + monkeypatch.setattr(conditioning, "get_local_torch_device", lambda: torch.device("cpu")) + batch = ForwardBatch(data_type="video", prompt="streaming parity") + output = stage.forward(batch, SimpleNamespace(text_encoder_cpu_offload=True)) + assert output.prompt_embeds[0].shape == (1, 2, 128) diff --git a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py index a92bb7eb7..e3d2e4921 100644 --- a/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py +++ b/fastvideo/tests/encoders/test_minimax_h3_qwen3_vl_checkpoint_nvfp4.py @@ -260,9 +260,13 @@ def test_runtime_preflight_gates_device_capability_parallelism_and_flashinfer(mo with pytest.raises(RuntimeError, match="requires a CUDA device"): config.validate_runtime(torch.device("cpu")) - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) - with pytest.raises(RuntimeError, match="sm100 or newer"): + monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 1) + # Below sm100 there is no FP4 GEMM: sm80+ de-quantizes each linear to bf16, older GPUs are refused. + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (7, 5)) + with pytest.raises(RuntimeError, match="sm80"): config.validate_runtime(torch.device("cuda")) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (8, 9)) + config.validate_runtime(torch.device("cuda")) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (12, 1)) monkeypatch.setattr(h3_nvfp4, "get_tp_world_size", lambda: 2) @@ -388,6 +392,7 @@ def test_apply_flattens_quantizes_and_restores_the_leading_dims(distributed_setu receipt.update(alpha=alpha, weight_packed=weight_packed, weight_scale=weight_scale) return torch.zeros(x_fp4.shape[0], weight_packed.shape[0], dtype=torch.bfloat16) + monkeypatch.setattr(h3_nvfp4, "_fp4_gemm_supported", lambda device: True) monkeypatch.setattr(h3_nvfp4, "_quantize_activation_nvfp4", fake_quantize) monkeypatch.setattr(h3_nvfp4, "_nvfp4_linear", fake_linear) output = MiniMaxH3SerializedNVFP4LinearMethod._apply_finalized(layer, torch.ones(1, 5, 128), layer.bias) @@ -402,6 +407,23 @@ def test_apply_flattens_quantizes_and_restores_the_leading_dims(distributed_setu assert torch.equal(output, torch.full((1, 5, 256), 3.0, dtype=torch.bfloat16)) +def test_apply_without_fp4_gemm_matches_the_dequantized_linear(distributed_setup, monkeypatch) -> None: + """Pre-Blackwell GPUs expand the packed weight to bf16 for the call; the result is x @ W_deq.T + b.""" + config = MiniMaxH3SerializedNVFP4Config.from_config(_checkpoint_quantization_config()) + layer = _language_linear(config, bias=True) + _fill_loaded(layer, global_scale=4.0) + layer.quant_method.process_weights_after_loading(layer) + layer.bias.data.fill_(3.0) + monkeypatch.setattr(h3_nvfp4, "_fp4_gemm_supported", lambda device: False) + x = torch.randn(1, 5, 128, dtype=torch.bfloat16) + output = MiniMaxH3SerializedNVFP4LinearMethod._apply_finalized(layer, x, layer.bias) + + weight = h3_nvfp4.dequantize_serialized_nvfp4(layer.weight_packed, layer.weight_scale, 4.0, torch.bfloat16) + # 0x11 packs two E2M1 codes of 0.5; E4M3 block scale 1; global scale 4 -> every weight is 0.125. + assert torch.equal(weight, torch.full_like(weight, 0.125)) + torch.testing.assert_close(output, torch.nn.functional.linear(x, weight, layer.bias.to(torch.bfloat16))) + + def _tiny_conditioner_config(keep_bf16: tuple[str, ...] = ("mlp.down_proj", )) -> MiniMaxH3Qwen3VLConfig: """One language layer whose linears satisfy the FlashInfer tile geometry, and a small vision tower.""" config = MiniMaxH3Qwen3VLConfig() diff --git a/fastvideo/tests/hooks/test_pinned_memory.py b/fastvideo/tests/hooks/test_pinned_memory.py new file mode 100644 index 000000000..9402397a8 --- /dev/null +++ b/fastvideo/tests/hooks/test_pinned_memory.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA host-registration lifetime and offload regressions (one GPU required).""" + +import gc +import weakref + +import pytest +import torch + +import fastvideo.envs as envs +from torch import nn + +from fastvideo.hooks.hooks import ModuleHookManager +from fastvideo.hooks.layerwise_offload import LayerwiseOffloadHook, LayerwiseOffloadState +from fastvideo.hooks.pinned_memory import PinnedTensorArena + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA host registration requires one GPU") + + +def test_arena_mixed_dtype_exact_copy_and_lifetime(): + sources = { + "weight": torch.randn(17, 33, device="cuda", dtype=torch.bfloat16), + "scale": torch.randn(17, 1, device="cuda", dtype=torch.float32), + "packed": torch.arange(513, device="cuda").to(torch.uint8), + "scalar": torch.tensor(3.0, device="cuda"), + "empty": torch.empty(0, 4, device="cuda"), + } + arena = PinnedTensorArena(sources.items()) + assert arena.buffer is not None, "This GPU must exercise registration, not fallback" + assert arena.buffer.numel() < sum(t.numel() * t.element_size() for t in sources.values()) + 4096 + 256 * len(sources) + hosts = {} + for name, source in sources.items(): + host = arena.empty_like(name, source) + host.copy_(source) + if host.numel(): + assert host.is_pinned() + assert (host.data_ptr() - arena.buffer.data_ptr()) % 256 == 0 or host.numel() == 0 + torch.testing.assert_close(host.to("cuda", non_blocking=True), source, rtol=0, atol=0) + hosts[name] = host + owner = weakref.ref(arena) + del arena + gc.collect() + assert owner() is not None, "Live typed views must retain the registration" + del hosts, host + gc.collect() + assert owner() is None + + +def test_registration_failure_uses_pinned_allocator(monkeypatch): + class RefusingRuntime: + + def cudaHostRegister(self, *_args): + return 1 + + monkeypatch.setattr(torch.cuda, "cudart", lambda: RefusingRuntime()) + source = torch.arange(27, dtype=torch.float32) + arena = PinnedTensorArena([("weight", source)]) + assert arena.buffer is None + host = arena.empty_like("weight", source) + host.copy_(source) + assert host.is_pinned() + torch.testing.assert_close(host, source, rtol=0, atol=0) + arena.close() + arena.close() + + +def test_offload_mutation_and_prefetched_detach(env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.override(True)) + module = nn.Linear(16, 16, device="cuda", dtype=torch.bfloat16) + module.register_buffer("packed", torch.arange(1 << 20, device="cuda").to(torch.uint8)) + expected = {name: tensor.clone() for name, tensor in list(module.named_parameters()) + list(module.named_buffers())} + state = LayerwiseOffloadState(torch.cuda.Stream(), torch.device("cuda")) + hook = LayerwiseOffloadHook(state) + manager = ModuleHookManager.get_from_or_default(module) + manager.append_forward_hook(hook) + old_buffer = state.cpu_arena.buffer + assert old_buffer.is_pinned() + with hook.mutate_params_scope(), torch.no_grad(): + module.weight.add_(1) + module.packed.add_(1) + assert not old_buffer.is_pinned(), "Reinitialization must unregister old storage" + expected["weight"].add_(1) + expected["packed"].add_(1) + state.prefetch_params() + new_buffer = state.cpu_arena.buffer + manager.remove_forward_hook(hook.name()) + assert not new_buffer.is_pinned(), "Detachment must unregister storage" + assert state.cpu_arena is None + assert not state.cpu_named_parameters and not state.gpu_named_parameters + for name, tensor in list(module.named_parameters()) + list(module.named_buffers()): + torch.testing.assert_close(tensor, expected[name], rtol=0, atol=0) + + +def test_h3_swap_reuses_host_storage_and_updates_buffers(): + from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import _pinned_swap + + module = nn.Linear(16, 16, device="cuda", dtype=torch.bfloat16) + module.register_buffer("cache", torch.ones(11, device="cuda")) + expected_weight = module.weight.detach().clone() + _pinned_swap(module, torch.device("cpu")) + hosts = module._pinned_host_tensors + pointers = {name: tensor.data_ptr() for name, tensor in hosts.items()} + assert all(tensor.is_pinned() for tensor in hosts.values()) + _pinned_swap(module, torch.device("cuda", torch.cuda.current_device())) + module.cache.add_(3) + _pinned_swap(module, torch.device("cpu")) + assert pointers == {name: tensor.data_ptr() for name, tensor in hosts.items()} + torch.testing.assert_close(module.cache, torch.full((11,), 4.0), rtol=0, atol=0) + torch.testing.assert_close(module.weight, expected_weight.cpu(), rtol=0, atol=0) diff --git a/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_activation_scale.py b/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_activation_scale.py new file mode 100644 index 000000000..7d34e346f --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_activation_scale.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU regression for the ModelOpt activation-scale export contract.""" +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.fixture +def converter(monkeypatch): + path = Path(__file__).resolve().parents[4] / 'scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py' + spec = importlib.util.spec_from_file_location('h3_modelopt_converter', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + packed = torch.zeros((2, 8), dtype=torch.uint8) + scales = torch.ones((2, 1), dtype=torch.float8_e4m3fn) + def quantize(*args, **kwargs): + return packed, scales + monkeypatch.setattr(module, '_flashinfer', lambda: ( + SimpleNamespace(layout_128x4=0), None, quantize, lambda value: value)) + return module + + +def convert(converter, input_scale): + return converter.convert_modelopt_linear( + torch.zeros((2, 8), dtype=torch.uint8), + torch.ones((2, 1), dtype=torch.float8_e4m3fn), + torch.tensor(0.25), 'cpu', input_scale=input_scale)[0] + + +def test_modelopt_activation_scale_is_preserved_as_reciprocal(converter): + buffers = convert(converter, torch.tensor(2.0)) + assert buffers['_nvfp4_input_global_sf'].dtype == torch.float32 + assert buffers['_nvfp4_input_global_sf'].shape == () + assert buffers['_nvfp4_input_global_sf'].item() == 0.5 + assert buffers['_nvfp4_alpha'].item() == 0.25 + + +def test_missing_modelopt_activation_scale_keeps_legacy_export(converter): + assert '_nvfp4_input_global_sf' not in convert(converter, None) + + +@pytest.mark.parametrize('value', [0.0, -1.0, float('nan'), float('inf')]) +def test_invalid_modelopt_activation_scale_is_rejected(converter, value): + with pytest.raises(ValueError, match='finite positive scalar'): + convert(converter, torch.tensor(value)) diff --git a/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_nvfp4_converter.py b/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_nvfp4_converter.py new file mode 100644 index 000000000..466731b11 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_minimax_h3_modelopt_nvfp4_converter.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Converter smoke test: a synthetic ModelOpt NVFP4 checkpoint, converted and loaded, matches its reference.""" + +from __future__ import annotations + +import importlib.util +import json +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from safetensors import safe_open +from safetensors.torch import save_file + +import fastvideo.layers.quantization.nvfp4_config as nv + +SCRIPT = Path(__file__).resolve().parents[4] / "scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py" +OUT, IN = 128, 128 +ACT_AMAX = 6.0 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), + reason="NVFP4 mm_fp4 needs a Blackwell GPU") + + +def _converter(): + 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 _modelopt_linear(prefix: str, gen: torch.Generator) -> dict[str, torch.Tensor]: + """Random E2M1 codes with positive E4M3 block scales: a valid ModelOpt unified-layout linear.""" + return { + f"{prefix}.weight": torch.randint(0, 256, (OUT, IN // 2), dtype=torch.uint8, generator=gen), + f"{prefix}.weight_scale": (torch.rand(OUT, IN // 16, generator=gen) + 0.5).to(torch.float8_e4m3fn), + f"{prefix}.weight_scale_2": torch.tensor(0.01, dtype=torch.float32), + f"{prefix}.input_scale": torch.tensor(1.0, dtype=torch.float32), + } + + +def _write_source(src: Path) -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(0) + tensors = { + **_modelopt_linear("transformer_blocks.0.ff.net.0.proj", gen), + # Already quantized by ModelOpt: --quantize-attention must keep it on the ModelOpt path. + **_modelopt_linear("transformer_blocks.0.attn.to_k", gen), + "transformer_blocks.0.attn.to_q.weight": (torch.randn(OUT, IN, generator=gen) * 0.05).to(torch.bfloat16), + "transformer_blocks.0.norm.weight": torch.ones(IN, dtype=torch.bfloat16), + } + src.mkdir() + save_file(tensors, str(src / "diffusion_pytorch_model-00001-of-00001.safetensors")) + weight_map = {k: "diffusion_pytorch_model-00001-of-00001.safetensors" for k in tensors} + (src / "diffusion_pytorch_model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + (src / "config.json").write_text(json.dumps({"quantization_config": {"quant_algo": "NVFP4"}})) + return tensors + + +class _Linear(nn.Module): + + def __init__(self, prefix: str) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(OUT, IN, dtype=torch.bfloat16), requires_grad=False) + self.quant_method = nv.NVFP4QuantizeMethod(layer_prefix=prefix) + + +def _model() -> nn.Module: + root = nn.Module() + block = nn.Module() + block.ff, block.attn = nn.Module(), nn.Module() + block.ff.fc_in = _Linear("transformer_blocks.0.ff.fc_in") + block.attn.to_q = _Linear("transformer_blocks.0.attn.to_q") + block.attn.to_k = _Linear("transformer_blocks.0.attn.to_k") + root.transformer_blocks = nn.ModuleList([block]) + return root + + +def test_converted_checkpoint_loads_and_matches_reference(tmp_path) -> None: + pytest.importorskip("flashinfer") + conv = _converter() + src, dst = tmp_path / "src", tmp_path / "dst" + source = _write_source(src) + amax = tmp_path / "amax.json" + amax.write_text(json.dumps({f"b0.{name}": {"all": ACT_AMAX} for name in ("ff.fc_in", "attn.to_q", "attn.to_k")})) + subprocess.run([sys.executable, str(SCRIPT), "--src", str(src), "--dst", str(dst), "--quantize-attention", + "--act-amax", str(amax)], check=True) + + export = dst / nv.H3_NVFP4_DIT_EXPORT_FILENAME + with safe_open(str(export), framework="pt", device="cpu") as reader: + keys = set(reader.keys()) + # ModelOpt bytes are carried over bit for bit, including the attention projection it already quantized. + for prefix, module in (("transformer_blocks.0.ff.net.0.proj", "transformer_blocks.0.ff.fc_in"), + ("transformer_blocks.0.attn.to_k", "transformer_blocks.0.attn.to_k")): + assert torch.equal(reader.get_tensor(f"{module}::_nvfp4_weight"), source[f"{prefix}.weight"]) + assert "transformer_blocks.0.attn.to_q::_nvfp4_weight" in keys + dense_index = json.loads((dst / "diffusion_pytorch_model.safetensors.index.json").read_text())["weight_map"] + assert set(dense_index) == {"transformer_blocks.0.norm.weight"} + assert "quantization_config" not in json.loads((dst / "config.json").read_text()) + + model = _model() + assert nv.load_minimax_h3_nvfp4_dit_export(model, str(export), device="cuda") == 3 + x = torch.randn(256, IN, device="cuda").clamp(-ACT_AMAX, ACT_AMAX).to(torch.bfloat16) + references = { + "ff.fc_in": conv.dequantize_modelopt(*(source[f"transformer_blocks.0.ff.net.0.proj.{s}"].cuda() + for s in ("weight", "weight_scale", "weight_scale_2"))), + "attn.to_k": conv.dequantize_modelopt(*(source[f"transformer_blocks.0.attn.to_k.{s}"].cuda() + for s in ("weight", "weight_scale", "weight_scale_2"))), + "attn.to_q": source["transformer_blocks.0.attn.to_q.weight"].cuda(), + } + block = model.transformer_blocks[0] + for name, reference in references.items(): + layer = block.get_submodule(name) + assert layer._nvfp4_input_global_sf.item() == pytest.approx(448.0 * 6.0 / ACT_AMAX) + out = layer.quant_method.apply(layer, x).float() + ref = x.float() @ reference.float().T + error = ((out - ref).norm() / ref.norm()).item() + assert error < 0.2, f"{name}: relative error {error:.3f}" diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_config.py b/fastvideo/tests/ops/quantization/test_nvfp4_config.py index 977528189..3a9fa0d4d 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_config.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_config.py @@ -44,10 +44,18 @@ def test_nvfp4config_layer_profile_round_trips_from_dict(): from fastvideo.layers.quantization.nvfp4_config import NVFP4Config config = NVFP4Config.from_config({"layer_profile": "base"}) assert config.layer_profile == "base" + config = NVFP4Config.from_config({"layer_profile": "h3_dit"}) + assert config.layer_profile == "h3_dit" config = NVFP4Config.from_config({}) assert config.layer_profile == "refine" +def test_nvfp4config_rejects_unknown_layer_profile(): + from fastvideo.layers.quantization.nvfp4_config import NVFP4Config + with pytest.raises(ValueError, match="h3_dit"): + NVFP4Config(layer_profile="h3-dit") + + def test_nvfp4_kernel_call_raises_clear_error_without_flashinfer(monkeypatch): """A call into the NVFP4 kernels must raise an actionable ImportError when flashinfer is missing, not a confusing diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_dequant.py b/fastvideo/tests/ops/quantization/test_nvfp4_dequant.py new file mode 100644 index 000000000..6eb9f6aa8 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_dequant.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused weight expansion against the independent serialized Torch decoder.""" +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton NVFP4 decoder") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("global_scale", [1.0, 2.7, 438.912]) +def test_fused_nvfp4_expands_all_codes_and_scale_tiles(dtype, global_scale): + from fastvideo.layers.quantization.nvfp4_dequant import dequantize_nvfp4_cuda + from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import dequantize_serialized_nvfp4 + + torch.manual_seed(23) + # Multiple row/column tiles distinguish the swizzle from a row-major decoder. + packed = torch.arange(256, device="cuda", dtype=torch.uint8).repeat(256, 1) + scales = torch.randint(0, 127, (256, 32), device="cuda", dtype=torch.uint8) + reference = dequantize_serialized_nvfp4(packed, scales, global_scale, dtype) + actual = dequantize_nvfp4_cuda(packed, scales, global_scale, dtype) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py b/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py new file mode 100644 index 000000000..32b9f84d1 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +from safetensors.torch import save_file + +import fastvideo.layers.quantization.nvfp4_config as nv + + +def _method(prefix: str) -> nv.NVFP4QuantizeMethod: + method = object.__new__(nv.NVFP4QuantizeMethod) + method.weight_fp4 = None + method.weight_scale = None + method.x_global_sf = torch.tensor(1.0, dtype=torch.float32) + method.layer_prefix = prefix + method._is_refine_only_layer = False + method._retain_original_weights = None + return method + + +class _ExportLinear(nn.Module): + + def __init__(self, prefix: str, out_dim: int = 8, in_dim: int = 16) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(out_dim, in_dim, dtype=torch.bfloat16), requires_grad=False) + self.quant_method = _method(prefix) + + +def _model() -> nn.Module: + root = nn.Module() + root.transformer_blocks = nn.ModuleList([nn.Module()]) + root.transformer_blocks[0].attn = nn.Module() + root.transformer_blocks[0].attn.to_q = _ExportLinear("transformer_blocks.0.attn.to_q") + return root + + +def test_dense_transformer_safetensors_drops_packed_export(tmp_path) -> None: + packed = tmp_path / nv.H3_NVFP4_DIT_EXPORT_FILENAME + shard = tmp_path / "diffusion_pytorch_model-00001-of-00008.safetensors" + packed.write_bytes(b"") + shard.write_bytes(b"") + kept = nv.dense_transformer_safetensors([str(packed), str(shard)]) + assert kept == [str(shard)] + assert nv.find_minimax_h3_nvfp4_dit_export([str(shard)]) == str(packed) + + +def test_load_minimax_h3_nvfp4_dit_export_overlays_and_purges_weight(tmp_path) -> None: + model = _model() + linear = model.transformer_blocks[0].attn.to_q + prefix = "transformer_blocks.0.attn.to_q" + tensors = { + f"{prefix}::_nvfp4_weight": torch.zeros(8, 8, dtype=torch.uint8), + f"{prefix}::_nvfp4_weight_scale": torch.zeros(8, 1, dtype=torch.uint8), + f"{prefix}::_nvfp4_alpha": torch.tensor(0.5, dtype=torch.float32), + f"{prefix}::_weight_global_sf": torch.tensor(2.0, dtype=torch.bfloat16), + } + path = tmp_path / nv.H3_NVFP4_DIT_EXPORT_FILENAME + save_file(tensors, str(path)) + + loaded = nv.load_minimax_h3_nvfp4_dit_export(model, str(path), device="cpu") + + assert loaded == 1 + assert linear.weight is None + assert linear._nvfp4_weight.dtype is torch.uint8 + assert tuple(linear._nvfp4_weight.shape) == (8, 8) + assert linear._nvfp4_alpha.item() == 0.5 + assert linear._weight_global_sf.item() == 2.0 + + +def test_load_minimax_h3_nvfp4_dit_export_rejects_untagged_linear(tmp_path) -> None: + model = _model() + model.transformer_blocks[0].attn.to_q.quant_method = object() + prefix = "transformer_blocks.0.attn.to_q" + tensors = { + f"{prefix}::_nvfp4_weight": torch.zeros(8, 8, dtype=torch.uint8), + f"{prefix}::_nvfp4_weight_scale": torch.zeros(8, 1, dtype=torch.uint8), + f"{prefix}::_nvfp4_alpha": torch.tensor(1.0, dtype=torch.float32), + f"{prefix}::_weight_global_sf": torch.tensor(1.0, dtype=torch.bfloat16), + } + path = tmp_path / nv.H3_NVFP4_DIT_EXPORT_FILENAME + save_file(tensors, str(path)) + + with pytest.raises(RuntimeError, match="layer_profile='h3_dit'"): + nv.load_minimax_h3_nvfp4_dit_export(model, str(path), device="cpu") diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py b/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py index d5529c04b..0e03d2daf 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py @@ -47,6 +47,35 @@ def test_is_minimax_h3_nvfp4_linear_prefix_non_main_ffn_linear(prefix: str) -> N assert not nvfp4.is_minimax_h3_nvfp4_linear_prefix(prefix) +@pytest.mark.parametrize( + "prefix", + [ + "transformer_blocks.0.attn.to_q", + "transformer_blocks.41.attn.to_out", + "minimax_h3.transformer_blocks.12.ff.fc_in", + "transformer_blocks.0.attn.to_k", + "transformer_blocks.0.attn.to_v", + "transformer_blocks.3.ff.fc_out", + ], +) +def test_is_minimax_h3_nvfp4_dit_linear_prefix_attn_and_ffn(prefix: str) -> None: + assert nvfp4.is_minimax_h3_nvfp4_dit_linear_prefix(prefix) + + +@pytest.mark.parametrize( + "prefix", + [ + "transformer_blocks.0.adaln_proj.linear", + "token_refiner.refiner_blocks.0.ff.fc_in", + "transformer_blocks.0.attn.to_gate_compress", + "proj_in", + "transformer_blocks.0.attn.to_q.weight", + ], +) +def test_is_minimax_h3_nvfp4_dit_linear_prefix_rejects_non_export_linears(prefix: str) -> None: + assert not nvfp4.is_minimax_h3_nvfp4_dit_linear_prefix(prefix) + + def test_nvfp4config_get_quant_method_minimax_h3_feed_forward() -> None: if not torch.cuda.is_available(): pytest.skip("NVFP4QuantizeMethod construction requires CUDA") diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py index 385cc4756..18b856634 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py +++ b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py @@ -114,6 +114,27 @@ def test_apply_out_dim_survives_purge(monkeypatch) -> None: assert out.shape == (2, 3, 8) +def test_convert_refuses_fsdp_sharded_bf16_purge(monkeypatch) -> None: + import torch.distributed.tensor as tdt + + class DummyDTensor: + def __init__(self, data: torch.Tensor) -> None: + self._data = data + + def to_local(self) -> torch.Tensor: + return self._data + + def float(self) -> torch.Tensor: + return self._data.float() + + monkeypatch.setattr(tdt, "DTensor", DummyDTensor) + model = _model(retain=False) + del model.always_fp4._parameters["weight"] + object.__setattr__(model.always_fp4, "weight", DummyDTensor(torch.randn(8, 16, dtype=torch.bfloat16))) + with pytest.raises(RuntimeError, match="FSDP-sharded"): + nv.convert_model_to_nvfp4(model) + + def test_dense_path_after_purge_raises_with_flag_named(monkeypatch) -> None: """Defensive guard: convert never purges dense-capable layers anymore, but a hand-purged module hitting the dense path must fail loudly.""" diff --git a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py index 02955b73f..6094be15a 100644 --- a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py +++ b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py @@ -5,8 +5,11 @@ from __future__ import annotations from contextlib import nullcontext from types import SimpleNamespace +import pytest import torch +import fastvideo.envs as envs + import fastvideo.pipelines.composed_pipeline_base as composed_pipeline_base from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.utils import FlexibleArgumentParser @@ -27,19 +30,36 @@ class _Profiler: def _stub_module(name: str) -> SimpleNamespace: if name in {"scheduler"}: - return SimpleNamespace(shift=12.0, name=name) - if name in {"audio_scheduler"}: - return SimpleNamespace(shift=3.0, name=name) - if name == "transformer": + module = SimpleNamespace(shift=12.0, name=name) + elif name in {"audio_scheduler"}: + module = SimpleNamespace(shift=3.0, name=name) + elif name == "transformer": # LoRAPipeline reads exclude_lora_layers off the DiT arch config. - return SimpleNamespace( + module = SimpleNamespace( name=name, config=SimpleNamespace(arch_config=SimpleNamespace(exclude_lora_layers=[])), ) - return SimpleNamespace(name=name) + else: + module = SimpleNamespace(name=name) + + def to(device): + module.device = device + module.moved_to.append(device) + return module + + module.device = None + module.moved_to = [] + module.to = to + return module -def _patch_pipeline_construction(monkeypatch, events: list) -> None: +@pytest.fixture(autouse=True) +def _no_pinned_swap(env_overrides): + # These contract tests use lightweight objects, not tensor-bearing modules. + env_overrides.enter_context(envs.FASTVIDEO_H3_PINNED_SWAP.override(False)) + + +def _patch_pipeline_construction(monkeypatch, events: list, *, unified_memory: bool = False) -> None: monkeypatch.setattr( composed_pipeline_base, "maybe_init_distributed_environment_and_model_parallel", @@ -49,7 +69,7 @@ def _patch_pipeline_construction(monkeypatch, events: list) -> None: monkeypatch.setattr(composed_pipeline_base, "get_world_group", lambda: SimpleNamespace(local_rank=0)) monkeypatch.setattr(composed_pipeline_base, "get_or_create_profiler", lambda trace_dir: _Profiler()) monkeypatch.setattr(composed_pipeline_base, "warmup_sequence_parallel_communication", lambda: None) - monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: False) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: unified_memory) monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False) @@ -103,6 +123,7 @@ def test_inference_defers_dit_and_vae_until_after_conditioning(monkeypatch) -> N monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + encoder = pipeline.get_module("text_encoder") batch = ForwardBatch(data_type="video", prompt="alpine dancer") out = pipeline.forward(batch, args) @@ -111,13 +132,204 @@ def test_inference_defers_dit_and_vae_until_after_conditioning(monkeypatch) -> N assert "transformer" in loads[1] assert "vae" in loads[1] assert "text_encoder" not in loads[1] - assert pipeline.get_module("text_encoder") is None - assert condition_stage.conditioner is None + assert pipeline.get_module("text_encoder") is encoder + assert encoder.moved_to[-1] == "cpu" + assert condition_stage.conditioner is not None assert pipeline.get_module("transformer") is not None assert pipeline._denoise_stages_ready is True + moves_before_second = len(encoder.moved_to) second = pipeline.forward(ForwardBatch(data_type="video", prompt="second clip"), args) assert second is not None + assert len(loads) == 2 + assert pipeline.get_module("text_encoder") is encoder + assert condition_stage.conditioner is not None + assert len(encoder.moved_to) > moves_before_second + assert encoder.moved_to[moves_before_second] == torch.device("cpu") + + +def test_sequential_skips_host_offload_for_dtensor_params(monkeypatch) -> None: + class _FakeDTensor: + pass + + monkeypatch.setattr( + "fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.DTensor", + _FakeDTensor, + ) + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def _dtensor_stub(name: str) -> SimpleNamespace: + module = _stub_module(name) + + def parameters(): + yield _FakeDTensor() + + module.parameters = parameters + return module + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + requested = list(self.required_config_modules) + loads.append(requested) + modules = dict(loaded_modules or {}) + for name in requested: + if name in modules: + continue + modules[name] = _dtensor_stub(name) if name == "text_encoder" else _stub_module(name) + return modules + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + args = FastVideoArgs( + model_path="unused/for-this-test", + enable_stage_verification=False, + h3_sequential_load=True, + ) + pipeline = MiniMaxH3Pipeline("unused/for-this-test", args) + pipeline.post_init() + + passthrough = lambda batch, _args: batch + monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) + monkeypatch.setattr(pipeline._stage_name_mapping["conditioning_stage"], "forward", passthrough) + original_add_denoise = pipeline._add_denoise_stages + + def fake_add_denoise(*, ref2va: bool) -> None: + original_add_denoise(ref2va=ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + monkeypatch.setattr(pipeline._stage_name_mapping[name], "forward", passthrough) + + monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + encoder = pipeline.get_module("text_encoder") + pipeline.forward(ForwardBatch(data_type="video", prompt="alpine dancer"), args) + pipeline.forward(ForwardBatch(data_type="video", prompt="second clip"), args) + + # A sharded (DTensor-params) encoder cannot be parked on the host, so after + # conditioning it is released (as on main) and reloaded for the next request. + assert len(loads) == 3 + assert loads[2] == ["text_encoder"] + assert encoder.moved_to == [] + transformer = pipeline.get_module("transformer") + assert transformer is not None + assert transformer.moved_to[-1] == torch.device("cpu") + + +def test_sequential_skips_host_offload_when_dense_params_precede_dtensors(monkeypatch) -> None: + class _FakeDTensor: + pass + + monkeypatch.setattr( + "fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline.DTensor", + _FakeDTensor, + ) + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def _mixed_stub(name: str) -> SimpleNamespace: + module = _stub_module(name) + + def parameters(): + yield torch.zeros(1) + yield _FakeDTensor() + + module.parameters = parameters + return module + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + requested = list(self.required_config_modules) + loads.append(requested) + modules = dict(loaded_modules or {}) + for name in requested: + if name in modules: + continue + modules[name] = _mixed_stub(name) if name == "text_encoder" else _stub_module(name) + return modules + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + args = FastVideoArgs( + model_path="unused/for-this-test", + enable_stage_verification=False, + h3_sequential_load=True, + ) + pipeline = MiniMaxH3Pipeline("unused/for-this-test", args) + pipeline.post_init() + + passthrough = lambda batch, _args: batch + monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) + monkeypatch.setattr(pipeline._stage_name_mapping["conditioning_stage"], "forward", passthrough) + original_add_denoise = pipeline._add_denoise_stages + + def fake_add_denoise(*, ref2va: bool) -> None: + original_add_denoise(ref2va=ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + monkeypatch.setattr(pipeline._stage_name_mapping[name], "forward", passthrough) + + monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + encoder = pipeline.get_module("text_encoder") + pipeline.forward(ForwardBatch(data_type="video", prompt="alpine dancer"), args) + assert encoder.moved_to == [] + + +def test_unified_memory_sequential_deletes_encoder_and_reloads(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events, unified_memory=True) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + requested = list(self.required_config_modules) + loads.append(requested) + modules = dict(loaded_modules or {}) + for name in requested: + modules.setdefault(name, _stub_module(name)) + return modules + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + args = FastVideoArgs( + model_path="unused/for-this-test", + enable_stage_verification=False, + h3_sequential_load=True, + lazy_module_load=False, + ) + pipeline = MiniMaxH3Pipeline("unused/for-this-test", args) + pipeline.post_init() + + condition_stage = pipeline._stage_name_mapping["conditioning_stage"] + passthrough = lambda batch, _args: batch + monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) + monkeypatch.setattr(condition_stage, "forward", passthrough) + original_add_denoise = pipeline._add_denoise_stages + + def fake_add_denoise(*, ref2va: bool) -> None: + original_add_denoise(ref2va=ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + monkeypatch.setattr(pipeline._stage_name_mapping[name], "forward", passthrough) + + monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + pipeline.forward(ForwardBatch(data_type="video", prompt="alpine dancer"), args) + + assert pipeline.get_module("text_encoder") is None + assert condition_stage.conditioner is None + assert len(loads) == 2 + + pipeline.forward(ForwardBatch(data_type="video", prompt="second clip"), args) assert len(loads) == 3 assert loads[2] == ["text_encoder"] assert pipeline.get_module("text_encoder") is None @@ -323,3 +535,15 @@ def test_resident_path_does_not_reread_encoder_on_later_request(monkeypatch) -> assert first is not None and second is not None assert len(loads) == 1 assert pipeline.get_module("text_encoder") is not None + + +@pytest.mark.parametrize("vae_offload", [True, False]) +def test_sequential_restore_keeps_offloaded_vaes_on_host_until_consumed(monkeypatch, vae_offload): + """Do not occupy denoise VRAM with decoders that stages load on demand.""" + pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline) + pipeline.modules = {name: _stub_module(name) for name in _DENOISE_MODULE_NAMES} + moved = [] + monkeypatch.setattr(composed_pipeline_base, "get_local_torch_device", lambda: torch.device("cuda", 0)) + monkeypatch.setattr(pipeline, "_move_module", lambda module, device: moved.append(module.name) or True) + pipeline._restore_denoise_modules(SimpleNamespace(vae_cpu_offload=vae_offload)) + assert moved == (["transformer"] if vae_offload else list(_DENOISE_MODULE_NAMES)) diff --git a/fastvideo/tests/stages/test_minimax_h3_vsa_guard.py b/fastvideo/tests/stages/test_minimax_h3_vsa_guard.py new file mode 100644 index 000000000..655909f34 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_vsa_guard.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from fastvideo.pipelines.basic.minimax_h3.vsa_guard import refuse_zero_initialized_h3_vsa +from fastvideo.pipelines.lazy_module import LazyModule + + +class _GateAttention(nn.Module): + + def __init__(self, trained: bool, *, has_weight: bool = True) -> None: + super().__init__() + if has_weight: + self.to_gate_compress = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + if trained: + self.to_gate_compress.weight.fill_(0.1) + else: + self.to_gate_compress.weight.zero_() + else: + self.to_gate_compress = SimpleNamespace(weight=None) + + +class _Block(nn.Module): + + def __init__(self, trained: bool, *, has_weight: bool = True) -> None: + super().__init__() + self.attn = _GateAttention(trained, has_weight=has_weight) + + +class _Transformer(nn.Module): + + def __init__(self, trained: bool, *, empty: bool = False, has_weight: bool = True) -> None: + super().__init__() + blocks = [] if empty else [_Block(trained, has_weight=has_weight)] + self.transformer_blocks = nn.ModuleList(blocks) + + +def test_refuse_zero_initialized_h3_vsa_raises() -> None: + with pytest.raises(RuntimeError, match="to_gate_compress"): + refuse_zero_initialized_h3_vsa(_Transformer(trained=False)) + + +def test_refuse_zero_initialized_h3_vsa_allows_trained_gates() -> None: + refuse_zero_initialized_h3_vsa(_Transformer(trained=True)) + + +def test_refuse_zero_initialized_h3_vsa_skips_dense_transformer() -> None: + refuse_zero_initialized_h3_vsa(SimpleNamespace()) + + +def test_refuse_zero_initialized_h3_vsa_skips_empty_blocks() -> None: + refuse_zero_initialized_h3_vsa(_Transformer(trained=False, empty=True)) + + +def test_refuse_zero_initialized_h3_vsa_skips_missing_gate_weight() -> None: + refuse_zero_initialized_h3_vsa(_Transformer(trained=False, has_weight=False)) + + +def test_refuse_zero_initialized_h3_vsa_skips_unmaterialized_lazy_module() -> None: + loads: list[int] = [] + + def loader() -> nn.Module: + loads.append(1) + return _Transformer(trained=False) + + refuse_zero_initialized_h3_vsa(LazyModule("transformer", loader)) + assert loads == [] + + +def _packed_gate_transformer(codes: int) -> SimpleNamespace: + """An h3_dit_vsa gate: ``weight`` purged, two E2M1 codes per ``_nvfp4_weight`` byte.""" + gate = SimpleNamespace(weight=None, _nvfp4_weight=torch.full((4, 2), codes, dtype=torch.uint8)) + return SimpleNamespace(transformer_blocks=[SimpleNamespace(attn=SimpleNamespace(to_gate_compress=gate))]) + + +def test_refuse_zero_initialized_h3_vsa_checks_packed_nvfp4_gates() -> None: + # 0x88 is -0 in both nibbles: still a zero gate. + for zero_codes in (0x00, 0x88): + with pytest.raises(RuntimeError, match="to_gate_compress"): + refuse_zero_initialized_h3_vsa(_packed_gate_transformer(zero_codes)) + refuse_zero_initialized_h3_vsa(_packed_gate_transformer(0x12)) diff --git a/fastvideo/tests/transformers/test_minimax_h3_tile_first.py b/fastvideo/tests/transformers/test_minimax_h3_tile_first.py new file mode 100644 index 000000000..fa2ac25fc --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_tile_first.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tile-first VSA parity on CUDA, including partial tiles and learned gates.""" +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +import fastvideo.envs as envs + +from fastvideo.layers.quantization.fp8_config import FP8Config, FP8QuantizeMethod +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_input_projections + + +def _install_fp8_buffers(layer): + weight = layer.weight.data.float() + if layer.quant_method.granularity == "channel": + scale = (weight.abs().amax(dim=1, keepdim=True) / 448).clamp_min(1e-6) + else: + scale = (weight.abs().amax().reshape(1) / 448).clamp_min(1e-6) + layer.register_buffer("_fp8_weight", (weight / scale).to(torch.float8_e4m3fn)) + layer.register_buffer("_fp8_weight_scale", scale) + layer.register_parameter("weight", None) + + +@pytest.mark.parametrize("granularity", ["tensor", "channel"]) +def test_shared_fp8_projections_match_independent_quantization(granularity): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9): + pytest.skip("sm89+ CUDA is required for FP8 GEMM") + from fastvideo.layers.linear import ReplicatedLinear + + torch.manual_seed(17) + layers = tuple(ReplicatedLinear(128, 256, bias=True, quant_config=FP8Config(granularity), + prefix=f"block.attn.to_{name}") for name in ("q", "k", "v")) + for layer in layers: + layer.to(device="cuda", dtype=torch.bfloat16) + layer.weight.data.normal_(std=0.1) + layer.bias.data.normal_(std=0.1) + _install_fp8_buffers(layer) + x = torch.randn(1, 272, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + reference = [layer(x)[0] for layer in layers] + with patch.object(FP8QuantizeMethod, "quantize_input", autospec=True, + side_effect=FP8QuantizeMethod.quantize_input) as quant: + actual = _shared_input_projections(layers, x) + assert quant.call_count == 1 + for expected, output in zip(reference, actual, strict=True): + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("kernel", ["original", "bf16", "int8"]) +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("gate_active", [False, True]) +@pytest.mark.parametrize("fused_rope", [False, True]) +def test_tile_first_matches_generic_vsa_with_partial_tiles(env_overrides, distributed_setup, tmp_path, + fp8, gate_active, fused_rope, kernel): + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("BF16 CUDA is required") + if fp8 and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("sm89+ CUDA is required for FP8 GEMM") + from fastvideo.attention.backends.video_sparse_attn_h3 import MiniMaxH3VSAMetadataBuilder + from fastvideo.forward_context import set_forward_context + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Attention + + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("VIDEO_SPARSE_ATTN_H3")) + env_overrides.enter_context(envs.override_external("FASTVIDEO_VSA_TRITON", "1")) + env_overrides.enter_context(envs.FASTVIDEO_VSA_SM100A.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_FP4.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_TILE_FIRST.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VSA_SM89_KERNEL.override("original")) + capture = kernel == "int8" and fp8 and gate_active and fused_rope + if capture: + env_overrides.enter_context(envs.FASTVIDEO_H3_CAPTURE_QKV.override(str(tmp_path))) + torch.manual_seed(21) + attn = MiniMaxH3Attention(256, 2, 128, 1e-5, (AttentionBackendEnum.VIDEO_SPARSE_ATTN_H3,), FP8Config("channel") if fp8 else None, + "transformer_blocks.0.attn", fuse_qknorm_rope=fused_rope) + attn.to(device="cuda", dtype=torch.bfloat16) + for parameter in attn.parameters(): + parameter.data.normal_(std=0.1) + if not gate_active: + attn.to_gate_compress.weight.data.zero_() + if fp8: + for layer in (attn.to_q, attn.to_k, attn.to_v, attn.to_out): + _install_fp8_buffers(layer) + meta = MiniMaxH3VSAMetadataBuilder().build(current_timestep=999, patch_size=(1, 1, 1), + VSA_sparsity=0.8, packed_segments=(65, 97, (4, 6, 10)), + device=torch.device("cuda"), tile_size=64) + length = meta.total_seq_length + x = torch.randn(1, length, 256, device="cuda", dtype=torch.bfloat16) + angles = torch.randn(length, 96, device="cuda") + rope = angles.cos(), angles.sin() + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta): + reference = attn(x, rope, length) + attn._vsa_tile_first = True + attn.distributed_attention.attn_impl._sm89_kernel = kernel + actual = attn(x, rope, length) + # Row order can choose a different GEMM reduction; neither attention nor + # the VSA selection/padding semantics are approximated by this route. + error = (actual.float() - reference.float()).norm() / reference.float().norm() + assert error < (0.02 if fp8 else 0.005), float(error) + torch.testing.assert_close(actual, reference, rtol=0.03, atol=0.05) + if capture: + data = torch.load(tmp_path / "layer-0.pt", weights_only=True) + torch.testing.assert_close(data["vbs"], meta.variable_block_sizes.cpu(), rtol=0, atol=0) + assert data["q"].shape == (1, 2, meta.variable_block_sizes.numel() * 64, 128) + assert data["mask"].dtype == torch.bool diff --git a/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py b/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py new file mode 100644 index 000000000..165ccc41e --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_int8_convrot.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +from fastvideo.models.vaes.minimax_h3_int8_convrot import ( + Int8ConvRotLinear, + _int8_linear_from_tensors, + dense_vae_safetensors, + overlay_minimax_h3_int8_convrot_decoder, + parse_comfy_quant_marker, + regular_hadamard, + rotate_activation, +) + + +def test_regular_hadamard_is_orthogonal() -> None: + hadamard = regular_hadamard(16, device=torch.device("cpu"), dtype=torch.float64) + identity = torch.eye(16, dtype=torch.float64) + torch.testing.assert_close(hadamard @ hadamard.T, identity, atol=1e-12, rtol=0.0) + + +def test_convrot_rotation_is_an_involution() -> None: + torch.manual_seed(0) + x = torch.randn(4, 256) + rotated = rotate_activation(x, 256) + twice = rotate_activation(rotated, 256) + torch.testing.assert_close(twice, x, atol=1e-5, rtol=1e-5) + + +def test_int8_gemm_scales_in_float32_not_fp16() -> None: + acc = torch.tensor([[100_000]], dtype=torch.int32) + x_scale = torch.tensor([[1.0 / 127.0]], dtype=torch.float16) + weight_scale = torch.tensor([[1.0 / 127.0]], dtype=torch.float16) + overflowed = acc.to(torch.float16) * x_scale * weight_scale.t() + scaled = Int8ConvRotLinear._dequant_int8_gemm(acc, x_scale, weight_scale) + assert torch.isinf(overflowed).all() + assert torch.isfinite(scaled).all() + torch.testing.assert_close(scaled, acc.float() * x_scale.float() * weight_scale.t().float()) + + +def test_int8_linear_rejects_incompatible_convrot_group() -> None: + marker = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256} + weight = torch.ones(8, 8, dtype=torch.int8) + scale = torch.ones(8, 1) + try: + _int8_linear_from_tensors(weight, scale, None, marker) + except ValueError as error: + assert "group_size 256" in str(error) + assert "in_features 8" in str(error) + else: + raise AssertionError("incompatible ConvRot overlay must be rejected") + + +def test_int8_linear_matches_dequantized_matmul() -> None: + torch.manual_seed(1) + layer = Int8ConvRotLinear(256, 32, bias=True, convrot=True, group_size=256) + layer.weight.copy_(torch.randint(-8, 8, (32, 256), dtype=torch.int8)) + layer.weight_scale.copy_(torch.linspace(0.01, 0.02, 32).unsqueeze(1)) + layer.bias.data.copy_(torch.randn(32)) + x = torch.randn(3, 256) + out = layer(x) + rotated = rotate_activation(x, 256) + expected = torch.nn.functional.linear(rotated, layer.weight.float() * layer.weight_scale, layer.bias) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_dense_vae_safetensors_drops_convrot_overlay() -> None: + kept = dense_vae_safetensors([ + "/tmp/vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "/tmp/vae/minimax_h3_video_vae_int8_convrot.safetensors", + "/tmp/vae/other.safetensors", + ]) + assert kept == [ + "/tmp/vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "/tmp/vae/other.safetensors", + ] + + +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) + blob[:len(payload)] = torch.tensor(list(payload), dtype=torch.uint8) + marker = parse_comfy_quant_marker(blob) + assert marker["format"] == "int8_tensorwise" + assert marker["convrot"] is True + assert marker["convrot_groupsize"] == 256 + + +class _FakeH3VAE(nn.Module): + + def __init__(self) -> None: + super().__init__() + decoder = nn.Module() + decoder.proj_in = nn.Linear(24, 8, bias=True) + decoder.register_tokens = nn.Parameter(torch.zeros(1, 4, 8), requires_grad=False) + decoder.norm_out = nn.LayerNorm(8) + decoder.proj_out = nn.Linear(8, 12, bias=True) + block = nn.Module() + block.norm1 = nn.RMSNorm(8, elementwise_affine=True) + block.norm2 = nn.RMSNorm(8, elementwise_affine=True) + block.scale1 = nn.Parameter(torch.zeros(8), requires_grad=False) + block.scale2 = nn.Parameter(torch.zeros(8), requires_grad=False) + attn = nn.Module() + attn.heads = 2 + attn.dim_head = 4 + attn.to_q = nn.Linear(8, 8, bias=True) + attn.to_k = nn.Linear(8, 8, bias=True) + attn.to_v = nn.Linear(8, 8, bias=True) + attn.to_out = nn.ModuleList([nn.Linear(8, 8, bias=True)]) + ff = nn.Module() + swiglu = nn.Module() + swiglu.proj = nn.Linear(8, 16, bias=True) + ff.net = nn.ModuleList([swiglu, nn.Dropout(0.0), nn.Linear(8, 8, bias=True)]) + block.attn = attn + block.ff = ff + decoder.transformer_blocks = nn.ModuleList([block]) + self.decoder = decoder + + +def _marker_tensor() -> torch.Tensor: + payload = json.dumps({"format": "int8_tensorwise", "convrot": False}).encode() + blob = torch.zeros(64, dtype=torch.uint8) + blob[:len(payload)] = torch.tensor(list(payload), dtype=torch.uint8) + return blob + + +def test_overlay_splits_fused_qkv_and_ffn(tmp_path) -> None: + tensors = { + "decoder.x_embedder.weight": torch.ones(8, 24), + "decoder.x_embedder.bias": torch.zeros(8), + "decoder.register_tokens": torch.ones(1, 4, 8), + "decoder.norm_out.weight": torch.ones(8), + "decoder.norm_out.bias": torch.zeros(8), + "decoder.proj_out.weight": torch.ones(12, 8), + "decoder.proj_out.bias": torch.zeros(12), + "decoder.transformer_blocks.0.attn.to_qkv.weight": torch.arange(24 * 8, dtype=torch.int8).reshape(24, 8), + "decoder.transformer_blocks.0.attn.to_qkv.weight_scale": torch.ones(24, 1), + "decoder.transformer_blocks.0.attn.to_qkv.bias": torch.zeros(24), + "decoder.transformer_blocks.0.attn.to_qkv.comfy_quant": _marker_tensor(), + "decoder.transformer_blocks.0.attn.to_out.weight": torch.ones(8, 8, dtype=torch.int8), + "decoder.transformer_blocks.0.attn.to_out.weight_scale": torch.ones(8, 1), + "decoder.transformer_blocks.0.attn.to_out.bias": torch.zeros(8), + "decoder.transformer_blocks.0.attn.to_out.comfy_quant": _marker_tensor(), + "decoder.transformer_blocks.0.ff.w1.weight": torch.cat( + [torch.ones(8, 8, dtype=torch.int8), torch.full((8, 8), 2, dtype=torch.int8)], dim=0), + "decoder.transformer_blocks.0.ff.w1.weight_scale": torch.cat( + [torch.ones(8, 1), torch.full((8, 1), 0.5)], dim=0), + "decoder.transformer_blocks.0.ff.w1.bias": torch.cat([torch.ones(8), torch.full((8, ), 3.0)], dim=0), + "decoder.transformer_blocks.0.ff.w1.comfy_quant": _marker_tensor(), + "decoder.transformer_blocks.0.ff.w2.weight": torch.ones(8, 8, dtype=torch.int8), + "decoder.transformer_blocks.0.ff.w2.weight_scale": torch.ones(8, 1), + "decoder.transformer_blocks.0.ff.w2.bias": torch.zeros(8), + "decoder.transformer_blocks.0.ff.w2.comfy_quant": _marker_tensor(), + "decoder.transformer_blocks.0.norm1.weight": torch.ones(8), + "decoder.transformer_blocks.0.norm2.weight": torch.ones(8), + "decoder.transformer_blocks.0.scale1": torch.ones(8), + "decoder.transformer_blocks.0.scale2": torch.ones(8), + } + path = tmp_path / "minimax_h3_video_vae_int8_convrot.safetensors" + save_file(tensors, path) + vae = _FakeH3VAE() + installed = overlay_minimax_h3_int8_convrot_decoder(vae, path) + assert installed == 6 + q = vae.decoder.transformer_blocks[0].attn.to_q + k = vae.decoder.transformer_blocks[0].attn.to_k + v = vae.decoder.transformer_blocks[0].attn.to_v + fused = tensors["decoder.transformer_blocks.0.attn.to_qkv.weight"].view(2, 3, 4, 8) + assert isinstance(q, Int8ConvRotLinear) + assert q.weight.shape == (8, 8) + assert torch.equal(q.weight, fused[:, 0].reshape(8, 8)) + assert torch.equal(k.weight, fused[:, 1].reshape(8, 8)) + assert torch.equal(v.weight, fused[:, 2].reshape(8, 8)) + proj = vae.decoder.transformer_blocks[0].ff.net[0].proj + assert isinstance(proj, Int8ConvRotLinear) + assert isinstance(vae.decoder.transformer_blocks[0].ff.net[2], Int8ConvRotLinear) + assert torch.equal(proj.weight[:8], torch.full((8, 8), 2, dtype=torch.int8)) + assert torch.equal(proj.weight[8:], torch.ones(8, 8, dtype=torch.int8)) + assert torch.equal(proj.weight_scale[:8], torch.full((8, 1), 0.5)) + assert torch.equal(proj.bias[:8], torch.full((8, ), 3.0)) + assert not hasattr(vae.decoder, "mask_token") + assert torch.equal(vae.decoder.proj_in.weight, torch.ones(8, 24)) diff --git a/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py b/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py new file mode 100644 index 000000000..87caab3ee --- /dev/null +++ b/fastvideo/tests/vaes/test_minimax_h3_int8_shared.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exact parity of shared VAE input preparation and strided INT8 weight GEMMs.""" +from unittest.mock import patch + +import pytest +import torch + +import fastvideo.envs as envs + +from fastvideo.models.vaes.minimax_h3_int8_convrot import Int8ConvRotLinear, shared_int8_projections + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for INT8 GEMM") +@pytest.mark.parametrize("rows", [3, 17, 129]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("convrot", [False, True]) +def test_shared_int8_and_transpose_views_are_exact(rows, dtype, convrot, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW.override(False)) + torch.manual_seed(73) + layers = tuple(Int8ConvRotLinear(256, out, bias=index != 1, convrot=convrot, group_size=256) + .to("cuda") for index, out in enumerate([128, 256, 64])) + for layer in layers: + layer.weight.random_(-127, 128) + layer.weight_scale.uniform_(0.0001, 0.03) + if layer.bias is not None: + layer.bias.data.normal_() + x = torch.randn(1, rows, 256, device="cuda", dtype=dtype) + x[0, 0].zero_() # clamp/padding semantics must also survive sharing + with torch.inference_mode(): + expected = tuple(layer(x) for layer in layers) + with patch.object(Int8ConvRotLinear, "quantize_input", autospec=True, + side_effect=Int8ConvRotLinear.quantize_input) as quant: + shared = shared_int8_projections(layers, x) + assert quant.call_count == 1 + for layer in layers: + layer._transpose_view = True + layer._fused_dequant = True + views = shared_int8_projections(layers, x) + for ref, actual, view in zip(expected, shared, views, strict=True): + assert torch.isfinite(ref).all() + torch.testing.assert_close(actual, ref, rtol=0, atol=0) + torch.testing.assert_close(view, ref, rtol=0, atol=0) + + +def test_shared_int8_keeps_cpu_fallback_exact(): + layers = tuple(Int8ConvRotLinear(16, 8, bias=False, convrot=False, group_size=16) for _ in range(3)) + for layer in layers: + layer.weight.fill_(1) + layer.weight_scale.fill_(0.01) + x = torch.ones(3, 16) + for actual, ref in zip(shared_int8_projections(layers, x), (layer(x) for layer in layers), strict=True): + torch.testing.assert_close(actual, ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for VAE attention parity") +def test_vae_attention_shares_quantized_projections_exactly(distributed_setup, env_overrides): + from fastvideo.models.vaes.minimax_h3_video import MiniMaxH3VideoAttention + + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_INT8_SHARED_QKV.override(False)) + env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW.override(False)) + torch.manual_seed(49) + attention = MiniMaxH3VideoAttention(256, 2, 128).to("cuda").eval() + for name in ("to_q", "to_k", "to_v"): + layer = Int8ConvRotLinear(256, 256, bias=True, convrot=True, group_size=256).to("cuda") + layer.weight.random_(-8, 9) + layer.weight_scale.fill_(0.01) + layer.bias.data.normal_(std=0.1) + setattr(attention, name, layer) + x = torch.randn(2, 33, 256, device="cuda") + with torch.inference_mode(): + expected = attention(x) + attention._share_int8_qkv = True + for layer in (attention.to_q, attention.to_k, attention.to_v): + layer._transpose_view = True + layer._fused_dequant = True + actual = attention(x) + torch.testing.assert_close(actual, expected, 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("has_bias", [False, True]) +def test_fused_int8_epilogue_large_accumulators_and_small_scales(dtype, has_bias): + from fastvideo.models.vaes.minimax_h3_int8_kernels import fused_int8_dequant_bias + + torch.manual_seed(19) + acc = torch.randint(-400_000_000, 400_000_000, (129, 264), device="cuda", dtype=torch.int32) + x_scale = torch.logspace(-30, -3, 129, device="cuda").view(-1, 1) + w_scale = torch.logspace(-6, -2, 264, device="cuda").view(-1, 1) + bias = torch.randn(264, device="cuda") if has_bias else None + expected = acc.float() * x_scale.float() * w_scale.t().float() + if bias is not None: + 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) diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index d0712091c..f4de6a30f 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -21,6 +21,38 @@ def _log_cuda_device_uuid(rank: int, device: torch.device) -> None: logger.info("Worker %d CUDA device UUID: GPU-%s", rank, device_uuid, local_main_process_only=False) +def _log_pipeline_memory(pipeline) -> None: + """Debug (FASTVIDEO_MEMORY_REPORT=1): bytes held per pipeline component, by device and dtype, plus the + largest tensors, so the resident footprint can be attributed before choosing offload placements.""" + gib = 1024**3 + for name, module in getattr(pipeline, "modules", {}).items(): + if not isinstance(module, torch.nn.Module): + continue + by_kind: dict[str, int] = {} + largest: list[tuple[int, str, str]] = [] + seen: set[int] = set() + for tname, t in list(module.named_parameters()) + list(module.named_buffers()): + if t is None or id(t) in seen: + continue + seen.add(id(t)) + nbytes = t.numel() * t.element_size() + key = f"{t.device.type}/{str(t.dtype).replace('torch.', '')}" + by_kind[key] = by_kind.get(key, 0) + nbytes + largest.append((nbytes, tname, key)) + largest.sort(reverse=True) + total = sum(by_kind.values()) + logger.info("MEMREPORT %s total=%.2f GiB %s", name, total / gib, { + k: round(v / gib, 2) + for k, v in sorted(by_kind.items(), key=lambda kv: -kv[1]) + }) + for nbytes, tname, key in largest[:8]: + logger.info("MEMREPORT %s %.3f GiB %s %s", name, nbytes / gib, key, tname) + if torch.cuda.is_available(): + logger.info("MEMREPORT cuda allocated=%.2f GiB reserved=%.2f GiB", + torch.cuda.memory_allocated() / gib, + torch.cuda.memory_reserved() / gib) + + class Worker: def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int, rank: int, distributed_init_method: str): @@ -71,6 +103,12 @@ class Worker: # Set the CUDA device BEFORE any CUDA calls if current_platform.is_cuda_alike(): torch.cuda.set_device(self.device) + # Debug: FASTVIDEO_CUDA_MEMORY_CAP_GIB emulates a smaller card by capping this process's allocator. + cap_gib = envs.FASTVIDEO_CUDA_MEMORY_CAP_GIB.get() + if cap_gib > 0: + total = torch.cuda.get_device_properties(self.device).total_memory + torch.cuda.set_per_process_memory_fraction(min(1.0, cap_gib * 1024**3 / total), self.device) + logger.info("Capped CUDA allocator at %s GiB of %.1f GiB", cap_gib, total / 1024**3) self.init_gpu_memory = torch.cuda.mem_get_info(self.device)[0] if current_platform.is_cuda(): _log_cuda_device_uuid(self.rank, self.device) @@ -90,6 +128,8 @@ class Worker: self.distributed_init_method) self.pipeline = build_pipeline(self.fastvideo_args) + if envs.FASTVIDEO_MEMORY_REPORT.get() and self.rank == 0: + _log_pipeline_memory(self.pipeline) def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) diff --git a/mkdocs.yml b/mkdocs.yml index a2d92045b..53ed86a51 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -184,6 +184,7 @@ nav: - Configuration: inference/configuration.md - Offloading: inference/offloading.md - Optimizations: inference/optimizations.md + - FastH3 on RTX PRO 6000: inference/fasth3_rtx_pro_6000.md - ComfyUI: inference/comfyui.md - Support Matrix: inference/support_matrix.md - Kandinsky 6 T2IVA: inference/kandinsky6.md diff --git a/scripts/benchmarks/minimax_h3_4090/README.md b/scripts/benchmarks/minimax_h3_4090/README.md new file mode 100644 index 000000000..11c8b75e8 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/README.md @@ -0,0 +1,422 @@ +# FastH3 on a single RTX 4090 + +Use the private `FastVideo/FastH3-Pruned-8Step-FP8-ckpt300` checkpoint. +Keep its `fastvideo_inference.json`: nine sigma-grid points produce eight +DMD forwards. Preserve VSA sparsity 0.8 and tile size 64. + +## Setup and validation + +The October 3, 2026 pod has one RTX 4090 (24,564 MiB), driver 580.126.20, +a 99,999,997,952-byte host cgroup limit, and 150 GB disk. Its runtime is +PyTorch 2.12.0+cu126, CUDA toolkit 12.6, and FlashInfer 0.7.1rc2. + +Exact-size host arenas replace the pinned allocator for layerwise blocks +and H3 module swaps. Each arena packs typed views at 256-byte offsets into +dedicated CUDA-registered pages. Live views retain their registration owner. +Mutation and hook detachment unregister old arenas. Registration failures +fall back to the ordinary pinned allocator with a warning. + +Validation on this pod: all 12 tests passed with the command below. +The tests cover repeated offloaded forwards, BF16, mixed-dtype exact copies, +owner lifetime, registration fallback, large-buffer mutation, detachment +after prefetch, and persistent H3 swaps with changing buffers. + +```bash +source /workspace/env.sh +source /workspace/venv/bin/activate +cd /workspace/fastvideo +python -P -m pytest fastvideo/tests/hooks/test_pinned_memory.py \ + fastvideo/tests/hooks/test_layerwise_offload.py -q +``` + +The supplied `handoff_4090/kernel_microbench/pinned_memory.py` measured +5.06 GiB extra cgroup usage for 2.87 GiB through `pin_memory()`, versus +2.87 GiB using direct host registration. Pinned H2D measured 25.9 GB/s; +pageable H2D measured 10.2 GB/s. These are microbenchmarks, not clip timings. + +The supplied `handoff_4090/kernel_microbench/t_fp8.py` measured: + +| K → N, M = 38,976 | Fused quantization | FP8 GEMM + scale epilogue | +| --- | --- | --- | +| 5376 → 5376 | 0.68 ms | 11.27 ms | +| 5376 → 28672 | 0.68 ms | 40.39 ms | +| 14336 → 5376 | 2.10 ms | 27.96 ms | + +Commands on the pod, run before clip benchmarks: + +```bash +cd /workspace +python -P kernel_microbench/pinned_memory.py +python -P kernel_microbench/t_fp8.py +``` + +## End-to-end baseline + +Run one warmup and at least two timed requests. The benchmark saves clips, +the exact Python command, runtime environment, sampling geometry, config, +source commit, each wall time, and the median to `results.json`. Stage logs +include GPU allocation peaks and conditioning, denoise, and decode timings. +Per-run host samples report both total cgroup memory and anonymous memory; +total usage includes checkpoint file cache. Both are sampled every 100 ms +and include all processes in the pod's cgroup. + +```bash +cd /workspace +export FASTVIDEO_SOURCE_COMMIT=ce877b200 +export FASTVIDEO_H3_PARK_MODULES=vae,audio_vae +export MAX_JOBS=4 +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/bench_pod.py \ + baseline-480p /workspace/vol/pruned_fp8_300 fp8 \ + --offload-buffers --lazy --no-vae-compile \ + --height 480 --width 832 --frames 243 --timed 2 +``` + +For the full-resolution baseline, change the name to `baseline-768p`, +height to 768, and width to 1344; retain 243 frames. Layerwise offload, +lazy component loading, and eager VAE decode are the starting configuration. +Do not compare these numbers with a different frame count or decoder. + +Completed baseline at source commit `ce877b200`, checkpoint revision +`f2ef54f9ff2091762ab8689b6514dcab5bc1d383`: + +| Configuration | Median e2e | Denoise stage | Video decode stage | Peak GPU allocated | Peak host anon | Peak total cgroup | +| --- | --- | --- | --- | --- | --- | --- | +| FP8, layerwise, lazy, eager H3 VAE, 832×480, 243 frames | 163.34 s | 91.72 s | 34.54 s | 17.47 GiB | 28.91 GiB | 76.95 GiB | + +Two timed requests took 163.67 s and 163.00 s after one warmup. +Stage times are medians and include deferred loading. Memory columns are +maxima across the timed requests. Total cgroup usage includes file cache. +The GPU peak occurs during conditioning. The saved clip contains 243 frames +at 24 fps (10.125 seconds) and an AAC audio track. A contact-sheet inspection +confirms a coherent pottery scene; speech and same-seed reference parity +still need review before claiming quality equivalence. + +Summarize a completed run while excluding warmup: + +```bash +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/summarize.py \ + /workspace/outputs/baseline-480p/results.json /workspace/baseline-480p.log +``` + +After a baseline works, measure FFN chunk sizes 16,384 and 8,192, then +increase resident DiT blocks within the measured GPU budget. Kernel or +decoder changes also need same-seed visual and auditory comparison. + +## Tile-first attention and full-resolution profiling + +Commit `c5d9f8132` shares compatible FP8 Q/K/V activation quantization and +releases dead block activations before residual modulation. The opt-in +`FASTVIDEO_H3_VSA_TILE_FIRST=1` scatters the attention input before projection, +then uses the existing BF16 VSA kernel. It retains tile-64 selection, partial +key validity and the learned compression branch. It supports eager, +single-rank inference; grad, compile, and multi-rank requests use the generic +path. Ten CUDA tests passed on the 4090, including mixed partial tiles, +active/zero gates, fused/unfused RoPE and FP8/nonquantized projections. + +The original 1344×768 baseline failed with a GPU OOM in post-attention +modulation before the FFN. The tile-first/FFN-16384 profiling run subsequently completed all three +768p clips without OOM: diagnostic median 349.76 s, denoise stage 244.07 s, +video decode stage 65.94 s, peak GPU allocated 17.55 GiB. A run without +profiling is still required for a release speed claim. +The profiling run used the following command. For SSH on macOS, +`-o UseKeychain=yes` retrieves the stored passphrase when the agent has no +loaded identities. Profiling/capture timings are +for diagnosis and must not be used as the final speed claim. + +```bash +FASTVIDEO_SOURCE_COMMIT=c5d9f8132 \ +FASTVIDEO_H3_PARK_MODULES=vae,audio_vae \ +FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384 \ +FASTVIDEO_H3_VSA_TILE_FIRST=1 \ +FASTVIDEO_H3_CAPTURE_QKV=/workspace/qkv-768p MAX_JOBS=4 \ +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/bench_pod.py \ + tile-first-768p-profile /workspace/vol/pruned_fp8_300 fp8 \ + --offload-buffers --lazy --no-vae-compile \ + --height 768 --width 1344 --frames 243 --timed 2 --profile +``` + +`FASTVIDEO_H3_CAPTURE_QKV` saves the first inputs from layers 0, 20 and 41, +two heads each, with full real sequences, masks, valid tile sizes and packed +row indices. Disable both capture and profiling for final clip timings. + +The separate `minimax_h3_sparse_int8.py` prototype uses INT8 QK and FP8 PV, +FP32 accumulators and the original 64-token mask. It has no automatic pipeline +route. Initial real-QKV tests found approximately 1.9× fine-kernel speedup +but 4.3–13.4% relative L2 error and a strict partial-tile elementwise test +failure. Do not select it for shipping. The microbenchmark now includes +BF16 QK/PV ablations to isolate that error. Offline compilation with Triton 3.8.0 for sm89 passed all four entry points +and emitted native INT8 and FP8 MMA instructions (20,480 bytes of shared +memory for the attention kernel). This is compilation evidence only. It must +pass CUDA tests and real-QKV/clip checks before integration. +Run its microbenchmark on an idle GPU: + +```bash +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/bench_sparse_qkv.py \ + /workspace/qkv-768p --output /workspace/sparse-qkv-results.json +``` + +SpargeAttn at `ae5b629ebb41e41f86b3ea2ab5a3283f13ac151a` built on the pod +with CUDA 12.8, `TORCH_CUDA_ARCH_LIST=8.9`, and `MAX_JOBS=4`. The upstream +`-Xcompiler -include,cassert` workaround was removed from `setup.py` to +avoid GCC 13 duplicate standard-library definitions. It is not selected by +the pipeline: its public 128-query/64-key adapter also needs correct masking +of partial H3 tiles before a meaningful parity comparison. + +## Cached-component 480p result + +At `4d9846573`, retain components between requests (omit `--lazy`), and set +`FASTVIDEO_H3_VSA_TILE_FIRST=1` and `FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384`. +Keep all other baseline settings, including the eager light VAE and original +BF16 attention kernel. After one warmup the two timed requests took 115.87 s +and 115.03 s, median **115.45 s** (29.3% less wall time than the lazy baseline). +Conditioning/denoise/video-decode stage medians were 10.85/72.62/24.98 s. +Peak GPU allocation was 19.02 GiB, host anon 43.96 GiB, total cgroup 92.72 GiB +including file cache. This recipe requires more host RAM than the 32 GB target; +its minimum RAM has not been tested under a smaller host limit. + +Decoded raw video and PCM audio SHA256 hashes match the baseline exactly for +both ceramics and harbor at seed 20260929. This establishes output identity +for these two prompts; it does not establish the checkpoint's BF16-reference +quality on other prompts. Raw results, clips and hash evidence are saved in +`output/fasth3-4090-20261003/` beside the workspace. + +## sm89 kernel precision choices + +`FASTVIDEO_H3_VSA_SM89_KERNEL=bf16` opts into the new entirely BF16 tile-64 +fine kernel. `int8` uses per-token INT8 QK with BF16 PV. `original` is the +unchanged default. Resolution happens when the backend is constructed; +unsupported devices, grad and compile requests retain the original route. +Both preserve the original tile selection, partial key masks and gated +compression. The rejected FP8-PV experiment is only exposed in the diagnostic +microbenchmark, never the pipeline route. + +Two-head real-QKV captures at 1344×768, layers 0/20/41, measured: + +| QK / PV | Fine-kernel speedup including input quantization | Relative L2 vs original BF16 | +| --- | --- | --- | +| BF16 / BF16 | 1.23× | 0.005–0.008% | +| INT8 / BF16 | 1.59× | 0.58–0.62% | +| INT8 / FP8 (rejected) | 1.91× | 4.3–13.4% | + +These are fine-kernel microbenchmarks, not end-to-end clip speedups. Same-seed +clip checks are required for the INT8 route. All 44 targeted CUDA/CPU checks +passed for native/tile-first routing, partial tiles, learned compression, +shared FP8 projections and sequential component restoration. + +At `9a8465ac4`, CPU-offloaded VAEs also remain on the host during denoising; +the encode/decode stages move them on demand. This frees room for resident +DiT blocks without changing any model arithmetic. + +## Six resident blocks and encoder priorities + +At `6d3c4cda5`, the opt-in BF16 fine kernel with six resident DiT blocks, +cached components and the settings below measured **111.27 s** median for +832×480, 243 frames. Timed requests were 111.61/110.92 s after one warmup. +Conditioning, denoise and video decode medians were 11.20/67.56/25.49 s. +Peak GPU allocation was 21.63 GiB, host anonymous memory 42.10 GiB and +total cgroup usage 90.94 GiB, including file cache. + +```bash +FASTVIDEO_SOURCE_COMMIT=6d3c4cda5 \ +FASTVIDEO_H3_PARK_MODULES=vae,audio_vae \ +FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384 \ +FASTVIDEO_H3_VSA_TILE_FIRST=1 FASTVIDEO_H3_VSA_SM89_KERNEL=bf16 \ +FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS=6 MAX_JOBS=4 \ +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/bench_pod.py \ + sm89-bf16-480p-resident6 /workspace/vol/pruned_fp8_300 fp8 \ + --offload-buffers --no-vae-compile \ + --height 480 --width 832 --frames 243 --timed 2 +``` + +The BF16 kernel changes floating-point reductions: the ceramics clip is +visually coherent in the sampled contact sheet but is not identical to the +original-kernel clip (decoded-video SSIM 0.597653). This is a speed candidate, +not proof of quality equivalence. The original kernel remains the default. + +The same cached/six-resident-block BF16 recipe at `26390848b`, with name +`sm89-bf16-480p-5s-resident6` and `--frames 124`, measured **65.27 s** median. +The legal frame count represents 5.167 s at 24 fps. Timed requests were +65.08/65.46 s after a 110.34 s warmup. Conditioning/denoise/video-decode +medians were 11.74/34.22/13.44 s. Peak GPU allocation was 21.62 GiB, +host anon 41.14 GiB and total cgroup 88.45 GiB. These generation wall times +include decode/export and exclude initial generator construction. No +profiling or QKV capture was enabled. + +At `26390848b`, per-key-tile V scaling reduced experimental INT8-QK/FP8-PV +real-tensor error to 0.84–1.54%, with 1.77× fine-kernel speedup. Dynamic +per-query, per-key-block P scaling also preserves contributions that would +underflow with the fixed P scale; it measured 0.80–1.52% error and 1.71× +speedup. All 30 focused CUDA kernel/routing tests passed. These FP8-PV routes +remain microbenchmark-only and need clip validation. + +The current text encoder is the trimmed 50-layer Qwen3-VL with serialized +NVFP4 weights, dequantized to BF16 per linear on sm89. The existing serialized +blockwise FP8 encoder requires sm100+ and FlashInfer's Blackwell GEMM; it +cannot run on the 4090 as written. An Ada FP8 implementation would also need +encoder streaming because its weights are larger. Reading the current +checkpoint tensor shapes gives 15.33 GiB total encoder weights, including +11.35 GiB packed values and 1.42 GiB block scales. Replacing those packed +values with FP8 while retaining the other tensors projects about 25.3 GiB +before activations (the FP8 block-scale overhead is small). This is a storage +estimate, not a measured FP8 encoder. First try fused NVFP4 +dequantization and avoid per-linear GPU scalar synchronization; then compare +a native sm89 FP8 encoder at equal prompts. The later streamed/fused conditioning stage is about 0.55 s, so a new +encoder export must be measured against that implementation. + +Remaining speed experiments include VAE compilation and tile-batch tuning, +INT8 decoder epilogue fusion, direct strided fine-attention reads, and fused +norm/activation quantization. See the completed smaller-VRAM results below. +The cached recipe's host peak does not establish a 32 GB system-RAM minimum. + + +## Streamed encoder and smaller VRAM caps + +At `753e560f6`, `FASTVIDEO_H3_ENCODER_LAYERWISE=1` streams the language +layers separately from DiT residency, retaining token embeddings and unused +vision modules on the CPU. This route currently supports text-only T2VA; +visual references fail explicitly. The pipeline preserves the streamed +placement. Twenty-eight offload/encoder/stage tests passed, including exact +repeated BF16 and NVFP4 parity. + +At `78540b635`, `FASTVIDEO_H3_ENCODER_FUSED_DEQUANT=1` expands packed NVFP4 +weights with one Triton pass. Fifteen strict tests passed across FP32, FP16, +BF16, swizzled scales and repeated encoder forwards. On three large stress +matrices the expansion was 25.3–25.9× faster and used 8× less temporary GPU +memory than Torch expansion. This is a dequantization microbenchmark; the +whole conditioning stage measured 0.55–0.62 seconds in the clip runs below. +The current encoder remains NVFP4 storage with BF16 GEMMs on Ada. + +All rows use 832×480, 243 frames, eight DMD forwards, sparsity 0.8, tile 64, +cached components, INT8 QK/BF16 PV and the eager light H3 VAE. Each median +has one warmup and two timed requests. Source is `78540b635` except the +16 GiB row (`753e560f6`, before fused dequantization). + +| 4090 configuration | Median e2e | Timed requests | Denoise | Video decode | Peak GPU allocated | Peak host anon | +| --- | --- | --- | --- | --- | --- | --- | +| 16 GiB cap, 6 resident | 104.03 s | 100.10 / 107.96 s | 72.42 s | 25.54 s | 11.28 GiB | 39.86 GiB | +| 12 GiB cap, 0 resident | 107.27 s | 111.02 / 103.52 s | 77.38 s | 25.48 s | 8.68 GiB | 42.35 GiB | +| Uncapped, 30 resident | 98.97 s | 98.88 / 99.05 s | 69.12 s | 25.47 s | 21.69 GiB | 29.44 GiB | + +Set `FASTVIDEO_CUDA_MEMORY_CAP_GIB=12` for the 12 GiB recipe; unset it for +the full card. Set resident blocks to the table value. Both fused rows use: + +```bash +FASTVIDEO_SOURCE_COMMIT=78540b635 \ +FASTVIDEO_H3_PARK_MODULES=vae,audio_vae \ +FASTVIDEO_H3_ENCODER_LAYERWISE=1 FASTVIDEO_H3_ENCODER_FUSED_DEQUANT=1 \ +FASTVIDEO_H3_VSA_TILE_FIRST=1 FASTVIDEO_H3_VSA_SM89_KERNEL=int8 \ +FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS=30 FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384 \ +FASTVIDEO_H3_VAE_TILE_BATCH=28 MAX_JOBS=4 \ +python -P /workspace/fastvideo/scripts/benchmarks/minimax_h3_4090/bench_pod.py \ + sm89-int8-480p-resident30-fused /workspace/vol/pruned_fp8_300 fp8 \ + --offload-buffers --no-vae-compile --height 480 --width 832 --frames 243 --timed 2 +``` + +The 12 GiB and 30-resident ceramics clips have identical decoded-video and +PCM audio hashes to the six-resident INT8 clip. These placement and dequant +changes preserve the candidate's output; quality equivalence of INT8 +attention to the original BF16 checkpoint still requires motion/speech review. +Allocator caps emulate available VRAM on a 4090, not another card's speed. +Host peaks include all pod processes; actual 32 GB host-limit support has +not been established. The 30-resident warmup reached 30.26 GiB anonymous +memory and the timed runs reached 77.50 GiB total cgroup usage including cache. + +The 8 GiB cap at `fcdba37fc` completed its warmup but OOMed on the timed +harbor prompt in fine attention. Do not report it as supported. A 34-resident +experiment completed denoising but OOMed during VAE INT8 epilogue allocation. +Both failures motivate subsequent memory work rather than speed claims. + + +## Consumer kernel memory work on the release core + +The consumer branch is rebased onto release core `a97d23f09` (fork PR #45). +Historical measured commits remain reachable through tag +`h3-consumer-fp8-measured-20261003`; rebase changes their branch commit IDs. + +`7a0d7d33b` lets the INT8-QK/BF16-PV fine kernel read BSHD-backed views +without retaining three full BF16 layout copies. Q/K quantization writes +contiguous INT8 arrays and the output stays BHSD; the K-mean reduction keeps +the reference's contiguous reduction order. Four CUDA tests require exact +output equality for partial tiles, empty selections, multiple batches and +partner padding, together with a lower peak allocation. + +`3c0668f6c` adds the opt-in eager +`FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT=1`. One Triton pass applies the INT32 +GEMM's row scale, output-channel scale and optional bias, then casts the +result. FP32 operations retain separate rounding steps (FP fusion disabled). +Strict tests cover zero rows, small input batches, FP32/FP16/BF16, bias, +large INT32 accumulators and tiny scales. Compiled and grad paths retain +the reference implementation. Combine it with shared QKV and transpose +views using `FASTVIDEO_H3_VAE_INT8_SHARED_QKV=1` and +`FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW=1`. + +After the core rebase, 120 focused tests passed (one unrelated GPU cudagraph +check excluded) and pre-commit passed. `87b22a5d8` adapts capture and tests +to packed-segment metadata while retaining the core's calibrated NVFP4 +activation-scale guard. `fb92af176` adds optional NVML total-device-memory +samples every 100 ms, alongside host samples. Summaries distinguish sampled +total GPU usage from PyTorch's allocated peak. Allocator caps omit driver +and external CUDA memory, so the 8 GiB total-budget experiment uses a +7.25 GiB allocator cap and must also satisfy the observed NVML budget. +That trial completed its warmup and first timed request, but sampled total +GPU usage reached about 8.28 GiB. It therefore does **not** meet a strict +8 GiB device target. A tighter allocator cap still needs validation. +These are 4090 simulations; real lower-VRAM and 30-series performance still +needs those devices. + + +## Warmed release-core clip results + +At `fb92af176`, after one warmup and two timed requests: + +| Clip | Median e2e | Timed requests | Conditioning | Denoise | Video decode | Peak allocated | Sampled total GPU | Peak host anon | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| 832×480, 124 frames / 5.167 s | **41.75 s** | 42.16 / 41.33 s | 0.55 s | 32.69 s | 6.63 s | 22.01 GiB | 23.983 GiB | 24.78 GiB | +| 832×480, 243 frames / 10.125 s | **79.67 s** | 79.64 / 79.70 s | 0.55 s | 63.51 s | 13.23 s | 20.29 GiB | 23.985 GiB | 27.02 GiB | + +Warmups took 84.40 / 125.16 s. These wall times include audio, frame export +and MP4 saving, and exclude initial generator construction. The 5 s recipe +keeps 34 DiT blocks resident; the 10 s recipe keeps 30. Both use the same +checkpoint revision and eight DMD forwards as the earlier rows. No profiling, +QKV capture, alternate decoder or sparsity increase is enabled. + +```bash +source /workspace/env.sh +source /workspace/venv/bin/activate +cd /workspace +export PYTHONPATH="/workspace/fastvideo-core:${PYTHONPATH:-}" +export FASTVIDEO_SOURCE_COMMIT=fb92af176 +export FASTVIDEO_H3_PARK_MODULES=vae,audio_vae +export FASTVIDEO_H3_ENCODER_LAYERWISE=1 FASTVIDEO_H3_ENCODER_FUSED_DEQUANT=1 +export FASTVIDEO_H3_VSA_TILE_FIRST=1 FASTVIDEO_H3_VSA_SM89_KERNEL=int8 +export FASTVIDEO_LAYERWISE_RESIDENT_BLOCKS=34 FASTVIDEO_H3_FFN_CHUNK_TOKENS=16384 +export FASTVIDEO_H3_VAE_TILE_BATCH=28 +export FASTVIDEO_H3_VAE_INT8_SHARED_QKV=1 FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW=1 +export FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT=1 MAX_JOBS=4 +unset FASTVIDEO_CUDA_MEMORY_CAP_GIB +python -P /workspace/fastvideo-core/scripts/benchmarks/minimax_h3_4090/bench_pod.py \ + sm89-int8-fast2-480p-5s /workspace/vol/pruned_fp8_300 fp8 \ + --offload-buffers --no-vae-compile --height 480 --width 832 --frames 124 --timed 2 +``` + +For 10 s, set resident blocks to 30, name to `sm89-int8-fast2-480p-10s`, +and frames to 243. Both decoded video and PCM audio hash-identically to the +previous INT8 candidate for ceramics and harbor at 10 s. This validates these +memory/decode changes on those prompts, while the INT8 attention candidate +still differs from the original BF16 attention clips and needs full quality +review. The sampled 5 s contact sheet is coherent. Raw clips, hashes and +results live in `output/fasth3-4090-20261003/` beside the worktree. + +The earlier unprofiled 1344×768, 243-frame run at `fcdba37fc` completed at +279.94 s median (289.06 / 270.82 s, 306.80 s warmup), with 12 resident blocks +and shared VAE QKV/transpose views, before direct-layout attention and fused +VAE epilogues. Its stage medians were 0.55 s conditioning, 227.21 s denoise, +44.82 s video decode and 0.48 s audio. Peak allocation was 21.19 GiB and host +anonymous memory 39.81 GiB. The updated 768p run was queued after the +7.25 GiB allocator-cap trial. SSH became unreachable before the final +results could be collected; updated 768p speed and strict 8 GiB support +remain unverified. + +Track B is staged in [draft PR #46](https://github.com/aryan5v/FastVideo/pull/46), +stacked on the shared release core in #45. Historical timing sources are +preserved by the `h3-consumer-fp8-measured-20261003` tag. diff --git a/scripts/benchmarks/minimax_h3_4090/bench_encoder_dequant.py b/scripts/benchmarks/minimax_h3_4090/bench_encoder_dequant.py new file mode 100644 index 000000000..798793e67 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/bench_encoder_dequant.py @@ -0,0 +1,59 @@ +"""Compare serialized NVFP4 weight expansion on an idle consumer GPU.""" +import argparse +import json +import os +import pathlib +import time + +import torch + +from fastvideo.layers.quantization.nvfp4_dequant import dequantize_nvfp4_cuda +from fastvideo.models.encoders.minimax_h3_checkpoint_nvfp4 import dequantize_serialized_nvfp4 + + +def measure(fn, packed, scales, global_scale): + for _ in range(3): + fn(packed, scales, global_scale) + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + wall = time.perf_counter() + start.record() + for _ in range(10): + fn(packed, scales, global_scale) + stop.record() + stop.synchronize() + return {"gpu_ms": start.elapsed_time(stop) / 10, + "wall_ms": (time.perf_counter() - wall) * 100, + "peak_extra_gib": (torch.cuda.max_memory_allocated() - baseline) / 2**30} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--output", type=pathlib.Path, required=True) + args = ap.parse_args() + rows = [] + global_scale = float(torch.tensor(317.224, dtype=torch.float32)) + torch.manual_seed(63) + for n, k in [(8192, 8192), (25600, 8192), (8192, 25600)]: + packed = torch.randint(0, 256, (n, k // 2), device="cuda", dtype=torch.uint8) + scales = torch.randint(0, 127, (n, k // 16), device="cuda", dtype=torch.uint8) + reference = dequantize_serialized_nvfp4(packed, scales, global_scale) + fused = dequantize_nvfp4_cuda(packed, scales, global_scale) + torch.testing.assert_close(fused, reference, rtol=0, atol=0) + del reference, fused + row = {"n": n, "k": k, "bf16_exact": True, + "torch": measure(dequantize_serialized_nvfp4, packed, scales, global_scale), + "fused": measure(dequantize_nvfp4_cuda, packed, scales, global_scale)} + row["speedup"] = row["torch"]["gpu_ms"] / row["fused"]["gpu_ms"] + print(json.dumps(row), flush=True) + rows.append(row) + del packed, scales + args.output.write_text(json.dumps({"gpu": torch.cuda.get_device_name(), "torch": torch.__version__, + "source_commit": os.environ.get("FASTVIDEO_SOURCE_COMMIT"), + "global_scale": global_scale, "rows": rows}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/minimax_h3_4090/bench_pod.py b/scripts/benchmarks/minimax_h3_4090/bench_pod.py new file mode 100644 index 000000000..b1f4dde62 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/bench_pod.py @@ -0,0 +1,187 @@ +"""Pruned FastH3 8-step benchmark on a single RTX 4090. + +usage: python bench_pod.py [--timed N] [--offload-buffers] [--no-layerwise] +Stage logs (FASTVIDEO_STAGE_LOGGING=1) carry per-stage time and memory peaks; results.json lands in outputs//. +""" +import argparse +import json +import os +import pathlib +import shlex +import statistics +import subprocess +import sys +import threading +import time + + +class HostMemoryPeak: + """Sample pod-wide cgroup usage; anon excludes cached checkpoint file pages.""" + + def __init__(self): + self.stop = threading.Event() + self.peak_bytes = 0 + self.peak_anon_bytes = 0 + self.peak_gpu_bytes = None + self.host_error: str | None = None + self._gpu_used = None + self._nvml_shutdown = None + self.thread = threading.Thread(target=self._sample, daemon=True) + try: + import pynvml + pynvml.nvmlInit() + self._nvml_shutdown = pynvml.nvmlShutdown + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + self._gpu_used = lambda: pynvml.nvmlDeviceGetMemoryInfo(handle).used + except Exception as exc: + print(f"GPU memory sampler unavailable: {exc}", flush=True) + + def _sample(self): + while not self.stop.is_set(): + try: + root = pathlib.Path("/sys/fs/cgroup") + self.peak_bytes = max(self.peak_bytes, int((root / "memory.current").read_text())) + stats = dict(line.split() for line in (root / "memory.stat").read_text().splitlines()) + self.peak_anon_bytes = max(self.peak_anon_bytes, int(stats["anon"])) + if self._gpu_used is not None: + self.peak_gpu_bytes = max(self.peak_gpu_bytes or 0, self._gpu_used()) + except (OSError, KeyError, ValueError) as exc: + self.host_error = f"{type(exc).__name__}: {exc}" + print(f"host memory sampling stopped: {self.host_error}", flush=True) + return + self.stop.wait(0.1) + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *_args): + self.stop.set() + self.thread.join() + if self._nvml_shutdown is not None: + self._nvml_shutdown() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("name"); ap.add_argument("model"); ap.add_argument("quant", choices=("nvfp4", "fp8", "bf16")) + ap.add_argument("--timed", type=int, default=2) + ap.add_argument("--offload-buffers", action="store_true") + ap.add_argument("--no-layerwise", action="store_true") + ap.add_argument("--resident-encoder", action="store_true") + ap.add_argument("--tile-batch", default=None) + ap.add_argument("--frames", type=int, default=243) + ap.add_argument("--height", type=int, default=768) + ap.add_argument("--width", type=int, default=1344) + ap.add_argument("--warmup", type=int, default=1) + ap.add_argument("--prompt-file", type=pathlib.Path, default=pathlib.Path(__file__).with_name("prompts_1k.json")) + ap.add_argument("--output-root", type=pathlib.Path, default=pathlib.Path("/workspace/outputs")) + ap.add_argument("--no-vae-compile", action="store_true") + ap.add_argument("--adaln-cache", action="store_true") + ap.add_argument("--profile", action="store_true") + ap.add_argument("--sparsity", type=float, default=0.8) + ap.add_argument("--decode", default="h3-vae") + ap.add_argument("--lazy", action="store_true", help="lazy_module_load: reload released modules per request") + ap.add_argument("--prompts", default=None, help="comma-separated prompt ids (default: both)") + a = ap.parse_args() + if a.timed < 2 or a.warmup < 1: + ap.error("Use at least one warmup and two timed runs") + if not pathlib.Path(a.model, "fastvideo_inference.json").is_file(): + ap.error("The model directory must contain fastvideo_inference.json for the 8-step DMD contract") + + os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1") + # NVFP4 is retained for other GPUs; the RTX 4090 DiT uses FP8. + if a.quant == "nvfp4": + os.environ.setdefault("FASTVIDEO_H3_VSA_FP4", "1") + os.environ.setdefault("FASTVIDEO_NVFP4_MM_BACKEND", "cutlass") + os.environ.setdefault("FASTVIDEO_MINIMAX_H3_FUSIONS", "all") + os.environ.setdefault("FASTVIDEO_H3_VAE_TILE_BATCH", "28") + os.environ.setdefault("FASTVIDEO_VSA_TRITON", "1") + os.environ.setdefault("FASTVIDEO_VSA_SM100A", "0") + os.environ.setdefault("FASTVIDEO_FA4", "0") + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + ap_table = os.environ.get("FASTVIDEO_H3_ADALN_TABLE") + if a.profile: + os.environ["FASTVIDEO_H3_SP_PROFILE"] = "1" + if a.adaln_cache and not ap_table: + os.environ["FASTVIDEO_H3_ADALN_CACHE"] = "1" + # Export the exact modulation tables (and inputs) for table-only loads and low-rank experiments. + os.environ.setdefault("FASTVIDEO_H3_ADALN_DUMP", f"/workspace/adaln_tables_{a.name}.pt") + if a.offload_buffers: + os.environ["FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS"] = "1" + if a.tile_batch: + os.environ["FASTVIDEO_H3_VAE_TILE_BATCH"] = a.tile_batch + import torch + from fastvideo import VideoGenerator + + texts = json.loads(a.prompt_file.read_text()) + layerwise = not a.no_layerwise + engine = {"num_gpus": 1, "use_fsdp_inference": False, + "parallelism": {"tp_size": 1, "sp_size": 1}, + "offload": {"dit": False, "dit_layerwise": layerwise, "text_encoder": not a.resident_encoder, + "vae": layerwise, "pin_cpu_memory": True, "lazy_module_load": a.lazy}, + "compile": {"enabled": False, "vae_enabled": not a.no_vae_compile}} + if a.quant == "nvfp4": + engine["quantization"] = {"transformer_quant": "NVFP4", "layer_profile": "h3_dit_vsa"} + elif a.quant == "fp8": + engine["quantization"] = {"transformer_quant": "FP8"} + experimental = {"attention_backend": "VIDEO_SPARSE_ATTN_H3", "VSA_sparsity": a.sparsity, "VSA_tile_size": 64, + "h3_sequential_load": not a.resident_encoder, "inference_torch_compile": False, + "vae_parallel_decode": False, "video_decode_backend": a.decode} + config = {"model_path": a.model, "engine": engine, "pipeline": {"experimental": experimental}} + out_dir = a.output_root / a.name + out_dir.mkdir(parents=True, exist_ok=True) + model_root = pathlib.Path(a.model) + revision_file = model_root / ".cache/huggingface/download/fastvideo_inference.json.metadata" + model_revision = revision_file.read_text().splitlines()[0] if revision_file.is_file() else None + hardware = subprocess.check_output( + ["nvidia-smi", "--query-gpu=name,memory.total,driver_version,pci.bus_id", "--format=csv,noheader"], text=True + ).strip() + sampling = {"seed": 20260929, "height": a.height, "width": a.width, "num_frames": a.frames, "fps": 24, + "num_inference_steps": 9, "guidance_scale": 1.0, "batch_cfg": False} + results = {"name": a.name, "quant": a.quant, "command": shlex.join([sys.executable, "-P", *sys.argv]), + "env": {k: v for k, v in os.environ.items() if k.startswith(("FASTVIDEO_", "PYTORCH_")) + or k in ("CUDA_VISIBLE_DEVICES", "MAX_JOBS")}, + "torch": torch.__version__, "cuda": torch.version.cuda, + "hardware": hardware, "model_revision": model_revision, + "model_contract": json.loads((model_root / "fastvideo_inference.json").read_text()), + "source_commit": os.environ.get("FASTVIDEO_SOURCE_COMMIT"), + "gpu": torch.cuda.get_device_name(0), "config": config, "sampling": sampling, "runs": []} + (out_dir / "results.json").write_text(json.dumps(results, indent=2)) + t0 = time.perf_counter() + generator = VideoGenerator.from_config(config) + results["load_s"] = round(time.perf_counter() - t0, 1) + ids = a.prompts.split(",") if a.prompts else list(texts) + order = [ids[i % len(ids)] for i in range(a.warmup + a.timed)] + try: + for i, pid in enumerate(order): + request = {"prompt": texts[pid], "negative_prompt": "", + "sampling": sampling, + "output": {"output_path": str(out_dir / f"{i:02d}_{pid}.mp4"), "save_video": True, + "return_frames": False}} + t = time.perf_counter() + with HostMemoryPeak() as host_peak: + generator.generate(request) + wall = round(time.perf_counter() - t, 2) + results["runs"].append({"prompt": pid, "warmup": i < a.warmup, "wall_s": wall, + "clip": request["output"]["output_path"], + "peak_gpu_used_gib": (round(host_peak.peak_gpu_bytes / 2**30, 3) + if host_peak.peak_gpu_bytes is not None else None), + # None when sampling failed: an unmeasured run must not read as 0 GiB. + "peak_host_cgroup_gib": (round(host_peak.peak_bytes / 2**30, 3) + if host_peak.host_error is None else None), + "peak_host_anon_gib": (round(host_peak.peak_anon_bytes / 2**30, 3) + if host_peak.host_error is None else None), + "host_memory_error": host_peak.host_error}) + timed = [run["wall_s"] for run in results["runs"] if not run["warmup"]] + if timed: + results["median_e2e_s"] = statistics.median(timed) + print("RUN", json.dumps(results["runs"][-1]), flush=True) + (out_dir / "results.json").write_text(json.dumps(results, indent=1)) + finally: + generator.shutdown() + print("DONE", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/minimax_h3_4090/bench_sparse_qkv.py b/scripts/benchmarks/minimax_h3_4090/bench_sparse_qkv.py new file mode 100644 index 000000000..3bf22496c --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/bench_sparse_qkv.py @@ -0,0 +1,57 @@ +"""Validate and time candidate kernels on captured real H3 Q/K/V. + +Run on an idle sm89 GPU, separately from clip benchmarks. Captures come from +FASTVIDEO_H3_CAPTURE_QKV on the tile-first path; they retain two full heads. +""" +import argparse +import json +import pathlib + +import torch +import triton + +from fastvideo.attention.backends.minimax_h3_sparse_int8 import sparse_sm89_attention +from fastvideo_kernel.block_sparse_attn import block_sparse_attn + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("captures", type=pathlib.Path) + parser.add_argument("--output", type=pathlib.Path, required=True) + args = parser.parse_args() + records = [] + with torch.inference_mode(): + for capture in sorted(args.captures.glob("layer-*.pt")): + state = torch.load(capture, map_location="cuda", weights_only=True) + q, k, v, mask, vbs = (state[key] for key in ("q", "k", "v", "mask", "vbs")) + def baseline(): + return block_sparse_attn(q, k, v, mask, vbs)[0] + expected = baseline() + for int8_qk in (True, False): + def candidate(): + return sparse_sm89_attention(q, k, v, mask, vbs, int8_qk=int8_qk) + output = candidate() + valid_rows = state["untile"] + ref = expected.index_select(2, valid_rows).float() + actual = output.index_select(2, valid_rows).float() + delta = actual - ref + reference_ms = triton.testing.do_bench(baseline) + candidate_ms = triton.testing.do_bench(candidate) + record = {"capture": capture.name, "shape": list(q.shape), "int8_qk": int8_qk, + "mask_density": float(mask.float().mean()), + "finite": bool(torch.isfinite(actual).all()), + "relative_l2": float(delta.norm() / ref.norm()), + "max_abs": float(delta.abs().max()), + "cosine": float(torch.nn.functional.cosine_similarity(actual.flatten(), ref.flatten(), dim=0)), + "bf16_ms": reference_ms, "int8_fp8_ms": candidate_ms, + "speedup": reference_ms / candidate_ms} + print(json.dumps(record), flush=True) + records.append(record) + if not records: + raise RuntimeError("No real Q/K/V captures found") + args.output.write_text(json.dumps({"gpu": torch.cuda.get_device_name(), "torch": torch.__version__, + "triton": triton.__version__, "records": records}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/minimax_h3_4090/prompts_1k.json b/scripts/benchmarks/minimax_h3_4090/prompts_1k.json new file mode 100644 index 000000000..553a5a7d6 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/prompts_1k.json @@ -0,0 +1,4 @@ +{ + "latency-ceramics-005": "In a quiet pottery studio an adult potter steadies a small spinning bowl while an adult apprentice watches. The apprentice asks, \"Is the rim ready?\" The potter says, \"One more gentle pass,\" and smooths the lip with a damp sponge. Begin with a close view of the hands, then make one restrained cut to a shoulder-level view showing both faces. The wheel hum, damp clay, a light splash and breathing form the soundscape. The movement is careful and unhurried, with no background music. The entire event is one finishing pass on the same bowl, not a demonstration of the whole pottery process.\nThe apprentice's apron pocket is indigo, with a plain surface. This detail belongs to the existing object, stays at its established location and remains subordinate to the main action. Preserve its material and appearance through the camera movement.\nThe wheel sits near a tall north-facing window on the left wall of the studio.\nIts circular splash pan is at the seated potter's waist.\nThe apprentice stands beyond the right side of the wheel, with enough room between their clothes and the wet rim to avoid accidental contact.\nBehind the potter, a shallow shelf holds a few dry cups.\nA sink is farther back on the right, beneath a small rack of towels.\nThese positions remain fixed through the cut.\nThe wider view reveals the same work area that surrounds the hands in the opening close view.\nThe bowl is modest in size, comfortably held between two hands.\nIts side rises from a low foot in a continuous shallow curve and ends in a slightly thickened lip.\nThe clay is warm gray with a brown undertone, darker where it is wet.\nA narrow spiral line below the rim records the potter's earlier touch.\nThat line rotates with the bowl, while a stationary highlight from the window travels over the passing surface.\nKeep these two motions distinguishable: the material marks belong to the clay, and the reflected light belongs to the relationship between the surface and the window.\nOpen with the camera low enough to see the bowl's interior without looking directly down onto the wheel.\nThe potter's left fingertips support the inside edge.\nThe right hand holds a small natural sponge outside the lip.\nNeither hand blocks the whole form.\nThe near edge of the splash pan occupies the bottom of the composition as a soft curved boundary.\nBeyond the hands, the apprentice's apron is initially out of focus.\nThis arrangement gives the close shot depth and prepares the later view of the two people without requiring another establishing shot.\nThe left hand is already stable when the clip begins.\nIts fingers form a loose supportive curve rather than squeezing the wall.\nThe sponge approaches the outside edge with only a small adjustment of the wrist.\nAs it contacts the clay, it compresses slightly and darkens where moisture gathers.\nThe bowl continues to rotate at a steady moderate speed.\nShow the finishing pass as a change in the surface's smoothness and the evenness of the lip, not as a large change in the bowl's overall shape.\nThe work is nearly finished before this moment begins.\nThe apprentice asks the question while looking at the rim, then briefly lifts their eyes toward the potter.\nTheir hands rest loosely together in front of the apron, safely away from the rotating work.\nThe question is curious and quiet, with the natural upward inflection of someone checking a detail.\nThe potter answers without stopping the wheel or turning their whole body.\nA small glance toward the apprentice is sufficient before attention returns to the clay.\nKeep the spoken words exactly as given, with no narrator explaining the technique and no extra exchange after the answer.\nCut once after the question has made the apprentice's presence clear.\nThe shoulder-level view places the potter to the left and the apprentice to the right, preserving the established relation to the wheel.\nThe bowl remains visible between them in the lower part of the frame.\nThe potter's right hand still holds the same sponge at the same point on the rim.\nContinue the wheel sound across the cut without a restart.\nThe change of view should feel like a closer understanding of the same instant, not a jump forward to a later stage of the work.\nThe potter wears a practical cotton work shirt with the sleeves rolled above the wrists.\nThe folds gather at the elbows and remain dry there.\nSmall clay marks on the forearms and apron are concentrated near the work area.\nThey do not spread or migrate during the pass.\nThe apprentice's apron is cleaner but shows a few dry pale smudges near one pocket.\nBoth garments have weight and ordinary creases.\nAvoid pristine costumes or exaggerated distressing; this is a used studio where people work carefully and clean their tools regularly.\nGive the potter a focused, patient expression.\nTheir mouth moves only for the brief reply, then settles while they feel the rim through the sponge.\nAllow the last small movement to settle within the established composition, with the environmental sound continuing around it. Preserve the quiet final composition.", + "latency-harbor-005": "## Harbor: the tide chart\nThe fabric cover around the folded tide chart is indigo, with a plain surface. This detail belongs to the existing object, stays at its established location and remains subordinate to the main action. Preserve its material and appearance through the camera movement.\nThe event takes place beside a small passenger ferry tied to a working harbor pier just before sunrise. An adult mechanic stands on the ferry's open side deck, and its captain stands beside the cabin entrance. The mechanic offers a folded tide chart and says, \"The channel is clear.\" The captain accepts it, answers, \"Then we can go,\" and looks out toward the harbor entrance. A slow lateral camera move reveals the channel beyond their shoulders. The boat remains moored throughout this brief exchange. Close voices, water against the hull, a loose halyard and a distant gull make the soundscape; there is no music.\nThe ferry is a practical coastal launch with a dark blue hull and a narrow cream band beneath its windows.\nIt carries a small enclosed cabin forward and an open passenger area behind it.\nThe camera is on the open deck, looking diagonally toward the cabin and the gap between the two people.\nThis angle places the pier along the left edge of the view and open water farther to the right.\nThe horizon is low enough that the upper part of the cabin has a clear silhouette against the pale sky.\nNothing in the composition suggests that the ferry is already underway.\nThe mechanic has finished a routine inspection rather than an emergency repair.\nTheir expression is alert but comfortable, with the slight tiredness of an early start.\nThey wear a plain work jacket over a warm shirt and carry no conspicuous badge or brand.\nA few old creases in the jacket show where the elbows bend.\nThe sleeve nearest the chart has a darker damp patch near its cuff from resting against the rail.\nKeep that patch in the same place as the arm moves.\nThe mechanic's free hand rests lightly on the top of a closed tool bag at hip level.\nThe captain is a different adult, dressed for a cool morning outside.\nA heavy knit sweater is visible beneath an open weatherproof coat.\nTheir hair is tidy but not freshly styled, and the light catches a few loose strands when they turn toward the water.\nTheir stance is balanced on the gently moving deck, with one foot slightly ahead of the other.\nThey are listening to the mechanic before the first line begins.\nThe captain does not interrupt or make a broad theatrical gesture.\nTheir reply is a small decision shared between people accustomed to working together.\nThe tide chart is a real paper object with several old folds.\nIt is partly folded into a rectangle that can be held in one hand, but one narrow flap remains loose.\nFaint printed lines and numbers are visible as a texture on its pale surface without becoming a readable title or a map inset.\nA soft graphite mark near one fold suggests that it has been used for planning.\nThe mechanic holds its lower edge between the thumb and fingers, keeping the paper clear of the damp rail.\nIts upper corner lifts slightly in the breeze before the captain takes it.\nBegin with both people already in the frame.\nThe mechanic's hand and the chart occupy the space between their bodies, below their faces.\nThis arrangement lets the first line and the handover belong to the same view.\nAs the mechanic speaks, the chart moves a short distance toward the captain.\nThe motion is neither a flourish nor an abrupt thrust.\nThe captain's receiving hand rises from beside the coat, touches the opposite edge and supports it before the mechanic releases their grip.\nThe paper bends a little between the two hands during that shared moment of support.\nThe mechanic's line is spoken in an ordinary low voice suitable for the quiet morning. The consonants remain clear, but the delivery does not sound like a public announcement. Their mouth and jaw make the small movements of the exact words, and their eyes remain on the captain. There is a slight release of breath after \"clear.\" The captain acknowledges the information with a very small nod before replying. The pause is long enough to register listening and short enough that the exchange feels familiar. Do not add another question, greeting or explanation of the voyage.\nWhen the captain says, \"Then we can go,\" the first part of the line is addressed to the mechanic. On the last words, their gaze begins to move toward the channel. The head follows the eyes through a modest turn, revealing more of the cheek nearest the exterior light. The chart settles against the front of the coat, still visibly held. The mechanic follows the captain's look with a quieter change of attention. Both remain in place. The ending is anticipation of departure, not departure itself: no engine surge, released rope or sudden movement of the ferry is needed.\nThe camera makes a restrained lateral movement toward the open-water side of the deck.\nKeep the final gesture restrained and preserve the surrounding atmosphere. Preserve the quiet final composition as the scene reaches its stated resolution." +} \ No newline at end of file diff --git a/scripts/benchmarks/minimax_h3_4090/summarize.py b/scripts/benchmarks/minimax_h3_4090/summarize.py new file mode 100644 index 000000000..fffecde50 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_4090/summarize.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Join benchmark results with per-request stage logs, excluding warmup runs.""" + +import argparse +import json +import re +import statistics +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("results", type=Path) + parser.add_argument("log", type=Path) + args = parser.parse_args() + results = json.loads(args.results.read_text()) + completed = [] + stages = {} + peaks = [] + for line in args.log.read_text().splitlines(): + timing = re.search(r"\[(\w+)_stage\|[^\]]+\] Execution completed in ([\d.]+) ms", line) + if timing: + stages[timing[1]] = float(timing[2]) / 1000 + peak = re.search(r"Memory peak_allocated=([\d.]+) GiB", line) + if peak: + peaks.append(float(peak[1])) + if line.startswith("RUN "): + run = json.loads(line[4:]) + run["stage_s"] = stages + run["peak_gpu_allocated_gib"] = max(peaks) if peaks else None + completed.append(run) + stages = {} + peaks = [] + if len(completed) != len(results["runs"]): + raise ValueError("Log and results.json have different completed run counts") + timed = [run for run in completed if not run["warmup"]] + if len(timed) < 2: + raise ValueError("At least two timed runs are required for a baseline summary") + for run in timed: + if not {"denoising", "video_decoding", "audio_decoding"}.issubset(run["stage_s"]): + raise ValueError("Missing stage timings for a completed run") + summary = { + "name": results["name"], + "sampling": results["sampling"], + "source_commit": results["source_commit"], + "timed_runs": len(timed), + "median_e2e_s": statistics.median(run["wall_s"] for run in timed), + "median_stage_s": {name: statistics.median(run["stage_s"][name] for run in timed) + for name in ("conditioning", "denoising", "video_decoding", "audio_decoding")}, + "peak_gpu_allocated_gib": max(run["peak_gpu_allocated_gib"] for run in timed), + "peak_gpu_used_gib": max((run["peak_gpu_used_gib"] for run in timed + if run.get("peak_gpu_used_gib") is not None), default=None), + "peak_host_anon_gib": max(run["peak_host_anon_gib"] for run in timed), + "peak_host_cgroup_gib": max(run["peak_host_cgroup_gib"] for run in timed), + "notes": "Stage times include deferred component loading. Host peaks are pod-wide samples every 100 ms.", + "runs": timed, + } + args.results.with_name(f"{args.results.stem}-summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/minimax_h3_pro6000/a2a8.py b/scripts/benchmarks/minimax_h3_pro6000/a2a8.py new file mode 100644 index 000000000..f2659517e --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/a2a8.py @@ -0,0 +1,72 @@ +"""All-to-all bandwidth on 8x RTX PRO 6000 at FastH3 Ulysses SP8 payload sizes (10 s, 1344x768: 73.6k tokens).""" +import modal + +image = (modal.Image.debian_slim(python_version="3.12") + .pip_install("torch==2.8.0", index_url="https://download.pytorch.org/whl/cu128")) +app = modal.App("h3-a2a8", image=image) + +TOKENS, WIDTH = 73642, 7168 +WORLD = 8 + + +def _worker(rank: int, env: dict, out_q) -> None: + import os + import time + import torch + import torch.distributed as dist + os.environ.update(env) + os.environ.update({"MASTER_ADDR": "127.0.0.1", "MASTER_PORT": "29511"}) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=WORLD) + rows = -(-TOKENS // WORLD) + payloads = { + "qkvg_bf16": 4 * rows * WIDTH * 2, # q, k, v, gate shards before attention + "qkv_fp4_plus_gate_bf16": 3 * rows * WIDTH * 9 // 16 + rows * WIDTH * 2, + "qkvg_fp4": 4 * rows * WIDTH * 9 // 16, # 4-bit values + one e4m3 scale per 16 + "out_bf16": rows * WIDTH * 2, # attention output back to sequence shards + } + res = {} + for name, nbytes in payloads.items(): + n = (nbytes // 2 // WORLD) * WORLD + send = torch.empty(n, dtype=torch.bfloat16, device="cuda") + recv = torch.empty_like(send) + for _ in range(3): + dist.all_to_all_single(recv, send) + torch.cuda.synchronize() + dist.barrier() + iters = 10 + t = time.perf_counter() + for _ in range(iters): + dist.all_to_all_single(recv, send) + torch.cuda.synchronize() + ms = (time.perf_counter() - t) / iters * 1e3 + res[name] = {"mb_per_rank": round(n * 2 / 1e6, 1), "ms": round(ms, 2), + "algbw_gbps": round(n * 2 / ms / 1e6, 1)} + if rank == 0: + out_q.put(res) + dist.destroy_process_group() + + +@app.function(gpu="RTX-PRO-6000:8", cpu=16, memory=65536, timeout=1800) +def bench() -> dict: + import subprocess + import torch.multiprocessing as mp + report = {"topo": subprocess.run(["nvidia-smi", "topo", "-m"], capture_output=True, text=True).stdout[-3000:]} + for label, env in (("p2p_default", {}), ("p2p_disabled", {"NCCL_P2P_DISABLE": "1"})): + ctx = mp.get_context("spawn") + q = ctx.Queue() + procs = [ctx.Process(target=_worker, args=(r, env, q)) for r in range(WORLD)] + for p in procs: + p.start() + report[label] = q.get(timeout=900) + for p in procs: + p.join() + return report + + +@app.local_entrypoint() +def main(): + import json + r = bench.remote() + print(r.pop("topo")) + print("A2A", json.dumps(r, indent=1)) diff --git a/scripts/benchmarks/minimax_h3_pro6000/app.py b/scripts/benchmarks/minimax_h3_pro6000/app.py new file mode 100644 index 000000000..40e087a63 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/app.py @@ -0,0 +1,613 @@ +"""FastH3 V2 NVFP4 on one RTX PRO 6000 (sm_120): build, convert, time generations.""" +import json +import os +import pathlib +import shlex +import subprocess +import time + +import modal + +_HERE = pathlib.Path(__file__).resolve() +# Inside the container this file is /root/app.py; the repo is already baked into the image there. +WORKTREE = _HERE.parents[3] if len(_HERE.parents) > 3 else pathlib.Path("/src/fastvideo") +CUTLASS_COMMIT = "e67e63c331d6e4b729047c95cf6b92c8454cba89" +volume = modal.Volume.from_name("h3-pro6000-weights", create_if_missing=True) + +image = ( + modal.Image.from_registry("nvidia/cuda:13.0.1-devel-ubuntu24.04", add_python="3.12") + .apt_install("git", "build-essential", "ffmpeg", "libgl1", "libglib2.0-0") + .pip_install("uv") + .add_local_dir(WORKTREE, "/src/fastvideo", copy=True, + ignore=[".git", "**/__pycache__", "fastvideo-kernel/include/cutlass/**", + "fastvideo-kernel/include/tk/**", "fastvideo/third_party/eval/**", "docs/**", + "assets/**", "comfyui/**", "apps/**", "**/*.mp4", "**/*.log"]) + .run_commands("cd /src/fastvideo && UV_TORCH_BACKEND=cu130 uv pip install --system -e . --no-sources") + .run_commands("uv pip install --system 'cmake==3.31.6' ninja 'scikit-build-core>=0.10' pybind11 hf_transfer") + .run_commands(f"git clone --filter=blob:none https://github.com/NVIDIA/cutlass.git /cutlass && " + f"git -C /cutlass checkout {CUTLASS_COMMIT}") + .env({"FLASHINFER_CUDA_ARCH_LIST": "12.0a", "FLASHINFER_WORKSPACE_BASE": "/vol/cache/flashinfer", + "TORCHINDUCTOR_CACHE_DIR": "/vol/cache/inductor", "TRITON_CACHE_DIR": "/vol/cache/triton", + "FASTVIDEO_VSA_SM100A": "0", "FASTVIDEO_FA4": "0", "FASTVIDEO_STAGE_LOGGING": "1", + "HF_HUB_ENABLE_HF_TRANSFER": "1"}) +) +app = modal.App("h3-pro6000-fastest", image=image) + + +def _sh(cmd: str, **kw) -> str: + # Fixed in-container commands that need cd, && and globs; quote any caller-supplied value. + proc = subprocess.run(cmd, shell=True, capture_output=True, text=True, **kw) + out = (proc.stdout + proc.stderr)[-8000:] + if proc.returncode != 0: + raise RuntimeError(f"command failed ({proc.returncode}): {cmd}\n{out}") + return out + + +@app.function(cpu=16, memory=32768, timeout=3600, volumes={"/vol": volume}) +def build_kernel() -> str: + _sh("rm -rf /src/fastvideo/fastvideo-kernel/include/cutlass && " + "ln -s /cutlass /src/fastvideo/fastvideo-kernel/include/cutlass && mkdir -p /vol/wheels/cu130") + _sh("rm -f /vol/wheels/cu130/*.whl") + env = dict(os.environ, TORCH_CUDA_ARCH_LIST="12.0a", MAX_JOBS="16", CC="gcc", CXX="g++", CUDAHOSTCXX="g++", + CMAKE_ARGS="-DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DFASTVIDEO_KERNEL_BUILD_TK=OFF " + "-DGPU_BACKEND=CUDA -DCMAKE_CUDA_ARCHITECTURES=120a") + _sh("cd /src/fastvideo/fastvideo-kernel && pip wheel . --no-build-isolation --no-deps -w /vol/wheels/cu130", + env=env) + volume.commit() + return _sh("ls -la /vol/wheels/cu130") + + +def _install_kernel(): + _sh("pip install --no-deps --force-reinstall /vol/wheels/cu130/*.whl") + + +def _build_light_int8_vae() -> str: + """26-block LynnReal light decoder: dense fp16 decoder + official encoder, Kijai int8-convrot overlay.""" + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + target = pathlib.Path("/vol/fv/vae_light_int8") + if (target / "config.json").exists(): + return "exists" + target.mkdir(parents=True, exist_ok=True) + overlay = hf_hub_download("Kijai/MiniMax-H3-experimental", "minimax_h3_lynnreal_light_vae_int8_convrot.safetensors", + local_dir="/vol/kijai") + official = pathlib.Path("/vol/official/vae") + weight_map = json.loads((official / "diffusion_pytorch_model.safetensors.index.json").read_text())["weight_map"] + tensors = {} + for shard in sorted(set(weight_map.values())): + with safe_open(str(official / shard), framework="pt") as reader: + for key in reader.keys(): + if not key.startswith("decoder.transformer_blocks."): + tensors[key] = reader.get_tensor(key) + with safe_open("/vol/light-vae/lynnreal_light_vae_decoder_fp16.safetensors", framework="pt") as reader: + light_keys = list(reader.keys()) + for key in light_keys: + tensors[key] = reader.get_tensor(key) + blocks = {int(k.split(".")[2]) for k in tensors if k.startswith("decoder.transformer_blocks.")} + assert blocks == set(range(26)), sorted(blocks) + save_file(tensors, str(target / "diffusion_pytorch_model.safetensors")) + config = json.loads((official / "config.json").read_text()) + config["decoder_num_layers"] = 26 + (target / "config.json").write_text(json.dumps(config, indent=2) + "\n") + (target / "minimax_h3_video_vae_int8_convrot.safetensors").symlink_to(overlay) + return f"light decoder keys={len(light_keys)} total={len(tensors)} blocks={len(blocks)}" + + +@app.function(gpu="RTX-PRO-6000", memory=131072, cpu=8, timeout=10800, volumes={"/vol": volume}) +def convert(minimal: bool = True) -> dict: + _install_kernel() + report = {} + os.makedirs("/vol/fv", exist_ok=True) + if not os.path.exists("/vol/fv/text_encoder_nvfp4/config.json"): + report["text_encoder"] = _sh( + "cd /src/fastvideo && python scripts/checkpoint_conversion/convert_minimax_h3_text_encoder_nvfp4.py " + "--src /vol/v2-nvfp4/text_encoder --dst /vol/fv/text_encoder_nvfp4")[-1500:] + volume.commit() + for name, src, flag in (("transformer_full", "v2-nvfp4", "--quantize-attention"), + ("transformer_ffn", "v2-nvfp4", ""), + ("transformer_vsa", "v2-nvfp4", "--quantize-attention --quantize-gate"), + ("v4_transformer_vsa", "v4-nvfp4", "--quantize-attention --quantize-gate")): + if not os.path.exists(f"/vol/{src}/transformer"): + continue + if minimal and name not in ("transformer_vsa", "v4_transformer_vsa"): + continue + if not os.path.exists(f"/vol/fv/{name}/config.json"): + report[name] = _sh( + "cd /src/fastvideo && python scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py " + f"--src /vol/{src}/transformer --dst /vol/fv/{name} {flag}")[-1500:] + volume.commit() + # VAE folders: official dense shards, plus Comfy's int8-convrot overlay variant. + for vae_name in (() if minimal else ("vae_dense", "vae_int8")): + target = pathlib.Path(f"/vol/fv/{vae_name}") + target.mkdir(parents=True, exist_ok=True) + for item in pathlib.Path("/vol/official/vae").iterdir(): + link = target / item.name + if not link.exists(): + link.symlink_to(item) + int8 = pathlib.Path("/vol/fv/vae_int8/minimax_h3_video_vae_int8_convrot.safetensors") + if not minimal and not int8.exists(): + int8.symlink_to("/vol/comfy/vae/minimax_h3_video_vae_int8_convrot.safetensors") + report["light_vae"] = _build_light_int8_vae() + # Model folders: V2 small components + chosen transformer / text encoder / VAE. + # 4-step VSA-0.9 model: its own manifest/schedulers, shared text encoder, VAEs and audio VAE. + root = pathlib.Path("/vol/fv/v4_vsa_light") + if pathlib.Path("/vol/v4-nvfp4/fastvideo_inference.json").exists(): + root.mkdir(parents=True, exist_ok=True) + for item in pathlib.Path("/vol/v4-nvfp4").iterdir(): + if item.name in ("transformer", "text_encoder", "vae", "audio_vae") or item.name.startswith("."): + continue + if not (root / item.name).exists(): + (root / item.name).symlink_to(item) + for comp, src in (("transformer", "/vol/fv/v4_transformer_vsa"), ("text_encoder", "/vol/fv/text_encoder_nvfp4"), + ("vae", "/vol/fv/vae_light_int8"), ("audio_vae", "/vol/v2-nvfp4/audio_vae")): + if not (root / comp).exists(): + (root / comp).symlink_to(src) + model_sets = (("v2_vsa_light", "transformer_vsa", "vae_light_int8"), + ("v2_full_light", "transformer_full", "vae_light_int8"), + ("v2_full_int8", "transformer_full", "vae_int8"), + ("v2_full_dense", "transformer_full", "vae_dense"), + ("v2_ffn_int8", "transformer_ffn", "vae_int8")) + for model, transformer, vae in (model_sets[:1] if minimal else model_sets): + root = pathlib.Path(f"/vol/fv/{model}") + root.mkdir(parents=True, exist_ok=True) + for item in pathlib.Path("/vol/v2-nvfp4").iterdir(): + if item.name in ("transformer", "text_encoder", "vae") or item.name.startswith("."): + continue + link = root / item.name + if not link.exists(): + link.symlink_to(item) + for comp, src in (("transformer", transformer), ("text_encoder", "text_encoder_nvfp4"), ("vae", vae)): + link = root / comp + if not link.exists(): + link.symlink_to(f"/vol/fv/{src}") + volume.commit() + report["models"] = _sh("ls -la /vol/fv /vol/fv/v2_vsa_light") + # Are V2's VSA compression gates trained (nonzero)? + import torch + from safetensors import safe_open + index = json.loads(pathlib.Path("/vol/v2-nvfp4/transformer/diffusion_pytorch_model.safetensors.index.json").read_text()) + gate_stats = {} + for blk in (0, 25, 49): + key = f"transformer_blocks.{blk}.attn.to_gate_compress.weight" + with safe_open(f"/vol/v2-nvfp4/transformer/{index['weight_map'][key]}", framework="pt") as reader: + w = reader.get_tensor(key).float() + gate_stats[blk] = {"abs_mean": w.abs().mean().item(), "nonzero_frac": (w != 0).float().mean().item()} + report["gate_stats"] = gate_stats + return report + + +PROMPTS = { + "kitesurf": ("A kite surfer carves hard across choppy bay water while the camera dives alongside; spray hisses off " + "the board edge, the sail flaps and snaps in the wind, and gulls cry overhead."), + "chef": ("(S1) In a bright home kitchen, a chef looks straight at the camera and says [English] Fold the eggs " + "gently and taste before you salt. A pot simmers behind her with soft bubbling and no music."), +} + + +@app.function(gpu="RTX-PRO-6000", memory=131072, cpu=8, timeout=5400, volumes={"/vol": volume}) +def run_variant(name: str, model: str, profile: str, attention: str, decode: str, vae_compile: bool, + height: int = 480, width: int = 832, warmups: int = 1, num_frames: int = 124, + env: dict | None = None, prompts: tuple = ("kitesurf", "chef"), sparsity: float = 0.8, + steps: int = 9, num_gpus: int = 1, parallel_decode: bool = False, pre_runs: tuple = (), + prompt_texts: dict | None = None, offload: dict | None = None, + experimental_extra: dict | None = None) -> dict: + os.environ.update(env or {}) + texts = {**PROMPTS, **(prompt_texts or {})} + _install_kernel() + import torch + from fastvideo import VideoGenerator + + if attention == "VIDEO_SPARSE_ATTN_H3": + os.environ["FASTVIDEO_VSA_TRITON"] = "1" + experimental = {"attention_backend": attention, "h3_sequential_load": False, "inference_torch_compile": False, + "vae_parallel_decode": parallel_decode, "video_decode_backend": decode} + if attention == "VIDEO_SPARSE_ATTN_H3": + experimental.update({"VSA_sparsity": sparsity, "VSA_tile_size": 64}) + experimental.update(experimental_extra or {}) + config = { + "model_path": f"/vol/fv/{model}", + "engine": {"num_gpus": num_gpus, "use_fsdp_inference": False, + "quantization": {"transformer_quant": "NVFP4", "layer_profile": profile}, + "parallelism": {"tp_size": 1, "sp_size": num_gpus}, + "offload": {"dit": False, "dit_layerwise": False, "text_encoder": False, "vae": False, + "pin_cpu_memory": num_gpus == 1, "lazy_module_load": False, **(offload or {})}, + "compile": {"enabled": False, "vae_enabled": vae_compile}}, + "pipeline": {"experimental": experimental}, + } + t0 = time.perf_counter() + generator = VideoGenerator.from_config(config) + load_s = time.perf_counter() - t0 + out_dir = pathlib.Path(f"/vol/outputs/{name}") + out_dir.mkdir(parents=True, exist_ok=True) + results = {"name": name, "load_s": round(load_s, 1), "env": env or {}, "shape": [height, width, num_frames], + "num_gpus": num_gpus, + "runs": []} + try: + # Optional runs at other shapes first (e.g. a 480p correctness clip), same loaded model. + for j, (ph, pw, pf, pid) in enumerate(pre_runs): + request = {"prompt": texts[pid], "negative_prompt": "", + "sampling": {"seed": 20260929, "height": ph, "width": pw, "num_frames": pf, "fps": 24, + "num_inference_steps": steps, "guidance_scale": 1.0, "batch_cfg": False}, + "output": {"output_path": str(out_dir / f"pre{j:02d}_{pid}_{ph}p.mp4"), "save_video": True, + "return_frames": False}} + t = time.perf_counter() + result = generator.generate(request) + results.setdefault("pre_runs", []).append({"shape": [ph, pw, pf], "prompt": pid, + "wall_s": round(time.perf_counter() - t, 2), + "video": getattr(result, "video_path", None)}) + print("PRE_RUN", json.dumps(results["pre_runs"][-1]), flush=True) + if pre_runs: + (out_dir / "results.json").write_text(json.dumps(results, indent=1)) + volume.commit() + # Warm every distinct prompt: prompt length changes the packed sequence, and shape-specialized + # compiled kernels would otherwise recompile inside the first timed run of each prompt. + distinct = list(dict.fromkeys(prompts)) + order = [distinct[i % len(distinct)] for i in range(warmups)] + list(prompts) + for i, pid in enumerate(order): + prompt = texts[pid] + request = {"prompt": prompt, "negative_prompt": "", + "sampling": {"seed": 20260929, "height": height, "width": width, "num_frames": num_frames, "fps": 24, + "num_inference_steps": steps, "guidance_scale": 1.0, "batch_cfg": False}, + "output": {"output_path": str(out_dir / f"{i:02d}_{pid}.mp4"), "save_video": True, + "return_frames": False}} + torch.cuda.synchronize() + t = time.perf_counter() + result = generator.generate(request) + torch.cuda.synchronize() + wall = time.perf_counter() - t + results["runs"].append({"prompt": pid, "warmup": i < warmups, "wall_s": round(wall, 2), + "generation_time_s": getattr(result, "generation_time", None), + "video": getattr(result, "video_path", None)}) + results["peak_mem_gb_device"] = _sh("nvidia-smi --query-gpu=memory.used --format=csv,noheader").strip() + finally: + generator.shutdown() + (out_dir / "results.json").write_text(json.dumps(results, indent=1)) + volume.commit() + return results + + +bench_image = image.add_local_file(pathlib.Path(__file__).parent / "bench_code.py", "/root/bench_code.py") + +SHAPES = [ + # name, prefix segments (text, audio rows), video latent tokens (t, h, w) after 1x2x2 patching + ("480p_124f", [256, 414], [37, 15, 26]), + ("768p_243f", [256, 810], [72, 24, 42]), +] + + +@app.function(gpu="RTX-PRO-6000", memory=65536, cpu=8, timeout=3600, volumes={"/vol": volume}, image=bench_image) +def bench_block() -> dict: + _install_kernel() + import sys + sys.path.insert(0, "/root") + import bench_code + report = {"check": bench_code.check_tile64()} + report.update(bench_code.run(SHAPES)) + return json.loads(json.dumps(report, default=str)) + + +@app.function(gpu="RTX-PRO-6000", memory=65536, cpu=8, timeout=1800, volumes={"/vol": volume}, image=bench_image) +def density_fn() -> dict: + _install_kernel() + import sys + sys.path.insert(0, "/root") + import bench_code + return json.loads(json.dumps(bench_code.density_study(), default=str)) + + +@app.function(gpu="RTX-PRO-6000", memory=65536, cpu=8, timeout=1800, volumes={"/vol": volume}, image=bench_image) +def kcheck_fn() -> dict: + _install_kernel() + import sys + sys.path.insert(0, "/root") + import bench_code + out = {"check": bench_code.check_tile64()} + out["density"] = bench_code.density_study() + return json.loads(json.dumps(out, default=str)) + + +@app.function(gpu="RTX-PRO-6000:2", memory=229376, cpu=16, timeout=5400, volumes={"/vol": volume}) +def run_variant2(*args, **kwargs) -> dict: + return run_variant.local(*args, **kwargs) + + +@app.function(gpu="RTX-PRO-6000:8", memory=196608, cpu=16, timeout=5400, volumes={"/vol": volume}) +def run_variant8(*args, **kwargs) -> dict: + return run_variant.local(*args, **kwargs) + + +@app.function(gpu="RTX-PRO-6000:4", memory=131072, cpu=8, timeout=5400, volumes={"/vol": volume}) +def run_variant4(*args, **kwargs) -> dict: + return run_variant.local(*args, **kwargs) + + +@app.function(gpu="RTX-PRO-6000:8", memory=196608, cpu=16, timeout=2 * 3600, volumes={"/vol": volume}) +def bench8_pair(prompt_texts: dict, timed: tuple, warmups: int, models: tuple = ("v2", "v2_profile")) -> list: + """8-GPU 10 s 768p V2 8-step benchmark, then a stage-profiled pass, in one container (shared compile caches).""" + specs = {"v2": ("sp8_v2_8step_768p10s_1k", "v2_vsa_light", 0.8, 9, {}, timed, warmups), + # One warm + one profiled generation; CUDA-event spans log per DiT forward as H3_STAGE_MS. + "v2_profile": ("sp8_v2_8step_profile", "v2_vsa_light", 0.8, 9, {"FASTVIDEO_H3_SP_PROFILE": "1"}, + timed[:1], 1)} + out = [] + for m in models: + name, model, sparsity, steps, extra_env, prompts, warm = specs[m] + r = run_variant.local(name, model, "h3_dit_vsa", attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", + vae_compile=True, height=768, width=1344, num_frames=243, warmups=warm, + env={**FAST_ENV, **extra_env}, prompts=prompts, sparsity=sparsity, steps=steps, + num_gpus=8, parallel_decode=True, prompt_texts=prompt_texts) + print("RESULT", json.dumps(r), flush=True) + out.append(r) + return out + + +@app.function(gpu="RTX-PRO-6000", memory=196608, cpu=16, timeout=3 * 3600, volumes={"/vol": volume}) +def memladder(prompt_texts: dict, pid: str, configs: dict | None = None) -> list: + """V2 8-step 10 s 768p on one GPU under memory placements; stage logs carry per-stage peaks. + + ``configs`` maps a name to (offload overrides, experimental overrides, extra env). The env can set + FASTVIDEO_CUDA_MEMORY_CAP_GIB to emulate a smaller card. + """ + seq = {"h3_sequential_load": True} + lw = {"text_encoder": True, "vae": True, "dit_layerwise": True} + configs = configs or { + "A_resident": ({}, {}, {}), + "B_seq_encoder_vae_offload": ({"text_encoder": True, "vae": True}, seq, {}), + "C_plus_dit_layerwise": (lw, seq, {}), + } + out = [] + for name, (offload, extra, env_extra) in configs.items(): + try: + r = run_variant.local(f"mem_{name}", "v2_vsa_light", "h3_dit_vsa", attention="VIDEO_SPARSE_ATTN_H3", + decode="h3-vae", vae_compile=False, height=768, width=1344, num_frames=243, + warmups=1, env={**FAST_ENV, **env_extra}, prompts=(pid,), sparsity=0.8, steps=9, + num_gpus=1, prompt_texts=prompt_texts, offload=offload, experimental_extra=extra) + except Exception as e: # keep the ladder going; record the failure + r = {"name": f"mem_{name}", "error": repr(e)[:3000]} + print("RESULT", json.dumps(r), flush=True) + out.append(r) + for k in env_extra: + os.environ.pop(k, None) + return out + + +@app.function(cpu=8, memory=32768, timeout=1800, volumes={"/vol": volume}) +def compare_videos(a: str, b: str) -> dict: + """Frame PSNR between two MP4s on the volume (same seed and prompt).""" + import imageio.v3 as iio + import numpy as np + fa = iio.imread(a, plugin="pyav").astype(np.float32) + fb = iio.imread(b, plugin="pyav").astype(np.float32) + n = min(len(fa), len(fb)) + mse = ((fa[:n] - fb[:n]) ** 2).reshape(n, -1).mean(axis=1) + psnr = 10 * np.log10(255.0**2 / np.maximum(mse, 1e-9)) + return {"frames": [len(fa), len(fb)], "psnr_mean": float(psnr.mean()), "psnr_min": float(psnr.min())} + + +FAST_ENV = {"FASTVIDEO_H3_VSA_FP4": "1", "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", "FASTVIDEO_NVFP4_MM_BACKEND": "cutlass", + "FASTVIDEO_H3_VAE_TILE_BATCH": "28"} + + +@app.local_entrypoint() +def main(step: str = "all", ladder: str = "base"): + if step == "prep_personal": + print("KERNEL", build_kernel.remote()) + print("CONVERT", json.dumps(convert.remote(minimal=True), indent=1)[:6000]) + return + if step == "sp2check": + common = dict(attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", vae_compile=False, height=480, width=832, + num_frames=124, warmups=0, env=FAST_ENV, prompts=("kitesurf",), sparsity=0.8, steps=9) + one = run_variant.spawn("sp1_480p", "v2_vsa_light", "h3_dit_vsa", **common) + two = run_variant2.spawn("sp2_480p", "v2_vsa_light", "h3_dit_vsa", num_gpus=2, **common) + r1, r2 = one.get(), two.get() + print("RESULT", json.dumps(r1)) + print("RESULT", json.dumps(r2)) + print("COMPARE", json.dumps(compare_videos.remote(r1["runs"][0]["video"], r2["runs"][0]["video"]))) + return + if step == "sp8": + r = run_variant8.remote( + "sp8_v2_8step", "v2_vsa_light", "h3_dit_vsa", attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", + vae_compile=True, height=768, width=1344, num_frames=243, warmups=1, env=FAST_ENV, + prompts=("kitesurf", "chef", "kitesurf", "chef"), sparsity=0.8, steps=9, num_gpus=8, parallel_decode=True, + pre_runs=((480, 832, 124, "kitesurf"), )) + print("RESULT", json.dumps(r)) + if r.get("pre_runs"): + print("COMPARE", json.dumps(compare_videos.remote("/vol/outputs/sp1_480p/00_kitesurf.mp4", + r["pre_runs"][0]["video"]))) + return + if step == "simfp8": + # SP=1 with and without the SP exchange's FP8 rounding, same settings as sp1_480p. + common = dict(attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", vae_compile=False, height=480, width=832, + num_frames=124, warmups=0, prompts=("kitesurf",), sparsity=0.8, steps=9) + plain = run_variant.spawn("sp1_480p_rerun", "v2_vsa_light", "h3_dit_vsa", env=FAST_ENV, **common) + sim = run_variant.spawn("sp1_480p_simfp8", "v2_vsa_light", "h3_dit_vsa", + env={**FAST_ENV, "FASTVIDEO_H3_SIM_SP_FP8": "1"}, **common) + r_plain, r_sim = plain.get(), sim.get() + print("RESULT", json.dumps(r_plain)) + print("RESULT", json.dumps(r_sim)) + return + if step == "memreport": + r = run_variant.remote("memreport", "v2_vsa_light", "h3_dit_vsa", attention="VIDEO_SPARSE_ATTN_H3", + decode="h3-vae", vae_compile=False, height=480, width=832, num_frames=124, warmups=0, + env={**FAST_ENV, "FASTVIDEO_MEMORY_REPORT": "1"}, prompts=("kitesurf",), sparsity=0.8, + steps=9, offload={"text_encoder": True, "vae": True}, + experimental_extra={"h3_sequential_load": False}) + print("RESULT", json.dumps(r)) + return + if step == "memladder": + rows = [json.loads(line) for line in open(WORKTREE.parent / "UniServe-sm120fp4" / "uniserve_eval" / "workloads" + / "fast_h3" / "latency.jsonl")] + pid = "latency-ceramics-005" + text = {r["id"]: r["prompt"] for r in rows if r["id"] == pid} + configs = None + if ladder == "caps": + seq = {"h3_sequential_load": True} + lw = {"text_encoder": True, "vae": True, "dit_layerwise": True} + configs = { + "D32_C_cap32": (lw, seq, {"FASTVIDEO_CUDA_MEMORY_CAP_GIB": "32"}), + "D24_buffers_cap24": (lw, seq, {"FASTVIDEO_CUDA_MEMORY_CAP_GIB": "24", + "FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS": "1"}), + "D16_buffers_tile8_cap16": (lw, seq, {"FASTVIDEO_CUDA_MEMORY_CAP_GIB": "16", + "FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS": "1", + "FASTVIDEO_H3_VAE_TILE_BATCH": "8"}), + } + for r in memladder.remote(text, pid, configs): + print("RESULT", json.dumps(r)) + return + if step == "bench8_pair": + # UniServe's 10 s / ~1K-token latency prompts, so the numbers line up with its published protocol. + rows = [json.loads(line) for line in open(WORKTREE.parent / "UniServe-sm120fp4" / "uniserve_eval" / "workloads" + / "fast_h3" / "latency.jsonl")] + ten = {r["id"]: r["prompt"] for r in rows if r["seconds"] == 10 and r["prompt_len"] == 1000} + ids = ("latency-ceramics-005", "latency-harbor-005") + texts = {i: ten[i] for i in ids} + for r in bench8_pair.remote(texts, (ids[0], ids[1], ids[0], ids[1]), 2): + timed = sorted(x["wall_s"] for x in r["runs"] if not x["warmup"]) + print("SUMMARY", r["name"], "timed", timed, "median", timed[len(timed) // 2] if len(timed) % 2 + else (timed[len(timed) // 2 - 1] + timed[len(timed) // 2]) / 2) + return + if step == "sp4": + r = run_variant4.remote( + "sp4_v2_8step", "v2_vsa_light", "h3_dit_vsa", attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", + vae_compile=True, height=768, width=1344, num_frames=243, warmups=1, env=FAST_ENV, + prompts=("kitesurf", "chef", "kitesurf", "chef"), sparsity=0.8, steps=9, num_gpus=4, parallel_decode=True, + pre_runs=((480, 832, 124, "kitesurf"), )) + print("RESULT", json.dumps(r)) + if r.get("pre_runs"): + print("COMPARE", json.dumps(compare_videos.remote("/vol/outputs/sp1_480p/00_kitesurf.mp4", + r["pre_runs"][0]["video"]))) + return + if step in ("bench8_v2", "bench8_v4"): + v2 = step == "bench8_v2" + r = run_variant8.remote( + f"sp8_{'v2_8step' if v2 else 'v4_4step'}_768p10s", "v2_vsa_light" if v2 else "v4_vsa_light", "h3_dit_vsa", + attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", vae_compile=True, height=768, width=1344, + num_frames=243, warmups=1, env=FAST_ENV, prompts=("kitesurf", "chef", "kitesurf", "chef"), + sparsity=0.8 if v2 else 0.9, steps=9 if v2 else 5, num_gpus=8, parallel_decode=True) + print("RESULT", json.dumps(r)) + return + if step == "build_bench": + print("KERNEL", build_kernel.remote()) + step = "bench" + if step in ("prep768", "e2e768"): + if step == "prep768": + print("KERNEL", build_kernel.remote()) + print("CONVERT", json.dumps(convert.remote(), indent=1)[:6000]) + fast = {"FASTVIDEO_H3_VSA_FP4": "1", "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", + "FASTVIDEO_NVFP4_MM_BACKEND": "cutlass", "FASTVIDEO_H3_VAE_TILE_BATCH": "28"} + common = dict(attention="VIDEO_SPARSE_ATTN_H3", decode="h3-vae", vae_compile=True, height=768, width=1344, + num_frames=243, warmups=1, env=fast, prompts=("kitesurf", "chef")) + calls = { + "v4_4step_vsa90_768p10s": run_variant.spawn("v4_4step_vsa90_768p10s", "v4_vsa_light", "h3_dit_vsa", + sparsity=0.9, steps=5, **common), + "v2_8step_vsa80_768p10s": run_variant.spawn("v2_8step_vsa80_768p10s", "v2_vsa_light", "h3_dit_vsa", + sparsity=0.8, steps=9, **common), + } + for name, call in calls.items(): + try: + print("RESULT", json.dumps(call.get())) + except Exception as exc: # noqa: BLE001 + print("FAILED", name, repr(exc)[:3000]) + return + if step == "kcheck": + print("KERNEL", build_kernel.remote()) + print("KCHECK", json.dumps(kcheck_fn.remote(), indent=1)) + return + if step == "density": + print("DENSITYRESULT", json.dumps(density_fn.remote(), indent=1)) + return + if step == "bench": + print("BENCHRESULT", json.dumps(bench_block.remote(), indent=1)) + return + if step in ("all", "build"): + print("KERNEL", build_kernel.remote()) + if step in ("all", "convert"): + print("CONVERT", json.dumps(convert.remote(), indent=1)[:6000]) + if step in ("all", "run"): + variants = [ + ("full_qatinfer_light", "v2_full_light", "h3_dit", "ATTN_QAT_INFER", "h3-vae", True), + ("full_vsa_light", "v2_full_light", "h3_dit", "VIDEO_SPARSE_ATTN_H3", "h3-vae", True), + ("full_qatinfer_int8", "v2_full_int8", "h3_dit", "ATTN_QAT_INFER", "h3-vae", True), + ("full_qatinfer_taeh3", "v2_full_light", "h3_dit", "ATTN_QAT_INFER", "taeh3", False), + ("ffn_vsa_int8", "v2_ffn_int8", "h3_dit_ffn", "VIDEO_SPARSE_ATTN_H3", "h3-vae", True), + ] + # One RTX PRO 6000 per variant, all at once; failures are returned, not raised. + for v, result in zip(variants, run_variant.starmap(variants, return_exceptions=True)): + if isinstance(result, Exception): + print("FAILED", v[0], repr(result)[:3000]) + else: + print("RESULT", json.dumps(result)) + + +# --- Headline benchmark (480p 5 s, 768p 10 s) on 1 / 4 / 8 GPUs from a FastVideo HF repo. --- +# MODAL_PROFILE=aryan5v modal run --detach app.py::headline --repo FastVideo/ --profile h3_dit_ffn --gpus 1,4,8 +HERE = pathlib.Path(__file__).resolve().parent +headline_image = (image.add_local_file(HERE / "bench_headline.py", "/root/bench_headline.py") + .add_local_file(HERE / "headline_prompts.json", "/root/headline_prompts.json") + .add_local_file(HERE / "showcase_prompts.json", "/root/showcase_prompts.json")) +SECRETS = [modal.Secret.from_name("hf-fastvideo")] + + +@app.function(cpu=8, memory=32768, timeout=3600, volumes={"/vol": volume}, secrets=SECRETS, image=headline_image) +def headline_fetch(repo: str) -> str: + from huggingface_hub import snapshot_download + local = f"/vol/models/{repo.split('/')[-1]}" + snapshot_download(repo, local_dir=local, token=os.environ["HF_TOKEN"], max_workers=16) + volume.commit() + return _sh(f"du -sh {shlex.quote(local)}/*") + + +def _headline(repo: str, gpus: int, profile: str, extra_env: dict | None, tag: str = "") -> dict: + _install_kernel() + model = f"/vol/models/{repo.split('/')[-1]}" + # The run name is a single path component locally and on the volume. + run_name = f"pro6000x{gpus}-{repo.split('/')[-1]}{tag}".replace("/", "-") + env = {**os.environ, **FAST_ENV, **(extra_env or {}), "HEADLINE_OUT": "/vol/outputs/headline", + "HEADLINE_DEVICE": f"{gpus}x RTX PRO 6000", "PYTHONPATH": "/src/fastvideo"} + proc = subprocess.run(["python", "/root/bench_headline.py", run_name, model, str(gpus), profile, + "--prompts", "/root/headline_prompts.json"], env=env, capture_output=True, text=True, + cwd="/root") + log_text = proc.stdout + proc.stderr + run_dir = pathlib.Path("/vol/outputs/headline") / run_name + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "run.log").write_text(log_text) + volume.commit() + errors = [line for line in log_text.splitlines() + if any(k in line for k in ("Error", "error:", "Traceback", "Killed", "OOM", "out of memory"))][-40:] + result_path = run_dir / "results.json" + results = json.loads(result_path.read_text()) if result_path.exists() else {} + return {"run": run_name, "returncode": proc.returncode, "results": results, "errors": errors, + "log_tail": log_text[-1500:]} + + +@app.function(image=headline_image, gpu="RTX-PRO-6000", memory=131072, cpu=8, timeout=2 * 3600, volumes={"/vol": volume}) +def headline1(repo: str, profile: str, extra_env: dict | None = None, tag: str = "") -> dict: + return _headline(repo, 1, profile, extra_env, tag) + + +@app.function(image=headline_image, gpu="RTX-PRO-6000:4", memory=196608, cpu=16, timeout=2 * 3600, volumes={"/vol": volume}) +def headline4(repo: str, profile: str, extra_env: dict | None = None, tag: str = "") -> dict: + return _headline(repo, 4, profile, extra_env, tag) + + +@app.function(image=headline_image, gpu="RTX-PRO-6000:8", memory=262144, cpu=32, timeout=2 * 3600, volumes={"/vol": volume}) +def headline8(repo: str, profile: str, extra_env: dict | None = None, tag: str = "") -> dict: + return _headline(repo, 8, profile, extra_env, tag) + + +@app.local_entrypoint() +def headline(repo: str, profile: str = "h3_dit_ffn", gpus: str = "1,4,8", skip_fetch: bool = False, + extra_env: str = "{}", tag: str = ""): + if not skip_fetch: + print(headline_fetch.remote(repo)) + fns = {"1": headline1, "4": headline4, "8": headline8} + calls = [fns[g].spawn(repo, profile, json.loads(extra_env), tag) for g in gpus.split(",")] + out = HERE / "headline_results" + out.mkdir(exist_ok=True) + for call in calls: + res = call.get() + print(json.dumps({k: v for k, v in res.items() if k != "log_tail"}, indent=1)[:3000]) + if res["returncode"]: + print(res["log_tail"]) + (out / f"{res['run']}.json").write_text(json.dumps(res, indent=1)) diff --git a/scripts/benchmarks/minimax_h3_pro6000/bench_code.py b/scripts/benchmarks/minimax_h3_pro6000/bench_code.py new file mode 100644 index 000000000..c45557b83 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/bench_code.py @@ -0,0 +1,338 @@ +"""Per-component timing of one MiniMax-H3 block on the current GPU. Imported inside the Modal container.""" +import json +import math +import time + +import torch +import torch.nn.functional as F + +HID, HEADS, HD, FFN = 5376, 56, 128, 14336 + + +def timed(fn, iters=5, warmup=2): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return round(start.elapsed_time(end) / iters, 3) + + +def smooth_tokens(prefix_len, shape, dev, gen, noise=0.5): + """Token features with spatial-temporal correlation (proxy for real q/k), packed [prefix | video].""" + t, h, w = shape + low = torch.randn(1, HEADS * HD, max(2, t // 6), max(2, h // 6), max(2, w // 6), device=dev, generator=gen) + vid = F.interpolate(low, size=(t, h, w), mode="trilinear", align_corners=False) + vid = vid.reshape(HEADS, HD, -1).permute(2, 0, 1) + vid = vid + noise * torch.randn(vid.shape, device=dev, generator=gen) + pre = torch.randn(prefix_len, HEADS, HD, device=dev, generator=gen) + return torch.cat([pre, vid]).to(torch.bfloat16)[None] # [1, L, H, D] + + +def run(shapes): + from flashinfer import SfLayout, mm_fp4, nvfp4_quantize + from fastvideo.attention.backends.video_sparse_attn_h3 import (_build_block_mask, _h3_tile_geometry, _pool_tiles) + from fastvideo_kernel.block_sparse_attn import block_sparse_attn as bsa64 + from fastvideo.models.dits.minimax_h3_fusions import (fused_qknorm_rope, fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, minimax_h3_swiglu) + import attn_qat_infer.api as fa + from attn_qat_infer.api import BLOCK_M, sageattn_blackwell, sageattn_blackwell_sparse + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(0) + report = {"device": torch.cuda.get_device_name(0), "torch": torch.__version__} + unit = torch.tensor(1.0, device=dev) + + def fp4_weight(n, k): + w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.02 + gsf = (448 * 6) / w.float().abs().max() + wq, ws = nvfp4_quantize(w, gsf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + return w, wq, ws, (1.0 / gsf).float() + + for name, prefix_segments, video_shape in shapes: + prefix_segments = tuple(prefix_segments) + video_shape = tuple(video_shape) + L = sum(prefix_segments) + math.prod(video_shape) + r = {"tokens": L} + torch.cuda.empty_cache() + + # ---------------- linears ---------------- + lin = {} + for lname, n, k in (("qkv_one", 7168, HID), ("out", HID, 7168), ("fc_in", 2 * FFN, HID), + ("fc_out", HID, FFN)): + w, wq, ws, alpha = fp4_weight(n, k) + x = torch.randn(L, k, device=dev, dtype=torch.bfloat16) + xq, xs = nvfp4_quantize(x, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + entry = {"quant_ms": timed(lambda: nvfp4_quantize(x, unit, sfLayout=SfLayout.layout_128x4, + do_shuffle=False))} + for backend in ("cudnn", "cutlass", "auto"): + try: + entry[f"mm_{backend}_ms"] = timed( + lambda: mm_fp4(xq, wq.T, xs, ws.T, alpha, torch.bfloat16, None, backend=backend)) + except Exception as exc: # noqa: BLE001 + entry[f"mm_{backend}_ms"] = f"ERR {type(exc).__name__}: {str(exc)[:120]}" + entry["bf16_ms"] = timed(lambda: F.linear(x, w)) + entry["tflops_fp4_auto"] = (round(2 * L * n * k / entry["mm_auto_ms"] / 1e9, 1) + if isinstance(entry["mm_auto_ms"], float) else None) + lin[lname] = entry + del w, wq, ws, x, xq, xs + r["linear"] = lin + + # ---------------- attention ---------------- + geom = _h3_tile_geometry(prefix_segments, video_shape, dev, (4, 4, 4)) + _, vbs, untile, n_prefix, n_video = geom + n_tiles = vbs.numel() + Lpad = n_tiles * 64 + att = {"tiles64": n_tiles, "prefix_tiles": n_prefix} + prefix_len = sum(prefix_segments) + q = smooth_tokens(prefix_len, video_shape, dev, gen) + k = (q + 0.5 * torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)).to(torch.bfloat16) + v = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16) + g = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16) * 0.1 + + def tile(x): + buf = torch.zeros((1, Lpad, HEADS, HD), device=dev, dtype=x.dtype) + buf[:, untile] = x + return buf + + qt, kt, vt, gt = tile(q), tile(k), tile(v), tile(g) + att["tile_scatter4_gather1_ms"] = timed(lambda: (tile(q), tile(k), tile(v), tile(g), qt[:, untile])) + + def vsa_full(sparsity=0.8, with_gate=True): + qp = _pool_tiles(qt, vbs, 64) + kp = _pool_tiles(kt, vbs, 64) + scores = torch.matmul(qp, kp.transpose(-2, -1)) / HD**0.5 + mask = _build_block_mask(scores, n_prefix, sparsity, True, ((n_prefix, n_prefix + n_video), ), (sparsity, )) + out, _ = bsa64(qt.transpose(1, 2).contiguous(), kt.transpose(1, 2).contiguous(), + vt.transpose(1, 2).contiguous(), mask, vbs) + out = out.transpose(1, 2).contiguous() + if with_gate: + vp = _pool_tiles(vt, vbs, 64) + oc = torch.matmul(torch.softmax(scores, dim=-1), vp).permute(0, 2, 1, 3).to(out.dtype) + out = (out.view(1, n_tiles, 64, HEADS, HD) + oc.unsqueeze(2) * gt.view(1, n_tiles, 64, HEADS, HD)).view( + 1, Lpad, HEADS, HD) + return out, mask + + _, mask = vsa_full() + att["tile_density"] = round(mask.float().mean().item(), 4) + att["vsa_triton_total_ms"] = timed(vsa_full) + qb, kb, vb = (x.transpose(1, 2).contiguous() for x in (qt, kt, vt)) + att["vsa_triton_kernel_only_ms"] = timed(lambda: bsa64(qb, kb, vb, mask, vbs)) + + # dense FP4 on the unpadded packed sequence + qd, kd, vd = (x.transpose(1, 2).contiguous() for x in (q, k, v)) + att["sage3_dense_fp4_ms"] = timed(lambda: sageattn_blackwell(qd, kd, vd), iters=3) + # same with delta_s cached + per_block_mean=False (no 9.5 GB memset) + zero_ds = {} + + def sage_nods(qx, kx, vx): + QL, KL = qx.size(2), kx.size(2) + qx, kx, vx = (F.pad(x, (0, 0, 0, (BLOCK_M - x.size(2) % BLOCK_M) % BLOCK_M)).contiguous() for x in (qx, kx, vx)) + key = (qx.shape[0], qx.shape[1], kx.shape[2]) + if key not in zero_ds: + zero_ds[key] = torch.zeros((qx.shape[0], qx.shape[1], 1, kx.shape[2]), device=dev, dtype=torch.float32) + ql, kl, vl = fa.scale_and_quant_fp4(qx), fa.scale_and_quant_fp4_permute(kx), fa.scale_and_quant_fp4_transpose(vx) + return fa.blockscaled_fp4_attn(ql, kl, vl, zero_ds[key], KL, False, False, True, True, None)[0][:, :, :QL] + + att["sage3_dense_fp4_no_deltas_ms"] = timed(lambda: sage_nods(qd, kd, vd), iters=3) + out_a = sageattn_blackwell(qd, kd, vd) + out_b = sage_nods(qd, kd, vd) + att["no_deltas_max_abs_diff"] = (out_a.float() - out_b.float()).abs().max().item() + Lq = qd.shape[2] + att["dense_fp4_tflops"] = round(4 * Lq * Lq * HD * HEADS / att["sage3_dense_fp4_no_deltas_ms"] / 1e9, 1) + if L < 30000: + att["sdpa_bf16_dense_ms"] = timed(lambda: F.scaled_dot_product_attention(qd, kd, vd), iters=3) + + # sparse FP4 on 128x128 blocks with 64x64 quadrant masks (exact VSA tile-64 semantics) + nt2 = n_tiles + (n_tiles % 2) + qs, ks, vs = (F.pad(x.transpose(1, 2), (0, 0, 0, nt2 * 64 - Lpad)).contiguous() for x in (qt, kt, vt)) + q2k_idx, q2k_num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs, validate=True) + att["block128_density"] = round(q2k_num.float().mean().item() / (nt2 // 2), 4) + att["sparse_fp4_quad_ms"] = timed(lambda: sageattn_blackwell_sparse(qs, ks, vs, q2k_idx, q2k_num, kvv, quad)) + att["mask_to_blocks_ms"] = timed(lambda: fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs)) + o_fp4 = sageattn_blackwell_sparse(qs, ks, vs, q2k_idx, q2k_num, kvv, quad)[:, :, :Lpad] + o_tri, _ = bsa64(qb, kb, vb, mask, vbs) + rows = untile # valid (non-pad) rows of the tile buffer + diff = (o_fp4[:, :, rows].float() - o_tri[:, :, rows].float()) + att["quad_vs_triton_rel_l2"] = round((diff.norm() / o_tri[:, :, rows].float().norm()).item(), 4) + od = sageattn_blackwell(qd, kd, vd) + from torch.nn.attention import SDPBackend, sdpa_kernel + if L < 30000: + ref_dense = F.scaled_dot_product_attention(qd, kd, vd) + att["dense_fp4_vs_bf16_rel_l2"] = round(((od.float() - ref_dense.float()).norm() / + ref_dense.float().norm()).item(), 4) + att["quad_finite"] = bool(torch.isfinite(o_fp4[:, :, rows]).all()) + r["attention"] = att + del q, k, v, g, qt, kt, vt, gt, qd, kd, vd, qs, ks, vs, qb, kb, vb, out_a, out_b, zero_ds, o_fp4, o_tri, od + + # ---------------- elementwise ---------------- + ew = {} + x = torch.randn(1, L, HID, device=dev, dtype=torch.bfloat16) + br = torch.randn_like(x) + nw = torch.randn(HID, device=dev, dtype=torch.bfloat16) + tab = [torch.randn(3, HID, device=dev, dtype=torch.bfloat16) * 0.1 for _ in range(6)] + idx = torch.randint(0, 3, (L,), device=dev) + norm = torch.nn.RMSNorm(HID, eps=1e-6, device=dev, dtype=torch.bfloat16) + + def eager_mod(): + n1 = norm(x) * (1.0 + tab[1].index_select(0, idx)) + tab[0].index_select(0, idx) + h = x + tab[2].index_select(0, idx) * br + n2 = norm(h) * (1.0 + tab[4].index_select(0, idx)) + tab[3].index_select(0, idx) + out = h + tab[5].index_select(0, idx) * br + return n1, n2, out + + def fused_mod(): + n1 = fused_rmsnorm_modulate(x, nw, tab[1], tab[0], idx, 1e-6) + h, n2 = fused_residual_gate_rmsnorm_modulate(x, br, tab[2], nw, tab[4], tab[3], idx, 1e-6) + out = h + tab[5].index_select(0, idx) * br + return n1, n2, out + + with torch.no_grad(): + ew["modulate_eager_ms"] = timed(eager_mod) + ew["modulate_fused_ms"] = timed(fused_mod) + ew["modulate_compiled_ms"] = timed(torch.compile(eager_mod)) + packed = torch.randn(1, L, 2 * FFN, device=dev, dtype=torch.bfloat16) + + def eager_swiglu(): + a, b = packed.chunk(2, dim=-1) + return a * F.silu(b) + + ew["swiglu_eager_ms"] = timed(eager_swiglu) + ew["swiglu_fused_ms"] = timed(lambda: minimax_h3_swiglu(packed)) + del packed + qq = torch.randn(1, L, HEADS, HD, device=dev, dtype=torch.bfloat16) + cos = torch.randn(L, 96, device=dev, dtype=torch.bfloat16) + sin = torch.randn(L, 96, device=dev, dtype=torch.bfloat16) + qn = torch.nn.RMSNorm(HD, eps=1e-6, device=dev, dtype=torch.bfloat16) + + def eager_rope(): + outs = [] + for t in (qq, qq): + t = qn(t) + rot, pas = t[..., :96], t[..., 96:] + c, s = cos[None, :, None, :], sin[None, :, None, :] + a, b = rot.chunk(2, dim=-1) + outs.append(torch.cat((rot * c + torch.cat((-b, a), -1) * s, pas), -1).contiguous()) + return outs + + ew["qknorm_rope_eager_ms"] = timed(eager_rope) + ew["qknorm_rope_fused_ms"] = timed( + lambda: (fused_qknorm_rope(qq, qn.weight, cos, sin, 1e-6), fused_qknorm_rope(qq, qn.weight, cos, sin, 1e-6))) + r["elementwise"] = ew + report[name] = r + print("BENCH", name, json.dumps(r), flush=True) + return report + + +def check_tile64(): + """Exact semantics check of the quadrant path on small H3-style layouts.""" + from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _h3_tile_geometry, _pool_tiles + import attn_qat_infer.api as fa + from attn_qat_infer.api import sageattn_blackwell, sageattn_blackwell_sparse + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(1) + out = {} + cases = [("odd_tiles", (77, 46), (9, 6, 10), 64), ("even_tiles", (64, 40), (8, 8, 8), 64), + ("tile256", (100, 70), (8, 8, 16), 256)] + for name, prefix, vshape, tt in cases: + H = 4 + shape = {64: (4, 4, 4), 256: (4, 8, 8)}[tt] + _, vbs, untile, n_prefix, n_video = _h3_tile_geometry(prefix, vshape, dev, shape) + n_tiles = vbs.numel() + Lpad = n_tiles * tt + L = sum(prefix) + math.prod(vshape) + q = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16) + k = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16) + v = torch.randn(1, L, H, HD, device=dev, generator=gen, dtype=torch.bfloat16) + + def tile(x): + buf = torch.zeros((1, Lpad, H, HD), device=dev, dtype=x.dtype) + buf[:, untile] = x + return buf.transpose(1, 2).contiguous() # BHSD + + qt, kt, vt = tile(q), tile(k), tile(v) + scores = torch.matmul(_pool_tiles(qt.transpose(1, 2), vbs, tt), _pool_tiles(kt.transpose(1, 2), vbs, tt).transpose(-2, -1)) + mask = _build_block_mask(scores, n_prefix, 0.8, True, ((n_prefix, n_prefix + n_video), ), (0.8, )) + q2k_idx, q2k_num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, tt, vbs, validate=True) + Lk = q2k_idx.shape[2] * 128 + pad = lambda x: F.pad(x, (0, 0, 0, Lk - Lpad)).contiguous() + o = sageattn_blackwell_sparse(pad(qt), pad(kt), pad(vt), q2k_idx, q2k_num, kvv, quad)[:, :, :Lpad] + tok_tile = torch.arange(n_tiles, device=dev).repeat_interleave(tt) + tok_valid = torch.zeros(Lpad, dtype=torch.bool, device=dev) + tok_valid[untile] = True + tm = mask[:, :, tok_tile][:, :, :, tok_tile] & tok_valid[None, None, None, :] + ref = F.scaled_dot_product_attention(qt.float(), kt.float(), vt.float(), attn_mask=tm) + rows = untile + err = (o.float() - ref)[:, :, rows] + # FP4 noise floor: dense FP4 vs dense fp32 on the same valid tokens + qv, kv_, vv = (x[:, :, rows].contiguous() for x in (qt, kt, vt)) + od = sageattn_blackwell(qv, kv_, vv) + rd = F.scaled_dot_product_attention(qv.float(), kv_.float(), vv.float()) + out[name] = { + "tiles": n_tiles, "density": round(mask.float().mean().item(), 3), + "rel_l2": round((err.norm() / ref[:, :, rows].norm()).item(), 4), + "dense_fp4_rel_l2_floor": round(((od.float() - rd).norm() / rd.norm()).item(), 4), + "cosine": round(F.cosine_similarity(o[:, :, rows].float().flatten(), ref[:, :, rows].flatten(), dim=0).item(), 5), + "finite": bool(torch.isfinite(o[:, :, rows]).all()), + } + if tt == 64: + from fastvideo_kernel.block_sparse_attn import block_sparse_attn as bsa64 + ot, _ = bsa64(qt, kt, vt, mask, vbs) + out[name]["triton_bf16_rel_l2"] = round(((ot.float() - ref)[:, :, rows].norm() / ref[:, :, rows].norm()).item(), 4) + print("CHECK", name, json.dumps(out[name]), flush=True) + return out + + +def density_study(prefix_segments=(256, 810), video_shape=(72, 24, 42)): + """Block densities each kernel granularity would compute, and FP4 sparse timing at 0.8 / 0.9.""" + from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _h3_tile_geometry, _pool_tiles + import attn_qat_infer.api as fa + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(0) + _, vbs, untile, n_prefix, n_video = _h3_tile_geometry(tuple(prefix_segments), tuple(video_shape), dev, (4, 4, 4)) + n_tiles = vbs.numel() + Lpad = n_tiles * 64 + q = smooth_tokens(sum(prefix_segments), tuple(video_shape), dev, gen) + k = (q + 0.5 * torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16)).to(torch.bfloat16) + v = torch.randn(q.shape, device=dev, generator=gen, dtype=torch.bfloat16) + + def tile(x): + buf = torch.zeros((1, Lpad, HEADS, HD), device=dev, dtype=x.dtype) + buf[:, untile] = x + return buf + + qt, kt, vt = tile(q), tile(k), tile(v) + scores = torch.matmul(_pool_tiles(qt, vbs, 64), _pool_tiles(kt, vbs, 64).transpose(-2, -1)) + nt2 = n_tiles + n_tiles % 2 + rows = nt2 * 64 + qs, ks, vs = (F.pad(x, (0, 0, 0, 0, 0, rows - Lpad)).contiguous() for x in (qt, kt, vt)) + out = {} + for sparsity in (0.8, 0.9): + mask = _build_block_mask(scores, n_prefix, sparsity, True, ((n_prefix, n_prefix + n_video), ), (sparsity, )) + m = F.pad(mask, (0, nt2 - n_tiles, 0, nt2 - n_tiles), value=False) + B, H = m.shape[:2] + d = {"tile64x64": m.float().mean().item()} + q2 = m.view(B, H, nt2 // 2, 2, nt2) # query pairs + d["q128_k64"] = q2.any(3).float().mean().item() + k2 = m.view(B, H, nt2, nt2 // 2, 2) + d["q64_k128"] = k2.any(4).float().mean().item() + d["q128_k128"] = m.view(B, H, nt2 // 2, 2, nt2 // 2, 2).any(5).any(3).float().mean().item() + idx, num, kvv, quad = fa.vsa_tile_mask_to_fp4_blocks(mask, 64, vbs) + def bshd(): + qh, kh, vh = (x.transpose(1, 2) for x in (qs, ks, vs)) + ds = fa._zero_delta_s(1, HEADS, rows, dev) + return fa.blockscaled_fp4_attn_sparse(fa.scale_and_quant_fp4(qh), fa.scale_and_quant_fp4_permute(kh), + fa.scale_and_quant_fp4_transpose(vh), ds, rows, idx, num, kvv, quad, + False, True, True, None)[0] + + d["sparse_fp4_quad_ms"] = timed(bshd, iters=3) + out[str(sparsity)] = {key: round(val, 4) for key, val in d.items()} + print("DENSITY", sparsity, json.dumps(out[str(sparsity)]), flush=True) + out["dense_fp4_bshd_ms"] = timed(lambda: fa.sageattn_blackwell(*(x.transpose(1, 2) for x in (qs, ks, vs))), iters=2) + return out diff --git a/scripts/benchmarks/minimax_h3_pro6000/bench_headline.py b/scripts/benchmarks/minimax_h3_pro6000/bench_headline.py new file mode 100644 index 000000000..0ed366cbb --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/bench_headline.py @@ -0,0 +1,141 @@ +"""Headline e2e benchmark for the FastH3 local release: 480p 5 s and 768p 10 s. + +usage: python bench_headline.py [--settings 480p5s,768p10s] + +Protocol (RELEASE_PLAN §5): prompts latency-ceramics-005 + latency-harbor-005, one untimed warmup per +setting, then each prompt timed twice; e2e = generate_video wall time (encode + DiT + decode + mp4 write). +Per-stage times come from FASTVIDEO_STAGE_LOGGING. Clips and results.json go to //, +and one W&B run per invocation (group headline-) when WANDB_PROJECT is set. +""" +import argparse +import json +import os +import statistics +import time + +SETTINGS = {"480p5s": (832, 480, 124), "768p5s": (1344, 768, 124), "768p10s": (1344, 768, 243)} # 17n+5 frames at 24 fps +PROMPT_IDS = ("latency-ceramics-005", "latency-harbor-005") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("run_name") + ap.add_argument("model_dir") + ap.add_argument("num_gpus", type=int) + ap.add_argument("nvfp4_profile") + ap.add_argument("--settings", default=os.environ.get("HEADLINE_SETTINGS", "480p5s,768p10s")) + ap.add_argument("--showcase", default=os.environ.get("HEADLINE_SHOWCASE"), + help="JSON {key: prompt}: after timing, render each prompt at 480p5s for every seed") + ap.add_argument("--showcase-seeds", default=os.environ.get("HEADLINE_SHOWCASE_SEEDS", "1234,42")) + ap.add_argument("--clip-prefix", default=os.environ.get("HEADLINE_CLIP_PREFIX", "model")) + ap.add_argument("--prompts", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "headline_prompts.json")) + ap.add_argument("--out", default=os.environ.get("HEADLINE_OUT", "headline")) + ap.add_argument("--timed", type=int, default=2) + ap.add_argument("--sparsity", type=float, default=None, help="default: the checkpoint contract's vsa_sparsity") + a = ap.parse_args() + if a.timed < 1: + ap.error("--timed must be at least 1") + + contract = json.load(open(os.path.join(a.model_dir, "fastvideo_inference.json"))) + steps = contract["dmd_denoising_steps"] + sparsity = a.sparsity if a.sparsity is not None else float(contract.get("vsa_sparsity", 0.8)) + texts = json.load(open(a.prompts)) + out_dir = os.path.join(a.out, a.run_name) + os.makedirs(out_dir, exist_ok=True) + + engine = {"num_gpus": a.num_gpus, "use_fsdp_inference": False, + "parallelism": {"tp_size": 1, "sp_size": a.num_gpus}, + "offload": {"dit": False, "dit_layerwise": False, "text_encoder": False, "vae": False, + "pin_cpu_memory": False, "lazy_module_load": False}, + "compile": {"enabled": False, "vae_enabled": os.environ.get("HEADLINE_VAE_COMPILE", "1") == "1"}} + if os.environ.get("HEADLINE_BACKEND"): + engine["execution_backend"] = os.environ["HEADLINE_BACKEND"] # "ray" for multi-node + if a.nvfp4_profile != "none": + engine["quantization"] = {"transformer_quant": "NVFP4", "layer_profile": a.nvfp4_profile} + experimental = {"attention_backend": "VIDEO_SPARSE_ATTN_H3", "VSA_sparsity": sparsity, "VSA_tile_size": 64, + "h3_sequential_load": False, "inference_torch_compile": False} + if os.environ.get("HEADLINE_VAE_PARALLEL") == "1" and a.num_gpus > 1: + # Decode VAE tiles on every GPU instead of rank 0 only. + experimental.update(vae_parallel_decode=True, vae_parallel_decode_strategy="gather") + # Placement overrides for memory-limited GPUs, e.g. on a 32 GB RTX 5090: + # HEADLINE_ENGINE_JSON='{"offload": {"text_encoder": true, "pin_cpu_memory": true}}' + # HEADLINE_EXPERIMENTAL_JSON='{"h3_sequential_load": true}' + for key, value in json.loads(os.environ.get("HEADLINE_ENGINE_JSON", "{}")).items(): + if isinstance(value, dict) and isinstance(engine.get(key), dict): + engine[key].update(value) + else: + engine[key] = value + experimental.update(json.loads(os.environ.get("HEADLINE_EXPERIMENTAL_JSON", "{}"))) + config = {"model_path": a.model_dir, "engine": engine, "pipeline": {"experimental": experimental}} + # HEADLINE_* switches (VAE_PARALLEL, ENGINE_JSON, ...) change the run, so record them with the rest. + env = {k: v for k, v in os.environ.items() if k.startswith(("FASTVIDEO_", "PYTORCH_CUDA", "HEADLINE_"))} + run = None + if os.environ.get("WANDB_PROJECT"): + import wandb + run = wandb.init(project=os.environ["WANDB_PROJECT"], entity=os.environ.get("WANDB_ENTITY"), + name=a.run_name, group=os.environ.get("HEADLINE_GROUP", "headline"), job_type="benchmark", + dir=out_dir, config={"model_dir": a.model_dir, "num_gpus": a.num_gpus, + "nvfp4_profile": a.nvfp4_profile, "dmd_steps": steps, + "vsa_sparsity": sparsity, "settings": a.settings, "engine": engine, + "experimental": experimental, + "env": env, "device": os.environ.get("HEADLINE_DEVICE", "")}) + + from fastvideo import VideoGenerator + t0 = time.perf_counter() + generator = VideoGenerator.from_config(config) + results = {"run_name": a.run_name, "model_dir": a.model_dir, "num_gpus": a.num_gpus, + "nvfp4_profile": a.nvfp4_profile, "load_s": round(time.perf_counter() - t0, 1), "env": env, + "engine": engine, "experimental": experimental, "settings": {}} + try: + for name in a.settings.split(","): + width, height, frames = SETTINGS[name] + runs = [] + plan = [(PROMPT_IDS[0], True)] + [(pid, False) for _ in range(a.timed) for pid in PROMPT_IDS] + for i, (pid, warmup) in enumerate(plan): + path = os.path.join(out_dir, f"{name}_{pid}_{'warmup' if warmup else i}.mp4") + t = time.perf_counter() + generator.generate_video(prompt=texts[pid], height=height, width=width, num_frames=frames, fps=24, + guidance_scale=1.0, num_inference_steps=len(steps) + 1, seed=1234, + output_path=path, save_video=True) + wall = round(time.perf_counter() - t, 2) + runs.append({"prompt": pid, "warmup": warmup, "e2e_s": wall, "path": path}) + print("RUN", name, pid, "warmup" if warmup else "timed", wall, flush=True) + if run is not None and not warmup: + import wandb + run.log({f"{name}/e2e_s": wall, f"{name}/{pid}": wandb.Video(path, fps=24, format="mp4")}) + timed = [r["e2e_s"] for r in runs if not r["warmup"]] + results["settings"][name] = {"width": width, "height": height, "frames": frames, + "e2e_median_s": statistics.median(timed), "e2e_min_s": min(timed), "runs": runs} + print("SETTING", name, json.dumps(results["settings"][name]), flush=True) + if run is not None: + run.summary[f"{name}_e2e_median_s"] = statistics.median(timed) + json.dump(results, open(os.path.join(out_dir, "results.json"), "w"), indent=1) + if a.showcase: + # Gallery candidates: one clip per (prompt, seed) at 480p 5 s, named -[-s].mp4. + showcase = json.load(open(a.showcase)) + width, height, frames = SETTINGS["480p5s"] + seeds = [int(x) for x in a.showcase_seeds.split(",")] + os.makedirs(os.path.join(out_dir, "showcase"), exist_ok=True) + results["showcase"] = [] + for seed in seeds: + for key, text in showcase.items(): + suffix = "" if seed == 1234 else f"-s{seed}" + path = os.path.join(out_dir, "showcase", f"{a.clip_prefix}-{key}{suffix}.mp4") + t = time.perf_counter() + generator.generate_video(prompt=text, height=height, width=width, num_frames=frames, fps=24, + guidance_scale=1.0, num_inference_steps=len(steps) + 1, seed=seed, + output_path=path, save_video=True) + wall = round(time.perf_counter() - t, 2) + results["showcase"].append({"prompt": key, "seed": seed, "e2e_s": wall, "path": path}) + print("SHOWCASE", key, seed, wall, flush=True) + json.dump(results, open(os.path.join(out_dir, "results.json"), "w"), indent=1) + finally: + generator.shutdown() + if run is not None: + run.summary["load_s"] = results["load_s"] + run.finish() + print("HEADLINE_DONE", a.run_name, flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/minimax_h3_pro6000/headline_prompts.json b/scripts/benchmarks/minimax_h3_pro6000/headline_prompts.json new file mode 100644 index 000000000..553a5a7d6 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/headline_prompts.json @@ -0,0 +1,4 @@ +{ + "latency-ceramics-005": "In a quiet pottery studio an adult potter steadies a small spinning bowl while an adult apprentice watches. The apprentice asks, \"Is the rim ready?\" The potter says, \"One more gentle pass,\" and smooths the lip with a damp sponge. Begin with a close view of the hands, then make one restrained cut to a shoulder-level view showing both faces. The wheel hum, damp clay, a light splash and breathing form the soundscape. The movement is careful and unhurried, with no background music. The entire event is one finishing pass on the same bowl, not a demonstration of the whole pottery process.\nThe apprentice's apron pocket is indigo, with a plain surface. This detail belongs to the existing object, stays at its established location and remains subordinate to the main action. Preserve its material and appearance through the camera movement.\nThe wheel sits near a tall north-facing window on the left wall of the studio.\nIts circular splash pan is at the seated potter's waist.\nThe apprentice stands beyond the right side of the wheel, with enough room between their clothes and the wet rim to avoid accidental contact.\nBehind the potter, a shallow shelf holds a few dry cups.\nA sink is farther back on the right, beneath a small rack of towels.\nThese positions remain fixed through the cut.\nThe wider view reveals the same work area that surrounds the hands in the opening close view.\nThe bowl is modest in size, comfortably held between two hands.\nIts side rises from a low foot in a continuous shallow curve and ends in a slightly thickened lip.\nThe clay is warm gray with a brown undertone, darker where it is wet.\nA narrow spiral line below the rim records the potter's earlier touch.\nThat line rotates with the bowl, while a stationary highlight from the window travels over the passing surface.\nKeep these two motions distinguishable: the material marks belong to the clay, and the reflected light belongs to the relationship between the surface and the window.\nOpen with the camera low enough to see the bowl's interior without looking directly down onto the wheel.\nThe potter's left fingertips support the inside edge.\nThe right hand holds a small natural sponge outside the lip.\nNeither hand blocks the whole form.\nThe near edge of the splash pan occupies the bottom of the composition as a soft curved boundary.\nBeyond the hands, the apprentice's apron is initially out of focus.\nThis arrangement gives the close shot depth and prepares the later view of the two people without requiring another establishing shot.\nThe left hand is already stable when the clip begins.\nIts fingers form a loose supportive curve rather than squeezing the wall.\nThe sponge approaches the outside edge with only a small adjustment of the wrist.\nAs it contacts the clay, it compresses slightly and darkens where moisture gathers.\nThe bowl continues to rotate at a steady moderate speed.\nShow the finishing pass as a change in the surface's smoothness and the evenness of the lip, not as a large change in the bowl's overall shape.\nThe work is nearly finished before this moment begins.\nThe apprentice asks the question while looking at the rim, then briefly lifts their eyes toward the potter.\nTheir hands rest loosely together in front of the apron, safely away from the rotating work.\nThe question is curious and quiet, with the natural upward inflection of someone checking a detail.\nThe potter answers without stopping the wheel or turning their whole body.\nA small glance toward the apprentice is sufficient before attention returns to the clay.\nKeep the spoken words exactly as given, with no narrator explaining the technique and no extra exchange after the answer.\nCut once after the question has made the apprentice's presence clear.\nThe shoulder-level view places the potter to the left and the apprentice to the right, preserving the established relation to the wheel.\nThe bowl remains visible between them in the lower part of the frame.\nThe potter's right hand still holds the same sponge at the same point on the rim.\nContinue the wheel sound across the cut without a restart.\nThe change of view should feel like a closer understanding of the same instant, not a jump forward to a later stage of the work.\nThe potter wears a practical cotton work shirt with the sleeves rolled above the wrists.\nThe folds gather at the elbows and remain dry there.\nSmall clay marks on the forearms and apron are concentrated near the work area.\nThey do not spread or migrate during the pass.\nThe apprentice's apron is cleaner but shows a few dry pale smudges near one pocket.\nBoth garments have weight and ordinary creases.\nAvoid pristine costumes or exaggerated distressing; this is a used studio where people work carefully and clean their tools regularly.\nGive the potter a focused, patient expression.\nTheir mouth moves only for the brief reply, then settles while they feel the rim through the sponge.\nAllow the last small movement to settle within the established composition, with the environmental sound continuing around it. Preserve the quiet final composition.", + "latency-harbor-005": "## Harbor: the tide chart\nThe fabric cover around the folded tide chart is indigo, with a plain surface. This detail belongs to the existing object, stays at its established location and remains subordinate to the main action. Preserve its material and appearance through the camera movement.\nThe event takes place beside a small passenger ferry tied to a working harbor pier just before sunrise. An adult mechanic stands on the ferry's open side deck, and its captain stands beside the cabin entrance. The mechanic offers a folded tide chart and says, \"The channel is clear.\" The captain accepts it, answers, \"Then we can go,\" and looks out toward the harbor entrance. A slow lateral camera move reveals the channel beyond their shoulders. The boat remains moored throughout this brief exchange. Close voices, water against the hull, a loose halyard and a distant gull make the soundscape; there is no music.\nThe ferry is a practical coastal launch with a dark blue hull and a narrow cream band beneath its windows.\nIt carries a small enclosed cabin forward and an open passenger area behind it.\nThe camera is on the open deck, looking diagonally toward the cabin and the gap between the two people.\nThis angle places the pier along the left edge of the view and open water farther to the right.\nThe horizon is low enough that the upper part of the cabin has a clear silhouette against the pale sky.\nNothing in the composition suggests that the ferry is already underway.\nThe mechanic has finished a routine inspection rather than an emergency repair.\nTheir expression is alert but comfortable, with the slight tiredness of an early start.\nThey wear a plain work jacket over a warm shirt and carry no conspicuous badge or brand.\nA few old creases in the jacket show where the elbows bend.\nThe sleeve nearest the chart has a darker damp patch near its cuff from resting against the rail.\nKeep that patch in the same place as the arm moves.\nThe mechanic's free hand rests lightly on the top of a closed tool bag at hip level.\nThe captain is a different adult, dressed for a cool morning outside.\nA heavy knit sweater is visible beneath an open weatherproof coat.\nTheir hair is tidy but not freshly styled, and the light catches a few loose strands when they turn toward the water.\nTheir stance is balanced on the gently moving deck, with one foot slightly ahead of the other.\nThey are listening to the mechanic before the first line begins.\nThe captain does not interrupt or make a broad theatrical gesture.\nTheir reply is a small decision shared between people accustomed to working together.\nThe tide chart is a real paper object with several old folds.\nIt is partly folded into a rectangle that can be held in one hand, but one narrow flap remains loose.\nFaint printed lines and numbers are visible as a texture on its pale surface without becoming a readable title or a map inset.\nA soft graphite mark near one fold suggests that it has been used for planning.\nThe mechanic holds its lower edge between the thumb and fingers, keeping the paper clear of the damp rail.\nIts upper corner lifts slightly in the breeze before the captain takes it.\nBegin with both people already in the frame.\nThe mechanic's hand and the chart occupy the space between their bodies, below their faces.\nThis arrangement lets the first line and the handover belong to the same view.\nAs the mechanic speaks, the chart moves a short distance toward the captain.\nThe motion is neither a flourish nor an abrupt thrust.\nThe captain's receiving hand rises from beside the coat, touches the opposite edge and supports it before the mechanic releases their grip.\nThe paper bends a little between the two hands during that shared moment of support.\nThe mechanic's line is spoken in an ordinary low voice suitable for the quiet morning. The consonants remain clear, but the delivery does not sound like a public announcement. Their mouth and jaw make the small movements of the exact words, and their eyes remain on the captain. There is a slight release of breath after \"clear.\" The captain acknowledges the information with a very small nod before replying. The pause is long enough to register listening and short enough that the exchange feels familiar. Do not add another question, greeting or explanation of the voyage.\nWhen the captain says, \"Then we can go,\" the first part of the line is addressed to the mechanic. On the last words, their gaze begins to move toward the channel. The head follows the eyes through a modest turn, revealing more of the cheek nearest the exterior light. The chart settles against the front of the coat, still visibly held. The mechanic follows the captain's look with a quieter change of attention. Both remain in place. The ending is anticipation of departure, not departure itself: no engine surge, released rope or sudden movement of the ferry is needed.\nThe camera makes a restrained lateral movement toward the open-water side of the deck.\nKeep the final gesture restrained and preserve the surrounding atmosphere. Preserve the quiet final composition as the scene reaches its stated resolution." +} \ No newline at end of file diff --git a/scripts/benchmarks/minimax_h3_pro6000/showcase_prompts.json b/scripts/benchmarks/minimax_h3_pro6000/showcase_prompts.json new file mode 100644 index 000000000..048f7d9d2 --- /dev/null +++ b/scripts/benchmarks/minimax_h3_pro6000/showcase_prompts.json @@ -0,0 +1,6 @@ +{ + "corgi-weather": "Live-action TV news parody, 16:9, bright studio lighting, 35mm lens, locked-off medium shot. A fluffy tan-and-white corgi with a tiny red bow tie sits upright behind a glossy news desk, a weather map of puffy clouds glowing on the screen behind it. [0-2 s] The corgi looks straight into the camera, ears perked, and says in a cheerful deep announcer voice, \"Tomorrow: one hundred percent chance of zoomies.\" [2-5 s] It tilts its head, its tongue flops out, and it gives one proud little bark as a short brassy news jingle plays. Sound: crisp studio voice, the jingle, the single bark. Only one corgi. No subtitles, no on-screen text, no watermark, no extra animals.", + "street-food-jingle": "Documentary-style phone video, 16:9, sunny afternoon, phone held still at chest height. A cheerful street-food vendor in a green apron and white bandana stands behind a stall of golden fried snacks. [0-2 s] She smiles at the camera and claps her hands twice to set a rhythm, swaying gently. [2-5 s] She sings a bouncy jingle, \"Crispy, crunchy, ready to munch!\", and finishes with a wink and a thumbs-up. Sound: her clear singing voice in a playful melody, two claps, sizzling oil and a soft market murmur. One vendor in focus, hands empty. No subtitles, no text, no watermark.", + "parrot-pirate": "Cinematic live-action, 16:9, golden hour on the deck of a wooden pirate ship, 50mm lens, locked-off medium shot. A weathered pirate captain with a braided beard, red bandana and long coat holds a rolled treasure map; a bright green parrot sits still on his shoulder. [0-2 s] He taps the map and declares in a gravelly voice, \"X marks the spot!\" [2-5 s] The parrot bobs its head and squawks back, \"Spot! Spot!\", and the captain slowly turns to give it a tired look. Sound: his gravelly voice, the parrot's squawk, creaking wood and soft waves. One captain and one parrot. No subtitles, no text, no watermark.", + "grandma-dj": "Cinematic live-action, 16:9, rooftop at night with neon pink and blue lights, 35mm lens, slow push-in. A grandmother with silver hair in a bun, round glasses and a sequined cardigan stands alone behind DJ turntables wearing big headphones. [0-2 s] She nods to a hip-hop beat and gives the vinyl a short, confident scratch. [2-5 s] She raises one hand, grins at the camera and the bass drops. Sound: record scratch, a punchy hip-hop beat with a heavy bass drop, a faint city hum. Only the grandmother, no crowd. No subtitles, no text, no watermark." +} diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py b/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py new file mode 100644 index 000000000..463a8371c --- /dev/null +++ b/scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Convert a ModelOpt-unified NVFP4 MiniMax-H3 transformer into FastVideo's packed export. + +ModelOpt's unified Hugging Face layout (e.g. ``FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4``) +stores each quantized linear as:: + + .weight uint8 [out, in // 2] two E2M1 values per byte + .weight_scale float8_e4m3fn [out, in // 16] linear layout + .weight_scale_2 float32 [] amax(|W|) / (6 * 448) + .input_scale float32 [] static activation scale + +FastVideo's packed NVFP4 DiT export (``nvfp4_weights.safetensors``, read by +``load_minimax_h3_nvfp4_dit_export``) stores ``::_nvfp4_weight`` (same +bytes), ``::_nvfp4_weight_scale`` (the same E4M3 bytes in FlashInfer's 128x4 +swizzled layout), ``::_nvfp4_alpha`` (= ``weight_scale_2``) and +``::_weight_global_sf`` (= 1 / ``weight_scale_2``), and +``::_nvfp4_input_global_sf`` (= 1 / ``input_scale``). The calibrated weight +bytes and activation scale are preserved. Dropping the activation scale would +replace its calibrated range with a unit global scale, clipping inputs above +2688. + +``--quantize-attention`` additionally quantizes the dense BF16 attention +projections (``attn.to_{q,k,v,out}``) of every main block exactly as +``convert_model_to_nvfp4`` would at runtime, producing the full +``layer_profile="h3_dit"`` set. ``--quantize-gate`` (with it) also quantizes +each block's VSA compression gate, for ``layer_profile="h3_dit_vsa"``. +Without either the export holds the FFN linears only and must be loaded with +``layer_profile="h3_dit_ffn"``. + +``--quantize-ffn`` takes the FFN linears from a dense BF16 source instead of a +ModelOpt export, with the same round-to-nearest weight math as the runtime +(ModelOpt's max calibration produces the same weight codes). ``--act-amax`` +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 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 +the dequantized weight must stay under ``--max-probe-error`` (genuine W4A4 +noise on random inputs is about 0.1; a wrong nibble order or scale layout +reads near 1.0). Needs a Blackwell GPU with FlashInfer. + +Usage:: + + python scripts/checkpoint_conversion/convert_minimax_h3_modelopt_nvfp4_dit.py \\ + --src /path/to/FastH3-V2-NVFP4/transformer --dst /path/to/out/transformer [--quantize-attention] +""" +from __future__ import annotations + +import argparse +import json +import re +import shutil +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +EXPORT_FILENAME = "nvfp4_weights.safetensors" +# ModelOpt keys not copied as dense tensors. ``input_scale`` (ModelOpt's static activation scale) +# is carried into the packed export as ``_nvfp4_input_global_sf`` = 1 / input_scale, unless +# ``--act-amax`` supplies a replacement calibration. +DROPPED_MODELOPT_SUFFIXES = ("input_scale",) +_BLOCK_ATTN = re.compile(r"^transformer_blocks\.\d+\.attn\.(?:to_q|to_k|to_v|to_out\.0)$") +_BLOCK_FFN = re.compile(r"^transformer_blocks\.\d+\.ff\.net\.(?:0\.proj|2)$") +_BLOCK_GATE = re.compile(r"^transformer_blocks\.\d+\.attn\.to_gate_compress$") +_RENAMES = ((re.compile(r"\.ff\.net\.0\.proj$"), ".ff.fc_in"), (re.compile(r"\.ff\.net\.2$"), ".ff.fc_out"), + (re.compile(r"\.attn\.to_out\.0$"), ".attn.to_out")) +_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + + +def fastvideo_module_name(diffusers_prefix: str) -> str: + name = diffusers_prefix + for pattern, replacement in _RENAMES: + name = pattern.sub(replacement, name) + return name + + +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) + low = lut[(packed & 0x0F).long()] + high = lut[(packed >> 4).long()] + values = torch.stack((low, high), dim=-1).reshape(packed.shape[0], packed.shape[1] * 2) + block = scale.to(torch.float32).repeat_interleave(16, dim=1) + return values * block * scale_2.to(torch.float32) + + +def _flashinfer(): + from flashinfer import SfLayout, mm_fp4, nvfp4_quantize + try: + from flashinfer import block_scale_interleave + except ImportError: + from flashinfer.fp4_quantization import block_scale_interleave + return SfLayout, mm_fp4, nvfp4_quantize, block_scale_interleave + + +def probe(buffers: dict[str, torch.Tensor], reference: torch.Tensor, rows: int = 256) -> float: + """Relative error of the loader's mm_fp4 path against BF16 x @ W_ref^T.""" + SfLayout, mm_fp4, nvfp4_quantize, _ = _flashinfer() + device = buffers["_nvfp4_weight"].device + x = torch.randn(rows, reference.shape[1], device=device, dtype=torch.bfloat16) + unit = torch.tensor(1.0, device=device, dtype=torch.float32) + x_fp4, x_scale = nvfp4_quantize(x, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + out = mm_fp4(x_fp4, buffers["_nvfp4_weight"].T, x_scale, buffers["_nvfp4_weight_scale"].T, + buffers["_nvfp4_alpha"] / unit, torch.bfloat16, None, backend="auto") + ref = x.float() @ reference.float().T + return ((out.float() - ref).norm() / ref.norm()).item() + + +def convert_modelopt_linear(weight, scale, scale_2, device, input_scale=None + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, float]: + """Carry the calibrated bytes over; return (buffers, dequantized weight, scale-byte agreement). + + The agreement compares the swizzled ModelOpt scales with the scales + ``nvfp4_quantize`` derives from the dequantized weight. It is a layout check + (identical swizzle and shape) and is near 1.0 when ModelOpt used max + calibration; blocks whose largest code is below 6 legitimately differ. + """ + SfLayout, _, nvfp4_quantize, block_scale_interleave = _flashinfer() + weight = weight.to(device) + scale = scale.to(device) + scale_2 = scale_2.to(device=device, dtype=torch.float32).reshape(()) + reference = dequantize_modelopt(weight, scale, scale_2) + _, layout_ref = nvfp4_quantize(reference.to(torch.bfloat16), 1.0 / scale_2, sfLayout=SfLayout.layout_128x4, + do_shuffle=False) + swizzled = block_scale_interleave(scale.view(torch.uint8).contiguous()).reshape(layout_ref.shape) + agreement = (swizzled.view(torch.uint8) == layout_ref.view(torch.uint8)).float().mean().item() + buffers = { + "_nvfp4_weight": weight.contiguous(), + "_nvfp4_weight_scale": swizzled.view(layout_ref.dtype).contiguous(), + "_nvfp4_alpha": scale_2.clone(), + "_weight_global_sf": (1.0 / scale_2).to(torch.bfloat16), + } + if input_scale is not None: + value = input_scale.to(dtype=torch.float32) + if value.numel() != 1 or not bool(torch.isfinite(value).all()) or not bool((value > 0).all()): + raise ValueError("ModelOpt input_scale must be a finite positive scalar") + buffers["_nvfp4_input_global_sf"] = value.reshape(()).reciprocal().to(device=device) + return buffers, reference, agreement + + +def quantize_dense_linear(weight: torch.Tensor, device) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Same math as ``nvfp4_config.convert_model_to_nvfp4``.""" + SfLayout, _, nvfp4_quantize, _ = _flashinfer() + weight = weight.to(device=device, dtype=torch.bfloat16) + global_sf = (448 * 6) / weight.float().abs().nan_to_num().max() + fp4_w, fp4_s = nvfp4_quantize(weight, global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) + global_sf = torch.as_tensor(global_sf, device=device, dtype=torch.float32) + buffers = {"_nvfp4_weight": fp4_w, "_nvfp4_weight_scale": fp4_s, "_nvfp4_alpha": (1.0 / global_sf).float(), + "_weight_global_sf": global_sf.to(torch.bfloat16)} + return buffers, weight + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--src", required=True, type=Path) + parser.add_argument("--dst", required=True, type=Path) + parser.add_argument("--quantize-attention", action="store_true") + 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("--max-probe-error", type=float, default=0.3) + parser.add_argument("--dense-shard-gb", type=float, default=5.0) + args = parser.parse_args() + device = torch.device("cuda") + args.dst.mkdir(parents=True, exist_ok=True) + + index = json.loads((args.src / "diffusion_pytorch_model.safetensors.index.json").read_text()) + weight_map: dict[str, str] = index["weight_map"] + modelopt = sorted(k[:-len(".weight_scale_2")] for k in weight_map if k.endswith(".weight_scale_2")) + modelopt_keys = {f"{p}.{s}" for p in modelopt for s in ("weight", "weight_scale", "weight_scale_2", *DROPPED_MODELOPT_SUFFIXES)} + if args.quantize_gate and not args.quantize_attention: + parser.error("--quantize-gate requires --quantize-attention") + selected = ([_BLOCK_ATTN.pattern] if args.quantize_attention else []) + ( + [_BLOCK_GATE.pattern] if args.quantize_gate else []) + ([_BLOCK_FFN.pattern] if args.quantize_ffn else []) + dense_pattern = re.compile("|".join(selected)) if selected else None + # Projections ModelOpt already quantized stay on the ModelOpt path; re-quantizing their packed + # uint8 weight as if it were dense BF16 would corrupt them. + 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 [] + 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 + if args.act_amax: + raw = json.loads(args.act_amax.read_text()) + amax_table = {k: float(v["all"] if isinstance(v, dict) else v) for k, v in raw.items()} + attention_keys = {f"{p}.weight" for p in attention} + + readers = {shard: safe_open(str(args.src / shard), framework="pt", device="cpu") for shard in set(weight_map.values())} + tensor = lambda key: readers[weight_map[key]].get_tensor(key) + + export: dict[str, torch.Tensor] = {} + worst = 0.0 + agreements: list[float] = [] + attention_set = set(attention) + for prefix in modelopt + attention: + if prefix in attention_set: + buffers, reference = quantize_dense_linear(tensor(f"{prefix}.weight"), device) + else: + buffers, reference, agreement = convert_modelopt_linear(tensor(f"{prefix}.weight"), + tensor(f"{prefix}.weight_scale"), + tensor(f"{prefix}.weight_scale_2"), device, + input_scale=tensor(f"{prefix}.input_scale") + if f"{prefix}.input_scale" in weight_map else None) + agreements.append(agreement) + error = probe(buffers, reference) + worst = max(worst, error) + if error > args.max_probe_error: + 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) + for name, value in buffers.items(): + export[f"{module}::{name}"] = value.cpu() + save_file(export, str(args.dst / EXPORT_FILENAME)) + + 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, [] + for key in dense_keys: + value = tensor(key) + shard[key] = value + shard_bytes += value.numel() * value.element_size() + if shard_bytes >= args.dense_shard_gb * 1e9: + shards.append(shard) + shard, shard_bytes = {}, 0 + if shard: + shards.append(shard) + new_map = {} + for i, part in enumerate(shards, 1): + name = f"diffusion_pytorch_model-{i:05d}-of-{len(shards):05d}.safetensors" + save_file(part, str(args.dst / name)) + new_map.update({k: name for k in part}) + (args.dst / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": new_map}, indent=2) + "\n") + config = json.loads((args.src / "config.json").read_text()) + config.pop("quantization_config", None) + (args.dst / "config.json").write_text(json.dumps(config, indent=2) + "\n") + for extra in args.src.iterdir(): + if extra.suffix not in (".safetensors", ".json"): + shutil.copy2(extra, args.dst / extra.name) + print(json.dumps({"exported_linears": len(modelopt) + len(attention), "modelopt_linears": len(modelopt), + "quantized_dense_linears": len(attention), "worst_probe_error": round(worst, 4), + "static_activation_scales": sum(k.endswith("::_nvfp4_input_global_sf") for k in export), + "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, + "dense_tensors": len(dense_keys), "dense_shards": len(shards)})) + + +if __name__ == "__main__": + main() From e235b6a332a2c7a855556cee89edab2b3af997f5 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:39:24 -0700 Subject: [PATCH 053/151] [docs] Announce FastH3 on consumer hardware and FastH3 Trim (#1923) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 69c94a5e4..76b8531f2 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ **FastVideo is a unified post-training and real-time inference framework for accelerated video generation.** ## NEWS +- `2026/10/06`: FastH3 V2 now runs on a single consumer machine: NVIDIA RTX 5090, RTX 4090 and RTX PRO 6000 GPUs, DGX Spark and Apple Silicon. We also release [FastH3 Trim](https://huggingface.co/FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4), an experimental pruned model that is 4.2× smaller than base H3 and runs in as little as 8 GB of GPU memory. Get the [models](https://huggingface.co/collections/FastVideo/fastvideo-fasth3) and read the [Blog](https://haoailab.com/blogs/fasth3-rtx/). - `2026/10/06`: FastVideo now supports [Kandinsky 6](https://x.com/kandinskylab_ai/status/2107374635218055345) from Kandinsky Lab: text- and image-to-video with synchronized audio (base and 10-step distilled pi-Flow checkpoints) plus video super-resolution up to 4x. See the [Kandinsky 6 recipes](https://haoailab.com/FastVideo/cookbook/kandinsky6/). - `2026/09/15`: Release [FastH3 8-Step V2](https://huggingface.co/FastVideo/FastVideo-FastH3-8-Step-V2), an eight-forward data-free DMD2 checkpoint distilled from MiniMax-H3 with 80% Video Sparse Attention. Run it with `examples/inference/basic/basic_fasth3_8step.py` or the [FastH3 8-Step V2 recipe](https://haoailab.com/FastVideo/cookbook/minimax-h3/). - `2026/09/01`: FastH3 now runs locally on Apple Silicon through MLX and on NVIDIA DGX Spark through CUDA 13, including two-Spark inference. Follow the [FastH3 recipes](https://haoailab.com/FastVideo/cookbook/minimax-h3/) and read the [Blog](https://haoailab.com/blogs/fasth3-local/). From 02a027c49f3fa099f1a0eb1305a6646415406d87 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:17:11 -0700 Subject: [PATCH 054/151] [feat]: FastH3 on DGX Spark and Apple Silicon (#1920) Co-authored-by: Aryan Kumar Co-authored-by: Aryan Kumar Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 --- .github/workflows/ci-macos-mlx.yml | 23 +- .gitignore | 1 + docs/contributing/env_vars.md | 1 + docs/cookbook/minimax-h3.md | 5 + docs/getting_started/installation/mlx.md | 85 ++++++ docs/getting_started/installation/spark.md | 3 + .../installation/spark_pair.md | 23 ++ .../installation/spark_performance.md | 131 ++++++++- .../basic_fasth3_spark_pair_pruned_nvfp4.yaml | 64 +++++ .../basic_fasth3_spark_pair_v2_nvfp4.yaml | 64 +++++ .../basic_fasth3_spark_pruned_nvfp4.yaml | 61 ++++ .../basic/basic_fasth3_spark_v2_nvfp4.yaml | 61 ++++ .../basic/benchmark_fasth3_spark_nvfp4.py | 149 ++++++++++ examples/inference/basic/mlx_fasth3.py | 59 ++-- fastvideo/envs.py | 4 + fastvideo/layers/fp4linear.py | 17 +- fastvideo/layers/quantization/nvfp4_config.py | 42 ++- fastvideo/mlx_runtime/minimax_h3.py | 43 ++- .../mlx_runtime/minimax_h3_conditioner.py | 246 ++++++++++++++-- fastvideo/mlx_runtime/minimax_h3_pipeline.py | 272 +++++++++++++++--- fastvideo/mlx_runtime/minimax_h3_vsa.py | 22 +- fastvideo/mlx_runtime/minimax_h3_vsa_simd.py | 9 + .../tests/test_minimax_h3_bf16_source.py | 17 ++ .../tests/test_minimax_h3_rank16_adaln.py | 72 +++++ .../tests/test_minimax_h3_resident_nvfp4.py | 176 ++++++++++++ .../test_minimax_h3_vsa_gather_budget.py | 24 ++ fastvideo/pipelines/stages/base.py | 4 + .../test_benchmark_fasth3_spark_nvfp4.py | 29 ++ .../tests/mlx/test_mlx_affine_dq_gemm.py | 6 +- .../mlx/test_mlx_minimax_h3_conditioner.py | 82 +++++- .../mlx/test_mlx_minimax_h3_fast_mode.py | 155 ++++++++++ .../mlx/test_mlx_minimax_h3_fast_spatial.py | 5 +- .../test_mlx_minimax_h3_vsa_regressions.py | 63 +++- .../quantization/test_nvfp4_spark_fence.py | 96 +++++++ .../tests/stages/kandinsky6_sr/conftest.py | 28 +- .../worker/test_ray_distributed_executor.py | 14 + .../convert_minimax_h3_mlx.py | 19 ++ 37 files changed, 2017 insertions(+), 158 deletions(-) create mode 100644 examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml create mode 100644 examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml create mode 100644 examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml create mode 100644 examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml create mode 100644 examples/inference/basic/benchmark_fasth3_spark_nvfp4.py create mode 100644 fastvideo/mlx_runtime/tests/test_minimax_h3_bf16_source.py create mode 100644 fastvideo/mlx_runtime/tests/test_minimax_h3_rank16_adaln.py create mode 100644 fastvideo/mlx_runtime/tests/test_minimax_h3_resident_nvfp4.py create mode 100644 fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py create mode 100644 fastvideo/tests/inference/test_benchmark_fasth3_spark_nvfp4.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml index 3ef19a311..29a3261eb 100644 --- a/.github/workflows/ci-macos-mlx.yml +++ b/.github/workflows/ci-macos-mlx.yml @@ -33,15 +33,15 @@ jobs: env: FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA TOKENIZERS_PARALLELISM: "false" - MASTER_ADDR: localhost + MASTER_ADDR: "127.0.0.1" MASTER_PORT: "29513" + GLOO_SOCKET_IFNAME: lo0 steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: "3.12" - cache: pip - uses: astral-sh/setup-uv@v3 @@ -49,9 +49,9 @@ jobs: run: | uv pip install --system \ --index-url https://download.pytorch.org/whl/cpu \ - torch==2.11.0 torchvision torchaudio + torch==2.12.0 torchvision torchaudio uv pip install --system \ - pytest numpy scipy pillow imageio einops cloudpickle filelock \ + pytest pytest-timeout numpy scipy pillow imageio einops cloudpickle filelock \ PyYAML diffusers huggingface_hub remote-pdb safetensors loguru mlx \ "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0" @@ -65,8 +65,8 @@ jobs: print("machine:", platform.machine()) print("processor:", platform.processor()) print("mlx default device:", mx.default_device()) - memory_size = mx.metal.device_info().get("memory_size") if mx.metal.is_available() else "metal unavailable" - print("mlx memory_size:", memory_size) + device_info = mx.metal.device_info() if mx.metal.is_available() else "metal unavailable" + print("mlx device_info:", device_info) print("torch:", torch.__version__) print("torch mps available:", torch.backends.mps.is_available()) PY @@ -74,6 +74,7 @@ jobs: - name: Run MLX smoke tests run: | python -m pytest \ + fastvideo/mlx_runtime/tests/ \ fastvideo/tests/mlx/test_dmd_sampling.py \ fastvideo/tests/mlx/test_memory_limits.py \ fastvideo/tests/mlx/test_quant_capability.py \ @@ -101,7 +102,7 @@ jobs: fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \ fastvideo/tests/platforms/test_mps_vsa_error.py \ fastvideo/tests/platforms/test_cpu_sdpa.py \ - -v -s -o faulthandler_timeout=120 + -v -s --timeout=120 -o faulthandler_timeout=120 # Same tests on MLX's CPU backend. Hosted macOS runners are scarce and # slower to schedule; this Linux job gives fast PR signal on the identical @@ -122,7 +123,6 @@ jobs: - uses: actions/setup-python@v5 with: python-version: "3.12" - cache: pip - uses: astral-sh/setup-uv@v3 @@ -130,15 +130,16 @@ jobs: run: | uv pip install --system \ --index-url https://download.pytorch.org/whl/cpu \ - torch==2.11.0 torchvision torchaudio + torch==2.12.0 torchvision torchaudio uv pip install --system \ - pytest numpy scipy pillow imageio einops cloudpickle filelock \ + pytest pytest-timeout numpy scipy pillow imageio einops cloudpickle filelock \ PyYAML diffusers huggingface_hub remote-pdb safetensors loguru "mlx[cpu]" \ "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0" - name: Run MLX smoke tests (CPU backend) run: | python -m pytest \ + fastvideo/mlx_runtime/tests/ \ fastvideo/tests/mlx/test_dmd_sampling.py \ fastvideo/tests/mlx/test_memory_limits.py \ fastvideo/tests/mlx/test_quant_capability.py \ @@ -166,4 +167,4 @@ jobs: fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \ fastvideo/tests/platforms/test_mps_vsa_error.py \ fastvideo/tests/platforms/test_cpu_sdpa.py \ - -v -s -o faulthandler_timeout=120 + -v -s --timeout=120 -o faulthandler_timeout=120 diff --git a/.gitignore b/.gitignore index 7764e2ea7..b65ab4377 100644 --- a/.gitignore +++ b/.gitignore @@ -134,6 +134,7 @@ openspec/ fastvideo/tests/ssim/reference_videos/** !fastvideo/tests/ssim/reference_videos/**/*.mp4 !fastvideo/tests/ssim/reference_videos/**/*.png +fastvideo/tests/ssim/.reference_videos_download.lock # Local H3 MLX kernel / exactness benches (JSON, logs, frames, videos) .kernel_bench/ diff --git a/docs/contributing/env_vars.md b/docs/contributing/env_vars.md index 883e2d8ac..3b85bfe18 100644 --- a/docs/contributing/env_vars.md +++ b/docs/contributing/env_vars.md @@ -192,6 +192,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry. | `FASTVIDEO_FA4` | bool | `0` | attention | The FLASH_ATTN backend uses FlashAttention-4 (flash_attn.cute) instead of FA3 or FA2. | | `FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN` | bool | `0` | attention | MiniMax-H3 dense DiT self-attention uses the FlashAttention-4 packed-varlen entry point. This changes the floating-point reduction order, so it is an inference-only opt-in. | | `FASTVIDEO_VSA_SM100A` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. fastvideo-kernel reads the same variable with the same rule. | +| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. | | `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. | | `FASTVIDEO_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. | diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index e01ee210a..0ccb0b3bb 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -13,6 +13,11 @@ is why V2 is the higher-quality FastH3. The V2 schedule contract is in **CompactH3** is the 42-block 20B NVFP4 H3 checkpoint for one Blackwell GPU (RTX 5090 or RTX PRO 6000). The FastH3 V1 and V2 recipes are unchanged. +The 42-block pruned checkpoint has an [MLX INT8/INT6 conversion and +eight-forward T2VA command](../getting_started/installation/mlx.md#pruned-eight-forward-checkpoint). +It reads `fastvideo_inference.json` for the trained schedule. The command +uses native 832x480 resolution and all requested frames. +
    All model families diff --git a/docs/getting_started/installation/mlx.md b/docs/getting_started/installation/mlx.md index 359218e58..fb94d3e9f 100644 --- a/docs/getting_started/installation/mlx.md +++ b/docs/getting_started/installation/mlx.md @@ -46,6 +46,91 @@ is the higher-quality FastH3. Recorded shapes and evidence live in the [support matrix](../../inference/support_matrix.md#apple-silicon-native-runtime). +## FastH3 V2 and Trim with INT6 + +The released sources are `FastVideo/FastVideo-FastH3-8-Step-V2` and +`FastVideo/FastVideo-FastH3-Trim-8-Step`. Trim has 42 transformer blocks and +rank-16 AdaLN. Both use eight denoising forwards, video/audio shifts of 10/3, +VSA sparsity 0.8, a native NVFP4 text encoder, and the 26-layer light video VAE. +Keep `fastvideo_inference.json` beside `transformer/`; conversion reads its +schedule to build the AdaLN cache. + +Convert the BF16 transformer to affine INT6 with its VSA gates: + +```bash +hf download FastVideo/FastVideo-FastH3-Trim-8-Step \ + --local-dir ./FastH3-Trim + +python scripts/checkpoint_conversion/convert_minimax_h3_mlx.py \ + --model-root ./FastH3-Trim/transformer \ + --out ./FastH3-Trim-MLX \ + --formats "int6" --include-vsa +``` + +For V2, use `FastVideo/FastVideo-FastH3-8-Step-V2` and separate source/output +directories. Preconverted release snapshots use the same names with the +`-MLX-INT6` suffix. Each snapshot includes the encoder in MLX layout, both VAEs, +and the trained schedule, so it does not require a second encoder download. + +The 36 GiB M4 Max release recipe uses phased placement. It loads the encoder, +DiT, and decoders in turn. Use reference attention and native output geometry: + +```python +from pathlib import Path +from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline + +root = Path("./FastH3-Trim") +pipeline = MiniMaxH3MLXPipeline( + model_root=root, + mlx_dit_checkpoint="./FastH3-Trim-MLX/int6", + conditioner_mode="nvfp4", + resident=False, + vae_dtype="fp16", + metal_wired_limit_gib=27, +) +try: + pipeline.generate( + "A corgi news anchor sits behind a desk and gives a cheerful bark.", + output_path="./outputs/trim-int6-corgi.mp4", + width=832, height=480, num_frames=124, seed=1234, + num_steps=8, vsa=True, vsa_sparsity=0.8, vsa_tile_size=64, + vsa_impl="reference", vae_tile_height=256, vae_tile_width=256, + ) +finally: + pipeline.close() +``` + +Set `FASTVIDEO_MLX_DQ_GEMM=1` before running this Python command. It selects the +validated affine dequantization followed by dense matrix multiplication. +124 frames at 24 fps is roughly five seconds. The recipe preserves all frames +and the requested resolution. + +### Native NVFP4 encoder and cache + +The MLX conditioner reads the released packed NVFP4 weights without +requantization. It retains the layers H3 reads and can cache the packed weights +in MLX layout. The cache is written in a staging directory and published by a +single rename. A cache hit changes storage layout, not encoder arithmetic. + +MLX uses BF16 embeddings and FP32 activations; CUDA uses quantized activations. +Generated video and audio must be reviewed before claiming cross-runtime +quality parity. MLX 0.32.2 supports the required operator on Apple Silicon. + +### Metal wired memory + +MLX's allocation limit and wired-memory limit are separate. The optional +`metal_wired_limit_gib` calls `mx.set_wired_limit` to keep selected Metal +allocations in physical memory. It does not add RAM. An explicit request fails +if the installed MLX build cannot apply it. `close()` restores the previous +wired limit. The phased pipeline also sets a 30 GiB maximum allocator guideline +and restores its previous value on close. Resident placement keeps the existing +allocator limit so larger Macs can hold all components. + +The tested 36 GiB M4 Max recipe uses 27 GiB and phased placement. Resident +placement also requires capacity for all components and peak activations; +wiring cannot make an oversized stack fit. Inspect `mx.device_info()` before +choosing a limit on another Mac. + ## Hardware - FastMetal 1.3B and 5B: 16 GB unified memory and up diff --git a/docs/getting_started/installation/spark.md b/docs/getting_started/installation/spark.md index 33888a853..570cf71f2 100644 --- a/docs/getting_started/installation/spark.md +++ b/docs/getting_started/installation/spark.md @@ -144,6 +144,9 @@ for which models are practical on the GB10, what makes them faster, and what won't help on this hardware (and why) — so you don't spend a night tuning knobs that can't move here. +For the eight-forward FastH3 V2 NVFP4 stack with a trimmed encoder and light +VAE, use the [one-Spark resident recipe](spark_performance.md#fasth3-v2-nvfp4-on-one-spark). + Two Sparks with QSFP cables: [Pair two NVIDIA DGX Sparks](spark_pair.md) for one FastH3 clip across both GPUs (`sp_size=2` over Ray). Copy-paste commands for one or two Sparks also live on the diff --git a/docs/getting_started/installation/spark_pair.md b/docs/getting_started/installation/spark_pair.md index fe8aee1aa..928bfe91b 100644 --- a/docs/getting_started/installation/spark_pair.md +++ b/docs/getting_started/installation/spark_pair.md @@ -149,6 +149,29 @@ 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 + +The public eight-forward V2 and Trim stacks include the NVFP4 encoder and +lightweight video VAE. Use the environment above plus the release kernel +settings, then run one of these configs from the head: + +```bash +export FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass +export FASTVIDEO_H3_VAE_TILE_BATCH=1 FASTVIDEO_VSA_TRITON=1 +export FASTVIDEO_VSA_SM100A=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_pruned_nvfp4.yaml for FastH3 Trim. +``` + +Both configs use their released Hugging Face model paths, `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 +to those larger stacks. Set 1344x768 for native 768p with the same 124 frames. +See [the resident release recipe](spark_performance.md#fasth3-v2-nvfp4-on-one-spark) +for checkpoint contents and the timing protocol. + ## FastH3 frame counts H3 is 24 fps. Legal `num_frames` values are `17n+5`. The pipeline rejects diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index e0f5cb4d1..c16dc9e9a 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -161,14 +161,16 @@ is power-cycled. To avoid it: on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains available because it partitions weights without parking them in a separate host pool. -- **MiniMax H3 / FastH3** still needs deferred loading on one GB10. The Qwen3-VL - conditioner is tens of gigabytes of BF16. If the DiT and VAEs load while that - encoder is still resident, the process is a typical `earlyoom` kill (Python is - preferred). On unified memory, `lazy_module_load` auto-enables and owns that +- **Older MiniMax H3 / FastH3 bf16 weights** need deferred loading on one GB10. + The full Qwen3-VL conditioner is tens of gigabytes of BF16. If the DiT and + VAEs load while that encoder is still resident, the process can be killed by + `earlyoom`. On unified memory, `lazy_module_load` auto-enables and owns that split (encoder, then DiT, then VAE; DiT can drop before decode). Sequential - load is the H3-only fallback when lazy is off; do not pass - `--no-lazy-module-load` here. Geometry scalars come from checkpoint - `config.json`, not live weights. See [Offloading](../../inference/offloading.md). + load is the H3-only fallback when lazy is off. Keep deferred loading for + those older checkpoints. The trimmed NVFP4 encoder and light VAE in the + [V2 resident recipe](#fasth3-v2-nvfp4-on-one-spark) are a different memory + profile. Geometry scalars come from checkpoint `config.json`, not live + weights. See [Offloading](../../inference/offloading.md). - **FastH3 TAEH3** (`--video-decode-backend taeh3`) is an opt-in preview decoder. T2VA never materializes the 9.7 GiB video VAE (DiT still loads after Qwen via sequential start). On this box, alpine 768×1344×124 decoded in **2.4 s** versus @@ -207,6 +209,121 @@ A few things that surprise people on this box (beyond the memory notes above): `Released MiniMax-H3 text encoder after conditioning` before `Loading MiniMax-H3 denoise modules`). +## FastH3 V2 NVFP4 on one Spark + +This recipe uses the full V2 eight-forward transformer, the 50-layer NVFP4 +Qwen3-VL encoder, and the light H3 video VAE. Its configuration keeps all +three resident on one GB10. This stack fits in the Spark's unified memory; +benchmark your installed runtime and review the clips before publishing a +speed claim. The earlier bf16 H3 memory guidance above concerns a larger checkpoint. + +Install FastVideo following [the Spark install guide](spark.md). The released +repositories are complete inference stacks: + +| Model | Repository | Packed DiT profile | +|---|---|---| +| FastH3 V2 | `FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer` | `h3_dit_vsa` | +| FastH3 Trim | `FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4` | `h3_dit_vsa` | + +Each ships its own trained schedule, NVFP4 transformer and 50-layer NVFP4 +text encoder, lightweight 26-layer video VAE with the INT8-weight overlay, +and audio VAE. Download the complete repository; the runtime selects these +components from its model index. The released Trim transformer also packs +attention and VSA gates, unlike the earlier FFN-only pruned export. Neither +release requires local checkpoint conversion or a separate encoder download. + +```bash +hf download FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer +hf download FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4 +``` + +Keep `fastvideo_inference.json` with the transformer if you stage the stack +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. + +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 +64-token tiles, and the light H3 VAE through the `h3-vae` decode backend. +It does not use frame dropping or spatial upscaling. + +```bash +FASTVIDEO_MINIMAX_H3_FUSIONS=all \ +FASTVIDEO_NVFP4_MM_BACKEND=cutlass \ +FASTVIDEO_H3_VAE_TILE_BATCH=1 \ +FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=0 FASTVIDEO_FA4=0 \ +FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 \ +FASTVIDEO_STAGE_LOGGING=1 \ +nice -n 19 fastvideo generate \ + --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml +``` + +The defaults produce a roughly five-second clip at 24 fps. For native 768p, +set `--request.sampling.width 1344` and `--request.sampling.height 768`, +keeping 124 frames. For the separate ten-second setting, use 243 frames. +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, +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 +gathering, with tile batch 1 on each worker. + +For release timing, create one generator per model/resolution. Run one +untimed ceramics warmup, then two timed ceramics calls and two timed harbor +calls in that same process, using the exact release prompt strings, seed +1234 and the settings above. Measure each `generate()` call through finished +MP4 output and report the median of the two timed calls per prompt. Keep the +warmup excluded. Record the model revision, code commit, command, environment +and peak-memory scope with the results. Review every clip's video and audio +before publishing a quality or speed claim. + +The benchmark helper requires a local stack so it can validate the schedule +before loading. For example: + +```bash +hf download FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer \ + --local-dir ./FastH3-V2-Consumer +python examples/inference/basic/benchmark_fasth3_spark_nvfp4.py \ + --config examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml \ + --model-path ./FastH3-V2-Consumer --frames 124 \ + --prompts /path/to/benchmark_prompts.json \ + --output-dir outputs/fasth3_spark_v2_nvfp4/benchmark-124 +``` + +Set the environment from the generation command above before benchmarking. +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. + +### Released model measurements + +The released Trim stack at revision `cae9ceb6feefe77d34a56640782cda3909363f19` +completed the native 832x480, 124-frame protocol on one Spark, seed 1234. +The tested code is `6ccdbc761e6b854003f472e39b826c06cb54de60`, using the +resident recipe above. Medians exclude warmups and cover finished MP4 output. + +| Released model | Resolution | One Spark, ceramics / harbor | Two Sparks | +|---|---|---:|---| +| FastH3 Trim NVFP4 | 832x480 | 124.090 / 127.519 s | Pending | +| FastH3 Trim NVFP4 | 1344x768 | Review pending | Pending | +| FastH3 V2 NVFP4-Consumer | 832x480 | Repeatability review pending | Pending | +| FastH3 V2 NVFP4-Consumer | 1344x768 | Review pending | Pending | + +The completed Trim 480p batch used one extra untimed harbor warmup in the +same process, six calls total. Both warmups are excluded from the medians. +Each prompt's three decoded videos and audio streams match exactly, and +sampled frames are coherent. These checks do not establish BF16 parity, +speech accuracy or lip sync. Pending cells are not measured substitutes +from older checkpoints. + ## Reproduce these numbers Two scripts under `examples/inference/optimizations/` reproduce the claims on diff --git a/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml new file mode 100644 index 000000000..5d59fcfb1 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_pair_pruned_nvfp4.yaml @@ -0,0 +1,64 @@ +# Start Ray on both Sparks and source spark_pair_env.sh first; see spark_pair.md. +# FastH3 Trim eight-forward video+audio on two DGX Sparks over QSFP RoCE. +# Download the complete checkpoint, including its trained schedule, as described in +# docs/getting_started/installation/spark_performance.md. +# +# GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. +# 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_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=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_pruned_nvfp4.yaml +generator: + model_path: FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4 + 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 + 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_pruned_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 new file mode 100644 index 000000000..4d1a0d55e --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_pair_v2_nvfp4.yaml @@ -0,0 +1,64 @@ +# Start Ray on both Sparks and source spark_pair_env.sh first; see spark_pair.md. +# FastH3 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. +# +# GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. +# 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_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=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_v2_nvfp4.yaml +generator: + model_path: FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer + 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 + 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_v2_nvfp4/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml new file mode 100644 index 000000000..5ad449512 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_pruned_nvfp4.yaml @@ -0,0 +1,61 @@ +# FastH3 Trim eight-forward video+audio on one DGX Spark. +# Download the complete checkpoint, including its trained schedule, as described in +# docs/getting_started/installation/spark_performance.md. +# +# GB10 uses Triton VSA. The release packs attention, FFN and VSA gate weights. +# 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_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=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 +generator: + model_path: FastVideo/FastVideo-FastH3-Trim-8-Step-NVFP4 + 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: false + pipeline: + workload_type: t2v + vae_tiling: true + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + 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_pruned_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 new file mode 100644 index 000000000..27531db0e --- /dev/null +++ b/examples/inference/basic/basic_fasth3_spark_v2_nvfp4.yaml @@ -0,0 +1,61 @@ +# FastH3 V2 eight-forward video+audio on one DGX Spark. +# 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. +# 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_VSA_TRITON=1 FASTVIDEO_VSA_SM100A=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 +generator: + model_path: FastVideo/FastVideo-FastH3-8-Step-V2-NVFP4-Consumer + 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: false + pipeline: + workload_type: t2v + vae_tiling: true + experimental: + attention_backend: VIDEO_SPARSE_ATTN_H3 + 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_v2_nvfp4/ + save_video: true + return_frames: false diff --git a/examples/inference/basic/benchmark_fasth3_spark_nvfp4.py b/examples/inference/basic/benchmark_fasth3_spark_nvfp4.py new file mode 100644 index 000000000..96778cf6a --- /dev/null +++ b/examples/inference/basic/benchmark_fasth3_spark_nvfp4.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Time a resident FastH3 eight-forward Spark recipe with the release prompts. + +One process loads the model, runs one excluded ceramics warmup, then times +each prompt at least twice. Every call writes a video. This script does not alter +the V2 schedule, VSA sparsity, or video resolution. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from copy import deepcopy +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api.parser import load_raw_config, parse_config +from fastvideo.api.schema import RunConfig + +PROMPT_IDS = ("latency-ceramics-005", "latency-harbor-005") + + +def _stage_metrics(result: object) -> dict[str, dict]: + logging_info = getattr(result, "logging_info", None) + stages = getattr(logging_info, "stages", None) + if isinstance(logging_info, dict): + stages = logging_info.get("stages", stages) + if not isinstance(stages, dict): + return {} + return {name: metrics for name, metrics in stages.items() if isinstance(metrics, dict)} + + +def _stage_seconds(metrics: dict[str, dict]) -> dict[str, float]: + return {name: float(stage["execution_time"]) for name, stage in metrics.items() + if stage.get("execution_time") is not None} + + +def _peak_mb(metrics: dict[str, dict], key: str) -> float | None: + values = [float(stage[key]) for stage in metrics.values() if stage.get(key) is not None] + return max(values) if values else None + + +def _stage_total(stages: dict[str, float], fragment: str, exclude: str | None = None) -> float | None: + matches = [ + seconds for name, seconds in stages.items() + if fragment in name.lower() and (exclude is None or exclude not in name.lower()) + ] + return sum(matches) if matches else None + + +def _request(base: RunConfig, prompt: str, frames: int, width: int, height: int, + output: Path): + request = deepcopy(base.request) + request.prompt = prompt + request.inputs.prompt_path = None + request.sampling.num_frames = frames + request.sampling.width = width + request.sampling.height = height + request.output.output_path = str(output) + request.output.save_video = True + request.output.return_frames = False + return request + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--prompts", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model-path", type=Path) + parser.add_argument("--frames", type=int, default=124) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--repeats", type=int, default=2) + args = parser.parse_args() + + if args.frames not in (124, 243): + parser.error("use 124 frames for roughly five seconds or 243 for roughly ten seconds") + if args.repeats < 2: + parser.error("the release protocol requires at least two timed calls") + + config = parse_config(RunConfig, load_raw_config(args.config)) + if args.model_path: + config.generator.model_path = str(args.model_path) + if config.request.sampling.num_inference_steps != 9: + parser.error("the V2 contract requires nine sigma points for eight DiT forwards") + if config.generator.engine.offload.lazy_module_load is not False: + parser.error("the resident recipe requires lazy_module_load: false") + contract = Path(config.generator.model_path) / "fastvideo_inference.json" + if not contract.is_file(): + parser.error(f"missing trained V2 schedule: {contract}") + inference = json.loads(contract.read_text()) + if inference.get("num_inference_steps") != 9 or inference.get("transformer_forwards") != 8: + parser.error("the checkpoint is not the trained V2 eight-forward schedule") + + prompts = json.loads(args.prompts.read_text()) + if any(prompt_id not in prompts for prompt_id in PROMPT_IDS): + parser.error(f"prompt JSON must contain {', '.join(PROMPT_IDS)}") + args.output_dir.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config(config.generator) + try: + for prompt_index, prompt_id in enumerate(PROMPT_IDS): + times = [] + first_index = 0 if prompt_index == 0 else 1 + for index in range(first_index, args.repeats + 1): + warmup = index == 0 + label = "warmup" if warmup else f"run-{index:02d}" + requested_path = args.output_dir / f"{prompt_id}-{args.width}x{args.height}-{args.frames}-{label}.mp4" + request = _request(config, prompts[prompt_id], args.frames, args.width, args.height, + requested_path) + started = time.perf_counter() + result = generator.generate(request) + wall = time.perf_counter() - started + output = Path(result.video_path) if result.video_path else requested_path + if not output.is_file(): + raise RuntimeError(f"generation returned without an MP4: {output}") + metrics = _stage_metrics(result) + stages = _stage_seconds(metrics) + row = { + "prompt_id": prompt_id, + "warmup": warmup, + "frames": args.frames, + "width": args.width, + "height": args.height, + "e2e_seconds": round(wall, 3), + "denoise_seconds": _stage_total(stages, "denois"), + "decode_seconds": _stage_total(stages, "decod", exclude="postdecode"), + "postprocess_seconds": _stage_total(stages, "postdecode"), + "peak_memory_mb": _peak_mb(metrics, "peak_allocated_mb"), + "peak_reserved_mb": _peak_mb(metrics, "peak_reserved_mb"), + "result_peak_memory_mb": result.peak_memory_mb, + "stages": stages, + "mp4": str(output), + } + print(json.dumps(row, sort_keys=True), flush=True) + if not warmup: + times.append(wall) + print(json.dumps({"prompt_id": prompt_id, "timed_runs": len(times), + "median_e2e_seconds": round(statistics.median(times), 3)}, + sort_keys=True), flush=True) + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_fasth3.py b/examples/inference/basic/mlx_fasth3.py index 62ec602ea..11277a200 100644 --- a/examples/inference/basic/mlx_fasth3.py +++ b/examples/inference/basic/mlx_fasth3.py @@ -190,34 +190,37 @@ def run(args: argparse.Namespace) -> None: taeh3_chunk_size=args.taeh3_chunk_size, prompt_cache_dir=args.prompt_cache_dir, ) - result = pipeline.generate( - args.prompt, - output_path=args.output_path, - height=args.height, - width=args.width, - num_frames=args.num_frames, - seed=args.seed, - num_steps=args.steps, - tiled_video_decode=args.tiled_video_decode, - vae_tile_height=args.vae_tile_height, - vae_tile_width=args.vae_tile_width, - inter_step_cooldown_s=args.inter_step_cooldown_s, - fast=args.fast, - fast_factor=args.fast_factor, - fast_sharpen=args.fast_sharpen, - rife_weights_dir=args.rife_weights_dir, - fast_spatial=args.fast_spatial, - fast_spatial_scale=args.fast_spatial_scale, - fast_spatial_upsample_mode=args.fast_spatial_upsample_mode, - fast_spatial_sharpen=args.fast_spatial_sharpen, - vsa=args.vsa, - vsa_sparsity=args.vsa_sparsity, - vsa_tile_size=args.vsa_tile_size, - vsa_prefix_mode=args.vsa_prefix_mode, - vsa_dense_first_n_steps=args.vsa_dense_first_n_steps, - vsa_dense_layers=args.vsa_dense_layers, - vsa_impl=args.vsa_impl, - ) + try: + result = pipeline.generate( + args.prompt, + output_path=args.output_path, + height=args.height, + width=args.width, + num_frames=args.num_frames, + seed=args.seed, + num_steps=args.steps, + tiled_video_decode=args.tiled_video_decode, + vae_tile_height=args.vae_tile_height, + vae_tile_width=args.vae_tile_width, + inter_step_cooldown_s=args.inter_step_cooldown_s, + fast=args.fast, + fast_factor=args.fast_factor, + fast_sharpen=args.fast_sharpen, + rife_weights_dir=args.rife_weights_dir, + fast_spatial=args.fast_spatial, + fast_spatial_scale=args.fast_spatial_scale, + fast_spatial_upsample_mode=args.fast_spatial_upsample_mode, + fast_spatial_sharpen=args.fast_spatial_sharpen, + vsa=args.vsa, + vsa_sparsity=args.vsa_sparsity, + vsa_tile_size=args.vsa_tile_size, + vsa_prefix_mode=args.vsa_prefix_mode, + vsa_dense_first_n_steps=args.vsa_dense_first_n_steps, + vsa_dense_layers=args.vsa_dense_layers, + vsa_impl=args.vsa_impl, + ) + finally: + pipeline.close() print(json.dumps({ "video_path": result.video_path, "timings_s": {k: round(v, 2) for k, v in result.timings.items()}, diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 77ba871a6..e9d4c3be8 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -331,6 +331,10 @@ FASTVIDEO_VSA_SM100A = EnvBool( category="attention", doc="VIDEO_SPARSE_ATTN_H3 sends no-grad tile-64 forwards to the data-center Blackwell (sm_100a) kernel. " "fastvideo-kernel reads the same variable with the same rule.") +FASTVIDEO_VSA_TRITON = EnvBool( + False, + category="attention", + doc="Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable.") FASTVIDEO_NVFP4_FA4 = EnvBool( False, category="attention", diff --git a/fastvideo/layers/fp4linear.py b/fastvideo/layers/fp4linear.py index 6b2b61368..82a2c7573 100644 --- a/fastvideo/layers/fp4linear.py +++ b/fastvideo/layers/fp4linear.py @@ -2,6 +2,8 @@ from typing import Any import torch +from fastvideo.layers.quantization.nvfp4_config import nvfp4_quantize_fenced + try: import flashinfer except ImportError: @@ -50,18 +52,9 @@ class _LinearFWD4BWD16Fn(torch.autograd.Function): global_sf_a = _global_sf(x2d) global_sf_b = _global_sf(weight_cast) - a_fp4, a_inv_s = flashinfer_mod.nvfp4_quantize( - x2d, - global_sf_a, - sfLayout=a_sf_layout, - do_shuffle=False, - ) - b_fp4, b_inv_s = flashinfer_mod.nvfp4_quantize( - weight_cast, - global_sf_b, - sfLayout=flashinfer_mod.SfLayout.layout_128x4, - do_shuffle=False, - ) + # Shared helper: carries the DGX Spark (GB10) ordering fence. + a_fp4, a_inv_s = nvfp4_quantize_fenced(x2d, global_sf_a, a_sf_layout.value) + b_fp4, b_inv_s = nvfp4_quantize_fenced(weight_cast, global_sf_b, flashinfer_mod.SfLayout.layout_128x4.value) alpha = 1.0 / (global_sf_a * global_sf_b) diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index 3e683f749..17460f46f 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/fastvideo/layers/quantization/nvfp4_config.py @@ -19,6 +19,7 @@ use time, with a clear error. """ from __future__ import annotations +import functools import logging import os import re @@ -171,6 +172,44 @@ def _get_ltx2_fp4_stage_profile(default: str = "refine") -> str: return default +@functools.cache +def _is_dgx_spark(device_index: int) -> bool: + return torch.cuda.get_device_capability(device_index) == (12, 1) + + +def nvfp4_quantize_fenced( + x: torch.Tensor, + global_sf: torch.Tensor, + sf_layout: int, + do_shuffle: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """FlashInfer NVFP4 quantization with the DGX Spark (GB10) ordering fence. + + Every FastVideo NVFP4 quantization goes through here, so the GB10 + workaround covers inference and the QAT straight-through linear alike. + """ + device_index = x.device.index if x.device.index is not None else torch.cuda.current_device() + spark = _is_dgx_spark(device_index) + if spark and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("NVFP4 activation quantization on DGX Spark requires a completion fence; " + "disable CUDA graph capture.") + SfLayout, _, nvfp4_quantize = _require_flashinfer() + if not spark: + return nvfp4_quantize(x, global_sf, sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle) + # FlashInfer's PDL kernel reads the global scale before its + # dependency wait. Fresh dynamic scales require normal ordering. + quantized, scales = nvfp4_quantize(x, + global_sf, + 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. + torch.cuda.current_stream(x.device).synchronize() + return quantized, scales + + _OPS_REGISTERED = False @@ -193,8 +232,7 @@ def _register_ops_once() -> None: sf_layout: int, do_shuffle: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - SfLayout, _, nvfp4_quantize = _require_flashinfer() - return nvfp4_quantize(x, global_sf, sfLayout=SfLayout(sf_layout), do_shuffle=do_shuffle) + return nvfp4_quantize_fenced(x, global_sf, sf_layout, do_shuffle) @_nvfp4_quantize_op.register_fake def _nvfp4_quantize_op_fake( diff --git a/fastvideo/mlx_runtime/minimax_h3.py b/fastvideo/mlx_runtime/minimax_h3.py index 58b677bd8..c31180de3 100644 --- a/fastvideo/mlx_runtime/minimax_h3.py +++ b/fastvideo/mlx_runtime/minimax_h3.py @@ -766,12 +766,12 @@ def _feed_forward(weights: dict[str, Any], x): return linear(value * silu(gate), weights["ff.net.2.weight"]) -def _adaln_tables(weights: dict[str, Any], temb): +def _adaln_tables(weights: dict[str, Any], temb, *, apply_silu: bool = True): """Six (n_t * 3, hidden) modulation tables from (n_t, time_embed_dim).""" import mlx.core as mx projected = linear( - silu(temb).astype(weight_dtype(weights["adaln_proj.linear.weight"])), + (silu(temb) if apply_silu else temb).astype(weight_dtype(weights["adaln_proj.linear.weight"])), weights["adaln_proj.linear.weight"], weights["adaln_proj.linear.bias"], ) @@ -914,6 +914,7 @@ class MLXMiniMaxH3DiT: self.qk_norm_eps = float(config["qk_norm_eps"]) self.final_norm_eps = float(config["final_norm_eps"]) self.patch_dim = self.in_channels * math.prod(self.patch_size) + self.adaln_rank = config.get("adaln_rank") self._adaln_cache: MiniMaxH3StepCache | None = None self.vsa_config = MiniMaxH3VSAConfig() self._vsa_geometry: MiniMaxH3VSAGeometry | None = None @@ -932,11 +933,15 @@ class MLXMiniMaxH3DiT: self.weights["time_embedder.linear_1.weight"], self.weights["time_embedder.linear_1.bias"], ) - return linear( + temb = linear( silu(temb), self.weights["time_embedder.linear_2.weight"], self.weights["time_embedder.linear_2.bias"], ) + if self.adaln_rank is not None: + temb = linear( + silu(temb).astype(weight_dtype(self.weights["adaln_basis.weight"])), self.weights["adaln_basis.weight"]) + return temb def refine_text(self, text_rows): hidden = linear( @@ -969,9 +974,10 @@ class MLXMiniMaxH3DiT: timesteps = np.unique(np.asarray(timesteps, dtype=np.float32)) temb = self.compute_temb(mx.array(timesteps)) - block_tables = [_adaln_tables(block, temb) for block in self.blocks] + block_tables = [_adaln_tables(block, temb, apply_silu=self.adaln_rank is None) for block in self.blocks] shift_scale = linear( - silu(temb).astype(weight_dtype(self.weights["norm_out.linear.weight"])), + (silu(temb) if self.adaln_rank is None else temb).astype( + weight_dtype(self.weights["norm_out.linear.weight"])), self.weights["norm_out.linear.weight"], self.weights["norm_out.linear.bias"], ) @@ -1107,7 +1113,7 @@ class MLXMiniMaxH3DiT: adaln_indices = (timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags).astype(mx.int32) for block_index, block in enumerate(self.blocks): - tables = _adaln_tables(block, temb) + tables = _adaln_tables(block, temb, apply_silu=self.adaln_rank is None) packed = _transformer_block( block, packed, @@ -1124,7 +1130,8 @@ class MLXMiniMaxH3DiT: mx.eval(packed) # per-block sync: see forward_with_cache note shift_scale = linear( - silu(temb).astype(weight_dtype(self.weights["norm_out.linear.weight"])), + (silu(temb) if self.adaln_rank is None else temb).astype( + weight_dtype(self.weights["norm_out.linear.weight"])), self.weights["norm_out.linear.weight"], self.weights["norm_out.linear.bias"], ) @@ -1353,10 +1360,10 @@ def mlx_h3_dit_from_diffusers_safetensors( for shard in _safetensors_shards(transformer_path): shard_arrays = mx.load(str(shard)) for key, source in shard_arrays.items(): - if not key.startswith("time_embedder."): + if not (key.startswith("time_embedder.") or key == "adaln_basis.weight"): continue keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES - target_dtype = mx.float32 if keep_fp32 else cast_dtype + target_dtype = mx.float32 if keep_fp32 else (mx.float16 if key == "adaln_basis.weight" else cast_dtype) assign(key, _load_array(source, target_dtype)) del shard_arrays required_time_keys = { @@ -1373,6 +1380,10 @@ def mlx_h3_dit_from_diffusers_safetensors( weight_dtype(weights["time_embedder.linear_1.weight"])) temb = linear(t_freq, weights["time_embedder.linear_1.weight"], weights["time_embedder.linear_1.bias"]) temb = linear(silu(temb), weights["time_embedder.linear_2.weight"], weights["time_embedder.linear_2.bias"]) + if config.get("adaln_rank") is not None: + if "adaln_basis.weight" not in weights: + raise KeyError("Rank-reduced AdaLN checkpoint is missing adaln_basis.weight") + temb = linear(silu(temb).astype(weight_dtype(weights["adaln_basis.weight"])), weights["adaln_basis.weight"]) mx.eval(temb) cached_block_tables = [None] * num_blocks @@ -1381,7 +1392,10 @@ def mlx_h3_dit_from_diffusers_safetensors( for key, source in shard_arrays.items(): if _is_ignored_dense_key(key, include_vsa=include_vsa): continue - if temb is not None and key.startswith("time_embedder."): + if source.dtype == mx.uint8 and key.endswith(".weight"): + raise ValueError(f"Packed transformer weight {key} is not a floating-point source. " + "Convert the released BF16 transformer to MLX INT6 instead.") + if temb is not None and (key.startswith("time_embedder.") or key == "adaln_basis.weight"): continue if key.startswith("transformer_blocks."): index = int(key.split(".")[1]) @@ -1390,7 +1404,9 @@ def mlx_h3_dit_from_diffusers_safetensors( if key.startswith("rope."): continue # non-persistent analytic buffer, rebuilt on the fly keep_fp32 = key.split(".", 1)[0] in FP32_MODULE_PREFIXES - target_dtype = mx.float32 if keep_fp32 else cast_dtype + factorized_adaln = config.get("adaln_rank") is not None and (".adaln_proj." in key or key.startswith( + ("norm_out.linear.", "adaln_basis."))) + target_dtype = mx.float32 if keep_fp32 else (mx.float16 if factorized_adaln else cast_dtype) array = _load_array(source, target_dtype) if temb is not None and ".adaln_proj.linear." in key: _, index_str, sub = key.split(".", 2) @@ -1398,7 +1414,7 @@ def mlx_h3_dit_from_diffusers_safetensors( block_pending = pending_adaln.setdefault(index, {}) block_pending[sub] = array if {"adaln_proj.linear.weight", "adaln_proj.linear.bias"} <= block_pending.keys(): - tables = _adaln_tables(block_pending, temb) + tables = _adaln_tables(block_pending, temb, apply_silu=config.get("adaln_rank") is None) mx.eval(tables) assert cached_block_tables is not None cached_block_tables[index] = tables @@ -1428,7 +1444,8 @@ def mlx_h3_dit_from_diffusers_safetensors( if missing_cache_blocks: raise KeyError(f"Missing AdaLN cache tables for blocks {missing_cache_blocks}") shift_scale = linear( - silu(temb).astype(weight_dtype(weights["norm_out.linear.weight"])), + (silu(temb) if config.get("adaln_rank") is None else temb).astype( + weight_dtype(weights["norm_out.linear.weight"])), weights["norm_out.linear.weight"], weights["norm_out.linear.bias"], ) diff --git a/fastvideo/mlx_runtime/minimax_h3_conditioner.py b/fastvideo/mlx_runtime/minimax_h3_conditioner.py index 464250a9c..12dcd8554 100644 --- a/fastvideo/mlx_runtime/minimax_h3_conditioner.py +++ b/fastvideo/mlx_runtime/minimax_h3_conditioner.py @@ -20,8 +20,12 @@ from __future__ import annotations import gc import json +import os +import shutil +import tempfile from dataclasses import dataclass from pathlib import Path +from typing import Any import numpy as np import mlx.core as mx @@ -77,9 +81,19 @@ class _ShardIndex: index_path = component_dir / "model.safetensors.index.json" self.key_to_shard: dict[str, str] = {} self._header_cache: dict[str, tuple[dict, int]] = {} + + def needed(key: str) -> bool: + if key == "model.language_model.embed_tokens.weight": + return True + prefix = "model.language_model.layers." + if not key.startswith(prefix): + return False + layer = key[len(prefix):].split(".", 1)[0] + return layer.isdigit() and int(layer) < TEXT_ENCODER_LAYER + if index_path.exists(): weight_map = json.loads(index_path.read_text())["weight_map"] - self.key_to_shard = {k: str(component_dir / s) for k, s in weight_map.items()} + self.key_to_shard = {k: str(component_dir / s) for k, s in weight_map.items() if needed(k)} else: single = component_dir / "model.safetensors" if not single.exists(): @@ -88,7 +102,7 @@ class _ShardIndex: with open(single, "rb") as handle: (header_len, ) = struct.unpack(" np.ndarray: @@ -205,8 +219,20 @@ def _rms_norm(x, weight, eps: float): return x / mx.sqrt(mx.mean(x * x, axis=-1, keepdims=True) + eps) * weight +@dataclass(frozen=True) +class NVFP4Matrix: + """MLX row-major E2M1/E4M3 weights and the export's inverse global scale.""" + + weight: mx.array + scales: mx.array + global_scale: float + + def matmul(self, x): + return mx.quantized_matmul(x, self.weight, self.scales, mode="nvfp4") / self.global_scale + + def _linear(x, weight, bias=None): - y = x @ weight.T + y = weight.matmul(x) if isinstance(weight, NVFP4Matrix) else x @ weight.T if bias is not None: y = y + bias return y @@ -238,7 +264,7 @@ class StreamedMiniMaxH3TextConditioner: def __init__(self, component_dir: str | Path, tokenizer_dir: str | Path | None = None): self.component_dir = Path(component_dir) self.config = ConditionerConfig.from_config_json(self.component_dir / "config.json") - self.index = _ShardIndex(self.component_dir) + self.index: Any = _ShardIndex(self.component_dir) self.tokenizer = self._load_tokenizer(tokenizer_dir) def _load_tokenizer(self, tokenizer_dir: str | Path | None): @@ -284,8 +310,22 @@ class StreamedMiniMaxH3TextConditioner: ]) cos, sin = _mrope_cos_sin(positions, cfg) - # Embedding rows gathered individually; the (151936, 5120) table is - # never fully materialized. + hidden = self._embed_tokens(token_ids) + + if cfg.num_layers < TEXT_ENCODER_LAYER: + raise ValueError(f"Conditioner needs at least {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.") + for layer in range(TEXT_ENCODER_LAYER): + hidden = self._decoder_layer(layer, hidden, cos, sin) + # Per-layer sync keeps the 50-layer activation graph bounded. + mx.eval(hidden) + gc.collect() + + tags = np.full((seq_len, ), 1, dtype=np.int64) # MINIMAX_H3_TEXT_TAG + return np.asarray(hidden).astype(np.float32), tags + + def _embed_tokens(self, token_ids: list[int]): + # Embedding rows gathered individually; the full table is never + # materialized by the BF16 streaming path. rows = [] for token in token_ids: key = "model.language_model.embed_tokens.weight" @@ -293,18 +333,7 @@ class StreamedMiniMaxH3TextConditioner: hidden = mx.array(np.stack(rows).astype(np.float32)) del rows gc.collect() - - if cfg.num_layers <= TEXT_ENCODER_LAYER: - raise ValueError(f"Conditioner needs > {TEXT_ENCODER_LAYER} layers, has {cfg.num_layers}.") - for layer in range(TEXT_ENCODER_LAYER): - hidden = self._decoder_layer(layer, hidden, cos, sin) - # Per-layer sync: without this the whole 50-layer graph accumulates - # and the machine runs out of memory (same failure mode as the DiT). - mx.eval(hidden) - gc.collect() - - tags = np.full((seq_len, ), 1, dtype=np.int64) # MINIMAX_H3_TEXT_TAG - return np.asarray(hidden).astype(np.float32), tags + return hidden # -- layers ---------------------------------------------------------- @@ -368,6 +397,185 @@ class StreamedMiniMaxH3TextConditioner: self.index.close() +def unswizzle_nvfp4_scales(scale: np.ndarray, rows: int, cols: int) -> np.ndarray: + """FlashInfer 128x4 scale bytes -> MLX row-major group-16 scale bytes.""" + pad_rows, pad_cols = -(-rows // 128) * 128, -(-cols // 4) * 4 + if scale.size < pad_rows * pad_cols: + raise ValueError(f"NVFP4 scales need {pad_rows * pad_cols} bytes, got {scale.size}.") + tiles = scale.reshape(-1)[:pad_rows * pad_cols].reshape(pad_rows // 128, pad_cols // 4, 32, 4, 4) + return np.ascontiguousarray(tiles.transpose(0, 3, 2, 1, 4).reshape(pad_rows, pad_cols)[:rows, :cols]) + + +MLX_NVFP4_ENCODER_MANIFEST = "mlx_h3_nvfp4_encoder.json" + + +def _read_nvfp4_encoder_config(component_dir: str | Path) -> dict[str, Any]: + raw = json.loads((Path(component_dir) / "config.json").read_text()) + # Same accepted spellings as the CUDA loader (MiniMaxH3SerializedNVFP4Config.from_config). + accepted = { + "quant_method": ("nvfp4", ), + "fmt": ("e2m1", "float4_e2m1fn", "nvfp4"), + "scale_fmt": ("e4m3", "float8_e4m3fn"), + "activation_scheme": ("dynamic", ), + } + quant = raw.get("quantization_config", {}) + group_size = quant.get("group_size") + if (any(str(quant.get(key, "")).lower() not in values for key, values in accepted.items()) + or isinstance(group_size, bool) or group_size != 16 or quant.get("scale_layout") != "128x4"): + raise ValueError("MLX NVFP4 conditioning requires the FastVideo group-16, 128x4 encoder export.") + return raw + + +def export_mlx_h3_nvfp4_encoder(component_dir: str | Path, output_dir: str | Path) -> Path: + """Cache the released packed encoder in MLX layout, without requantization. + + Keep original packed nibbles, row-major scale bytes, global scales and + embedding/norm values. Later loads skip CPU scale unswizzling and staging. + """ + raw = _read_nvfp4_encoder_config(component_dir) + final_dir = Path(output_dir) + if final_dir.exists() and any(final_dir.iterdir()): + raise FileExistsError(f"Encoder cache output must be empty: {final_dir}") + # Build in a sibling staging directory and publish it with one rename, so an + # out-of-memory or out-of-disk failure never leaves a partial cache behind. + final_dir.parent.mkdir(parents=True, exist_ok=True) + output_dir = Path(tempfile.mkdtemp(prefix=f".{final_dir.name}.partial-", dir=final_dir.parent)) + index = _ResidentNVFP4Index(_ShardIndex(Path(component_dir))) + published = False + try: + arrays = {} + matrices = {} + dense_keys = [] + for key, value in index.weights.items(): + if isinstance(value, NVFP4Matrix): + arrays[key] = value.weight + arrays[key + ".scales"] = value.scales + matrices[key] = {"global_scale": value.global_scale} + else: + arrays[key] = value + dense_keys.append(key) + mx.save_safetensors(str(output_dir / "model.safetensors"), arrays) + (output_dir / "config.json").write_text(json.dumps(raw, indent=2) + "\n") + manifest = { + "format_version": 1, + "language_layers": TEXT_ENCODER_LAYER, + "matrices": matrices, + "dense_keys": dense_keys, + "source_dir": str(Path(component_dir).resolve()) + } + (output_dir / MLX_NVFP4_ENCODER_MANIFEST).write_text(json.dumps(manifest, indent=2) + "\n") + umask = os.umask(0) + os.umask(umask) + output_dir.chmod(0o777 & ~umask) # mkdtemp is owner-only; shared model dirs need the umask mode + if final_dir.exists(): + final_dir.rmdir() # empty, checked above + output_dir.rename(final_dir) + published = True + finally: + index.close() + if not published: + shutil.rmtree(output_dir, ignore_errors=True) + return final_dir + + +class _ResidentNVFP4Index: + """Load the released 50-layer encoder without expanding packed matrices.""" + + def __init__(self, source: _ShardIndex): + self.weights: dict[str, mx.array | NVFP4Matrix] = {} + for key in sorted(source.key_to_shard): + if key.endswith(".weight_packed"): + prefix = key.removesuffix(".weight_packed") + packed = np.array(source.get(key), copy=True) + if packed.dtype != np.uint8 or packed.ndim != 2 or packed.shape[1] % 4: + raise ValueError(f"Invalid packed NVFP4 matrix {key}: {packed.shape}, {packed.dtype}") + rows, cols = packed.shape[0], packed.shape[1] * 2 + if cols % 16: + raise ValueError(f"NVFP4 input width must be divisible by 16: {key}") + scales = unswizzle_nvfp4_scales(source.get(prefix + ".weight_scale"), rows, cols // 16) + global_scale = float(source.get(prefix + ".weight_global_scale").reshape(-1)[0]) + if not np.isfinite(global_scale) or global_scale <= 0: + raise ValueError(f"Invalid NVFP4 global scale for {prefix}: {global_scale}") + weight = mx.array(packed).view(mx.uint32) + scale_bytes = mx.array(scales) + if not bool(mx.all(mx.isfinite(mx.from_fp8(scale_bytes, dtype=mx.float32)))): + raise ValueError(f"Non-finite NVFP4 block scales for {prefix}") + mx.eval(weight, scale_bytes) + self.weights[prefix + ".weight"] = NVFP4Matrix(weight, scale_bytes, global_scale) + elif key.endswith(".weight"): + if key == "model.language_model.embed_tokens.weight": + shard = source.key_to_shard[key] + header, data_start = source._cache_header(shard) + if header[key]["dtype"] == "BF16": + value = mx.array(_read_bf16_words(shard, key, header, data_start)).view(mx.bfloat16) + else: + value = mx.array(source.get(key)) + else: + value = source.get_mlx(key) + mx.eval(value) + self.weights[key] = value + source.close() + + @classmethod + def from_mlx_checkpoint(cls, component_dir: str | Path): + component_dir = Path(component_dir) + manifest = json.loads((component_dir / MLX_NVFP4_ENCODER_MANIFEST).read_text()) + if manifest.get("format_version") != 1 or manifest.get("language_layers") != TEXT_ENCODER_LAYER: + raise ValueError("Unsupported native MLX NVFP4 encoder cache") + arrays = mx.load(str(component_dir / "model.safetensors")) + matrices = manifest["matrices"] + expected = set(manifest["dense_keys"]) | set(matrices) | {key + ".scales" for key in matrices} + if set(arrays) != expected: + raise ValueError("Native MLX encoder arrays do not match the manifest") + index = cls.__new__(cls) + index.weights = {key: arrays[key] for key in manifest["dense_keys"]} + for key, info in matrices.items(): + weight, scales = arrays[key], arrays[key + ".scales"] + factor = float(info["global_scale"]) + if (weight.ndim != 2 or weight.dtype != mx.uint32 or scales.dtype != mx.uint8 + or scales.shape != (weight.shape[0], weight.shape[1] // 2) or weight.shape[1] % 2 + or not np.isfinite(factor) or factor <= 0): + raise ValueError(f"Invalid native MLX NVFP4 matrix: {key}") + index.weights[key] = NVFP4Matrix(weight, scales, factor) + mx.eval(list(arrays.values())) + return index + + def get_mlx(self, key: str): + return self.weights[key] + + def close(self) -> None: + self.weights.clear() + gc.collect() + + +class ResidentNVFP4MiniMaxH3TextConditioner(StreamedMiniMaxH3TextConditioner): + """Released NVFP4 encoder weights with floating-point MLX activations. + + The packed weights and embedding table stay resident. CUDA quantizes + activations to FP4; this path keeps FP32 activations, so hidden states are + not expected to be bit-exact with the CUDA encoder. + """ + + def __init__(self, component_dir: str | Path, tokenizer_dir: str | Path | None = None): + _read_nvfp4_encoder_config(component_dir) + # Fail on an older MLX before reading the encoder's large shards. + try: + packed, scales = mx.quantize(mx.ones((1, 64)), mode="nvfp4") + mx.eval(mx.quantized_matmul(mx.ones((1, 64)), packed, scales, mode="nvfp4")) + except (ValueError, RuntimeError) as error: + raise RuntimeError("Native NVFP4 conditioning requires an MLX build with nvfp4 matmul support.") from error + super().__init__(component_dir, tokenizer_dir) + if (Path(component_dir) / MLX_NVFP4_ENCODER_MANIFEST).exists(): + self.index.close() + self.index = _ResidentNVFP4Index.from_mlx_checkpoint(component_dir) + else: + self.index = _ResidentNVFP4Index(self.index) + + def _embed_tokens(self, token_ids: list[int]): + table = self.index.get_mlx("model.language_model.embed_tokens.weight") + return table[mx.array(token_ids, dtype=mx.int32)].astype(mx.float32) + + def _apply_mrope(q_or_k, cos, sin): """q_or_k: (S, H, D); cos/sin: (S, 1, D).""" half = q_or_k.shape[-1] // 2 diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index c6ca112f3..0fc805d13 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -49,6 +49,7 @@ from fastvideo.mlx_runtime.minimax_h3 import ( MINIMAX_H3_MIN_DURATION, MINIMAX_H3_VIDEO_SHIFT, MiniMaxH3SchedulerState, + _eval_value, adaln_timestep_union, align_num_frames, audio_latent_num_frames, @@ -64,6 +65,7 @@ from fastvideo.mlx_runtime.minimax_h3 import ( video_latent_num_frames, ) from fastvideo.mlx_runtime.minimax_h3_vsa import MiniMaxH3VSAConfig +from fastvideo.mlx_runtime.prompt_cache import fingerprint_digest, text_encoder_fingerprint logger = init_logger(__name__) @@ -231,7 +233,7 @@ def _cleanup_mlx() -> None: def _default_metal_wired_limit_gib(mx) -> float: - """Keep the default below both physical memory and the tested 30 GiB cap.""" + """Legacy helper for allocator capacity, not the wired-residency setting.""" metal = getattr(mx, "metal", None) if metal is None: return 30.0 @@ -244,6 +246,64 @@ def _default_metal_wired_limit_gib(mx) -> float: return min(30.0, 0.84 * total_bytes / 2**30) +def _configure_metal_memory_limits(mx, + wired_limit_gib: float | None, + *, + resident: bool = False) -> tuple[int | None, int | None]: + """Set Metal limits and return their previous process-wide values.""" + if wired_limit_gib is not None and (not math.isfinite(wired_limit_gib) or wired_limit_gib <= 0): + raise ValueError("metal_wired_limit_gib must be finite and positive") + set_wired = getattr(mx, "set_wired_limit", None) + if set_wired is None and hasattr(mx, "metal"): + set_wired = getattr(mx.metal, "set_wired_limit", None) + if wired_limit_gib is not None and set_wired is None: + raise RuntimeError("This MLX build cannot set the requested wired-memory limit") + + previous_memory = None + set_memory = getattr(mx, "set_memory_limit", None) + if set_memory is None and hasattr(mx, "metal"): + set_memory = getattr(mx.metal, "set_memory_limit", None) + if set_memory is not None and not resident: + # An explicit wired request must fit under the allocator cap, or it pins memory MLX cannot allocate. + memory_limit_gib = max(_default_metal_wired_limit_gib(mx), wired_limit_gib or 0.0) + try: + previous_memory = int(set_memory(int(memory_limit_gib * 2**30))) + except Exception as error: # noqa: BLE001 - older MLX best effort + logger.info("Could not set the Metal allocation limit: %s", error) + if wired_limit_gib is None: + return previous_memory, None + # Explicit requests must succeed; do not silently benchmark an unwired model. + assert set_wired is not None + try: + previous = set_wired(int(wired_limit_gib * 2**30)) + except Exception: + if previous_memory is not None and set_memory is not None: + set_memory(previous_memory) + raise + logger.info("MLX wired limit %.2f GiB (previous %.2f GiB)", wired_limit_gib, previous / 2**30) + return previous_memory, int(previous) + + +def _restore_metal_wired_limit(mx, previous_bytes: int | None) -> None: + if previous_bytes is None: + return + set_wired = getattr(mx, "set_wired_limit", None) + if set_wired is None and hasattr(mx, "metal"): + set_wired = getattr(mx.metal, "set_wired_limit", None) + if set_wired is not None: + set_wired(previous_bytes) + + +def _restore_metal_memory_limit(mx, previous_bytes: int | None) -> None: + if previous_bytes is None: + return + set_memory = getattr(mx, "set_memory_limit", None) + if set_memory is None and hasattr(mx, "metal"): + set_memory = getattr(mx.metal, "set_memory_limit", None) + if set_memory is not None: + set_memory(previous_bytes) + + MINIMAX_H3_PROMPT_CACHE_VERSION = "v2-attention-layout" @@ -274,6 +334,27 @@ def _cached_adaln_timesteps(checkpoint_dir: str | Path) -> np.ndarray | None: return np.asarray(cache_info["timesteps"], dtype=np.float32) +def _load_resident_dit(checkpoint_dir: str | Path): + """Load an H3 DiT and materialize every weight before timed requests.""" + import mlx.core as mx + + dit = load_mlx_h3_checkpoint(checkpoint_dir) + for group in [dit.weights, *dit.blocks, *dit.refiner]: + for value in group.values(): + # With the converter's AdaLN cache, dropped AdaLN projection weights are None. + if value is not None: + _eval_value(value) + cache = dit._adaln_cache + if cache is not None: + mx.eval(cache.block_tables, cache.norm_out_shift, cache.norm_out_scale) + return dit + + +def _adaln_weights_dropped(dit: Any) -> bool: + key = "adaln_proj.linear.weight" + return any(key in block and block[key] is None for block in getattr(dit, "blocks", ())) + + def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int, *, @@ -315,6 +396,10 @@ def _preflight_media_dependencies(*, class MiniMaxH3MLXPipeline: """Text-to-video-with-audio generation through the native MLX runtime.""" + # Defaults for pipelines built without __init__ (unit tests use __new__). + resident = False + conditioner_mode = "auto" + def __init__( self, *, @@ -328,21 +413,19 @@ class MiniMaxH3MLXPipeline: video_decode_backend: str = "h3-vae", taeh3_checkpoint: str | Path | None = None, taeh3_chunk_size: int = 5, + conditioner_mode: str = "auto", + resident: bool = False, ) -> None: import mlx.core as mx - set_limit = getattr(mx, "set_memory_limit", None) - if set_limit is None and hasattr(mx, "metal"): - set_limit = getattr(mx.metal, "set_memory_limit", None) - if set_limit is not None: - # Keep large resident models inside a predictable wired budget. - try: - if metal_wired_limit_gib is None: - metal_wired_limit_gib = _default_metal_wired_limit_gib(mx) - set_limit(int(metal_wired_limit_gib * 2**30)) - except Exception as error: # noqa: BLE001 - best effort on older MLX - logger.info("Could not raise the Metal wired limit: %s", error) self.model_root = Path(model_root) + if conditioner_mode not in ("auto", "streamed", "nvfp4"): + raise ValueError(f"Unknown H3 conditioner mode: {conditioner_mode}") + if resident and video_decode_backend != "h3-vae": + raise ValueError("Resident H3 generation requires the H3 video VAE.") + self.conditioner_mode = conditioner_mode + self.resident = resident + self._resident_components: dict[str, Any] = {} self.dit_checkpoint = Path(mlx_dit_checkpoint) self.vae_dtype = vae_dtype if video_decode_backend not in ("h3-vae", "taeh3"): @@ -367,6 +450,9 @@ class MiniMaxH3MLXPipeline: self._dit_in_channels = int(dit_config["in_channels"]) self.last_dit_forward_s = 0.0 self.last_vsa_stats: dict[str, Any] | None = None + self._previous_memory_limit, self._previous_wired_limit = _configure_metal_memory_limits(mx, + metal_wired_limit_gib, + resident=resident) # -- input validation (before anything heavy loads) ------------------- @@ -414,20 +500,70 @@ class MiniMaxH3MLXPipeline: # -- phase 1: conditioning ------------------------------------------- + def prepare_resident(self) -> None: + """Load the encoder, DiT, and both decoders once, before timed requests.""" + if not self.resident or self._resident_components: + return + from fastvideo.mlx_runtime.minimax_h3_conditioner import ResidentNVFP4MiniMaxH3TextConditioner + from fastvideo.mlx_runtime.minimax_h3_audio_vae import mlx_h3_audio_vae_from_dir + from fastvideo.mlx_runtime.minimax_h3_video_vae import mlx_h3_video_vae_from_dir + + import mlx.core as mx + + try: + conditioner = self._load_conditioner() + if not isinstance(conditioner, ResidentNVFP4MiniMaxH3TextConditioner): + conditioner.close() + raise ValueError("All-resident generation requires the packed NVFP4 text encoder.") + self._resident_components["conditioner"] = conditioner + self._resident_components["dit"] = _load_resident_dit(self.dit_checkpoint) + self._resident_components["video_vae"] = mlx_h3_video_vae_from_dir(self.model_root / "vae", + include_encoder=False, + storage_dtype=self.vae_dtype) + self._resident_components["audio_vae"] = mlx_h3_audio_vae_from_dir(self.model_root / "audio_vae", + include_encoder=False) + mx.eval(list(self._resident_components["audio_vae"].weights.values())) + logger.info("H3 components resident: %.2f GiB active MLX memory", mx.get_active_memory() / 2**30) + except Exception: + self.close() + raise + + def close(self) -> None: + conditioner = self._resident_components.get("conditioner") + if conditioner is not None: + conditioner.close() + self._resident_components.clear() + _cleanup_mlx() + previous = getattr(self, "_previous_wired_limit", None) + if previous is not None: + import mlx.core as mx + + _restore_metal_wired_limit(mx, previous) + self._previous_wired_limit = None + previous_memory = getattr(self, "_previous_memory_limit", None) + if previous_memory is not None: + import mlx.core as mx + + _restore_metal_memory_limit(mx, previous_memory) + self._previous_memory_limit = None + def encode_prompt(self, prompt: str) -> tuple[np.ndarray, np.ndarray]: """Returns (hidden states (S, hidden), token tags). Uses the cache or the streamed conditioner.""" - cache_key = None - if self.prompt_cache_dir is not None: - cache_key = prompt_cache_path(self.prompt_cache_dir, self.model_root, prompt) - if cache_key.exists(): - data = np.load(cache_key) - logger.info("Loaded prompt embeddings from cache %s", cache_key) - return data["hidden_states"], data["token_tags"] + cache_key = self._prompt_cache_key(prompt) + if cache_key is not None and cache_key.exists(): + data = np.load(cache_key) + logger.info("Loaded prompt embeddings from cache %s", cache_key) + return data["hidden_states"], data["token_tags"] - conditioner = self._load_conditioner() + if self.resident: + self.prepare_resident() + conditioner = self._resident_components["conditioner"] + else: + conditioner = self._load_conditioner() hidden, tags = conditioner.encode_prompt(prompt) - conditioner.close() + if not self.resident: + conditioner.close() _cleanup_mlx() if cache_key is not None: cache_key.parent.mkdir(parents=True, exist_ok=True) @@ -440,6 +576,19 @@ class MiniMaxH3MLXPipeline: tmp_cache.unlink(missing_ok=True) return hidden, tags + def _prompt_cache_key(self, prompt: str) -> Path | None: + """Cache path bound to the effective encoder and its files; None skips the cache.""" + if self.prompt_cache_dir is None: + return None + fingerprint = { + "conditioner": self._conditioner_kind(), + "text_encoder": text_encoder_fingerprint(self.conditioner_dir), + "tokenizer": text_encoder_fingerprint(self.tokenizer_dir), + } + if not (fingerprint["text_encoder"]["complete"] and fingerprint["tokenizer"]["complete"]): + return None + return prompt_cache_path(self.prompt_cache_dir, fingerprint_digest(fingerprint), prompt) + def load_prompt_cache(self, path: str | Path) -> tuple[np.ndarray, np.ndarray]: data = np.load(path) return data["hidden_states"], data["token_tags"] @@ -450,10 +599,25 @@ class MiniMaxH3MLXPipeline: return marker.exists() or single.exists() def _load_conditioner(self): - from fastvideo.mlx_runtime.minimax_h3_conditioner import StreamedMiniMaxH3TextConditioner + from fastvideo.mlx_runtime.minimax_h3_conditioner import ( + ResidentNVFP4MiniMaxH3TextConditioner, + StreamedMiniMaxH3TextConditioner, + ) + if self._conditioner_kind() == "nvfp4": + return ResidentNVFP4MiniMaxH3TextConditioner(self.conditioner_dir, self.tokenizer_dir) return StreamedMiniMaxH3TextConditioner(self.conditioner_dir, self.tokenizer_dir) + def _conditioner_kind(self) -> str: + """The encoder that conditioner_mode selects for these weights: 'nvfp4' or 'streamed'.""" + config = json.loads((self.conditioner_dir / "config.json").read_text()) + packed = str(config.get("quantization_config", {}).get("quant_method", "")).lower() == "nvfp4" + if self.conditioner_mode == "nvfp4" or (self.conditioner_mode == "auto" and packed): + return "nvfp4" + if packed: + raise ValueError("The streamed conditioner requires BF16 weights; use conditioner_mode='nvfp4'.") + return "streamed" + # -- phase 2: denoise -------------------------------------------------- def denoise( @@ -478,12 +642,46 @@ class MiniMaxH3MLXPipeline: geometry = self.resolve_geometry(height, width, num_frames, enforce_duration=audio_num_frames is None) audio_frames = geometry["num_frames"] if audio_num_frames is None else align_num_frames(audio_num_frames) + resident_dit = dit is None and self.resident + if resident_dit: + _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps, model_root=self.model_root) + self.prepare_resident() + dit = self._resident_components["dit"] owned_dit = dit is None if owned_dit: _validate_checkpoint_step_ladder(self.dit_checkpoint, num_steps, model_root=self.model_root) t0 = time.perf_counter() dit = load_mlx_h3_checkpoint(self.dit_checkpoint) logger.info("Loaded MLX H3 DiT from %s in %.1fs", self.dit_checkpoint, time.perf_counter() - t0) + + schedule = resolve_h3_denoise_schedule(self.model_root, num_steps) + video_scheduler = schedule.video + audio_scheduler = schedule.audio + union = schedule.adaln_timesteps + num_steps = schedule.num_steps + # The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA + # conditioning rows; those modes recompute the ladder before denoise. + + cache = getattr(dit, "_adaln_cache", None) + if cache is None or not np.array_equal(cache.timesteps.astype(np.float32), union): + if cache is not None: + extra = np.setdiff1d(union, cache.timesteps) + logger.info("Recomputing AdaLN cache for %d-step ladder (extra timesteps %s).", num_steps, extra) + if _adaln_weights_dropped(dit): + # An earlier request on this DiT already released the AdaLN projections. + if not resident_dit: + raise ValueError("This H3 DiT dropped its AdaLN weights for another step ladder; " + "pass a freshly loaded DiT for a different num_steps.") + logger.info("Reloading the resident H3 DiT for the %d-step ladder.", num_steps) + cache = dit = None + self._resident_components.pop("dit", None) + _cleanup_mlx() + dit = self._resident_components["dit"] = _load_resident_dit(self.dit_checkpoint) + dit.precompute_adaln(union, drop_weights=True) + + if resident_dit: + # The resident DiT keeps the previous request's VSA mode unless reset to the default. + vsa_config = vsa_config or MiniMaxH3VSAConfig() if vsa_config is not None: dit.configure_vsa(vsa_config) if hasattr(dit, "reset_vsa_stats"): @@ -502,22 +700,6 @@ class MiniMaxH3MLXPipeline: if getattr(dit, "vsa_config", None) is not None and dit.vsa_config.enabled: dit.prepare_vsa_geometry(layout) - schedule = resolve_h3_denoise_schedule(self.model_root, num_steps) - video_scheduler = schedule.video - audio_scheduler = schedule.audio - union = schedule.adaln_timesteps - num_steps = schedule.num_steps - # The keyframe-noise timestep (0.999) is only exercised by FL2VA/Ref2VA - # conditioning rows; those modes recompute the ladder before denoise. - - cache = getattr(dit, "_adaln_cache", None) - if cache is None: - dit.precompute_adaln(union, drop_weights=True) - elif not np.array_equal(cache.timesteps.astype(np.float32), union): - extra = np.setdiff1d(union, cache.timesteps) - logger.info("Recomputing AdaLN cache for %d-step ladder (extra timesteps %s).", num_steps, extra) - dit.precompute_adaln(union, drop_weights=True) - video_key, audio_key = mx.random.split(mx.random.key(seed)) target_video_rows = int(layout.video_indices.shape[0] - layout.num_condition_video_rows) target_audio_rows = int(layout.audio_indices.shape[0] - layout.num_condition_audio_rows) @@ -624,7 +806,13 @@ class MiniMaxH3MLXPipeline: raise RuntimeError(f"TAEH3 produced unexpected frame shape: {frames.shape}") _cleanup_mlx() return frames - vae = mlx_h3_video_vae_from_dir(self.model_root / "vae", include_encoder=False, storage_dtype=self.vae_dtype) + if self.resident: + self.prepare_resident() + vae = self._resident_components["video_vae"] + else: + vae = mlx_h3_video_vae_from_dir(self.model_root / "vae", + include_encoder=False, + storage_dtype=self.vae_dtype) expected_height = height // vae.spatial_compression_ratio expected_width = width // vae.spatial_compression_ratio if (geometry["latent_height"], geometry["latent_width"]) != (expected_height, expected_width): @@ -666,7 +854,11 @@ class MiniMaxH3MLXPipeline: num_audio_latents = audio_latent_num_frames(align_num_frames(num_frames)) latents = unpack_audio_tokens(audio_rows, num_audio_latents) - vae = mlx_h3_audio_vae_from_dir(self.model_root / "audio_vae", include_encoder=False) + if self.resident: + self.prepare_resident() + vae = self._resident_components["audio_vae"] + else: + vae = mlx_h3_audio_vae_from_dir(self.model_root / "audio_vae", include_encoder=False) z = vae.denormalize_latents(mx.array(latents)) waveform = np.asarray(vae.decode(z))[:, 0, :] # (B, 1, S) -> (B, S) del vae, z diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa.py b/fastvideo/mlx_runtime/minimax_h3_vsa.py index 3c49969e5..833bd139f 100644 --- a/fastvideo/mlx_runtime/minimax_h3_vsa.py +++ b/fastvideo/mlx_runtime/minimax_h3_vsa.py @@ -29,7 +29,7 @@ from __future__ import annotations import math from dataclasses import dataclass, field -from functools import cached_property +from functools import cached_property, lru_cache from numbers import Integral from typing import Any, Literal @@ -544,13 +544,27 @@ def _key_valid_mask(block_idx, variable_block_sizes, tile_elems: int): return offsets[None, None, None, :] < selected_sizes[:, :, :, None] -_REFERENCE_GATHER_TARGET_BYTES = 2 * 1024**3 +_REFERENCE_GATHER_MIN_BYTES = 256 * 1024**2 +_REFERENCE_GATHER_MAX_BYTES = 2 * 1024**3 + + +@lru_cache(maxsize=1) +def _reference_gather_target_bytes() -> int: + """1/128 of unified memory in 256 MiB-2 GiB: 32 GB Macs keep room for resident weights.""" + import mlx.core as mx + + device_info = getattr(mx, "device_info", None) or mx.metal.device_info + try: + total = int(device_info().get("memory_size", 0)) + except (AttributeError, RuntimeError, TypeError, ValueError): + total = 0 + return min(_REFERENCE_GATHER_MAX_BYTES, max(_REFERENCE_GATHER_MIN_BYTES, total // 128)) def _reference_gather_query_chunk(heads: int, dim: int, k_sel: int, tile_elems: int, n_q: int) -> int: - """Batch as many query tiles as fit in ~2 GiB of gathered BF16 K/V.""" + """Bound gathered BF16 K/V by the device's gather budget.""" bytes_per_query = 4 * heads * max(k_sel, 1) * tile_elems * dim - chunk = min(n_q, max(1, _REFERENCE_GATHER_TARGET_BYTES // max(bytes_per_query, 1))) + chunk = min(n_q, max(1, _reference_gather_target_bytes() // max(bytes_per_query, 1))) return int(chunk) diff --git a/fastvideo/mlx_runtime/minimax_h3_vsa_simd.py b/fastvideo/mlx_runtime/minimax_h3_vsa_simd.py index bb21b1912..910384ac9 100644 --- a/fastvideo/mlx_runtime/minimax_h3_vsa_simd.py +++ b/fastvideo/mlx_runtime/minimax_h3_vsa_simd.py @@ -242,6 +242,15 @@ def _simd_kernel() -> Any | None: if not mx.metal.is_available(): _SIMD_KERNEL_ERROR = "Metal is not available in this MLX build" return None + # The SIMD kernel compiles but fails BF16 parity on the 7 GiB virtual + # Metal device in GitHub's M1 runner. Physical M1 devices also use the + # reference path until this kernel is validated on that GPU family. + device_info = mx.metal.device_info() + device_name = str(device_info.get("device_name", "")) + memory_size = int(device_info.get("memory_size", 0)) + if device_name.startswith("Apple M1") or 0 < memory_size <= 8 * 2**30: + _SIMD_KERNEL_ERROR = "SIMD VSA BF16 parity is unvalidated on Apple M1 or 8 GiB Metal devices" + return None if not hasattr(mx.fast, "metal_kernel"): _SIMD_KERNEL_ERROR = "mx.fast.metal_kernel is not available in this MLX build" return None diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_bf16_source.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_bf16_source.py new file mode 100644 index 000000000..9c9646b7c --- /dev/null +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_bf16_source.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reject packed source bytes before they can become invalid affine weights.""" +import pytest + +mx = pytest.importorskip("mlx.core") +from fastvideo.mlx_runtime.minimax_h3 import mlx_h3_dit_from_diffusers_safetensors + + +def test_packed_transformer_source_is_rejected(tmp_path): + mx.save_safetensors(str(tmp_path / "diffusion_pytorch_model.safetensors"), { + "transformer_blocks.0.attn.to_q.weight": mx.zeros((128, 128), dtype=mx.uint8), + }) + with pytest.raises(ValueError, match="released BF16 transformer"): + mlx_h3_dit_from_diffusers_safetensors( + tmp_path, config={"num_layers": 1, "num_refiner_layers": 0}, + quantization="int6", include_vsa=True, + ) diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_rank16_adaln.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_rank16_adaln.py new file mode 100644 index 000000000..cd954981b --- /dev/null +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_rank16_adaln.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Check the pruned model's shared rank-16 modulation against NumPy math.""" +import numpy as np +import pytest + +mx = pytest.importorskip('mlx.core') +from fastvideo.mlx_runtime.fastwan import timestep_embedding +from fastvideo.mlx_runtime.minimax_h3 import ( + MLXMiniMaxH3DiT, + MiniMaxH3SchedulerState, + adaln_timestep_union, + load_mlx_h3_checkpoint, + save_mlx_h3_checkpoint, +) + + +def test_rank16_cache_has_one_silu_before_shared_basis(tmp_path): + rng = np.random.default_rng(2026) + hidden, rank = 32, 16 + + def array(shape): + return mx.array(rng.normal(0, 0.1, shape).astype(np.float32)) + + config = dict(hidden_size=hidden, num_attention_heads=1, attention_head_dim=hidden, ffn_dim=64, + in_channels=24, audio_in_channels=24, patch_size=[1, 1, 1], text_dim=hidden, + freq_dim=hidden, time_embed_dim=hidden, rope_freq_dim=4, rope_theta=10000., + norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5, adaln_rank=rank, num_layers=42) + weights = { + 'time_embedder.linear_1.weight': array((hidden, hidden)), + 'time_embedder.linear_1.bias': array((hidden,)), + 'time_embedder.linear_2.weight': array((hidden, hidden)), + 'time_embedder.linear_2.bias': array((hidden,)), + 'adaln_basis.weight': array((rank, hidden)), + 'norm_out.linear.weight': array((2 * hidden, rank)), + 'norm_out.linear.bias': array((2 * hidden,)), + } + blocks = [{'attn.to_q.weight': array((hidden, hidden)), + 'adaln_proj.linear.weight': array((18 * hidden, rank)), + 'adaln_proj.linear.bias': array((18 * hidden,))} for _ in range(42)] + dit = MLXMiniMaxH3DiT(weights, blocks, [], config) + rungs = [999, 874, 749, 624, 500, 375, 250, 125] + timesteps = adaln_timestep_union(MiniMaxH3SchedulerState.from_dmd_steps(10, rungs), + MiniMaxH3SchedulerState.from_dmd_steps(3, rungs)) + + def project(x, weight, bias=None): + result = x @ np.array(weight).T + return result if bias is None else result + np.array(bias) + + def silu(x): + return x / (1 + np.exp(-x)) + + features = np.array(timestep_embedding(mx.array(timesteps), hidden)) + first = project(features, weights['time_embedder.linear_1.weight'], weights['time_embedder.linear_1.bias']) + second = project(silu(first), weights['time_embedder.linear_2.weight'], weights['time_embedder.linear_2.bias']) + expected_basis = project(silu(second), weights['adaln_basis.weight']) + np.testing.assert_allclose(np.array(dit.compute_temb(mx.array(timesteps))), expected_basis, atol=1e-6) + expected_blocks = [project(expected_basis, block['adaln_proj.linear.weight'], + block['adaln_proj.linear.bias']).reshape(-1, 6 * hidden) for block in blocks] + expected_out = project(expected_basis, weights['norm_out.linear.weight'], weights['norm_out.linear.bias']) + cache = dit.precompute_adaln(timesteps) + for tables, expected in zip(cache.block_tables, expected_blocks, strict=True): + np.testing.assert_allclose(np.concatenate([np.array(t) for t in tables], axis=-1), expected, atol=1e-6) + np.testing.assert_allclose(np.array(cache.norm_out_shift), expected_out[:, :hidden], atol=1e-6) + np.testing.assert_allclose(np.array(cache.norm_out_scale), expected_out[:, hidden:], atol=1e-6) + assert all(block['adaln_proj.linear.weight'] is None for block in blocks) + + save_mlx_h3_checkpoint(dit, tmp_path) + loaded = load_mlx_h3_checkpoint(tmp_path) + assert loaded.adaln_rank == rank + assert len(loaded.blocks) == 42 + np.testing.assert_array_equal(loaded._adaln_cache.timesteps, timesteps) + np.testing.assert_array_equal(np.array(loaded._adaln_cache.norm_out_scale), np.array(cache.norm_out_scale)) diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_resident_nvfp4.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_resident_nvfp4.py new file mode 100644 index 000000000..288b3bf46 --- /dev/null +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_resident_nvfp4.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native MLX NVFP4 encoder storage and residency regression checks.""" +import json +from types import SimpleNamespace + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") +from fastvideo.mlx_runtime.minimax_h3_conditioner import ( + NVFP4Matrix, ResidentNVFP4MiniMaxH3TextConditioner, _ResidentNVFP4Index, + _ShardIndex, export_mlx_h3_nvfp4_encoder, unswizzle_nvfp4_scales, +) +from fastvideo.mlx_runtime.minimax_h3_pipeline import MiniMaxH3MLXPipeline + + +def _swizzle(values): + # Independent coordinate mapping of FlashInfer's 128-row, four-column tiles. + rows, cols = values.shape + padded = np.zeros((-(-rows // 128) * 128, -(-cols // 4) * 4), np.uint8) + padded[:rows, :cols] = values + output = np.empty(padded.size, np.uint8) + for r in range(padded.shape[0]): + for c in range(padded.shape[1]): + address = ((((r // 128) * (padded.shape[1] // 4) + c // 4) * 32 + + r % 32) * 4 + (r % 128) // 32) * 4 + c % 4 + output[address] = padded[r, c] + return output + + +def _decode_e4m3(values): + sign = np.where(values & 128, -1.0, 1.0) + exponent = (values >> 3) & 15 + fraction = values & 7 + return sign * np.where(exponent == 0, fraction * 2.0**-9, + (1.0 + fraction / 8.0) * 2.0**(exponent.astype(int) - 7)) + + +def test_padded_scale_layout_round_trip(): + rng = np.random.default_rng(11) + values = rng.integers(0, 127, (140, 7), dtype=np.uint8) + np.testing.assert_array_equal(unswizzle_nvfp4_scales(_swizzle(values), 140, 7), values) + with pytest.raises(ValueError, match="bytes"): + unswizzle_nvfp4_scales(np.zeros(1, np.uint8), 140, 7) + + +@pytest.mark.parametrize("global_scale", [0.5, 4.0]) +@pytest.mark.parametrize("native_cache", [False, True]) +def test_serialized_encoder_linear_matches_independent_fp4_reference(tmp_path, global_scale, native_cache): + from safetensors.numpy import save_file + + rng = np.random.default_rng(3) + packed = rng.integers(0, 256, (128, 64), dtype=np.uint8) + scales = rng.integers(24, 96, (128, 8), dtype=np.uint8) + prefix = "model.language_model.layers.0.self_attn.q_proj" + save_file({prefix + ".weight_packed": packed, + prefix + ".weight_scale": _swizzle(scales), + prefix + ".weight_global_scale": np.array([global_scale], np.float32)}, + tmp_path / "model.safetensors") + index = _ResidentNVFP4Index(_ShardIndex(tmp_path)) + if native_cache: + _write_encoder_config(tmp_path) + cache_dir = export_mlx_h3_nvfp4_encoder(tmp_path, tmp_path / "cache") + cached = _ResidentNVFP4Index.from_mlx_checkpoint(cache_dir) + for key, original in index.weights.items(): + value = cached.get_mlx(key) + np.testing.assert_array_equal(np.array(value.weight), np.array(original.weight)) + np.testing.assert_array_equal(np.array(value.scales), np.array(original.scales)) + assert value.global_scale == original.global_scale + index.close() + index = cached + weight = index.get_mlx(prefix + ".weight") + assert isinstance(weight, NVFP4Matrix) + assert weight.weight.dtype == mx.uint32 + assert weight.scales.dtype == mx.uint8 + lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, 0, -.5, -1, -1.5, -2, -3, -4, -6], np.float32) + dense = np.stack((lut[packed & 15], lut[packed >> 4]), axis=-1).reshape(128, 128) + dense *= np.repeat(_decode_e4m3(scales), 16, axis=1) / global_scale + x = rng.standard_normal((3, 128)).astype(np.float32) + np.testing.assert_allclose(np.array(weight.matmul(mx.array(x))), x @ dense.T, rtol=3e-5, atol=1e-3) + index.close() + assert not index.weights + + +def _write_encoder_config(path): + (path / "config.json").write_text(json.dumps({"quantization_config": { + "quant_method": "nvfp4", "fmt": "e2m1", "group_size": 16, + "scale_fmt": "e4m3", "scale_layout": "128x4", "activation_scheme": "dynamic", + }})) + + +@pytest.mark.parametrize("native_cache", [False, True]) +def test_resident_embedding_keeps_bf16_storage(tmp_path, native_cache): + torch = pytest.importorskip("torch") + from safetensors.torch import save_file + + key = "model.language_model.embed_tokens.weight" + table = torch.arange(60).reshape(10, 6).to(torch.bfloat16) + save_file({key: table}, tmp_path / "model.safetensors") + if native_cache: + _write_encoder_config(tmp_path) + cache_dir = export_mlx_h3_nvfp4_encoder(tmp_path, tmp_path / "cache") + index = _ResidentNVFP4Index.from_mlx_checkpoint(cache_dir) + with pytest.raises(FileExistsError, match="empty"): + export_mlx_h3_nvfp4_encoder(tmp_path, cache_dir) + else: + index = _ResidentNVFP4Index(_ShardIndex(tmp_path)) + assert index.get_mlx(key).dtype == mx.bfloat16 + conditioner = ResidentNVFP4MiniMaxH3TextConditioner.__new__(ResidentNVFP4MiniMaxH3TextConditioner) + conditioner.index = index + np.testing.assert_array_equal(np.array(conditioner._embed_tokens([7, 1])), table[[7, 1]].float().numpy()) + conditioner.close() + + +def _pipeline(): + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.resident = True + pipeline._resident_components = {} + pipeline.dit_checkpoint = "tiny" + pipeline.model_root = __import__("pathlib").Path("tiny") + pipeline.vae_dtype = "fp16" + return pipeline + + +def test_resident_preload_reuses_models_and_evaluates_audio(monkeypatch): + import fastvideo.mlx_runtime.minimax_h3_pipeline as module + import fastvideo.mlx_runtime.minimax_h3_audio_vae as audio + import fastvideo.mlx_runtime.minimax_h3_video_vae as video + + pipeline = _pipeline() + conditioner = ResidentNVFP4MiniMaxH3TextConditioner.__new__(ResidentNVFP4MiniMaxH3TextConditioner) + conditioner.index = SimpleNamespace(close=lambda: None) + monkeypatch.setattr(pipeline, "_load_conditioner", lambda: conditioner) + calls = [] + def load_dit(path): + calls.append(path) + return SimpleNamespace(weights={"x": mx.ones((1,))}, blocks=[], refiner=[], _adaln_cache=None) + monkeypatch.setattr(module, "load_mlx_h3_checkpoint", load_dit) + monkeypatch.setattr(video, "mlx_h3_video_vae_from_dir", lambda *a, **k: object()) + decoder = SimpleNamespace(weights={"x": mx.ones((2,))}) + monkeypatch.setattr(audio, "mlx_h3_audio_vae_from_dir", lambda *a, **k: decoder) + pipeline.prepare_resident() + pipeline.prepare_resident() + assert calls == ["tiny"] + assert set(pipeline._resident_components) == {"conditioner", "dit", "video_vae", "audio_vae"} + pipeline.close() + assert not pipeline._resident_components + + +def test_failed_preload_releases_encoder(monkeypatch): + import fastvideo.mlx_runtime.minimax_h3_pipeline as module + + pipeline = _pipeline() + conditioner = ResidentNVFP4MiniMaxH3TextConditioner.__new__(ResidentNVFP4MiniMaxH3TextConditioner) + closed = [] + conditioner.index = SimpleNamespace(close=lambda: closed.append(True)) + monkeypatch.setattr(pipeline, "_load_conditioner", lambda: conditioner) + def fail(path): + raise RuntimeError("out of memory") + monkeypatch.setattr(module, "load_mlx_h3_checkpoint", fail) + with pytest.raises(RuntimeError, match="out of memory"): + pipeline.prepare_resident() + assert closed == [True] + assert not pipeline._resident_components + + +def test_single_shard_omits_unused_language_layers_and_vision(tmp_path): + from safetensors.numpy import save_file + + kept = "model.language_model.layers.49.input_layernorm.weight" + dropped = "model.language_model.layers.50.input_layernorm.weight" + save_file({kept: np.ones(8, np.float32), dropped: np.ones(8, np.float32), + "model.visual.weight": np.ones((8, 8), np.float32)}, tmp_path / "model.safetensors") + index = _ResidentNVFP4Index(_ShardIndex(tmp_path)) + assert set(index.weights) == {kept} + index.close() diff --git a/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py b/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py new file mode 100644 index 000000000..b35bd5942 --- /dev/null +++ b/fastvideo/mlx_runtime/tests/test_minimax_h3_vsa_gather_budget.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Changing the gather memory budget must preserve selected-tile attention.""" +import numpy as np +import pytest + +mx = pytest.importorskip('mlx.core') +import fastvideo.mlx_runtime.minimax_h3_vsa as vsa + + +def test_gather_budget_preserves_attention_with_partial_tiles(monkeypatch): + geom = vsa.build_h3_tile_geometry((7, 67), (6, 8, 8), 64) + rng = np.random.default_rng(2026) + shape = (geom.padded_length, 2, 128) + q, k, value = [mx.array(rng.normal(size=shape).astype(np.float32)).astype(mx.bfloat16) for _ in range(3)] + # Reverse video order also exercises the selected-key order, not a dense mask. + selected = np.array([0, 1, geom.num_tiles - 1, geom.num_prefix_tiles], dtype=np.int32) + idx = mx.array(np.broadcast_to(selected, (2, geom.num_video_tiles, selected.size)).copy()) + monkeypatch.setattr(vsa, '_reference_gather_target_bytes', lambda: 2 * 1024**3) + expected = vsa._reference_gather_sdpa(q, k, value, idx, geom, 128**-0.5) + mx.eval(expected) + monkeypatch.setattr(vsa, '_reference_gather_target_bytes', lambda: 1) + actual = vsa._reference_gather_sdpa(q, k, value, idx, geom, 128**-0.5) + mx.eval(actual) + np.testing.assert_array_equal(np.array(actual.astype(mx.float32)), np.array(expected.astype(mx.float32))) diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index aa086a7d8..71d9dd013 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -195,6 +195,10 @@ class PipelineStage(ABC): torch.cuda.max_memory_allocated() / gib, torch.cuda.memory_reserved() / gib, torch.cuda.memory_allocated() / gib) + batch.logging_info.add_stage_metric(stage_key, "peak_allocated_mb", + torch.cuda.max_memory_allocated() / 1024**2) + batch.logging_info.add_stage_metric(stage_key, "peak_reserved_mb", + torch.cuda.max_memory_reserved() / 1024**2) torch.cuda.reset_peak_memory_stats() batch.logging_info.add_stage_execution_time(stage_key, execution_time) batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name) diff --git a/fastvideo/tests/inference/test_benchmark_fasth3_spark_nvfp4.py b/fastvideo/tests/inference/test_benchmark_fasth3_spark_nvfp4.py new file mode 100644 index 000000000..5fd58c056 --- /dev/null +++ b/fastvideo/tests/inference/test_benchmark_fasth3_spark_nvfp4.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of examples/inference/basic/benchmark_fasth3_spark_nvfp4.py stage accounting.""" +from __future__ import annotations + +import importlib.util +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "benchmark_fasth3_spark_nvfp4.py" + + +def _load_benchmark(): + spec = importlib.util.spec_from_file_location("benchmark_fasth3_spark_nvfp4", EXAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_decode_seconds_exclude_post_decode_frame_processing() -> None: + benchmark = _load_benchmark() + stages = { + "MiniMaxH3DenoisingStage": 9.0, + "MiniMaxH3VideoDecodingStage": 2.0, + "MiniMaxH3AudioDecodingStage": 0.5, + "PostDecodeFrameProcessStage": 1.0, + } + assert benchmark._stage_total(stages, "decod", exclude="postdecode") == 2.5 + assert benchmark._stage_total(stages, "postdecode") == 1.0 + assert benchmark._stage_total(stages, "denois") == 9.0 diff --git a/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py b/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py index 9c8d3813a..81fade2f2 100644 --- a/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py +++ b/fastvideo/tests/mlx/test_mlx_affine_dq_gemm.py @@ -73,7 +73,11 @@ def test_dq_gemm_matches_qmm_for_supported_bit_widths(bits: int, group_size: int ref_np = np.asarray(ref.astype(mx.float32)) got_np = np.asarray(got.astype(mx.float32)) scale = max(float(np.max(np.abs(ref_np))), 1e-3) - assert rel < 2e-2, rel + # MLX's CPU backend accumulates the BF16 dense GEMM less precisely than Metal; at group size 128 + # (512 inputs) its error against qmm reaches about 2.6% for every bit width, including 8, so the + # gap is accumulation rather than quantization. + rel_limit = 3e-2 if mx.default_device() == mx.cpu else 2e-2 + assert rel < rel_limit, rel assert float(np.max(np.abs(got_np - ref_np))) / scale < 0.08 diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py index 4fb261ce7..a94d3b7a6 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_conditioner.py @@ -4,6 +4,7 @@ from __future__ import annotations import inspect +import json import numpy as np import pytest @@ -133,6 +134,33 @@ def test_prompt_cache_is_versioned_for_conditioner_layout() -> None: assert current.name != f"prompt_embeds_{legacy_digest}.npz" +def _cache_pipeline(tmp_path, quant_method=None): + encoder = tmp_path / "model" / "text_encoder" + tokenizer = tmp_path / "model" / "tokenizer" + encoder.mkdir(parents=True, exist_ok=True) + tokenizer.mkdir(parents=True, exist_ok=True) + config = {"quantization_config": {"quant_method": quant_method}} if quant_method else {} + (encoder / "config.json").write_text(json.dumps(config)) + (encoder / "model.safetensors").write_bytes(b"weights") + (tokenizer / "tokenizer.json").write_text("{}") + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.prompt_cache_dir = tmp_path / "cache" + pipeline.model_root = tmp_path / "model" + pipeline.conditioner_dir = encoder + pipeline.tokenizer_dir = tokenizer + return pipeline + + +def test_prompt_cache_key_tracks_effective_encoder_and_files(tmp_path) -> None: + pipeline = _cache_pipeline(tmp_path, quant_method="nvfp4") + auto_key = pipeline._prompt_cache_key("prompt") + pipeline.conditioner_mode = "nvfp4" + assert pipeline._prompt_cache_key("prompt") == auto_key + (pipeline.conditioner_dir / "model.safetensors").write_bytes(b"replaced weights") + assert pipeline._prompt_cache_key("prompt") != auto_key + assert pipeline._prompt_cache_key("other prompt") != pipeline._prompt_cache_key("prompt") + + def test_prompt_cache_write_is_atomic(tmp_path, monkeypatch) -> None: class FakeConditioner: def encode_prompt(self, _prompt): @@ -141,9 +169,7 @@ def test_prompt_cache_write_is_atomic(tmp_path, monkeypatch) -> None: def close(self): return None - pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) - pipeline.prompt_cache_dir = tmp_path - pipeline.model_root = tmp_path / "model" + pipeline = _cache_pipeline(tmp_path) monkeypatch.setattr(MiniMaxH3MLXPipeline, "_load_conditioner", lambda _self: FakeConditioner()) def interrupted_save(handle, **_arrays): @@ -151,7 +177,7 @@ def test_prompt_cache_write_is_atomic(tmp_path, monkeypatch) -> None: raise RuntimeError("interrupted") monkeypatch.setattr(np, "savez", interrupted_save) - cache_path = prompt_cache_path(tmp_path, pipeline.model_root, "prompt") + cache_path = pipeline._prompt_cache_key("prompt") with pytest.raises(RuntimeError, match="interrupted"): pipeline.encode_prompt("prompt") @@ -173,3 +199,51 @@ def test_video_vae_decode_defaults_to_reference_spatial_tiles() -> None: def test_h3_audio_duration_rounds_up_to_cover_last_video_packet() -> None: assert _audio_sample_count(124) == 165334 + + +@pytest.mark.parametrize("fmt,scale_fmt,activation", [("e2m1", "e4m3", "dynamic"), + ("nvfp4", "float8_e4m3fn", "Dynamic"), + ("FLOAT4_E2M1FN", "E4M3", "dynamic")]) +def test_nvfp4_encoder_config_accepts_cuda_loader_spellings(tmp_path, fmt, scale_fmt, activation) -> None: + from fastvideo.mlx_runtime.minimax_h3_conditioner import _read_nvfp4_encoder_config + + quant = {"quant_method": "NVFP4", "fmt": fmt, "group_size": 16, "scale_fmt": scale_fmt, + "scale_layout": "128x4", "activation_scheme": activation} + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": quant})) + assert _read_nvfp4_encoder_config(tmp_path)["quantization_config"] == quant + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": {**quant, "group_size": 32}})) + with pytest.raises(ValueError, match="group-16"): + _read_nvfp4_encoder_config(tmp_path) + + +def test_unswizzle_nvfp4_scales_ignores_trailing_padding() -> None: + from fastvideo.mlx_runtime.minimax_h3_conditioner import unswizzle_nvfp4_scales + + exact = np.arange(128 * 4, dtype=np.uint8) + padded = np.concatenate([exact, np.full(64, 255, dtype=np.uint8)]) + np.testing.assert_array_equal(unswizzle_nvfp4_scales(padded, 100, 3), unswizzle_nvfp4_scales(exact, 100, 3)) + with pytest.raises(ValueError, match="need 512 bytes"): + unswizzle_nvfp4_scales(exact[:-1], 100, 3) + + +def test_encoder_cache_is_published_with_umask_permissions(tmp_path, monkeypatch) -> None: + import os + + from fastvideo.mlx_runtime import minimax_h3_conditioner as conditioner + + class FakeIndex: + def __init__(self, _source): + self.weights = {"model.norm.weight": mx.ones((4, ))} + + def close(self): + pass + + monkeypatch.setattr(conditioner, "_read_nvfp4_encoder_config", lambda _dir: {}) + monkeypatch.setattr(conditioner, "_ShardIndex", lambda path: path) + monkeypatch.setattr(conditioner, "_ResidentNVFP4Index", FakeIndex) + old_umask = os.umask(0o022) + try: + out = conditioner.export_mlx_h3_nvfp4_encoder(tmp_path / "src", tmp_path / "cache") + finally: + os.umask(old_umask) + assert out.stat().st_mode & 0o777 == 0o755 diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py index e00950554..62e91df61 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py @@ -21,12 +21,16 @@ from fastvideo.mlx_runtime.minimax_h3_pipeline import ( # noqa: E402 MiniMaxH3MLXPipeline, _adaln_schedule_union, _center_crop_frames, + _configure_metal_memory_limits, _default_metal_wired_limit_gib, _preflight_media_dependencies, + _restore_metal_memory_limit, + _restore_metal_wired_limit, _validate_checkpoint_step_ladder, plan_fast_temporal, ) from fastvideo.mlx_runtime import rife_interp # noqa: E402 +from fastvideo.mlx_runtime.minimax_h3_vsa import MiniMaxH3VSAConfig # noqa: E402 def test_fast_plan_keeps_full_audio_and_reduces_only_video() -> None: @@ -143,3 +147,154 @@ def test_mux_cleans_temporary_files_after_ffmpeg_failure(tmp_path, monkeypatch: assert not output.with_suffix(".tmp.mp4").exists() assert not output.with_suffix(".tmp.wav").exists() + + +def test_explicit_wired_limit_uses_wired_api_separately_from_allocator(): + calls = [] + fake = SimpleNamespace( + metal=SimpleNamespace(device_info=lambda: {"memory_size": 36 * 2**30}), + set_memory_limit=lambda size: calls.append(("allocator", size)) or 40 * 2**30, + set_wired_limit=lambda size: calls.append(("wired", size)) or 0, + ) + previous_memory, previous_wired = _configure_metal_memory_limits(fake, 27.0) + assert calls == [("allocator", 30 * 2**30), ("wired", 27 * 2**30)] + assert (previous_memory, previous_wired) == (40 * 2**30, 0) + _restore_metal_wired_limit(fake, previous_wired) + _restore_metal_memory_limit(fake, previous_memory) + assert calls[-2:] == [("wired", 0), ("allocator", 40 * 2**30)] + + +def test_resident_placement_keeps_the_existing_allocator_limit(): + calls = [] + fake = SimpleNamespace( + set_memory_limit=lambda size: calls.append(("allocator", size)) or 40 * 2**30, + set_wired_limit=lambda size: calls.append(("wired", size)) or 0, + ) + assert _configure_metal_memory_limits(fake, 36.0, resident=True) == (None, 0) + assert calls == [("wired", 36 * 2**30)] + + +def test_explicit_wired_limit_failure_is_not_silently_ignored(): + def reject(size): + raise ValueError("exceeds system wired limit") + fake = SimpleNamespace(set_wired_limit=reject) + with pytest.raises(ValueError, match="system wired limit"): + _configure_metal_memory_limits(fake, 31.0) + with pytest.raises(ValueError, match="finite and positive"): + _configure_metal_memory_limits(fake, float("nan")) + with pytest.raises(RuntimeError, match="cannot set"): + _configure_metal_memory_limits(SimpleNamespace(), 27.0) + + +def test_wired_limit_failure_restores_allocator_limit(): + calls = [] + + def set_memory(size): + calls.append(("allocator", size)) + return 40 * 2**30 + + def reject(size): + calls.append(("wired", size)) + raise ValueError("exceeds system wired limit") + + fake = SimpleNamespace(metal=SimpleNamespace(device_info=lambda: {"memory_size": 36 * 2**30}), + set_memory_limit=set_memory, + set_wired_limit=reject) + with pytest.raises(ValueError, match="system wired limit"): + _configure_metal_memory_limits(fake, 27.0) + assert calls == [("allocator", 30 * 2**30), ("wired", 27 * 2**30), ("allocator", 40 * 2**30)] + + +def test_invalid_pipeline_options_do_not_change_process_memory_limits(monkeypatch): + calls = [] + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._configure_metal_memory_limits", + lambda *args, **kwargs: calls.append((args, kwargs))) + with pytest.raises(ValueError, match="Unknown H3 conditioner mode"): + MiniMaxH3MLXPipeline(model_root="missing", mlx_dit_checkpoint="missing", conditioner_mode="invalid", + metal_wired_limit_gib=27.0) + assert calls == [] + + +def test_explicit_wired_limit_raises_the_allocator_cap(): + calls = [] + fake = SimpleNamespace( + set_memory_limit=lambda size: calls.append(("allocator", size)) or 40 * 2**30, + set_wired_limit=lambda size: calls.append(("wired", size)) or 0, + ) + _configure_metal_memory_limits(fake, 48.0) + assert calls == [("allocator", 48 * 2**30), ("wired", 48 * 2**30)] + + +class _StopAfterPrologue(Exception): + pass + + +class _FakeAdaLNDiT: + patch_size = (1, 2, 2) + + def __init__(self): + self.blocks = [{"adaln_proj.linear.weight": np.ones(1)}] + self._adaln_cache = None + self.vsa_config = MiniMaxH3VSAConfig() + + def precompute_adaln(self, timesteps, *, drop_weights=True): + assert all(block["adaln_proj.linear.weight"] is not None for block in self.blocks) + self._adaln_cache = SimpleNamespace(timesteps=np.asarray(timesteps, dtype=np.float32)) + if drop_weights: + for block in self.blocks: + block["adaln_proj.linear.weight"] = None + + def configure_vsa(self, config): + self.vsa_config = config + + def reset_vsa_stats(self): + pass + + +def _denoise_prologue(pipeline, monkeypatch, **kwargs): + import fastvideo.mlx_runtime.minimax_h3_pipeline as pipeline_mod + + def stop(*_args, **_kwargs): + raise _StopAfterPrologue + + monkeypatch.setattr(pipeline_mod, "build_packed_layout", stop) + with pytest.raises(_StopAfterPrologue): + pipeline.denoise(np.zeros((2, 4), dtype=np.float32), np.zeros(2, dtype=np.int64), height=256, width=256, + num_frames=124, audio_num_frames=124, seed=0, **kwargs) + + +def test_resident_dit_reloads_for_a_new_step_ladder_and_resets_vsa(tmp_path, monkeypatch) -> None: + import fastvideo.mlx_runtime.minimax_h3_pipeline as pipeline_mod + + first, reloaded = _FakeAdaLNDiT(), _FakeAdaLNDiT() + monkeypatch.setattr(pipeline_mod, "_validate_checkpoint_step_ladder", lambda *a, **k: None) + monkeypatch.setattr(pipeline_mod, "_load_resident_dit", lambda _path: reloaded) + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.resident = True + pipeline.model_root = tmp_path + pipeline.dit_checkpoint = tmp_path + pipeline._resident_components = {"dit": first} + + sparse = MiniMaxH3VSAConfig(enabled=True, sparsity=0.8) + _denoise_prologue(pipeline, monkeypatch, num_steps=4, vsa_config=sparse) + assert first.vsa_config == sparse + _denoise_prologue(pipeline, monkeypatch, num_steps=4) + assert pipeline._resident_components["dit"] is first + assert first.vsa_config == MiniMaxH3VSAConfig() + + _denoise_prologue(pipeline, monkeypatch, num_steps=8) + assert pipeline._resident_components["dit"] is reloaded + np.testing.assert_array_equal(reloaded._adaln_cache.timesteps, _adaln_schedule_union(8)) + + +def test_caller_dit_with_dropped_adaln_weights_rejects_a_new_ladder(tmp_path, monkeypatch) -> None: + import fastvideo.mlx_runtime.minimax_h3_pipeline as pipeline_mod + + monkeypatch.setattr(pipeline_mod, "_validate_checkpoint_step_ladder", lambda *a, **k: None) + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.model_root = tmp_path + dit = _FakeAdaLNDiT() + _denoise_prologue(pipeline, monkeypatch, num_steps=4, dit=dit) + with pytest.raises(ValueError, match="dropped its AdaLN weights"): + pipeline.denoise(np.zeros((2, 4), dtype=np.float32), np.zeros(2, dtype=np.int64), height=256, width=256, + num_frames=124, audio_num_frames=124, seed=0, num_steps=8, dit=dit) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py index 370aa1c54..b8a46a7f3 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py @@ -93,9 +93,10 @@ def _generate_with_mocked_phases(monkeypatch, tmp_path, **generate_kwargs): pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) pipeline.video_decode_backend = "h3-vae" pipeline.dit_checkpoint = tmp_path + pipeline.model_root = tmp_path monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._validate_checkpoint_step_ladder", - lambda _checkpoint, _steps: None) + lambda _checkpoint, _steps, **_kwargs: None) monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._preflight_media_dependencies", lambda **_kwargs: None) monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline.mlx_h3_checkpoint_vsa_capable", @@ -109,7 +110,7 @@ def _generate_with_mocked_phases(monkeypatch, tmp_path, **generate_kwargs): calls["denoise"] = kwargs return np.zeros((4, 4), dtype=np.float32), np.zeros((4, 4), dtype=np.float32) - def fake_decode_video(_rows, *, height, width, num_frames, tiled): + def fake_decode_video(_rows, *, height, width, num_frames, tiled, **_kwargs): events.append("decode_video") calls["decode_video"] = {"height": height, "width": width, "num_frames": num_frames} return np.zeros((num_frames, height, width, 3), dtype=np.uint8) diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py index ca3fe3d11..433b007b7 100644 --- a/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py @@ -166,7 +166,8 @@ def test_dense_generate_ignores_unused_vsa_parameters(tmp_path, monkeypatch): pipeline = object.__new__(pipeline_mod.MiniMaxH3MLXPipeline) pipeline.video_decode_backend = "h3-vae" pipeline.dit_checkpoint = tmp_path - monkeypatch.setattr(pipeline_mod, "_validate_checkpoint_step_ladder", lambda *a: None) + pipeline.model_root = tmp_path + monkeypatch.setattr(pipeline_mod, "_validate_checkpoint_step_ladder", lambda *a, **k: None) def stop_at_media_preflight(**kwargs): raise RuntimeError("reached media preflight") monkeypatch.setattr(pipeline_mod, "_preflight_media_dependencies", stop_at_media_preflight) @@ -261,6 +262,26 @@ def test_no_metal_probe_is_cached(monkeypatch): assert not simd.simd_kernel_available() +def test_m1_simd_falls_back_before_compilation(monkeypatch): + monkeypatch.setattr(simd, "_SIMD_KERNEL", None) + monkeypatch.setattr(simd, "_SIMD_KERNEL_ERROR", None) + monkeypatch.setattr(mx.metal, "is_available", lambda: True) + monkeypatch.setattr(mx.metal, "device_info", lambda: {"device_name": "Apple M1"}) + monkeypatch.setattr(mx.fast, "metal_kernel", lambda **kwargs: pytest.fail("M1 must not compile SIMD VSA")) + assert not simd.simd_kernel_available() + assert "BF16 parity is unvalidated" in simd.simd_kernel_error() + + +def test_small_virtual_metal_simd_falls_back_before_compilation(monkeypatch): + monkeypatch.setattr(simd, "_SIMD_KERNEL", None) + monkeypatch.setattr(simd, "_SIMD_KERNEL_ERROR", None) + monkeypatch.setattr(mx.metal, "is_available", lambda: True) + monkeypatch.setattr(mx.metal, "device_info", lambda: {"device_name": "Virtual GPU", "memory_size": 7 * 2**30}) + monkeypatch.setattr(mx.fast, "metal_kernel", lambda **kwargs: pytest.fail("small GPU must not compile SIMD VSA")) + assert not simd.simd_kernel_available() + assert "8 GiB Metal devices" in simd.simd_kernel_error() + + def _require_metal(): if not mx.metal.is_available() or mx.default_device() != mx.gpu: pytest.skip("requires actual Metal execution") @@ -321,10 +342,48 @@ def test_converter_continues_past_mismatched_existing_format(tmp_path, monkeypat (existing / h3.H3_MANIFEST_FILENAME).write_text(json.dumps({"vsa": {"capable": False}})) (existing / h3.H3_WEIGHTS_FILENAME).write_bytes(b"existing") monkeypatch.setattr(converter, "parse_args", lambda: argparse.Namespace( - formats="int8 int6", out=tmp_path, model_root="unused", include_vsa=True)) + formats="int8 int6", out=tmp_path, model_root="unused", include_vsa=True, + nvfp4_conditioner_root=None, nvfp4_conditioner_out=None)) monkeypatch.setattr(converter, "mlx_h3_dit_from_diffusers_safetensors", lambda *a, **k: _dit()) saved = [] monkeypatch.setattr(converter, "save_mlx_h3_checkpoint", lambda model, path: saved.append(path.name)) converter.main() assert saved == ["int6"] assert (existing / h3.H3_WEIGHTS_FILENAME).read_bytes() == b"existing" + + +def test_converter_rerun_skips_published_encoder_cache(tmp_path, monkeypatch): + path = Path(__file__).resolve().parents[3] / "scripts/checkpoint_conversion/convert_minimax_h3_mlx.py" + spec = importlib.util.spec_from_file_location("h3_converter_encoder_rerun", path) + converter = importlib.util.module_from_spec(spec) + spec.loader.exec_module(converter) + from fastvideo.mlx_runtime import minimax_h3_conditioner as conditioner + + existing = tmp_path / "int8" + existing.mkdir() + (existing / h3.H3_MANIFEST_FILENAME).write_text(json.dumps({"vsa": {"capable": False}})) + (existing / h3.H3_WEIGHTS_FILENAME).write_bytes(b"existing") + encoder = tmp_path / "nvfp4-encoder" + encoder.mkdir() + (encoder / conditioner.MLX_NVFP4_ENCODER_MANIFEST).write_text("{}") + (encoder / "model.safetensors").write_bytes(b"cached") + monkeypatch.setattr(converter, "parse_args", lambda: argparse.Namespace( + formats="int8", out=tmp_path, model_root="unused", include_vsa=False, + nvfp4_conditioner_root=tmp_path / "packed", nvfp4_conditioner_out=None)) + exported = [] + monkeypatch.setattr(conditioner, "export_mlx_h3_nvfp4_encoder", lambda *args: exported.append(args)) + converter.main() + assert exported == [] + assert (encoder / "model.safetensors").read_bytes() == b"cached" + + +@pytest.mark.parametrize("memory_gib,budget_mib", [(16, 256), (32, 256), (64, 512), (128, 1024), (512, 2048)]) +def test_reference_gather_budget_scales_with_unified_memory(monkeypatch, memory_gib, budget_mib): + info = lambda: {"memory_size": memory_gib * 2**30} # noqa: E731 + monkeypatch.setattr(mx, "device_info", info, raising=False) + monkeypatch.setattr(mx.metal, "device_info", info) + vsa._reference_gather_target_bytes.cache_clear() + try: + assert vsa._reference_gather_target_bytes() == budget_mib * 2**20 + finally: + vsa._reference_gather_target_bytes.cache_clear() diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py new file mode 100644 index 000000000..fda1aedc8 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_spark_fence.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The DGX Spark (GB10) NVFP4 quantization fence covers inference and the QAT linear. + +GB10 is emulated by patching ``_is_dgx_spark``; any Blackwell GPU runs the kernels. +""" +from __future__ import annotations + +import pytest +import torch + +flashinfer = pytest.importorskip("flashinfer") + +from fastvideo.layers import fp4linear # noqa: E402 +from fastvideo.layers.quantization import nvfp4_config as nv # noqa: E402 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + reason="NVFP4 kernels need a Blackwell GPU") + +LAYOUT_128X4 = flashinfer.SfLayout.layout_128x4 + + +def _activation(): + torch.manual_seed(0) + x = torch.randn(200, 256, device="cuda", dtype=torch.bfloat16) + return x, (448.0 * 6.0) / x.float().abs().amax() + + +def _pad_rows(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.pad(x, (0, 0, 0, (-x.shape[0]) % 128)) + + +def _spy_quantize(monkeypatch): + real = flashinfer.nvfp4_quantize + pdl = [] + + def spy(*args, **kwargs): + pdl.append(kwargs.get("enable_pdl")) + return real(*args, **kwargs) + + monkeypatch.setattr(flashinfer, "nvfp4_quantize", spy) + return real, pdl + + +def _spy_sync(monkeypatch): + real = torch.cuda.Stream.synchronize + syncs = [] + monkeypatch.setattr(torch.cuda.Stream, "synchronize", lambda self: syncs.append(self) or real(self)) + return syncs + + +def test_fenced_quantize_matches_flashinfer_without_a_fence_off_spark(monkeypatch) -> None: + real, pdl = _spy_quantize(monkeypatch) + syncs = _spy_sync(monkeypatch) + x, global_sf = _activation() + quantized, scales = nv.nvfp4_quantize_fenced(x, global_sf, LAYOUT_128X4.value) + expected = real(x, global_sf, sfLayout=LAYOUT_128X4, do_shuffle=False) + assert pdl == [None] and syncs == [] + assert torch.equal(quantized, expected[0]) and torch.equal(scales, expected[1]) + + +def test_spark_inference_op_disables_pdl_and_fences(monkeypatch) -> None: + real, pdl = _spy_quantize(monkeypatch) + syncs = _spy_sync(monkeypatch) + monkeypatch.setattr(nv, "_is_dgx_spark", lambda _index: True) + x, global_sf = _activation() + quantized, scales = nv._nvfp4_quantize(x, global_sf, sfLayout=LAYOUT_128X4) + expected = real(_pad_rows(x), global_sf, sfLayout=LAYOUT_128X4, do_shuffle=False) + assert pdl == [False] and len(syncs) == 1 + assert torch.equal(quantized, expected[0][:x.shape[0]]) and torch.equal(scales, expected[1]) + + +def test_spark_rejects_cuda_graph_capture(monkeypatch) -> None: + monkeypatch.setattr(nv, "_is_dgx_spark", lambda _index: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + x, global_sf = _activation() + with pytest.raises(RuntimeError, match="completion fence"): + nv.nvfp4_quantize_fenced(x, global_sf, LAYOUT_128X4.value) + + +def test_qat_linear_quantizes_through_the_fence_with_unchanged_numerics(monkeypatch) -> None: + torch.manual_seed(1) + x = torch.randn(4, 64, 256, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(512, 256, device="cuda", dtype=torch.bfloat16) / 16 + bias = torch.randn(512, device="cuda", dtype=torch.bfloat16) + off_spark = fp4linear._LinearFWD4BWD16Fn.apply(x, weight, bias, "cutlass", 16, True) + + _, pdl = _spy_quantize(monkeypatch) + syncs = _spy_sync(monkeypatch) + monkeypatch.setattr(nv, "_is_dgx_spark", lambda _index: True) + on_spark = fp4linear._LinearFWD4BWD16Fn.apply(x, weight, bias, "cutlass", 16, True) + assert pdl == [False, False] and len(syncs) == 2 + assert torch.equal(on_spark, off_spark) + + reference = torch.nn.functional.linear(x.float(), weight.float(), bias.float()) + relative = (on_spark.float() - reference).norm() / reference.norm() + assert relative < 0.15 diff --git a/fastvideo/tests/stages/kandinsky6_sr/conftest.py b/fastvideo/tests/stages/kandinsky6_sr/conftest.py index a629b055c..ae0559e7c 100644 --- a/fastvideo/tests/stages/kandinsky6_sr/conftest.py +++ b/fastvideo/tests/stages/kandinsky6_sr/conftest.py @@ -4,31 +4,33 @@ from __future__ import annotations import os import sys +from contextlib import ExitStack from pathlib import Path import pytest import torch +import fastvideo.envs as envs sys.path.insert(0, str(Path(__file__).resolve().parent)) import k6_sr_tiny # noqa: E402 -_DIST_ENV_DEFAULTS = { - "MASTER_ADDR": "localhost", - "MASTER_PORT": "29531", - "RANK": "0", - "WORLD_SIZE": "1", - "LOCAL_RANK": "0", -} - - @pytest.fixture(autouse=True) -def _single_process_dist_env(monkeypatch): +def _single_process_dist_env(): """torch.distributed needs a rendezvous even for one process. Fill only what is missing (never overwrite what a CI runner leased) and undo it after the test so the settings do not leak into other tests of the same session.""" - for name, value in _DIST_ENV_DEFAULTS.items(): - if name not in os.environ: - monkeypatch.setenv(name, value) + with ExitStack() as stack: + if "MASTER_ADDR" not in os.environ: + stack.enter_context(envs.override_external("MASTER_ADDR", "localhost")) + if "MASTER_PORT" not in os.environ: + stack.enter_context(envs.override_external("MASTER_PORT", "29531")) + if "RANK" not in os.environ: + stack.enter_context(envs.override_external("RANK", "0")) + if "WORLD_SIZE" not in os.environ: + stack.enter_context(envs.override_external("WORLD_SIZE", "1")) + if "LOCAL_RANK" not in os.environ: + stack.enter_context(envs.override_external("LOCAL_RANK", "0")) + yield @pytest.fixture(scope="session") diff --git a/fastvideo/tests/worker/test_ray_distributed_executor.py b/fastvideo/tests/worker/test_ray_distributed_executor.py index 0bf8b4d1e..ffb3e4842 100644 --- a/fastvideo/tests/worker/test_ray_distributed_executor.py +++ b/fastvideo/tests/worker/test_ray_distributed_executor.py @@ -36,3 +36,17 @@ def test_ray_log_queue_stays_on_the_driver() -> None: executor.clear_log_queue() assert executor._log_queue is None assert "log_queue" in signature(Executor.set_log_queue).parameters + + +def test_ray_carries_h3_performance_switches() -> None: + import fastvideo.envs as envs + from fastvideo.worker.ray_env import get_env_vars_to_copy + + with (envs.FASTVIDEO_H3_VAE_TILE_BATCH.override(8), + envs.FASTVIDEO_NVFP4_MM_BACKEND.override("cutlass"), + envs.FASTVIDEO_VSA_TRITON.override(True)): + copied = get_env_vars_to_copy() + assert {"FASTVIDEO_H3_VAE_TILE_BATCH", "FASTVIDEO_NVFP4_MM_BACKEND", "FASTVIDEO_VSA_TRITON"} <= copied + assert envs.FASTVIDEO_H3_VAE_TILE_BATCH.get() == 8 + assert envs.FASTVIDEO_NVFP4_MM_BACKEND.get() == "cutlass" + assert envs.FASTVIDEO_VSA_TRITON.get() is True diff --git a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py index db96f413c..688e2f326 100644 --- a/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py +++ b/scripts/checkpoint_conversion/convert_minimax_h3_mlx.py @@ -90,6 +90,10 @@ def parse_args() -> argparse.Namespace: help=("retain and quantize transformer_blocks.*.attn.to_gate_compress.weight " "(required for MLX VSA inference; omitted by dense conversion)"), ) + parser.add_argument("--nvfp4-conditioner-root", type=Path, + help="also cache the released packed encoder in native MLX layout, without requantization") + parser.add_argument("--nvfp4-conditioner-out", type=Path, + help="empty encoder-cache output directory; defaults to OUT/nvfp4-encoder") return parser.parse_args() @@ -144,6 +148,21 @@ def main() -> None: if hasattr(mx, "clear_cache"): mx.clear_cache() + if args.nvfp4_conditioner_root is not None: + from fastvideo.mlx_runtime.minimax_h3_conditioner import ( + MLX_NVFP4_ENCODER_MANIFEST, + export_mlx_h3_nvfp4_encoder, + ) + + encoder_out = args.nvfp4_conditioner_out or out_base / "nvfp4-encoder" + if (encoder_out / MLX_NVFP4_ENCODER_MANIFEST).exists() and (encoder_out / "model.safetensors").exists(): + print(f"[skip] NVFP4 encoder already cached at {encoder_out}", flush=True) + else: + started = time.perf_counter() + export_mlx_h3_nvfp4_encoder(args.nvfp4_conditioner_root, encoder_out) + print(f"[encoder] cached packed NVFP4 encoder in {time.perf_counter() - started:.1f}s at {encoder_out}", + flush=True) + if __name__ == "__main__": main() From 33d81730e9fe423ab3c33dbc0451973267b80d92 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 6 Oct 2026 17:09:53 -0700 Subject: [PATCH 055/151] [feat]: Kandinsky 6 Lite: register the 3.2B checkpoints and add docs + cookbook recipes (#1927) --- docs/assets/cookbook-recipes.json | 34 ++++++++++++++++++- docs/cookbook/cosmos.md | 2 +- docs/cookbook/flux.md | 2 +- docs/cookbook/glm-image.md | 2 +- docs/cookbook/hunyuan.md | 2 +- docs/cookbook/kandinsky5.md | 2 +- docs/cookbook/kandinsky6.md | 4 +-- docs/cookbook/longcat.md | 2 +- docs/cookbook/ltx.md | 2 +- docs/cookbook/matrix-game.md | 2 +- docs/cookbook/minimax-h3.md | 2 +- docs/cookbook/mmaudio.md | 2 +- docs/cookbook/stable-audio.md | 2 +- docs/cookbook/stable-diffusion.md | 2 +- docs/cookbook/turbodiffusion.md | 2 +- docs/cookbook/wan.md | 2 +- docs/cookbook/z-image.md | 2 +- docs/inference/kandinsky6.md | 18 ++++++---- docs/inference/support_matrix.md | 8 ++--- fastvideo/registry.py | 6 +++- .../tests/stages/test_kandinsky6_registry.py | 18 +++++++--- tests/local_tests/test_cookbook_serving.py | 19 +++++++---- 22 files changed, 96 insertions(+), 41 deletions(-) diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 63a0cdbf5..8e8cc4576 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1,5 +1,5 @@ { - "version": 14, + "version": 15, "recipes": [ { "id": "fastwan21-t2v", @@ -364,6 +364,38 @@ "expected_artifact": "MP4 video with synchronized audio under video_samples_kandinsky6_ti2va/", "related": ["kandinsky6-ti2va-base", "kandinsky6-vsr-distilled"] }, + { + "id": "kandinsky6-ti2va-lite", + "family": "kandinsky6", + "stage": "inference", + "task": "Text or image to video with audio", + "label": "Kandinsky 6 Lite TI2VA", + "summary": "Generate a five-second, 512x768 video with synchronized audio from text or an image with the 3.2B Lite checkpoint in fifty inference steps and guidance 5.0.", + "model": "kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_ti2va.py", + "command": "KANDINSKY6_MODEL_PATH=kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers python examples/inference/basic/basic_kandinsky6_ti2va.py", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "MP4 video with synchronized audio under video_samples_kandinsky6_ti2va/", + "related": ["kandinsky6-ti2va-lite-piflow", "kandinsky6-ti2va-base"] + }, + { + "id": "kandinsky6-ti2va-lite-piflow", + "family": "kandinsky6", + "stage": "inference", + "task": "Distilled text or image to video with audio", + "label": "Kandinsky 6 Lite pi-Flow", + "summary": "Generate a five-second video with synchronized audio using the 3.2B Lite distilled pi-Flow checkpoint in ten inference steps and guidance 1.0.", + "model": "kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers", + "source": "examples/inference/basic/basic_kandinsky6_ti2va.py", + "command": "KANDINSKY6_MODEL_PATH=kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers python examples/inference/basic/basic_kandinsky6_ti2va.py", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "MP4 video with synchronized audio under video_samples_kandinsky6_ti2va/", + "related": ["kandinsky6-ti2va-lite", "kandinsky6-ti2va-piflow"] + }, { "id": "kandinsky6-vsr", "family": "kandinsky6", diff --git a/docs/cookbook/cosmos.md b/docs/cookbook/cosmos.md index 4d9c8c0fb..89ad17313 100644 --- a/docs/cookbook/cosmos.md +++ b/docs/cookbook/cosmos.md @@ -5,7 +5,7 @@ hide: # Cosmos recipes -
    +
    All model families
    diff --git a/docs/cookbook/flux.md b/docs/cookbook/flux.md index 160c56ca5..bcb59a2fe 100644 --- a/docs/cookbook/flux.md +++ b/docs/cookbook/flux.md @@ -5,7 +5,7 @@ hide: # FLUX recipes -
    +
    All model families
    diff --git a/docs/cookbook/glm-image.md b/docs/cookbook/glm-image.md index 93f8b5cf2..64bcdae21 100644 --- a/docs/cookbook/glm-image.md +++ b/docs/cookbook/glm-image.md @@ -5,7 +5,7 @@ hide: # GLM-Image recipes -
    +
    All model families
    diff --git a/docs/cookbook/hunyuan.md b/docs/cookbook/hunyuan.md index dede8af94..52c11cd56 100644 --- a/docs/cookbook/hunyuan.md +++ b/docs/cookbook/hunyuan.md @@ -5,7 +5,7 @@ hide: # Hunyuan recipes -
    +
    All model families
    diff --git a/docs/cookbook/kandinsky5.md b/docs/cookbook/kandinsky5.md index d899dd977..2389f1bd6 100644 --- a/docs/cookbook/kandinsky5.md +++ b/docs/cookbook/kandinsky5.md @@ -5,7 +5,7 @@ hide: # Kandinsky 5 recipes -
    +
    All model families
    diff --git a/docs/cookbook/kandinsky6.md b/docs/cookbook/kandinsky6.md index fc9f32dff..9fd289c66 100644 --- a/docs/cookbook/kandinsky6.md +++ b/docs/cookbook/kandinsky6.md @@ -5,7 +5,7 @@ hide: # Kandinsky 6 recipes -
    +
    All model families
    @@ -15,7 +15,7 @@ hide:

    Maintained family · Inference

    Kandinsky 6 inference recipes

    -

    Kandinsky 6 from the Kandinsky Lab generates five-second video with synchronized audio from text or an image, with base and distilled pi-Flow checkpoints, and upscales existing clips with base or distilled video super-resolution.

    +

    Kandinsky 6 from the Kandinsky Lab generates five-second video with synchronized audio from text or an image, in Pro (30.1B) and Lite (3.2B) sizes with base and distilled pi-Flow checkpoints, and upscales existing clips with base or distilled video super-resolution.

    diff --git a/docs/cookbook/longcat.md b/docs/cookbook/longcat.md index 28f303deb..33addd0f4 100644 --- a/docs/cookbook/longcat.md +++ b/docs/cookbook/longcat.md @@ -5,7 +5,7 @@ hide: # LongCat recipes -
    +
    All model families
    diff --git a/docs/cookbook/ltx.md b/docs/cookbook/ltx.md index f751661c9..d8878be7b 100644 --- a/docs/cookbook/ltx.md +++ b/docs/cookbook/ltx.md @@ -5,7 +5,7 @@ hide: # LTX recipes -
    +
    All model families
    diff --git a/docs/cookbook/matrix-game.md b/docs/cookbook/matrix-game.md index b6d4e16cd..10f0eb4ca 100644 --- a/docs/cookbook/matrix-game.md +++ b/docs/cookbook/matrix-game.md @@ -5,7 +5,7 @@ hide: # Matrix Game recipes -
    +
    All model families
    diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 0ccb0b3bb..265a508d6 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -18,7 +18,7 @@ eight-forward T2VA command](../getting_started/installation/mlx.md#pruned-eight- It reads `fastvideo_inference.json` for the trained schedule. The command uses native 832x480 resolution and all requested frames. -
    +
    All model families
    diff --git a/docs/cookbook/mmaudio.md b/docs/cookbook/mmaudio.md index 544241d44..8cd584cfa 100644 --- a/docs/cookbook/mmaudio.md +++ b/docs/cookbook/mmaudio.md @@ -5,7 +5,7 @@ hide: # MMAudio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-audio.md b/docs/cookbook/stable-audio.md index 5fa15b299..64f52ce32 100644 --- a/docs/cookbook/stable-audio.md +++ b/docs/cookbook/stable-audio.md @@ -5,7 +5,7 @@ hide: # Stable Audio recipes -
    +
    All model families
    diff --git a/docs/cookbook/stable-diffusion.md b/docs/cookbook/stable-diffusion.md index 845a30594..9bd005ea0 100644 --- a/docs/cookbook/stable-diffusion.md +++ b/docs/cookbook/stable-diffusion.md @@ -5,7 +5,7 @@ hide: # Stable Diffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/turbodiffusion.md b/docs/cookbook/turbodiffusion.md index df3dac64f..0c72c4fea 100644 --- a/docs/cookbook/turbodiffusion.md +++ b/docs/cookbook/turbodiffusion.md @@ -5,7 +5,7 @@ hide: # TurboDiffusion recipes -
    +
    All model families
    diff --git a/docs/cookbook/wan.md b/docs/cookbook/wan.md index 7139d02fb..d23766ab3 100644 --- a/docs/cookbook/wan.md +++ b/docs/cookbook/wan.md @@ -5,7 +5,7 @@ hide: # Wan recipes -
    +
    All model families
    diff --git a/docs/cookbook/z-image.md b/docs/cookbook/z-image.md index 4db21f991..e912e932a 100644 --- a/docs/cookbook/z-image.md +++ b/docs/cookbook/z-image.md @@ -5,7 +5,7 @@ hide: # Z-Image recipes -
    +
    All model families
    diff --git a/docs/inference/kandinsky6.md b/docs/inference/kandinsky6.md index da4307f7b..f623b1400 100644 --- a/docs/inference/kandinsky6.md +++ b/docs/inference/kandinsky6.md @@ -7,12 +7,16 @@ mp4 automatically. To upscale a generated clip, see [Kandinsky 6 Video SR](kandi ## Models -Both variants are official Diffusers repos, loaded directly through their `model_index.json`: +Kandinsky 6 comes in two sizes, Pro (30.1B-parameter DiT) and Lite (3.2B), each with a base and a pi-Flow distilled +checkpoint. All four are official Diffusers repos, loaded directly through their `model_index.json`; Lite shares Pro's +architecture, text encoders, VAEs and schedulers, with a narrower and shallower DiT. -| Variant | Hub repo | Scheduler | Steps | Guidance | Example | -|---|---|---|---|---|---| -| T2IVA | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | `FlowMatchEulerDiscreteScheduler` (shift 5.0) | 50 | 5.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | -| T2IVA distilled | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | `PiflowScheduler` (`n_grid` 10, shift 5.0) | 10 | 1.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| Size | Variant | Hub repo | Scheduler | Steps | Guidance | Example | +|---|---|---|---|---|---|---| +| Pro | T2IVA | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers` | `FlowMatchEulerDiscreteScheduler` (shift 5.0) | 50 | 5.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| Pro | T2IVA distilled | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | `PiflowScheduler` (`n_grid` 10, shift 5.0) | 10 | 1.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| Lite | T2IVA | `kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers` | `FlowMatchEulerDiscreteScheduler` (shift 5.0) | 50 | 5.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| Lite | T2IVA distilled | `kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers` | `PiflowScheduler` (`n_grid` 10, shift 5.0) | 10 | 1.0 | [`basic_kandinsky6_ti2va.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | The steps and guidance columns are the defaults of the preset the registry selects for each repo id. Everything else is shared: 512x768, 121 frames (5 s at 24 fps) and the Diffusers default negative prompt (only used when @@ -101,7 +105,7 @@ nothing: ## Memory -The Pro DiT has 30.1B parameters, about 60 GB in bf16 (`dit_precision` defaults to `bf16`), and the Qwen2.5-VL text -encoder adds 16.6 GB. FastVideo enables `dit_cpu_offload` by default; the examples turn it off (`dit_cpu_offload=False`) +The Pro DiT has 30.1B parameters, about 60 GB in bf16 (`dit_precision` defaults to `bf16`); the Lite DiT has 3.2B, +about 6.4 GB. Both use the same Qwen2.5-VL text encoder, which adds 16.6 GB. FastVideo enables `dit_cpu_offload` by default; the examples turn it off (`dit_cpu_offload=False`) to keep the DiT resident on the GPU and offload the text encoder instead (`text_encoder_cpu_offload=True`). See [Offloading](offloading.md) for the memory knobs. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index c89be5ce8..16f7f10d9 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -54,8 +54,8 @@ column links a runnable script in `examples/inference/basic/` where one exists. | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Lite-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-sft-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | | kandinsky5 | `kandinskylab/Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers` | I2V | [basic_kandinsky5_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky5_i2v.py) | -| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`
    `kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | -| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`
    `kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers`
    `kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | +| kandinsky6 | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers`
    `kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers` | T2V, I2V | [basic_kandinsky6_ti2va.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_ti2va.py) | | kandinsky6_sr | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`
    `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | — | [basic_kandinsky6_sr.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_kandinsky6_sr.py) | | lingbot_video | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) | | lingbot_video | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) | @@ -109,8 +109,8 @@ reference videos as sparse VSA regions) run through [basic_fasth3_omniref_pdd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_fasth3_omniref_pdd.py); see [FastH3 distilled checkpoint schedules](fasth3-distilled.md#ref2va-pdd-students). -**Note (Kandinsky 6)**: the two `kandinsky6` IDs generate video with audio from -text, optionally plus an image (see the [T2IVA guide](kandinsky6.md)); the +**Note (Kandinsky 6)**: the `kandinsky6` IDs (Pro 30.1B and Lite 3.2B, each base +and distilled) generate video with audio from text, optionally plus an image (see the [T2IVA guide](kandinsky6.md)); the `kandinsky6_sr` IDs upscale an existing video (see the [Video SR guide](kandinsky6_sr.md)). diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 8a8e2f068..6d2c9e7f8 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -814,7 +814,10 @@ def _register_configs() -> None: sampling_param_cls=None, pipeline_config_cls=Kandinsky6TI2VAConfig, workload_types=(WorkloadType.T2V, WorkloadType.I2V), - hf_model_paths=["kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers"], + hf_model_paths=[ + "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", + "kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers", + ], model_detectors=[ _is_kandinsky6_distilled, ], @@ -829,6 +832,7 @@ def _register_configs() -> None: hf_model_paths=[ "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers", "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", + "kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers", ], model_detectors=[ _is_kandinsky6, diff --git a/fastvideo/tests/stages/test_kandinsky6_registry.py b/fastvideo/tests/stages/test_kandinsky6_registry.py index 4dfd6c809..b81383421 100644 --- a/fastvideo/tests/stages/test_kandinsky6_registry.py +++ b/fastvideo/tests/stages/test_kandinsky6_registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Registry routing and presets of the Kandinsky6 TI2VA checkpoints: base (Pro-sft) and pi-Flow distilled (Pro-distill). +"""Registry routing and presets of the Kandinsky6 TI2VA checkpoints: base and pi-Flow distilled, Pro and Lite sizes. Both official Diffusers repos declare the same model_index ``_class_name`` (``Kandinsky6TI2VAPipeline``), so they are told apart by repo id or directory name, laid out like the LTX-2 distilled/base pair: the distilled entry is registered @@ -28,6 +28,8 @@ from fastvideo.pipelines.basic.kandinsky6.presets import ( BASE_ID = "kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers" OLD_BASE_ID = "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers" DISTILLED_ID = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" +LITE_ID = "kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers" +LITE_DISTILLED_ID = "kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers" BASE_PRESET = "kandinsky6_ti2va" DISTILLED_PRESET = "kandinsky6_ti2va_distilled" @@ -48,7 +50,8 @@ def _matching_presets(text: str) -> set[str]: @pytest.mark.parametrize("repo_id,preset", [(BASE_ID, BASE_PRESET), (OLD_BASE_ID, BASE_PRESET), - (DISTILLED_ID, DISTILLED_PRESET)]) + (DISTILLED_ID, DISTILLED_PRESET), (LITE_ID, BASE_PRESET), + (LITE_DISTILLED_ID, DISTILLED_PRESET)]) def test_official_repo_ids_resolve_to_their_presets(repo_id, preset): assert repo_id in registry.get_registered_model_paths() assert registry.get_pipeline_config_cls_from_name(repo_id) is Kandinsky6TI2VAConfig @@ -60,7 +63,7 @@ def test_official_repo_ids_resolve_to_their_presets(repo_id, preset): assert info.workload_types == (WorkloadType.T2V, WorkloadType.I2V) -@pytest.mark.parametrize("repo_id", [BASE_ID, DISTILLED_ID]) +@pytest.mark.parametrize("repo_id", [BASE_ID, DISTILLED_ID, LITE_ID, LITE_DISTILLED_ID]) def test_get_model_info_resolves_both_repo_ids_to_the_ti2va_pipeline(monkeypatch, repo_id): model_index = {"_class_name": "Kandinsky6TI2VAPipeline", "_diffusers_version": "0.41.0.dev0"} monkeypatch.setattr(registry, "maybe_download_model_index", lambda *_args, **_kwargs: model_index) @@ -71,13 +74,15 @@ def test_get_model_info_resolves_both_repo_ids_to_the_ti2va_pipeline(monkeypatch def test_registered_model_listing_offers_both_repo_ids_for_text_and_image_workloads(): listed = {model["id"]: model["workload_types"] for model in registry.get_registered_models_with_workloads()} - assert listed[BASE_ID] == listed[DISTILLED_ID] == ["t2v", "i2v"] + assert listed[BASE_ID] == listed[DISTILLED_ID] == listed[LITE_ID] == listed[LITE_DISTILLED_ID] == ["t2v", "i2v"] @pytest.mark.parametrize("name,preset", [ ("Kandinsky-6.0-Pro-5s-Diffusers", BASE_PRESET), ("Kandinsky-6.0-Pro-sft-5s-Diffusers", BASE_PRESET), ("Kandinsky-6.0-Pro-distill-5s-Diffusers", DISTILLED_PRESET), + ("Kandinsky-6.0-Lite-5s-Diffusers", BASE_PRESET), + ("Kandinsky-6.0-Lite-distill-5s-Diffusers", DISTILLED_PRESET), ("kandinsky6-distilled-export", DISTILLED_PRESET), ("kandinsky6-export", BASE_PRESET), ("my_export", BASE_PRESET), @@ -102,6 +107,8 @@ def test_distill_in_a_parent_directory_does_not_make_a_base_checkpoint_distilled ("/models/kandinsky6-distilled-export", {DISTILLED_PRESET}), ("/models/kandinsky6-distilled-export/", {DISTILLED_PRESET}), ("kandinsky-6.0-pro-sft-5s-diffusers", {BASE_PRESET}), + ("kandinsky-6.0-lite-5s-diffusers", {BASE_PRESET}), + ("kandinsky-6.0-lite-distill-5s-diffusers", {DISTILLED_PRESET}), ("kandinsky6ti2vapipeline", {BASE_PRESET}), ("/distill_experiments/kandinsky6-export", {BASE_PRESET}), ]) @@ -141,7 +148,8 @@ def test_denoise_stage_overrides_accept_steps_and_guidance(preset): validate_preset_selection(preset, "kandinsky6", stage_overrides=overrides) -@pytest.mark.parametrize("repo_id,steps,guidance", [(BASE_ID, 50, 5.0), (DISTILLED_ID, 10, 1.0)]) +@pytest.mark.parametrize("repo_id,steps,guidance", [(BASE_ID, 50, 5.0), (DISTILLED_ID, 10, 1.0), (LITE_ID, 50, 5.0), + (LITE_DISTILLED_ID, 10, 1.0)]) def test_sampling_param_defaults_follow_the_repo_id(repo_id, steps, guidance): sampling_param = SamplingParam.from_pretrained(repo_id) assert (sampling_param.num_inference_steps, sampling_param.guidance_scale) == (steps, guidance) diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index f90b17c9d..a2ff32014 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -265,9 +265,12 @@ def test_kandinsky6_cookbook_recipes_match_maintained_examples(): assert [recipe["id"] for recipe in family] == [ "kandinsky6-ti2va-base", "kandinsky6-ti2va-piflow", + "kandinsky6-ti2va-lite", + "kandinsky6-ti2va-lite-piflow", "kandinsky6-vsr", "kandinsky6-vsr-distilled", ] + by_id = {recipe["id"]: recipe for recipe in family} assert all((ROOT / recipe["source"]).is_file() for recipe in family) assert all(recipe["evidence"] == "Source-backed" for recipe in family) assert all(recipe["hardware"] == { @@ -277,10 +280,14 @@ def test_kandinsky6_cookbook_recipes_match_maintained_examples(): } for recipe in family) assert all("accelerator" not in recipe["hardware"] and "peak_memory" not in recipe["hardware"] for recipe in family) - assert family[1]["model"].endswith("Pro-distill-5s-Diffusers") - # The distilled recipe reruns the shared example with KANDINSKY6_MODEL_PATH; the repo id selects the 10-step preset. - assert family[1]["command"].startswith("KANDINSKY6_MODEL_PATH=" + family[1]["model"] + " ") - assert "KANDINSKY6_MODEL_PATH" in (ROOT / family[1]["source"]).read_text() - assert all("INPUT_VIDEO" in recipe["command"] for recipe in family[2:]) - assert "VSR-distilled2steps" in family[3]["command"] + # The distilled and Lite recipes rerun the shared example with KANDINSKY6_MODEL_PATH; the repo id selects the preset. + for recipe_id, suffix in [("kandinsky6-ti2va-piflow", "Pro-distill-5s-Diffusers"), + ("kandinsky6-ti2va-lite", "Lite-5s-Diffusers"), + ("kandinsky6-ti2va-lite-piflow", "Lite-distill-5s-Diffusers")]: + recipe = by_id[recipe_id] + assert recipe["model"].endswith(suffix) + assert recipe["command"].startswith("KANDINSKY6_MODEL_PATH=" + recipe["model"] + " ") + assert "KANDINSKY6_MODEL_PATH" in (ROOT / by_id["kandinsky6-ti2va-piflow"]["source"]).read_text() + assert all("INPUT_VIDEO" in by_id[recipe_id]["command"] for recipe_id in ("kandinsky6-vsr", "kandinsky6-vsr-distilled")) + assert "VSR-distilled2steps" in by_id["kandinsky6-vsr-distilled"]["command"] validate_cookbook() From d8ca702dfdb5140ab649252bafaff34db43cb02f Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 7 Oct 2026 16:28:53 +0800 Subject: [PATCH 056/151] [bugfix] registry: register magi_human presets (#1751) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- fastvideo/registry.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 6d2c9e7f8..b1e1486d7 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -1273,6 +1273,8 @@ def _register_presets() -> None: ALL_PRESETS as LONGCAT_PRESETS, ) from fastvideo.pipelines.basic.ltx2.presets import ( ALL_PRESETS as LTX2_PRESETS, ) + from fastvideo.pipelines.basic.magi_human.presets import ( + ALL_PRESETS as MAGI_HUMAN_PRESETS, ) from fastvideo.pipelines.basic.matrixgame2.presets import ( ALL_PRESETS as MATRIXGAME2_PRESETS, ) from fastvideo.pipelines.basic.matrixgame3.presets import ( @@ -1311,6 +1313,7 @@ def _register_presets() -> None: LINGBOTWORLD2_PRESETS, LONGCAT_PRESETS, LTX2_PRESETS, + MAGI_HUMAN_PRESETS, MATRIXGAME2_PRESETS, MATRIXGAME3_PRESETS, MINIMAX_H3_PRESETS, From 38966056d0907a1b5b6216442694278562eab3a5 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:29:42 -0700 Subject: [PATCH 057/151] [ci]: make merge-gate startup reliable (#1762) Co-authored-by: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- .github/workflows/ci-slash-commands.yml | 26 ++- .github/workflows/ci-trigger-full-suite.yml | 164 ++++++++++++++---- docs/contributing/ci_architecture.md | 3 +- docs/contributing/pull_requests.md | 3 +- .../tests/contract/test_ci_test_collection.py | 85 ++++++++- 5 files changed, 241 insertions(+), 40 deletions(-) diff --git a/.github/workflows/ci-slash-commands.yml b/.github/workflows/ci-slash-commands.yml index 2755a7bb6..495bc7f70 100644 --- a/.github/workflows/ci-slash-commands.yml +++ b/.github/workflows/ci-slash-commands.yml @@ -32,7 +32,7 @@ jobs: } core.setOutput('has_write', String(hasWrite)); - - name: Add ready label and react + - name: Add ready label if: steps.perm.outputs.has_write == 'true' uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 with: @@ -40,14 +40,34 @@ jobs: const owner = context.repo.owner; const repo = context.repo.repo; const prNumber = context.payload.issue.number; - try { await github.rest.issues.removeLabel({ owner, repo, issue_number: prNumber, name: 'ready' }); } catch {} await github.rest.issues.addLabels({ owner, repo, issue_number: prNumber, labels: ['ready'] }); + + - name: React to comment + if: steps.perm.outputs.has_write == 'true' + continue-on-error: true + uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + with: + script: | await github.rest.reactions.createForIssueComment({ - owner, repo, + owner: context.repo.owner, + repo: context.repo.repo, comment_id: context.payload.comment.id, content: 'rocket', }); + trigger-merge-gate: + needs: handle-merge + if: needs.handle-merge.result == 'success' + permissions: + actions: read + contents: read + pull-requests: read + uses: ./.github/workflows/ci-trigger-full-suite.yml + with: + pr_number: ${{ github.event.issue.number }} + secrets: + BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} + parse-command: if: >- github.event.issue.pull_request != null diff --git a/.github/workflows/ci-trigger-full-suite.yml b/.github/workflows/ci-trigger-full-suite.yml index e96ad9c81..31102f683 100644 --- a/.github/workflows/ci-trigger-full-suite.yml +++ b/.github/workflows/ci-trigger-full-suite.yml @@ -3,74 +3,178 @@ name: Trigger Merge Gate on: pull_request_target: types: [labeled, synchronize] + workflow_call: + inputs: + pr_number: + description: Pull request number to enter into the merge gate + required: true + type: number + secrets: + BUILDKITE_API_TOKEN: + required: true permissions: contents: read pull-requests: read actions: read -concurrency: - group: merge-gate-${{ github.event.pull_request.number }} - cancel-in-progress: false - jobs: trigger: if: >- - (github.event.action == 'labeled' && github.event.label.name == 'ready') + inputs.pr_number > 0 + || (github.event.action == 'labeled' && github.event.label.name == 'ready') || github.event.action == 'synchronize' runs-on: ubuntu-latest + # Job-level concurrency: only this guarded job acquires the group, so an + # unrelated `labeled` event (which skips the job) cannot cancel an in-flight + # gate and then skip its replacement. The newest real trigger (`ready`, + # push, or `/merge`) supersedes the in-flight run, whose Buildkite build the + # cancel step below replaces. + concurrency: + group: merge-gate-${{ inputs.pr_number || github.event.pull_request.number }} + cancel-in-progress: true # Gate below may wait for cheap checks (up to MAX_WAIT_SECS = 25 min). timeout-minutes: 35 steps: - name: Check ready label id: check uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + env: + CALLED_PR_NUMBER: ${{ inputs.pr_number }} with: script: | + const eventPrNumber = context.payload.pull_request?.number; + const calledPrNumber = Number(process.env.CALLED_PR_NUMBER); + const prNumber = eventPrNumber ?? calledPrNumber; + if (!Number.isSafeInteger(prNumber) || prNumber <= 0) { + core.setFailed(`Invalid pull request number: ${process.env.CALLED_PR_NUMBER}`); + return; + } const { data: pr } = await github.rest.pulls.get({ owner: context.repo.owner, repo: context.repo.repo, - pull_number: context.payload.pull_request.number, + pull_number: prNumber, }); + if (pr.state !== 'open') { + core.setFailed(`PR #${prNumber} is not open.`); + return; + } + if (pr.base.repo.full_name !== context.payload.repository.full_name + || pr.base.ref !== context.payload.repository.default_branch) { + core.setFailed(`PR #${prNumber} does not target this repository's default branch.`); + return; + } const hasReady = pr.labels.some(l => l.name === 'ready'); core.setOutput('has_ready', String(hasReady)); core.setOutput('changed_files', String(pr.changed_files)); + core.setOutput('pr_number', String(pr.number)); + core.setOutput('head_sha', pr.head.sha); + core.setOutput('head_ref', pr.head.ref); + core.setOutput('base_sha', pr.base.sha); + core.setOutput('title', pr.title); if (!hasReady) core.info('No ready label — skipping merge-gate trigger.'); - name: Cancel previous Buildkite builds + # Cancelling stale builds only saves agent time. If it cannot run, the + # merge gate must still be triggered by the steps below, so a failure + # here is reported and stepped over rather than ending the job. + continue-on-error: true + timeout-minutes: 3 if: steps.check.outputs.has_ready == 'true' env: + BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }} + BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }} BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} - PR_BRANCH: ${{ github.event.pull_request.head.ref }} - PR_NUMBER: ${{ github.event.pull_request.number }} + PR_BRANCH: ${{ steps.check.outputs.head_ref }} + PR_NUMBER: ${{ steps.check.outputs.pr_number }} run: | - # Match both branch and PR number: forks can reuse the same branch name. - builds=$(curl -sS --get -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \ - --data-urlencode "branch=$PR_BRANCH" \ - --data-urlencode "state=running,scheduled" \ - "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds" \ - | jq -r --arg pr_number "$PR_NUMBER" \ - '.[] | select((.env.TEST_SCOPE? == "merge") and (.env.PR_NUMBER? == $pr_number)) | .number') - for build_num in $builds; do - echo "Cancelling Buildkite build #$build_num" - curl -sS -X PUT -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \ - "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel" - done + set -euo pipefail + response_file=$(mktemp) + builds_file=$(mktemp) + trap 'rm -f "$response_file" "$builds_file"' EXIT - # Check out the immutable BASE SHA: pull_request_target must never run a - # planner or gate script from the untrusted PR head. + if [[ ! "$PR_NUMBER" =~ ^[1-9][0-9]*$ ]]; then + echo "::warning::Invalid pull request number; stale Buildkite builds may continue." + exit 1 + fi + + if ! curl -sS --fail-with-body --connect-timeout 5 --max-time 20 --get \ + -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \ + --data-urlencode "branch=$PR_BRANCH" \ + --data-urlencode "state[]=running" \ + --data-urlencode "state[]=scheduled" \ + --data-urlencode "state[]=failing" \ + --data-urlencode "exclude_jobs=true" \ + --data-urlencode "exclude_pipeline=true" \ + --output "$response_file" \ + "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds"; then + echo "::warning::Could not list Buildkite builds; stale merge-gate builds may continue." + exit 1 + fi + + if ! jq -e ' + if type != "array" then false + else all(.[]; + if type != "object" then false + else + (.number | if type == "number" then . > 0 and floor == . else false end) + and ( + (.env? | if . == null then {} else . end) as $env + | if ($env | type) != "object" then false + else + ($env.TEST_SCOPE? | . == null or type == "string") + and ($env.PR_NUMBER? | . == null or type == "string") + end + ) + end + ) + end + ' "$response_file" >/dev/null 2>&1; then + # Do not print the response body: it is remote data and may contain + # multiline values that would be interpreted as workflow commands. + echo "::warning::Buildkite returned an invalid build list; stale merge-gate builds may continue." + exit 1 + fi + + # Match both branch and PR number: forks can reuse the same branch name. + if ! jq -r --arg pr_number "$PR_NUMBER" ' + .[] + | select((.env.TEST_SCOPE? == "merge") and (.env.PR_NUMBER? == $pr_number)) + | .number + ' "$response_file" > "$builds_file"; then + echo "::warning::Could not select stale Buildkite builds; stale merge-gate builds may continue." + exit 1 + fi + + cancellation_failed=0 + while IFS= read -r build_num; do + echo "Cancelling Buildkite build #$build_num" + if ! curl -sS --fail-with-body --connect-timeout 5 --max-time 20 -o /dev/null -X PUT \ + -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \ + "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds/${build_num}/cancel"; then + echo "::warning::Could not cancel Buildkite build #$build_num; trying remaining builds." + cancellation_failed=1 + fi + done < "$builds_file" + + if (( cancellation_failed != 0 )); then + exit 1 + fi + + # Check out the immutable BASE SHA: neither pull_request_target nor the + # privileged slash-command call may run code from the untrusted PR head. - name: Checkout trusted merge planner if: steps.check.outputs.has_ready == 'true' uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 with: - ref: ${{ github.event.pull_request.base.sha }} + ref: ${{ steps.check.outputs.base_sha }} persist-credentials: false - name: Collect changed paths if: steps.check.outputs.has_ready == 'true' env: GH_TOKEN: ${{ github.token }} - PR_NUMBER: ${{ github.event.pull_request.number }} + PR_NUMBER: ${{ steps.check.outputs.pr_number }} EXPECTED_CHANGED_FILES: ${{ steps.check.outputs.changed_files }} run: | set -euo pipefail @@ -105,18 +209,18 @@ jobs: if: steps.check.outputs.has_ready == 'true' env: GH_TOKEN: ${{ github.token }} - PR_SHA: ${{ github.event.pull_request.head.sha }} - PR_NUMBER: ${{ github.event.pull_request.number }} + PR_SHA: ${{ steps.check.outputs.head_sha }} + PR_NUMBER: ${{ steps.check.outputs.pr_number }} run: bash .github/scripts/gate_full_suite.sh - name: Trigger Buildkite merge gate if: steps.check.outputs.has_ready == 'true' env: BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }} - PR_SHA: ${{ github.event.pull_request.head.sha }} - PR_BRANCH: ${{ github.event.pull_request.head.ref }} - PR_NUMBER: ${{ github.event.pull_request.number }} - PR_TITLE: ${{ github.event.pull_request.title }} + PR_SHA: ${{ steps.check.outputs.head_sha }} + PR_BRANCH: ${{ steps.check.outputs.head_ref }} + PR_NUMBER: ${{ steps.check.outputs.pr_number }} + PR_TITLE: ${{ steps.check.outputs.title }} BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }} BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }} MERGE_TEST_PLAN: ${{ steps.plan.outputs.merge_test_plan }} diff --git a/docs/contributing/ci_architecture.md b/docs/contributing/ci_architecture.md index d4868e957..2dd2d06d5 100644 --- a/docs/contributing/ci_architecture.md +++ b/docs/contributing/ci_architecture.md @@ -211,7 +211,8 @@ The Buildkite agent and Slurm have deliberately separate responsibilities: ```text /merge PR comment - -> GitHub verifies write permission and refreshes the `ready` label + -> GitHub verifies write permission, adds the `ready` label, and directly + calls the trusted base-branch merge-gate workflow -> base-branch `ci-trigger-full-suite` workflow fetches the PR file list, computes MERGE_TEST_PLAN plus focused golden/SSIM basenames, and gates on cheap checks diff --git a/docs/contributing/pull_requests.md b/docs/contributing/pull_requests.md index 02f682916..b581b4edf 100644 --- a/docs/contributing/pull_requests.md +++ b/docs/contributing/pull_requests.md @@ -67,7 +67,8 @@ filters, and workflow files. 3. Fix pre-commit failures locally with `pre-commit run --all-files`. 4. Wait for at least one approving review. 5. When the PR is approved and ready, comment `/merge`. -6. `/merge` adds `ready`, waits for cheap checks, and triggers the minimal +6. `/merge` adds `ready`, directly calls the trusted merge-gate workflow, + waits for cheap checks, and triggers the minimal path-relevant integration lanes for the PR branch. 7. If all required checks pass, Mergify squash-merges the PR to `main`. 8. If the merge gate fails, fix the regression, push again, and re-run diff --git a/fastvideo/tests/contract/test_ci_test_collection.py b/fastvideo/tests/contract/test_ci_test_collection.py index 60c1fc4dc..92398e5cc 100644 --- a/fastvideo/tests/contract/test_ci_test_collection.py +++ b/fastvideo/tests/contract/test_ci_test_collection.py @@ -145,23 +145,98 @@ def test_all_gpu_ci_routes_use_the_trusted_slurm_dispatcher(): def test_merge_comment_has_one_change_aware_trigger_path(): slash_commands = (REPO_ROOT / ".github/workflows/ci-slash-commands.yml").read_text() - merge_job = slash_commands.split(" parse-command:", maxsplit=1)[0] + merge_jobs = slash_commands.split(" parse-command:", maxsplit=1)[0] ready_workflow = (REPO_ROOT / ".github/workflows/ci-trigger-full-suite.yml").read_text() - assert "labels: ['ready']" in merge_job - assert "api.buildkite.com" not in merge_job - assert "BUILDKITE_API_TOKEN" not in merge_job + assert "labels: ['ready']" in merge_jobs + assert "removeLabel" not in merge_jobs + assert "continue-on-error: true" in merge_jobs + assert "needs: handle-merge" in merge_jobs + assert "if: needs.handle-merge.result == 'success'" in merge_jobs + assert "uses: ./.github/workflows/ci-trigger-full-suite.yml" in merge_jobs + assert "pr_number: ${{ github.event.issue.number }}" in merge_jobs + assert "BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}" in merge_jobs + assert "actions: write" not in merge_jobs + assert "api.buildkite.com" not in merge_jobs + assert "workflow_call:" in ready_workflow + # Both entry points share one concurrency group, and the newest trigger wins: + # a `/merge` call must not queue a second serialized gate run. The group is + # job-level so an unrelated `labeled` event (which skips the guarded job) + # cannot cancel an in-flight gate and then skip its replacement. + workflow_header = ready_workflow.split("\njobs:", maxsplit=1)[0] + assert "concurrency:" not in workflow_header + trigger_job = ready_workflow.split("\n trigger:", maxsplit=1)[1] + assert "concurrency:" in trigger_job + assert "group: merge-gate-${{ inputs.pr_number || github.event.pull_request.number }}" in trigger_job + assert "cancel-in-progress: true" in trigger_job + assert "BUILDKITE_API_TOKEN:\n required: true" in ready_workflow + assert "CALLED_PR_NUMBER: ${{ inputs.pr_number }}" in ready_workflow + assert "Number.isSafeInteger(prNumber)" in ready_workflow + assert "pull_number: prNumber" in ready_workflow + assert "pr.state !== 'open'" in ready_workflow + assert "pr.base.repo.full_name !== context.payload.repository.full_name" in ready_workflow + assert "pr.base.ref !== context.payload.repository.default_branch" in ready_workflow + assert "core.setOutput('head_sha', pr.head.sha)" in ready_workflow + assert "core.setOutput('base_sha', pr.base.sha)" in ready_workflow assert "api.buildkite.com" in ready_workflow assert 'jq -r --arg pr_number "$PR_NUMBER"' in ready_workflow assert '.env.PR_NUMBER? == $pr_number' in ready_workflow assert 'TEST_SCOPE: "merge"' in ready_workflow assert 'FULL_SUITE: "true"' in ready_workflow assert "plan_merge_ci.py" in ready_workflow - assert "github.event.pull_request.base.sha" in ready_workflow + assert "ref: ${{ steps.check.outputs.base_sha }}" in ready_workflow + assert "PR_SHA: ${{ steps.check.outputs.head_sha }}" in ready_workflow + assert "PR_NUMBER: ${{ steps.check.outputs.pr_number }}" in ready_workflow assert "MERGE_TEST_PLAN" in ready_workflow assert "MERGE_GOLDEN_TESTS" in ready_workflow assert "MERGE_SSIM_TESTS" in ready_workflow assert "__FASTVIDEO_CI_PLAN_ALL__" in ready_workflow + downstream_steps = ready_workflow.split(" - name: Cancel previous Buildkite builds", maxsplit=1)[1] + assert "github.event.pull_request" not in downstream_steps + + +def test_merge_gate_buildkite_cancellation_is_best_effort_and_strict(): + workflow = yaml.safe_load((REPO_ROOT / ".github/workflows/ci-trigger-full-suite.yml").read_text()) + steps = workflow["jobs"]["trigger"]["steps"] + cancel_step = next(step for step in steps if step.get("name") == "Cancel previous Buildkite builds") + trigger_step = next(step for step in steps if step.get("name") == "Trigger Buildkite merge gate") + cancel_script = cancel_step["run"] + + # Stale-build cleanup is best effort; creating the replacement gate remains + # a hard failure so merge protection can never pass without a test build. + assert cancel_step["continue-on-error"] is True + assert cancel_step["timeout-minutes"] == 3 + assert trigger_step.get("continue-on-error") is not True + assert "curl -sS --fail-with-body -X POST" in trigger_step["run"] + + assert cancel_script.splitlines()[0] == "set -euo pipefail" + assert cancel_script.count("--connect-timeout 5") == 2 + assert cancel_script.count("--max-time 20") == 2 + assert "curl -sS --fail-with-body --connect-timeout 5 --max-time 20 --get" in cancel_script + assert cancel_script.count('--data-urlencode "state[]=running"') == 1 + assert cancel_script.count('--data-urlencode "state[]=scheduled"') == 1 + assert cancel_script.count('--data-urlencode "state[]=failing"') == 1 + assert cancel_script.count('--data-urlencode "exclude_jobs=true"') == 1 + assert cancel_script.count('--data-urlencode "exclude_pipeline=true"') == 1 + assert 'state=running,scheduled' not in cancel_script + + # Only a validated array of build records can reach the URL construction. + assert 'if type != "array" then false' in cancel_script + assert 'if type != "object" then false' in cancel_script + assert 'if type == "number" then . > 0 and floor == .' in cancel_script + assert 'if ($env | type) != "object" then false' in cancel_script + assert '$env.TEST_SCOPE?' in cancel_script + assert '$env.PR_NUMBER?' in cancel_script + + # API bodies are kept out of annotations, every cancellation is checked, + # and one failure does not stop attempts for the remaining build numbers. + assert '--output "$response_file"' in cancel_script + assert 'cat "$response_file"' not in cancel_script + assert ("if ! curl -sS --fail-with-body --connect-timeout 5 --max-time 20 " + "-o /dev/null -X PUT") in cancel_script + assert "cancellation_failed=1" in cancel_script + assert cancel_script.index("cancellation_failed=1") < cancel_script.index('done < "$builds_file"') + assert cancel_script.index('done < "$builds_file"') < cancel_script.index("if (( cancellation_failed != 0 ))") def test_full_ssim_has_a_weekly_slurm_schedule(): From 274b922d399880f617017372bc8130e4373cd165 Mon Sep 17 00:00:00 2001 From: Kyle Hu <160547796+KyleNeverGivesUp@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:29:56 -0700 Subject: [PATCH 058/151] [misc]: surface GenerationResult.peak_memory_mb in the MiniMax H3 examples (#1785) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- .../installation/spark_performance.md | 10 ++- examples/inference/basic/README.md | 2 +- examples/inference/basic/basic_fasth3.py | 5 ++ .../inference/basic/basic_minimax_h3_t2v.py | 7 +++ .../inference/test_basic_fasth3_profile.py | 61 +++++++++++++++++-- 5 files changed, 77 insertions(+), 8 deletions(-) diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index c16dc9e9a..9f1a87dd7 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -79,9 +79,13 @@ The GB10 has **no separate VRAM** — CPU and GPU share one 128 GB LPDDR5X pool - **`nvidia-smi` reports memory as `[N/A]`** on the GB10, and the system "used" figure conflates CPU + GPU + cache, so it's only a soft upper bound — treat the - whole 128 GB as one shared budget. For a per-run figure, use FastVideo's own - `peak_memory_mb` (reported on the generation result and by the performance - benchmark), which is measured inside the worker that runs the model. + whole 128 GB as one shared budget. For a worker-measured figure, use FastVideo's + own `peak_memory_mb` (reported on the generation result and by the performance + benchmark), which is measured inside the worker that runs the model. It is the + worker's allocator high-water mark, never reset between generations, so it + includes model load and every prior run rather than a single run's delta. The + `mp` backend populates it; the Ray backend leaves `peak_memory_mb` unset, so + `--execution-backend ray` prints no peak-memory figure. - **The 128 GB is a *working-set* ceiling, not storage** — the model cache lives on the NVMe (3.7 TB, ample). What has to fit in 128 GB is the weights, activations, and KV cache — and, critically, the **VAE decode buffers**, which diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index fcc5b0696..099a0fcf9 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -122,7 +122,7 @@ python examples/inference/basic/basic_fasth3.py \ --warmup-seed 999 ``` -`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. `--h3-sequential-load` / `--no-h3-sequential-load` override the auto split that releases Qwen3-VL before DiT/VAE load (on by default on GB10, off on discrete GPUs). The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median. +`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. `--h3-sequential-load` / `--no-h3-sequential-load` override the auto split that releases Qwen3-VL before DiT/VAE load (on by default on GB10, off on discrete GPUs). The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time, the worker's peak memory, and a warmup-excluded median. `Peak memory` is the worker's allocator high-water mark: it is never reset between generations, so it includes model load and every prior run, and only rank 0's value is reported. It is populated by the `mp` backend only: the Ray backend leaves `peak_memory_mb` unset, so `--execution-backend ray` prints no `Peak memory` line. One script covers each validated duration; regional compile is the fastest measured DiT route for all three: diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 02f1e918b..368684599 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -383,6 +383,11 @@ def run(args: argparse.Namespace) -> list[float]: generation_time = getattr(result, "generation_time", None) if generation_time is not None: print(f"Generation time: {float(generation_time):.3f}s") + # Worker-lifetime allocator high-water mark: torch never resets it + # between generations, so it includes model load and every prior run. + peak_memory_mb = getattr(result, "peak_memory_mb", None) + if peak_memory_mb is not None: + print(f"Peak memory: {float(peak_memory_mb):.1f} MB") denoise_time = _denoise_seconds(result) if denoise_time is not None: measured_denoise_times.append(denoise_time) diff --git a/examples/inference/basic/basic_minimax_h3_t2v.py b/examples/inference/basic/basic_minimax_h3_t2v.py index d4224f116..30e4db6cd 100644 --- a/examples/inference/basic/basic_minimax_h3_t2v.py +++ b/examples/inference/basic/basic_minimax_h3_t2v.py @@ -139,10 +139,17 @@ def main() -> None: # machine-readable: benchmark harnesses parse this line to separate # generation from model-load time (last occurrence = steady state) print(f"Generation time: {result.generation_time:.2f}s") + # Worker-lifetime allocator high-water mark: torch never resets it + # between generations, so it includes model load and every prior run. + if result.peak_memory_mb is not None: + print(f"Peak memory: {result.peak_memory_mb:.1f} MB") for _ in range(args.repeats - 1): result = generator.generate(request) if result.generation_time is not None: print(f"Generation time: {result.generation_time:.2f}s") + # Same worker-lifetime peak as above; it does not reset per repeat. + if result.peak_memory_mb is not None: + print(f"Peak memory: {result.peak_memory_mb:.1f} MB") finally: generator.shutdown() diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index 7fa7c7431..0098bbc8a 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util from pathlib import Path +import sys from types import SimpleNamespace import pytest @@ -12,17 +13,19 @@ from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler REPO_ROOT = Path(__file__).resolve().parents[3] EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_fasth3.py" +MINIMAX_EXAMPLE_PATH = REPO_ROOT / "examples" / "inference" / "basic" / "basic_minimax_h3_t2v.py" -def _load_example(): - spec = importlib.util.spec_from_file_location("basic_fasth3_contract", EXAMPLE_PATH) +def _load_example(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module -fasth3 = _load_example() +fasth3 = _load_example("basic_fasth3_contract", EXAMPLE_PATH) +minimax_h3_t2v = _load_example("basic_minimax_h3_t2v_contract", MINIMAX_EXAMPLE_PATH) def _args(*overrides: str): @@ -204,7 +207,8 @@ def test_selected_fast_profile_requires_its_optional_routes(monkeypatch): fasth3.validate_profile_dependencies(_args()) -def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp_path, capsys): +@pytest.mark.parametrize("peak_memory_mb", (42.0, None)) +def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp_path, capsys, peak_memory_mb): calls = [] class FakeGenerator: @@ -216,6 +220,7 @@ def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp return SimpleNamespace( video_path=request.output.output_path, generation_time=1.25, + peak_memory_mb=peak_memory_mb, logging_info=SimpleNamespace(stages={"denoising": {"execution_time": 2.5}}), ) @@ -256,11 +261,59 @@ def test_run_excludes_warmup_and_uses_distinct_measured_outputs(monkeypatch, tmp assert f"Warmup output written to: {tmp_path / '_fasth3_warmup.mp4'}" in output assert output.count("Output written to:") == 3 assert output.count("Denoising time: 2.500s") == 3 + if peak_memory_mb is None: + # The Ray backend leaves peak_memory_mb unset: stay silent, never print "None MB". + assert "Peak memory:" not in output + else: + assert output.count("Peak memory: 42.0 MB") == 3 assert "Measured E2E wall times (n=3, warmup excluded): [6.0, 7.0, 8.0]" in output assert "Median E2E wall time: 7.000s" in output assert "Median denoising time: 2.500s" in output +@pytest.mark.parametrize("peak_memory_mb", (42.0, None)) +def test_minimax_h3_t2v_reports_peak_memory_only_when_present(monkeypatch, tmp_path, capsys, peak_memory_mb): + generated = [] + + class FakeGenerator: + + def generate(self, request): + generated.append(request) + return SimpleNamespace( + video_path=request.output.output_path, + generation_time=1.25, + peak_memory_mb=peak_memory_mb, + ) + + def shutdown(self): + pass + + class FakeVideoGenerator: + + @classmethod + def from_config(cls, config): + return FakeGenerator() + + monkeypatch.setattr(minimax_h3_t2v, "VideoGenerator", FakeVideoGenerator) + monkeypatch.setattr( + sys, + "argv", + ["basic_minimax_h3_t2v.py", "--prompt", "a test prompt", "--output", + str(tmp_path), "--repeats", "2"], + ) + + minimax_h3_t2v.main() + + assert len(generated) == 2 + output = capsys.readouterr().out + assert output.count("Output written to:") == 1 + assert output.count("Generation time: 1.25s") == 2 + if peak_memory_mb is None: + assert "Peak memory:" not in output + else: + assert output.count("Peak memory: 42.0 MB") == 2 + + def test_taeh3_backend_is_opt_in_experimental(): default = fasth3.build_generator_config(_args()) taeh3 = fasth3.build_generator_config(_args("--video-decode-backend", "taeh3")) From abe98e87f529f6ae44f3aadb7020a8a5145f9566 Mon Sep 17 00:00:00 2001 From: Jack Chen Date: Wed, 7 Oct 2026 16:30:11 +0800 Subject: [PATCH 059/151] [feat] Add phase-aware modular training performance metrics (#1826) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- docs/training/train_infra.md | 79 ++++ examples/train/configs/example.yaml | 7 + .../train/trainer/test_performance_wiring.py | 172 ++++++++ fastvideo/tests/train/utils/test_config.py | 17 + .../tests/train/utils/test_performance.py | 252 ++++++++++++ fastvideo/train/trainer.py | 41 +- fastvideo/train/utils/config.py | 15 + fastvideo/train/utils/performance.py | 382 ++++++++++++++++++ fastvideo/train/utils/training_config.py | 11 + 9 files changed, 975 insertions(+), 1 deletion(-) create mode 100644 fastvideo/tests/train/trainer/test_performance_wiring.py create mode 100644 fastvideo/tests/train/utils/test_performance.py create mode 100644 fastvideo/train/utils/performance.py diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 766d261d2..a37689636 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -151,6 +151,10 @@ training: project_name: my_project run_name: my_run + performance: + enabled: true + peak_tflops_per_gpu: null # optional dense BF16 peak for one GPU + model: weighting_scheme: uniform # uniform, logit_normal, mode precondition_outputs: false @@ -177,6 +181,81 @@ The repeat count duplicates that dataset's parquet file list before shuffling/sa See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab, including SwanLab installation and authentication. +### Training performance metrics + +The modular trainer logs low-overhead performance metrics on every optimizer +step. The Transformer boundary is instrumented separately for each model role, +so the accounting includes all forwards performed by the selected method: + +- bidirectional SFT records its dense or VSA student forward; +- causal SFT records block-causal attention geometry; +- streaming/self-forcing records every rollout chunk and KV-cache update; +- DMD records student, critic, and teacher forwards independently, including + repeated rollouts and the conditional/unconditional teacher passes. + +The common metrics are: + +| Metric | Meaning | +|--------|---------| +| `step_time_sec` | Training wall time for one optimizer step, excluding checkpoint and validation callbacks | +| `perf/steps_per_sec` | Reciprocal of `step_time_sec` | +| `perf/samples_per_sec` | Configured global samples processed per second, including gradient accumulation and data-parallel replicas | +| `perf/model_forward_calls` | Actual Transformer invocations during the optimizer step | +| `perf/causal_chunks` | Block-causal chunks represented by those invocations | +| `perf/query_latent_frames_per_sec` | Latent query frames processed across all roles and rollouts | +| `perf/query_tokens_per_sec` | Patch tokens processed across all roles and rollouts | +| `perf/attention_density` | Effective self-attention pairs divided by dense attention pairs | +| `perf/estimated_tflops_per_gpu` | Estimated useful model FLOP/s per GPU | +| `perf/estimated_mfu` | Estimated useful FLOP/s divided by the dense BF16 peak (ratio, not percent) | +| `perf/peak_tflops_per_gpu` | Dense BF16 peak used for the MFU ratio (configured or inferred) | +| `perf/role//*` | Forward count, grad-carrying forward count, causal chunks, and estimated TFLOP/s for one role | + +For DMD2, combine these metrics with the existing `update_student` metric to +separate the expensive generator-update steps from critic-only steps. + +`perf/samples_per_sec` uses the configured `train_batch_size`, gradient +accumulation, and data-parallel replica count. Methods that manage their own +optimization (for example DiffusionNFT, whose outer step consumes +`num_batches_per_epoch` dataloader batches) report the throughput of a single +batch in the denominator, so use the method's own sample counter +(`nft/num_sampled`) for those runs. The FLOP-derived metrics still aggregate +every forward the method performs. + +MFU is an analytic Transformer-core estimate. Each no-grad forward contributes +`1F`; a forward whose result carries autograd contributes `3F` (forward plus an +approximately `2F` backward). Activation-checkpoint recomputation is excluded, +as expected for model FLOPs utilization. Wan VSA uses the kernel's clamped +tile-top-k density and includes its gate projection and pooled-attention +overhead. VMOBA runs are modeled as dense self-attention: the MoBA top-k +selection is not estimated, so `perf/attention_density` stays at `1.0` and the +attention-dependent FLOPs are an upper bound. Causal Wan uses the configured +chunk size, local window, and actual streaming cache position; when +`local_attn_size` is unset, the modeled window is the transformer's 21-frame +compatibility cap, and `pipeline.dit_config.sliding_window_num_frames` only +sizes the streaming KV cache. MatrixGame2's 15-frame compatibility window is +not modeled, so its attention density is an upper bound. Embedding, +normalization, optimizer, communication, MatrixGame action modules, and other +non-core work are not included, so MFU is an estimate rather than a +hardware-profiler measurement. + +`perf/estimated_tflops_per_gpu` and `perf/estimated_mfu` aggregate every role +(student, teacher, critic, EMA) that ran during the step, so they are not +directly comparable to a single-model MFU number. Replica scaling assumes +`world_size = data_parallel x sp_size`; `tp_size > 1` is not modeled and would +inflate the per-GPU estimates. + +Forward counts and wall-clock throughput work for every modular model. The +FLOP-derived metrics currently require a Wan-style architecture exposing +`hidden_size`, `ffn_dim`, `num_layers`, and a three-axis `patch_size`; they are +omitted for other architectures instead of reporting a misleading estimate. + +Known NVIDIA accelerators use an inferred dense BF16 peak. Set +`training.performance.peak_tflops_per_gpu` explicitly for a different board +form factor or clock. If the device is unknown and no peak is configured, the +trainer still logs throughput, call counts, attention density, and estimated +TFLOP/s, but omits `perf/estimated_mfu`. Set `enabled: false` to disable the +forward hooks and all `perf/*` metrics. + ### `callbacks` — Pluggable hooks Callbacks run at specific points in the training loop (before/after optimizer diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 51dc03e66..af64fd45a 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -130,6 +130,13 @@ training: project_name: my_project # default: "fastvideo" run_name: my_run # default: "" + # --- training.performance [TYPED] -> PerformanceConfig --- + performance: + enabled: true # default: true + # Optional dense BF16 peak for one GPU. Known NVIDIA device names are + # inferred automatically; set this for non-standard clocks/form factors. + peak_tflops_per_gpu: null # e.g. 989.5 for H100 SXM + # --- training.vsa [TYPED] -> VSAConfig --- vsa: sparsity: 0.0 # default: 0.0 (0.0 = disabled) diff --git a/fastvideo/tests/train/trainer/test_performance_wiring.py b/fastvideo/tests/train/trainer/test_performance_wiring.py new file mode 100644 index 000000000..5fd3ff3b7 --- /dev/null +++ b/fastvideo/tests/train/trainer/test_performance_wiring.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for Trainer performance-metric wiring.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch + +from fastvideo.train.trainer import Trainer +from fastvideo.train.utils.performance import TrainingPerformanceMonitor +from fastvideo.train.utils.training_config import TrainingConfig + + +class _FakeTransformer(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace( + arch_config=SimpleNamespace( + hidden_size=8, + ffn_dim=16, + num_layers=2, + patch_size=(1, 2, 2), + )) + self.weight = torch.nn.Parameter(torch.ones(())) + + def forward( + self, + *, + hidden_states: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + return hidden_states * self.weight + + +class _FakeRoleModel: + + def __init__(self) -> None: + self.transformer = _FakeTransformer() + + +class _DummyTracker: + + def __init__(self) -> None: + self.logs: list[tuple[dict[str, float], int]] = [] + + def log(self, metrics: dict[str, float], step: int) -> None: + self.logs.append((metrics, step)) + + def finish(self) -> None: + pass + + +class _PerformanceMethod: + + def __init__(self) -> None: + self._role_models = {"student": _FakeRoleModel()} + self.weight = torch.nn.Parameter(torch.tensor(1.0)) + self.tracker = None + self.pre_loop_forwards = 0 + + def set_tracker(self, tracker: Any) -> None: + self.tracker = tracker + + def on_train_start(self) -> None: + # Stands in for teacher-cache generation or step-0 validation. + self._forward() + self.pre_loop_forwards += 1 + + def manages_optimization(self) -> bool: + return False + + def _forward(self) -> torch.Tensor: + hidden = torch.ones(1, 3, 2, 4, 4) + return self._role_models["student"].transformer(hidden_states=hidden) + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]: + del batch, iteration + output = self._forward() + return {"total_loss": output.sum() * 0.0 + 1.0}, {}, {} + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int, + ) -> None: + del outputs + (loss_map["total_loss"] / grad_accum_rounds).backward() + + def optimizers_schedulers_step(self, iteration: int) -> None: + del iteration + + def optimizers_zero_grad(self, iteration: int) -> None: + del iteration + self.weight.grad = None + + +def _run_trainer( + monkeypatch: Any, + *, + performance_enabled: bool, +) -> tuple[_PerformanceMethod, _DummyTracker]: + tracker = _DummyTracker() + group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1) + + monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group) + monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group) + monkeypatch.setattr( + "fastvideo.train.trainer.build_tracker", + lambda *args, **kwargs: tracker, + ) + + cfg = TrainingConfig() + cfg.tracker.project_name = "" + cfg.loop.gradient_accumulation_steps = 1 + cfg.performance.enabled = performance_enabled + trainer = Trainer(cfg) + method = _PerformanceMethod() + + trainer.run( + method, + dataloader=[{ + "sample": "x" + }], + max_steps=1, + ) + return method, tracker + + +def test_pre_loop_forwards_are_not_counted(monkeypatch) -> None: + method, tracker = _run_trainer(monkeypatch, performance_enabled=True) + + assert method.pre_loop_forwards == 1 + metrics = tracker.logs[-1][0] + # Only the forward inside the training step is measured; the forward run + # during on_train_start() was dropped by the pre-loop reset. + assert metrics["perf/model_forward_calls"] == 1.0 + assert metrics["perf/role/student/forward_calls"] == 1.0 + + +def test_pre_loop_forwards_are_dropped_before_the_first_step(monkeypatch) -> None: + observed: list[int] = [] + original_reset = TrainingPerformanceMonitor.reset + + def recording_reset(self: TrainingPerformanceMonitor) -> None: + observed.append(len(self._work)) + original_reset(self) + + monkeypatch.setattr(TrainingPerformanceMonitor, "reset", recording_reset) + method, _ = _run_trainer(monkeypatch, performance_enabled=True) + + assert method.pre_loop_forwards == 1 + # The monitor must already be empty when the first training step starts, so + # a long pre-loop phase (teacher cache generation, step-0 validation) + # cannot accumulate per-forward records. + assert observed[-1] == 0 + + +def test_disabled_performance_logs_no_perf_metrics(monkeypatch) -> None: + _, tracker = _run_trainer(monkeypatch, performance_enabled=False) + + metrics = tracker.logs[-1][0] + assert not any(key.startswith("perf/") for key in metrics) diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index 8471eea8c..a15c52027 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -81,6 +81,9 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.tracker.project_name == "fastvideo" assert t.tracker.run_name == "" + assert t.performance.enabled is True + assert t.performance.peak_tflops_per_gpu is None + assert t.model.weighting_scheme == "uniform" assert t.model.precondition_outputs is False assert t.model.moba_config == {} @@ -134,6 +137,10 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "project_name": "myproj", "run_name": "myrun", }, + "performance": { + "enabled": True, + "peak_tflops_per_gpu": 989.5, + }, "vsa": { "sparsity": 0.5 }, @@ -172,6 +179,9 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.tracker.trackers == ["wandb"] assert t.tracker.project_name == "myproj" + assert t.performance.enabled is True + assert t.performance.peak_tflops_per_gpu == pytest.approx(989.5) + assert t.vsa_sparsity == pytest.approx(0.5) assert t.model.weighting_scheme == "logit_normal" assert t.model.precondition_outputs is True @@ -223,6 +233,13 @@ def test_missing_config_file_raises(tmp_path: Path) -> None: load_run_config(str(tmp_path / "does_not_exist.yaml")) +def test_non_positive_peak_tflops_raises(tmp_path: Path) -> None: + data = _minimal_yaml() + data["training"] = {"performance": {"peak_tflops_per_gpu": 0}} + with pytest.raises(ValueError, match="peak_tflops_per_gpu"): + load_run_config(_write_yaml(tmp_path, data)) + + # --------------------------------------------------------------------------- # Special parsing # --------------------------------------------------------------------------- diff --git a/fastvideo/tests/train/utils/test_performance.py b/fastvideo/tests/train/utils/test_performance.py new file mode 100644 index 000000000..0d9d08022 --- /dev/null +++ b/fastvideo/tests/train/utils/test_performance.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for modular training performance accounting.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.train.utils.performance import ( + TrainingPerformanceMonitor, + _blockwise_causal_frame_pairs, + _teacher_forcing_frame_pairs, + estimate_transformer_forward, + infer_peak_bf16_tflops, +) + + +def _arch() -> SimpleNamespace: + return SimpleNamespace( + hidden_size=8, + ffn_dim=16, + num_layers=2, + patch_size=(1, 2, 2), + ) + + +class _FakeTransformer(torch.nn.Module): + + def __init__(self, *, causal: bool = False) -> None: + super().__init__() + self.config = SimpleNamespace(arch_config=_arch()) + self.weight = torch.nn.Parameter(torch.ones(())) + if causal: + self.num_frame_per_block = 2 + self.local_attn_size = 3 + self.text_len = 4 + + def forward( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + del encoder_hidden_states, kwargs + return hidden_states * self.weight + + +class _UnsupportedTransformer(torch.nn.Module): + + def forward(self, *, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +def _kwargs(*, frames: int = 4) -> dict[str, torch.Tensor]: + return { + "hidden_states": torch.ones(2, 3, frames, 4, 4), + "encoder_hidden_states": torch.ones(2, 3, 8), + } + + +def test_dense_forward_flops_and_backward_factor() -> None: + module = _FakeTransformer() + output = torch.ones(2, 3, 4, 4, 4, requires_grad=True) + + work = estimate_transformer_forward( + module, + _kwargs(), + output, + role="student", + attention_metadata=None, + ) + + assert work is not None + # B=2, L=16, C=3, d=8, ffn=16, layers=2. + expected_per_layer = ( + 8 * 2 * 16 * 8**2 + + 4 * 2 * 16 * 8 * 16 + + 4 * 2 * 16 * 8**2 + + 4 * 2 * 3 * 8**2 + + 4 * 2 * 16 * 3 * 8 + + 4 * 2 * 16**2 * 8 + ) + assert work.forward_flops == expected_per_layer * 2 + assert work.useful_flops == work.forward_flops * 3 + assert work.query_frames == 8 + assert work.query_tokens == 32 + assert work.attention_pairs == 2 * 16**2 + + +def test_causal_chunk_and_teacher_forcing_pair_counts() -> None: + assert _blockwise_causal_frame_pairs(5, 2, -1) == 17 + assert _blockwise_causal_frame_pairs(5, 2, 3) == 13 + assert _teacher_forcing_frame_pairs(4, 2) == 24 + + +def test_streaming_causal_forward_uses_cache_window() -> None: + module = _FakeTransformer(causal=True) + kwargs = _kwargs(frames=2) + # Four spatial tokens/frame, starting after two frames. The three-frame + # local window gives 8 query tokens x 12 key tokens. + kwargs.update({"kv_cache": [{}], "current_start": 8}) + + work = estimate_transformer_forward( + module, + kwargs, + torch.ones(2, 3, 2, 4, 4), + role="student", + attention_metadata=None, + ) + + assert work is not None + assert work.causal_chunks == 1 + assert work.attention_pairs == 2 * 8 * 12 + assert work.dense_attention_pairs == 2 * 8 * 16 + + +def test_vsa_uses_kernel_topk_density_and_overhead() -> None: + module = _FakeTransformer() + metadata = SimpleNamespace( + variable_block_sizes=torch.tensor([4, 4, 4, 4]), + VSA_sparsity=0.75, + ) + + work = estimate_transformer_forward( + module, + _kwargs(), + torch.ones(2, 3, 4, 4, 4), + role="student", + attention_metadata=metadata, + ) + + assert work is not None + assert work.attention_pairs == pytest.approx(2 * 16**2 * 0.25) + dense_work = estimate_transformer_forward( + module, + _kwargs(), + torch.ones(2, 3, 4, 4, 4), + role="student", + attention_metadata=None, + ) + assert dense_work is not None + # VSA includes its gate projection and pooled dense attention overhead. + assert work.forward_flops > dense_work.forward_flops - ( + dense_work.attention_pairs - work.attention_pairs) * 4 * 8 * 2 + + +def test_streaming_cross_attention_cache_skips_kv_projections() -> None: + module = _FakeTransformer(causal=True) + kwargs = _kwargs(frames=2) + kwargs.update({"kv_cache": [{}], "current_start": 8}) + output = torch.ones(2, 3, 2, 4, 4) + + uncached = estimate_transformer_forward( + module, + kwargs, + output, + role="student", + attention_metadata=None, + ) + cached = estimate_transformer_forward( + module, + kwargs, + output, + role="student", + attention_metadata=None, + cross_attention_cached=True, + ) + + assert uncached is not None and cached is not None + expected_saved = 4 * 2 * 4 * 8**2 * 2 + assert uncached.forward_flops - cached.forward_flops == expected_saved + + +def test_monitor_counts_multiple_rollouts_and_roles_without_double_counting_accumulation() -> None: + models = { + role: SimpleNamespace(transformer=_FakeTransformer()) + for role in ("student", "teacher", "critic") + } + monitor = TrainingPerformanceMonitor() + monitor.attach(models) + monitor.reset() + + with set_forward_context(current_timestep=0, attn_metadata=None): + models["student"].transformer(**_kwargs()) + models["student"].transformer(**_kwargs()) + models["critic"].transformer(**_kwargs()) + with torch.no_grad(): + models["teacher"].transformer(**_kwargs()) + models["teacher"].transformer(**_kwargs()) + + metrics = monitor.metrics( + step_time_sec=2.0, + local_batch_size=2, + grad_accum=4, + world_size=8, + sp_size=4, + peak_tflops_per_gpu=100.0, + ) + monitor.close() + + assert metrics["perf/model_forward_calls"] == 5 + assert metrics["perf/causal_chunks"] == 0 + assert metrics["perf/role/student/forward_calls"] == 2 + assert metrics["perf/role/teacher/forward_calls"] == 2 + assert metrics["perf/role/teacher/backward_forwards"] == 0 + assert metrics["perf/role/critic/backward_forwards"] == 1 + # 2 local samples x DP=2 x grad-accum=4 / 2 seconds. + assert metrics["perf/samples_per_sec"] == 8 + assert metrics["perf/estimated_mfu"] > 0 + + +def test_monitor_counts_forward_when_flop_estimator_does_not_support_model() -> None: + monitor = TrainingPerformanceMonitor() + model = SimpleNamespace(transformer=_UnsupportedTransformer()) + monitor.attach({"student": model}) + monitor.reset() + + model.transformer(hidden_states=torch.ones(1)) + metrics = monitor.metrics( + step_time_sec=1.0, + local_batch_size=1, + grad_accum=1, + world_size=1, + sp_size=1, + peak_tflops_per_gpu=None, + ) + monitor.close() + + assert metrics["perf/model_forward_calls"] == 1 + assert metrics["perf/role/student/forward_calls"] == 1 + assert "perf/estimated_tflops_per_gpu" not in metrics + + +@pytest.mark.parametrize( + ("device_name", "expected"), + [ + ("NVIDIA A40", 149.7), + ("NVIDIA L40S", 362.05), + ("NVIDIA H100 80GB HBM3", 989.5), + ("NVIDIA H200", 989.5), + ("NVIDIA B200", 2250.0), + ("NVIDIA GB200", 2500.0), + ("Some Future GPU", None), + ], +) +def test_infer_peak_bf16_tflops(device_name: str, expected: float | None) -> None: + assert infer_peak_bf16_tflops(device_name) == expected diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index be6054f77..adb5cd7e5 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -13,6 +13,10 @@ from tqdm.auto import tqdm from fastvideo.distributed import get_sp_group, get_world_group from fastvideo.train.callbacks.callback import CallbackDict from fastvideo.train.methods.base import LogScalar, TrainingMethod +from fastvideo.train.utils.performance import ( + TrainingPerformanceMonitor, + infer_peak_bf16_tflops, +) from fastvideo.train.utils.tracking import build_tracker if TYPE_CHECKING: @@ -72,6 +76,7 @@ class Trainer: callback_configs or {}, training_config, ) + self.performance_monitor = TrainingPerformanceMonitor() def _iter_dataloader(self, dataloader: Any) -> Iterator[dict[str, Any]]: data_iter = iter(dataloader) @@ -113,6 +118,13 @@ class Trainer: int(tc.loop.gradient_accumulation_steps or 1), ) + performance_enabled = bool(tc.performance.enabled) + peak_tflops = tc.performance.peak_tflops_per_gpu + if performance_enabled: + if peak_tflops is None and torch.cuda.is_available(): + peak_tflops = infer_peak_bf16_tflops(torch.cuda.get_device_name(torch.cuda.current_device())) + self.performance_monitor.attach(getattr(method, "_role_models", {}), ) + method.set_tracker(self.tracker) method.on_train_start() self.callbacks.on_train_start( @@ -141,6 +153,11 @@ class Trainer: # have advanced the RNG as a side-effect. if (checkpoint_manager is not None and resume_from_checkpoint): checkpoint_manager.load_rng_snapshot(resume_from_checkpoint, ) + # Pre-loop forwards (teacher cache generation, step-0 validation) are + # not part of a training step; drop them so a long cache build cannot + # accumulate per-forward records. + if performance_enabled: + self.performance_monitor.reset() progress = tqdm( range(start_step + 1, max_steps + 1), initial=start_step, @@ -151,6 +168,8 @@ class Trainer: method_manages_optimization = bool(method.manages_optimization()) for step in progress: t0 = time.perf_counter() + if performance_enabled: + self.performance_monitor.reset() # Accumulate on GPU during grad-accum; materialise # to CPU once per step right before logging. @@ -219,8 +238,27 @@ class Trainer: divisor = 1 if method_manages_optimization else grad_accum metrics = {k: float(v) / divisor for k, v in loss_sums.items()} metrics.update({k: float(v) / divisor for k, v in metric_sums.items()}) - metrics["step_time_sec"] = (time.perf_counter() - t0) + step_time_sec = time.perf_counter() - t0 + metrics["step_time_sec"] = step_time_sec metrics["vsa_sparsity"] = float(tc.vsa_sparsity) + if performance_enabled: + metrics.update( + self.performance_monitor.metrics( + step_time_sec=step_time_sec, + local_batch_size=tc.data.train_batch_size, + grad_accum=(1 if method_manages_optimization else grad_accum), + world_size=getattr( + self.world_group, + "world_size", + tc.distributed.num_gpus, + ), + sp_size=getattr( + self.sp_group, + "world_size", + tc.distributed.sp_size, + ), + peak_tflops_per_gpu=peak_tflops, + )) if self.global_rank == 0 and metrics: self.tracker.log(metrics, step) @@ -251,4 +289,5 @@ class Trainer: if checkpoint_manager is not None: checkpoint_manager.save_final(max_steps) + self.performance_monitor.close() self.tracker.finish() diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index d396b0258..c34fcb0a0 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -18,6 +18,7 @@ from fastvideo.train.utils.training_config import ( DistributedConfig, ModelTrainingConfig, OptimizerConfig, + PerformanceConfig, TrackerConfig, TrainingLoopConfig, ) @@ -351,6 +352,7 @@ def _build_training_config( lo = dict(t.get("loop", {}) or {}) ck = dict(t.get("checkpoint", {}) or {}) tr = dict(t.get("tracker", {}) or {}) + perf = dict(t.get("performance", {}) or {}) vs = dict(t.get("vsa", {}) or {}) m = dict(t.get("model", {}) or {}) @@ -382,6 +384,15 @@ def _build_training_config( "{'t2v', 't2va', 'text_only'}, got " f"{preprocessed_data_type!r}") + peak_tflops = get_optional_float( + perf, + "peak_tflops_per_gpu", + where="training.performance.peak_tflops_per_gpu", + ) + if peak_tflops is not None and peak_tflops <= 0.0: + raise ValueError("training.performance.peak_tflops_per_gpu must be " + f"> 0, got {peak_tflops}") + return TrainingConfig( distributed=DistributedConfig( num_gpus=num_gpus, @@ -429,6 +440,10 @@ def _build_training_config( project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"), run_name=str(tr.get("run_name", "") or ""), ), + performance=PerformanceConfig( + enabled=bool(perf.get("enabled", True)), + peak_tflops_per_gpu=peak_tflops, + ), vsa_sparsity=float(vs.get("sparsity", 0.0) or 0.0), vsa_cache_tile_buf=bool(vs.get("cache_tile_buf", False) or False), model=ModelTrainingConfig( diff --git a/fastvideo/train/utils/performance.py b/fastvideo/train/utils/performance.py new file mode 100644 index 000000000..bef74566f --- /dev/null +++ b/fastvideo/train/utils/performance.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Low-overhead performance accounting for modular training. + +The monitor counts logical Transformer invocations at the role boundary. This +is important for distillation: one optimizer step may contain several student +rollouts plus critic and teacher forwards. FLOPs are an analytic estimate of +the Transformer core; measured step and sample throughput remain exact wall +clock metrics. +""" + +from __future__ import annotations + +import math +from collections import defaultdict +from dataclasses import dataclass +from typing import Any + +import torch + +from fastvideo.attention.backends.video_sparse_attn import compute_topk +from fastvideo.forward_context import get_forward_context + +_FLOPS_PER_TFLOP = 1.0e12 + + +@dataclass(slots=True) +class _ForwardWork: + role: str + forward_flops: float + useful_flops: float + causal_chunks: int + query_frames: int + query_tokens: int + attention_pairs: float + dense_attention_pairs: float + + +def infer_peak_bf16_tflops(device_name: str) -> float | None: + """Infer dense BF16 tensor-core peak TFLOP/s from a CUDA device name. + + Values intentionally exclude structured sparsity. Explicit configuration + should be used when board form factor or clocks differ from the common + data-center variant. + """ + name = device_name.upper() + # Order matters: GB200 contains B200. + peaks = ( + ("GB200", 2500.0), + ("B200", 2250.0), + ("H200", 989.5), + ("H100", 989.5), + ("L40S", 362.05), + ("A100", 312.0), + ("A40", 149.7), + ) + return next((peak for token, peak in peaks if token in name), None) + + +def _arch_config(module: torch.nn.Module) -> Any | None: + config = getattr(module, "config", None) + return getattr(config, "arch_config", config) + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, list | tuple): + return next((item for item in value if isinstance(item, torch.Tensor)), None) + return None + + +def _output_requires_grad(output: Any) -> bool: + if isinstance(output, torch.Tensor): + return bool(output.requires_grad) + if isinstance(output, dict): + return any(_output_requires_grad(value) for value in output.values()) + if isinstance(output, list | tuple): + return any(_output_requires_grad(value) for value in output) + return False + + +def _blockwise_causal_frame_pairs( + num_frames: int, + frames_per_block: int, + local_attn_size: int, +) -> int: + """Frame-level pairs used by full-sequence block-causal training.""" + total = 0 + block = max(1, frames_per_block) + for start in range(0, num_frames, block): + end = min(start + block, num_frames) + attended_start = 0 if local_attn_size < 0 else max(0, end - local_attn_size) + total += (end - start) * (end - attended_start) + return total + + +def _teacher_forcing_frame_pairs( + num_frames: int, + frames_per_block: int, +) -> int: + """Pairs for causal Wan's concatenated ``[clean | noisy]`` mask. + + Clean and noisy query blocks each attend every frame, which is the same + count as a full block-causal mask with no local window. + """ + return 2 * _blockwise_causal_frame_pairs(num_frames, frames_per_block, -1) + + +def _vsa_attention_pairs( + seq_len: int, + metadata: Any, +) -> tuple[float, int]: + """Return nominal sparse token pairs and tile count. + + The kernel keeps a clamped top-k of key tiles per query tile; the nominal + pair count scales the dense count by that tile fraction. + """ + block_sizes = getattr(metadata, "variable_block_sizes", None) + num_blocks = int(block_sizes.numel()) if block_sizes is not None else 0 + if num_blocks <= 0: + return float(seq_len * seq_len), 0 + sparsity = float(getattr(metadata, "VSA_sparsity", 0.0)) + topk = compute_topk(sparsity, num_blocks) + return float(seq_len * seq_len) * topk / num_blocks, num_blocks + + +def estimate_transformer_forward( + module: torch.nn.Module, + kwargs: dict[str, Any], + output: Any, + *, + role: str, + attention_metadata: Any | None, + cross_attention_cached: bool = False, +) -> _ForwardWork | None: + """Estimate one Wan-style Transformer invocation from its logical input.""" + hidden_states = _first_tensor(kwargs.get("hidden_states")) + if hidden_states is None or hidden_states.ndim != 5: + return None + + arch = _arch_config(module) + if arch is None: + return None + required = ("hidden_size", "ffn_dim", "num_layers", "patch_size") + if any(not hasattr(arch, name) for name in required): + return None + + batch_size = int(hidden_states.shape[0]) + raw_frames = int(hidden_states.shape[2]) + raw_height = int(hidden_states.shape[3]) + raw_width = int(hidden_states.shape[4]) + patch_size = arch.patch_size + patch: tuple[int, + ...] = ((1, patch_size, + patch_size) if isinstance(patch_size, int) else tuple(int(value) for value in patch_size)) + if len(patch) != 3 or any(value <= 0 for value in patch): + return None + + frames = raw_frames // patch[0] + spatial_tokens = (raw_height // patch[1]) * (raw_width // patch[2]) + seq_len = frames * spatial_tokens + if seq_len <= 0: + return None + + hidden_size = int(arch.hidden_size) + ffn_dim = int(arch.ffn_dim) + num_layers = int(arch.num_layers) + + context = _first_tensor(kwargs.get("encoder_hidden_states")) + context_tokens = int(context.shape[-2]) if context is not None and context.ndim >= 2 else 0 + image_context = _first_tensor(kwargs.get("encoder_hidden_states_image")) + if image_context is not None and image_context.ndim >= 2: + context_tokens += int(image_context.shape[-2]) + + is_causal = hasattr(module, "num_frame_per_block") + teacher_forcing = is_causal and _first_tensor(kwargs.get("clean_x")) is not None + query_tokens = seq_len + dense_pairs = float(seq_len * seq_len) + vsa_tiles = 0 + causal_chunks = 0 + + if is_causal: + frames_per_block = int(getattr(module, "num_frame_per_block", 1)) + causal_chunks = math.ceil(frames / max(1, frames_per_block)) + local_attn_size = int(getattr(module, "local_attn_size", -1)) + kv_cache = kwargs.get("kv_cache") + if kv_cache is not None: + current_start = int(kwargs.get("current_start", 0)) + if local_attn_size >= 0: + max_frames = local_attn_size + else: + # Causal Wan caps the KV window at + # GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES (21) frames when + # local_attn_size is unset; sliding_window_num_frames only + # sizes the streaming KV cache. MatrixGame2's 15-frame + # compatibility window is not modeled here. + max_frames = 21 + key_tokens = min(current_start + seq_len, max_frames * spatial_tokens) + attention_pairs = float(seq_len * key_tokens) + dense_pairs = float(seq_len * (current_start + seq_len)) + elif teacher_forcing: + query_tokens = 2 * seq_len + attention_pairs = float(_teacher_forcing_frame_pairs(frames, frames_per_block) * spatial_tokens**2) + dense_pairs = float(query_tokens * query_tokens) + else: + attention_pairs = float( + _blockwise_causal_frame_pairs(frames, frames_per_block, local_attn_size) * spatial_tokens**2) + else: + attention_pairs, vsa_tiles = _vsa_attention_pairs(seq_len, attention_metadata) + + # Causal Wan pads text to ``text_len`` before every cross-attention block. + if is_causal: + context_tokens = max(context_tokens, int(getattr(module, "text_len", 0))) + + b = batch_size + seq_length = query_tokens + dim = hidden_size + # Four self-attention projections, two MLP projections, cross-attention + # Q/K/V/out projections, and the two attention matmuls. + per_layer = (8.0 * b * seq_length * dim * dim + 4.0 * b * seq_length * dim * ffn_dim + + 4.0 * b * seq_length * dim * dim + + (0.0 if cross_attention_cached else 4.0 * b * context_tokens * dim * dim) + + 4.0 * b * seq_length * context_tokens * dim + 4.0 * b * attention_pairs * dim) + + if vsa_tiles: + # Wan VSA adds one gate projection and dense pooled QK/AV matmuls. + per_layer += 2.0 * b * seq_length * dim * dim + per_layer += 4.0 * b * vsa_tiles * vsa_tiles * dim + + forward_flops = per_layer * num_layers + backward_expected = _output_requires_grad(output) + # Standard MFU counts model-algorithm FLOPs, excluding activation + # checkpoint recompute: training is approximately F + 2B = 3F. + useful_flops = forward_flops * (3.0 if backward_expected else 1.0) + return _ForwardWork( + role=role, + forward_flops=forward_flops, + useful_flops=useful_flops, + causal_chunks=causal_chunks, + query_frames=b * (2 * frames if teacher_forcing else frames), + query_tokens=b * query_tokens, + attention_pairs=b * attention_pairs, + dense_attention_pairs=b * dense_pairs, + ) + + +class TrainingPerformanceMonitor: + """Count per-role Transformer work for one optimizer step at a time.""" + + def __init__(self) -> None: + self._handles: list[Any] = [] + self._work: list[_ForwardWork] = [] + self._forward_calls: dict[str, int] = defaultdict(int) + self._backward_forwards: dict[str, int] = defaultdict(int) + + def attach(self, role_models: dict[str, Any]) -> None: + self.close() + for role, model in role_models.items(): + transformer = getattr(model, "transformer", None) + if not isinstance(transformer, torch.nn.Module): + continue + + cached_cross_attention: list[bool] = [] + + def pre_hook( + module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + cache_state: list[bool] = cached_cross_attention, + ) -> None: + del module, args + cache = kwargs.get("crossattn_cache") + initialized = (isinstance(cache, list) and bool(cache) + and all(isinstance(item, dict) and bool(item.get("is_init", False)) for item in cache)) + cache_state.append(initialized) + + def hook( + module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + role_name: str = role, + cache_state: list[bool] = cached_cross_attention, + ) -> None: + del args + self._forward_calls[role_name] += 1 + if _output_requires_grad(output): + self._backward_forwards[role_name] += 1 + try: + metadata = get_forward_context().attn_metadata + except AssertionError: + metadata = None + try: + work = estimate_transformer_forward( + module, + kwargs, + output, + role=role_name, + attention_metadata=metadata, + cross_attention_cached=(cache_state.pop() if cache_state else False), + ) + except Exception: + # Metrics must never abort a training run. + work = None + if work is not None: + self._work.append(work) + + self._handles.append(transformer.register_forward_pre_hook(pre_hook, with_kwargs=True)) + self._handles.append(transformer.register_forward_hook(hook, with_kwargs=True)) + + def reset(self) -> None: + self._work.clear() + self._forward_calls.clear() + self._backward_forwards.clear() + + def close(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + def metrics( + self, + *, + step_time_sec: float, + local_batch_size: int, + grad_accum: int, + world_size: int, + sp_size: int, + peak_tflops_per_gpu: float | None, + ) -> dict[str, float]: + if step_time_sec <= 0.0: + return {} + world_size = max(1, int(world_size)) + sp_size = max(1, int(sp_size)) + data_parallel_size = max(1, world_size // sp_size) + accumulation = max(1, int(grad_accum)) + sample_scale_to_world = data_parallel_size * accumulation + work_scale_to_world = data_parallel_size + + metrics = { + "perf/steps_per_sec": 1.0 / step_time_sec, + "perf/samples_per_sec": (max(0, int(local_batch_size)) * sample_scale_to_world / step_time_sec), + "perf/model_forward_calls": float(sum(self._forward_calls.values())), + } + + for role, calls in self._forward_calls.items(): + prefix = f"perf/role/{role}" + metrics[f"{prefix}/forward_calls"] = float(calls) + metrics[f"{prefix}/backward_forwards"] = float(self._backward_forwards[role]) + + if not self._work: + return metrics + + useful_flops = sum(item.useful_flops for item in self._work) * work_scale_to_world + query_frames = sum(item.query_frames for item in self._work) * work_scale_to_world + query_tokens = sum(item.query_tokens for item in self._work) * work_scale_to_world + attention_pairs = sum(item.attention_pairs for item in self._work) + dense_pairs = sum(item.dense_attention_pairs for item in self._work) + achieved_tflops_per_gpu = useful_flops / step_time_sec / world_size / _FLOPS_PER_TFLOP + + metrics.update({ + "perf/causal_chunks": float(sum(item.causal_chunks for item in self._work)), + "perf/query_latent_frames_per_sec": query_frames / step_time_sec, + "perf/query_tokens_per_sec": query_tokens / step_time_sec, + "perf/attention_density": (attention_pairs / dense_pairs if dense_pairs > 0.0 else 1.0), + "perf/estimated_tflops_per_gpu": achieved_tflops_per_gpu, + }) + if peak_tflops_per_gpu is not None: + metrics["perf/peak_tflops_per_gpu"] = peak_tflops_per_gpu + metrics["perf/estimated_mfu"] = achieved_tflops_per_gpu / peak_tflops_per_gpu + + by_role: dict[str, list[_ForwardWork]] = defaultdict(list) + for item in self._work: + by_role[item.role].append(item) + for role, items in by_role.items(): + prefix = f"perf/role/{role}" + metrics[f"{prefix}/causal_chunks"] = float(sum(item.causal_chunks for item in items)) + role_flops = sum(item.useful_flops for item in items) * work_scale_to_world + metrics[f"{prefix}/estimated_tflops_per_gpu"] = (role_flops / step_time_sec / world_size / _FLOPS_PER_TFLOP) + return metrics diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 25a546cef..85e838ec8 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -68,6 +68,16 @@ class TrackerConfig: run_name: str = "" +@dataclass(slots=True) +class PerformanceConfig: + """Low-overhead per-step training performance metrics.""" + + enabled: bool = True + # Dense BF16 tensor-core peak for one GPU. When unset, known NVIDIA + # accelerators are inferred from their device name. + peak_tflops_per_gpu: float | None = None + + @dataclass(slots=True) class ModelTrainingConfig: weighting_scheme: str = "uniform" @@ -87,6 +97,7 @@ class TrainingConfig: loop: TrainingLoopConfig = field(default_factory=TrainingLoopConfig) checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) tracker: TrackerConfig = field(default_factory=TrackerConfig) + performance: PerformanceConfig = field(default_factory=PerformanceConfig) vsa_sparsity: float = 0.0 # Reuse the per-step padded VSA tile buffer across attention layers. # Defaults to False for training: under full activation checkpointing the From 69a6215e7f8173bf7122360c92c9ff6c1a603132 Mon Sep 17 00:00:00 2001 From: jayzou <133576930+jayzou3773@users.noreply.github.com> Date: Wed, 7 Oct 2026 04:30:26 -0400 Subject: [PATCH 060/151] [feat] Add Dreamverse multimodal inputs and H3 routing (#1835) Co-authored-by: SolitaryThinker Co-authored-by: Claude Opus 5.5 (1M context) --- apps/dreamverse/dreamverse/assets.py | 184 ++++++++++++ apps/dreamverse/dreamverse/config.py | 50 +++- .../dreamverse/generation_contracts.py | 3 + .../dreamverse/generation_inputs.py | 102 +++++++ .../dreamverse/generation_worker.py | 3 + apps/dreamverse/dreamverse/gpu_pool.py | 30 +- apps/dreamverse/dreamverse/ltx2_generation.py | 4 + apps/dreamverse/dreamverse/main.py | 9 + .../dreamverse/minimax_h3_generation.py | 130 ++++++-- apps/dreamverse/dreamverse/mock_server.py | 43 ++- apps/dreamverse/dreamverse/routes/assets.py | 72 +++++ .../dreamverse/session/controller.py | 53 +++- .../dreamverse/tests/test_config.py | 12 + .../tests/test_generation_inputs.py | 248 ++++++++++++++++ .../tests/test_generation_mode_routing.py | 214 ++++++++++++++ .../tests/test_generation_session_inputs.py | 278 ++++++++++++++++++ .../dreamverse/tests/test_import_contract.py | 38 +++ .../dreamverse/tests/test_session_logging.py | 1 + apps/dreamverse/dreamverse/worker_ipc.py | 3 + apps/dreamverse/scripts/slurm/README.md | 107 +++++++ .../scripts/slurm/launch_backend.sh | 49 +++ .../web/e2e/generation-modes.spec.ts | 126 ++++++++ apps/dreamverse/web/next.config.ts | 14 + .../web/src/app/page.integration.test.tsx | 37 +++ apps/dreamverse/web/src/app/page.tsx | 127 +++++++- .../web/src/components/AssetList.test.tsx | 41 +++ .../web/src/components/AssetList.tsx | 154 ++++++++++ .../web/src/components/ChatBar.test.tsx | 155 ++++++++++ .../dreamverse/web/src/components/ChatBar.tsx | 133 +++++++-- .../src/components/SessionTimeoutModal.tsx | 2 +- .../components/devtools/DevtoolsComposer.tsx | 61 +++- .../devtools/DevtoolsShell.test.tsx | 1 + .../src/components/devtools/DevtoolsShell.tsx | 15 + .../web/src/hooks/useAssetLibrary.test.ts | 55 ++++ .../web/src/hooks/useAssetLibrary.ts | 136 +++++++++ .../hooks/useGenerationCapabilities.test.ts | 21 ++ .../src/hooks/useGenerationCapabilities.ts | 37 +++ .../web/src/lib/generationMode.test.ts | 58 ++++ apps/dreamverse/web/src/lib/generationMode.ts | 114 +++++++ apps/dreamverse/web/src/lib/projectStorage.ts | 7 + assets/images/dreamverse-ref2va-demo.png | Bin 0 -> 173871 bytes docs/contributing/dreamverse-development.md | 89 ++++++ fastvideo/api/__init__.py | 17 ++ 43 files changed, 2942 insertions(+), 91 deletions(-) create mode 100644 apps/dreamverse/dreamverse/assets.py create mode 100644 apps/dreamverse/dreamverse/generation_inputs.py create mode 100644 apps/dreamverse/dreamverse/routes/assets.py create mode 100644 apps/dreamverse/dreamverse/tests/test_generation_inputs.py create mode 100644 apps/dreamverse/dreamverse/tests/test_generation_mode_routing.py create mode 100644 apps/dreamverse/dreamverse/tests/test_generation_session_inputs.py create mode 100644 apps/dreamverse/scripts/slurm/README.md create mode 100644 apps/dreamverse/scripts/slurm/launch_backend.sh create mode 100644 apps/dreamverse/web/e2e/generation-modes.spec.ts create mode 100644 apps/dreamverse/web/src/components/AssetList.test.tsx create mode 100644 apps/dreamverse/web/src/components/AssetList.tsx create mode 100644 apps/dreamverse/web/src/components/ChatBar.test.tsx create mode 100644 apps/dreamverse/web/src/hooks/useAssetLibrary.test.ts create mode 100644 apps/dreamverse/web/src/hooks/useAssetLibrary.ts create mode 100644 apps/dreamverse/web/src/hooks/useGenerationCapabilities.test.ts create mode 100644 apps/dreamverse/web/src/hooks/useGenerationCapabilities.ts create mode 100644 apps/dreamverse/web/src/lib/generationMode.test.ts create mode 100644 apps/dreamverse/web/src/lib/generationMode.ts create mode 100644 assets/images/dreamverse-ref2va-demo.png diff --git a/apps/dreamverse/dreamverse/assets.py b/apps/dreamverse/dreamverse/assets.py new file mode 100644 index 000000000..3b5af4734 --- /dev/null +++ b/apps/dreamverse/dreamverse/assets.py @@ -0,0 +1,184 @@ +"""Bounded, runtime-local media library shared by the HTTP and generation APIs.""" + +from __future__ import annotations + +import json +import math +import os +import re +import shutil +import subprocess +import tempfile +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image, UnidentifiedImageError + +IMAGE_LIMIT = 15 * 1024 * 1024 +MEDIA_LIMIT = 100 * 1024 * 1024 +STORE_LIMIT = 2 * 1024 * 1024 * 1024 +ASSET_LIMIT = 100 +MAX_MEDIA_SECONDS = 30 +MIME_TYPES = { + "image/png": ("image", ".png"), + "image/jpeg": ("image", ".jpg"), + "image/webp": ("image", ".webp"), + "video/mp4": ("video", ".mp4"), + "video/quicktime": ("video", ".mov"), + "video/webm": ("video", ".webm"), + "audio/mpeg": ("audio", ".mp3"), + "audio/mp4": ("audio", ".m4a"), + "audio/x-m4a": ("audio", ".m4a"), + "audio/wav": ("audio", ".wav"), + "audio/x-wav": ("audio", ".wav"), + "audio/flac": ("audio", ".flac"), + "audio/x-flac": ("audio", ".flac"), + "audio/ogg": ("audio", ".ogg"), + "audio/webm": ("audio", ".webm"), +} + + +@dataclass(frozen=True) +class StoredAsset: + asset_id: str + kind: str + path: str + name: str + mime_type: str + size: int + + def public(self) -> dict: + return { + "asset_id": self.asset_id, + "kind": self.kind, + "name": self.name, + "mime_type": self.mime_type, + "size": self.size, + "url": f"/assets/{self.asset_id}", + } + + +def validate_media(path: Path, mime_type: str) -> None: + """Inspect content, not filenames; refuse playlists and non-media uploads.""" + kind = MIME_TYPES[mime_type][0] + if kind == "image": + try: + with Image.open(path) as img: + expected = {"image/png": "PNG", "image/jpeg": "JPEG", "image/webp": "WEBP"}[mime_type] + if img.format != expected: + raise ValueError("The image content does not match its file type.") + if img.width * img.height > 16_777_216: + raise ValueError("Images must contain at most 16 megapixels.") + if getattr(img, "is_animated", False): + raise ValueError("Use a still image or upload the animation as a video.") + img.verify() + except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc: + raise ValueError("The image could not be decoded. Use PNG, JPEG, or WebP.") from exc + return + + probe = shutil.which(os.getenv("FASTVIDEO_FFPROBE_BIN", "ffprobe")) + if not probe: + raise ValueError("This runtime needs ffprobe installed to accept video and audio assets.") + try: + result = subprocess.run( + [ + probe, "-v", "error", "-protocol_whitelist", "file,pipe", "-format_whitelist", + "mov,matroska,webm,mp3,wav,flac,ogg", "-show_format", "-show_streams", "-of", "json", + str(path) + ], + check=True, + capture_output=True, + timeout=15, + ) + info = json.loads(result.stdout) + formats = set(info.get("format", {}).get("format_name", "").split(",")) + if not formats.intersection({"mov", "mp4", "matroska", "webm", "mp3", "wav", "flac", "ogg"}): + raise ValueError("Upload a media file, not a playlist or external reference.") + streams = [stream for stream in info.get("streams", []) if stream.get("codec_type") == kind] + if not streams: + raise ValueError(f"The file contains no {kind} stream.") + for stream in info.get("streams", []): + if stream.get("codec_type") == "audio" and int(stream.get("channels", 0)) not in (1, 2): + raise ValueError("H3 references require mono or stereo audio, including video soundtracks.") + duration = float(info.get("format", {}).get("duration", "nan")) + if not math.isfinite(duration) or not 0 < duration <= MAX_MEDIA_SECONDS: + raise ValueError(f"Reference video and audio must be between 0 and {MAX_MEDIA_SECONDS} seconds long.") + for stream in streams: + if kind == "video" and int(stream.get("width", 0)) * int(stream.get("height", 0)) > 8_294_400: + raise ValueError("Reference videos must be 4K or smaller.") + except (subprocess.SubprocessError, json.JSONDecodeError, OSError) as exc: + raise ValueError("The media file could not be decoded. Check its format and try again.") from exc + + +class AssetStore: + """Assets live until deletion or runtime exit; pinned generation inputs cannot be deleted.""" + + def __init__(self) -> None: + self._directory: tempfile.TemporaryDirectory | None = None + self._assets: dict[str, StoredAsset] = {} + self._pins: dict[str, int] = {} + self._lock = threading.RLock() + + def staging_path(self, mime_type: str) -> Path: + with self._lock: + if mime_type not in MIME_TYPES: + raise ValueError("Unsupported media type. Use PNG/JPEG/WebP, MP4/WebM/MOV, or WAV/MP3/M4A/FLAC/OGG.") + if len(self._assets) >= ASSET_LIMIT or sum(item.size for item in self._assets.values()) >= STORE_LIMIT: + raise ValueError("The runtime asset library is full. Remove unused assets before uploading more.") + if self._directory is None: + self._directory = tempfile.TemporaryDirectory(prefix="dreamverse-assets-") + return Path(self._directory.name) / f"{uuid.uuid4().hex}{MIME_TYPES[mime_type][1]}" + + def add(self, path: Path, name: str, mime_type: str) -> StoredAsset: + validate_media(path, mime_type) + size = path.stat().st_size + if size == 0 or size > (IMAGE_LIMIT if MIME_TYPES[mime_type][0] == "image" else MEDIA_LIMIT): + raise ValueError("The asset is empty or exceeds its upload size limit.") + with self._lock: + if len(self._assets) >= ASSET_LIMIT or size + sum(item.size + for item in self._assets.values()) > STORE_LIMIT: + raise ValueError("The runtime asset library is full. Remove unused assets before uploading more.") + if self._directory is None or path.parent != Path(self._directory.name): + raise ValueError("The asset must be uploaded to this runtime.") + display_name = re.sub(r"[\x00-\x1f\x7f/\\]", "_", name).strip()[:200] or "Untitled asset" + asset = StoredAsset(path.stem, MIME_TYPES[mime_type][0], str(path), display_name, mime_type, size) + self._assets[asset.asset_id] = asset + return asset + + def get(self, asset_id: str) -> StoredAsset: + with self._lock: + if not isinstance(asset_id, str) or not re.fullmatch(r"[a-f0-9]{32}", asset_id): + raise ValueError("Invalid asset ID. Upload or select an asset from the library.") + asset = self._assets.get(asset_id) + if asset is None or not Path(asset.path).is_file(): + raise ValueError("An asset is no longer available. Upload it again and reselect it.") + return asset + + def pin(self, asset_ids: list[str]) -> None: + with self._lock: + for asset_id in asset_ids: + self.get(asset_id) + for asset_id in asset_ids: + self._pins[asset_id] = self._pins.get(asset_id, 0) + 1 + + def release(self, asset_ids: list[str]) -> None: + with self._lock: + for asset_id in asset_ids: + count = self._pins.get(asset_id, 0) + if count > 1: + self._pins[asset_id] = count - 1 + else: + self._pins.pop(asset_id, None) + + def delete(self, asset_id: str) -> None: + with self._lock: + asset = self.get(asset_id) + if self._pins.get(asset_id, 0): + raise ValueError("This asset is in use by a generation session. End the session before deleting it.") + Path(asset.path).unlink(missing_ok=True) + del self._assets[asset_id] + + +asset_store = AssetStore() diff --git a/apps/dreamverse/dreamverse/config.py b/apps/dreamverse/dreamverse/config.py index 0c65f70c4..e46a038ee 100644 --- a/apps/dreamverse/dreamverse/config.py +++ b/apps/dreamverse/dreamverse/config.py @@ -84,6 +84,19 @@ MODEL_REGISTRY = { "num_inference_steps": 5, "seed": 1000, }, + "full-h3": { + "name": "MiniMax H3 (Full)", + "generation_backend": "minimax_h3", + "default_sp_size": 4, + "model_path": "MiniMaxAI/MiniMax-H3", + "attention_backend": "FLASH_ATTN", + "height": 768, + "width": 1344, + "num_frames": 124, + "num_inference_steps": 50, + "seed": 1000, + "full_checkpoint": True, + }, } DEFAULT_MODEL_ID = "fast-ltx2" @@ -95,22 +108,6 @@ if ACTIVE_MODEL_ID not in MODEL_REGISTRY: # Active model configuration MODEL_CONFIG = MODEL_REGISTRY[ACTIVE_MODEL_ID] -# Generation limits -SESSION_TIMEOUT_SECONDS = 300 - -# Frame settings -NUM_FRAMES = 121 -FRAME_HEIGHT = 1088 -FRAME_WIDTH = 1920 -NUM_INFERENCE_STEPS = 5 -JPEG_QUALITY = 100 -BATCH_SIZE = 3 - -# Streaming mode: -# - legacy_jpeg: send frame_batch JSON payloads with base64 JPEGs -# - av_fmp4: send muxed fMP4 binary chunks over WebSocket -STREAM_MODE = os.getenv("STREAM_MODE", "av_fmp4").strip().lower() - def _env_int(name: str, default: int) -> int: value = os.getenv(name) @@ -187,6 +184,27 @@ def _optional_env(*names: str) -> str | None: return None +# Generation limits +# Full H3 loads and generates substantially longer than the Preview adapter. +# This also covers a base/ref pipeline reload inside a retained session. +# Values below 60 seconds are floored so a single segment cannot outlast the session. +SESSION_TIMEOUT_SECONDS = max( + 60, _env_int("DREAMVERSE_SESSION_TIMEOUT_SECONDS", 7200 if ACTIVE_MODEL_ID == "full-h3" else 300)) + +# Frame settings +NUM_FRAMES = 121 +FRAME_HEIGHT = 1088 +FRAME_WIDTH = 1920 +NUM_INFERENCE_STEPS = 5 +JPEG_QUALITY = 100 +BATCH_SIZE = 3 + +# Streaming mode: +# - legacy_jpeg: send frame_batch JSON payloads with base64 JPEGs +# - av_fmp4: send muxed fMP4 binary chunks over WebSocket +STREAM_MODE = os.getenv("STREAM_MODE", "av_fmp4").strip().lower() + + DEVTOOLS_ENABLED = _env_bool("FASTVIDEO_ENABLE_DEVTOOLS", False) PROMPT_SAFETY_ENABLED = _env_bool("FASTVIDEO_ENABLE_PROMPT_SAFETY", False) DREAMVERSE_MAX_AUTOTUNE = _env_bool("DREAMVERSE_MAX_AUTOTUNE", True) diff --git a/apps/dreamverse/dreamverse/generation_contracts.py b/apps/dreamverse/dreamverse/generation_contracts.py index 782b23304..30928044a 100644 --- a/apps/dreamverse/dreamverse/generation_contracts.py +++ b/apps/dreamverse/dreamverse/generation_contracts.py @@ -5,6 +5,8 @@ from __future__ import annotations from dataclasses import dataclass from typing import Any, Protocol +from dreamverse.generation_inputs import GenerationInputs + @dataclass class StepResult: @@ -36,6 +38,7 @@ class GenerationBackend(Protocol): segment_idx: int, image_path: str | None, reset_conditioning: bool, + generation_inputs: GenerationInputs | None = None, ) -> StepResult: ... diff --git a/apps/dreamverse/dreamverse/generation_inputs.py b/apps/dreamverse/dreamverse/generation_inputs.py new file mode 100644 index 000000000..0e75109be --- /dev/null +++ b/apps/dreamverse/dreamverse/generation_inputs.py @@ -0,0 +1,102 @@ +"""GPU-independent validation for generation modes and ordered asset handles.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from PIL import Image, UnidentifiedImageError + +from dreamverse.assets import asset_store + +GENERATION_MODES = ("t2va", "fl2va", "ref2va") + + +@dataclass(frozen=True) +class GenerationAsset: + asset_id: str + kind: str + path: str + role: str + + +@dataclass(frozen=True) +class GenerationInputs: + mode: str | None = None + assets: tuple[GenerationAsset, ...] = () + + @property + def first_frame_path(self) -> str | None: + return next((asset.path for asset in self.assets if asset.role == "first_frame"), None) + + @property + def last_frame_path(self) -> str | None: + return next((asset.path for asset in self.assets if asset.role == "last_frame"), None) + + @property + def references(self) -> tuple[GenerationAsset, ...]: + return tuple(asset for asset in self.assets if asset.role == "reference") + + +def supported_generation_modes(model_id: str) -> tuple[str, ...]: + return GENERATION_MODES if model_id in ("full-h3", "mock") else ("t2va", ) + + +def resolve_generation_inputs(payload: dict, model_id: str) -> GenerationInputs: + mode = payload.get("generation_mode") + raw_assets = payload.get("conditioning_assets", []) + if mode is None and "generation_mode" not in payload: + if raw_assets: + raise ValueError("Select a generation mode before attaching conditioning assets.") + return GenerationInputs() + if not isinstance(mode, str) or mode not in GENERATION_MODES: + raise ValueError("Unknown generation mode. Choose T2VA, FL2VA, or Ref2VA.") + if mode not in supported_generation_modes(model_id): + raise ValueError(f"{mode.upper()} requires the Full H3 runtime. This runtime is running {model_id}.") + if payload.get("initial_image") is not None: + raise ValueError("Use asset IDs for generation modes; do not combine them with the legacy initial_image field.") + if not isinstance(raw_assets, list) or len(raw_assets) > 12: + raise ValueError("conditioning_assets must be an ordered list with at most 12 assets.") + if mode == "t2va" and raw_assets: + raise ValueError("T2VA accepts text only. Remove conditioning assets or choose another mode.") + assets: list[GenerationAsset] = [] + for item in raw_assets: + if not isinstance(item, dict) or set(item) != {"asset_id", "role"}: + raise ValueError("Each conditioning asset must contain only asset_id and role.") + role = item["role"] + if role not in ("first_frame", "last_frame", "reference"): + raise ValueError("Asset role must be first_frame, last_frame, or reference.") + stored = asset_store.get(item["asset_id"]) + assets.append(GenerationAsset(stored.asset_id, stored.kind, stored.path, role)) + if mode == "fl2va": + if any(asset.kind != "image" or asset.role == "reference" for asset in assets): + raise ValueError("FL2VA accepts only first-frame and last-frame images.") + if sum(asset.role == "first_frame" for asset in assets) != 1: + raise ValueError("FL2VA requires exactly one first-frame image.") + if sum(asset.role == "last_frame" for asset in assets) > 1: + raise ValueError("FL2VA accepts at most one last-frame image.") + elif mode == "ref2va": + if not assets or any(asset.role != "reference" for asset in assets): + raise ValueError("Ref2VA requires an ordered list of reference assets, without keyframe roles.") + if not any(asset.kind in ("image", "video") for asset in assets): + raise ValueError("Ref2VA requires at least one image or video; audio alone is not supported.") + for kind, limit in (("image", 9), ("video", 3), ("audio", 3)): + if sum(asset.kind == kind for asset in assets) > limit: + raise ValueError(f"Ref2VA accepts at most {limit} {kind} references.") + for asset in assets: + if asset.kind == "image": + try: + with Image.open(asset.path) as image: + if image.width > 4 * image.height or image.height > 4 * image.width: + raise ValueError( + "Ref2VA image aspect ratios must be between 1:4 and 4:1. Crop this image first.") + except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc: + raise ValueError("A selected reference image could not be decoded. Upload it again.") from exc + return GenerationInputs(mode, tuple(assets)) + + +def pin_generation_inputs(inputs: GenerationInputs) -> None: + asset_store.pin([asset.asset_id for asset in inputs.assets]) + + +def release_generation_inputs(inputs: GenerationInputs) -> None: + asset_store.release([asset.asset_id for asset in inputs.assets]) diff --git a/apps/dreamverse/dreamverse/generation_worker.py b/apps/dreamverse/dreamverse/generation_worker.py index 1585940ae..3ff18c34d 100644 --- a/apps/dreamverse/dreamverse/generation_worker.py +++ b/apps/dreamverse/dreamverse/generation_worker.py @@ -4,6 +4,7 @@ from __future__ import annotations from dreamverse.config import MODEL_CONFIG from dreamverse.generation_contracts import GenerationBackend, StepResult +from dreamverse.generation_inputs import GenerationInputs def _create_generation_backend(backend_name: str, gpu_id: int) -> GenerationBackend: @@ -80,6 +81,7 @@ class VideoGenerationWorker: segment_idx: int, image_path: str | None, reset_conditioning: bool, + generation_inputs: GenerationInputs | None = None, ) -> StepResult: """Generate one segment through the selected model backend.""" return self._require_backend().generate_step( @@ -87,6 +89,7 @@ class VideoGenerationWorker: segment_idx, image_path, reset_conditioning, + generation_inputs=generation_inputs, ) def warmup(self, prompt: str) -> dict[str, float]: diff --git a/apps/dreamverse/dreamverse/gpu_pool.py b/apps/dreamverse/dreamverse/gpu_pool.py index a524fe16c..5ecef8317 100644 --- a/apps/dreamverse/dreamverse/gpu_pool.py +++ b/apps/dreamverse/dreamverse/gpu_pool.py @@ -29,6 +29,7 @@ from dreamverse.av_streaming import ( generate_stream_id, stream_fmp4, ) +from dreamverse.generation_inputs import GenerationInputs, pin_generation_inputs, release_generation_inputs from dreamverse.worker_ipc import ( CommandPayload, InitAck, @@ -189,6 +190,7 @@ def gpu_worker_process( segment_idx, image_path=payload.image_path, reset_conditioning=payload.reset_conditioning, + generation_inputs=payload.generation_inputs, ) head_trim_frames = step_result.head_trim_frames head_trim_audio_frames = step_result.head_trim_audio_frames @@ -432,6 +434,7 @@ class GPUSlot: self.connected_users: set[str] = set() self._pending_futures: dict[str, asyncio.Future] = {} self._stream_queues: dict[str, asyncio.Queue] = {} + self._step_asset_inputs: dict[str, GenerationInputs] = {} self._response_reader_task: asyncio.Task | None = None self._active: bool = False self._reader_lock: asyncio.Lock | None = None @@ -663,6 +666,9 @@ class GPUSlot: if isinstance(event, (StepComplete, WarmupComplete)): event.timings["ipc_get_done_ns"] = time.time_ns() + if isinstance(event, (StepComplete, WorkerError)) and event.user_id is not None: + self._release_step_assets(event.user_id) + user_id = event.user_id if user_id and user_id in self._pending_futures: future = self._pending_futures.pop(user_id) @@ -753,6 +759,7 @@ class GPUSlot: segment_idx: int = 1, image_path: str | None = None, reset_conditioning: bool = False, + generation_inputs: GenerationInputs | None = None, ) -> dict[str, float]: """Execute a generation step for a specific user. @@ -766,9 +773,19 @@ class GPUSlot: segment_idx=segment_idx, image_path=image_path, reset_conditioning=bool(reset_conditioning), + generation_inputs=generation_inputs, ) + if generation_inputs is not None: + if user_id in self._step_asset_inputs: + raise RuntimeError("The previous generation is still using this project's assets.") + pin_generation_inputs(generation_inputs) + self._step_asset_inputs[user_id] = generation_inputs + # Pins intentionally survive a waiter timeout/cancellation: the GPU + # command keeps running. The response reader releases them when the + # worker actually completes (even if that response is now unmatched). response = await self._send_command_tagged(Command(CommandType.USER_STEP, payload=payload, user_id=user_id), timeout=1800.0) + self._release_step_assets(user_id) match response: case StepComplete(timings=timings): return timings @@ -778,6 +795,11 @@ class GPUSlot: raise RuntimeError(f"Unexpected step response for {user_id[:8]}: " f"{type(response).__name__}") + def _release_step_assets(self, user_id: str) -> None: + inputs = self._step_asset_inputs.pop(user_id, None) + if inputs is not None: + release_generation_inputs(inputs) + async def apply_lora_stack( self, stack: list[tuple[str, float]], @@ -800,7 +822,9 @@ class GPUSlot: async def leave_user(self, user_id: str) -> None: """Remove a user from this GPU.""" try: - await self._send_command_tagged(Command(CommandType.USER_LEAVE, user_id=user_id), timeout=30.0) + response = await self._send_command_tagged(Command(CommandType.USER_LEAVE, user_id=user_id), timeout=30.0) + if isinstance(response, LeaveAck): + self._release_step_assets(user_id) except Exception as e: print(f"[GPU {self.gpu_id}] Leave user error: {e}") finally: @@ -837,6 +861,10 @@ class GPUSlot: except Exception: pass + if self.process is None or not self.process.is_alive(): + for user_id in list(self._step_asset_inputs): + self._release_step_assets(user_id) + for q in (self.command_queue, self.response_queue): if q is not None: try: diff --git a/apps/dreamverse/dreamverse/ltx2_generation.py b/apps/dreamverse/dreamverse/ltx2_generation.py index 59f74680c..4b45e92c7 100644 --- a/apps/dreamverse/dreamverse/ltx2_generation.py +++ b/apps/dreamverse/dreamverse/ltx2_generation.py @@ -33,6 +33,7 @@ from dreamverse.config import ( _resolve_lora_spec, ) from dreamverse.generation_contracts import StepResult +from dreamverse.generation_inputs import GenerationInputs # Multi-frame decoded continuation defaults from # examples/inference/basic/basic_ltx2_distilled_video_continuation.py. @@ -454,8 +455,11 @@ class LTX2GenerationBackend: segment_idx: int, image_path: str | None, reset_conditioning: bool, + generation_inputs: GenerationInputs | None = None, ) -> StepResult: """Execute one generation step; snapshot state for the next segment.""" + if generation_inputs is not None and (generation_inputs.mode not in (None, "t2va") or generation_inputs.assets): + raise ValueError("LTX supports text generation only through the generation mode API.") timings: dict = {} prompt = self._inject_style_trigger(prompt) diff --git a/apps/dreamverse/dreamverse/main.py b/apps/dreamverse/dreamverse/main.py index 65d4f9964..6768e9934 100644 --- a/apps/dreamverse/dreamverse/main.py +++ b/apps/dreamverse/dreamverse/main.py @@ -15,6 +15,7 @@ from dreamverse.gpu_pool import GPUPool, get_available_gpus from dreamverse.session_logger import SessionEventLogger from dreamverse.config import ( + ACTIVE_MODEL_ID, AVAILABLE_LORAS, DEVTOOLS_ENABLED, FRONTEND_STATIC_DIR_CANDIDATES, @@ -34,6 +35,8 @@ from dreamverse.routes.presets import ( curated_presets_router, ) from dreamverse.session.controller import SessionController +from dreamverse.generation_inputs import supported_generation_modes +from dreamverse.routes.assets import router as asset_router class _HeartbeatAccessLogFilter(logging.Filter): @@ -92,10 +95,16 @@ app.add_middleware( app.include_router(build_health_router(lambda: runtime.gpu_pool)) app.include_router(internal_monitor_router) app.include_router(prompt_config_router) +app.include_router(asset_router) if DEVTOOLS_ENABLED: app.include_router(curated_presets_router) +@app.get("/generation-capabilities") +async def generation_capabilities() -> dict: + return {"model_id": ACTIVE_MODEL_ID, "modes": supported_generation_modes(ACTIVE_MODEL_ID), "mock": False} + + @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): controller = SessionController( diff --git a/apps/dreamverse/dreamverse/minimax_h3_generation.py b/apps/dreamverse/dreamverse/minimax_h3_generation.py index 6c99b40b7..4861ef2bc 100644 --- a/apps/dreamverse/dreamverse/minimax_h3_generation.py +++ b/apps/dreamverse/dreamverse/minimax_h3_generation.py @@ -1,4 +1,4 @@ -"""FastH3 model lifecycle and first-frame continuation for DreamVerse.""" +"""Full/Preview H3 lifecycle, conditioning and per-project pipeline selection.""" from __future__ import annotations @@ -12,6 +12,7 @@ import torch from dreamverse.config import DREAMVERSE_SP_SIZE from dreamverse.generation_contracts import StepResult +from dreamverse.generation_inputs import GenerationInputs if TYPE_CHECKING: from PIL.Image import Image @@ -26,13 +27,14 @@ def _required_config_str(model_config: dict, field_name: str) -> str: class MiniMaxH3GenerationBackend: - """Run the VSA data-free FastH3 adapter and retain one continuation frame.""" + """Own one H3 pipeline at a time and retain base-pipeline continuation.""" def __init__(self, gpu_id: int): self.gpu_id = gpu_id self.generator: Any | None = None self.model_config: dict = {} self.continuation_image: Image | None = None + self.pipeline_mode = "base" def _gpu_mem(self) -> str: allocated_gib = torch.cuda.memory_allocated() / 1024**3 @@ -51,32 +53,37 @@ class MiniMaxH3GenerationBackend: os.environ.pop("FASTVIDEO_INFERENCE_TORCH_COMPILE", None) def initialize(self, model_config: dict | None = None) -> None: - """Download the fixed Preview adapter and load the FastH3 generator. - - The model profile owns the base checkpoint, adapter file, attention - backend, and generation geometry. The backend translates that profile - into FastVideo's typed generator configuration. - """ + """Load the profile's base pipeline; Ref2VA is loaded on first use.""" if model_config is not None: self.model_config = dict(model_config) if not self.model_config: raise ValueError("FastH3 initialization requires a model configuration.") + self._load_pipeline("base") + def _load_pipeline(self, pipeline_mode: str) -> None: + """Unload the old executor before loading a base or reference transformer. + + GPU worker commands are serialized, so a project boundary never swaps + weights while another request is using them. Keeping one executor also + avoids simultaneously retaining two large H3 transformers in VRAM. A + failed load leaves no executor behind so the next step retries it. + """ + full_checkpoint = bool(self.model_config.get("full_checkpoint", False)) + if pipeline_mode == "ref2va" and not full_checkpoint: + raise ValueError("Ref2VA requires the full-h3 model profile.") if self.generator is not None: - self.generator.shutdown() + previous_generator = self.generator self.generator = None + previous_generator.shutdown() + del previous_generator gc.collect() torch.cuda.empty_cache() self.clear_conditioning() model_path = _required_config_str(self.model_config, "model_path") - adapter_repo = _required_config_str(self.model_config, "adapter_repo") - adapter_filename = _required_config_str(self.model_config, "adapter_filename") attention_backend = _required_config_str(self.model_config, "attention_backend") self._configure_environment(attention_backend) - from huggingface_hub import hf_hub_download - from fastvideo import VideoGenerator from fastvideo.api import ( CompileConfig, @@ -88,10 +95,20 @@ class MiniMaxH3GenerationBackend: PipelineSelection, ) - adapter_path = hf_hub_download(repo_id=adapter_repo, filename=adapter_filename) + components = ComponentConfig() + if not full_checkpoint: + from huggingface_hub import hf_hub_download + + adapter_repo = _required_config_str(self.model_config, "adapter_repo") + adapter_filename = _required_config_str(self.model_config, "adapter_filename") + components.lora_path = hf_hub_download(repo_id=adapter_repo, filename=adapter_filename) + components.lora_strength = 1.0 + print(f"[GPU {self.gpu_id}] FastH3 adapter: {adapter_repo}/{adapter_filename}") + if pipeline_mode == "ref2va": + components.override_pipeline_cls_name = "MiniMaxH3Ref2VAModularPipeline" experimental = { "attention_backend": attention_backend, - "inference_torch_compile": attention_backend == "FLASH_ATTN", + "inference_torch_compile": not full_checkpoint and attention_backend == "FLASH_ATTN", "vae_parallel_decode": True, "vae_parallel_decode_strategy": "gather", } @@ -103,7 +120,8 @@ class MiniMaxH3GenerationBackend: generator_config = GeneratorConfig( model_path=model_path, pipeline=PipelineSelection( - components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0), + workload_type="i2v" if pipeline_mode == "ref2va" else None, + components=components, experimental=experimental, ), engine=EngineConfig( @@ -115,17 +133,23 @@ class MiniMaxH3GenerationBackend: text_encoder=True, image_encoder=True, vae=True, - pin_cpu_memory=True, + pin_cpu_memory=not full_checkpoint, ), compile=CompileConfig(enabled=False, vae_enabled=True), - use_fsdp_inference=False, + use_fsdp_inference=full_checkpoint and DREAMVERSE_SP_SIZE > 1, ), ) - print(f"[GPU {self.gpu_id}] Loading FastH3 model: {model_path}") - print(f"[GPU {self.gpu_id}] FastH3 adapter: {adapter_repo}/{adapter_filename}") + print(f"[GPU {self.gpu_id}] Loading H3 model: {model_path} ({pipeline_mode})") print(f"[GPU {self.gpu_id}] Before model load: {self._gpu_mem()}") - self.generator = VideoGenerator.from_config(generator_config) + try: + self.generator = VideoGenerator.from_config(generator_config) + except Exception: + # The old executor is already gone; leaving no executor behind lets + # the next step retry this load instead of stranding the GPU slot. + self.generator = None + raise + self.pipeline_mode = pipeline_mode print(f"[GPU {self.gpu_id}] FastH3 loaded: {self._gpu_mem()} (warmup pending)") def shutdown(self) -> None: @@ -166,14 +190,28 @@ class MiniMaxH3GenerationBackend: return self._load_rgb_image(image_path), False return None, False - def _build_request(self, prompt: str, conditioning_image: Image | None): + def _build_request( + self, + prompt: str, + conditioning_image: Image | None, + last_image: Image | None = None, + generation_inputs: GenerationInputs | None = None, + ): """Build the typed FastVideo request owned by the FastH3 profile.""" from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig + references = None + if generation_inputs is not None and generation_inputs.mode == "ref2va": + from fastvideo.api import MiniMaxH3Reference + + references = [ + MiniMaxH3Reference(source=str(asset.path), media_type=asset.kind) + for asset in generation_inputs.references + ] return GenerationRequest( prompt=prompt, negative_prompt="", - inputs=InputConfig(pil_image=conditioning_image), + inputs=InputConfig(pil_image=conditioning_image, last_image=last_image, references=references), sampling=SamplingConfig( height=int(self.model_config["height"]), width=int(self.model_config["width"]), @@ -200,6 +238,7 @@ class MiniMaxH3GenerationBackend: segment_idx: int, image_path: str | None, reset_conditioning: bool, + generation_inputs: GenerationInputs | None = None, ) -> StepResult: """Generate one synchronized FastH3 segment and retain its last frame. @@ -207,20 +246,46 @@ class MiniMaxH3GenerationBackend: conditioned frame and its matching audio duration are trimmed before streaming so adjacent segments do not duplicate media. """ - if self.generator is None: - raise RuntimeError("FastH3 generator is not initialized.") - conditioning_image, uses_continuation = self._select_conditioning_image( - segment_idx, - image_path, - reset_conditioning, - ) - request = self._build_request(prompt, conditioning_image) + mode = generation_inputs.mode if generation_inputs is not None else None + if mode not in (None, "t2va", "fl2va", "ref2va"): + raise ValueError(f"Unsupported H3 generation mode: {mode!r}.") + if mode in ("fl2va", "ref2va") and not self.model_config.get("full_checkpoint", False): + raise ValueError(f"{mode.upper()} requires the full-h3 model profile.") + pipeline_mode = "ref2va" if mode == "ref2va" else "base" + if self.generator is None or self.pipeline_mode != pipeline_mode: + # A failed switch leaves no executor behind; reload here so the + # slot recovers on the next step instead of staying broken. + if segment_idx > 1 and not reset_conditioning: + raise ValueError("Generation mode cannot change in the middle of a project.") + self._load_pipeline(pipeline_mode) + + conditioning_image = None + last_image = None + uses_continuation = False + if mode == "ref2va": + # The reference pipeline rejects first/last-frame inputs. Preserve + # all original references for every clip and do not trim overlap. + self.clear_conditioning() + else: + if mode == "fl2va" and generation_inputs is not None: + image_path = generation_inputs.first_frame_path + conditioning_image, uses_continuation = self._select_conditioning_image( + segment_idx, + image_path, + reset_conditioning, + ) started = time.perf_counter() try: + if (mode == "fl2va" and segment_idx == 1 and generation_inputs is not None + and generation_inputs.last_frame_path): + last_image = self._load_rgb_image(generation_inputs.last_frame_path) + request = self._build_request(prompt, conditioning_image, last_image, generation_inputs) result = self.generator.generate(request) finally: if conditioning_image is not None: conditioning_image.close() + if last_image is not None: + last_image.close() torch.cuda.synchronize() generation_ms = (time.perf_counter() - started) * 1000.0 @@ -235,7 +300,8 @@ class MiniMaxH3GenerationBackend: raise RuntimeError("FastH3 returned audio without an audio sample rate.") save_started = time.perf_counter() - self._save_continuation_frame(frames) + if mode != "ref2va": + self._save_continuation_frame(frames) save_conditioning_ms = (time.perf_counter() - save_started) * 1000.0 timings = { "generation_ms": generation_ms, diff --git a/apps/dreamverse/dreamverse/mock_server.py b/apps/dreamverse/dreamverse/mock_server.py index 1d788c88e..201b57ac0 100644 --- a/apps/dreamverse/dreamverse/mock_server.py +++ b/apps/dreamverse/dreamverse/mock_server.py @@ -32,6 +32,14 @@ from fastapi.staticfiles import StaticFiles from dreamverse._deps import require_dreamverse_runtime_deps from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image +from dreamverse.generation_inputs import ( + GenerationInputs, + pin_generation_inputs, + release_generation_inputs, + resolve_generation_inputs, + supported_generation_modes, +) +from dreamverse.routes.assets import router as asset_router LATENCY_MS = 200 SESSION_TIMEOUT_SECONDS = 300 @@ -170,6 +178,12 @@ app.add_middleware( allow_methods=["*"], allow_headers=["*"], ) +app.include_router(asset_router) + + +@app.get("/generation-capabilities") +async def generation_capabilities(): + return {"model_id": "mock", "modes": supported_generation_modes("mock"), "mock": True} @app.get("/healthz") @@ -290,6 +304,7 @@ async def websocket_endpoint(websocket: WebSocket): send_lock = asyncio.Lock() stop_event = asyncio.Event() session_init_image = None + generation_inputs = GenerationInputs() async def ws_send_json(payload: dict) -> None: async with send_lock: @@ -347,10 +362,13 @@ async def websocket_endpoint(websocket: WebSocket): generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0) try: + generation_inputs = resolve_generation_inputs(init_data, "mock") + pin_generation_inputs(generation_inputs) session_init_image = persist_session_init_image(init_data.get("initial_image")) except ValueError as exc: await ws_send_json({ "type": "error", + "error_code": "invalid_generation_input", "message": str(exc), }) await websocket.close(code=1003, reason="Invalid initial image") @@ -362,6 +380,8 @@ async def websocket_endpoint(websocket: WebSocket): "type": "gpu_assigned", "gpu_id": 0, "session_timeout": SESSION_TIMEOUT_SECONDS, + "generation_mode": generation_inputs.mode, + "mock": True, }) raw_prompt_queue: asyncio.Queue[PromptSubmission] = asyncio.Queue() @@ -486,6 +506,7 @@ async def websocket_endpoint(websocket: WebSocket): }) async def apply_project_init_payload(payload: dict[str, object], ) -> bool: + nonlocal generation_inputs nonlocal preset_id nonlocal preset_label nonlocal initial_rollout_prompt @@ -519,10 +540,19 @@ async def websocket_endpoint(websocket: WebSocket): ] try: - replace_session_image(payload.get("initial_image")) + next_inputs = resolve_generation_inputs(payload, "mock") + pin_generation_inputs(next_inputs) + try: + replace_session_image(payload.get("initial_image")) + except ValueError: + release_generation_inputs(next_inputs) + raise + release_generation_inputs(generation_inputs) + generation_inputs = next_inputs except ValueError as exc: await ws_send_json({ "type": "error", + "error_code": "invalid_generation_input", "message": str(exc), }) return False @@ -568,6 +598,7 @@ async def websocket_endpoint(websocket: WebSocket): return drained async def enter_project_idle() -> None: + nonlocal generation_inputs nonlocal seed_prompt_memory nonlocal curated_prompts nonlocal curated_idx @@ -587,6 +618,8 @@ async def websocket_endpoint(websocket: WebSocket): dropped_raw = drain_queue_nowait(raw_prompt_queue) dropped_ready = drain_queue_nowait(ready_prompt_queue) + release_generation_inputs(generation_inputs) + generation_inputs = GenerationInputs() seed_prompt_memory = [] curated_prompts = [] curated_idx = 0 @@ -767,10 +800,17 @@ async def websocket_endpoint(websocket: WebSocket): continue try: + if generation_inputs.mode is not None and data.get("initial_image") is not None: + raise ValueError("Choose conditioning assets when starting a project; legacy initial_image " + "cannot replace generation mode inputs.") + if "generation_mode" in data or "conditioning_assets" in data: + raise ValueError( + "simple_generate cannot change the mode; use project_init_v1.") replace_session_image(data.get("initial_image")) except ValueError as exc: await ws_send_json({ "type": "error", + "error_code": "invalid_generation_input", "message": str(exc), }) continue @@ -1182,6 +1222,7 @@ async def websocket_endpoint(websocket: WebSocket): finally: stop_event.set() cleanup_session_init_image(session_init_image) + release_generation_inputs(generation_inputs) for static_dir in FRONTEND_STATIC_DIR_CANDIDATES: diff --git a/apps/dreamverse/dreamverse/routes/assets.py b/apps/dreamverse/dreamverse/routes/assets.py new file mode 100644 index 000000000..4663d13bc --- /dev/null +++ b/apps/dreamverse/dreamverse/routes/assets.py @@ -0,0 +1,72 @@ +"""Raw, bounded media uploads keep large binary data out of websocket messages.""" + +from __future__ import annotations + +import asyncio +from urllib.parse import unquote + +from fastapi import APIRouter, HTTPException, Request, Response +from fastapi.responses import FileResponse +from starlette.concurrency import run_in_threadpool + +from dreamverse.assets import IMAGE_LIMIT, MEDIA_LIMIT, MIME_TYPES, asset_store + +router = APIRouter() +_upload_lock = asyncio.Lock() + + +@router.post("/assets", status_code=201) +async def upload_asset(request: Request) -> dict: + mime_type = request.headers.get("content-type", "").split(";", 1)[0].lower() + if mime_type not in MIME_TYPES: + raise HTTPException(415, "Unsupported asset type. Select a supported image, video, or audio file.") + limit = IMAGE_LIMIT if MIME_TYPES[mime_type][0] == "image" else MEDIA_LIMIT + try: + if int(request.headers.get("content-length", "0")) > limit: + raise HTTPException(413, f"Asset exceeds the {limit // (1024 * 1024)} MB upload limit.") + except ValueError as exc: + raise HTTPException(400, "Invalid Content-Length.") from exc + async with _upload_lock: + try: + path = asset_store.staging_path(mime_type) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + try: + size = 0 + with path.open("xb") as handle: + async for chunk in request.stream(): + size += len(chunk) + if size > limit: + raise HTTPException(413, f"Asset exceeds the {limit // (1024 * 1024)} MB upload limit.") + await run_in_threadpool(handle.write, chunk) + asset = await run_in_threadpool(asset_store.add, path, + unquote(request.headers.get("x-asset-name", "Untitled asset")), mime_type) + return asset.public() + except ValueError as exc: + path.unlink(missing_ok=True) + raise HTTPException(400, str(exc)) from exc + except BaseException: + path.unlink(missing_ok=True) + raise + + +@router.api_route("/assets/{asset_id}", methods=["GET", "HEAD"]) +async def get_asset(asset_id: str) -> FileResponse: + try: + asset = asset_store.get(asset_id) + except ValueError as exc: + raise HTTPException(404, str(exc)) from exc + return FileResponse(asset.path, media_type=asset.mime_type, headers={"X-Content-Type-Options": "nosniff"}) + + +@router.delete("/assets/{asset_id}", status_code=204) +async def delete_asset(asset_id: str) -> Response: + try: + asset_store.get(asset_id) + except ValueError as exc: + raise HTTPException(404, str(exc)) from exc + try: + asset_store.delete(asset_id) + except ValueError as exc: + raise HTTPException(409, str(exc)) from exc + return Response(status_code=204) diff --git a/apps/dreamverse/dreamverse/session/controller.py b/apps/dreamverse/dreamverse/session/controller.py index e3985938d..42aa43b84 100644 --- a/apps/dreamverse/dreamverse/session/controller.py +++ b/apps/dreamverse/dreamverse/session/controller.py @@ -26,6 +26,12 @@ from typing import TYPE_CHECKING from fastapi import WebSocket, WebSocketDisconnect from dreamverse.gpu_pool import GPUSlot +from dreamverse.generation_inputs import ( + GenerationInputs, + pin_generation_inputs, + release_generation_inputs, + resolve_generation_inputs, +) from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image from dreamverse.worker_ipc import MediaChunk, MediaComplete, MediaInit @@ -156,6 +162,7 @@ class SessionController: prompt_worker_task: asyncio.Task | None = None rewrite_seed_prompts_task: asyncio.Task | None = None session_init_image = None + generation_inputs: GenerationInputs | None = None async def session_timeout(): """Close the session after timeout.""" @@ -191,6 +198,18 @@ class SessionController: init_data = {} init_type = init_data.get("type") + try: + next_generation_inputs = resolve_generation_inputs(init_data, ACTIVE_MODEL_ID) + pin_generation_inputs(next_generation_inputs) + generation_inputs = next_generation_inputs + except ValueError as exc: + await ws_send_json({ + "type": "error", + "error_code": "invalid_generation_input", + "message": str(exc), + }) + await websocket.close(code=1008, reason="Invalid generation inputs") + return preset_id = init_data.get("preset_id") preset_label = str(init_data.get("preset_label") or "").strip() initial_rollout_prompt = str(init_data.get("initial_rollout_prompt") or "").strip() @@ -271,6 +290,7 @@ class SessionController: "type": "gpu_assigned", "gpu_id": gpu_id, "session_timeout": SESSION_TIMEOUT_SECONDS, + "generation_mode": generation_inputs.mode, }) await log_event( "gpu_assigned", @@ -361,10 +381,16 @@ class SessionController: return try: + if generation_inputs.mode is not None and payload.get("initial_image") is not None: + raise ValueError("Choose conditioning assets when starting a project; legacy initial_image " + "cannot replace generation mode inputs.") + if "generation_mode" in payload or "conditioning_assets" in payload: + raise ValueError("simple_generate cannot change the mode; use project_init_v1.") replace_session_init_image(payload.get("initial_image")) except ValueError as exc: await ws_send_json({ "type": "error", + "error_code": "invalid_generation_input", "message": str(exc), }) return @@ -417,6 +443,7 @@ class SessionController: }) async def apply_project_init_payload(payload: dict[str, object]) -> bool: + nonlocal generation_inputs nonlocal preset_id nonlocal preset_label nonlocal initial_rollout_prompt @@ -496,15 +523,26 @@ class SessionController: }) return False + next_generation_inputs = None + next_inputs_pinned = False try: + next_generation_inputs = resolve_generation_inputs(payload, ACTIVE_MODEL_ID) + pin_generation_inputs(next_generation_inputs) + next_inputs_pinned = True replace_session_init_image(payload.get("initial_image")) except ValueError as exc: + if next_inputs_pinned: + release_generation_inputs(next_generation_inputs) await ws_send_json({ "type": "error", + "error_code": "invalid_generation_input", "message": str(exc), }) return False + release_generation_inputs(generation_inputs) + generation_inputs = next_generation_inputs + initial_rollout_prompt = next_initial_rollout_prompt enhancement_enabled = next_enhancement_enabled auto_extension_enabled = next_auto_extension_enabled @@ -1171,6 +1209,7 @@ class SessionController: return drained async def enter_project_idle() -> None: + nonlocal generation_inputs nonlocal curated_prompts nonlocal seed_prompt_memory nonlocal curated_idx @@ -1221,6 +1260,9 @@ class SessionController: project_active = False pending_project_end = False + release_generation_inputs(generation_inputs) + generation_inputs = GenerationInputs() + if project_stream_started: project_stream_started = False await ws_send_json({"type": "ltx2_stream_complete"}) @@ -1627,6 +1669,7 @@ class SessionController: segment_idx=segment_idx, image_path=step_image_path, reset_conditioning=step_reset_conditioning, + generation_inputs=generation_inputs, )) segment_generation_active = True try: @@ -1681,10 +1724,10 @@ class SessionController: print(f"[GPU {gpu_id}] Unknown AV event: " f"{type(event).__name__}") - if not step_task.done(): - step_task.cancel() - else: - timings = await step_task + # A GPU command cannot be cancelled by cancelling its + # asyncio waiter. Await completion before releasing pinned + # asset files or making this GPU available to a new user. + timings = await step_task finally: segment_generation_active = False if not step_task.done(): @@ -1808,3 +1851,5 @@ class SessionController: await self.gpu_pool.release(client_id) finally: cleanup_session_init_image(session_init_image) + if generation_inputs is not None: + release_generation_inputs(generation_inputs) diff --git a/apps/dreamverse/dreamverse/tests/test_config.py b/apps/dreamverse/dreamverse/tests/test_config.py index 382794a28..50b1e9d2a 100644 --- a/apps/dreamverse/dreamverse/tests/test_config.py +++ b/apps/dreamverse/dreamverse/tests/test_config.py @@ -186,3 +186,15 @@ def test_config_uses_fasth3_sequence_parallel_default(monkeypatch): assert module.ACTIVE_MODEL_ID == "fast-h3" assert module.MODEL_CONFIG["generation_backend"] == "minimax_h3" assert module.DREAMVERSE_SP_SIZE == 4 + + +def test_full_h3_profile_has_no_preview_adapter_and_longer_session(monkeypatch): + monkeypatch.setenv("DREAMVERSE_MODEL_ID", "full-h3") + monkeypatch.delenv("DREAMVERSE_SP_SIZE", raising=False) + monkeypatch.delenv("DREAMVERSE_SESSION_TIMEOUT_SECONDS", raising=False) + module = _load_config_module() + assert module.MODEL_CONFIG["full_checkpoint"] is True + assert "adapter_repo" not in module.MODEL_CONFIG + assert module.MODEL_CONFIG["num_inference_steps"] == 50 + assert module.DREAMVERSE_SP_SIZE == 4 + assert module.SESSION_TIMEOUT_SECONDS == 7200 diff --git a/apps/dreamverse/dreamverse/tests/test_generation_inputs.py b/apps/dreamverse/dreamverse/tests/test_generation_inputs.py new file mode 100644 index 000000000..50c3cde81 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_generation_inputs.py @@ -0,0 +1,248 @@ +"""Contract regressions independent of CUDA and model weights.""" + +import io +import asyncio +import shutil +import subprocess + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from PIL import Image + +from dreamverse import assets, generation_inputs +from dreamverse.generation_inputs import resolve_generation_inputs +from dreamverse.routes import assets as asset_routes +from dreamverse.tests.test_mock_server import _FakeWebSocket + + +@pytest.fixture +def library(monkeypatch): + store = assets.AssetStore() + monkeypatch.setattr(asset_routes, "asset_store", store) + monkeypatch.setattr(generation_inputs, "asset_store", store) + app = FastAPI() + app.include_router(asset_routes.router) + with TestClient(app) as client: + yield store, client + + +def upload_image(client, color="red"): + image_bytes = io.BytesIO() + Image.new("RGB", (32, 32), color).save(image_bytes, format="PNG") + response = client.post("/assets", content=image_bytes.getvalue(), + headers={"Content-Type": "image/png", "X-Asset-Name": "frame.png"}) + assert response.status_code == 201, response.text + return response.json() + + +def conditioning(asset, role): + return {"asset_id": asset["asset_id"], "role": role} + + +def test_assets_validate_content_and_support_head_range_and_delete(library): + store, client = library + asset = upload_image(client) + assert set(asset) == {"asset_id", "kind", "name", "mime_type", "size", "url"} + assert client.head(asset["url"]).status_code == 200 + response = client.get(asset["url"], headers={"Range": "bytes=0-7"}) + assert response.status_code == 206 + assert response.content == b"\x89PNG\r\n\x1a\n" + assert client.post("/assets", content=b"not an image", headers={"Content-Type": "image/png"}).status_code == 400 + assert client.post("/assets", content=b"", headers={"Content-Type": "image/svg+xml"}).status_code == 415 + assert client.post("/assets", content=b"", headers={"Content-Type": "image/png", + "Content-Length": str(assets.IMAGE_LIMIT + 1)}).status_code == 413 + with pytest.raises(ValueError, match="Invalid asset ID"): + store.get("../../etc/passwd") + assert client.delete(asset["url"]).status_code == 204 + assert client.head(asset["url"]).status_code == 404 + + +def test_generation_pin_prevents_deletion_until_session_releases(library): + _, client = library + asset = upload_image(client) + inputs = resolve_generation_inputs({"generation_mode": "fl2va", "conditioning_assets": [ + conditioning(asset, "first_frame") + ]}, "full-h3") + generation_inputs.pin_generation_inputs(inputs) + generation_inputs.pin_generation_inputs(inputs) + assert client.delete(asset["url"]).status_code == 409 + generation_inputs.release_generation_inputs(inputs) + assert client.delete(asset["url"]).status_code == 409 + generation_inputs.release_generation_inputs(inputs) + assert client.delete(asset["url"]).status_code == 204 + + +def test_legacy_init_remains_compatible_but_explicit_t2va_is_text_only(library): + _, client = library + assert resolve_generation_inputs({"initial_image": {"old": "payload"}}, "fast-ltx2").mode is None + assert resolve_generation_inputs({"generation_mode": "t2va"}, "fast-ltx2").mode == "t2va" + with pytest.raises(ValueError, match="legacy initial_image"): + resolve_generation_inputs({"generation_mode": "t2va", "initial_image": {}}, "full-h3") + asset = upload_image(client) + with pytest.raises(ValueError, match="text only"): + resolve_generation_inputs({"generation_mode": "t2va", "conditioning_assets": [ + conditioning(asset, "reference") + ]}, "full-h3") + + +@pytest.mark.parametrize("mode", ["unknown", None, 3, [], {}]) +def test_unknown_mode_fails_before_assets_are_resolved(mode): + with pytest.raises(ValueError, match="Unknown generation mode"): + resolve_generation_inputs({"generation_mode": mode}, "full-h3") + + +@pytest.mark.parametrize("model_id", ["fast-h3", "fast-ltx2", "fast-ltx23"]) +def test_preview_and_ltx_cannot_advertise_full_h3_modes(model_id): + with pytest.raises(ValueError, match="Full H3"): + resolve_generation_inputs({"generation_mode": "ref2va"}, model_id) + + +def test_fl2va_first_required_last_optional_and_roles_unique(library): + _, client = library + first = upload_image(client) + last = upload_image(client, "blue") + payload = {"generation_mode": "fl2va", "conditioning_assets": [conditioning(first, "first_frame")]} + inputs = resolve_generation_inputs(payload, "full-h3") + assert inputs.first_frame_path.endswith(".png") + assert inputs.last_frame_path is None + payload["conditioning_assets"].append(conditioning(last, "last_frame")) + assert resolve_generation_inputs(payload, "full-h3").last_frame_path is not None + payload["conditioning_assets"].append(conditioning(first, "first_frame")) + with pytest.raises(ValueError, match="exactly one first-frame"): + resolve_generation_inputs(payload, "full-h3") + with pytest.raises(ValueError, match="exactly one first-frame"): + resolve_generation_inputs({"generation_mode": "fl2va", "conditioning_assets": [ + conditioning(last, "last_frame") + ]}, "full-h3") + + +def test_ref_order_is_preserved_and_limits_are_enforced(library): + _, client = library + first, second = upload_image(client), upload_image(client, "blue") + refs = [conditioning(second, "reference"), conditioning(first, "reference")] + payload = {"generation_mode": "ref2va", "conditioning_assets": refs} + inputs = resolve_generation_inputs(payload, "full-h3") + assert [asset.asset_id for asset in inputs.references] == [second["asset_id"], first["asset_id"]] + with pytest.raises(ValueError, match="at most 9 image"): + resolve_generation_inputs({**payload, "conditioning_assets": refs * 5}, "full-h3") + with pytest.raises(ValueError, match="without keyframe roles"): + resolve_generation_inputs({**payload, "conditioning_assets": [conditioning(first, "first_frame")]}, "full-h3") + with pytest.raises(ValueError, match="at most 12"): + resolve_generation_inputs({**payload, "conditioning_assets": refs * 7}, "full-h3") + + +def test_ref_audio_requires_visual_reference(library, monkeypatch): + store, _ = library + monkeypatch.setattr(store, "get", lambda asset_id: assets.StoredAsset(asset_id, "audio", "/audio.wav", "audio", + "audio/wav", 100)) + with pytest.raises(ValueError, match="audio alone"): + resolve_generation_inputs({"generation_mode": "ref2va", "conditioning_assets": [ + {"asset_id": "a" * 32, "role": "reference"} + ]}, "full-h3") + + +def test_ref_rejects_extreme_image_aspect_before_gpu(library): + _, client = library + content = io.BytesIO() + Image.new("RGB", (500, 50), "blue").save(content, format="PNG") + response = client.post("/assets", content=content.getvalue(), headers={"Content-Type": "image/png"}) + assert response.status_code == 201 + with pytest.raises(ValueError, match="aspect ratios"): + resolve_generation_inputs({"generation_mode": "ref2va", "conditioning_assets": [ + conditioning(response.json(), "reference") + ]}, "full-h3") + + +def test_ref_reports_undecodable_image_as_invalid_input(library, monkeypatch, tmp_path): + store, _ = library + broken = tmp_path / "broken.png" + broken.write_bytes(b"not an image") + monkeypatch.setattr(store, "get", lambda asset_id: assets.StoredAsset(asset_id, "image", str(broken), "broken.png", + "image/png", 11)) + with pytest.raises(ValueError, match="could not be decoded"): + resolve_generation_inputs({"generation_mode": "ref2va", "conditioning_assets": [ + {"asset_id": "a" * 32, "role": "reference"} + ]}, "full-h3") + + +def test_audio_upload_rejects_surround_sound(library, monkeypatch): + import json + _, client = library + monkeypatch.setattr(assets.shutil, "which", lambda name: "/usr/bin/ffprobe") + info = {"format": {"format_name": "wav", "duration": "1"}, + "streams": [{"codec_type": "audio", "channels": 6}]} + monkeypatch.setattr(assets.subprocess, "run", lambda *args, **kwargs: subprocess.CompletedProcess( + [], 0, stdout=json.dumps(info).encode(), stderr=b"")) + response = client.post("/assets", content=b"surround wav", headers={"Content-Type": "audio/wav"}) + assert response.status_code == 400 + assert "mono or stereo" in response.json()["detail"] + + +@pytest.mark.parametrize("mime,format_name", [("audio/x-m4a", "mov,mp4,m4a,3gp,3g2,mj2"), ("audio/x-flac", "flac")]) +def test_legacy_audio_mime_aliases_are_accepted(library, monkeypatch, mime, format_name): + """Browsers report x- variants for the M4A and FLAC formats the docs promise.""" + import json + _, client = library + monkeypatch.setattr(assets.shutil, "which", lambda name: "/usr/bin/ffprobe") + info = {"format": {"format_name": format_name, "duration": "1"}, + "streams": [{"codec_type": "audio", "channels": 2}]} + monkeypatch.setattr(assets.subprocess, "run", lambda *args, **kwargs: subprocess.CompletedProcess( + [], 0, stdout=json.dumps(info).encode(), stderr=b"")) + response = client.post("/assets", content=b"audio bytes", headers={"Content-Type": mime}) + assert response.status_code == 201, response.text + assert response.json()["kind"] == "audio" + assert response.json()["mime_type"] == mime + + +@pytest.mark.parametrize("entries", [None, {}, "x", [{"path": "/etc/passwd", "role": "reference"}], + [{"asset_id": "x", "role": "unknown"}]]) +def test_malformed_conditioning_is_rejected(entries): + with pytest.raises(ValueError): + resolve_generation_inputs({"generation_mode": "ref2va", "conditioning_assets": entries}, "full-h3") + + +@pytest.mark.parametrize("mode", ["t2va", "fl2va", "ref2va"]) +def test_mock_streams_all_valid_modes_and_releases_assets(library, monkeypatch, mode): + from dreamverse import mock_server + _, client = library + monkeypatch.setattr(mock_server, "MOCK_SEGMENT_BYTES", b"mock-fmp4") + monkeypatch.setattr(mock_server, "LATENCY_MS", 1) + image = upload_image(client) + refs = [] if mode == "t2va" else [conditioning(image, "first_frame" if mode == "fl2va" else "reference")] + ws = _FakeWebSocket([ + (0, {"type": "session_init_v2", "generation_mode": mode, "conditioning_assets": refs, + "curated_prompts": ["A bird flies over a lake."], "single_clip_mode": True, + "enhancement_enabled": False}), + (0.15, {"type": "leave"}), + ]) + asyncio.run(mock_server.websocket_endpoint(ws)) + assert not [event for event in ws.sent_json if event["type"] == "error"] + assert any(event["type"] == "media_segment_complete" for event in ws.sent_json) + assert ws.sent_bytes + assert client.delete(image["url"]).status_code == 204 + + +def test_mock_rejects_invalid_mode_before_gpu_assignment(library): + from dreamverse import mock_server + ws = _FakeWebSocket([(0, {"type": "session_init_v2", "generation_mode": "fl2va"})]) + asyncio.run(mock_server.websocket_endpoint(ws)) + assert not any(event["type"] == "gpu_assigned" for event in ws.sent_json) + errors = [event for event in ws.sent_json if event["type"] == "error"] + assert errors[0]["error_code"] == "invalid_generation_input" + assert "first-frame" in errors[0]["message"] + + +@pytest.mark.skipif(not shutil.which("ffmpeg") or not shutil.which("ffprobe"), reason="ffmpeg + ffprobe required") +@pytest.mark.parametrize("kind,mime,suffix", [("video", "video/mp4", ".mp4"), ("audio", "audio/wav", ".wav")]) +def test_actual_video_and_audio_upload_validation(library, tmp_path, kind, mime, suffix): + _, client = library + media_path = tmp_path / f"sample{suffix}" + source = "testsrc2=size=64x64:rate=24" if kind == "video" else "sine=frequency=440:sample_rate=24000" + command = [shutil.which("ffmpeg"), "-v", "error", "-f", "lavfi", "-i", source, "-t", "0.5", str(media_path)] + subprocess.run(command, check=True, capture_output=True, timeout=30) + response = client.post("/assets", content=media_path.read_bytes(), headers={"Content-Type": mime}) + assert response.status_code == 201, response.text + assert response.json()["kind"] == kind + response = client.post("/assets", content=b"#EXTM3U\nhttp://example.com/stream", headers={"Content-Type": mime}) + assert response.status_code == 400 diff --git a/apps/dreamverse/dreamverse/tests/test_generation_mode_routing.py b/apps/dreamverse/dreamverse/tests/test_generation_mode_routing.py new file mode 100644 index 000000000..482391a42 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_generation_mode_routing.py @@ -0,0 +1,214 @@ +"""CPU contract tests; fake executors do not validate generated-media quality.""" + +from __future__ import annotations + +import importlib.util +import pickle +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +from PIL import Image + +from dreamverse.config import MODEL_REGISTRY +from dreamverse.generation_inputs import GenerationAsset, GenerationInputs +from dreamverse.generation_worker import VideoGenerationWorker +from dreamverse.minimax_h3_generation import MiniMaxH3GenerationBackend +from dreamverse.worker_ipc import UserStepPayload + + +@pytest.fixture +def fastvideo_api(monkeypatch): + """Use the actual lightweight API schema with only GPU execution replaced.""" + schema_path = Path(__file__).resolve().parents[4] / "fastvideo/api/schema.py" + spec = importlib.util.spec_from_file_location("dreamverse_test_api_schema", schema_path) + assert spec is not None and spec.loader is not None + schema = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, schema) + spec.loader.exec_module(schema) + package = ModuleType("fastvideo") + package.__path__ = [] + package.VideoGenerator = SimpleNamespace(from_config=Mock()) + monkeypatch.setitem(sys.modules, "fastvideo", package) + monkeypatch.setitem(sys.modules, "fastvideo.api", schema) + schema.MiniMaxH3Reference = lambda **kwargs: SimpleNamespace(**kwargs) + monkeypatch.setattr("dreamverse.minimax_h3_generation.torch.cuda.synchronize", lambda: None) + monkeypatch.setattr("dreamverse.minimax_h3_generation.torch.cuda.empty_cache", lambda: None) + return package.VideoGenerator.from_config + + +class RecordingGenerator: + def __init__(self): + self.requests = [] + self.images = [] + self.closed = False + + def shutdown(self): + self.closed = True + + def generate(self, request): + self.requests.append(request) + self.images.append(tuple(None if image is None else np.asarray(image).copy() + for image in (request.inputs.pil_image, request.inputs.last_image))) + return SimpleNamespace( + frames=[np.full((2, 3, 3), 7, dtype=np.uint8), np.full((2, 3, 3), 29, dtype=np.uint8)], + audio=np.zeros((2, 16), dtype=np.float32), + audio_sample_rate=44100, + generation_time=0.1, + ) + + +def prepared_backend(monkeypatch): + backend = MiniMaxH3GenerationBackend(0) + backend.model_config = dict(MODEL_REGISTRY["full-h3"]) + backend.generator = RecordingGenerator() + monkeypatch.setattr(backend, "_gpu_mem", lambda: "fake executor") + return backend + + +def test_ipc_preserves_immutable_ordered_references(): + inputs = GenerationInputs("ref2va", ( + GenerationAsset("second", "video", "/assets/second.mp4", "reference"), + GenerationAsset("first", "image", "/assets/first.png", "reference"), + )) + payload = UserStepPayload("follow the references", 1, None, True, inputs) + restored = pickle.loads(pickle.dumps(payload)) + assert restored == payload + assert [asset.asset_id for asset in restored.generation_inputs.references] == ["second", "first"] + + +def test_worker_passes_conditioning_to_selected_backend(): + inputs = GenerationInputs("t2va") + worker = VideoGenerationWorker(0) + worker.backend = Mock() + worker.generate_step("prompt", 1, None, True, inputs) + worker.backend.generate_step.assert_called_once_with("prompt", 1, None, True, generation_inputs=inputs) + + +def test_full_h3_uses_full_weights_without_preview_lora(monkeypatch, fastvideo_api): + backend = prepared_backend(monkeypatch) + old_generator = backend.generator + fastvideo_api.return_value = RecordingGenerator() + monkeypatch.setattr("dreamverse.minimax_h3_generation.DREAMVERSE_SP_SIZE", 4) + backend.initialize(MODEL_REGISTRY["full-h3"]) + config = fastvideo_api.call_args.args[0] + assert old_generator.closed + assert config.pipeline.components.lora_path is None + assert config.pipeline.components.override_pipeline_cls_name is None + assert config.engine.use_fsdp_inference + assert config.engine.num_gpus == 4 + assert not config.pipeline.experimental["inference_torch_compile"] + + +def test_fl2va_maps_endpoints_only_on_initial_segment(monkeypatch, fastvideo_api, tmp_path): + first = tmp_path / "first.png" + last = tmp_path / "last.png" + Image.new("RGB", (3, 2), (10, 20, 30)).save(first) + Image.new("RGB", (3, 2), (40, 50, 60)).save(last) + inputs = GenerationInputs("fl2va", ( + GenerationAsset("first", "image", str(first), "first_frame"), + GenerationAsset("last", "image", str(last), "last_frame"), + )) + backend = prepared_backend(monkeypatch) + first_result = backend.generate_step("first", 1, None, True, inputs) + later_result = backend.generate_step("later", 2, None, False, inputs) + assert backend.generator.images[0][0][0, 0].tolist() == [10, 20, 30] + assert backend.generator.images[0][1][0, 0].tolist() == [40, 50, 60] + assert backend.generator.images[1][0][0, 0].tolist() == [29, 29, 29] + assert backend.generator.images[1][1] is None + assert first_result.head_trim_frames == 0 + assert later_result.head_trim_frames == 1 + assert backend.generator.requests[0].sampling.num_inference_steps == 50 + + +def test_ref2va_switches_pipeline_and_preserves_reference_order(monkeypatch, fastvideo_api): + inputs = GenerationInputs("ref2va", ( + GenerationAsset("video", "video", "/assets/reference.mp4", "reference"), + GenerationAsset("audio", "audio", "/assets/reference.wav", "reference"), + GenerationAsset("image", "image", "/assets/reference.png", "reference"), + )) + backend = prepared_backend(monkeypatch) + base_generator = backend.generator + reference_generator = RecordingGenerator() + + def load(config): + assert base_generator.closed, "Old executor must release memory before loading reference weights" + assert config.pipeline.components.override_pipeline_cls_name == "MiniMaxH3Ref2VAModularPipeline" + assert config.pipeline.workload_type == "i2v" + assert config.pipeline.components.lora_path is None + return reference_generator + + fastvideo_api.side_effect = load + backend.generate_step("first", 1, None, True, inputs) + result = backend.generate_step("second", 2, None, False, inputs) + assert fastvideo_api.call_count == 1 + for request in reference_generator.requests: + assert [(reference.media_type, reference.source) for reference in request.inputs.references] == [ + ("video", "/assets/reference.mp4"), ("audio", "/assets/reference.wav"), ("image", "/assets/reference.png") + ] + assert request.inputs.pil_image is None + assert request.inputs.last_image is None + assert result.head_trim_frames == result.head_trim_audio_frames == 0 + assert backend.continuation_image is None + + fastvideo_api.side_effect = None + fastvideo_api.return_value = RecordingGenerator() + backend.generate_step("new project", 1, None, True, GenerationInputs("t2va")) + assert reference_generator.closed + config = fastvideo_api.call_args.args[0] + assert config.pipeline.components.override_pipeline_cls_name is None + assert backend.pipeline_mode == "base" + + +def test_ref2va_pipeline_switch_failure_drops_unloaded_executor(monkeypatch, fastvideo_api): + backend = prepared_backend(monkeypatch) + old_generator = backend.generator + fastvideo_api.side_effect = RuntimeError("checkpoint unavailable") + with pytest.raises(RuntimeError, match="checkpoint unavailable"): + backend.generate_step("prompt", 1, None, True, GenerationInputs("ref2va")) + assert old_generator.closed + assert backend.generator is None + + +def test_failed_pipeline_switch_reloads_on_the_next_step(monkeypatch, fastvideo_api): + """A failed base<->ref2va switch must not strand the slot for later steps.""" + backend = prepared_backend(monkeypatch) + fastvideo_api.side_effect = RuntimeError("checkpoint unavailable") + with pytest.raises(RuntimeError, match="checkpoint unavailable"): + backend.generate_step("prompt", 1, None, True, GenerationInputs("ref2va")) + + fastvideo_api.side_effect = None + fastvideo_api.return_value = RecordingGenerator() + backend.generate_step("retry", 1, None, True, GenerationInputs("ref2va")) + + assert fastvideo_api.call_count == 2 + assert backend.pipeline_mode == "ref2va" + assert backend.generator is not None + + +def test_mode_cannot_switch_mid_project(monkeypatch, fastvideo_api): + backend = prepared_backend(monkeypatch) + with pytest.raises(ValueError, match="middle of a project"): + backend.generate_step("prompt", 2, None, False, GenerationInputs("ref2va")) + fastvideo_api.assert_not_called() + + +@pytest.mark.parametrize("mode", ["fl2va", "ref2va"]) +def test_preview_rejects_unsupported_generation_modes(monkeypatch, fastvideo_api, mode): + backend = prepared_backend(monkeypatch) + backend.model_config = dict(MODEL_REGISTRY["fast-h3"]) + with pytest.raises(ValueError, match="full-h3"): + backend.generate_step("prompt", 1, None, True, GenerationInputs(mode)) + assert backend.generator.requests == [] + + +def test_legacy_h3_continuation_is_preserved(monkeypatch, fastvideo_api): + backend = prepared_backend(monkeypatch) + backend.generate_step("first", 1, None, False) + result = backend.generate_step("second", 2, None, False) + assert backend.generator.requests[0].inputs.pil_image is None + assert backend.generator.images[1][0][0, 0].tolist() == [29, 29, 29] + assert result.head_trim_frames == 1 diff --git a/apps/dreamverse/dreamverse/tests/test_generation_session_inputs.py b/apps/dreamverse/dreamverse/tests/test_generation_session_inputs.py new file mode 100644 index 000000000..f13f2c35f --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_generation_session_inputs.py @@ -0,0 +1,278 @@ +"""Session-mode validation and IPC handoff without a GPU worker process.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest + +from dreamverse.generation_inputs import GenerationAsset, GenerationInputs +from dreamverse.worker_ipc import MediaChunk, MediaComplete, MediaInit + + +@pytest.fixture +def controller_module(monkeypatch): + gpu_pool = ModuleType("dreamverse.gpu_pool") + gpu_pool.GPUSlot = object + monkeypatch.setitem(sys.modules, "dreamverse.gpu_pool", gpu_pool) + path = Path(__file__).resolve().parents[1] / "session/controller.py" + spec = importlib.util.spec_from_file_location("dreamverse_test_session_controller", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "ACTIVE_MODEL_ID", "full-h3") + monkeypatch.setattr(module, "pin_generation_inputs", Mock()) + monkeypatch.setattr(module, "release_generation_inputs", Mock()) + return module + + +class Socket: + def __init__(self): + self.incoming = asyncio.Queue() + self.outgoing = asyncio.Queue() + self.messages = [] + self.closed = False + + async def accept(self): + pass + + async def receive_json(self): + return await self.incoming.get() + + async def send_json(self, payload): + self.messages.append(payload) + await self.outgoing.put(payload) + + async def send_bytes(self, payload): + pass + + async def close(self, **kwargs): + self.closed = True + + async def wait_for(self, kind): + while True: + payload = await asyncio.wait_for(self.outgoing.get(), 3) + if payload["type"] == kind: + return payload + + +class Slot: + def __init__(self): + self.shared_stream_buffer = None + self.queue = asyncio.Queue() + self.calls = [] + + async def join_user(self, *args, **kwargs): + pass + + def register_stream_queue(self, client_id): + return self.queue + + def unregister_stream_queue(self, client_id): + pass + + async def user_step(self, client_id, **kwargs): + self.calls.append(kwargs) + segment_idx = kwargs["segment_idx"] + await self.queue.put(MediaInit(client_id, segment_idx, "test", "video/mp4", False)) + await self.queue.put(MediaChunk(client_id, segment_idx, "test", chunk=b"test")) + await self.queue.put(MediaComplete(client_id, segment_idx, "test", 1)) + return {"e2e_latency_ms": 1.0} + + +class Pool: + def __init__(self): + self.slot = Slot() + self.acquire_count = 0 + + def get_status(self): + return {"queue_size": 0, "available_gpus": 1, "total_gpus": 1} + + async def acquire(self, *args): + self.acquire_count += 1 + return 0, self.slot + + async def release(self, *args): + pass + + +def start_controller(module, socket, pool): + enhancer = SimpleNamespace( + resolve_rewrite_model=lambda value: "test-model", + resolve_rewrite_system_prompt=lambda value: "test-system", + resolve_rewrite_temperature=lambda value: 1.0, + ) + controller = module.SessionController(socket, pool, enhancer, None, None) + return asyncio.create_task(controller.run()) + + +def test_invalid_initial_mode_does_not_acquire_gpu(controller_module): + async def scenario(): + socket, pool = Socket(), Pool() + await socket.incoming.put({"type": "session_init_v2", "generation_mode": "unknown"}) + await asyncio.wait_for(start_controller(controller_module, socket, pool), 3) + error = next(message for message in socket.messages if message["type"] == "error") + assert error["error_code"] == "invalid_generation_input" + assert pool.acquire_count == 0 + assert socket.closed + controller_module.pin_generation_inputs.assert_not_called() + + asyncio.run(scenario()) + + +def test_new_project_replaces_conditioning_and_passes_it_to_gpu(controller_module, monkeypatch): + first = GenerationInputs("t2va") + second = GenerationInputs("fl2va", (GenerationAsset("first", "image", "/assets/first.png", "first_frame"),)) + monkeypatch.setattr(controller_module, "resolve_generation_inputs", Mock(side_effect=[first, second])) + + async def scenario(): + socket, pool = Socket(), Pool() + await socket.incoming.put({ + "type": "session_init_v2", "generation_mode": "t2va", "curated_prompts": ["first prompt"], + "enhancement_enabled": False, + }) + task = start_controller(controller_module, socket, pool) + try: + await socket.wait_for("media_segment_complete") + await socket.incoming.put({"type": "end_project_keep_session"}) + await socket.wait_for("project_idle") + assert first in [call.args[0] for call in controller_module.release_generation_inputs.call_args_list] + await socket.incoming.put({ + "type": "project_init_v1", "generation_mode": "fl2va", "curated_prompts": ["second prompt"], + "enhancement_enabled": False, + }) + await socket.wait_for("media_segment_complete") + assert [call["generation_inputs"] for call in pool.slot.calls] == [first, second] + assert pool.slot.calls[1]["segment_idx"] == 1 + assert pool.slot.calls[1]["reset_conditioning"] + await socket.incoming.put({"type": "leave"}) + await asyncio.wait_for(task, 3) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert [call.args[0] for call in controller_module.pin_generation_inputs.call_args_list] == [first, second] + assert second in [call.args[0] for call in controller_module.release_generation_inputs.call_args_list] + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("injection", [ + {"initial_image": {"data_url": "not allowed"}}, + {"generation_mode": "ref2va"}, + {"conditioning_assets": []}, +]) +def test_simple_generate_cannot_replace_locked_inputs(controller_module, injection): + async def scenario(): + socket, pool = Socket(), Pool() + await socket.incoming.put({ + "type": "session_init_v2", "generation_mode": "t2va", "single_clip_mode": True, + "enhancement_enabled": False, + }) + task = start_controller(controller_module, socket, pool) + try: + await socket.wait_for("gpu_assigned") + await socket.incoming.put({"type": "simple_generate", "prompt": "prompt", **injection}) + error = await socket.wait_for("error") + assert error["error_code"] == "invalid_generation_input" + assert "project" in error["message"] + assert pool.slot.calls == [] + await socket.incoming.put({"type": "leave"}) + await asyncio.wait_for(task, 3) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(scenario()) + + +def test_disconnect_waits_for_worker_before_releasing_assets(controller_module, monkeypatch): + async def scenario(): + socket, pool = Socket(), Pool() + worker_started = asyncio.Event() + worker_finished = asyncio.Event() + proceed = asyncio.Event() + + async def slow_step(client_id, **kwargs): + worker_started.set() + await proceed.wait() + worker_finished.set() + return {"e2e_latency_ms": 1.0} + + pool.slot.user_step = slow_step + await socket.incoming.put({ + "type": "session_init_v2", "generation_mode": "t2va", "curated_prompts": ["prompt"], + "enhancement_enabled": False, + }) + task = start_controller(controller_module, socket, pool) + try: + await asyncio.wait_for(worker_started.wait(), 3) + await socket.incoming.put({"type": "leave"}) + await asyncio.sleep(0.07) + assert not task.done() + controller_module.release_generation_inputs.assert_not_called() + proceed.set() + await asyncio.wait_for(task, 3) + assert worker_finished.is_set() + controller_module.release_generation_inputs.assert_called_once() + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(scenario()) + + +@pytest.fixture +def gpu_pool_module(monkeypatch): + streaming = ModuleType("dreamverse.av_streaming") + for name in ("StreamChunk", "StreamComplete", "StreamEvent", "StreamInit", "generate_stream_id", "stream_fmp4"): + setattr(streaming, name, object) + streaming.SHARED_STREAM_BUFFER_BYTES = 1024 + streaming.USE_SHARED_STREAM_BUFFER = False + monkeypatch.setitem(sys.modules, "dreamverse.av_streaming", streaming) + path = Path(__file__).resolve().parents[1] / "gpu_pool.py" + spec = importlib.util.spec_from_file_location("dreamverse_test_gpu_pool", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "pin_generation_inputs", Mock()) + monkeypatch.setattr(module, "release_generation_inputs", Mock()) + return module + + +def test_gpu_step_timeout_keeps_assets_pinned_until_late_worker_completion(gpu_pool_module): + from dreamverse.worker_ipc import StepComplete + + async def scenario(): + slot = gpu_pool_module.GPUSlot(0, "0") + inputs = GenerationInputs("ref2va", (GenerationAsset("ref", "image", "/assets/ref.png", "reference"),)) + + async def timeout(command, timeout): + assert command.payload.generation_inputs == inputs + raise asyncio.TimeoutError + + slot._send_command_tagged = timeout + with pytest.raises(asyncio.TimeoutError): + await slot.user_step("user", "prompt", generation_inputs=inputs) + gpu_pool_module.pin_generation_inputs.assert_called_once_with(inputs) + gpu_pool_module.release_generation_inputs.assert_not_called() + + def late_response(timeout): + slot._active = False + return StepComplete("user", 1, {}) + + slot.response_queue = SimpleNamespace(get=late_response) + slot._active = True + await slot._response_reader() + gpu_pool_module.release_generation_inputs.assert_called_once_with(inputs) + assert slot._step_asset_inputs == {} + + asyncio.run(scenario()) diff --git a/apps/dreamverse/dreamverse/tests/test_import_contract.py b/apps/dreamverse/dreamverse/tests/test_import_contract.py index 584083548..fff0bd4a0 100644 --- a/apps/dreamverse/dreamverse/tests/test_import_contract.py +++ b/apps/dreamverse/dreamverse/tests/test_import_contract.py @@ -1,4 +1,6 @@ import ast +import subprocess +import sys from pathlib import Path ALLOWED_PREFIXES = ( @@ -48,3 +50,39 @@ def test_dreamverse_server_imports_only_public_fastvideo_surfaces() -> None: bad.append((str(path.relative_to(root)), getattr(node, "lineno", 0), name)) assert bad == [], f"Forbidden internal imports: {bad}" + + +def test_h3_reference_public_export_is_lazy_and_preserves_type_identity() -> None: + """Only explicit reference usage should load H3's optional GPU dependencies.""" + repo_root = Path(__file__).resolve().parents[4] + # Isolate the import graph: keep the real public API implementation/schema, + # substituting only the unrelated legacy sampling module and heavy H3 leaf. + script = r''' +import importlib +import sys +from pathlib import Path +from types import ModuleType + +root = Path(sys.argv[1]) +fastvideo = ModuleType("fastvideo") +fastvideo.__path__ = [str(root / "fastvideo")] +sys.modules["fastvideo"] = fastvideo +sampling = ModuleType("fastvideo.api.sampling_param") +sampling.SamplingParam = type("SamplingParam", (), {}) +sys.modules[sampling.__name__] = sampling + +api = importlib.import_module("fastvideo.api") +assert "MiniMaxH3Reference" in api.__all__ +assert "MiniMaxH3Reference" not in vars(api) +assert not any(name.startswith("fastvideo.pipelines") for name in sys.modules) + +internal = ModuleType("fastvideo.pipelines.basic.minimax_h3.reference") +internal.MiniMaxH3Reference = type("MiniMaxH3Reference", (), {}) +sys.modules[internal.__name__] = internal +from fastvideo.api import MiniMaxH3Reference +assert MiniMaxH3Reference is internal.MiniMaxH3Reference +assert api.MiniMaxH3Reference is internal.MiniMaxH3Reference +assert not hasattr(api, "UnknownReference") +''' + result = subprocess.run([sys.executable, "-c", script, str(repo_root)], capture_output=True, text=True, timeout=30) + assert result.returncode == 0, result.stderr diff --git a/apps/dreamverse/dreamverse/tests/test_session_logging.py b/apps/dreamverse/dreamverse/tests/test_session_logging.py index ead39b924..db4b05733 100644 --- a/apps/dreamverse/dreamverse/tests/test_session_logging.py +++ b/apps/dreamverse/dreamverse/tests/test_session_logging.py @@ -105,6 +105,7 @@ class _FakeSlot: segment_idx: int, reset_conditioning: bool, image_path: str | None = None, + generation_inputs=None, ): self.calls.append({ "client_id": client_id, diff --git a/apps/dreamverse/dreamverse/worker_ipc.py b/apps/dreamverse/dreamverse/worker_ipc.py index cd01867c2..2fb80ba3e 100644 --- a/apps/dreamverse/dreamverse/worker_ipc.py +++ b/apps/dreamverse/dreamverse/worker_ipc.py @@ -15,6 +15,8 @@ from __future__ import annotations from dataclasses import dataclass +from dreamverse.generation_inputs import GenerationInputs + # ---- User-scoped events (carry user_id) ------------------------------------ @@ -147,6 +149,7 @@ class UserStepPayload: segment_idx: int image_path: str | None reset_conditioning: bool + generation_inputs: GenerationInputs | None = None @dataclass(frozen=True) diff --git a/apps/dreamverse/scripts/slurm/README.md b/apps/dreamverse/scripts/slurm/README.md new file mode 100644 index 000000000..ff9efd0f0 --- /dev/null +++ b/apps/dreamverse/scripts/slurm/README.md @@ -0,0 +1,107 @@ +# Dreamverse on Slurm + +Run Full H3 inside a one-node, four-GPU allocation. The maintained H3 examples +default to four GPUs; this is a starting configuration, not a measured minimum. +The full checkpoint supports T2VA, FL2VA, and Ref2VA. The FastH3 Preview profile +is a separate T2VA configuration. + +`launch_backend.sh` checks that it is inside an `srun` step, preserves +`CUDA_VISIBLE_DEVICES`, and replaces itself with the backend process. It does +not allocate GPUs, kill existing processes, or source a personal credentials +file. The local `dreamverse-deploy` helper is not suitable for a shared Slurm +cluster because it kills processes by physical GPU and port. + +## Prepare and allocate + +Keep the checkout, weights, outputs, and logs on storage visible to the compute +node. Source installation is documented in the [GPU guide](../../../../docs/getting_started/installation/gpu.md). +On ARM64 GB200 use CUDA 13, a matching PyTorch build, and kernels built for +`sm_100`; the DGX Spark `sm_121` kernel image is not the GB200 image. + +The repository's image workflow publishes an ARM64 GB200 variant under +`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0-sm100-latest`. +Resolve that tag to a digest for reproducible runs. If your compute nodes use +Pyxis/Enroot, pass the approved image or a prepared SquashFS file to +`srun --container-image`, with explicit mounts for your checkout and model cache. +The Dreamverse-specific Docker images are currently AMD64-only. + +For the Slinky customer partition, a bounded allocation is: + +```bash +salloc --account=customer --qos=normal --partition=hpc-rack-1 \ + --nodes=1 --ntasks=1 --cpus-per-task=72 --gres=gpu:nvidia_gb200:4 \ + --mem=800G --time=02:00:00 --job-name=dreamverse +srun --ntasks=1 --pty bash +``` + +Wait for Slurm to grant the allocation before entering the compute step. A +successful SSH login does not grant GPU resources. Inspect pending capacity +with `squeue -u "$USER" --start`; do not attach to another user's job. + +The checkpoint includes duplicate release layouts. Download the diffusers +components needed by both base and reference pipelines, rather than the whole +repository (about 210 GB versus about 498 GB at revision +`42ed227ee7df40d41602854ae760620d6eb651fe`): + +```bash +hf download MiniMaxAI/MiniMax-H3 \ + --revision 42ed227ee7df40d41602854ae760620d6eb651fe \ + --include model_index.json --include modular_model_index.json \ + --include 'audio_scheduler/*' --include 'audio_vae/*' \ + --include 'processor/*' --include 'scheduler/*' \ + --include 'text_encoder/*' --include 'tokenizer/*' \ + --include 'transformer/*' --include 'transformer_ref/*' --include 'vae/*' \ + --local-dir /path/to/models/MiniMax-H3 +``` + +The GPU environment needs `fastvideo[dreamverse]`, the Dreamverse workspace +package, and FFmpeg with H.264/AAC encoders. In a prepared FastVideo image, +install the checked-out code and its Dreamverse dependencies in that image's +Python environment. Keep its matching CUDA/PyTorch/kernel stack intact. + +## Start and connect + +From the checked-out repository inside the allocated step: + +```bash +export DREAMVERSE_PYTHON=/path/to/environment/bin/python +export DREAMVERSE_MODEL_PATH=/path/to/models/MiniMax-H3 +export FASTVIDEO_DREAMVERSE_HOME=/path/to/persistent/dreamverse-state +bash apps/dreamverse/scripts/slurm/launch_backend.sh +``` + +The default backend binds port 8009 on the private compute node. Connect through +the login node from your laptop, replacing `COMPUTE_NODE_IP` with the allocated +node's `NodeAddr` from `scontrol show node`: + +```bash +ssh -N -L 8009:COMPUTE_NODE_IP:8009 USER@LOGIN_NODE +``` + +In another laptop terminal, run the frontend from your local checkout: + +```bash +cd apps/dreamverse/web +BACKEND_HOST=127.0.0.1 BACKEND_PORT=8009 npm run dev +``` + +Open `http://localhost:5299`. `/healthz` reports the server process; `/readyz` +reports model readiness. Full H3 loads and generates more slowly than the +Preview adapter. Keep prompt enhancement disabled in the UI unless the +runtime has the selected provider's credentials. + +## Verify and stop + +Check all three modes with small, valid user-owned assets. Capture the selected +mode and assets, WebSocket errors or completion events, the generated video and +audio, and GPU memory usage. Also verify actionable validation errors and +backward compatibility with clients that omit `generation_mode`. + +Use the frontend Playwright instructions in the +[Dreamverse development guide](../../../../docs/contributing/dreamverse-development.md) +against the forwarded backend. A mock-server demo validates UI and protocol +behavior; it is not evidence of GPU generation. + +Stop the backend with Ctrl-C, exit the compute step, and release your allocation. +For a detached allocation, use `scancel YOUR_JOB_ID`. Cancel a pending demo job +when it is no longer needed; do not leave an unattended reservation queued. diff --git a/apps/dreamverse/scripts/slurm/launch_backend.sh b/apps/dreamverse/scripts/slurm/launch_backend.sh new file mode 100644 index 000000000..2d3e3fd79 --- /dev/null +++ b/apps/dreamverse/scripts/slurm/launch_backend.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# Run inside an existing Slurm step. Slurm owns the GPU visibility and lifetime. +set -euo pipefail + +if [[ -z "${SLURM_JOB_ID:-}" || -z "${SLURM_STEP_ID:-}" ]]; then + echo "Run this launcher inside an allocated Slurm step (srun), not on the login node." >&2 + exit 2 +fi + +script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repo_root="$(cd -- "${script_dir}/../../../.." && pwd)" +python_bin="${DREAMVERSE_PYTHON:-${repo_root}/.venv/bin/python}" +if [[ ! -x "${python_bin}" ]]; then + echo "Set DREAMVERSE_PYTHON to a Python environment with fastvideo[dreamverse] installed." >&2 + exit 2 +fi + +export DREAMVERSE_MODEL_ID="${DREAMVERSE_MODEL_ID:-full-h3}" +export DREAMVERSE_SP_SIZE="${DREAMVERSE_SP_SIZE:-4}" +export FASTVIDEO_GPU_COUNT="${FASTVIDEO_GPU_COUNT:-${DREAMVERSE_SP_SIZE}}" +export FASTVIDEO_ENABLE_STARTUP_WARMUP="${FASTVIDEO_ENABLE_STARTUP_WARMUP:-0}" +export ENABLE_TORCH_COMPILE="${ENABLE_TORCH_COMPILE:-0}" +export STREAM_MODE="${STREAM_MODE:-av_fmp4}" +export PYTHONPATH="${repo_root}/apps/dreamverse:${repo_root}${PYTHONPATH:+:${PYTHONPATH}}" +export PYTHONUNBUFFERED=1 + +"${python_bin}" - <<'PY' +import os +import shutil + +import torch + +expected = int(os.environ["DREAMVERSE_SP_SIZE"]) +visible = torch.cuda.device_count() +if expected < 1 or visible < expected: + raise SystemExit(f"The Slurm step exposes {visible} GPUs; DREAMVERSE_SP_SIZE requires {expected}.") +ffmpeg = os.environ.get("FASTVIDEO_FFMPEG_BIN", "ffmpeg") +if not shutil.which(ffmpeg): + raise SystemExit("FFmpeg is missing; install it in the compute environment or set FASTVIDEO_FFMPEG_BIN.") +print(f"Slurm job {os.environ['SLURM_JOB_ID']}: {visible} visible GPUs; using {expected} per worker") +for index in range(expected): + properties = torch.cuda.get_device_properties(index) + print(f" GPU {index}: {properties.name}, {properties.total_memory / 2**30:.1f} GiB") +PY + +cd "${repo_root}" +exec "${python_bin}" -m dreamverse.server_entry \ + --host "${DREAMVERSE_BIND_HOST:-0.0.0.0}" \ + --port "${DREAMVERSE_BACKEND_PORT:-8009}" "$@" diff --git a/apps/dreamverse/web/e2e/generation-modes.spec.ts b/apps/dreamverse/web/e2e/generation-modes.spec.ts new file mode 100644 index 000000000..9b0359ea9 --- /dev/null +++ b/apps/dreamverse/web/e2e/generation-modes.spec.ts @@ -0,0 +1,126 @@ +import { execFileSync } from "node:child_process"; +import { readFile } from "node:fs/promises"; +import path from "node:path"; +import { test, expect } from "@playwright/test"; + +const imagePath = path.resolve("public/k2.png"); +const framePrompt = "A paper fox walks through a sunlit forest, gentle birdsong."; + +function makeAudio(sampleRate = 8000, seconds = 1): Buffer { + const sampleCount = sampleRate * seconds; + const bytes = Buffer.alloc(44 + sampleCount * 2); + bytes.write("RIFF", 0); bytes.writeUInt32LE(bytes.length - 8, 4); bytes.write("WAVEfmt ", 8); + bytes.writeUInt32LE(16, 16); bytes.writeUInt16LE(1, 20); bytes.writeUInt16LE(1, 22); + bytes.writeUInt32LE(sampleRate, 24); bytes.writeUInt32LE(sampleRate * 2, 28); + bytes.writeUInt16LE(2, 32); bytes.writeUInt16LE(16, 34); bytes.write("data", 36); + bytes.writeUInt32LE(sampleCount * 2, 40); + for (let i = 0; i < sampleCount; i++) bytes.writeInt16LE(Math.round(Math.sin(i * 440 * 2 * Math.PI / sampleRate) * 1000), 44 + i * 2); + return bytes; +} + +test.describe("generation modes through the mock runtime", () => { + for (const mode of ["t2va", "fl2va", "ref2va"] as const) { + test(`${mode} sends validated assets and plays a clearly labeled sample`, async ({ page, request }, testInfo) => { + const response = await request.get("/generation-capabilities"); + const capabilities = response.ok() ? await response.json() : {}; + test.skip(capabilities.mock !== true, "This test uses the CPU mock runtime; it must not silently allocate a real GPU."); + const sent: Record[] = []; + const received: Record[] = []; + page.on("websocket", (socket) => { + socket.on("framesent", ({ payload }) => { if (typeof payload === "string") { try { sent.push(JSON.parse(payload)); } catch {} } }); + socket.on("framereceived", ({ payload }) => { if (typeof payload === "string") { try { received.push(JSON.parse(payload)); } catch {} } }); + }); + await page.goto("/"); + await expect(page.getByText(/Demo runtime · Sample playback only/)).toBeVisible(); + const modeSelect = page.getByRole("combobox", { name: "Generation mode" }); + const modeLabel = mode === "ref2va" ? "Ref2VA" : mode.toUpperCase(); + await modeSelect.click(); + await page.getByRole("option", { name: modeLabel, exact: true }).click(); + await expect(modeSelect).toHaveText(modeLabel); + await page.getByLabel("Continuation prompt").fill(framePrompt); + const uploadedIds: string[] = []; + page.on("response", async (uploadResponse) => { + if (uploadResponse.request().method() === "POST" && uploadResponse.url().endsWith("/assets") && uploadResponse.ok()) { + const asset = await uploadResponse.json().catch(() => null); + if (asset?.asset_id) uploadedIds.push(asset.asset_id); + } + }); + try { + if (mode === "fl2va") { + await expect(page.getByRole("button", { name: "Generate", exact: true })).toBeDisabled(); + await page.locator('input[type="file"]').setInputFiles([ + { name: "first-frame.png", mimeType: "image/png", buffer: await readFile(imagePath) }, + { name: "last-frame.png", mimeType: "image/png", buffer: await readFile(imagePath) }, + ]); + await expect(page.getByRole("option", { name: "first-frame.png", exact: true }).first()).toBeAttached(); + await page.getByRole("combobox", { name: "First frame", exact: true }).selectOption({ label: "first-frame.png" }); + await expect(page.getByRole("button", { name: "Generate", exact: true })).toBeEnabled(); + await page.getByRole("combobox", { name: "Last frame", exact: true }).selectOption({ label: "last-frame.png" }); + } + if (mode === "ref2va") { + const video = execFileSync(process.env.FASTVIDEO_FFMPEG_BIN || "ffmpeg", ["-v", "error", "-f", "lavfi", "-i", "color=c=royalblue:s=64x64:r=8", "-t", "1", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "frag_keyframe+empty_moov", "-f", "mp4", "pipe:1"]); + await page.locator('input[type="file"]').setInputFiles([ + { name: "subject.png", mimeType: "image/png", buffer: await readFile(imagePath) }, + { name: "motion.mp4", mimeType: "video/mp4", buffer: video }, + { name: "sound.wav", mimeType: "audio/wav", buffer: makeAudio() }, + ]); + await expect(page.getByRole("button", { name: "Add sound.wav as reference" })).toBeEnabled(); + await page.getByRole("button", { name: "Add sound.wav as reference" }).click(); + await expect(page.getByRole("button", { name: "Generate", exact: true })).toBeDisabled(); + await page.getByRole("button", { name: "Add subject.png as reference" }).click(); + await page.getByRole("button", { name: "Add motion.mp4 as reference" }).click(); + await page.getByRole("button", { name: "Move sound.wav down" }).click(); + const names = await page.getByRole("list", { name: "Ordered references" }).locator("li p.font-medium").allTextContents(); + expect(names).toEqual(["subject.png", "sound.wav", "motion.mp4"]); + } + await page.screenshot({ path: testInfo.outputPath(`${mode}-inputs.png`), fullPage: true }); + await page.getByRole("button", { name: "Generate", exact: true }).click(); + await expect.poll(() => sent.find((item) => item.type === "session_init_v2")?.generation_mode).toBe(mode); + const init = sent.find((item) => item.type === "session_init_v2")!; + expect(init.conditioning_assets.map((item: any) => item.role)).toEqual(mode === "t2va" ? [] : mode === "fl2va" ? ["first_frame", "last_frame"] : ["reference", "reference", "reference"]); + if (mode === "ref2va") expect(init.conditioning_assets.map((item: any) => item.asset_id)).toEqual([uploadedIds[0], uploadedIds[2], uploadedIds[1]]); + await expect.poll(() => received.find((item) => item.type === "gpu_assigned")?.generation_mode).toBe(mode); + await expect.poll(() => received.some((item) => item.type === "media_segment_complete")).toBe(true); + await expect(page.getByText(/Demo runtime · Sample playback only/)).toBeVisible(); + await expect(modeSelect).toHaveCount(0); + await expect.poll(async () => page.locator("video:visible").first().evaluate((element: HTMLVideoElement) => element.readyState)).toBeGreaterThanOrEqual(2); + await page.screenshot({ path: testInfo.outputPath(`${mode}-playback.png`), fullPage: true }); + if (mode === "ref2va") { + await page.getByRole("button", { name: "Toggle sidebar" }).click(); + await page.getByRole("button", { name: "New project", exact: true }).click(); + await expect(modeSelect).toHaveText("T2VA"); + await modeSelect.click(); + await page.getByRole("option", { name: "FL2VA", exact: true }).click(); + await expect(modeSelect).toHaveText("FL2VA"); + await page.getByRole("combobox", { name: "First frame", exact: true }).selectOption({ label: "subject.png" }); + await expect(page.getByRole("combobox", { name: "Last frame", exact: true })).toHaveValue(""); + await page.getByLabel("Continuation prompt").fill("The paper fox explores a new scene."); + await page.getByRole("button", { name: "Generate", exact: true }).click(); + await expect.poll(() => sent.find((item) => item.type === "project_init_v1")?.generation_mode).toBe("fl2va"); + const secondProject = sent.find((item) => item.type === "project_init_v1")!; + expect(secondProject.conditioning_assets).toEqual([{ asset_id: uploadedIds[0], role: "first_frame" }]); + expect(sent.filter((item) => item.type === "session_init_v2")).toHaveLength(1); + await expect.poll(() => received.filter((item) => item.type === "media_segment_complete").length).toBeGreaterThan(1); + } + } finally { + await page.close(); + for (const id of uploadedIds) await request.delete(`/assets/${id}`); + } + }); + } + + test("proxies a media upload larger than Next's default 10 MiB body limit", async ({ request }) => { + const response = await request.get("/generation-capabilities"); + const capabilities = response.ok() ? await response.json() : {}; + test.skip(capabilities.mock !== true, "Requires the local mock runtime."); + const audio = makeAudio(192000, 29); + expect(audio.length).toBeGreaterThan(10 * 1024 * 1024); + const upload = await request.post("/assets", { + headers: { "Content-Type": "audio/wav", "X-Asset-Name": "large-proxy-check.wav" }, + data: audio, + }); + expect(upload.status()).toBe(201); + const asset = await upload.json(); + try { expect(asset.size).toBe(audio.length); } finally { await request.delete(`/assets/${asset.asset_id}`); } + }); +}); diff --git a/apps/dreamverse/web/next.config.ts b/apps/dreamverse/web/next.config.ts index 6a7e42e3f..06d399f7a 100644 --- a/apps/dreamverse/web/next.config.ts +++ b/apps/dreamverse/web/next.config.ts @@ -9,6 +9,8 @@ const configDir = path.dirname(fileURLToPath(import.meta.url)); const staticExport = process.env.NEXT_OUTPUT_EXPORT === '1'; const nextConfig: NextConfig = { + // Next 15.5 name for the dev rewrite-proxy body limit; Next 16 renames it to `proxyClientMaxBodySize`. + experimental: { middlewareClientMaxBodySize: 100 * 1024 * 1024 }, ...(staticExport ? { output: 'export' as const } : {}), ...(staticExport ? { images: { unoptimized: true } } : {}), outputFileTracingRoot: path.join(configDir, '..', '..', '..'), @@ -38,6 +40,18 @@ const nextConfig: NextConfig = { source: '/router/:path*', destination: `${backendUrl}/router/:path*` }, + { + source: '/generation-capabilities', + destination: `${backendUrl}/generation-capabilities`, + }, + { + source: '/assets', + destination: `${backendUrl}/assets`, + }, + { + source: '/assets/:path*', + destination: `${backendUrl}/assets/:path*`, + }, { source: '/prompt-system-config', destination: `${backendUrl}/prompt-system-config`, diff --git a/apps/dreamverse/web/src/app/page.integration.test.tsx b/apps/dreamverse/web/src/app/page.integration.test.tsx index 6a0b0440f..8bc3555a9 100644 --- a/apps/dreamverse/web/src/app/page.integration.test.tsx +++ b/apps/dreamverse/web/src/app/page.integration.test.tsx @@ -589,6 +589,7 @@ describe.skip('App websocket integration', () => { }); const initMessage = outbound.find((message) => message.type === 'session_init_v2'); + expect(initMessage.generation_mode).toBe('t2va'); expect(initMessage.preset_id).toBe('test_preset'); expect(initMessage.curated_prompts).toEqual(['segment one', 'segment two']); expect(initMessage.enhancement_enabled).toBe(true); @@ -597,6 +598,42 @@ describe.skip('App websocket integration', () => { expect(initMessage.initial_rollout_prompt).toBe(''); }); + it('sends the selected generation mode and locks it after session start', async () => { + const outbound: any[] = []; + server.on('connection', (socket) => { + socket.on('message', (rawMessage) => { + outbound.push(JSON.parse(rawMessage as string)); + }); + }); + + const user = userEvent.setup(); + render(); + + const modeSelect = await screen.findByRole('combobox', { name: 'Generation mode' }); + expect(modeSelect).toHaveTextContent('T2VA'); + + await user.click(modeSelect); + await user.click(await screen.findByRole('option', { name: 'FL2VA' })); + expect(modeSelect).toHaveTextContent('FL2VA'); + expect(modeSelect).toHaveAttribute( + 'title', + 'First/last frames to video + audio. Start from a first frame image. Add an optional last frame to guide the ending.', + ); + + const generateButton = await screen.findByRole('button', { name: 'Generate' }); + await waitFor(() => expect(generateButton).toBeEnabled()); + await user.click(generateButton); + + await waitFor(() => { + expect(outbound.some((message) => message.type === 'session_init_v2')).toBe(true); + }); + + const initMessage = outbound.find((message) => message.type === 'session_init_v2'); + expect(initMessage.generation_mode).toBe('fl2va'); + expect(screen.queryByRole('combobox', { name: 'Generation mode' })) + .not.toBeInTheDocument(); + }); + it('starts a streaming session from a custom initial prompt without using curated prompts', async () => { const outbound: any[] = []; server.on('connection', (socket) => { diff --git a/apps/dreamverse/web/src/app/page.tsx b/apps/dreamverse/web/src/app/page.tsx index 0af6176f0..8a472d503 100644 --- a/apps/dreamverse/web/src/app/page.tsx +++ b/apps/dreamverse/web/src/app/page.tsx @@ -5,6 +5,7 @@ import { Download, Share2 } from "lucide-react"; import DevtoolsShell from "@/components/devtools/DevtoolsShell"; import MonitorPage from "@/components/MonitorPage"; import ChatBar from "@/components/ChatBar"; +import AssetList from "@/components/AssetList"; import SessionTimeoutModal from "@/components/SessionTimeoutModal"; import Sidebar from "@/components/Sidebar"; import Header from "@/components/Header"; @@ -13,10 +14,13 @@ import Workspace from "@/components/Workspace"; import { saveProject, saveProjectMetadata, listProjects, loadProjectClips, deleteProject, pruneOldProjects, type StoredProject, type StoredClip } from "@/lib/projectStorage"; import { isInfrastructureError } from "@/lib/ws/reducer"; import { useStore } from "@/hooks/useStore"; +import { useAssetLibrary } from "@/hooks/useAssetLibrary"; +import { useGenerationCapabilities } from "@/hooks/useGenerationCapabilities"; import { resolveDevtoolsMode } from "@/lib/devtoolsMode"; import { createAvPipeline, DEFAULT_AV_MIME } from "@/lib/media/avPipeline"; import { remuxArchivedFmp4Segments } from "@/lib/media/fmp4Remux"; import { DEFAULT_CUSTOM_PRESET_ID, parseStoryPresets, sanitizePresetId } from "@/lib/presets"; +import { DEFAULT_GENERATION_MODE, buildGenerationInitFields, validateGenerationInputs, type GenerationMode, type GenerationInitFields, type GenerationAsset } from "@/lib/generationMode"; import { buildRewritePromptWindowSnapshot, buildRewritePromptWindowSnapshotFromPrompts, @@ -341,6 +345,19 @@ export default function Page() { const [isMobileShareCapable, setIsMobileShareCapable] = useState(false); const [videoMuted, setVideoMuted] = useState(true); const [timeoutModalOpen, setTimeoutModalOpen] = useState(false); + const [generationMode, setGenerationMode] = useState(DEFAULT_GENERATION_MODE); + const assetLibrary = useAssetLibrary(); + const { capabilities, capabilityNotice, refreshCapabilities } = useGenerationCapabilities(); + const joiningRef = useRef(false); + const activeGenerationRef = useRef<{ fields: GenerationInitFields; assets: GenerationAsset[]; mock: boolean } | null>(null); + const generationInputError = validateGenerationInputs(generationMode, assetLibrary.conditioningAssets, assetLibrary.assets); + const generationSupported = capabilities.modes.includes(generationMode); + const generationInputsValid = !generationInputError && generationSupported && !assetLibrary.uploading; + function changeGenerationMode(mode: GenerationMode) { + if (sessionStore.get().sessionStarted || joiningRef.current || !capabilities.modes.includes(mode)) return; + setGenerationMode(mode); + assetLibrary.clearConditioning(); + } useEffect(() => { setIsMobileShareCapable(typeof navigator.canShare === "function" && window.matchMedia("(pointer: coarse)").matches); }, []); @@ -391,7 +408,7 @@ export default function Page() { // --- Derived values --- - const canStartSession = !projectResetPending && (canJoinSession || Boolean(normalizeInitialPrompt(livePromptDraft as string))); + const canStartSession = generationInputsValid && !projectResetPending && (canJoinSession || Boolean(normalizeInitialPrompt(livePromptDraft as string))); const currentClipLabel = useMemo(() => { if ((activeClip as Record)?.label) return (activeClip as Record).label; @@ -747,7 +764,11 @@ export default function Page() { function recoverFailedSessionStart(notice: string) { const restoredDraft = normalizeInitialPrompt(pendingInitialPromptRef.current); - resetToLobbyState(); + if (wsRef.current) { + detachAndCloseWebSocket(wsRef.current); + wsRef.current = null; + } + resetToLobbyState({ preserveSessionNotice: true }); clearPendingProjectPointers(); pendingInitialPromptRef.current = ""; sessionStore.patch({ @@ -1702,6 +1723,10 @@ export default function Page() { function resetToLobbyState({ preserveSessionNotice = false, preservePlayback = false } = {}) { setVideoMuted(true); + if (!preserveSessionNotice) { + setGenerationMode(DEFAULT_GENERATION_MODE); + assetLibrary.clearConditioning(); + } clearCountdownInterval(); pendingInitialPromptRef.current = ""; sessionStore.patch({ @@ -1736,6 +1761,8 @@ export default function Page() { function resetToProjectLobbyState() { setVideoMuted(true); + setGenerationMode(DEFAULT_GENERATION_MODE); + assetLibrary.clearConditioning(); pendingInitialPromptRef.current = ""; sessionStore.patch({ sessionStarted: false, @@ -1764,6 +1791,7 @@ export default function Page() { setSeedPrompts(segmentPrompts); return { type, + ...(activeGenerationRef.current?.fields || buildGenerationInitFields(generationMode, assetLibrary.conditioningAssets, assetLibrary.assets)), preset_id: getInitialPresetId(), preset_label: getInitialPresetLabel(), curated_prompts: segmentPrompts, @@ -1812,6 +1840,12 @@ export default function Page() { return; } if (decoded.kind !== "json") return; + if (decoded.data?.type === "error" && sessionStore.get().sessionStarted + && (!sessionStore.get().gpuAssigned || decoded.data.error_code === "invalid_generation_input")) { + const message = typeof decoded.data.message === "string" ? decoded.data.message : "The generation inputs were rejected. Check the mode and selected assets."; + recoverFailedSessionStart(message); + return; + } if (decoded.data?.type === "error" && isInfrastructureError(decoded.data)) { const message = typeof decoded.data?.message === "string" && decoded.data.message.trim() ? decoded.data.message.trim() @@ -1937,9 +1971,15 @@ export default function Page() { } } - function beginProjectLocally({ force = false } = {}) { + function beginProjectLocally({ force = false, mockRuntime = capabilities.mock === true } = {}) { if (!force && !canStartSession) return; + if (!generationInputsValid) return false; if (sessionStore.get().sessionStarted || sessionStore.get().projectResetPending) return false; + activeGenerationRef.current = { + fields: buildGenerationInitFields(generationMode, assetLibrary.conditioningAssets, assetLibrary.assets), + assets: assetLibrary.assets.filter((asset) => assetLibrary.conditioningAssets.some((item) => item.asset_id === asset.asset_id)), + mock: mockRuntime, + }; setTimeoutModalOpen(false); // Unmute during the user gesture so iOS Safari permits audio playback. setVideoMuted(false); @@ -1999,12 +2039,43 @@ export default function Page() { } async function joinSession({ force = false } = {}) { + if (joiningRef.current || sessionStore.get().sessionStarted) return; + if (generationInputError || assetLibrary.uploading) { + showPreSessionNotice(generationInputError || "Wait for the asset upload to finish."); + return; + } + joiningRef.current = true; + try { + await startGenerationSession({ force }); + } finally { + joiningRef.current = false; + } + } + + async function startGenerationSession({ force = false } = {}) { + sessionStore.patch({ sessionNotice: "" }); + streamStore.patch({ loadingAnimation: true }); + const currentCapabilities = await refreshCapabilities(); + if (!currentCapabilities.modes.includes(generationMode)) { + streamStore.patch({ loadingAnimation: false }); + showPreSessionNotice(`${generationMode.toUpperCase()} is unavailable on this runtime. Connect a full H3 runtime or choose a supported mode.`); + return; + } + const assetProblem = await assetLibrary.verifySelectedAssets(); + if (assetProblem) { + streamStore.patch({ loadingAnimation: false }); + showPreSessionNotice(assetProblem); + return; + } if ( wsRef.current && wsRef.current.readyState === WebSocket.OPEN && sessionStore.get().connected ) { - if (!beginProjectLocally({ force })) return; + if (!beginProjectLocally({ force, mockRuntime: currentCapabilities.mock === true })) { + streamStore.patch({ loadingAnimation: false }); + return; + } sendProjectInitMessage(); return; } @@ -2017,7 +2088,7 @@ export default function Page() { showPreSessionNotice(probe.notice); return; } - if (!beginProjectLocally({ force })) { + if (!beginProjectLocally({ force, mockRuntime: currentCapabilities.mock === true })) { streamStore.patch({ loadingAnimation: false }); sessionStore.patch({ connecting: false }); return; @@ -2044,6 +2115,10 @@ export default function Page() { createdAt: currentProjectCreatedAtRef.current || Date.now(), lastThumbnail: currentThumbnail, promptEvents: [...(rewriteStore.get().promptEvents as Record[])], + generationMode: activeGenerationRef.current?.fields.generation_mode || DEFAULT_GENERATION_MODE, + conditioningAssets: activeGenerationRef.current?.fields.conditioning_assets || [], + assets: activeGenerationRef.current?.assets || [], + mock: activeGenerationRef.current?.mock === true, }; const clips: StoredClip[] = (streamStore.get().completedClips as any[]) .filter((clip: any) => clip?.blob instanceof Blob) @@ -2478,6 +2553,24 @@ export default function Page() { // --- Render --- + const conditioningPanel = generationMode !== "t2va" && !sessionStarted && !sessionExpired ? ( + + ) : null; + if (!runtimeReady) { return null; } @@ -2501,13 +2594,17 @@ export default function Page() { enhancementEnabled={enhancementEnabled as boolean} autoExtensionEnabled={autoExtensionEnabled as boolean} loopGenerationEnabled={loopGenerationEnabled as boolean} - canJoinSession={canJoinSession as boolean} + canJoinSession={canStartSession} canSubmitContinuation={canSubmitContinuation} editableMode={editableMode as boolean} demoMode={demoMode as boolean} editableCanJoin={editableCanJoin as boolean} curatedPromptLimit={curatedPromptLimit as number} maxCuratedPromptCount={maxCuratedPromptCount as number} + generationMode={generationMode} + supportedGenerationModes={capabilities.modes} + conditioningPanel={conditioningPanel} + onGenerationModeChange={changeGenerationMode} onPresetChange={handlePresetSelectionChange} onEnhancementToggle={handleEnhancementToggle} onCuratedPromptLimitChange={handleCuratedPromptLimitChange} @@ -2641,7 +2738,12 @@ export default function Page() { />
    setSidebarOpen((prev) => !prev)} /> -
    +
    {isViewingMode && ( <> {viewingSelectedClip && ( @@ -2680,6 +2782,8 @@ export default function Page() { /> + {viewingProject?.project.mock &&

    Demo sample · This saved clip was not generated by an AI model.

    } + {viewingProject?.project.generationMode &&

    {viewingProject.project.generationMode.toUpperCase()} · {viewingProject.project.conditioningAssets?.length || 0} saved references. Uploaded originals may expire; your saved video remains available.

    }     @@ -2759,7 +2863,7 @@ export default function Page() { - {!showActiveProject && ( + {!showActiveProject && generationMode === "t2va" && ( diff --git a/apps/dreamverse/web/src/components/AssetList.test.tsx b/apps/dreamverse/web/src/components/AssetList.test.tsx new file mode 100644 index 000000000..1db71480a --- /dev/null +++ b/apps/dreamverse/web/src/components/AssetList.test.tsx @@ -0,0 +1,41 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import AssetList from "./AssetList"; +import type { GenerationAsset } from "@/lib/generationMode"; + +const frame: GenerationAsset = { asset_id: "first", kind: "image", name: "frame.png", mime_type: "image/png", size: 2000, url: "/assets/first" }; +const video: GenerationAsset = { asset_id: "video", kind: "video", name: "motion.mp4", mime_type: "video/mp4", size: 2000, url: "/assets/video" }; +function props() { + return { assets: [frame, video], onUpload: vi.fn(), onAssign: vi.fn(), onRemove: vi.fn(), onUnselect: vi.fn(), onMove: vi.fn(), onMissing: vi.fn() }; +} + +describe("Asset List", () => { + it("uploads to the library and assigns images to endpoint roles", async () => { + const callbacks = props(); + const user = userEvent.setup(); + render(); + await user.selectOptions(screen.getByRole("combobox", { name: "First frame" }), "first"); + expect(callbacks.onAssign).toHaveBeenCalledWith("first", "first_frame"); + expect(screen.getByRole("combobox", { name: "Last frame" })).toHaveValue(""); + expect(screen.queryByRole("option", { name: "motion.mp4" })).not.toBeInTheDocument(); + const file = new File(["image"], "new.png", { type: "image/png" }); + await user.upload(screen.getByLabelText("Upload assets", { selector: "input" }), file); + expect(callbacks.onUpload).toHaveBeenCalledWith([file]); + }); + it("exposes accessible ordering and removal controls for multimodal references", async () => { + const callbacks = props(); + const user = userEvent.setup(); + render(); + await user.click(screen.getByRole("button", { name: "Move motion.mp4 up" })); + expect(callbacks.onMove).toHaveBeenCalledWith(1, 0); + await user.click(screen.getByRole("button", { name: "Unselect frame.png" })); + expect(callbacks.onUnselect).toHaveBeenCalledWith(0); + expect(screen.getByRole("button", { name: "Move frame.png up" })).toBeDisabled(); + }); + it("locks uploads and assignments while starting generation", () => { + render(); + expect(screen.getByRole("button", { name: "Upload assets" })).toBeDisabled(); + expect(screen.getByRole("combobox", { name: "First frame" })).toBeDisabled(); + }); +}); diff --git a/apps/dreamverse/web/src/components/AssetList.tsx b/apps/dreamverse/web/src/components/AssetList.tsx new file mode 100644 index 000000000..b27d2b5ba --- /dev/null +++ b/apps/dreamverse/web/src/components/AssetList.tsx @@ -0,0 +1,154 @@ +"use client"; + +import { useRef, useState } from "react"; +import { ArrowDown, ArrowUp, AudioLines, Check, GripVertical, ImagePlus, Plus, Trash2, Upload, X } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { NativeSelect } from "@/components/ui/native-select"; +import { cn } from "@/lib/utils"; +import type { ConditioningAsset, ConditioningRole, GenerationAsset, GenerationMode } from "@/lib/generationMode"; + +interface AssetListProps { + mode: GenerationMode; + assets: GenerationAsset[]; + conditioning: ConditioningAsset[]; + locked?: boolean; + uploading?: boolean; + error?: string; + validationNotice?: string | null; + onUpload: (files: File[]) => void; + onAssign: (assetId: string, role: ConditioningRole) => void; + onRemove: (assetId: string) => void; + onUnselect: (index: number) => void; + onMove: (from: number, to: number) => void; + onMissing: (assetId: string) => void; +} + +function AssetPreview({ asset, onMissing, compact = false }: { + asset: GenerationAsset; + onMissing: (assetId: string) => void; + compact?: boolean; +}) { + const [previewFailed, setPreviewFailed] = useState(false); + function previewError() { + setPreviewFailed(true); + onMissing(asset.asset_id); + } + const className = cn("h-full w-full object-cover", asset.missing && "opacity-25"); + if (asset.missing) return Upload again; + if (previewFailed) return Preview unavailable; + if (asset.kind === "image") { + return {asset.name}; + } + if (asset.kind === "video") { + return