flashinfer.diffusion_ops.prepare_minimax_h3_fc2_weight_fp8

flashinfer.diffusion_ops.prepare_minimax_h3_fc2_weight_fp8(fc2_weight: Tensor, chunk_rows: int = 1024) → Tuple[Tensor, Tensor]

Quantize the FC2 weight for minimax_h3_mlp_fp8_sm120().

fc2_weight is the BF16 nn.Linear matrix [5376, 14336] (output channels x FFN). Each output channel (row) is quantized to E4M3 with scale = RN(max(amax(row), 1e-12) / 448) (true IEEE division) and q = RN_sat(row / scale); the row order is unchanged.

Parameters:

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

  • chunk_rows (int) – Number of output channels quantized per chunk (default: 1024).

Returns:

(fc2_weight_q, fc2_weight_scale): float8_e4m3fn [5376, 14336] weights and float32 [5376] per-row scales.

Return type:

Tuple[torch.Tensor, torch.Tensor]