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
[gpt-oss] Enhance error msg on attention sink init (#22335)
Signed-off-by: simon-mo <xmo@berkeley.edu> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu>
Yongye Zhu committed
Aug 6, 2025 at 11:41 UTC
31f5dc5b2a5da18bc17240c7a67e8770d00901d8
1 file changed
+9
-5
vllm/v1/attention/backends/flashinfer.py
+9
-5
@@ -638,11 +638,15 @@ class FlashInferImpl(AttentionImpl):
638
639
self.sinks: Optional[torch.Tensor] = None
640
if sinks is not None:
641
- assert sinks.shape[0] == num_heads, (
642
- "Sinks must have the same number of heads "
643
- "as the number of heads in the layer"
644
- )
645
- assert sinks.dtype == torch.float32, "Sinks must be of type float32"
641
+ if sinks.shape[0] != num_heads:
642
+ raise ValueError(
643
+ "Sinks must have the same number of heads as the number of "
644
+ f"heads in the layer. Expected {num_heads}, but got "
645
+ f"{sinks.shape[0]}."
646
+ )
647
+ if sinks.dtype != torch.float32:
648
+ raise ValueError("Sinks must be of type float32, but got "
649
+ f"{sinks.dtype}.")
650
self.sinks = sinks
651
652
def forward(