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 6, 14:30 UTC
Bugfix for offloading+prefetch for GLM-4.7-FP8 (#37178)
Signed-off-by: Benjamin Merkel <benjamin.merkel@tngtech.com> Co-authored-by: Benjamin Merkel <benjamin.merkel@tngtech.com>
sfbemerk committed
Mar 17, 2026 at 14:22 UTC
2660b9289c1f9e26ae65a247ceac2b9add52fa90
1 file changed
+42
-1
vllm/model_executor/offloader/prefetch.py
+42
-1
@@ -431,10 +431,32 @@ class _ModuleOffloader:
431
432
Called after process_weights_after_loading to ensure _cpu_storage
433
contains the final processed weights, not stale pre-loading data.
434
+
435
+ Parameters whose underlying nn.Parameter was deleted by
436
+ process_weights_after_loading (e.g. transient KV-cache scale params)
437
+ are pruned from self._param_offloaders so they do not participate in
438
+ buffer-pool allocation or prefetching.
439
"""
440
for param_offloader in self._param_offloaders.values():
441
param_offloader.sync_cpu_storage()
442
443
+ # Remove offloaders whose parameter was deleted during
444
+ # process_weights_after_loading (e.g. k_scale / v_scale).
445
+ deleted = [
446
+ name
447
+ for name, offloader in self._param_offloaders.items()
448
+ if getattr(offloader, "_param_deleted", False)
449
+ ]
450
+ if deleted:
451
+ logger.debug(
452
+ "Pruning %d transient offloaded param(s) that were deleted "
453
+ "by process_weights_after_loading: %s",
454
+ len(deleted),
455
+ deleted,
456
+ )
457
+ for name in deleted:
458
+ del self._param_offloaders[name]
459
+
460
def get_param_infos(self) -> list[ParamInfo]:
461
"""Get parameter metadata for buffer pool allocation.
462
@@ -590,6 +612,11 @@ class _CpuParamOffloader(_BaseParamOffloader):
612
super().__init__(module, param_name)
613
self._cpu_storage: torch.Tensor | None = None
614
self._gpu_buffer: torch.Tensor | None = None # Store reference to GPU buffer
615
+ # Set to True if the underlying nn.Parameter was deleted by
616
+ # process_weights_after_loading (e.g. transient KV-cache scale params
617
+ # such as k_scale/v_scale created by BaseKVCacheMethod.create_weights
618
+ # and deleted after copying into permanent _k_scale buffers).
619
+ self._param_deleted: bool = False
620
621
# Offload to CPU immediately to free GPU memory during model loading
622
self._offload_to_cpu_internal()
@@ -696,8 +723,22 @@ class _CpuParamOffloader(_BaseParamOffloader):
723
1. process_weights_after_loading may transform weights (quantization)
724
2. device_loading_context creates NEW CPU tensors when moving back
725
3. Our old _cpu_storage would have pre-processed or stale data
726
+
727
+ If the parameter no longer exists on the module (e.g. transient
728
+ KV-cache scale parameters such as k_scale/v_scale that are created
729
+ by BaseKVCacheMethod.create_weights() and then deleted by
730
+ process_weights_after_loading() after copying their values into
731
+ permanent _k_scale buffers), the offloader marks itself as deleted
732
+ and skips the sync. The caller (_ModuleOffloader.sync_cpu_storage)
733
+ is responsible for removing these stale entries.
734
"""
700
- self._update_cpu_storage_from_param()
735
+ try:
736
+ self._update_cpu_storage_from_param()
737
+ except AttributeError:
738
+ # The parameter was deleted by process_weights_after_loading.
739
+ # Drop the now-stale CPU storage so this offloader can be pruned.
740
+ self._param_deleted = True
741
+ self._cpu_storage = None
742
743
def post_init(self):
744
"""No-op: offloading done in offload_to_cpu/assign_static_buffer."""