@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3f53e2138f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Refactor] Remove `Fp8OnlineLinearMethod` as scheduled (#45463)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Wentao Ye committed Jun 16, 2026 at 07:35 UTC 3f53e2138f09843c5bddb4611df58e11bed01500

5 files changed +25 -82

docs/training/layerwise.md

+3 -3

@@ -28,9 +28,9 @@ For more information on implementation, see [Low Level `layerwise` API](#low-lev

28 Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this:

30 ```python

31 -class Fp8OnlineLinearMethod(Fp8LinearMethod):

32 - """Online version of Fp8LinearMethod which loads a full precision checkpoint

33 - and quantizes weights during loading."""

31 +class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):

32 + """Online version of FP8 per-tensor quantization which loads a full

33 + precision checkpoint and quantizes weights during loading."""

35 uses_meta_device: bool = True

tests/quantization/test_fp8.py

+4 -1

@@ -21,6 +21,9 @@ from vllm.model_executor.layers.quantization.fp8 import (

21 Fp8LinearMethod,

22 Fp8MoEMethod,

23 )

24 +from vllm.model_executor.layers.quantization.online.fp8 import (

25 + Fp8PerTensorOnlineLinearMethod,

26 +)

27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

28 from vllm.platforms import current_platform

@@ -164,7 +167,7 @@ def test_online_quantization(

167

168 def check_model(model):

169 fc1 = model.model.decoder.layers[0].fc1

167 - assert isinstance(fc1.quant_method, Fp8LinearMethod)

170 + assert isinstance(fc1.quant_method, Fp8PerTensorOnlineLinearMethod)

171 if kv_cache_dtype == "fp8":

172 attn = model.model.decoder.layers[0].self_attn.attn

173 assert isinstance(attn.quant_method, Fp8KVCacheMethod)

vllm/model_executor/layers/quantization/fp8.py

+5 -77

@@ -80,7 +80,6 @@ from vllm.model_executor.model_loader.reload.layerwise import (

80 )

81 from vllm.model_executor.parameter import (

82 BlockQuantScaleParameter,

83 - ModelWeightParameter,

83 PerTensorScaleParameter,

84 )

85 from vllm.model_executor.utils import replace_parameter, set_weight_attrs

@@ -188,7 +187,11 @@ class Fp8Config(QuantizationConfig):

187 ):

188 return UnquantizedLinearMethod()

189 if not self.is_checkpoint_fp8_serialized:

191 - online_method = Fp8OnlineLinearMethod(self)

190 + from vllm.model_executor.layers.quantization.online.fp8 import (

191 + Fp8PerTensorOnlineLinearMethod,

192 + )

193 +

194 + online_method = Fp8PerTensorOnlineLinearMethod()

195 online_method.marlin_input_dtype = get_marlin_input_dtype(prefix)

196 return online_method

197 else:

@@ -489,81 +492,6 @@ class Fp8LinearMethod(LinearMethodBase):

492 return self.fp8_linear.apply_weights(layer, x, bias)

493

494

492 -# TODO(future PR): remove this class in favor of

493 -# online/fp8.py::Fp8PerTensorOnlineLinearMethod

494 -class Fp8OnlineLinearMethod(Fp8LinearMethod):

495 - """Online version of Fp8LinearMethod which loads a full precision checkpoint

496 - and quantizes weights during loading."""

497 -

498 - uses_meta_device: bool = True

499 -

500 - def create_weights(

501 - self,

502 - layer: torch.nn.Module,

503 - input_size_per_partition: int,

504 - output_partition_sizes: list[int],

505 - input_size: int,

506 - output_size: int,

507 - params_dtype: torch.dtype,

508 - **extra_weight_attrs,

509 - ):

510 - output_size_per_partition = sum(output_partition_sizes)

511 - weight_loader = extra_weight_attrs.get("weight_loader")

512 - layer.logical_widths = output_partition_sizes

513 - layer.input_size_per_partition = input_size_per_partition

514 - layer.output_size_per_partition = output_size_per_partition

515 - layer.orig_dtype = params_dtype

516 - layer.weight_block_size = None

517 -

518 - weight = ModelWeightParameter(

519 - data=torch.empty(

520 - output_size_per_partition,

521 - input_size_per_partition,

522 - device="meta", # materialized and processed during loading

523 - dtype=params_dtype,

524 - ),

525 - input_dim=1,

526 - output_dim=0,

527 - weight_loader=weight_loader,

528 - )

529 - layer.register_parameter("weight", weight)

530 -

531 - initialize_online_processing(layer)

532 -

533 - self.fp8_linear = init_fp8_linear_kernel(

534 - activation_quant_key=self.activation_quant_key,

535 - weight_quant_key=self.weight_quant_key,

536 - weight_shape=layer.weight.shape,

537 - input_dtype=self.input_dtype,

538 - out_dtype=self.out_dtype,

539 - module_name=self.__class__.__name__,

540 - )

541 - self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)

542 -

543 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

544 - if getattr(layer, "_already_called_process_weights_after_loading", False):

545 - return

546 -

547 - # TODO(future): support block_quant in online quant path

548 - assert not self.block_quant

549 -

550 - layer.input_scale = None

551 - qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)

552 -

553 - # Update layer with new values. Canonicalize to (K, N) for the kernel.

554 - replace_parameter(layer, "weight", qweight.t().data)

555 - replace_parameter(layer, "weight_scale", weight_scale.data)

556 -

557 - if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):

558 - # Only Marlin kernels support `marlin_input_dtype`; guard to avoid

559 - # AttributeError if backend selection changes.

560 - self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype

561 - self.fp8_linear.process_weights_after_loading(layer)

562 -

563 - # Prevent duplicate processing (e.g., during weight reload)

564 - layer._already_called_process_weights_after_loading = True

565 -

566 -

495 class Fp8MoEMethod(FusedMoEMethodBase):

496 """MoE method for FP8.

497 Supports loading FP8 checkpoints with static weight scale and

vllm/model_executor/layers/quantization/online/fp8.py

@@ -109,6 +109,10 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):

109 def __init__(self):

110 super().__init__()

111

112 + self.block_quant = False

113 + self.use_deep_gemm = False

114 + self.use_marlin = False

115 + self.marlin_input_dtype = None

116 self.weight_quant_key = kFp8StaticTensorSym

117 # Use per-token quantization for better perf if dynamic and cutlass

118 if cutlass_fp8_supported():

@@ -144,6 +148,7 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):

148 out_dtype=self.out_dtype,

149 module_name=self.__class__.__name__,

150 )

151 + self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)

152

153 def process_weights_after_loading(self, layer: Module) -> None:

154 if getattr(layer, "_already_called_process_weights_after_loading", False):

@@ -156,6 +161,8 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):

161 replace_parameter(layer, "weight", qweight.t().data)

162 replace_parameter(layer, "weight_scale", weight_scale.data)

163

164 + if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):

165 + self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype

166 self.fp8_linear.process_weights_after_loading(layer)

167

168 # Prevent duplicate processing (e.g., during weight reload)

vllm/model_executor/layers/quantization/utils/quant_utils.py

+6 -1

@@ -401,6 +401,9 @@ def get_and_maybe_dequant_weights(

401 """Return layer's unquantized weights in [out, in] layout"""

402 from vllm.model_executor.layers.linear import UnquantizedLinearMethod

403 from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod

404 + from vllm.model_executor.layers.quantization.online.fp8 import (

405 + Fp8PerTensorOnlineLinearMethod,

406 + )

407

408 # LoRA linear wrappers store quantization metadata on `base_layer`.

409 # Unwrap here so callers can pass either a raw linear layer or its LoRA

@@ -418,7 +421,9 @@ def get_and_maybe_dequant_weights(

421

422 # Simple Fp8 case: rescale with tensor or block weight scales

423 if (

421 - isinstance(layer.quant_method, Fp8LinearMethod)

424 + isinstance(

425 + layer.quant_method, (Fp8LinearMethod, Fp8PerTensorOnlineLinearMethod)

426 + )

427 and not layer.quant_method.use_marlin

428 # DeepGEMM transforms the scales using `transform_sf_into_required_layout` into

429 # a layout that is not compatible with `scaled_dequantize`.