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 14, 05:30 UTC
[Feature] Add process_weights_after_loading to AttentionImpl (#26870)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
Oct 16, 2025 at 23:02 UTC
5afd3276dfd70397a08ba16ee8eb246ddb3d13ef
3 files changed
+9
-10
vllm/attention/backends/abstract.py
+3
@@ -207,6 +207,9 @@ class AttentionImpl(ABC, Generic[T]):
207
"""
208
return False
209
210
+ def process_weights_after_loading(self, act_dtype: torch.dtype):
211
+ pass
212
+
213
214
class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
215
@abstractmethod
vllm/attention/layer.py
+1
-10
@@ -404,16 +404,7 @@ class Attention(nn.Module, AttentionLayerBase):
404
return s
405
406
def process_weights_after_loading(self, act_dtype: torch.dtype):
407
- if hasattr(self.impl, "process_weights_after_loading"):
408
- self.impl.process_weights_after_loading(act_dtype)
409
-
410
- # FlashInfer requires attention sinks to be float32
411
- if self.backend == _Backend.FLASHINFER and hasattr(self.impl, "sinks"):
412
- from vllm.v1.attention.backends.flashinfer import FlashInferImpl
413
-
414
- assert isinstance(self.impl, FlashInferImpl)
415
- if self.impl.sinks is not None and self.impl.sinks.dtype != torch.float32:
416
- self.impl.sinks = self.impl.sinks.to(torch.float32)
407
+ self.impl.process_weights_after_loading(act_dtype)
408
409
def get_attn_backend(self) -> type[AttentionBackend]:
410
return self.attn_backend
vllm/v1/attention/backends/flashinfer.py
+5
@@ -833,6 +833,11 @@ class FlashInferImpl(AttentionImpl):
833
834
return self.support_trtllm_attn
835
836
+ # FlashInfer requires attention sinks to be float32
837
+ def process_weights_after_loading(self, act_dtype: torch.dtype):
838
+ if self.sinks is not None and self.sinks.dtype != torch.float32:
839
+ self.sinks = self.sinks.to(torch.float32)
840
+
841
def forward(
842
self,
843
layer: torch.nn.Module,