A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 23:00 UTC
fp8 online quant: split out Fp8OnlineLinearMethod (#32189)
Vasiliy Kuznetsov committed
Jan 20, 2026 at 18:13 UTC
d2389c12622090a3f5985f2c424afddf5e447d53
2 files changed
+143
-110
tests/quantization/test_fp8.py
+4
-1
@@ -133,7 +133,7 @@ def test_kv_cache_model_load_and_run(
133
@pytest.mark.parametrize(
134
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
135
)
136
-def test_load_fp16_model(
136
+def test_online_quantization(
137
vllm_runner,
138
kv_cache_dtype: str,
139
force_marlin: bool,
@@ -191,6 +191,9 @@ def test_load_fp16_model(
191
192
llm.apply_model(check_model)
193
194
+ outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
195
+ print(outputs[0][1])
196
+
197
198
@pytest.mark.skipif(
199
not is_quant_method_supported("fp8"),
vllm/model_executor/layers/quantization/fp8.py
+139
-109
@@ -230,9 +230,14 @@ class Fp8Config(QuantizationConfig):
230
fused_mapping=self.packed_modules_mapping,
231
):
232
return UnquantizedLinearMethod()
233
- quant_method = Fp8LinearMethod(self)
234
- quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
235
- return quant_method
233
+ if not self.is_checkpoint_fp8_serialized:
234
+ online_method = Fp8OnlineLinearMethod(self)
235
+ online_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
236
+ return online_method
237
+ else:
238
+ offline_method = Fp8LinearMethod(self)
239
+ offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
240
+ return offline_method
241
elif isinstance(layer, FusedMoE):
242
if is_layer_skipped(
243
prefix=prefix,
@@ -295,13 +300,8 @@ class Fp8LinearMethod(LinearMethodBase):
300
Supports loading FP8 checkpoints with static weight scale and
301
dynamic/static activation scale.
302
298
- Also supports loading quantized FP16/BF16 model checkpoints with dynamic
299
- activation scaling. The weight scaling factor will be initialized after
300
- the model weights are loaded.
301
-
303
Limitations:
303
- 1. Only support per-tensor quantization due to torch._scaled_mm support.
304
- 2. Only support float8_e4m3fn data type due to the limitation of
304
+ 1. Only support float8_e4m3fn data type due to the limitation of
305
torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
306
307
Args:
@@ -388,89 +388,43 @@ class Fp8LinearMethod(LinearMethodBase):
388
self.weight_block_size,
389
)
390
391
- # WEIGHT
392
- if self.quant_config.is_checkpoint_fp8_serialized:
393
- weight = create_fp8_weight_parameter(
394
- output_size_per_partition, input_size_per_partition, weight_loader
391
+ weight = create_fp8_weight_parameter(
392
+ output_size_per_partition, input_size_per_partition, weight_loader
393
+ )
394
+ layer.register_parameter("weight", weight)
395
+
396
+ # WEIGHT SCALE
397
+ if not self.block_quant:
398
+ scale = create_fp8_scale_parameter(
399
+ PerTensorScaleParameter,
400
+ output_partition_sizes,
401
+ input_size_per_partition,
402
+ None,
403
+ weight_loader,
404
)
405
+ set_weight_attrs(scale, {"scale_type": "weight_scale"})
406
+ layer.register_parameter("weight_scale", scale)
407
else:
397
-
398
- def patched_weight_loader(param, loaded_weight, *args, **kwargs):
399
- # track how many elements we have updated
400
- if not hasattr(layer, "_loaded_numel"):
401
- layer._loaded_numel = 0
402
-
403
- # load the current weight chunk
404
- copy_numel_counter = CopyNumelCounter()
405
- with copy_numel_counter:
406
- res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
407
- layer._loaded_numel += copy_numel_counter.copied_numel
408
-
409
- # if we have loaded all of the elements, call
410
- # process_weights_after_loading
411
- target_loaded_numel = layer.weight.numel()
412
- if layer._loaded_numel == target_loaded_numel:
413
- self.process_weights_after_loading(layer)
414
-
415
- # Delete the bookkeeping
416
- del layer._loaded_numel
417
- # Prevent the usual `process_weights_after_loading` call from doing
418
- # anything
419
- layer._already_called_process_weights_after_loading = True
420
-
421
- return res
422
-
423
- # For non-serialized checkpoints, use original dtype
424
- weight = ModelWeightParameter(
425
- data=torch.empty(
426
- output_size_per_partition,
427
- input_size_per_partition,
428
- dtype=params_dtype,
429
- ),
430
- input_dim=1,
431
- output_dim=0,
432
- weight_loader=patched_weight_loader,
408
+ assert not self.act_q_static
409
+ assert self.weight_block_size is not None
410
+ scale = create_fp8_scale_parameter(
411
+ BlockQuantScaleParameter,
412
+ output_partition_sizes,
413
+ input_size_per_partition,
414
+ self.weight_block_size,
415
+ weight_loader,
416
)
434
- layer.register_parameter("weight", weight)
417
+ set_weight_attrs(scale, {"scale_type": "weight_scale"})
418
+ # The weight_scale_inv name is intentional for deepseekv3
419
+ layer.register_parameter("weight_scale_inv", scale)
420
436
- # If checkpoint is serialized fp8, load them.
437
- # Otherwise, wait until process_weights_after_loading.
438
- if self.quant_config.is_checkpoint_fp8_serialized:
439
- # WEIGHT SCALE
440
- if not self.block_quant:
441
- scale = create_fp8_scale_parameter(
442
- PerTensorScaleParameter,
443
- output_partition_sizes,
444
- input_size_per_partition,
445
- None,
446
- weight_loader,
447
- )
448
- set_weight_attrs(scale, {"scale_type": "weight_scale"})
449
- layer.register_parameter("weight_scale", scale)
450
- else:
451
- assert not self.act_q_static
452
- assert self.weight_block_size is not None
453
- scale = create_fp8_scale_parameter(
454
- BlockQuantScaleParameter,
455
- output_partition_sizes,
456
- input_size_per_partition,
457
- self.weight_block_size,
458
- weight_loader,
459
- )
460
- set_weight_attrs(scale, {"scale_type": "weight_scale"})
461
- # The weight_scale_inv name is intentional for deepseekv3
462
- layer.register_parameter("weight_scale_inv", scale)
463
-
464
- # INPUT ACTIVATION SCALE
465
- if self.act_q_static:
466
- scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
467
- set_weight_attrs(scale, {"scale_type": "input_scale"})
468
- layer.register_parameter("input_scale", scale)
421
+ # INPUT ACTIVATION SCALE
422
+ if self.act_q_static:
423
+ scale = create_fp8_input_scale(output_partition_sizes, weight_loader)
424
+ set_weight_attrs(scale, {"scale_type": "input_scale"})
425
+ layer.register_parameter("input_scale", scale)
426
427
def process_weights_after_loading(self, layer: Module) -> None:
471
- if getattr(layer, "_already_called_process_weights_after_loading", False):
472
- return
473
-
428
size_k_first = True
429
input_scale = None
430
# TODO(rob): refactor block quant into separate class.
@@ -488,31 +442,24 @@ class Fp8LinearMethod(LinearMethodBase):
442
443
# If checkpoint not serialized fp8, quantize the weights.
444
else:
491
- if not self.quant_config.is_checkpoint_fp8_serialized:
492
- qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
493
- weight = qweight.t()
494
-
445
# If checkpoint is fp8 per-tensor, handle that there are N scales for N
446
# shards in a fused module
497
- else:
498
- weight = layer.weight
499
- weight_scale = layer.weight_scale
500
-
501
- # If using w8a8, torch._scaled_mm needs per tensor, so
502
- # requantize the logical shards as a single weight.
503
- if not self.use_marlin:
504
- weight, weight_scale, input_scale = (
505
- process_fp8_weight_tensor_strategy(
506
- weight,
507
- weight_scale,
508
- layer.logical_widths,
509
- getattr(layer, "input_scale", None),
510
- )
511
- )
512
- if self.act_q_static:
513
- assert input_scale is not None
514
- input_scale = input_scale.max()
515
- weight = weight.t()
447
+ weight = layer.weight
448
+ weight_scale = layer.weight_scale
449
+
450
+ # If using w8a8, torch._scaled_mm needs per tensor, so
451
+ # requantize the logical shards as a single weight.
452
+ if not self.use_marlin:
453
+ weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
454
+ weight,
455
+ weight_scale,
456
+ layer.logical_widths,
457
+ getattr(layer, "input_scale", None),
458
+ )
459
+ if self.act_q_static:
460
+ assert input_scale is not None
461
+ input_scale = input_scale.max()
462
+ weight = weight.t()
463
464
# Update layer with new values.
465
replace_parameter(layer, "weight", weight.data)
@@ -607,6 +554,89 @@ class Fp8LinearMethod(LinearMethodBase):
554
return self.fp8_linear.apply_weights(layer, x, bias)
555
556
557
+class Fp8OnlineLinearMethod(Fp8LinearMethod):
558
+ """Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint
559
+ and quantized the weights during loading."""
560
+
561
+ def create_weights(
562
+ self,
563
+ layer: torch.nn.Module,
564
+ input_size_per_partition: int,
565
+ output_partition_sizes: list[int],
566
+ input_size: int,
567
+ output_size: int,
568
+ params_dtype: torch.dtype,
569
+ **extra_weight_attrs,
570
+ ):
571
+ output_size_per_partition = sum(output_partition_sizes)
572
+ weight_loader = extra_weight_attrs.get("weight_loader")
573
+ layer.logical_widths = output_partition_sizes
574
+ layer.input_size_per_partition = input_size_per_partition
575
+ layer.output_size_per_partition = output_size_per_partition
576
+ layer.orig_dtype = params_dtype
577
+ layer.weight_block_size = None
578
+
579
+ # WEIGHT
580
+ def patched_weight_loader(param, loaded_weight, *args, **kwargs):
581
+ # track how many elements we have updated
582
+ if not hasattr(layer, "_loaded_numel"):
583
+ layer._loaded_numel = 0
584
+
585
+ # load the current weight chunk
586
+ copy_numel_counter = CopyNumelCounter()
587
+ with copy_numel_counter:
588
+ res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
589
+ layer._loaded_numel += copy_numel_counter.copied_numel
590
+
591
+ # if we have loaded all of the elements, call
592
+ # process_weights_after_loading
593
+ target_loaded_numel = layer.weight.numel()
594
+ if layer._loaded_numel == target_loaded_numel:
595
+ self.process_weights_after_loading(layer)
596
+
597
+ # Delete the bookkeeping
598
+ del layer._loaded_numel
599
+ # Prevent the usual `process_weights_after_loading` call from doing
600
+ # anything
601
+ layer._already_called_process_weights_after_loading = True
602
+
603
+ return res
604
+
605
+ weight = ModelWeightParameter(
606
+ data=torch.empty(
607
+ output_size_per_partition,
608
+ input_size_per_partition,
609
+ dtype=params_dtype,
610
+ ),
611
+ input_dim=1,
612
+ output_dim=0,
613
+ weight_loader=patched_weight_loader,
614
+ )
615
+ layer.register_parameter("weight", weight)
616
+
617
+ def process_weights_after_loading(self, layer: Module) -> None:
618
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
619
+ return
620
+
621
+ # TODO(future): support block_quant in online quant path
622
+ assert not self.block_quant
623
+
624
+ layer.input_scale = None
625
+ qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
626
+ weight = qweight.t()
627
+
628
+ # Update layer with new values.
629
+ replace_parameter(layer, "weight", weight.data)
630
+ replace_parameter(layer, "weight_scale", weight_scale.data)
631
+
632
+ if self.use_marlin:
633
+ size_k_first = True
634
+ prepare_fp8_layer_for_marlin(
635
+ layer, size_k_first, input_dtype=self.marlin_input_dtype
636
+ )
637
+ # Activations not quantized for marlin.
638
+
639
+
640
class Fp8MoEMethod(FusedMoEMethodBase):
641
"""MoE method for FP8.
642
Supports loading FP8 checkpoints with static weight scale and