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 11, 21:15 UTC
online fp8 quant with streaming weight post-processing (#29196)
Signed-off-by: vasiliy <vasiliy@fb.com>
Vasiliy Kuznetsov committed
Dec 8, 2025 at 15:15 UTC
0d402d2600490bac17bc5d079e89b1136fe37eda
1 file changed
+66
-1
vllm/model_executor/layers/quantization/fp8.py
+66
-1
@@ -465,6 +465,30 @@ class Fp8LinearMethod(LinearMethodBase):
465
output_size_per_partition, input_size_per_partition, weight_loader
466
)
467
else:
468
+
469
+ def patched_weight_loader(param, loaded_weight, *args, **kwargs):
470
+ # load the current weight chunk
471
+ res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
472
+
473
+ # track how many elements we have updated
474
+ if not hasattr(layer, "_loaded_numel"):
475
+ layer._loaded_numel = 0
476
+ layer._loaded_numel += loaded_weight.numel()
477
+
478
+ # if we have loaded all of the elements, call
479
+ # process_weights_after_loading
480
+ target_loaded_numel = layer.weight.numel()
481
+ if layer._loaded_numel == target_loaded_numel:
482
+ self.process_weights_after_loading(layer)
483
+
484
+ # Delete the bookkeeping
485
+ del layer._loaded_numel
486
+ # Prevent the usual `process_weights_after_loading` call from doing
487
+ # anything
488
+ layer._already_called_process_weights_after_loading = True
489
+
490
+ return res
491
+
492
# For non-serialized checkpoints, use original dtype
493
weight = ModelWeightParameter(
494
data=torch.empty(
@@ -474,7 +498,7 @@ class Fp8LinearMethod(LinearMethodBase):
498
),
499
input_dim=1,
500
output_dim=0,
477
- weight_loader=weight_loader,
501
+ weight_loader=patched_weight_loader,
502
)
503
layer.register_parameter("weight", weight)
504
@@ -515,6 +539,9 @@ class Fp8LinearMethod(LinearMethodBase):
539
layer.register_parameter("input_scale", None)
540
541
def process_weights_after_loading(self, layer: Module) -> None:
542
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
543
+ return
544
+
545
size_k_first = True
546
input_scale = None
547
# TODO(rob): refactor block quant into separate class.
@@ -738,6 +765,41 @@ class Fp8MoEMethod(FusedMoEMethodBase):
765
f"weight quantization block_k = {block_k}."
766
)
767
768
+ # if we are doing online quantization, patch the weight
769
+ # loaded to call `process_weights_after_loading` in a streaming fashion
770
+ # as soon as the last weight chunk is loaded
771
+ if not self.quant_config.is_checkpoint_fp8_serialized:
772
+ weight_loader = extra_weight_attrs["weight_loader"]
773
+ # create a new holder to prevent modifying behavior of any other
774
+ # objects which might depend on the old one
775
+ new_extra_weight_attrs = extra_weight_attrs
776
+
777
+ def patched_weight_loader(param, loaded_weight, *args, **kwargs):
778
+ # load the current weight chunk
779
+ res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]
780
+
781
+ # add a counter to track how many elements we have updated
782
+ if not hasattr(layer, "_loaded_numel"):
783
+ layer._loaded_numel = 0
784
+ layer._loaded_numel += loaded_weight.numel()
785
+
786
+ # if we have loaded all of the elements, call
787
+ # process_weights_after_loading
788
+ target_loaded_numel = layer.w13_weight.numel() + layer.w2_weight.numel()
789
+ if layer._loaded_numel == target_loaded_numel:
790
+ self.process_weights_after_loading(layer)
791
+
792
+ # Delete the bookkeeping
793
+ del layer._loaded_numel
794
+ # Prevent the usual `process_weights_after_loading` call
795
+ # from doing anything
796
+ layer._already_called_process_weights_after_loading = True
797
+
798
+ return res
799
+
800
+ new_extra_weight_attrs["weight_loader"] = patched_weight_loader
801
+ extra_weight_attrs = new_extra_weight_attrs
802
+
803
# WEIGHTS
804
w13_weight = torch.nn.Parameter(
805
torch.empty(
@@ -839,6 +901,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
901
self.rocm_aiter_moe_enabled = False
902
903
def process_weights_after_loading(self, layer: Module) -> None:
904
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
905
+ return
906
+
907
# Lazy import to avoid importing triton too early.
908
909
self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()