vlm: streamline vision sdpa reshapes (#34991)

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

@@ -401,7 +401,9 @@ class VisionSdpaAttention(nn.Module):
         else:
             attention_mask = attention_mask.to(device=q.device)
 
-        q, k, v = [rearrange(x, "(b s) h d -> b h s d", b=bsz) for x in [q, k, v]]
+        q = q.reshape(bsz, s, self.num_heads, self.head_size).transpose(1, 2)
+        k = k.reshape(bsz, s, self.num_kv_heads, self.head_size).transpose(1, 2)
+        v = v.reshape(bsz, s, self.num_kv_heads, self.head_size).transpose(1, 2)
 
         if self.softmax_in_single_precision:
             k = rearrange(k, "b h s d -> b h d s")
@@ -434,7 +436,7 @@ class VisionSdpaAttention(nn.Module):
             )
 
         # [b, h, s, head_size] --> [b * s, h, head_size]
-        output = rearrange(output, "b h s d -> (b s) h d")
+        output = output.transpose(1, 2).reshape(bsz * s, self.num_heads, self.head_size)
 
         return output
 
@@ -1477,7 +1479,7 @@ class VisionAttention(nn.Module):
 
         if self.use_qkv_parallel:
             # [b * s, h, head_size] --> [b, s, h * head_size]
-            output = rearrange(output, "(b s) ... h d -> b s ... (h d)", b=bsz)
+            output = output.reshape(bsz, s, -1)
 
             # [b, s, h * head_size] --> [b, s, h * head_size]
             output, _ = self.proj(output)

test/registered/unit/layers/attention/test_vision_backend_selection.py

@@ -52,6 +52,31 @@ def test_npu_backend_selection_priority(
     assert backend == expected
 
 
+def test_sdpa_preserves_flattened_batch_layout():
+    torch.manual_seed(0)
+    bsz, seq_len, num_heads, head_dim = 3, 5, 2, 8
+    q, k, v = [torch.randn(bsz * seq_len, num_heads, head_dim) for _ in range(3)]
+    backend = vision.VisionSdpaAttention(
+        head_dim=head_dim,
+        num_heads=num_heads,
+        num_kv_heads=num_heads,
+    )
+
+    output = backend(q=q, k=k, v=v, bsz=bsz)
+    q_ref, k_ref, v_ref = [
+        x.reshape(bsz, seq_len, num_heads, head_dim).transpose(1, 2) for x in (q, k, v)
+    ]
+    expected = F.scaled_dot_product_attention(
+        q_ref,
+        k_ref,
+        v_ref,
+        scale=backend.scale,
+    )
+    expected = expected.transpose(1, 2).reshape(bsz * seq_len, num_heads, head_dim)
+
+    torch.testing.assert_close(output, expected)
+
+
 @pytest.mark.parametrize("mask_kind", ["causal", "padding"])
 def test_ascend_attention_masked_inputs_fall_back_to_sdpa(
     monkeypatch,