Add an assertion to enhance the robustness of the operator (#5736)

Author: liwenju0

python/sglang/srt/layers/attention/vision.py

@@ -271,6 +271,8 @@ def forward(
         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