Class Fp8BlockMarlinLinear

java.lang.Object
smile.llm.quant.Fp8BlockMarlinLinear
All Implemented Interfaces:
AutoCloseable, LinearOp

public final class Fp8BlockMarlinLinear extends Object implements LinearOp, AutoCloseable
Block-scaled FP8 linear for GPUs without native FP8 (Ampere, Ada): DeepSeek / Qwen fine-grained checkpoints (weight_block_size=[128,128] with float32 weight_scale_inv) keep their FP8 weights in device memory and run through the weight-only Marlin kernel with bfloat16 activations. Each 128x128 scale is expanded over its 128 output columns, which is the kernel's group-128 layout, so the result equals dequantizing the weights (up to the rounding of the scales to bfloat16) and multiplying.

The layer takes ownership of the tensors passed to the constructor and replaces them by the repacked weight.

  • Constructor Details

    • Fp8BlockMarlinLinear

      public Fp8BlockMarlinLinear(Tensor weight, Tensor weightScaleInv, Tensor bias)
      Repacks a block-FP8 weight.
      Parameters:
      weight - FP8 E4M3 [out, in] on CUDA (in % 128 == 0, out % 128 == 0).
      weightScaleInv - float32 [out/128, in/128], W = fp8 * scale per tile.
      bias - optional bias, or null.
      Throws:
      IllegalStateException - if the native library lacks the Marlin FP8 kernels.
      IllegalArgumentException - on shapes the kernel cannot take.
  • Method Details

    • inFeatures

      public int inFeatures()
      Returns the number of input features.
      Returns:
      input features.
    • outFeatures

      public int outFeatures()
      Returns the number of output features.
      Returns:
      output features.
    • 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