[perf] MiniMax H3: faster FP4 attention padding, uint8 VAE pixel path and cuDNN FP4 GEMMs (#2046)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
This commit is contained in:
Aryan Kumar
co-authored by
Aryan Kumar
Aryan Kumar
parent
166b9b8a07
commit
d954fad9ec
@@ -196,6 +196,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `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_MINIMAX_H3_DENSE_ATTN` | one of flash, cudnn | `flash` | attention | Kernel for MiniMax-H3 dense DiT self-attention: flash keeps the selected FlashAttention backend; with FLASH_ATTN, cudnn sends the no-mask, no-grad dense call through cuDNN's fused SDPA. cuDNN has a different reduction order, so this is an inference-only opt-in. |
|
||||
| `FASTVIDEO_VSA_PTX_BLACKWELL` | bool | `0` | attention | VIDEO_SPARSE_ATTN_H3 sends tile-64 forwards to the data-center Blackwell (sm_100a / sm_103a) kernel; grad-tracking forwards need its autograd operator, else Triton. fastvideo-kernel reads the same variable with the same rule. Deprecated names: `FASTVIDEO_VSA_SM100A`. |
|
||||
| `FASTVIDEO_VSA_TRITON` | bool | `0` | attention | Force the Triton MiniMax-H3 sparse attention kernel. fastvideo-kernel reads the same variable. |
|
||||
| `FASTVIDEO_NVFP4_FA4` | bool | `0` | attention | FlashAttention-4 quantizes Q and K to NVFP4. An explicit nvfp4_fa4 attention implementation argument takes precedence. |
|
||||
@@ -269,7 +270,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_H3_FP8_ATTENTION` | bool | `0` | performance | With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention projections in FP8. |
|
||||
| `FASTVIDEO_H3_FP8_GRANULARITY` | one of tensor, channel | `tensor` | performance | FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION. |
|
||||
| `FASTVIDEO_LORA_SIDE_BRANCH` | bool | `0` | performance | Apply inference LoRA adapters as a bf16 side branch next to every target linear instead of merging them, so the adapter strength can change at runtime. Packed NVFP4 or FP8 linears without their original dense weight always use the side branch. |
|
||||
| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. |
|
||||
| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears: auto (cudnn for MiniMax-H3 DiT profiles on sm_100), cutlass or cudnn. |
|
||||
| `FASTVIDEO_NVFP4_MM_ROW_SPLIT` | one of off, table, auto | `off` | performance | Run large cuDNN NVFP4 GEMMs over row chunks where its plan heuristic is slow: off, table (fixed per-shape chunk counts) or auto (time 1/2/4 chunks once per shape; not deterministic across runs). |
|
||||
| `FASTVIDEO_NVFP4_ACT_AMAX` | path | unset | performance | JSON of calibrated NVFP4 input amax per linear, keyed b<block>.<sub> or full prefix; sets a static activation scale. |
|
||||
| `FASTVIDEO_NVFP4_DYNAMIC_ACT` | str | `""` | performance | NVFP4 linears that derive the activation scale per call: all, or comma-separated layer-name suffixes such as ff.fc_out. |
|
||||
@@ -290,6 +291,9 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_H3_VAE_INT8_OVERLAY` | bool | `1` | performance | Apply the INT8 ConvRot decoder overlay when the MiniMax-H3 VAE folder ships one; 0 decodes with the dense weights. |
|
||||
| `FASTVIDEO_H3_VAE_INT8_SHARED_QKV` | bool | `0` | performance | Share the INT8 activation rotation and quantization across the MiniMax-H3 VAE Q/K/V projections. |
|
||||
| `FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW` | bool | `0` | performance | Use transposed weight views in the MiniMax-H3 VAE INT8 projections. |
|
||||
| `FASTVIDEO_H3_VAE_COMPILE_BLOCKS` | bool | `0` | performance | Run the MiniMax-H3 video VAE decoder transformer blocks through one regional torch.compile graph: Inductor fuses the fp32 residual scale/add pairs, the RMSNorm+cast and the SwiGLU silu\*mul. |
|
||||
| `FASTVIDEO_H3_VAE_UINT8_PIXELS` | bool | `1` | performance | Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per chunk and stream uint8 into the host buffer; frames are bitwise identical to the fp32 path. 0 streams fp32. |
|
||||
| `FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE encoder call (reference and keyframe encodes); 1 encodes per tile. Batched tiles differ from per-tile encoding only in reduction order. |
|
||||
| `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. |
|
||||
|
||||
@@ -114,6 +114,24 @@ the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point
|
||||
reduction order relative to fixed-length FA4, so treat it as a speed/quality
|
||||
evaluation option rather than an exact-parity mode.
|
||||
|
||||
### cuDNN SDPA for H3 dense attention (opt-in)
|
||||
|
||||
With the `FLASH_ATTN` backend, MiniMax-H3's dense, unmasked, inference-only DiT
|
||||
self-attention can use cuDNN's fused SDPA engine instead of FlashAttention:
|
||||
|
||||
```bash
|
||||
export FASTVIDEO_MINIMAX_H3_DENSE_ATTN=cudnn
|
||||
```
|
||||
|
||||
It keeps the same precision contract (bf16 operands, fp32 softmax/accumulation,
|
||||
bf16 output) but uses a different reduction order, so treat it as a
|
||||
speed/quality option rather than an exact-parity mode. Training calls, compiled
|
||||
regions, head sizes above 128, non-contiguous inputs, and the sparse VSA blocks
|
||||
stay on their existing paths. Standalone kernel timings at 37.7k and 109k tokens
|
||||
are within measurement noise of FA4; an in-model A/B on one GB200 measured a
|
||||
2.2% faster DiT forward. Requires CUDA and a torch build with
|
||||
`aten::_scaled_dot_product_cudnn_attention` (cuDNN 9.x).
|
||||
|
||||
### FP4 Flash Attention 4 (Blackwell only)
|
||||
|
||||
**`FLASH_ATTN`** with **`--nvfp4_fa4`**
|
||||
|
||||
@@ -6,6 +6,7 @@ import torch.nn.functional as F
|
||||
from dataclasses import dataclass
|
||||
|
||||
import fastvideo.envs as envs
|
||||
from fastvideo.attention.utils.cudnn_sdpa import cudnn_sdpa_bshd, cudnn_sdpa_supported
|
||||
from fastvideo.attention.utils.flash_attn_default import (
|
||||
fa_version,
|
||||
flash_attn_func_compilable,
|
||||
@@ -104,7 +105,11 @@ def _nvfp4_quantize_for_fa4_impl(tensor_4d: torch.Tensor) -> tuple[torch.Tensor,
|
||||
tile_m = 128
|
||||
seqlen_padded = (seqlen + tile_m - 1) // tile_m * tile_m
|
||||
if seqlen_padded != seqlen:
|
||||
tensor_4d = F.pad(tensor_4d, (0, 0, 0, 0, 0, seqlen_padded - seqlen))
|
||||
# F.pad fills the whole padded tensor before copying; only the tail rows need zeros.
|
||||
padded = tensor_4d.new_empty((batch, seqlen_padded, nheads, headdim))
|
||||
padded[:, :seqlen] = tensor_4d
|
||||
padded[:, seqlen:].zero_()
|
||||
tensor_4d = padded
|
||||
|
||||
# Quantize with nheads squashed into K dimension so M=batch*seqlen (divisible by 128)
|
||||
# and K=nheads*headdim. This ensures 128-row SF tiles align with seqlen boundaries.
|
||||
@@ -286,6 +291,11 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
# identical to before the knob; kwargs win over the
|
||||
# FASTVIDEO_FA4_PV_MODE env bridge (same pattern as nvfp4_fa4 above).
|
||||
self.fa4_pv_mode = _resolve_fa4_pv_mode(extra_impl_args)
|
||||
# Explicit impl argument (set by the H3 model) so other models cannot inherit the cuDNN route.
|
||||
self.cudnn_sdpa = bool(extra_impl_args.get("cudnn_sdpa", False))
|
||||
if self.cudnn_sdpa and self.nvfp4_fa4:
|
||||
raise ValueError("cudnn_sdpa and nvfp4_fa4 select different Q/K precision; enable only one of "
|
||||
"FASTVIDEO_MINIMAX_H3_DENSE_ATTN=cudnn and FASTVIDEO_NVFP4_FA4=1")
|
||||
if self.nvfp4_fa4:
|
||||
cap = torch.cuda.get_device_capability()
|
||||
assert cap in [(10, 0), (10, 3)], (f"NVFP4 FA4 requires Blackwell (sm100a/sm103a), got sm{cap[0]}{cap[1]}")
|
||||
@@ -480,6 +490,10 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
elif self.nvfp4_fa4:
|
||||
output = self._forward_nvfp4(query, key, value)
|
||||
|
||||
elif self.cudnn_sdpa and cudnn_sdpa_supported(query, key, value):
|
||||
# The ATen cuDNN op silently misreads contiguous BSHD; the helper passes BHSD strided views.
|
||||
output = cudnn_sdpa_bshd(query, key, value, scale=self.softmax_scale, causal=self.causal)
|
||||
|
||||
elif (self.fa4_packed_varlen and fa_version == "4" and not torch.is_grad_enabled() and query.shape[0] == 1
|
||||
and query.shape[1] == key.shape[1] == value.shape[1]):
|
||||
# FA4's packed-varlen entry point is materially faster for H3's
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""cuDNN fused scaled-dot-product attention for dense ``[B, S, H, D]`` tensors.
|
||||
|
||||
Drop-in alternative to ``flash_attn_func`` for the no-mask, inference-only dense
|
||||
self-attention call. It keeps the same precision contract (bf16/fp16 operands,
|
||||
fp32 softmax/accumulation, bf16/fp16 output) but uses cuDNN's attention engine,
|
||||
which has a different reduction order than FlashAttention.
|
||||
|
||||
The ATen entry point has no shape metadata: it always reads its inputs as
|
||||
``[B, H, S, D]``. Handing it the caller's contiguous ``[B, S, H, D]`` tensors is
|
||||
silently misread -- when ``H`` and ``S`` swap into the same output shape the
|
||||
call even returns a correctly-shaped tensor -- so every call here goes through
|
||||
``transpose(1, 2)`` to present a BHSD *strided view*, exactly like the
|
||||
documented PyTorch usage. Never pass contiguous BSHD to the op directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
# cuDNN's fused attention kernels take fp16/bf16 only; the H3 DiT runs bf16.
|
||||
_CUDNN_SDPA_DTYPES = (torch.float16, torch.bfloat16)
|
||||
# cuDNN 9.x fused SDPA supports head sizes up to 128.
|
||||
_CUDNN_SDPA_MAX_HEAD_DIM = 128
|
||||
|
||||
# Older torch builds may not register the op.
|
||||
_ATEN_CUDNN_SDPA = getattr(torch.ops.aten, "_scaled_dot_product_cudnn_attention", None)
|
||||
if _ATEN_CUDNN_SDPA is not None:
|
||||
_ATEN_CUDNN_SDPA = _ATEN_CUDNN_SDPA.default
|
||||
|
||||
|
||||
def cudnn_sdpa_supported(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
) -> bool:
|
||||
"""Whether the dense cuDNN route can take this Q/K/V triple unchanged.
|
||||
|
||||
Grad-enabled calls stay on the training path, and calls inside a compiled
|
||||
region fall through to the backend's traceable FlashAttention route so the
|
||||
opt-in never changes what ``torch.compile`` can capture.
|
||||
"""
|
||||
if _ATEN_CUDNN_SDPA is None:
|
||||
return False
|
||||
if torch.is_grad_enabled() or torch.compiler.is_compiling():
|
||||
return False
|
||||
if not (query.is_cuda and key.is_cuda and value.is_cuda):
|
||||
return False
|
||||
if query.dtype not in _CUDNN_SDPA_DTYPES or key.dtype != query.dtype or value.dtype != query.dtype:
|
||||
return False
|
||||
if query.ndim != 4 or key.shape != query.shape or value.shape != query.shape:
|
||||
return False
|
||||
if query.shape[-1] > _CUDNN_SDPA_MAX_HEAD_DIM:
|
||||
return False
|
||||
# The BHSD view below only describes the intended (B, S, H, D) layout when
|
||||
# the caller's heads are head-dim contiguous; require that explicitly.
|
||||
return query.is_contiguous() and key.is_contiguous() and value.is_contiguous()
|
||||
|
||||
|
||||
def cudnn_sdpa_bshd(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
*,
|
||||
scale: float,
|
||||
causal: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Dense scaled-dot-product attention on contiguous ``[B, S, H, D]`` inputs."""
|
||||
if _ATEN_CUDNN_SDPA is None:
|
||||
raise RuntimeError("torch build does not provide aten::_scaled_dot_product_cudnn_attention")
|
||||
output = _ATEN_CUDNN_SDPA(
|
||||
query.transpose(1, 2),
|
||||
key.transpose(1, 2),
|
||||
value.transpose(1, 2),
|
||||
None,
|
||||
False,
|
||||
0.0,
|
||||
causal,
|
||||
False,
|
||||
scale=scale,
|
||||
)[0]
|
||||
return output.transpose(1, 2).contiguous()
|
||||
@@ -1041,10 +1041,16 @@ class VideoGenerator:
|
||||
# `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
|
||||
]
|
||||
src_frames = output_batch.output
|
||||
if (isinstance(src_frames, torch.Tensor) and src_frames.dim() == 5 and src_frames.dtype == torch.uint8
|
||||
and src_frames.shape[0] == 1 and src_frames.shape[2] > 1):
|
||||
frames = list(src_frames[0].permute(1, 2, 3, 0).contiguous().cpu().numpy())
|
||||
else:
|
||||
vid_u8 = _quantize_video_frames_to_uint8(src_frames)
|
||||
frames = [
|
||||
torchvision.utils.make_grid(x, nrow=6).permute(1, 2, 0).squeeze(-1).contiguous().numpy()
|
||||
for x in vid_u8
|
||||
]
|
||||
postprocess_time = time.perf_counter() - postprocess_start
|
||||
logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time)
|
||||
if logging_info is not None:
|
||||
|
||||
+24
-1
@@ -355,6 +355,13 @@ FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN = EnvBool(
|
||||
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_MINIMAX_H3_DENSE_ATTN = EnvChoice(
|
||||
"flash",
|
||||
choices=("flash", "cudnn"),
|
||||
category="attention",
|
||||
doc="Kernel for MiniMax-H3 dense DiT self-attention: flash keeps the selected FlashAttention backend; with "
|
||||
"FLASH_ATTN, cudnn sends the no-mask, no-grad dense call through cuDNN's fused SDPA. cuDNN has a different "
|
||||
"reduction order, so this is an inference-only opt-in.")
|
||||
FASTVIDEO_VSA_PTX_BLACKWELL = EnvBool(
|
||||
False,
|
||||
category="attention",
|
||||
@@ -683,7 +690,8 @@ FASTVIDEO_LORA_SIDE_BRANCH = EnvBool(False,
|
||||
"weight always use the side branch.")
|
||||
FASTVIDEO_NVFP4_MM_BACKEND = EnvStr("auto",
|
||||
category="performance",
|
||||
doc="FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass.")
|
||||
doc="FlashInfer mm_fp4 backend for NVFP4 linears: auto (cudnn for MiniMax-H3 DiT "
|
||||
"profiles on sm_100), cutlass or cudnn.")
|
||||
FASTVIDEO_NVFP4_MM_ROW_SPLIT = EnvChoice(
|
||||
"off",
|
||||
choices=("off", "table", "auto"),
|
||||
@@ -774,6 +782,21 @@ 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_COMPILE_BLOCKS = EnvBool(False,
|
||||
category="performance",
|
||||
doc="Run the MiniMax-H3 video VAE decoder transformer blocks through one "
|
||||
"regional torch.compile graph: Inductor fuses the fp32 residual scale/add "
|
||||
"pairs, the RMSNorm+cast and the SwiGLU silu*mul.")
|
||||
FASTVIDEO_H3_VAE_UINT8_PIXELS = EnvBool(True,
|
||||
category="performance",
|
||||
doc="Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per "
|
||||
"chunk and stream uint8 into the host buffer; frames are bitwise identical "
|
||||
"to the fp32 path. 0 streams fp32.")
|
||||
FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH = EnvInt(1,
|
||||
category="performance",
|
||||
doc="Spatial tiles per MiniMax-H3 video VAE encoder call (reference and "
|
||||
"keyframe encodes); 1 encodes per tile. Batched tiles differ from "
|
||||
"per-tile encoding only in reduction order.")
|
||||
FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT = EnvBool(False,
|
||||
category="performance",
|
||||
doc="Fused dequantization epilogue for the MiniMax-H3 VAE INT8 "
|
||||
|
||||
@@ -509,14 +509,27 @@ def _mm_fp4_row_split(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_
|
||||
return _mm_fp4_chunked(*args, _SPLIT_CHOICE[key])
|
||||
|
||||
|
||||
def _mm_fp4_backend() -> str:
|
||||
_H3_LAYER_PROFILES = ("h3_dit", "h3_dit_ffn", "h3_dit_vsa")
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _is_sm100() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
|
||||
|
||||
|
||||
def _mm_fp4_backend(layer_profile: str | None = None) -> 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.
|
||||
73k rows on an RTX PRO 6000); short sequences are unaffected. On sm_100,
|
||||
``auto`` resolves to ``cudnn`` for the MiniMax-H3 DiT profiles, where it is
|
||||
bitwise identical to the default kernel and faster at H3 shapes.
|
||||
"""
|
||||
return envs.FASTVIDEO_NVFP4_MM_BACKEND.get()
|
||||
backend = envs.FASTVIDEO_NVFP4_MM_BACKEND.get()
|
||||
if backend == "auto" and layer_profile in _H3_LAYER_PROFILES and _is_sm100():
|
||||
return "cudnn"
|
||||
return backend
|
||||
|
||||
|
||||
def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor:
|
||||
@@ -552,6 +565,7 @@ 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
|
||||
_layer_profile: str | None = None
|
||||
|
||||
def __init__(self, layer_prefix: str = ""):
|
||||
super().__init__()
|
||||
@@ -708,7 +722,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
|
||||
alpha,
|
||||
torch.bfloat16,
|
||||
None,
|
||||
backend=_mm_fp4_backend(),
|
||||
backend=_mm_fp4_backend(self._layer_profile),
|
||||
)
|
||||
|
||||
if bias is not None:
|
||||
@@ -804,6 +818,7 @@ class NVFP4Config(QuantizationConfig):
|
||||
if tagged:
|
||||
method = NVFP4QuantizeMethod(layer_prefix=prefix)
|
||||
method._retain_original_weights = self.retain_original_weights
|
||||
method._layer_profile = self.layer_profile
|
||||
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):
|
||||
|
||||
@@ -222,6 +222,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
prefix: str,
|
||||
fuse_qknorm_rope: bool = False,
|
||||
fa4_packed_varlen: bool = False,
|
||||
cudnn_sdpa: bool = False,
|
||||
exact_rope: bool = False,
|
||||
mk_qkv: bool = False,
|
||||
vsa_fp4: dict[str, Any] | None = None,
|
||||
@@ -280,6 +281,7 @@ class MiniMaxH3Attention(nn.Module):
|
||||
supported_attention_backends=supported_attention_backends,
|
||||
prefix=prefix,
|
||||
fa4_packed_varlen=fa4_packed_varlen,
|
||||
cudnn_sdpa=cudnn_sdpa,
|
||||
)
|
||||
# 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.
|
||||
@@ -692,6 +694,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
fuse_qknorm_rope: bool = False,
|
||||
fuse_swiglu: bool = False,
|
||||
fa4_packed_varlen: bool = False,
|
||||
cudnn_sdpa: bool = False,
|
||||
exact_kernels: frozenset[str] = frozenset(),
|
||||
vsa_fp4: dict[str, Any] | None = None,
|
||||
mk_tier: bool = False,
|
||||
@@ -708,6 +711,7 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
prefix=f"{prefix}.attn",
|
||||
fuse_qknorm_rope=fuse_qknorm_rope,
|
||||
fa4_packed_varlen=fa4_packed_varlen,
|
||||
cudnn_sdpa=cudnn_sdpa,
|
||||
exact_rope="rope" in exact_kernels,
|
||||
vsa_fp4=vsa_fp4,
|
||||
mk_qkv=mk_tier,
|
||||
@@ -990,6 +994,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
|
||||
fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions,
|
||||
fuse_swiglu="swiglu" in self.enabled_fusions,
|
||||
fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN.get(),
|
||||
cudnn_sdpa=envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "cudnn",
|
||||
exact_kernels=self.enabled_exact_kernels,
|
||||
vsa_fp4=config.vsa_fp4,
|
||||
mk_tier=self.use_mk_kernels,
|
||||
|
||||
@@ -482,6 +482,15 @@ class MiniMaxH3VideoTransformerBlock(nn.Module):
|
||||
return hidden_states + self.ff(norm_hidden_states) * self.scale2
|
||||
|
||||
|
||||
def _run_decoder_block(
|
||||
block: nn.Module,
|
||||
hidden_states: torch.Tensor,
|
||||
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
# Takes the block as an argument so one compiled graph serves every block (parameters are graph inputs).
|
||||
return block(hidden_states, rotary_emb)
|
||||
|
||||
|
||||
class MiniMaxH3VideoViTDecoder3d(nn.Module):
|
||||
|
||||
def __init__(
|
||||
@@ -522,6 +531,25 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module):
|
||||
self.gradient_checkpointing = False
|
||||
self.fused_blocks_forward: Callable[[nn.ModuleList, torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
|
||||
torch.Tensor | None] | None = None
|
||||
self._compiled_block_forward: Callable[[nn.Module, torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
|
||||
torch.Tensor] | None = None
|
||||
if envs.FASTVIDEO_H3_VAE_COMPILE_BLOCKS.get():
|
||||
self.install_compiled_blocks()
|
||||
|
||||
def install_compiled_blocks(self) -> int:
|
||||
"""Compile the shared decoder-block region; returns the number of blocks it now serves (0 if installed)."""
|
||||
if self._compiled_block_forward is not None:
|
||||
return 0
|
||||
self._compiled_block_forward = torch.compile(_run_decoder_block, dynamic=False)
|
||||
return len(self.transformer_blocks)
|
||||
|
||||
def set_compiled_blocks(self, enabled: bool) -> int:
|
||||
"""Switch between the compiled region and the eager loop; returns the number of blocks switched."""
|
||||
if enabled:
|
||||
return self.install_compiled_blocks()
|
||||
switched = self._compiled_block_forward is not None
|
||||
self._compiled_block_forward = None
|
||||
return len(self.transformer_blocks) if switched else 0
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
"""Decode one latent spatial input through the H3 video transformer."""
|
||||
@@ -551,9 +579,12 @@ class MiniMaxH3VideoViTDecoder3d(nn.Module):
|
||||
rotary_emb) if self.fused_blocks_forward is not None else None
|
||||
if fused is not None:
|
||||
hidden_states = fused
|
||||
compiled_block_forward = self._compiled_block_forward
|
||||
for block in self.transformer_blocks if fused is None else ():
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
hidden_states = checkpoint(block, hidden_states, rotary_emb, use_reentrant=False)
|
||||
elif compiled_block_forward is not None:
|
||||
hidden_states = compiled_block_forward(block, hidden_states, rotary_emb)
|
||||
else:
|
||||
hidden_states = block(hidden_states, rotary_emb)
|
||||
|
||||
@@ -842,8 +873,21 @@ class AutoencoderKLMiniMaxH3(nn.Module):
|
||||
self.tile_sample_min_width,
|
||||
self.tile_sample_min_overlap_width,
|
||||
)
|
||||
per_call = envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.get()
|
||||
rows = []
|
||||
for y_position, y_length in zip(y_indices, y_lengths):
|
||||
if per_call > 1 and len(set(y_lengths)) == 1 and len(set(x_lengths)) == 1:
|
||||
tiles = [
|
||||
x[..., y_position:y_position + y_length, x_position:x_position + x_length]
|
||||
for y_position, y_length in zip(y_indices, y_lengths)
|
||||
for x_position, x_length in zip(x_indices, x_lengths)
|
||||
]
|
||||
encoded: list[torch.Tensor] = []
|
||||
for start in range(0, len(tiles), per_call):
|
||||
batch = torch.cat(tiles[start:start + per_call], dim=0)
|
||||
encoded.extend(self.quant_conv(self.encoder(batch)).split(x.shape[0], dim=0))
|
||||
columns = len(x_indices)
|
||||
rows = [encoded[index:index + columns] for index in range(0, len(encoded), columns)]
|
||||
for y_position, y_length in zip(y_indices, y_lengths) if not rows else ():
|
||||
row = []
|
||||
for x_position, x_length in zip(x_indices, x_lengths):
|
||||
tile = x[..., y_position:y_position + y_length, x_position:x_position + x_length]
|
||||
@@ -1077,6 +1121,8 @@ class AutoencoderKLMiniMaxH3(nn.Module):
|
||||
before releasing the buffer.
|
||||
"""
|
||||
target = output[:, :, frame_start:frame_start + pixels.shape[2]]
|
||||
if output.dtype == torch.uint8:
|
||||
pixels = (pixels * 255).clamp_(0, 255).to(torch.uint8)
|
||||
if pixels.device.type == "cuda":
|
||||
pixels = pixels.contiguous()
|
||||
for batch_index in range(pixels.shape[0]):
|
||||
@@ -1195,11 +1241,17 @@ class AutoencoderKLMiniMaxH3(nn.Module):
|
||||
return DecoderOutput(sample=decoded)
|
||||
|
||||
def decode_to_pixels(self, z: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
|
||||
"""Stream decoded ``[0, 1]`` FP32 pixels into a caller-owned CPU buffer."""
|
||||
"""Stream decoded ``[0, 1]`` pixels into a caller-owned CPU buffer.
|
||||
|
||||
``output`` is float32, or uint8 when the caller asks for the quantized
|
||||
stream (``FASTVIDEO_H3_VAE_UINT8_PIXELS``); uint8 quantizes each chunk on
|
||||
the GPU with the fp32 host path's exact mul/clamp/truncate arithmetic.
|
||||
"""
|
||||
expected_shape = self.decoded_pixel_shape(z.shape)
|
||||
if output.device.type != "cpu" or output.dtype != torch.float32 or tuple(output.shape) != expected_shape:
|
||||
if (output.device.type != "cpu" or output.dtype not in (torch.float32, torch.uint8)
|
||||
or tuple(output.shape) != expected_shape):
|
||||
raise ValueError(
|
||||
"`output` must be a CPU float32 tensor with shape "
|
||||
"`output` must be a CPU float32 or uint8 tensor with shape "
|
||||
f"{expected_shape}, got device={output.device}, dtype={output.dtype}, shape={tuple(output.shape)}.")
|
||||
try:
|
||||
if self.use_slicing and z.shape[0] > 1:
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
import fastvideo.envs as envs
|
||||
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
|
||||
@@ -68,10 +69,14 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
|
||||
def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None:
|
||||
super().__init__()
|
||||
self.vae = vae
|
||||
self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None
|
||||
self._output_buffer: tuple[tuple[tuple[int, ...], bool, torch.dtype], torch.Tensor] | None = None
|
||||
self._unified_memory: bool | None = None
|
||||
|
||||
def _pixel_buffer(self, shape: tuple[int, ...], pin_memory: bool, device: torch.device) -> torch.Tensor:
|
||||
def _pixel_buffer(self,
|
||||
shape: tuple[int, ...],
|
||||
pin_memory: bool,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype = torch.float32) -> torch.Tensor:
|
||||
"""The CPU FP32 pixel buffer, reused across requests of the same geometry on unified memory.
|
||||
|
||||
The stage returns a uint8 copy, so the buffer never escapes. Reuse matters on
|
||||
@@ -84,11 +89,11 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
|
||||
if self._unified_memory is None:
|
||||
self._unified_memory = current_platform.has_unified_memory(device.index or 0)
|
||||
if not self._unified_memory:
|
||||
return allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory)
|
||||
key = (tuple(shape), bool(pin_memory))
|
||||
return allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory)
|
||||
key = (tuple(shape), bool(pin_memory), dtype)
|
||||
if self._output_buffer is None or self._output_buffer[0] != key:
|
||||
self._output_buffer = None
|
||||
buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=torch.float32, pin_memory=pin_memory)
|
||||
buffer = allocate_cpu_tensor_with_pin_fallback(shape, dtype=dtype, pin_memory=pin_memory)
|
||||
if not buffer.is_pinned():
|
||||
buffer.zero_()
|
||||
self._output_buffer = (key, buffer)
|
||||
@@ -170,8 +175,10 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
|
||||
|
||||
output = None
|
||||
if is_output_rank:
|
||||
pixel_dtype = torch.uint8 if envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.get(
|
||||
) and not parallel else torch.float32
|
||||
output = self._pixel_buffer(tuple(self.vae.decoded_pixel_shape(latents.shape)),
|
||||
fastvideo_args.pin_cpu_memory, device)
|
||||
fastvideo_args.pin_cpu_memory, device, pixel_dtype)
|
||||
# Attribute the streamed decoder computation while retaining
|
||||
# per-chunk device-to-host transfer and pinned-buffer reuse.
|
||||
# fp16 for the release decoders; NVFP4 decoders were trained under bf16.
|
||||
@@ -194,6 +201,8 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
|
||||
# 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)
|
||||
elif self._output_buffer is not None and output is self._output_buffer[1]:
|
||||
output = output.clone()
|
||||
batch.output = output if is_output_rank else placeholder
|
||||
return batch
|
||||
finally:
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax-H3's opt-in cuDNN SDPA route for dense self-attention.
|
||||
|
||||
The ATen cuDNN attention op carries no shape metadata and always reads its
|
||||
inputs as ``[B, H, S, D]``; a contiguous ``[B, S, H, D]`` tensor is silently
|
||||
misread. These tests pin the BHSD-strided-view contract and the opt-in flag.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo import envs
|
||||
from fastvideo.attention.utils.cudnn_sdpa import cudnn_sdpa_bshd, cudnn_sdpa_supported
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuDNN SDPA requires CUDA")
|
||||
|
||||
|
||||
def _qkv(seq_len: int, *, heads: int = 8, head_dim: int = 64, dtype=torch.bfloat16):
|
||||
generator = torch.Generator(device="cuda").manual_seed(0)
|
||||
return [
|
||||
torch.randn(1, seq_len, heads, head_dim, device="cuda", dtype=dtype, generator=generator) for _ in range(3)
|
||||
]
|
||||
|
||||
|
||||
def _rel_l2(actual: torch.Tensor, expected: torch.Tensor) -> float:
|
||||
return ((actual.float() - expected.float()).norm() / expected.float().norm()).item()
|
||||
|
||||
|
||||
def test_cudnn_route_presents_bhsd_strided_views():
|
||||
"""Contiguous BSHD is silently misread; the BHSD view is the supported form."""
|
||||
query, key, value = _qkv(257)
|
||||
scale = 1.0 / math.sqrt(query.shape[-1])
|
||||
with torch.no_grad():
|
||||
output = cudnn_sdpa_bshd(query, key, value, scale=scale)
|
||||
# The trap: hand the op the caller's contiguous BSHD tensors directly.
|
||||
misread = torch.ops.aten._scaled_dot_product_cudnn_attention.default(
|
||||
query, key, value, None, False, 0.0, False, False, scale=scale)[0]
|
||||
|
||||
# Same shape survives the misread, which is why it is silent.
|
||||
assert misread.shape == output.shape
|
||||
assert _rel_l2(misread, output) > 1.0
|
||||
assert output.is_contiguous()
|
||||
assert output.dtype is torch.bfloat16
|
||||
assert torch.isfinite(output).all()
|
||||
|
||||
|
||||
def test_cudnn_route_is_bit_identical_to_cudnn_dispatched_sdpa():
|
||||
query, key, value = _qkv(2048)
|
||||
scale = 1.0 / math.sqrt(query.shape[-1])
|
||||
previous = torch.backends.cuda.cudnn_sdp_enabled()
|
||||
torch.backends.cuda.enable_cudnn_sdp(True)
|
||||
try:
|
||||
with torch.no_grad():
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), scale=scale).transpose(1, 2)
|
||||
finally:
|
||||
torch.backends.cuda.enable_cudnn_sdp(previous)
|
||||
|
||||
with torch.no_grad():
|
||||
output = cudnn_sdpa_bshd(query, key, value, scale=scale)
|
||||
assert torch.equal(output, reference.contiguous())
|
||||
|
||||
|
||||
def test_cudnn_route_tracks_eager_math_sdpa():
|
||||
"""Same precision contract (bf16 in, fp32 softmax, bf16 out), different reduction order."""
|
||||
query, key, value = _qkv(257)
|
||||
scale = 1.0 / math.sqrt(query.shape[-1])
|
||||
with torch.no_grad():
|
||||
output = cudnn_sdpa_bshd(query, key, value, scale=scale)
|
||||
with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), scale=scale).transpose(1, 2)
|
||||
assert _rel_l2(output, reference) < 5e-3
|
||||
assert (output.float() - reference.float()).abs().max().item() < 0.02
|
||||
|
||||
|
||||
def test_cudnn_route_guard_rejects_training_unsupported_dtypes_and_shapes():
|
||||
query, key, value = _qkv(64, heads=4, head_dim=128)
|
||||
with torch.no_grad():
|
||||
assert cudnn_sdpa_supported(query, key, value)
|
||||
assert not cudnn_sdpa_supported(query.float(), key.float(), value.float())
|
||||
assert not cudnn_sdpa_supported(query.transpose(1, 2).transpose(1, 2).transpose(1, 2), key, value)
|
||||
wide, _, _ = _qkv(64, heads=2, head_dim=256)
|
||||
assert not cudnn_sdpa_supported(wide, wide, wide)
|
||||
# A grad-enabled training call stays on the FlashAttention path.
|
||||
assert torch.is_grad_enabled()
|
||||
assert not cudnn_sdpa_supported(query, key, value)
|
||||
|
||||
|
||||
def test_dense_attention_env_default_is_flash_and_accepts_cudnn():
|
||||
assert envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "flash"
|
||||
with envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.override("cudnn"):
|
||||
assert envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get() == "cudnn"
|
||||
with pytest.raises(ValueError):
|
||||
with envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.override("cutlass"):
|
||||
envs.FASTVIDEO_MINIMAX_H3_DENSE_ATTN.get()
|
||||
|
||||
|
||||
def test_flash_attn_impl_routes_dense_calls_to_cudnn():
|
||||
pytest.importorskip("flash_attn")
|
||||
from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
|
||||
|
||||
query, key, value = _qkv(257)
|
||||
scale = 1.0 / math.sqrt(query.shape[-1])
|
||||
impl = FlashAttentionImpl(num_heads=query.shape[2],
|
||||
head_size=query.shape[3],
|
||||
causal=False,
|
||||
softmax_scale=scale,
|
||||
cudnn_sdpa=True)
|
||||
assert impl.cudnn_sdpa
|
||||
with torch.no_grad():
|
||||
routed = impl.forward(query, key, value, None)
|
||||
expected = cudnn_sdpa_bshd(query, key, value, scale=scale)
|
||||
assert torch.equal(routed, expected)
|
||||
|
||||
# The cuDNN opt-in wins over FA4 packed-varlen when both are configured.
|
||||
packed = FlashAttentionImpl(num_heads=query.shape[2],
|
||||
head_size=query.shape[3],
|
||||
causal=False,
|
||||
softmax_scale=scale,
|
||||
cudnn_sdpa=True,
|
||||
fa4_packed_varlen=True)
|
||||
with torch.no_grad():
|
||||
assert torch.equal(packed.forward(query, key, value, None), expected)
|
||||
|
||||
# Default construction keeps the existing FlashAttention behavior.
|
||||
default = FlashAttentionImpl(num_heads=query.shape[2],
|
||||
head_size=query.shape[3],
|
||||
causal=False,
|
||||
softmax_scale=scale)
|
||||
assert not default.cudnn_sdpa
|
||||
|
||||
|
||||
def test_flash_attn_impl_rejects_cudnn_and_nvfp4_together():
|
||||
pytest.importorskip("flash_attn")
|
||||
from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
|
||||
|
||||
with pytest.raises(ValueError, match="cudnn_sdpa and nvfp4_fa4"):
|
||||
FlashAttentionImpl(num_heads=4, head_size=64, causal=False, softmax_scale=0.125, cudnn_sdpa=True,
|
||||
nvfp4_fa4=True)
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 quantization needs CUDA")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seqlen", [77, 300])
|
||||
def test_fa4_quantize_pads_like_zero_padding(seqlen: int) -> None:
|
||||
flash_attn = pytest.importorskip("fastvideo.attention.backends.flash_attn")
|
||||
try:
|
||||
flash_attn._resolve_fa4_quant_ops()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
pytest.skip(f"NVFP4 quantize ops unavailable: {exc}")
|
||||
torch.manual_seed(0)
|
||||
x = torch.randn(1, seqlen, 4, 128, device="cuda", dtype=torch.bfloat16)
|
||||
padded = F.pad(x, (0, 0, 0, 0, 0, (-seqlen) % 128))
|
||||
|
||||
fp4, sf = flash_attn._nvfp4_quantize_for_fa4_impl(x)
|
||||
ref_fp4, ref_sf = flash_attn._nvfp4_quantize_for_fa4_impl(padded)
|
||||
|
||||
assert torch.equal(fp4.view(torch.uint8), ref_fp4.view(torch.uint8))
|
||||
assert torch.equal(sf, ref_sf)
|
||||
@@ -0,0 +1,25 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import fastvideo.envs as envs
|
||||
import pytest
|
||||
|
||||
from fastvideo.layers.quantization import nvfp4_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("profile", "sm100", "expected"), [
|
||||
("h3_dit", True, "cudnn"),
|
||||
("h3_dit_ffn", True, "cudnn"),
|
||||
("h3_dit_vsa", True, "cudnn"),
|
||||
("refine", True, "auto"),
|
||||
(None, True, "auto"),
|
||||
("h3_dit", False, "auto"),
|
||||
])
|
||||
def test_auto_backend_resolution(monkeypatch, profile, sm100, expected, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_MM_BACKEND.override(None))
|
||||
monkeypatch.setattr(nvfp4_config, "_is_sm100", lambda: sm100)
|
||||
assert nvfp4_config._mm_fp4_backend(profile) == expected
|
||||
|
||||
|
||||
def test_explicit_backend_wins(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_MM_BACKEND.override("cutlass"))
|
||||
monkeypatch.setattr(nvfp4_config, "_is_sm100", lambda: True)
|
||||
assert nvfp4_config._mm_fp4_backend("h3_dit") == "cutlass"
|
||||
@@ -4,6 +4,7 @@ from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
|
||||
from fastvideo.pipelines.basic.minimax_h3 import encoder_split
|
||||
@@ -74,7 +75,8 @@ def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
|
||||
assert rows[0].shape == (7 * 4 * 4, 4)
|
||||
|
||||
|
||||
def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
|
||||
def test_decode_stage_uses_cpu_output_buffer(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(False))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
latents = torch.randn(latent_shape)
|
||||
rows = patchify_video_latents(latents, (1, 1, 1))
|
||||
@@ -128,7 +130,8 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
|
||||
assert torch.equal(result.output, expected)
|
||||
|
||||
|
||||
def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
|
||||
def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(False))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
latents = torch.randn(latent_shape)
|
||||
rows = patchify_video_latents(latents, (1, 1, 1))
|
||||
@@ -177,10 +180,52 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
|
||||
}
|
||||
|
||||
|
||||
def test_decode_stage_streams_uint8_by_default(monkeypatch, env_overrides) -> None:
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
rows = patchify_video_latents(torch.randn(latent_shape), (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):
|
||||
observed["dtype"] = output.dtype
|
||||
output.fill_(63)
|
||||
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(None))
|
||||
monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
|
||||
result = MiniMaxH3VideoDecodingStage(VAE()).forward(
|
||||
batch,
|
||||
SimpleNamespace(
|
||||
output_type="pil",
|
||||
pin_cpu_memory=False,
|
||||
vae_cpu_offload=False,
|
||||
vae_parallel_decode=False,
|
||||
pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))),
|
||||
),
|
||||
)
|
||||
|
||||
assert observed["dtype"] == torch.uint8
|
||||
assert torch.equal(result.output, torch.full((1, 3, 5, 16, 16), 63, dtype=torch.uint8))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("uint8_pixels", [False, True])
|
||||
@pytest.mark.parametrize("unified_memory", [True, False])
|
||||
def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory) -> None:
|
||||
def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory,
|
||||
uint8_pixels, env_overrides) -> None:
|
||||
"""On unified memory, same-geometry requests reuse one pre-touched CPU buffer; returned
|
||||
frames never alias it. Discrete GPUs allocate per request and hold nothing between them."""
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(uint8_pixels))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1))
|
||||
allocations = []
|
||||
@@ -200,12 +245,13 @@ def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch
|
||||
return (1, 3, 5, 16, 16)
|
||||
|
||||
def decode_to_pixels(self, decoded_latents, output):
|
||||
buffers.append((output, bool(torch.isnan(output).any())))
|
||||
output.fill_(VAE.fill)
|
||||
poisoned = bool(torch.isnan(output).any()) if output.is_floating_point() else bool((output == 255).any())
|
||||
buffers.append((output, poisoned))
|
||||
output.fill_(VAE.fill if output.is_floating_point() else round(VAE.fill * 255 - 0.5))
|
||||
|
||||
def fake_allocate(size, *, dtype=None, pin_memory=False):
|
||||
allocations.append(size)
|
||||
return torch.full(size, float("nan"), dtype=dtype)
|
||||
return torch.full(size, float("nan") if dtype.is_floating_point else 255, 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)
|
||||
@@ -235,6 +281,7 @@ def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch
|
||||
assert buffers[0][0] is not buffers[1][0]
|
||||
assert stage._output_buffer is None
|
||||
assert torch.equal(first, torch.full_like(first, 63)) and torch.equal(second, torch.full_like(second, 127))
|
||||
assert first.data_ptr() != second.data_ptr()
|
||||
|
||||
|
||||
def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None:
|
||||
|
||||
@@ -11,6 +11,7 @@ from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
@@ -20,6 +21,10 @@ from fastvideo.models.vaes.minimax_h3_audio import (
|
||||
MiniMaxH3AudioBigVGANDecoder,
|
||||
MiniMaxH3AudioVAE,
|
||||
)
|
||||
from fastvideo.configs.models.vaes.minimax_h3_video import (
|
||||
MiniMaxH3VideoVAEArchConfig,
|
||||
MiniMaxH3VideoVAEConfig,
|
||||
)
|
||||
from fastvideo.models.vaes.minimax_h3_video import (
|
||||
AutoencoderKLMiniMaxH3,
|
||||
MiniMaxH3VideoAttention,
|
||||
@@ -502,3 +507,79 @@ def test_tiled_decode_and_encode_survive_cudagraph_buffer_reuse_on_cuda() -> Non
|
||||
assert not eager_vae._tile_helpers_compiled
|
||||
torch.testing.assert_close(decoded_first, eager_vae.decode(z).sample, atol=2e-4, rtol=2e-4)
|
||||
torch.testing.assert_close(encoded_first, eager_vae.encode(pixels).latent_dist.parameters, atol=2e-4, rtol=2e-4)
|
||||
|
||||
|
||||
def test_regional_block_compile_installs_one_graph_and_toggles() -> None:
|
||||
"""FASTVIDEO_H3_VAE_COMPILE_BLOCKS compiles one block region shared by every block."""
|
||||
decoder = _empty_typed_module(MiniMaxH3VideoViTDecoder3d)
|
||||
decoder.transformer_blocks = nn.ModuleList([nn.Identity(), nn.Identity(), nn.Identity()])
|
||||
decoder._compiled_block_forward = None
|
||||
compiled_block = Mock(name="compiled_block")
|
||||
|
||||
with patch("fastvideo.models.vaes.minimax_h3_video.torch.compile", return_value=compiled_block) as compile_mock:
|
||||
assert decoder.install_compiled_blocks() == 3
|
||||
compile_mock.assert_called_once()
|
||||
assert compile_mock.call_args.args[0].__name__ == "_run_decoder_block"
|
||||
assert compile_mock.call_args.kwargs == {"dynamic": False}
|
||||
# One graph for the whole decoder: a second install is a no-op.
|
||||
assert decoder.install_compiled_blocks() == 0
|
||||
compile_mock.assert_called_once()
|
||||
# Disabling restores the eager loop; enabling reinstalls the same graph.
|
||||
assert decoder.set_compiled_blocks(False) == 3
|
||||
assert decoder._compiled_block_forward is None
|
||||
assert decoder.set_compiled_blocks(True) == 3
|
||||
assert decoder._compiled_block_forward is compiled_block
|
||||
assert decoder.set_compiled_blocks(False) == 3
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_decoder_forward_uses_the_compiled_block_region() -> None:
|
||||
"""The decoder routes every transformer block through the compiled region when installed."""
|
||||
torch.manual_seed(20260818)
|
||||
decoder = _tiny_h3_decoder()
|
||||
latents = torch.randn(1, 4, 12, 4, 4)
|
||||
with torch.no_grad():
|
||||
expected = decoder(latents)
|
||||
|
||||
seen: list[nn.Module] = []
|
||||
compiled_block = Mock(name="compiled_block")
|
||||
compiled_block.side_effect = lambda block, hidden, rotary: (seen.append(block), block(hidden, rotary))[1]
|
||||
decoder._compiled_block_forward = compiled_block
|
||||
|
||||
with torch.no_grad():
|
||||
actual = decoder(latents)
|
||||
|
||||
assert len(seen) == len(decoder.transformer_blocks)
|
||||
assert seen == list(decoder.transformer_blocks)
|
||||
assert torch.equal(actual, expected)
|
||||
|
||||
|
||||
def test_regional_block_compile_flag_installs_on_construction(monkeypatch, env_overrides) -> None:
|
||||
"""The env flag is read when the decoder is built, like the other VAE load-time knobs."""
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_COMPILE_BLOCKS.override(True))
|
||||
compiled_block = Mock(name="compiled_block")
|
||||
with patch("fastvideo.models.vaes.minimax_h3_video.torch.compile", return_value=compiled_block) as compile_mock:
|
||||
decoder = _tiny_h3_decoder()
|
||||
compile_mock.assert_called_once()
|
||||
assert decoder._compiled_block_forward is compiled_block
|
||||
|
||||
|
||||
def _tiny_h3_decoder() -> MiniMaxH3VideoViTDecoder3d:
|
||||
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=2,
|
||||
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,
|
||||
)
|
||||
vae = AutoencoderKLMiniMaxH3(
|
||||
MiniMaxH3VideoVAEConfig(arch_config=arch, use_tiling=False, use_temporal_tiling=False,
|
||||
use_parallel_tiling=False))
|
||||
return vae.decoder
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
from torch.testing import assert_close
|
||||
|
||||
@@ -108,6 +109,40 @@ def test_decode_to_pixels_matches_decode(latent_frames: int) -> None:
|
||||
assert_close(actual, expected, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="needs CUDA"))])
|
||||
@torch.inference_mode()
|
||||
def test_decode_to_pixels_uint8_matches_fp32_quantization(device: str) -> None:
|
||||
torch.manual_seed(20260810)
|
||||
vae = _tiny_vae().to(device)
|
||||
latents = torch.randn(1, 4, 12, 4, 4, device=device)
|
||||
fp32 = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
|
||||
streamed = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.uint8)
|
||||
|
||||
vae.decode_to_pixels(latents, fp32)
|
||||
vae.decode_to_pixels(latents, streamed)
|
||||
|
||||
assert torch.equal(streamed, fp32.mul_(255).clamp_(0, 255).to(torch.uint8))
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_batched_tile_encode_matches_per_tile(monkeypatch, env_overrides) -> None:
|
||||
torch.manual_seed(20260810)
|
||||
vae = _tiny_vae()
|
||||
vae.enable_tiling(tile_sample_min_height=16,
|
||||
tile_sample_min_width=16,
|
||||
tile_sample_min_overlap_height=8,
|
||||
tile_sample_min_overlap_width=8)
|
||||
pixels = torch.rand(1, 3, 1, 32, 32)
|
||||
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.override(1))
|
||||
per_tile = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.override(8))
|
||||
batched = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters
|
||||
|
||||
assert_close(batched, per_tile, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
# 3: one chunk with pad tokens; 6: pad hitting the intra-clip tail;
|
||||
# 12: two blended chunks without padding; 13: three chunks plus pad trim.
|
||||
@pytest.mark.parametrize("latent_frames", (3, 6, 12, 13))
|
||||
|
||||
Reference in New Issue
Block a user