flashinfer.diffusion_ops.minimax_h3_mlp_nvfp4_sm120

flashinfer.diffusion_ops.minimax_h3_mlp_nvfp4_sm120(x: Tensor, x_norm_weight: Tensor, adaln_scale: Tensor, adaln_shift: Tensor, adaln_index: Tensor, gate: Tensor, residual: Tensor, a_global_scale: Tensor, fc1_weight_q: Tensor, fc1_scale_tiles: Tensor, fc1_alpha: float | Tensor, y_global_scale: Tensor, fc2_weight_q: Tensor, fc2_weight_sf: Tensor, fc2_alpha: float | Tensor, *, out: Tensor | None = None, workspace_a_q: Tensor | None = None, workspace_a_sf: Tensor | None = None, workspace_y_q: Tensor | None = None, workspace_y_sf: Tensor | None = None, workspace_flags: Tensor | None = None, eps: float = 1e-05) → Tensor

Fused NVFP4 (W4A4) MiniMax-H3 MLP block for SM120 (RTX 5090 / RTX PRO 6000 Blackwell).

The norm kernel computes the BF16 modulated activation a (as in minimax_h3_mlp_fp8_sm120()) and quantizes it with FlashInfer’s nvfp4_quantize() recipe (per 16 consecutive K elements sf = E4M3_RN(g * absmax * rcp(6)) saturating at 448 with g = a_global_scale, codes E2M1_RN_saturate(a * rcp(sf * rcp(g))); an all-zero block writes sf = 0 and codes 0). The FC1 GEMM accumulates (a_q * a_sf) . (w_q * w_sf) in FP32, applies fc1_alpha before the BF16 round of h, forms y = bf16(bf16(silu(h_gate)) * h_up) in registers and quantizes y in the same epilogue with the same recipe under g = y_global_scale (block-16 E2M1 codes + dense UE4M3 scales); the FC2 GEMM consumes those operands directly:

h   = bf16(fc1_alpha * ((a_q * a_sf) @ (fc1_weight_q * fc1_scale_tiles)^T))
y   = bf16(bf16(silu(h_gate)) * h_up)
o   = bf16(fc2_alpha * ((y_q * y_sf) @ (fc2_weight_q * fc2_weight_sf)^T))
out = bf16(residual + bf16(gate[i] * o))
Parameters:

  • x – As in minimax_h3_mlp_fp8_sm120().

  • x_norm_weight – As in minimax_h3_mlp_fp8_sm120().

  • adaln_scale – As in minimax_h3_mlp_fp8_sm120().

  • adaln_shift – As in minimax_h3_mlp_fp8_sm120().

  • adaln_index – As in minimax_h3_mlp_fp8_sm120().

  • gate – As in minimax_h3_mlp_fp8_sm120().

  • residual – As in minimax_h3_mlp_fp8_sm120().

  • eps – As in minimax_h3_mlp_fp8_sm120().

  • a_global_scale (torch.Tensor) – float32 [1] CUDA tensors: the activation global scales of a and y (448 * 6 / absmax convention, typically calibrated); read on the device.

  • y_global_scale (torch.Tensor) – float32 [1] CUDA tensors: the activation global scales of a and y (448 * 6 / absmax convention, typically calibrated); read on the device.

  • fc1_weight_q (torch.Tensor) – Packed E2M1 uint8 [28672, 2688] weights and the flat uint8 [28672 * 336] swizzled block scales from prepare_minimax_h3_fc1_weight_nvfp4_sm120().

  • fc1_scale_tiles (torch.Tensor) – Packed E2M1 uint8 [28672, 2688] weights and the flat uint8 [28672 * 336] swizzled block scales from prepare_minimax_h3_fc1_weight_nvfp4_sm120().

  • fc1_alpha (float or torch.Tensor) – 1 / (a_global_scale * w1_global_scale) and 1 / (y_global_scale * w2_global_scale) (a Python float or a one-element tensor; a device tensor is read with a host synchronization).

  • fc2_alpha (float or torch.Tensor) – 1 / (a_global_scale * w1_global_scale) and 1 / (y_global_scale * w2_global_scale) (a Python float or a one-element tensor; a device tensor is read with a host synchronization).

  • fc2_weight_q (torch.Tensor) – Packed E2M1 uint8 [5376, 7168] weights and the flat uint8 [5376 * 896] swizzled block scales from prepare_minimax_h3_fc2_weight_nvfp4_sm120().

  • fc2_weight_sf (torch.Tensor) – Packed E2M1 uint8 [5376, 7168] weights and the flat uint8 [5376 * 896] swizzled block scales from prepare_minimax_h3_fc2_weight_nvfp4_sm120().

  • out (Optional[torch.Tensor]) – Optional bfloat16 [M, 5376] output; may be residual or x.

  • workspace_a_q (Optional[torch.Tensor]) – Optional caller-owned uint8 [M, 2688] / uint8 [M, 336] buffers (packed codes and dense row-major block scales of a).

  • workspace_a_sf (Optional[torch.Tensor]) – Optional caller-owned uint8 [M, 2688] / uint8 [M, 336] buffers (packed codes and dense row-major block scales of a).

  • workspace_y_q (Optional[torch.Tensor]) – Optional caller-owned uint8 [M, 7168] / uint8 [M, 896] buffers (packed codes and dense row-major block scales of y).

  • workspace_y_sf (Optional[torch.Tensor]) – Optional caller-owned uint8 [M, 7168] / uint8 [M, 896] buffers (packed codes and dense row-major block scales of y).

  • workspace_flags (Optional[torch.Tensor]) – Optional caller-owned int32 buffer of minimax_h3_mlp_fc2_flags_sm120() (M) elements.

Returns:

bfloat16 [M, 5376] (out).

Return type:

torch.Tensor