Class MxFp8MarlinLinear

java.lang.Object
smile.llm.quant.MxFp8MarlinLinear
All Implemented Interfaces:
AutoCloseable, LinearOp, MxFp8Layer
public class MxFp8MarlinLinear extends Object implements MxFp8Layer
Linear layer whose weight stays in MXFP8 (FP8 E4M3 elements, E8M0 scale per 32 elements) in the Marlin tile layout and is multiplied with bfloat16 activations by the weight-only Marlin kernel. It works on GPUs without native FP8 (Ampere). The weight takes one byte per element plus 1/32 byte of scale.

The layer is created empty and receives its weight from the checkpoint through setWeight(Tensor, Tensor); it owns no module parameters, so LayerBlock.to(device, dtype) never converts it.

  • Constructor Details

    • MxFp8MarlinLinear

      public MxFp8MarlinLinear(int inFeatures, int outFeatures)
      Constructor.
      Parameters:
      inFeatures - reduction size K; a multiple of 128.
      outFeatures - output size N; a multiple of 64.

  • Method Details

    • setWeight

      public void setWeight(Tensor weight, Tensor scale)
      Installs the weight from checkpoint tensors, repacking them for the kernel.
      Specified by:
      setWeight in interface MxFp8Layer
      Parameters:
      weight - FP8 E4M3 [N, K] on the target CUDA device.
      scale - E8M0 scales as uint8 [N, K / 32] on the same device.

    • fuse

      public static MxFp8MarlinLinear fuse(MxFp8MarlinLinear... parts)
      Concatenates the outputs of several loaded layers that share their input into one layer, so a single kernel launch computes all of them. Marlin tiles are 64 output columns wide, so packed weights and permuted scales concatenate along the output dimension. The parts are closed.
      Parameters:
      parts - loaded layers with equal inFeatures; their outputs appear in this order.
      Returns:
      the fused layer.

    • outFeatures

      public int outFeatures()
      Returns the output size N.
      Returns:
      the output size N.

    • isLoaded

      public boolean isLoaded()
      Returns whether the weight has been installed.
      Specified by:
      isLoaded in interface MxFp8Layer
      Returns:
      true once 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