From 4e52f47d1ef5a8ecb1d33f6d3c45f3fad03d5853 Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:08:58 -0700 Subject: [PATCH] [feat] add MXFP8 support on H3 (#1796) --- .../basic/basic_fasth3_simplified.py | 147 ++++++++++ .../basic/basic_fasth3_simplified.sh | 10 + .../basic/basic_fasth3_simplified_profile.py | 126 ++++++++ .../basic/basic_fasth3_simplified_profile.sh | 53 ++++ fastvideo/layers/mxfp8linear.py | 276 ++++++++++++++++++ fastvideo/layers/quantization/__init__.py | 13 +- fastvideo/layers/quantization/mxfp8_config.py | 147 ++++++++++ fastvideo/layers/quantization/nvfp4_config.py | 31 +- fastvideo/models/dits/minimax_h3.py | 72 +++-- fastvideo/models/loader/fsdp_load.py | 50 +++- fastvideo/pipelines/lora_pipeline.py | 67 +++++ .../test_fsdp_inference_quantization.py | 52 ++++ .../tests/ops/quantization/test_mxfp8.py | 120 ++++++++ .../test_nvfp4_minimax_h3_wiring.py | 234 +++++++++++++++ pyproject.toml | 1 + 15 files changed, 1350 insertions(+), 49 deletions(-) create mode 100644 examples/inference/basic/basic_fasth3_simplified.py create mode 100644 examples/inference/basic/basic_fasth3_simplified.sh create mode 100644 examples/inference/basic/basic_fasth3_simplified_profile.py create mode 100755 examples/inference/basic/basic_fasth3_simplified_profile.sh create mode 100644 fastvideo/layers/mxfp8linear.py create mode 100644 fastvideo/layers/quantization/mxfp8_config.py create mode 100644 fastvideo/tests/loader/test_fsdp_inference_quantization.py create mode 100644 fastvideo/tests/ops/quantization/test_mxfp8.py create mode 100644 fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py diff --git a/examples/inference/basic/basic_fasth3_simplified.py b/examples/inference/basic/basic_fasth3_simplified.py new file mode 100644 index 000000000..b53bd655c --- /dev/null +++ b/examples/inference/basic/basic_fasth3_simplified.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate 5-second FastH3 videos with a supported Preview adapter.""" + +from __future__ import annotations + +import argparse +import os +import statistics +import sys +import time +from pathlib import Path + +from huggingface_hub import hf_hub_download + +FASTVIDEO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(FASTVIDEO_ROOT)) + +from fastvideo import VideoGenerator # noqa: E402 +from fastvideo.api import ( # noqa: E402 + CompileConfig, + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + SamplingConfig, +) + +# FastVideo exposes these three backend switches only through environment variables. +os.environ.update({ + "FASTVIDEO_FA4": "1", + "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", + "FASTVIDEO_VSA_SM100A": "0", +}) +os.environ.pop("FASTVIDEO_INFERENCE_TORCH_COMPILE", None) + +VARIANT_BACKENDS = { + "dense-datafree": "FLASH_ATTN", + "vsa-datafree": "VIDEO_SPARSE_ATTN_H3", + "vsa-synthetic-step1300": "VIDEO_SPARSE_ATTN_H3", + "vsa-synthetic-step1900": "VIDEO_SPARSE_ATTN_H3", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("variant", choices=VARIANT_BACKENDS) + parser.add_argument("--prompt", required=True) + return parser.parse_args() + + +def main() -> None: + """Run one compile warmup and three measured generations with one fixed recipe.""" + args = parse_args() + + attention_backend = VARIANT_BACKENDS[args.variant] + experimental = { + "attention_backend": attention_backend, + "inference_torch_compile": attention_backend == "FLASH_ATTN", + "vae_parallel_decode": True, + "vae_parallel_decode_strategy": "gather", + } + if attention_backend == "VIDEO_SPARSE_ATTN_H3": + experimental.update({ + "VSA_sparsity": 0.9, + "VSA_tile_size": 64, + }) + + adapter_path = hf_hub_download( + repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + filename=f"{args.variant}/adapter_model.safetensors", + ) + output_dir = FASTVIDEO_ROOT / "outputs/fasth3_lora_preview" / args.variant + output_dir.mkdir(parents=True, exist_ok=True) + + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path="MiniMaxAI/MiniMax-H3", + pipeline=PipelineSelection( + components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0), + experimental=experimental, + ), + engine=EngineConfig( + num_gpus=4, + parallelism=ParallelismConfig(tp_size=1, sp_size=4), + offload=OffloadConfig(dit=False, dit_layerwise=False), + compile=CompileConfig(vae_enabled=True), + ), + ) + ) + + generation_count = 4 + warmup_count = 1 + measured_count = generation_count - warmup_count + measured_seconds: list[float] = [] + + print(f"Variant: {args.variant} ({attention_backend})") + print(f"Output directory: {output_dir}") + try: + for generation_index in range(generation_count): + measured = generation_index >= warmup_count + if measured: + measured_index = generation_index - warmup_count + 1 + label = f"measured {measured_index}/{measured_count}" + output_path = output_dir / f"fasth3_all_run_{measured_index:02d}.mp4" + else: + warmup_index = generation_index + 1 + label = f"warmup {warmup_index}/{warmup_count}" + output_path = output_dir / f"_fasth3_warmup_{warmup_index:02d}.mp4" + + request = GenerationRequest( + prompt=args.prompt, + negative_prompt="", + sampling=SamplingConfig( + height=768, + width=1344, + num_frames=345, # <-- Change video length: 5sec: 124; 10sec: 243; 15sec: 345. + fps=24, + num_inference_steps=5, + guidance_scale=1.0, + batch_cfg=False, + seed=1000, + ), + output=OutputConfig( + output_path=str(output_path), + save_video=True, + return_frames=False, + ), + ) + started = time.perf_counter() + result = generator.generate(request) + elapsed = time.perf_counter() - started + if measured: + measured_seconds.append(elapsed) + suffix = "" if measured else " (excluded from median)" + print(f"[{label}] {result.video_path or output_path}: {elapsed:.3f}s{suffix}") + finally: + generator.shutdown() + + print(f"Median E2E wall time: {statistics.median(measured_seconds):.3f}s") + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_fasth3_simplified.sh b/examples/inference/basic/basic_fasth3_simplified.sh new file mode 100644 index 000000000..625530e3a --- /dev/null +++ b/examples/inference/basic/basic_fasth3_simplified.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash + +PROMPT="integrated_multimodal_description: A red fox runs through fresh snow at dawn. overall_soundscape: Fast pawsteps in snow, winter wind, and distant birds." + +# Options: dense-datafree, vsa-datafree, vsa-synthetic-step1300, vsa-synthetic-step1900 +VARIANT="vsa-datafree" + +exec python "$(dirname "$0")/basic_fasth3_simplified.py" \ + "$VARIANT" \ + --prompt "$PROMPT" diff --git a/examples/inference/basic/basic_fasth3_simplified_profile.py b/examples/inference/basic/basic_fasth3_simplified_profile.py new file mode 100644 index 000000000..c33b9de23 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_simplified_profile.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Profile one FastH3 Preview generation with NVTX after warmup.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch +from huggingface_hub import hf_hub_download + +from fastvideo import VideoGenerator +from fastvideo.api import ( + CompileConfig, + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + ParallelismConfig, + PipelineSelection, + QuantizationConfig, + SamplingConfig, +) +from fastvideo.profiler import nvtx_range + +VARIANT_BACKENDS = { + "dense-datafree": "FLASH_ATTN", + "vsa-datafree": "VIDEO_SPARSE_ATTN_H3", + "vsa-synthetic-step1300": "VIDEO_SPARSE_ATTN_H3", + "vsa-synthetic-step1900": "VIDEO_SPARSE_ATTN_H3", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("variant", choices=VARIANT_BACKENDS) + parser.add_argument("--prompt", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--num-frames", type=int, default=345) + parser.add_argument("--warmup-runs", type=int, default=3) + return parser.parse_args() + + +def main() -> None: + """Warm the selected FastH3 recipe, then profile one identical generation.""" + args = parse_args() + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + attention_backend = VARIANT_BACKENDS[args.variant] + experimental = { + "attention_backend": attention_backend, + "inference_torch_compile": attention_backend == "FLASH_ATTN", + "vae_parallel_decode": True, + "vae_parallel_decode_strategy": "gather", + } + if attention_backend == "VIDEO_SPARSE_ATTN_H3": + experimental.update({ + "VSA_sparsity": 0.9, + "VSA_tile_size": 64, + }) + + adapter_path = hf_hub_download( + repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + filename=f"{args.variant}/adapter_model.safetensors", + ) + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path="MiniMaxAI/MiniMax-H3", + pipeline=PipelineSelection( + components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0), + experimental=experimental, + ), + engine=EngineConfig( + num_gpus=4, + parallelism=ParallelismConfig(tp_size=1, sp_size=4), + offload=OffloadConfig(dit=False, dit_layerwise=False), + compile=CompileConfig(vae_enabled=True), + quantization=QuantizationConfig(transformer_quant="MXFP8"), + ), + ) + ) + + request = GenerationRequest( + prompt=args.prompt, + negative_prompt="", + sampling=SamplingConfig( + height=768, + width=1344, + num_frames=args.num_frames, + fps=24, + num_inference_steps=5, + guidance_scale=1.0, + batch_cfg=False, + seed=1000, + ), + output=OutputConfig( + output_path=str(output_dir / "fasth3_profile.mp4"), + save_video=True, + return_frames=False, + ), + ) + + try: + for warmup_index in range(args.warmup_runs): + warmup_result = generator.generate(request) + print(f"Warmup {warmup_index + 1}/{args.warmup_runs}: {warmup_result.video_path}") + + torch.cuda.profiler.start() + try: + with nvtx_range("fasth3.profiled_generation"): + measured_result = generator.generate(request) + finally: + torch.cuda.profiler.stop() + + print(f"Profiled output: {measured_result.video_path}") + if measured_result.generation_time is not None: + print(f"Generation time: {measured_result.generation_time:.2f}s") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_fasth3_simplified_profile.sh b/examples/inference/basic/basic_fasth3_simplified_profile.sh new file mode 100755 index 000000000..c20981cd3 --- /dev/null +++ b/examples/inference/basic/basic_fasth3_simplified_profile.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# Capture one FastH3 Preview generation with CUDA and FastVideo NVTX ranges. + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +WORKTREE_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." && pwd)" +WORKSPACE_ROOT="$(cd -- "${WORKTREE_ROOT}/.." && pwd)" +PYTHON_BIN="${WORKSPACE_ROOT}/.venv-fv/bin/python" +NSYS_BIN=/usr/local/cuda/bin/nsys +PROFILE_SCRIPT="${SCRIPT_DIR}/basic_fasth3_simplified_profile.py" + +export CUDA_VISIBLE_DEVICES=0,1,2,3 +export FASTVIDEO_FA4=1 +export FASTVIDEO_INFERENCE_TORCH_COMPILE=0 +export FASTVIDEO_MINIMAX_H3_FUSIONS=all +export FASTVIDEO_NVTX_PROFILE=1 +export FASTVIDEO_VSA_SM100A=0 +export PYTHONUNBUFFERED=1 +export PYTHONPATH="${WORKTREE_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +PROMPT="integrated_multimodal_description: A red fox runs through fresh snow at dawn. overall_soundscape: Fast pawsteps in snow, winter wind, and distant birds." + +# Options: dense-datafree, vsa-datafree, vsa-synthetic-step1300, vsa-synthetic-step1900 +VARIANT="vsa-datafree" +NUM_FRAMES=345 +WARMUP_RUNS=3 + +PROFILE_ID="${VARIANT}_4gpu_sp4_${NUM_FRAMES}frames" +RUN_ID="$(date -u +%Y%m%dT%H%M%SZ)-$$" +RESULT_DIR="${WORKTREE_ROOT}/runs/fasth3_lora_preview_profile/${PROFILE_ID}/${RUN_ID}" +MEDIA_DIR="${RESULT_DIR}/media" +RUN_LOG="${RESULT_DIR}/run.log" +mkdir -p "${MEDIA_DIR}" + +"${NSYS_BIN}" profile \ + --trace=cuda,nvtx \ + --sample=none \ + --cpuctxsw=none \ + --capture-range=cudaProfilerApi \ + --capture-range-end=stop \ + --output="${RESULT_DIR}/${PROFILE_ID}" \ + "${PYTHON_BIN}" "${PROFILE_SCRIPT}" \ + "${VARIANT}" \ + --num-frames "${NUM_FRAMES}" \ + --warmup-runs "${WARMUP_RUNS}" \ + --prompt "${PROMPT}" \ + --output "${MEDIA_DIR}" \ + 2>&1 | tee "${RUN_LOG}" + +printf 'RESULT_DIR=%s\n' "${RESULT_DIR}" +printf 'NSYS_REPORT=%s\n' "${RESULT_DIR}/${PROFILE_ID}.nsys-rep" +printf 'MEDIA_DIR=%s\n' "${MEDIA_DIR}" diff --git a/fastvideo/layers/mxfp8linear.py b/fastvideo/layers/mxfp8linear.py new file mode 100644 index 000000000..6892161c0 --- /dev/null +++ b/fastvideo/layers/mxfp8linear.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 block quantization and linear operations for Blackwell inference.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from quack.mx_utils import to_blocked, to_mx + +MXFP8_BLOCK_SIZE = 32 +_BLOCKS_PER_PROGRAM = 16 + + +def _validate_mxfp8_matrix(matrix: torch.Tensor) -> None: + """Validate one row-major matrix before MXFP8 quantization.""" + if matrix.ndim != 2: + raise ValueError(f"MXFP8 quantization requires a 2D tensor, got shape {tuple(matrix.shape)}.") + if matrix.dtype not in (torch.bfloat16, torch.float32): + raise TypeError(f"MXFP8 quantization requires BF16 or FP32 input, got {matrix.dtype}.") + if matrix.shape[1] % MXFP8_BLOCK_SIZE: + raise ValueError(f"MXFP8 reduction dimension must be divisible by {MXFP8_BLOCK_SIZE}, got {matrix.shape[1]}.") + + +@torch.compile(dynamic=True, fullgraph=True) +def _quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + values, natural_scales = to_mx(matrix.contiguous(), MXFP8_BLOCK_SIZE) + return values, to_blocked(natural_scales) + + +def quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Prequantize a weight with Quack and return hardware-blocked scales.""" + _validate_mxfp8_matrix(matrix) + return _quantize_mxfp8_weight_blockwise(matrix) + + +@triton.jit +def _store_mxfp8_blocks( + matrix, + quantized_ptr, + scale_ptr, + row, + first_block, + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM: tl.constexpr, +): + """Quantize row blocks and store scales in cuBLAS's 32x4x4 layout.""" + element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32) + columns = first_block * 32 + element_offsets + valid_elements = (row < row_count) & (columns < column_count) + matrix = tl.reshape(matrix, (BLOCKS_PER_PROGRAM, 32)) + + # FLOOR scaling encodes exponent(max_abs) - 8 as one E8M0 byte per 32 values. + max_abs = tl.max(tl.abs(matrix), axis=1) + exponent = (max_abs.to(tl.int32, bitcast=True) >> 23) & 0xFF + scale_biased = tl.maximum(tl.minimum(exponent - 8, 255), 0) + scale_biased = tl.where(max_abs != max_abs, 255, scale_biased) + scale_bits = scale_biased.to(tl.int32) << 23 + scale = tl.maximum(scale_bits.to(tl.float32, bitcast=True), 1.1754943508222875e-38) + + scaled_matrix = tl.maximum(tl.minimum(matrix / scale[:, None], 448.0), -448.0) + tl.store( + quantized_ptr + row * column_count + columns, + tl.reshape(scaled_matrix.to(tl.float8e4nv), (BLOCKS_PER_PROGRAM * 32, )), + mask=valid_elements, + ) + + blocks = first_block + tl.arange(0, BLOCKS_PER_PROGRAM) + # Map natural [row, K / 32] coordinates into cuBLAS's padded 128x4 tiles. + row_block = row // 128 + row_in_block = row % 128 + column_block = blocks // 4 + column_in_block = blocks % 4 + scale_offset = ((row_block * (scale_column_count // 4) + column_block) * 512 + (row_in_block % 32) * 16 + + (row_in_block // 32) * 4 + column_in_block) + tl.store(scale_ptr + scale_offset, scale_biased.to(tl.uint8), mask=blocks < scale_column_count) + + +@triton.jit +def _quantize_mxfp8_kernel( + matrix_ptr, + quantized_ptr, + scale_ptr, + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM: tl.constexpr, +): + """Load BF16 activation blocks for direct MXFP8 quantization.""" + row = tl.program_id(0) + first_block = tl.program_id(1) * BLOCKS_PER_PROGRAM + element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32) + columns = first_block * 32 + element_offsets + matrix = tl.load( + matrix_ptr + row * column_count + columns, + mask=(row < row_count) & (columns < column_count), + other=0.0, + ).to(tl.float32) + _store_mxfp8_blocks( + matrix, + quantized_ptr, + scale_ptr, + row, + first_block, + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM, + ) + + +@triton.jit +def _swiglu_quantize_mxfp8_kernel( + preactivation_ptr, + quantized_ptr, + scale_ptr, + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM: tl.constexpr, +): + """Apply value-first SwiGLU and directly quantize its BF16 result.""" + row = tl.program_id(0) + first_block = tl.program_id(1) * BLOCKS_PER_PROGRAM + element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32) + columns = first_block * 32 + element_offsets + valid_elements = (row < row_count) & (columns < column_count) + packed_row_offset = row * (2 * column_count) + values = tl.load( + preactivation_ptr + packed_row_offset + columns, + mask=valid_elements, + other=0.0, + ).to(tl.float32) + gates = tl.load( + preactivation_ptr + packed_row_offset + column_count + columns, + mask=valid_elements, + other=0.0, + ).to(tl.float32) + postactivation = (values * gates * tl.sigmoid(gates)).to(tl.bfloat16).to(tl.float32) + _store_mxfp8_blocks( + postactivation, + quantized_ptr, + scale_ptr, + row, + first_block, + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM, + ) + + +def _allocate_mxfp8_outputs(matrix: torch.Tensor, column_count: int) -> tuple[torch.Tensor, torch.Tensor, int]: + """Allocate contiguous values and a complete padded blocked-scale buffer.""" + row_count = matrix.shape[0] + scale_column_count = triton.cdiv(column_count // MXFP8_BLOCK_SIZE, 4) * 4 + blocked_scale_size = triton.cdiv(row_count, 128) * 128 * scale_column_count + quantized = torch.empty((row_count, column_count), dtype=torch.float8_e4m3fn, device=matrix.device) + scale_storage = torch.empty(blocked_scale_size, dtype=torch.uint8, device=matrix.device) + return quantized, scale_storage.view(torch.float8_e8m0fnu), scale_column_count + + +def quantize_mxfp8_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize activation rows and write hardware-blocked scales directly.""" + _validate_mxfp8_matrix(matrix) + matrix = matrix.contiguous() + row_count, column_count = matrix.shape + quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(matrix, column_count) + grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM)) + _quantize_mxfp8_kernel[grid]( + matrix, + quantized, + blocked_scales.view(torch.uint8), + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM, + num_warps=8, + ) + return quantized, blocked_scales + + +def swiglu_quantize_mxfp8_blockwise(preactivation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Apply H3 value-first SwiGLU and quantize the BF16 result to MXFP8.""" + if preactivation.ndim != 2: + raise ValueError(f"MXFP8 SwiGLU requires a 2D tensor, got shape {tuple(preactivation.shape)}.") + if preactivation.dtype != torch.bfloat16: + raise TypeError(f"MXFP8 SwiGLU requires BF16 input, got {preactivation.dtype}.") + if preactivation.shape[1] % (2 * MXFP8_BLOCK_SIZE): + raise ValueError("MXFP8 SwiGLU requires each packed half to be divisible by " + f"{MXFP8_BLOCK_SIZE}, got packed width {preactivation.shape[1]}.") + preactivation = preactivation.contiguous() + row_count = preactivation.shape[0] + column_count = preactivation.shape[1] // 2 + quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(preactivation, column_count) + grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM)) + _swiglu_quantize_mxfp8_kernel[grid]( + preactivation, + quantized, + blocked_scales.view(torch.uint8), + row_count, + column_count, + scale_column_count, + BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM, + num_warps=8, + ) + return quantized, blocked_scales + + +def mxfp8_scaled_mm( + activation_values: torch.Tensor, + activation_scales: torch.Tensor, + weight_values: torch.Tensor, + weight_scales: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + """Multiply two MXFP8 matrices and return a BF16 matrix.""" + return F.scaled_mm( + mat_a=activation_values, + mat_b=weight_values.mT, + scale_a=activation_scales, + scale_recipe_a=F.ScalingType.BlockWise1x32, + scale_b=weight_scales, + scale_recipe_b=F.ScalingType.BlockWise1x32, + swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, + swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, + bias=bias, + output_dtype=torch.bfloat16, + ) + + +def _resolve_merged_linear(linear: torch.nn.Module) -> torch.nn.Module: + """Return a linear whose weight includes every active inference adapter.""" + base_layer = getattr(linear, "base_layer", linear) + if base_layer is linear: + return base_layer + if not getattr(linear, "merged", False) and not getattr(linear, "disable_lora", False): + raise RuntimeError("MXFP8 feed-forward requires active LoRA weights to be merged before inference.") + return base_layer + + +def mxfp8_swiglu_feed_forward( + hidden_states: torch.Tensor, + fc_in: torch.nn.Module, + fc_out: torch.nn.Module, +) -> torch.Tensor: + """Run the MiniMax-H3 feed-forward network with MXFP8 GEMMs.""" + fc_in_base = _resolve_merged_linear(fc_in) + fc_out_base = _resolve_merged_linear(fc_out) + fc_in_method = fc_in_base.quant_method + fc_out_method = fc_out_base.quant_method + + preactivation = fc_in_method.apply(fc_in_base, hidden_states, fc_in_base.bias) + output_shape = (*hidden_states.shape[:-1], fc_out_base.output_size) + preactivation_2d = preactivation.reshape(-1, preactivation.shape[-1]) + activation_values, activation_scales = swiglu_quantize_mxfp8_blockwise(preactivation_2d) + output_2d = fc_out_method.apply_quantized( + fc_out_base, + activation_values, + activation_scales, + fc_out_base.bias, + ) + return output_2d.reshape(output_shape) + + +__all__ = [ + "MXFP8_BLOCK_SIZE", + "mxfp8_scaled_mm", + "mxfp8_swiglu_feed_forward", + "quantize_mxfp8_blockwise", + "quantize_mxfp8_weight_blockwise", + "swiglu_quantize_mxfp8_blockwise", +] diff --git a/fastvideo/layers/quantization/__init__.py b/fastvideo/layers/quantization/__init__.py index 5e805fa3a..83f5076f1 100644 --- a/fastvideo/layers/quantization/__init__.py +++ b/fastvideo/layers/quantization/__init__.py @@ -2,7 +2,16 @@ from typing import Literal, get_args from fastvideo.layers.quantization.base_config import QuantizationConfig -QuantizationMethods = Literal[None, "AbsMaxFP8", "FP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train", "fp8_qat_train"] +QuantizationMethods = Literal[ + None, + "AbsMaxFP8", + "FP8", + "MXFP8", + "NVFP4", + "nvfp4_qat", + "nvfp4_qat_train", + "fp8_qat_train", +] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -52,6 +61,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # lazy import to avoid triggering `torch.compile` too early from .absmax_fp8 import AbsMaxFP8Config from .fp8_config import FP8Config + from .mxfp8_config import MXFP8Config from .nvfp4_config import NVFP4Config from .nvfp4_qat_config import NVFP4QATConfig from .nvfp4_qat_train_config import NVFP4QATTrainConfig @@ -60,6 +70,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: method_to_config: dict[str, type[QuantizationConfig]] = { "AbsMaxFP8": AbsMaxFP8Config, "FP8": FP8Config, + "MXFP8": MXFP8Config, "NVFP4": NVFP4Config, "nvfp4_qat": NVFP4QATConfig, "nvfp4_qat_train": NVFP4QATTrainConfig, diff --git a/fastvideo/layers/quantization/mxfp8_config.py b/fastvideo/layers/quantization/mxfp8_config.py new file mode 100644 index 000000000..a450ede22 --- /dev/null +++ b/fastvideo/layers/quantization/mxfp8_config.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 quantization for MiniMax-H3 transformer-block feed-forward layers.""" + +from __future__ import annotations + +import re +from typing import Any + +import torch +from torch.nn.parameter import Parameter + +from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase +from fastvideo.logger import init_logger +from fastvideo.models.utils import set_weight_attrs + +logger = init_logger(__name__) + +_MINIMAX_H3_FF_PREFIX = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.ff\.(?:fc_in|fc_out)$") + + +class MXFP8QuantizeMethod(QuantizeMethodBase): + """Dynamically quantize activations against prequantized MXFP8 weights.""" + + 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: Any, + ) -> None: + """Create the BF16 checkpoint weight and non-persistent MXFP8 buffers.""" + del input_size, output_size + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter("weight", weight) + set_weight_attrs(weight, extra_weight_attrs) + layer.register_buffer("_mxfp8_weight", None, persistent=False) + layer.register_buffer("_mxfp8_weight_scale", None, persistent=False) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Prequantize one adapter-merged BF16 linear weight.""" + from fastvideo.layers.mxfp8linear import quantize_mxfp8_weight_blockwise + + quantized_weight, blocked_scales = quantize_mxfp8_weight_blockwise(layer.weight.detach()) + layer._mxfp8_weight = quantized_weight + layer._mxfp8_weight_scale = blocked_scales + + def apply( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Quantize one BF16 activation and apply the prequantized linear.""" + from fastvideo.layers.mxfp8linear import quantize_mxfp8_blockwise + + original_shape = hidden_states.shape + hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1]) + activation_values, activation_scales = quantize_mxfp8_blockwise(hidden_states_2d) + output_2d = self.apply_quantized(layer, activation_values, activation_scales, bias) + return output_2d.reshape(*original_shape[:-1], layer.output_size) + + def apply_quantized( + self, + layer: torch.nn.Module, + activation_values: torch.Tensor, + activation_scales: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Apply one linear to an activation that is already in MXFP8.""" + from fastvideo.layers.mxfp8linear import mxfp8_scaled_mm + + if layer._mxfp8_weight is None or layer._mxfp8_weight_scale is None: + raise RuntimeError(f"MXFP8 weight buffers are not initialized for {layer.prefix}.") + return mxfp8_scaled_mm( + activation_values, + activation_scales, + layer._mxfp8_weight, + layer._mxfp8_weight_scale, + bias, + ) + + +class MXFP8Config(QuantizationConfig): + """Select MXFP8 for the main MiniMax-H3 transformer-block FFN linears.""" + + def get_name(self) -> str: + return "MXFP8" + + def get_supported_act_dtypes(self) -> 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]) -> MXFP8Config: + del config + return cls() + + def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None: + from fastvideo.layers.linear import LinearBase + + if isinstance(layer, LinearBase) and _MINIMAX_H3_FF_PREFIX.search(prefix): + return MXFP8QuantizeMethod() + return None + + +def convert_model_to_mxfp8(model: torch.nn.Module) -> int: + """Prequantize every MXFP8-tagged weight and return the converted count.""" + converted_count = 0 + with torch.no_grad(): + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if not isinstance(quant_method, MXFP8QuantizeMethod): + continue + if converted_count == 0 and (not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10): + raise RuntimeError("MXFP8 inference requires an NVIDIA Blackwell GPU with compute capability 10.0+.") + quant_method.process_weights_after_loading(module) + converted_count += 1 + + if converted_count: + logger.info("Prequantized %d MiniMax-H3 feed-forward linear weights to MXFP8", converted_count) + torch.cuda.empty_cache() + return converted_count + + +__all__ = [ + "MXFP8Config", + "MXFP8QuantizeMethod", + "convert_model_to_mxfp8", +] diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index ba90fa4a7..284622d59 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/fastvideo/layers/quantization/nvfp4_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""LTX-2 NVFP4 quantization (FlashInfer-backed). +"""NVFP4 linear quantization for supported transformer layer sets. NVFP4 is NVIDIA's block-scaled FP4 format (e2m1 mantissa, fp32 alpha, ``layout_128x4`` scale layout, group size 16) — distinct from @@ -8,9 +8,8 @@ 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). -Upstreamed from ``FastVideo-internal`` so consumers that load LTX-2 -weights with NVFP4 quantization can drive the public package -end-to-end. +The registered config targets the curated LTX-2 deployment set and the +main MiniMax-H3 transformer-block FFN linears. `flashinfer` is imported lazily inside the call paths that need it. This keeps ``import fastvideo`` cheap on hosts where flashinfer is @@ -20,6 +19,7 @@ use time, with a clear error. from __future__ import annotations import logging +import re from typing import Any import torch @@ -75,6 +75,7 @@ _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES = ( _LTX2_NVFP4_LINEAR_PREFIXES = frozenset(f"ltx2.blocks.{block_idx}.{suffix}" for block_idx in range(48) 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)$") def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool: @@ -82,6 +83,11 @@ def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool: return prefix in _LTX2_NVFP4_LINEAR_PREFIXES +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_ltx2_refine_only_prefix(prefix: str) -> bool: return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES) @@ -421,13 +427,12 @@ class NVFP4QuantizeMethod(QuantizeMethodBase): class NVFP4Config(QuantizationConfig): - """LTX-2-specific NVFP4 quantization configuration. + """Select NVFP4 for the supported LTX-2 and MiniMax-H3 linear sets. NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha, - ``layout_128x4`` scale layout, group size 16). Today this class - hardcodes the LTX-2 layer paths it covers. When a second model - wants NVFP4, lift the layer-path list into a config field - instead of hardcoding it here. + ``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. """ def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None): @@ -468,9 +473,10 @@ class NVFP4Config(QuantizationConfig): def get_quant_method(self, layer: torch.nn.Module, prefix: str): from fastvideo.layers.linear import LinearBase - # Use the superset at build/load time, then switch active subset - # dynamically in NVFP4QuantizeMethod.apply based on stage profile. - if isinstance(layer, LinearBase) and is_ltx2_nvfp4_linear_prefix(prefix): + # 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)): method = NVFP4QuantizeMethod(layer_prefix=prefix) method._retain_original_weights = self.retain_original_weights return method @@ -552,4 +558,5 @@ __all__ = [ "NVFP4QuantizeMethod", "convert_model_to_nvfp4", "is_ltx2_nvfp4_linear_prefix", + "is_minimax_h3_nvfp4_linear_prefix", ] diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 3dff637fe..370d11d47 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -21,6 +21,7 @@ from fastvideo.distributed.communication_op import ( ) from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod from fastvideo.layers.mlp import MLP from fastvideo.layers.quantization import QuantizationConfig from fastvideo.layers.visual_embedding import Timesteps @@ -119,8 +120,14 @@ class MiniMaxH3FeedForward(nn.Module): prefix=f"{prefix}.fc_out", ) self.fuse_swiglu = fuse_swiglu + self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance( + self.fc_out.quant_method, MXFP8QuantizeMethod) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_mxfp8: + from fastvideo.layers.mxfp8linear import mxfp8_swiglu_feed_forward + + return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out) hidden_states, _ = self.fc_in(hidden_states) if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states): hidden_states = minimax_h3_swiglu(hidden_states) @@ -514,41 +521,48 @@ class MiniMaxH3TransformerBlock(nn.Module): rotary_emb: tuple[torch.Tensor, torch.Tensor], original_seq_len: int, ) -> torch.Tensor: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - t.to(hidden_states.dtype) for t in self.adaln_proj(temb)) + with nvtx_range("minimax_h3.transformer_block.adaln_projection"): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + t.to(hidden_states.dtype) for t in self.adaln_proj(temb)) use_modulate_fusion = self.fuse_modulate and _can_run_minimax_h3_fusion(hidden_states) if use_modulate_fusion: - norm_hidden_states = fused_rmsnorm_modulate( - hidden_states, - self.norm1.weight, - scale_msa, - shift_msa, - adaln_indices, - self.norm1.eps, - ) + with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): + norm_hidden_states = fused_rmsnorm_modulate( + hidden_states, + self.norm1.weight, + scale_msa, + shift_msa, + adaln_indices, + self.norm1.eps, + ) else: - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * ( - 1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices) - attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len) + with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"): + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * ( + 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) if use_modulate_fusion: - hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( - hidden_states, - attention_output, - gate_msa, - self.norm2.weight, - scale_mlp, - shift_mlp, - adaln_indices, - self.norm2.eps, - ) + with nvtx_range("minimax_h3.transformer_block.modulate_fusion"): + hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate( + hidden_states, + attention_output, + gate_msa, + self.norm2.weight, + scale_mlp, + shift_mlp, + adaln_indices, + self.norm2.eps, + ) else: - hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output - 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) - feed_forward_output = self.ff(norm_hidden_states) + with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"): + hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output + 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"): + feed_forward_output = self.ff(norm_hidden_states) return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 2028d1b26..e076359ee 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -36,14 +36,14 @@ 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 _maybe_quantize_model(model: nn.Module) -> None: - """Quantize NVFP4- or FP8-tagged linear layers in-place after weights are loaded. +def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lora_merge: bool = False) -> None: + """Quantize inference linear weights after checkpoint loading. - Walks the module tree once, looking for layers whose ``quant_method`` - is an :class:`NVFP4QuantizeMethod` or :class:`FP8QuantizeMethod` (attached - at construction time by the respective ``get_quant_method``). When at least + Walks the module tree once, looking for layers whose ``quant_method`` is an + inference quantization method attached at construction time. When at least one such layer exists, calls the matching conversion function to register - quantized weight buffers on each targeted layer. + quantized weight buffers on each targeted layer. NVFP4 and MXFP8 conversion + can be deferred until a configured inference LoRA is merged. The walk returns on the first quantized layer found so unquantized callers pay only an ``isinstance`` check per module. Both imports are deferred so @@ -70,12 +70,19 @@ def _maybe_quantize_model(model: nn.Module) -> None: FP8QuantizeMethod, convert_model_to_fp8, ) + from fastvideo.layers.quantization.mxfp8_config import ( + MXFP8QuantizeMethod, + convert_model_to_mxfp8, + ) 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 @@ -87,6 +94,13 @@ def _maybe_quantize_model(model: nn.Module) -> None: logger.info("Converting loaded model weights for FP8 linear layers") convert_model_to_fp8(model) return + if isinstance(qm, MXFP8QuantizeMethod): + if defer_weight_conversion_until_lora_merge: + logger.info("Deferring MXFP8 weight conversion until the inference LoRA merge completes") + return + logger.info("Converting loaded model weights for MXFP8 linear layers") + convert_model_to_mxfp8(model) + return # QAT-train configs are mutually exclusive with the inference schemes # above (one quant_config per model), so when they're active the loop # always runs to completion and the counts below are model-wide. @@ -144,6 +158,26 @@ def _prepare_model_for_compile(model: nn.Module, *, regional: bool) -> str | Non return unsupported if isinstance(unsupported, str) and unsupported else None +def _validate_fsdp_inference_quantization(init_params: dict[str, Any], fsdp_inference: bool) -> None: + """Enforce the transformer-quantization allowlist for FSDP inference.""" + transformer_config = init_params.get("config") + quant_config = getattr(transformer_config, "quant_config", None) + if not fsdp_inference or quant_config is None: + return + + from fastvideo.layers.quantization.nvfp4_qat_train_config import ( + NVFP4QATTrainConfig, ) + + if isinstance(quant_config, NVFP4QATTrainConfig): + return + + # TODO: (David) Currently reject every FSDP inference quantization config except NVFP4QATTrainConfig. + # Support FSDP inference with precomputed quantized weights. + raise NotImplementedError( + "FSDP inference supports unquantized transformers and NVFP4QATTrainConfig only; " + f"got {type(quant_config).__name__}.") + + # Supports optional torch.compile for FSDP-wrapped models during training def maybe_load_fsdp_model( model_cls: type[nn.Module], @@ -178,6 +212,8 @@ def maybe_load_fsdp_model( base checkpoint does not contain, which has to happen while the tensor is still unsharded. """ + _validate_fsdp_inference_quantization(init_params, fsdp_inference) + # 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) @@ -287,7 +323,7 @@ def maybe_load_fsdp_model( # responsibility is just to materialize the quantized weight buffers # from the freshly-loaded bf16 weights. No-op when no quantized layers # are present (lazy imports inside the helper). - _maybe_quantize_model(model) + _maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=lora_path is not None) compile_in_loader = enable_torch_compile and training_mode if compile_in_loader: diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index 3dbeae645..412cb2278 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -32,6 +32,36 @@ from fastvideo.utils import maybe_download_lora logger = init_logger(__name__) +def _has_quantized_mxfp8_weights(transformer_modules: dict[str, nn.Module]) -> bool: + """Return whether any transformer contains an initialized MXFP8 weight.""" + from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod + + return any( + isinstance(getattr(module, "quant_method", None), MXFP8QuantizeMethod) + and getattr(module, "_mxfp8_weight", None) is not None for transformer in transformer_modules.values() + for module in transformer.modules()) + + +def _has_nvfp4_weights_without_bf16(transformer_modules: dict[str, nn.Module]) -> bool: + """Return whether NVFP4 quantization removed any transformer's BF16 weight.""" + from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod + + return any( + isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod) + and getattr(module, "_nvfp4_weight", None) is not None and getattr(module, "weight", None) is None + for transformer in transformer_modules.values() for module in transformer.modules()) + + +def _has_quantized_nvfp4_weights(transformer_modules: dict[str, nn.Module]) -> bool: + """Return whether any transformer contains an initialized NVFP4 weight.""" + from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod + + return any( + isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod) + and getattr(module, "_nvfp4_weight", None) is not None for transformer in transformer_modules.values() + for module in transformer.modules()) + + def _get_hook_ctx(module: nn.Module | None): if module is None: return nullcontext() @@ -43,6 +73,24 @@ def _get_hook_ctx(module: nn.Module | None): return nullcontext() +def _convert_quantized_weights_after_lora_merge(transformer_modules: dict[str, nn.Module]) -> None: + """Pack deferred NVFP4 or MXFP8 weights after all LoRA layers are merged.""" + from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod + from fastvideo.layers.quantization.mxfp8_config import convert_model_to_mxfp8 + from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod + from fastvideo.layers.quantization.nvfp4_config import convert_model_to_nvfp4 + + for transformer_module in transformer_modules.values(): + for module in transformer_module.modules(): + quant_method = getattr(module, "quant_method", None) + if isinstance(quant_method, NVFP4QuantizeMethod): + convert_model_to_nvfp4(transformer_module) + break + if isinstance(quant_method, MXFP8QuantizeMethod): + convert_model_to_mxfp8(transformer_module) + break + + def _named_module_by_prefix(module: nn.Module, prefixes: list[str]) -> list[tuple[str | None, list[tuple[str, nn.Module]]]]: none_list: list[tuple[str, nn.Module]] = [] @@ -389,6 +437,13 @@ class LoRAPipeline(ComposedPipelineBase): if exact_current_adapter: return + if not self._setting_constructor_adapter and _has_nvfp4_weights_without_bf16( + self.trainable_transformer_modules): + # TODO(David): Restore the BF16 weights and requantize them after an NVFP4 LoRA adapter change. + raise RuntimeError( + "Runtime LoRA adapter changes are unsupported after NVFP4 quantization removed the BF16 weights. " + "Create a new VideoGenerator with the desired LoRA adapter.") + if not self._setting_constructor_adapter: if self._constructor_dense_lora_path is not None: raise RuntimeError( @@ -532,6 +587,7 @@ class LoRAPipeline(ComposedPipelineBase): lora_path) if unmatched: logger.warning("LoRA adapter %s: %d weights did not reach a layer", lora_path, len(unmatched)) + _convert_quantized_weights_after_lora_merge(self.trainable_transformer_modules) def merge_lora_weights(self) -> None: for ( @@ -547,6 +603,17 @@ class LoRAPipeline(ComposedPipelineBase): layer.merge_lora_weights() def unmerge_lora_weights(self) -> None: + """Unmerge LoRA weights when the transformer's quantized weights remain valid.""" + if _has_quantized_mxfp8_weights(self.trainable_transformer_modules): + # TODO(David): Requantize MXFP8 weights after LoRA unmerge before enabling this operation. + raise RuntimeError( + "LoRA unmerge is unsupported after MXFP8 weight quantization because the quantized weights still " + "contain the merged LoRA adapter.") + if _has_quantized_nvfp4_weights(self.trainable_transformer_modules): + # TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge. + raise RuntimeError( + "LoRA unmerge is unsupported after NVFP4 weight quantization because the quantized weights would " + "not reflect the unmerged LoRA adapter.") for ( transformer_name, transformer_lora_layers, diff --git a/fastvideo/tests/loader/test_fsdp_inference_quantization.py b/fastvideo/tests/loader/test_fsdp_inference_quantization.py new file mode 100644 index 000000000..30ee5548d --- /dev/null +++ b/fastvideo/tests/loader/test_fsdp_inference_quantization.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FSDP inference accepts unquantized transformers and NVFP4QATTrainConfig.""" + +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.layers.quantization.fp8_qat_train_config import FP8QATTrainConfig +from fastvideo.layers.quantization.mxfp8_config import MXFP8Config +from fastvideo.layers.quantization.nvfp4_config import NVFP4Config +from fastvideo.layers.quantization.nvfp4_qat_train_config import ( + NVFP4QATTrainConfig, ) +from fastvideo.models.loader.fsdp_load import ( + _validate_fsdp_inference_quantization, + maybe_load_fsdp_model, +) + + +def _init_params(quant_config: object | None) -> dict[str, object]: + return {"config": SimpleNamespace(quant_config=quant_config)} + + +@pytest.mark.parametrize("quant_config", [MXFP8Config(), NVFP4Config(), FP8QATTrainConfig()]) +def test_maybe_load_fsdp_model_unsupported_quantization_rejected_before_construction(quant_config: object) -> None: + """Reject unsupported quantization before model construction or FSDP sharding.""" + with pytest.raises(NotImplementedError, match=type(quant_config).__name__): + maybe_load_fsdp_model( + model_cls=torch.nn.Module, + init_params=_init_params(quant_config), + weight_dir_list=[], + device=torch.device("cpu"), + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + default_dtype=torch.bfloat16, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + fsdp_inference=True, + training_mode=False, + ) + + +def test_validate_fsdp_inference_quantization_nvfp4_qat_training_allowed() -> None: + _validate_fsdp_inference_quantization(_init_params(NVFP4QATTrainConfig()), fsdp_inference=True) + + +def test_validate_fsdp_inference_quantization_unquantized_allowed() -> None: + _validate_fsdp_inference_quantization(_init_params(None), fsdp_inference=True) + + +def test_validate_fsdp_inference_quantization_without_fsdp_allowed() -> None: + _validate_fsdp_inference_quantization(_init_params(MXFP8Config()), fsdp_inference=False) diff --git a/fastvideo/tests/ops/quantization/test_mxfp8.py b/fastvideo/tests/ops/quantization/test_mxfp8.py new file mode 100644 index 000000000..20c215fa5 --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_mxfp8.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell numerical parity tests for FastVideo MXFP8 operations.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from quack.mx_utils import to_blocked, to_mx + +from fastvideo.layers.mxfp8linear import ( + quantize_mxfp8_blockwise, + quantize_mxfp8_weight_blockwise, + swiglu_quantize_mxfp8_blockwise, +) + + +def _require_blackwell() -> None: + """Require hardware that can execute MXFP8 block-scaled GEMMs.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: + pytest.skip("MXFP8 numerical parity requires an NVIDIA Blackwell GPU") + + +def _quack_mxfp8(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize one matrix with the independent Quack reference operations.""" + values, natural_scales = to_mx(matrix.contiguous(), 32) + return values, to_blocked(natural_scales) + + +def _assert_same_mxfp8( + actual: tuple[torch.Tensor, torch.Tensor], + expected: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Require identical FP8 values and E8M0 scale bytes.""" + actual_values, actual_scales = actual + expected_values, expected_scales = expected + torch.testing.assert_close(actual_values.float(), expected_values.float(), rtol=0, atol=0) + torch.testing.assert_close( + actual_scales.view(torch.uint8), + expected_scales.view(torch.uint8), + rtol=0, + atol=0, + ) + + +def test_mxfp8_weight_quantization_matches_quack() -> None: + """Compare prequantized weight values and scales with Quack.""" + _require_blackwell() + generator = torch.Generator().manual_seed(20260905) + weight = torch.randn(129, 256, generator=generator, dtype=torch.bfloat16).cuda() + + _assert_same_mxfp8(quantize_mxfp8_weight_blockwise(weight), _quack_mxfp8(weight)) + + +def test_mxfp8_activation_quantization_matches_quack() -> None: + """Compare fused activation quantization and scale swizzling with Quack.""" + _require_blackwell() + generator = torch.Generator().manual_seed(20260905) + activation = torch.randn(129, 256, generator=generator, dtype=torch.bfloat16).cuda() + + _assert_same_mxfp8(quantize_mxfp8_blockwise(activation), _quack_mxfp8(activation)) + + +def test_mxfp8_swiglu_quantization_matches_bf16_quack() -> None: + """Compare fused SwiGLU quantization with BF16 SwiGLU followed by Quack.""" + _require_blackwell() + generator = torch.Generator().manual_seed(20260905) + preactivation = torch.randn(129, 512, generator=generator, dtype=torch.bfloat16).cuda() + values, gates = preactivation.chunk(2, dim=-1) + bf16_swiglu = (values.float() * gates.float() * torch.sigmoid(gates.float())).to(torch.bfloat16) + + _assert_same_mxfp8( + swiglu_quantize_mxfp8_blockwise(preactivation), + _quack_mxfp8(bf16_swiglu), + ) + + +def test_minimax_h3_mxfp8_feed_forward_matches_bf16() -> None: + """Compare the complete MXFP8 H3 feed-forward output with BF16.""" + _require_blackwell() + from fastvideo.layers.quantization.mxfp8_config import MXFP8Config, convert_model_to_mxfp8 + from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward + + previous_default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + dense = MiniMaxH3FeedForward(128, 256, prefix="minimax_h3.transformer_blocks.0.ff") + quantized = MiniMaxH3FeedForward( + 128, + 256, + quant_config=MXFP8Config(), + prefix="minimax_h3.transformer_blocks.0.ff", + ) + finally: + torch.set_default_dtype(previous_default_dtype) + + generator = torch.Generator().manual_seed(20260905) + with torch.no_grad(): + for parameter in dense.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + quantized.load_state_dict(dense.state_dict(), strict=True) + dense = dense.cuda().eval() + quantized = quantized.cuda().eval() + assert convert_model_to_mxfp8(quantized) == 2 + + hidden_states = torch.randn(2, 128, 128, generator=generator, dtype=torch.bfloat16).cuda() + with torch.inference_mode(): + dense_output = dense(hidden_states) + quantized_output = quantized(hidden_states) + + dense_output_flat = dense_output.float().flatten() + quantized_output_flat = quantized_output.float().flatten() + similarity = F.cosine_similarity(dense_output_flat, quantized_output_flat, dim=0) + relative_l2_error = torch.linalg.vector_norm(dense_output_flat - quantized_output_flat) / torch.linalg.vector_norm( + dense_output_flat + ) + assert dense_output.dtype == torch.bfloat16 + assert quantized_output.dtype == torch.bfloat16 + assert similarity > 0.995 + assert relative_l2_error < 0.10 diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py b/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py new file mode 100644 index 000000000..d5529c04b --- /dev/null +++ b/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3 NVFP4 targeting and inference-LoRA lifecycle tests.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +import fastvideo.layers.quantization.mxfp8_config as mxfp8 +import fastvideo.layers.quantization.nvfp4_config as nvfp4 + + +class _TaggedModel(nn.Module): + """Minimal model whose child carries one quantization method.""" + + def __init__(self, quant_method: object) -> None: + super().__init__() + self.linear = nn.Module() + self.linear.quant_method = quant_method + + +@pytest.mark.parametrize( + "prefix", + [ + "transformer_blocks.0.ff.fc_in", + "transformer_blocks.49.ff.fc_out", + "minimax_h3.transformer_blocks.12.ff.fc_in", + ], +) +def test_is_minimax_h3_nvfp4_linear_prefix_main_ffn_linear(prefix: str) -> None: + assert nvfp4.is_minimax_h3_nvfp4_linear_prefix(prefix) + + +@pytest.mark.parametrize( + "prefix", + [ + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.adaln_proj.linear", + "token_refiner.refiner_blocks.0.ff.fc_in", + "transformer_blocks.0.ff.extra", + ], +) +def test_is_minimax_h3_nvfp4_linear_prefix_non_main_ffn_linear(prefix: str) -> None: + assert not nvfp4.is_minimax_h3_nvfp4_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") + + from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward + + feed_forward = MiniMaxH3FeedForward( + hidden_size=16, + ffn_dim=32, + quant_config=nvfp4.NVFP4Config(), + prefix="minimax_h3.transformer_blocks.0.ff", + ) + + assert isinstance(feed_forward.fc_in.quant_method, nvfp4.NVFP4QuantizeMethod) + assert isinstance(feed_forward.fc_out.quant_method, nvfp4.NVFP4QuantizeMethod) + + +def test_minimax_h3_feed_forward_nvfp4_cuda_similarity() -> None: + """Run both H3 FFN GEMMs through NVFP4 and compare with BF16.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: + pytest.skip("MiniMax-H3 NVFP4 inference requires an NVIDIA Blackwell GPU") + pytest.importorskip("flashinfer") + + from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward + + previous_default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + dense = MiniMaxH3FeedForward(128, 256, prefix="minimax_h3.transformer_blocks.0.ff") + quantized = MiniMaxH3FeedForward( + 128, + 256, + quant_config=nvfp4.NVFP4Config(), + prefix="minimax_h3.transformer_blocks.0.ff", + ) + finally: + torch.set_default_dtype(previous_default_dtype) + + generator = torch.Generator().manual_seed(20260901) + with torch.no_grad(): + for parameter in dense.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02) + quantized.load_state_dict(dense.state_dict(), strict=True) + dense = dense.cuda().eval() + quantized = quantized.cuda().eval() + nvfp4.convert_model_to_nvfp4(quantized) + + inputs = torch.randn(2, 128, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + dense_output = dense(inputs) + quantized_output = quantized(inputs) + + similarity = torch.nn.functional.cosine_similarity( + dense_output.float().flatten(), + quantized_output.float().flatten(), + dim=0, + ) + assert dense_output.dtype == torch.bfloat16 + assert quantized_output.dtype == torch.bfloat16 + assert similarity > 0.97 + assert quantized.fc_in.weight is None + assert quantized.fc_out.weight is None + + +def test_maybe_quantize_model_lora_defers_nvfp4(monkeypatch: pytest.MonkeyPatch) -> None: + from fastvideo.models.loader.fsdp_load import _maybe_quantize_model + + model = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod)) + converted: list[nn.Module] = [] + monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", converted.append) + + _maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=True) + assert converted == [] + + _maybe_quantize_model(model) + assert converted == [model] + + +def test_convert_quantized_weights_after_lora_merge_dispatches_nvfp4( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastvideo.pipelines.lora_pipeline import _convert_quantized_weights_after_lora_merge + + model = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod)) + nvfp4_converted: list[nn.Module] = [] + mxfp8_converted: list[nn.Module] = [] + monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", nvfp4_converted.append) + monkeypatch.setattr(mxfp8, "convert_model_to_mxfp8", mxfp8_converted.append) + + _convert_quantized_weights_after_lora_merge({"transformer": model}) + + assert nvfp4_converted == [model] + assert mxfp8_converted == [] + + +def test_unmerge_lora_weights_rejects_quantized_mxfp8() -> None: + """Reject LoRA unmerge while the MXFP8 weight still contains the merged adapter.""" + from fastvideo.pipelines.lora_pipeline import LoRAPipeline + + transformer = _TaggedModel(object.__new__(mxfp8.MXFP8QuantizeMethod)) + transformer.linear.register_buffer("_mxfp8_weight", torch.empty(1)) + pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer}) + + with pytest.raises(RuntimeError, match="LoRA unmerge is unsupported after MXFP8 weight quantization"): + LoRAPipeline.unmerge_lora_weights(pipeline) + + +@pytest.mark.parametrize("retain_bf16_weight", [False, True]) +def test_unmerge_lora_weights_rejects_quantized_nvfp4(retain_bf16_weight: bool) -> None: + """Reject NVFP4 LoRA unmerge whether the BF16 weight was removed or retained.""" + from fastvideo.pipelines.lora_pipeline import LoRAPipeline + + transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod)) + transformer.linear.register_parameter( + "weight", + nn.Parameter(torch.empty(1), requires_grad=False) if retain_bf16_weight else None, + ) + transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1)) + pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer}) + + with pytest.raises(RuntimeError, match="LoRA unmerge is unsupported after NVFP4 weight quantization"): + LoRAPipeline.unmerge_lora_weights(pipeline) + + +def test_set_lora_adapter_rejects_nvfp4_without_bf16_weight() -> None: + """Reject runtime adapter changes after NVFP4 removes the BF16 weight.""" + from fastvideo.pipelines.lora_pipeline import LoRAPipeline + + transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod)) + transformer.linear.register_parameter("weight", None) + transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1)) + pipeline = SimpleNamespace( + _setting_constructor_adapter=False, + cur_adapter_name="first-adapter", + cur_adapter_path="/first-adapter", + cur_adapter_strength=1.0, + lora_adapter_paths={}, + trainable_transformer_modules={"transformer": transformer}, + ) + + with pytest.raises(RuntimeError, match="Runtime LoRA adapter changes are unsupported after NVFP4 quantization"): + LoRAPipeline.set_lora_adapter(pipeline, "second-adapter", "/second-adapter") + + +def test_set_lora_adapter_mxfp8_requantizes_after_adapter_merge(monkeypatch: pytest.MonkeyPatch) -> None: + """Requantize MXFP8 weights after merging a replacement LoRA adapter.""" + from fastvideo.pipelines.lora_pipeline import LoRAPipeline + + events: list[str] = [] + + class _RecordingLoRALayer: + + def set_lora_weights(self, lora_a: torch.Tensor, lora_b: torch.Tensor, **kwargs) -> None: + del lora_a, lora_b, kwargs + events.append("adapter merged") + + transformer = _TaggedModel(object.__new__(mxfp8.MXFP8QuantizeMethod)) + transformer.linear.register_buffer("_mxfp8_weight", torch.empty(1)) + transformer_lora_layers = SimpleNamespace( + lora_layers_by_block=lambda: iter([(None, {"linear": _RecordingLoRALayer()})]) + ) + pipeline = SimpleNamespace( + _setting_constructor_adapter=False, + _constructor_dense_lora_path=None, + cur_adapter_name="first-adapter", + cur_adapter_path=None, + cur_adapter_strength=1.0, + lora_adapter_paths={}, + lora_adapters={ + "second-adapter": { + "linear.lora_A": torch.empty(1, 1), + "linear.lora_B": torch.empty(1, 1), + } + }, + lora_initialized=True, + lora_layers={"transformer": transformer_lora_layers}, + fastvideo_args=SimpleNamespace(training_mode=False), + trainable_transformer_modules={"transformer": transformer}, + ) + monkeypatch.setattr("fastvideo.pipelines.lora_pipeline.dist.get_rank", lambda: 0) + monkeypatch.setattr(mxfp8, "convert_model_to_mxfp8", lambda model: events.append("weights requantized")) + + LoRAPipeline.set_lora_adapter(pipeline, "second-adapter") + + assert events == ["adapter merged", "weights requantized"] diff --git a/pyproject.toml b/pyproject.toml index cdaaf2a00..da3eadaea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -188,6 +188,7 @@ dev = [ "fastvideo[lint]", "fastvideo[test]", ] fasth3 = [ "flash-attn-4", "fastvideo-kernel==0.3.5; sys_platform == 'linux'", + "quack-kernels==0.5.3", ] prompt-safety = [