vllm.model_executor.layers.quantization.modelopt ¶

@dataclass(frozen=True)
class CkptCtx:
    """Per-checkpoint facts a QuantKey cannot carry."""

    group_size: int | None = None
    scale_block_size: tuple[int, int] | None = None

class FormatScheme:
    """Optional per-format hooks that compose *around* the QuantKey schemes.

    Extension seam for residue that belongs to a format as a whole rather than a
    single QuantKey — an extra parameter the weight/activation schemes don't
    cover, or a tweak before/after their ``process``. All hooks default to
    no-ops; most formats need none (they are a pure ``(weight, activation)`` key
    pair). To add one: subclass, override the hooks you need, and return an
    instance from the format's ``resolve()`` branch — no change to
    ``ModelOptLinearMethod`` itself.
    """

    def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None:
        """Register format-level params (after the key schemes' weights)."""

    def kernel_weight_shape(self, layer) -> tuple[int, int]:
        """Return the weight shape the selected kernel will receive."""
        return tuple(layer.weight.shape)

    def pre_process(self, layer) -> None:
        """Run before the key schemes' ``process``."""

    def post_process(self, layer) -> None:
        """Run after the key schemes' ``process``, before the kernel's."""

    def apply(self, layer, x, bias, kernel_apply):
        """Wrap the kernel's forward. ``kernel_apply(layer, x, bias) -> Tensor``.

        Default: delegate unchanged. A format whose residue lives at compute
        time (e.g. run the GEMM on a padded weight then slice the output back
        to its logical width, #53132-style) overrides this and calls
        ``kernel_apply`` in the middle.
        """
        return kernel_apply(layer, x, bias)

class KDynamicNoParam(QuantKeyScheme):
    """Dynamic activation with no stored scale (W8A8): quantized at runtime in
    the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a
    non-None activation key. Activation-role only. Serves the fp8 per-token, fp8
    per-block, and mxfp8 dynamic activation keys."""

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not ACT:
            self.reject(role)
        # dynamic -> nothing stored

    def process(self, layer, role) -> None:
        if role is not ACT:
            self.reject(role)

class KFp8Block128(QuantKeyScheme):
    """128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt
    exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D.
    No transpose (block kernel keeps [out,in])."""

    key = kFp8Static128BlockSym

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not WEIGHT:
            self.reject(role)
        out = shapes.output_size_per_partition
        in_ = shapes.input_size_per_partition
        # Scale sized by cdiv to match the checkpoint's block grid (the last
        # output block may be partial). Weight is registered at its logical
        # size; a non-128 output width is padded to a block boundary by the
        # _Fp8PbWoPartialBlock format scheme. Input must be 128-aligned (same as
        # #53132, which only pads the output).
        if in_ % 128 != 0:
            raise ValueError(f"FP8_PB_WO requires in divisible by 128, got {in_}")
        self.register_params(
            layer,
            "weight",
            (out, in_),
            torch.float8_e4m3fn,
            ModelWeightParameter,
            wl,
            input_dim=1,
            output_dim=0,
        )
        ob, ib = cdiv(out, 128), in_ // 128
        self.register_params(
            layer,
            "weight_scale",
            (ob, 1, ib, 1),
            torch.float32,
            BlockQuantScaleParameter,
            wl,
            input_dim=2,
            output_dim=0,
            init=FP8_SCALE_SENTINEL,
        )
        layer.weight_block_size = [128, 128]

    def process(self, layer, role) -> None:
        if role is not WEIGHT:
            self.reject(role)
        layer.weight = Parameter(layer.weight.data, requires_grad=False)
        s = layer.weight_scale
        if s.dim() == 4:
            s = s.squeeze(1).squeeze(-1)  # [ob,1,ib,1] -> [ob,ib]
        elif s.dim() != 2:
            raise ValueError(
                f"Unexpected FP8_PB_WO weight_scale shape {tuple(s.shape)}"
            )
        layer.weight_scale = Parameter(s.contiguous(), requires_grad=False)

class KFp8StaticChannel(QuantKeyScheme):
    """Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only —
    there is no static per-channel *activation* today."""

    key = kFp8StaticTokenSym

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not WEIGHT:
            self.reject(role)
        self.register_params(
            layer,
            "weight",
            (shapes.output_size_per_partition, shapes.input_size_per_partition),
            torch.float8_e4m3fn,
            ModelWeightParameter,
            wl,
            input_dim=1,
            output_dim=0,
        )
        self.register_params(
            layer,
            "weight_scale",
            (shapes.output_size_per_partition,),
            torch.float32,
            ChannelQuantScaleParameter,
            wl,
            output_dim=0,
            init=FP8_SCALE_SENTINEL,
        )

    def process(self, layer, role) -> None:
        if role is not WEIGHT:
            self.reject(role)
        weight, weight_scale, _ = process_fp8_weight_channel_strategy(
            layer.weight, layer.weight_scale.data
        )
        layer.weight = Parameter(weight.t(), requires_grad=False)
        layer.weight_scale = Parameter(weight_scale, requires_grad=False)

class KFp8StaticTensor(QuantKeyScheme):
    """Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and
    the activation slot (W8A8). One key in both QuantSpec slots."""

    key = kFp8StaticTensorSym

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is WEIGHT:
            self.register_params(
                layer,
                "weight",
                (shapes.output_size_per_partition, shapes.input_size_per_partition),
                torch.float8_e4m3fn,
                ModelWeightParameter,
                wl,
                input_dim=1,
                output_dim=0,
            )
            layer.orig_dtype = shapes.params_dtype
            self.register_params(
                layer,
                "weight_scale",
                (shapes.num_partitions,),
                torch.float32,
                PerTensorScaleParameter,
                wl,
                init=FP8_SCALE_SENTINEL,
            )
        elif role is ACT:
            self.register_params(
                layer,
                "input_scale",
                (shapes.num_partitions,),
                torch.float32,
                PerTensorScaleParameter,
                wl,
                init=FP8_SCALE_SENTINEL,
            )
        else:
            self.reject(role)

    def process(self, layer, role) -> None:
        if role is WEIGHT:
            weight = layer.weight
            max_w_scale = layer.weight_scale.max()
            if not (layer.weight_scale == layer.weight_scale[0]).all():
                max_w_scale, weight = requantize_with_max_scale(
                    layer.weight, layer.weight_scale, layer.logical_widths
                )
            # Transposed here rather than in the kernel: ModelOpt maps each
            # fp8 key to exactly one kernel, so the layout is key-determined.
            layer.weight = Parameter(weight.t(), requires_grad=False)
            # The plain Parameter drops ModelWeightParameter's dim attributes;
            # restore them for the transposed layout, which Humming reads.
            layer.weight.input_dim = 0
            layer.weight.output_dim = 1
            layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
        elif role is ACT:
            if torch.unique(layer.input_scale).numel() != 1:
                logger.warning_once(
                    "In FP8 linear, the static input scale differs across "
                    "parallel layers (e.g. q_proj, k_proj, v_proj). Collapsing "
                    "them to the max will likely reduce accuracy. Consider a "
                    "checkpoint with the same input scale for fused weights."
                )
            layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
        else:
            self.reject(role)

class KMxfp8Static(QuantKeyScheme):
    """MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale.
    Weight-role only. process is validate-only plus an idempotency guard."""

    key = kMxfp8Static

    @staticmethod
    def get_scale_weight_loader(weight_loader: Callable, ctx: CkptCtx) -> Callable:
        block_rows, block_cols = ctx.scale_block_size or (1, MXFP8_BLOCK_SIZE)
        if block_rows < 1 or block_cols != MXFP8_BLOCK_SIZE:
            raise NotImplementedError(
                f"MXFP8 checkpoint scale block {ctx.scale_block_size} is unsupported"
            )

        def scaled_loader(param, loaded_weight, *args, **kwargs):
            assert loaded_weight.dtype in (torch.uint8, torch.float8_e8m0fnu)
            loaded_weight = loaded_weight.view(torch.uint8).repeat_interleave(
                block_rows, dim=0
            )
            return weight_loader(param, loaded_weight, *args, **kwargs)

        return scaled_loader if block_rows > 1 else weight_loader

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not WEIGHT:
            self.reject(role)
        if shapes.input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
            raise ValueError(
                f"MXFP8 requires input size divisible by {MXFP8_BLOCK_SIZE}, "
                f"got {shapes.input_size_per_partition}"
            )
        self.register_params(
            layer,
            "weight",
            (shapes.output_size_per_partition, shapes.input_size_per_partition),
            MXFP8_VALUE_DTYPE,
            ModelWeightParameter,
            wl,
            input_dim=1,
            output_dim=0,
        )
        scale_loader = self.get_scale_weight_loader(wl, ctx)
        self.register_params(
            layer,
            "weight_scale",
            (
                shapes.output_size_per_partition,
                shapes.input_size_per_partition // MXFP8_BLOCK_SIZE,
            ),
            MXFP8_SCALE_DTYPE,
            ModelWeightParameter,
            scale_loader,
            input_dim=1,
            output_dim=0,
        )
        layer.weight_block_size = [1, MXFP8_BLOCK_SIZE]

    def process(self, layer, role) -> None:
        if role is not WEIGHT:
            self.reject(role)
        # Idempotency: the emulation kernel's post-load may replace the 1-byte
        # MXFP8 weight with BF16 (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD, default
        # on). On a weight reload process runs again after that swap, so skip
        # the fp8 asserts when the weight is already >=2-byte.
        if layer.weight.element_size() >= 2:
            return
        if getattr(layer, "is_bmm", False) and layer.weight.ndim == 3:
            return
        assert layer.weight.ndim == 2 and layer.weight.dtype == MXFP8_VALUE_DTYPE
        assert layer.weight_scale.ndim == 2
        assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE

class KNvfp4Dynamic(QuantKeyScheme):
    """NVFP4 activation scheme (W4A4). Has a static global input scale on disk;
    the per-group scale is computed at runtime inside the kernel."""

    key = kNvfp4Dynamic

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not ACT:
            self.reject(role)
        self.register_params(
            layer,
            "input_scale",
            (shapes.num_partitions,),
            torch.float32,
            PerTensorScaleParameter,
            wl,
        )

    def process(self, layer, role) -> None:
        if role is not ACT:
            self.reject(role)
        if torch.unique(layer.input_scale).numel() != 1:
            logger.warning_once(
                "In NVFP4 linear, the global input scale differs across "
                "parallel layers (e.g. q_proj, k_proj, v_proj). This will "
                "likely reduce accuracy. Consider a checkpoint with the same "
                "global NVFP4 scale for fused weights."
            )
        input_global_scale = layer.input_scale.max().to(torch.float32)
        layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
        layer.input_global_scale_inv = Parameter(
            (1.0 / layer.input_global_scale).to(torch.float32),
            requires_grad=False,
        )
        del layer.input_scale

class KNvfp4Static(QuantKeyScheme):
    """NVFP4 weight scheme (W4A4 and W4A16 share it). Weight-role only today."""

    key = kNvfp4Static

    def create_weights(self, layer, role, ctx, shapes, wl) -> None:
        if role is not WEIGHT:
            self.reject(role)
        if shapes.input_size_per_partition % 16 != 0:
            raise ValueError(
                "Unsupported model when in features size is not multiple of 16"
            )
        # Packed NVFP4 weight: 2 fp4 items per byte along the input dim.
        self.register_params(
            layer,
            "weight",
            (shapes.output_size_per_partition, shapes.input_size_per_partition // 2),
            torch.uint8,
            ModelWeightParameter,
            wl,
            input_dim=1,
            output_dim=0,
        )
        # Per-tensor global weight scale.
        self.register_params(
            layer,
            "weight_scale_2",
            (shapes.num_partitions,),
            torch.float32,
            PerTensorScaleParameter,
            wl,
        )
        # Per-block (group_size) weight scale.
        self.register_params(
            layer,
            "weight_scale",
            (
                shapes.output_size_per_partition,
                shapes.input_size_per_partition // ctx.group_size,
            ),
            torch.float8_e4m3fn,
            ModelWeightParameter,
            wl,
            input_dim=1,
            output_dim=0,
        )

    def process(self, layer, role) -> None:
        if role is not WEIGHT:
            self.reject(role)
        if torch.unique(layer.weight_scale_2).numel() != 1:
            logger.warning_once(
                "In NVFP4 linear, the global weight scale differs across "
                "parallel layers (e.g. q_proj, k_proj, v_proj). This will "
                "likely reduce accuracy. Consider a checkpoint with the same "
                "global NVFP4 scale for fused weights."
            )
        # Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688.
        weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
        layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
        del layer.weight_scale_2

class ModelOptFp8Config(ModelOptQuantConfigBase):
    """Config class for ModelOpt FP8."""

    def __init__(
        self,
        quant_method: str,
        is_checkpoint_fp8_serialized: bool,
        kv_cache_quant_method: str | None,
        exclude_modules: list[str],
    ) -> None:
        super().__init__(exclude_modules)
        self.quant_method = quant_method
        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
        self.kv_cache_quant_method = kv_cache_quant_method

        # Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch.
        supported = algos_owned_by("modelopt")
        if self.quant_method not in supported:
            raise ValueError(
                "Unsupported ModelOpt FP8 quant_algo for vLLM: "
                f"{self.quant_method}. Supported: {' / '.join(supported)}."
            )

    def has_blocked_weights(self) -> bool:
        # Gates "+quant_fp8" in VllmConfig. FP8_PB_WO runs the block-scaled
        # kernels, whose DeepGEMM path needs UE8M0-packed group scales; without
        # the op enabled QuantFP8 falls back to forward_native, which emits
        # unpacked fp32 scales and the GEMM fails to launch.
        return self.quant_method == "FP8_PB_WO"

    def get_name(self) -> QuantizationMethods:
        return "modelopt"

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.bfloat16, torch.half]

    @classmethod
    def get_min_capability(cls) -> int:
        return 80

    @classmethod
    def override_quantization_method(
        cls, hf_quant_cfg, user_quant, hf_config=None
    ) -> QuantizationMethods | None:
        algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)
        if algo is not None and algo == "FP8":
            return "modelopt"
        return None

    @classmethod
    def _from_config(
        cls,
        *,
        quant_method: str,
        kv_cache_quant_method: str | None,
        exclude_modules: list[str],
        original_config: dict[str, Any],
        **kwargs: Any,
    ) -> "ModelOptFp8Config":
        is_checkpoint_fp8_serialized = "FP8" in quant_method

        return cls(
            quant_method,
            is_checkpoint_fp8_serialized,
            kv_cache_quant_method,
            exclude_modules,
        )

class ModelOptFp8MoEMethod(FusedMoEMethodBase):
    """MoE method for ModelOpt FP8.
    Supports loading FP8 checkpoints with static weight scale and
    activation scale.
    Args:
        quant_config: The ModelOpt quantization config.
    """

    def __init__(
        self,
        quant_config: ModelOptFp8Config,
        moe_config: FusedMoEConfig,
    ) -> None:
        super().__init__(moe_config)
        self.quant_config = quant_config
        assert self.quant_config.is_checkpoint_fp8_serialized

        # Select Fp8 MoE backend
        self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
            config=self.moe,
            weight_key=kFp8StaticTensorSym,
            activation_key=kFp8StaticTensorSym,
        )

    def create_weights(
        self,
        layer: RoutedExperts,
        num_experts: int,
        hidden_size: int,
        intermediate_size_per_partition: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ):
        layer.orig_dtype = params_dtype
        layer.num_experts = num_experts

        # Use FP8 dtype if checkpoint is serialized
        weight_dtype = (
            torch.float8_e4m3fn
            if self.quant_config.is_checkpoint_fp8_serialized
            else params_dtype
        )
        weight_loader = extra_weight_attrs.get("weight_loader")

        w13_num_shards = 2 if self.moe.is_act_and_mul else 1

        w13_weight = ModelWeightParameter(
            data=torch.empty(
                num_experts,
                w13_num_shards * intermediate_size_per_partition,
                hidden_size,
                dtype=weight_dtype,
            ),
            input_dim=2,
            output_dim=1,
            weight_loader=weight_loader,
        )
        layer.register_parameter("w13_weight", w13_weight)

        w2_weight = ModelWeightParameter(
            data=torch.empty(
                num_experts,
                hidden_size,
                intermediate_size_per_partition,
                dtype=weight_dtype,
            ),
            input_dim=2,
            output_dim=1,
            weight_loader=weight_loader,
        )
        layer.register_parameter("w2_weight", w2_weight)

        # WEIGHT SCALES - Per-tensor scaling for ModelOpts
        # For gated MoE, allocate 2 scales for w1 and w3 respectively.
        # They will be combined to a single scale after weight loading.
        # For non-gated MoE, allocate 1 scale for w13.
        w13_weight_scale = PerTensorScaleParameter(
            data=torch.full(
                (num_experts, w13_num_shards),
                1.0,
                dtype=torch.float32,
            ),
            weight_loader=weight_loader,
        )
        w2_weight_scale = PerTensorScaleParameter(
            data=torch.full((num_experts,), 1.0, dtype=torch.float32),
            weight_loader=weight_loader,
        )
        layer.register_parameter("w13_weight_scale", w13_weight_scale)
        layer.register_parameter("w2_weight_scale", w2_weight_scale)

        # INPUT SCALES - Per-tensor scaling for ModelOpt
        w13_input_scale = PerTensorScaleParameter(
            data=torch.full((num_experts,), 1.0, dtype=torch.float32),
            weight_loader=weight_loader,
        )
        w2_input_scale = PerTensorScaleParameter(
            data=torch.full((num_experts,), 1.0, dtype=torch.float32),
            weight_loader=weight_loader,
        )
        layer.register_parameter("w13_input_scale", w13_input_scale)
        layer.register_parameter("w2_input_scale", w2_input_scale)

    def _setup_kernel(
        self,
        layer: RoutedExperts,
        w13: torch.Tensor,
        w2: torch.Tensor,
        w13_scale: torch.Tensor,
        w2_scale: torch.Tensor,
        w13_input_scale: torch.Tensor,
        w2_input_scale: torch.Tensor,
    ):
        w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
            fp8_backend=self.fp8_backend,
            layer=layer,
            w13=w13,
            w2=w2,
            w13_scale=w13_scale,
            w2_scale=w2_scale,
            w13_input_scale=w13_input_scale,
            w2_input_scale=w2_input_scale,
        )

        # Replace parameters with updated versions. Note that this helper
        # function ensures the replacement is compatible with RL weight reloads.
        replace_parameter(layer, "w13_weight", w13)
        replace_parameter(layer, "w2_weight", w2)
        replace_parameter(layer, "w13_weight_scale", w13_scale)
        replace_parameter(layer, "w2_weight_scale", w2_scale)

        # Setup modular kernel.
        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
        assert self.experts_cls is not None
        self.moe_kernel = make_fp8_moe_kernel(
            moe_quant_config=self.moe_quant_config,
            moe_config=self.moe,
            fp8_backend=self.fp8_backend,
            experts_cls=self.experts_cls,
            routing_tables=layer._expert_routing_tables(),
        )

    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
        w13 = layer.w13_weight
        w2 = layer.w2_weight
        w13_scale = layer.w13_weight_scale
        w2_scale = layer.w2_weight_scale
        w13_input_scale = layer.w13_input_scale
        w2_input_scale = layer.w2_input_scale

        # Per tensor kernels require single activation scale. Use the max.
        w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
            w13_input_scale,
            w2_input_scale,
            layer.moe_config.moe_parallel_config.enable_eplb,
        )
        replace_parameter(layer, "w13_input_scale", w13_input_scale)
        replace_parameter(layer, "w2_input_scale", w2_input_scale)

        # Per tensor kernels require single weight scale for w13 per expert, but
        # on disk there is a scale for w1 and w3. Use the max to requantize.
        shard_size = layer.intermediate_size_per_partition
        w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
            w13,
            w13_scale,
            shard_size,
            num_experts=layer.w13_weight.shape[0],
            is_act_and_mul=self.moe.is_act_and_mul,
        )

        # Shuffle weights to runtime format and setup kernel.
        self._setup_kernel(
            layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
        )

    def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
        w1_scale = layer.w13_weight_scale
        w2_scale = layer.w2_weight_scale
        a1_scale = layer.w13_input_scale
        a2_scale = layer.w2_input_scale

        return make_fp8_moe_quant_config(
            fp8_backend=self.fp8_backend,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            swiglu_limit=getattr(layer, "swiglu_limit", None),
            layer=layer,
        )

    def apply_monolithic(
        self,
        layer: RoutedExperts,
        x: torch.Tensor,
        router_logits: torch.Tensor,
        input_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert self.is_monolithic
        assert self.moe_kernel is not None
        return self.moe_kernel.apply_monolithic(
            x,
            layer.w13_weight,
            layer.w2_weight,
            router_logits,
            activation=layer.activation,
            global_num_experts=layer.global_num_experts,
            expert_map=layer.expert_map,
            apply_router_weight_on_input=layer.apply_router_weight_on_input,
            num_expert_group=layer.num_expert_group,
            topk_group=layer.topk_group,
            e_score_correction_bias=layer.e_score_correction_bias,
            routed_scaling_factor=layer.routed_scaling_factor,
        )

    def apply(
        self,
        layer: RoutedExperts,
        x: torch.Tensor,
        topk_weights: torch.Tensor,
        topk_ids: torch.Tensor,
        shared_experts: SharedExperts | None,
        shared_experts_input: torch.Tensor | None,
    ) -> torch.Tensor:
        assert not self.is_monolithic
        assert self.moe_kernel is not None
        return self.moe_kernel.apply(
            x,
            layer.w13_weight,
            layer.w2_weight,
            topk_weights,
            topk_ids,
            activation=layer.activation,
            global_num_experts=layer.global_num_experts,
            expert_map=layer.expert_map,
            apply_router_weight_on_input=layer.apply_router_weight_on_input,
            shared_experts=shared_experts,
            shared_experts_input=shared_experts_input,
        )

class ModelOptKVCacheMethod(BaseKVCacheMethod):
    """
    Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints.
    """

    def __init__(self, quant_config: "ModelOptQuantConfigBase"):
        super().__init__(quant_config)

@register_weight_loader_v2_supported_method
class ModelOptLinearMethod(LinearMethodBase):
    """Generic, format-agnostic ModelOpt linear method. Holds a weight scheme +
    an activation scheme (from the QuantSpec pair) and an optional
    ``FormatScheme``, runs a fixed lifecycle, selects the kernel from the pair,
    and applies.

    Registered for weight_loader_v2 (like every ModelOpt linear method) so the
    ``BasevLLMParameter`` params route through the v2 fused loader, not the
    legacy shape-assert path.
    """

    def __init__(
        self,
        spec: QuantSpec,
        ctx: CkptCtx,
        format_scheme=None,
    ) -> None:
        self.spec = spec
        self.ctx = ctx
        self.fmt = format_scheme or FormatScheme()
        self.wkey = SCHEME_FOR[spec.weight]
        self.akey = None if spec.activation is None else SCHEME_FOR[spec.activation]
        # Only the fp8/mxfp8 kernels read input_dtype; nvfp4 ignores it. During
        # real serving model_config is always set; fall back defensively when
        # there is no config context (bare unit-test dispatch).
        model_config = getattr(get_current_vllm_config_or_none(), "model_config", None)
        dtype = (
            model_config.dtype
            if model_config is not None
            else torch.get_default_dtype()
        )
        self.out_dtype = dtype
        self.input_dtype = dtype
        self.marlin_input_dtype = None
        # Kernel/backend are chosen in create_weights (after get_quant_method),
        # so the front-end marlin poke stays dormant here — same as the old
        # NVFP4 methods.
        self.kernel: Any = None

    @property
    def supports_pre_processed_weights(self) -> bool:  # type: ignore[override]
        # TODO(Isotr0py): support fp8/mxfp8 ModelOpt kernels transpose/repack.
        w = self.spec.weight
        return isinstance(w, QuantKey) and w.dtype == FP4_DTYPE

    def create_weights(
        self,
        layer,
        input_size_per_partition: int,
        output_partition_sizes: list[int],
        input_size: int,
        output_size: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ) -> None:
        del input_size, output_size
        weight_loader = extra_weight_attrs.get("weight_loader")
        layer.logical_widths = output_partition_sizes
        layer.input_size_per_partition = input_size_per_partition
        layer.output_size_per_partition = sum(output_partition_sizes)
        # Humming reads both off the layer in
        # prepare_humming_linear_layer_config. LinearBase sets them itself;
        # ParallelLMHead does not, so supply them here.
        layer.output_partition_sizes = output_partition_sizes
        if not hasattr(layer, "has_bias"):
            layer.has_bias = getattr(layer, "bias", None) is not None
        shapes = Shapes(output_partition_sizes, input_size_per_partition, params_dtype)

        self.wkey.create_weights(layer, WEIGHT, self.ctx, shapes, weight_loader)
        if self.akey:
            self.akey.create_weights(layer, ACT, self.ctx, shapes, weight_loader)
        self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader)

        rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype)
        self.kernel = select_linear_kernel(
            self.spec,
            layer,
            rt,
            weight_shape=self.fmt.kernel_weight_shape(layer),
        )
        expose_input_quant_key(layer, self.kernel)

    def process_weights_after_loading(self, layer) -> None:
        if is_weights_pre_processed():
            return
        self.fmt.pre_process(layer)
        self.wkey.process(layer, WEIGHT)
        if self.akey:
            self.akey.process(layer, ACT)
        maybe_fuse_global_scales(layer)
        self.fmt.post_process(layer)
        layer.is_w4a16_nvfp4 = (
            self.spec.weight == kNvfp4Static and self.spec.activation is None
        )
        if getattr(layer, "_retain_weight_for_gather", False):
            if not layer.is_w4a16_nvfp4:
                raise NotImplementedError(
                    "Gathered projection is only supported for ModelOpt W4A16 "
                    "NVFP4 weights."
                )
            assert self.ctx.group_size is not None
            layer.register_buffer(
                "_nvfp4_weight_for_gather", layer.weight.detach(), persistent=False
            )
            layer.register_buffer(
                "_nvfp4_weight_scale_for_gather",
                layer.weight_scale.detach(),
                persistent=False,
            )
            layer.register_buffer(
                "_nvfp4_weight_global_scale_for_gather",
                layer.weight_global_scale.detach(),
                persistent=False,
            )
            layer._nvfp4_group_size_for_gather = self.ctx.group_size
        if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
            self.kernel = init_mxfp8_linear_kernel(bmm_batch_size=layer.bmm_batch_size)
        self.kernel.process_weights_after_loading(layer)

    def apply(self, layer, x, bias=None):
        return self.fmt.apply(
            layer,
            x,
            bias,
            lambda lyr, inp, b: self.kernel.apply_weights(layer=lyr, x=inp, bias=b),
        )