Enable Online FP8&NF4 on Minimax-H3

docs/en/quantization.md

@@ -79,6 +79,26 @@ TorchAO and PyTorch must be version-compatible. Check the
 [TorchAO release compatibility table](https://github.com/pytorch/ao/releases) instead of installing the newest
 TorchAO release blindly; an import warning or failure means that configuration has not been validated.
 
+### tf-kernel FP8: online W8A8
+
+MiniMax H3 can use TeleFuser's tf-kernel FP8 GEMM wrapper for online W8A8 inference directly from the original BF16
+checkpoint. Activations are quantized per token at each forward, weights are quantized per output channel and cached
+on first use, and `tf_kernel.fp8_scaled_mm` produces BF16 output. This path requires a tf-kernel wheel built for the
+runtime's PyTorch/CUDA ABI and GPU architecture; on H100, build the SM90 wheel from `tf-kernel/` with the Makefile.
+
+```python
+quant_config = QuantConfig(
+    enabled=True,
+    quant_type=QuantType.FP8,
+    kernel_backend=QuantKernelBackend.TF_KERNEL,
+)
+```
+
+For MiniMax H3, use `quantization="tf-kernel-fp8"` with
+`examples/minimax_h3/minimax_h3_fl2va_h100.py`. This backend is single-GPU only and keeps the FP8
+weights resident after first-use conversion. It is distinct from the scaled-FP8 checkpoint path below: the latter
+expects weights and scales already serialized in the checkpoint.
+
 ### bitsandbytes NF4: W4A16
 
 NF4 uses a non-uniform 4-bit codebook designed for approximately normal weight distributions. TeleFuser replaces selected Linear layers with `bitsandbytes.nn.Linear4bit`, uses BF16 compute, and enables compressed quantization statistics.

examples/minimax_h3/README.md

@@ -291,6 +291,71 @@ The Ulysses degree must divide 56 attention heads. Scripts must run from their g
 processes can spawn safely. H100 examples request packed FlashAttention 4 and fall back to packed PyTorch SDPA when
 FlashAttention 4 is unavailable.
 
+## Online DiT Quantization
+
+MiniMax H3 supports three single-GPU online quantization backends for the DiT transformer Linear layers:
+
+| CLI value | Backend | Weight/activation path |
+|---|---|---|
+| torchao-fp8 | TorchAO | FP8 dynamic activation and FP8 weight when supported, otherwise TorchAO's FP8 weight-only path |
+| tf-kernel-fp8 | TeleFuser tf-kernel | Per-token activation and per-output-channel weight FP8 (W8A8), BF16 output |
+| bnb-nf4 | bitsandbytes | NF4 weight-only with BF16 compute |
+
+All three paths convert the 258 Linear layers in the main and token-refiner transformer blocks. The FP32 video/audio
+patch projections, timestep embedding, output projections, text encoder, and VAEs retain their reference dtypes.
+The BF16 DiT is loaded from the original shards, moved to CUDA after text encoding, quantized on first denoising use,
+and then kept resident for the pipeline lifetime. This ordering avoids a simultaneous BF16 text encoder and DiT on
+one GPU and avoids unsupported CPU transfers of quantized tensor subclasses.
+TorchAO and tf-kernel conversion have a transient memory peak near the BF16 footprint; use the full 80 GB device without colocated workloads.
+
+Use the single FL2VA example and choose the quantization backend with `--quantization`:
+
+~~~bash
+python examples/minimax_h3/minimax_h3_fl2va_h100.py \
+  --mode t2va \
+  --quantization torchao-fp8 \
+  --duration 5 \
+  --output outputs/minimax_h3_torchao_fp8.mp4
+python examples/minimax_h3/minimax_h3_fl2va_h100.py \
+  --mode t2va \
+  --quantization bnb-nf4 \
+  --duration 5 \
+  --output outputs/minimax_h3_bnb_nf4.mp4
+python examples/minimax_h3/minimax_h3_fl2va_h100.py \
+  --mode t2va \
+  --quantization tf-kernel-fp8 \
+  --duration 5 \
+  --output outputs/minimax_h3_tf_kernel_fp8.mp4
+~~~
+
+The FL2VA CLI accepts `--quantization` with `torchao-fp8`, `tf-kernel-fp8`, or `bnb-nf4`; omit it for BF16. The Python
+loader accepts the same names:
+
+~~~python
+from examples.minimax_h3.common import load_minimax_h3_pipeline
+
+pipeline = load_minimax_h3_pipeline(
+    "/path/to/MiniMaxAI_MiniMax-H3",
+    partition="FL2VA",
+    quantization="tf-kernel-fp8",
+)
+~~~
+
+Online quantization currently requires ulysses_degree=1, tp_degree=1, and FSDP disabled. Quantizing before TP/FSDP
+would invalidate those wrappers' BF16 parameter-sharding contract, so unsupported combinations fail before checkpoint
+loading.
+
+For matched BF16/TorchAO-FP8/tf-kernel-FP8/NF4 profiling, use the validation benchmark. It writes the synchronized MP4 plus a JSON report
+containing load time, end-to-end generation time, stage timings, and denoising allocator peaks:
+
+~~~bash
+python tools/validation/benchmark_minimax_h3_quantization.py \
+  --backend tf-kernel-fp8 \
+  --duration 5 \
+  --steps 50 \
+  --output outputs/minimax_h3_tf_kernel_fp8_50step.mp4
+~~~
+
 For multi-GPU resident profiles, `WorkerTensorChannel` transports text conditioning, visual condition rows, and the
 final video latent directly between worker groups. CUDA intermediates therefore do not stage through the parent
 process or CPU. The pipeline reports media, text, condition VAE, denoising, video/audio decode, allocator peak, and
@@ -314,10 +379,10 @@ The standard four-GPU profile already uses Ulysses2 x TP2 and therefore leaves F
 `load_minimax_h3_pipeline` directly to construct another supported combination; the product of Ulysses and TP degrees
 must be 1, 2, or 4.
 
-Ring attention, CFG parallelism, pipeline parallelism, sparse attention, quantization, and `torch.compile` are not
-enabled for H3. Video-VAE parallelism is spatial tiling over the existing TP process group, not parameter tensor
-parallelism. The dedicated service manifests expose the pipeline without adding framework-level configuration fields
-or changing the shared request schema.
+Ring attention, CFG parallelism, pipeline parallelism, sparse attention, and `torch.compile` are not enabled for H3.
+Video-VAE parallelism is spatial tiling over the existing TP process group, not parameter tensor parallelism. The
+dedicated service manifests expose the pipeline without adding framework-level configuration fields or changing the
+shared request schema.
 
 ## Four-GPU Regression
 

examples/minimax_h3/common.py

@@ -19,6 +19,9 @@
     ModelRuntimeConfig,
     OffloadConfig,
     ParallelConfig,
+    QuantConfig,
+    QuantKernelBackend,
+    QuantType,
     WeightOffloadType,
 )
 from telefuser.core.module_manager import ModuleManager
@@ -127,6 +130,36 @@ def _checkpoint_shards(component: Path) -> list[str]:
     return shards
 
 
+def minimax_h3_quant_config(quantization: str | QuantType | None) -> QuantConfig:
+    """Resolve a public MiniMax H3 online-quantization name to runtime config."""
+    if quantization is None:
+        return QuantConfig()
+    if isinstance(quantization, str):
+        normalized = quantization.strip().lower().replace("_", "-")
+        names = {
+            "torchao-fp8": QuantType.TORCHAO_FP8,
+            "bnb-nf4": QuantType.BNB_NF4,
+            "tf-kernel-fp8": QuantType.FP8,
+        }
+        try:
+            quant_type = names[normalized]
+        except KeyError as exc:
+            raise ValueError("quantization must be 'torchao-fp8', 'tf-kernel-fp8', 'bnb-nf4', or None") from exc
+    elif isinstance(quantization, QuantType):
+        quant_type = quantization
+    else:
+        raise TypeError("quantization must be a string, QuantType, or None")
+
+    backends = {
+        QuantType.TORCHAO_FP8: QuantKernelBackend.TORCHAO,
+        QuantType.BNB_NF4: QuantKernelBackend.BITSANDBYTES,
+        QuantType.FP8: QuantKernelBackend.TF_KERNEL,
+    }
+    if quant_type not in backends:
+        raise ValueError(f"MiniMax H3 does not support online quantization type {quant_type.name}")
+    return QuantConfig(enabled=True, quant_type=quant_type, kernel_backend=backends[quant_type])
+
+
 def load_minimax_h3_pipeline(
     model_root: str | Path,
     *,
@@ -141,6 +174,7 @@ def load_minimax_h3_pipeline(
     feature_cache_config: FeatureCacheConfig | None = None,
     adaln_cache_path: str | Path | None = None,
     online_adaln_cache: bool = False,
+    quantization: str | QuantType | None = None,
 ) -> MiniMaxH3Pipeline:
     if adaln_cache_path is not None and online_adaln_cache:
         raise ValueError("Choose either adaln_cache_path or online_adaln_cache, not both.")
@@ -162,6 +196,11 @@ def load_minimax_h3_pipeline(
         raise ValueError("enable_fsdp requires multi-GPU sequence parallelism without tensor parallelism")
     if (adaln_cache_path is not None or online_adaln_cache) and resolved_enable_fsdp:
         raise ValueError("AdaLN cache modes do not yet support FSDP deployment.")
+    quant_config = minimax_h3_quant_config(quantization)
+    if quant_config.enabled and world_size != 1:
+        raise ValueError("MiniMax H3 online quantization currently requires a single-GPU profile")
+    if quant_config.enabled and resolved_enable_fsdp:
+        raise ValueError("MiniMax H3 online quantization cannot be combined with FSDP")
     if isinstance(attn_impl, str):
         try:
             attn_impl = AttnImplType[attn_impl]
@@ -171,13 +210,20 @@ def load_minimax_h3_pipeline(
     if not component_root.is_dir():
         raise FileNotFoundError(f"MiniMax H3 partition not found: {component_root}")
     runtime_device = torch.device(device)
+    if quant_config.enabled and runtime_device.type != "cuda":
+        raise ValueError("MiniMax H3 online quantization requires a CUDA device")
     use_resident_modules = world_size > 1 or resolved_enable_fsdp
     resident_offload = OffloadConfig(
         offload_type=(
             WeightOffloadType.NO_CPU_OFFLOAD if use_resident_modules else WeightOffloadType.MODEL_CPU_OFFLOAD
         ),
         pin_cpu_memory=False,
     )
+    dit_offload = (
+        OffloadConfig(offload_type=WeightOffloadType.NO_CPU_OFFLOAD, pin_cpu_memory=False)
+        if quant_config.enabled
+        else resident_offload
+    )
     text_parallel = (
         ParallelConfig(
             device_ids=list(range(resolved_encoder_tp)),
@@ -203,9 +249,10 @@ def load_minimax_h3_pipeline(
         device_type=runtime_device.type,
         device_id=runtime_device.index or 0,
         torch_dtype=torch.bfloat16,
-        offload_config=resident_offload,
+        offload_config=dit_offload,
         attention_config=AttentionConfig.dense_attention(attn_impl),
         feature_cache_config=feature_cache_config or FeatureCacheConfig(),
+        quant_config=quant_config,
         parallel_config=ParallelConfig(
             device_ids=list(range(world_size)),
             sp_ulysses_degree=ulysses_degree,
@@ -344,6 +391,7 @@ def save_generation(result: MiniMaxH3Generation, output_path: str | Path) -> Non
     "load_minimax_h3_pipeline",
     "load_minimax_h3_request",
     "minimax_h3_adaln_cache_timesteps",
+    "minimax_h3_quant_config",
     "partition_for_minimax_h3_request",
     "run_minimax_h3_request",
     "save_generation",

examples/minimax_h3/minimax_h3_fl2va_h100.py

@@ -38,6 +38,7 @@
     "feature_cache_model_type": "MiniMax-H3-Base",
     "feature_cache_n_derivatives": 1,
     "feature_cache_taylor_threshold": 2,
+    "quantization": None,
 }
 
 
@@ -88,6 +89,7 @@ def get_pipeline(
     feature_cache_model_type: str = PPL_CONFIG["feature_cache_model_type"],
     feature_cache_n_derivatives: int = PPL_CONFIG["feature_cache_n_derivatives"],
     feature_cache_taylor_threshold: int = PPL_CONFIG["feature_cache_taylor_threshold"],
+    quantization: str | None = PPL_CONFIG["quantization"],
 ) -> MiniMaxH3Pipeline:
     """Load the FL2VA checkpoint partition for one, two, or four GPUs."""
     tp_degree = 2 if parallelism == 4 else 1
@@ -108,6 +110,7 @@ def get_pipeline(
             n_derivatives=feature_cache_n_derivatives,
             taylor_threshold=feature_cache_taylor_threshold,
         ),
+        quantization=quantization,
     )
 
 
@@ -241,7 +244,7 @@ def run_with_file(
     return {"output_path": str(Path(output_path))}
 
 
-def main() -> None:
+def _main(default_quantization: str | None = PPL_CONFIG["quantization"]) -> None:
     parser = argparse.ArgumentParser(description="Generate MiniMax H3 T2VA/FL2VA audio-video on H100 GPUs")
     parser.add_argument("--model-root", default=PPL_CONFIG["model_root"])
     parser.add_argument("--mode", choices=("t2va", "first-frame", "last-frame", "first-last"))
@@ -265,6 +268,12 @@ def main() -> None:
     parser.add_argument("--flow-shift", type=float, default=PPL_CONFIG["flow_shift"])
     parser.add_argument("--audio-flow-shift", type=float, default=PPL_CONFIG["audio_flow_shift"])
     parser.add_argument("--device", default=PPL_CONFIG["device"])
+    parser.add_argument(
+        "--quantization",
+        choices=("torchao-fp8", "tf-kernel-fp8", "bnb-nf4"),
+        default=default_quantization,
+        help="Online DiT Linear quantization backend (single GPU only).",
+    )
     parser.add_argument("--gpu-num", "--ulysses-degree", dest="gpu_num", type=int, choices=(1, 2, 4), default=1)
     parser.add_argument(
         "--attn-impl",
@@ -319,6 +328,7 @@ def main() -> None:
         feature_cache_model_type=args.feature_cache_model_type,
         feature_cache_n_derivatives=args.feature_cache_n_derivatives,
         feature_cache_taylor_threshold=args.feature_cache_taylor_threshold,
+        quantization=args.quantization,
     )
     try:
         result = run_with_file(
@@ -339,5 +349,9 @@ def main() -> None:
         pipeline.stop()
 
 
+def main() -> None:
+    _main()
+
+
 if __name__ == "__main__":
     main()

examples/minimax_h3/minimax_h3_ref2va_h100.py

@@ -36,6 +36,7 @@
     "device": "cuda:0",
     "enable_fsdp": None,
     "online_adaln_cache": True,
+    "quantization": None,
 }
 
 PIPELINE_MANIFEST = build_pipeline_manifest(
@@ -91,6 +92,7 @@ def get_pipeline(
     num_inference_steps: int = PPL_CONFIG["num_inference_steps"],
     enable_fsdp: bool | None = PPL_CONFIG["enable_fsdp"],
     online_adaln_cache: bool = PPL_CONFIG["online_adaln_cache"],
+    quantization: str | None = PPL_CONFIG["quantization"],
 ) -> MiniMaxH3Pipeline:
     """Load the Ref2VA checkpoint partition for one, two, or four GPUs."""
     tp_degree = 2 if parallelism == 4 else 1
@@ -104,6 +106,7 @@ def get_pipeline(
         text_encoder_tp_degree=parallelism,
         enable_fsdp=enable_fsdp,
         online_adaln_cache=online_adaln_cache,
+        quantization=quantization,
     )
 
 
@@ -260,6 +263,7 @@ def main() -> None:
     parser.add_argument("--flow-shift", type=float, default=PPL_CONFIG["flow_shift"])
     parser.add_argument("--audio-flow-shift", type=float, default=PPL_CONFIG["audio_flow_shift"])
     parser.add_argument("--device", default=PPL_CONFIG["device"])
+    parser.add_argument("--quantization", choices=("torchao-fp8", "tf-kernel-fp8", "bnb-nf4"))
     parser.add_argument("--gpu-num", "--ulysses-degree", dest="gpu_num", type=int, choices=(1, 2, 4), default=1)
     fsdp_group = parser.add_mutually_exclusive_group()
     fsdp_group.add_argument("--enable-fsdp", dest="enable_fsdp", action="store_true")
@@ -299,6 +303,7 @@ def main() -> None:
         num_inference_steps=request_steps,
         enable_fsdp=args.enable_fsdp,
         online_adaln_cache=online_adaln_cache,
+        quantization=args.quantization,
     )
     try:
         if args.request is not None: