Skip to content

vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel

Classes:

Mxfp8LinearKernel

Bases: ABC

Base class for MXFP8 quantized linear kernels.

Each subclass implements a specific GEMM backend (FlashInfer CUTLASS, Marlin, emulation).

Methods:

  • input_quant_key –

    Return the input quantization key supported by this kernel. If the kernel

Attributes:

Source code in vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class Mxfp8LinearKernel(ABC):
    """Base class for MXFP8 quantized linear kernels.

    Each subclass implements a specific GEMM backend (FlashInfer CUTLASS,
    Marlin, emulation).
    """

    supports_pre_processed_weights: bool = False
    """True if ``process_weights_after_loading`` only rewrites parameters, so
    weights exported by the weight cache daemon can be used as-is."""

    def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
        assert self.can_implement(c)[0]
        assert self.is_supported()[0]
        self.config = c

    def input_quant_key(self) -> QuantKey | None:
        """Return the input quantization key supported by this kernel. If the kernel
        does not support input quantization outside of the kernel, return None.
        """
        return None

    @classmethod
    @abstractmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        raise NotImplementedError

    @classmethod
    @abstractmethod
    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
        raise NotImplementedError

    @abstractmethod
    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        raise NotImplementedError

    @abstractmethod
    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

supports_pre_processed_weights = False class-attribute instance-attribute

True if process_weights_after_loading only rewrites parameters, so weights exported by the weight cache daemon can be used as-is.

input_quant_key()

Return the input quantization key supported by this kernel. If the kernel does not support input quantization outside of the kernel, return None.

Source code in vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py 
44
45
46
47
48
def input_quant_key(self) -> QuantKey | None:
    """Return the input quantization key supported by this kernel. If the kernel
    does not support input quantization outside of the kernel, return None.
    """
    return None

Mxfp8LinearLayerConfig dataclass

Configuration for an MXFP8 linear layer.

All MXFP8 layers share the same structure: FP8-E4M3 weights with uint8 (E8M0) per-block scales at block size 32.

Attributes:

  • weight_shape (tuple[int, int]) –

    The layer's (out_features, in_features), i.e. (N, K).

Source code in vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
@dataclass
class Mxfp8LinearLayerConfig:
    """Configuration for an MXFP8 linear layer.

    All MXFP8 layers share the same structure: FP8-E4M3 weights with
    uint8 (E8M0) per-block scales at block size 32.

    Attributes:
        weight_shape: The layer's `(out_features, in_features)`, i.e. `(N, K)`.

    """

    weight_shape: tuple[int, int]
    bmm_batch_size: int | None = None