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
[Bugfix gpt-oss] Fix float32 convert for flashinfer sink support (#23016)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 16, 2025 at 14:16 UTC
000cceca8c329d5b5d99e0186fbd444a390384cd
2 files changed
+9
-3
vllm/attention/layer.py
+9
@@ -308,6 +308,15 @@ class Attention(nn.Module):
308
if hasattr(self.impl, "process_weights_after_loading"):
309
self.impl.process_weights_after_loading(act_dtype)
310
311
+ # FlashInfer requires attention sinks to be float32
312
+ if (self.backend == _Backend.FLASHINFER_VLLM_V1
313
+ and hasattr(self.impl, 'sinks')):
314
+ from vllm.v1.attention.backends.flashinfer import FlashInferImpl
315
+ assert isinstance(self.impl, FlashInferImpl)
316
+ if (self.impl.sinks is not None
317
+ and self.impl.sinks.dtype != torch.float32):
318
+ self.impl.sinks = self.impl.sinks.to(torch.float32)
319
+
320
def get_attn_backend(self) -> type[AttentionBackend]:
321
return self.attn_backend
322
vllm/v1/attention/backends/flashinfer.py
-3
@@ -642,9 +642,6 @@ class FlashInferImpl(AttentionImpl):
642
f"heads in the layer. Expected {num_heads}, but got "
643
f"{sinks.shape[0]}."
644
)
645
- # Cast sinks to float32 if needed (FlashInfer requirement)
646
- if sinks.dtype != torch.float32:
647
- sinks = sinks.to(torch.float32)
645
self.sinks = sinks
646
647
def forward(