flashinfer.diffusion_ops.prepare_minimax_h3_fc2_weight_nvfp4_sm120

flashinfer.diffusion_ops.prepare_minimax_h3_fc2_weight_nvfp4_sm120(fc2_weight: Tensor, w_global_scale: float | Tensor) → Tuple[Tensor, Tensor]

Quantize the FC2 weight for minimax_h3_mlp_nvfp4_sm120().

fc2_weight is the BF16 [5376, 14336] matrix and w_global_scale its float32 global scale (448 * 6 / absmax). The rows are quantized with FlashInfer’s nvfp4_quantize() (sfLayout=SfLayout.layout_128x4, do_shuffle=False; per 16 consecutive K elements: UE4M3 scale = E4M3_RN(g * absmax / 6) saturating at 448, E2M1 round-to-nearest with saturation of w * g / scale).

Parameters:

  • fc2_weight (torch.Tensor) – CUDA bfloat16 [5376, 14336] FC2 weight matrix.

  • w_global_scale (float or torch.Tensor) – Global weight scale (448 * 6 / absmax), supplied as a scalar or one-element tensor.

Returns:

(fc2_weight_q, fc2_weight_sf): packed E2M1 uint8 [5376, 7168] weights (even element in the low nibble) and a flat uint8 tensor of 5376 * 896 bytes holding the block scales in the FlashInfer 128x4 swizzled layout (the GEMM streams them as 42 output tiles of 448 rows x 256 bytes).

Return type:

Tuple[torch.Tensor, torch.Tensor]