From 4d1e52abea8277e69de281cb23634edb723fcd85 Mon Sep 17 00:00:00 2001 From: liwenju0 Date: Sun, 27 Apr 2025 09:09:12 +0800 Subject: [PATCH] Add an assertion to enhance the robustness of the operator (#5736) --- python/sglang/srt/layers/attention/vision.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 860994913..d65104beb 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -271,6 +271,8 @@ class VisionSdpaAttention(nn.Module): Returns: [b * s, h, head_size] """ + if self.flatten_batch: + assert bsz == 1, "flatten_batch is True, bsz must be 1" s = q.shape[0] // bsz