flashinfer.diffusion_ops.minimax_h3_mlp_fp8_sm120

flashinfer.diffusion_ops.minimax_h3_mlp_fp8_sm120(x: Tensor, x_norm_weight: Tensor, adaln_scale: Tensor, adaln_shift: Tensor, adaln_index: Tensor, gate: Tensor, residual: Tensor, fc1_weight_q: Tensor, fc1_weight_scale: Tensor, fc2_weight_q: Tensor, fc2_weight_scale: Tensor, *, out: Tensor | None = None, workspace_a_q: Tensor | None = None, workspace_a_scale: Tensor | None = None, workspace_y: Tensor | None = None, workspace_y_q: Tensor | None = None, workspace_y_scale: Tensor | None = None, workspace_flags: Tensor | None = None, eps: float = 1e-05, fp8_mma_form: int = -1) → Tensor

Fused FP8 (W8A8) MiniMax-H3 MLP block for SM120 (RTX 5090 / RTX PRO 6000 Blackwell): RMSNorm + indexed AdaLN + FC1 + SwiGLU + FC2 + gated residual.

Computes, for each of the M rows (batch 1, no sequence parallelism):

n   = bf16(rmsnorm(x, eps) * x_norm_weight)                        # FP32 sum of squares, rsqrt
a   = bf16(adaln_shift[i] + n * bf16(1 + adaln_scale[i]))         # i = adaln_index[row]
a[i outside [0, rows)] = 0                                          # device-side guard
a_q = e4m3(a / s_a),  s_a = RN(amax_row(|a|) / 448)                # per-token activation scale
h   = bf16(a_q @ fc1_weight_q^T * s_a * fc1_weight_scale)          # FP32 accumulation, [M, 28672]
y   = bf16(bf16(silu(h_gate)) * h_up)                               # [M, 14336]
y_q = e4m3(y / s_y),  s_y = RN(amax_row(|y|) / 448)                # per-token, inside the FC2 launch
o   = bf16(y_q @ fc2_weight_q^T * s_y * fc2_weight_scale)          # FP32 accumulation, [M, 5376]
out = bf16(residual + bf16(gate[i] * o))

Three kernels run on the current stream (plus one memset of the FC2 tile flags): the one-CTA-per-row norm/AdaLN/quantization kernel, the persistent FC1 GEMM with the SwiGLU epilogue, and the persistent FC2 GEMM whose producer warps quantize y tile by tile and whose epilogue applies the gate and the residual. On GeForce GB202 boards the FP8 mma.sync is issued in the kind::mxf8f6f4 form with unit UE8M0 scales (twice the FP32-accumulating issue rate of the dense form there); the result is bitwise identical to the legacy form used on the RTX PRO 6000.

Parameters:

  • x (torch.Tensor) – Contiguous bfloat16 [M, 5376] hidden states, 1 <= M <= 149796.

  • x_norm_weight (torch.Tensor) – bfloat16 [5376] RMSNorm weight.

  • adaln_scale (torch.Tensor) – bfloat16 [rows, 5376] table views sharing one row count (unit column stride, 16-byte-aligned base and row pitch; adaln_scale and adaln_shift share one row stride). Column chunks of a wider [rows, k * 5376] modulation projection qualify without a copy.

  • adaln_shift (torch.Tensor) – bfloat16 [rows, 5376] table views sharing one row count (unit column stride, 16-byte-aligned base and row pitch; adaln_scale and adaln_shift share one row stride). Column chunks of a wider [rows, k * 5376] modulation projection qualify without a copy.

  • gate (torch.Tensor) – bfloat16 [rows, 5376] table views sharing one row count (unit column stride, 16-byte-aligned base and row pitch; adaln_scale and adaln_shift share one row stride). Column chunks of a wider [rows, k * 5376] modulation projection qualify without a copy.

  • adaln_index (torch.Tensor) – int64 [M] table row per activation row; rows with an index outside [0, rows) get a zero activation and a zero gate.

  • residual (torch.Tensor) – Contiguous bfloat16 [M, 5376] residual stream.

  • fc1_weight_q (torch.Tensor) – float8_e4m3fn [28672, 5376] and float32 [28672] from prepare_minimax_h3_fc1_weight_fp8() (SM120 prepacked row order).

  • fc1_weight_scale (torch.Tensor) – float8_e4m3fn [28672, 5376] and float32 [28672] from prepare_minimax_h3_fc1_weight_fp8() (SM120 prepacked row order).

  • fc2_weight_q (torch.Tensor) – float8_e4m3fn [5376, 14336] and float32 [5376] from prepare_minimax_h3_fc2_weight_fp8().

  • fc2_weight_scale (torch.Tensor) – float8_e4m3fn [5376, 14336] and float32 [5376] from prepare_minimax_h3_fc2_weight_fp8().

  • out (Optional[torch.Tensor]) – Optional bfloat16 [M, 5376] output (allocated when omitted); it may be residual or x (the epilogue reads a row’s residual before writing it).

  • workspace_a_q (Optional[torch.Tensor]) – Optional caller-owned float8_e4m3fn [M, 5376] / float32 [M] buffers that receive the quantized activation.

  • workspace_a_scale (Optional[torch.Tensor]) – Optional caller-owned float8_e4m3fn [M, 5376] / float32 [M] buffers that receive the quantized activation.

  • workspace_y (Optional[torch.Tensor]) – Optional caller-owned bfloat16 [M, 14336] / float8_e4m3fn [M, 14336] / float32 [M] buffers for y and its quantization.

  • workspace_y_q (Optional[torch.Tensor]) – Optional caller-owned bfloat16 [M, 14336] / float8_e4m3fn [M, 14336] / float32 [M] buffers for y and its quantization.

  • workspace_y_scale (Optional[torch.Tensor]) – Optional caller-owned bfloat16 [M, 14336] / float8_e4m3fn [M, 14336] / float32 [M] buffers for y and its quantization.

  • workspace_flags (Optional[torch.Tensor]) – Optional caller-owned int32 buffer of minimax_h3_mlp_fc2_flags_sm120() (M) elements (the FC2 tile ready flags; zeroed by the operator).

  • eps (float) – RMSNorm epsilon (positive; the operator is validated at 1e-5).

  • fp8_mma_form (int) – -1 selects the board’s FP8 instruction form, 0 forces the legacy form, 2 the kind::mxf8f6f4 form with unit scales.

Returns:

bfloat16 [M, 5376] (out).

Return type:

torch.Tensor