Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

minimax_h3_nvfp4_fused

Bit-exact fused inference kernels for the NVFP4 MiniMax-H3 decoder blocks.

The eager NVFP4 block (NVFP4DecoderLinear inside MiniMaxH3VideoTransformerBlock) spends most of its time in small elementwise kernels around the FP4 GEMMs: bias adds, RMSNorms, casts, RoPE, SiLU, residuals and three separate quantizations of the shared q/k/v input. These Triton kernels fuse them while reproducing every rounding step of the eager path, so the decoder output is bit-identical:

  • rms_norm follows ATen's vectorized_layer_norm_kernel reduction order (128 threads, 4-wide vectors, sequential FMA per thread, warp shuffle-down tree, then a 4-warp tree), so the fp32 statistics match bit for bit;
  • NVFP4 quantization follows FlashInfer's cvt_warp_fp16_to_fp4 (approximate reciprocals, E4M3 block scales, layout_128x4 swizzle, zeroed padding rows);
  • every bf16 rounding of the op-by-op eager path is kept, and products that eager rounds before an add use mul.rn so they are never contracted into an FMA.

fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py checks each kernel and whole decodes against the eager path with torch.equal.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear

NVFP4DecoderLinear(weight: torch.Tensor, bias: torch.Tensor | None, rotation_group: int | None, compute_dtype: torch.dtype, act_scale: str = 'dynamic')

Drop-in for an nn.Linear in the H3 decoder blocks.

weight is the full-precision master (already rotated when rotation_group is set). With gradients enabled the forward is the STE; otherwise the cached packed weight is used, so call invalidate() after any weight update (the trainer does after each optimizer step).

act_scale="static" uses the persistent input_amax buffer. Set calibrating = True to run dynamically while folding each call's abs-max into it.

freeze() turns the layer inference-only: the packed weight becomes a buffer and the master weight is dropped (weight is None).

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.from_linear

from_linear(linear: nn.Linear, *, rotation_group: int | None, compute_dtype: torch.dtype, act_scale: str = 'dynamic') -> NVFP4DecoderLinear

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.from_packed

from_packed(packed: torch.Tensor, inv_scale: torch.Tensor, global_sf_w: torch.Tensor, bias: torch.Tensor | None, input_amax: torch.Tensor, *, in_features: int, rotation_group: int | None, compute_dtype: torch.dtype, act_scale: str) -> NVFP4DecoderLinear

A frozen layer built from already packed tensors (freeze()'s buffers), with no master weight.

packed holds two FP4 codes per byte, so its shape is [out_features, in_features // 2]. Nothing full-precision is allocated: the constructor's placeholder weight lives on meta.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.invalidate

invalidate() -> None

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.freeze

freeze() -> None

Pack the weight once and drop the full-precision master (inference only, irreversible).

Resident memory falls to the packed FP4 weight and its scales. Outputs are unchanged: the same packed tensors feed the eager GEMM and fused_state(). Buffers are non-persistent, so a frozen layer's state_dict() holds only bias and input_amax.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.fused_state

fused_state() -> NVFP4FusedState | None

Cached inference state for the fused decoder kernels, or None when they cannot reproduce this layer.

The fused kernels quantize with a precomputed global scale, so a dynamic (per-call abs-max) scale, calibration, and the activation rotation keep the eager path.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.clear_offload_caches

clear_offload_caches() -> None

Drop the packed and fused caches, which are plain attributes and so invisible to offload.

pinned_offload moves parameters and buffers only; these caches hold device tensors, so without this they stay on the GPU after an unload and accumulate across clips. The data-pointer keys re-create them on the next forward.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.extra_repr

extra_repr() -> str

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4DecoderLinear.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.vaes.minimax_h3_nvfp4_fused.NVFP4FusedState

NVFP4FusedState()

Everything the fused inference kernels need from one NVFP4DecoderLinear.

Each tensor is computed with the same expression as nvfp4_linear_inference, so a GEMM fed from here matches the eager linear bit for bit.

fastvideo.models.vaes.minimax_h3_nvfp4_fused.new_quantized

new_quantized(m: int, k: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]

Buffers shaped like FlashInfer's nvfp4_quantize output for an [m, k] activation.

Source
def new_quantized(m: int, k: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    """Buffers shaped like FlashInfer's ``nvfp4_quantize`` output for an [m, k] activation."""
    packed = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
    scales = torch.empty((_padded_rows(m), k // NVFP4_BLOCK), dtype=torch.uint8, device=device)
    return packed, scales

fastvideo.models.vaes.minimax_h3_nvfp4_fused.residual_norm_quantize

residual_norm_quantize(h: torch.Tensor, residual: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None, norm: tuple[torch.Tensor, float, torch.Tensor] | None, rows_per_program: int = 4) -> tuple[torch.Tensor, torch.Tensor] | None

In place h += bf16(o + bias) * scale and/or NVFP4 of bf16(rms_norm(h, gamma, eps)).

h is a contiguous fp32 [M, D] tensor; residual is (o, bias_bf16, scale_fp32); norm is (gamma_fp32, eps, global_sf). Returns the quantized activation when norm is set.

Source
def residual_norm_quantize(h: torch.Tensor, residual: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None,
                           norm: tuple[torch.Tensor, float, torch.Tensor] | None,
                           rows_per_program: int = 4) -> tuple[torch.Tensor, torch.Tensor] | None:
    """In place ``h += bf16(o + bias) * scale`` and/or NVFP4 of ``bf16(rms_norm(h, gamma, eps))``.

``h`` is a contiguous fp32 [M, D] tensor; ``residual`` is ``(o, bias_bf16, scale_fp32)``;
``norm`` is ``(gamma_fp32, eps, global_sf)``. Returns the quantized activation when ``norm`` is set.
"""
m, d = h.shape
if d % 512:
    raise ValueError(f"hidden size must be a multiple of 512, got {d}")
quantized = new_quantized(m, d, h.device) if norm is not None else (h, h)
o, bias, scale = residual if residual is not None else (h, h, h)
gamma, eps, gsf = norm if norm is not None else (h, 0.0, h)
rows = _padded_rows(m) if norm is not None else m
grid = (triton.cdiv(rows, rows_per_program), )
_residual_norm_quant_kernel[grid](h,
                                  o,
                                  bias,
                                  scale,
                                  gamma,
                                  gsf,
                                  quantized[0],
                                  quantized[1],
                                  m,
                                  rows,
                                  eps,
                                  D=d,
                                  R=rows_per_program,
                                  HAS_RESIDUAL=residual is not None,
                                  HAS_NORM=norm is not None,
                                  num_warps=4)
return quantized if norm is not None else None</code></pre></details>


fastvideo.models.vaes.minimax_h3_nvfp4_fused.quantize


quantize(x: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 8, block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]


FlashInfer-exact NVFP4 quantization (layout_128x4) of a contiguous bf16 [M, K] tensor.


Source
def quantize(x: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 8,
             block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]:
    """FlashInfer-exact NVFP4 quantization (``layout_128x4``) of a contiguous bf16 [M, K] tensor."""
    m, k = x.shape
    quantized = new_quantized(m, k, x.device)
    grid = (triton.cdiv(_padded_rows(m), rows_per_program), triton.cdiv(k, block_n))
    _quantize_kernel[grid](x, gsf, quantized[0], quantized[1], m, _padded_rows(m), K=k, R=rows_per_program, BN=block_n)
    return quantized


fastvideo.models.vaes.minimax_h3_nvfp4_fused.swiglu_quantize


swiglu_quantize(o: torch.Tensor, bias: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 2, block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]


NVFP4 of the H3 SwiGLU of a value-first packed projection o [M, 2F] (bias not yet added).


Source
def swiglu_quantize(o: torch.Tensor, bias: torch.Tensor, gsf: torch.Tensor, rows_per_program: int = 2,
                    block_n: int = 512) -> tuple[torch.Tensor, torch.Tensor]:
    """NVFP4 of the H3 SwiGLU of a value-first packed projection ``o`` [M, 2F] (bias not yet added)."""
    m, two_f = o.shape
    f = two_f // 2
    quantized = new_quantized(m, f, o.device)
    grid = (triton.cdiv(_padded_rows(m), rows_per_program), triton.cdiv(f, block_n))
    _swiglu_quantize_kernel[grid](o,
                                  bias,
                                  _silu_table(o.device.index),
                                  gsf,
                                  quantized[0],
                                  quantized[1],
                                  m,
                                  _padded_rows(m),
                                  F=f,
                                  R=rows_per_program,
                                  BN=block_n,
                                  num_warps=2)
    return quantized


fastvideo.models.vaes.minimax_h3_nvfp4_fused.qkv_epilogue


qkv_epilogue(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, biases: tuple[torch.Tensor, torch.Tensor, torch.Tensor], cos: torch.Tensor, sin: torch.Tensor, heads: int, eps: float, rows_per_program: int = 2) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]


Projected q/k/v [M, H*HD] bf16 (bias not yet added) -> roped q, k and biased v, each [M, H, HD].


cos/sin are bf16 [M, ROT] (the eager path's cos.to(query.dtype)).


Source
def qkv_epilogue(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, biases: tuple[torch.Tensor, torch.Tensor,
                                                                               torch.Tensor], cos: torch.Tensor,
                 sin: torch.Tensor, heads: int, eps: float,
                 rows_per_program: int = 2) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Projected q/k/v [M, H*HD] bf16 (bias not yet added) -> roped q, k and biased v, each [M, H, HD].

``cos``/``sin`` are bf16 &#91;M, ROT&#93; (the eager path&#x27;s ``cos.to(query.dtype)``).
&quot;&quot;&quot;
m, d = q.shape
head_dim = d // heads
outs = &#91;torch.empty((m, heads, head_dim), dtype=torch.bfloat16, device=q.device) for _ in range(3)&#93;
grid = (triton.cdiv(m, rows_per_program), 3)
_qkv_epilogue_kernel&#91;grid&#93;(q,
                           k,
                           v,
                           *biases,
                           cos,
                           sin,
                           *outs,
                           m,
                           eps,
                           H=heads,
                           HD=head_dim,
                           ROT=cos.shape&#91;-1&#93;,
                           R=rows_per_program,
                           num_warps=4)
return outs&#91;0&#93;, outs&#91;1&#93;, outs&#91;2&#93;</code></pre></details>


fastvideo.models.vaes.minimax_h3_nvfp4_fused.fused_nvfp4_blocks_forward


fused_nvfp4_blocks_forward(blocks: nn.ModuleList, hidden: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None) -> torch.Tensor | None


Run every decoder block through the fused kernels; None (caller runs the blocks) when not applicable.


Applies at inference (no grad) to fp32 residual streams on CUDA when every block's linears are NVFP4
with a precomputed activation scale. Bit-identical to the eager block loop. torch.compile treats it
as a boundary: a compiled decoder compiles the code around the block stack and runs the stack here.
Calls of at most GRAPH_MAX_ROWS tokens replay a per-shape CUDA graph
(FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH).


Source
@torch.compiler.disable
def fused_nvfp4_blocks_forward(blocks: nn.ModuleList, hidden: torch.Tensor,
                               rotary_emb: tuple[torch.Tensor, torch.Tensor] | None) -> torch.Tensor | None:
    """Run every decoder block through the fused kernels; None (caller runs the blocks) when not applicable.

Applies at inference (no grad) to fp32 residual streams on CUDA when every block&#x27;s linears are NVFP4
with a precomputed activation scale. Bit-identical to the eager block loop. ``torch.compile`` treats it
as a boundary: a compiled decoder compiles the code around the block stack and runs the stack here.
Calls of at most ``GRAPH_MAX_ROWS`` tokens replay a per-shape CUDA graph
(``FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH``).
&quot;&quot;&quot;
if (torch.compiler.is_compiling() or torch.is_grad_enabled() or not envs.FASTVIDEO_H3_VAE_NVFP4_FUSED.get()
        or rotary_emb is None or len(blocks) == 0 or hidden.dtype != torch.float32 or not hidden.is_cuda):
    return None
plans = &#91;_block_plan(block) for block in blocks&#93;
if any(plan is None for plan in plans):
    return None
if (hidden.shape&#91;0&#93; * hidden.shape&#91;1&#93; &lt;= GRAPH_MAX_ROWS and envs.FASTVIDEO_H3_VAE_NVFP4_CUDA_GRAPH.get()
        and not torch.cuda.is_current_stream_capturing()):
    return _graphed_blocks(blocks, plans, hidden, rotary_emb)
return _run_blocks(plans, hidden, rotary_emb)</code></pre></details>

Was this page helpful?