mxfp8linear ¶

def mxfp8_scaled_mm(
    activation_values: torch.Tensor,
    activation_scales: torch.Tensor,
    weight_values: torch.Tensor,
    weight_scales: torch.Tensor,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    """Multiply two MXFP8 matrices and return a BF16 matrix."""
    return F.scaled_mm(
        mat_a=activation_values,
        mat_b=weight_values.mT,
        scale_a=activation_scales,
        scale_recipe_a=F.ScalingType.BlockWise1x32,
        scale_b=weight_scales,
        scale_recipe_b=F.ScalingType.BlockWise1x32,
        swizzle_a=F.SwizzleType.SWIZZLE_32_4_4,
        swizzle_b=F.SwizzleType.SWIZZLE_32_4_4,
        bias=bias,
        output_dtype=torch.bfloat16,
    )

def mxfp8_swiglu_feed_forward(
    hidden_states: torch.Tensor,
    fc_in: torch.nn.Module,
    fc_out: torch.nn.Module,
) -> torch.Tensor:
    """Run the MiniMax-H3 feed-forward network with MXFP8 GEMMs."""
    fc_in_base = _resolve_merged_linear(fc_in)
    fc_out_base = _resolve_merged_linear(fc_out)
    fc_in_method = fc_in_base.quant_method
    fc_out_method = fc_out_base.quant_method

    preactivation = fc_in_method.apply(fc_in_base, hidden_states, fc_in_base.bias)
    output_shape = (*hidden_states.shape[:-1], fc_out_base.output_size)
    preactivation_2d = preactivation.reshape(-1, preactivation.shape[-1])
    activation_values, activation_scales = swiglu_quantize_mxfp8_blockwise(preactivation_2d)
    output_2d = fc_out_method.apply_quantized(
        fc_out_base,
        activation_values,
        activation_scales,
        fc_out_base.bias,
    )
    return output_2d.reshape(output_shape)

def quantize_mxfp8_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize activation rows and write hardware-blocked scales directly."""
    _validate_mxfp8_matrix(matrix)
    matrix = matrix.contiguous()
    row_count, column_count = matrix.shape
    quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(matrix, column_count)
    grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM))
    _quantize_mxfp8_kernel[grid](
        matrix,
        quantized,
        blocked_scales.view(torch.uint8),
        row_count,
        column_count,
        scale_column_count,
        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
        num_warps=8,
    )
    return quantized, blocked_scales

def quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Prequantize a weight with Quack and return hardware-blocked scales."""
    _validate_mxfp8_matrix(matrix)
    return _quantize_mxfp8_weight_blockwise(matrix)

def swiglu_quantize_mxfp8_blockwise(preactivation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply H3 value-first SwiGLU and quantize the BF16 result to MXFP8."""
    if preactivation.ndim != 2:
        raise ValueError(f"MXFP8 SwiGLU requires a 2D tensor, got shape {tuple(preactivation.shape)}.")
    if preactivation.dtype != torch.bfloat16:
        raise TypeError(f"MXFP8 SwiGLU requires BF16 input, got {preactivation.dtype}.")
    if preactivation.shape[1] % (2 * MXFP8_BLOCK_SIZE):
        raise ValueError("MXFP8 SwiGLU requires each packed half to be divisible by "
                         f"{MXFP8_BLOCK_SIZE}, got packed width {preactivation.shape[1]}.")
    preactivation = preactivation.contiguous()
    row_count = preactivation.shape[0]
    column_count = preactivation.shape[1] // 2
    quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(preactivation, column_count)
    grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM))
    _swiglu_quantize_mxfp8_kernel[grid](
        preactivation,
        quantized,
        blocked_scales.view(torch.uint8),
        row_count,
        column_count,
        scale_column_count,
        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
        num_warps=8,
    )
    return quantized, blocked_scales