Class MxFp8BlockLinear

java.lang.Object

smile.llm.quant.MxFp8BlockLinear

All Implemented Interfaces:

AutoCloseable, LinearOp, MxFp8Layer

public class MxFp8BlockLinear
extends Object
implements MxFp8Layer

Linear layer for MXFP8 weights on the native block-scaled FP8 GEMM of Hopper. The checkpoint's 1x32 E8M0 scales are folded into 128x128 float32 tile scales when the weight is installed (Native.mxfp8ToBlockFp8(Tensor, Tensor)), then every call is a W8A8 Fp8BlockGemm.

Unlike the Marlin path (FP8 weights, bfloat16 activations) the activations are rounded to FP8 here, so results differ slightly from W8A16 execution.

Constructor Summary

Constructors

Constructor	Description
MxFp8BlockLinear(int inFeatures, int outFeatures, boolean emulate)	Constructor.

Method Summary

Modifier and Type	Method	Description
void	close()
Tensor	forward(Tensor input)	Applies the linear transformation.
boolean	isLoaded()	Returns whether the weight has been installed.
void	setWeight(Tensor w, Tensor scale)	Installs the weight from checkpoint tensors.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

MxFp8BlockLinear

public MxFp8BlockLinear(int inFeatures,
 int outFeatures,
 boolean emulate)

Constructor.

Parameters:

inFeatures - reduction size K; a multiple of 128.

outFeatures - output size N; a multiple of 128.

emulate - true to emulate the FP8 tensor-core GEMM (runs on any CUDA GPU).

Method Details

setWeight

public void setWeight(Tensor w,
 Tensor scale)

Description copied from interface: MxFp8Layer

Installs the weight from checkpoint tensors.

Specified by:

setWeight in interface MxFp8Layer

Parameters:

w - FP8 E4M3 [N, K] on the target CUDA device.

scale - E8M0 scales as uint8 [N, K / 32] on the same device.

isLoaded

public boolean isLoaded()

Description copied from interface: MxFp8Layer

Returns whether the weight has been installed.

Specified by:

isLoaded in interface MxFp8Layer

Returns:

true once MxFp8Layer.setWeight(Tensor, Tensor) has run.

forward

public Tensor forward(Tensor input)

Description copied from interface: LinearOp

Applies the linear transformation.

Specified by:

forward in interface LinearOp

Parameters:

input - activations (caller retains ownership).

Returns:

owned output tensor (caller must close).

close

public void close()

Specified by:

close in interface AutoCloseable

Specified by:

close in interface MxFp8Layer