From 968b355f129b491c507a5371a85c5141c22b7de7 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 16 Aug 2026 19:05:53 +0800 Subject: [PATCH] vlm: streamline vision sdpa reshapes (#34991) --- python/sglang/srt/layers/attention/vision.py | 8 +++--- .../test_vision_backend_selection.py | 25 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 0d69dea7d..5008fad4e 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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) diff --git a/test/registered/unit/layers/attention/test_vision_backend_selection.py b/test/registered/unit/layers/attention/test_vision_backend_selection.py index 3ed5a4d64..1b894f548 100644 --- a/test/registered/unit/layers/attention/test_vision_backend_selection.py +++ b/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,