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

`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 [M, ROT] (the eager path's ``cos.to(query.dtype)``).
    """
    m, d = q.shape
    head_dim = d // heads
    outs = [torch.empty((m, heads, head_dim), dtype=torch.bfloat16, device=q.device) for _ in range(3)]
    grid = (triton.cdiv(m, rows_per_program), 3)
    _qkv_epilogue_kernel[grid](q,
                               k,
                               v,
                               *biases,
                               cos,
                               sin,
                               *outs,
                               m,
                               eps,
                               H=heads,
                               HD=head_dim,
                               ROT=cos.shape[-1],
                               R=rows_per_program,
                               num_warps=4)
    return outs[0], outs[1], outs[2]

`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'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``).
    """
    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 = [_block_plan(block) for block in blocks]
    if any(plan is None for plan in plans):
        return None
    if (hidden.shape[0] * hidden.shape[1] <= 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)