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 8, 22:45 UTC
Implement zero-copy GQA for multimodal and CPU (#33732)
Signed-off-by: Taeksang Kim <ts.kim@hyperaccel.ai>
Taeksang Kim committed
Feb 5, 2026 at 05:11 UTC
6e98f6d8b64984d5e3a8a9b323321f1095bac8b3
4 files changed
+18
-32
vllm/model_executor/layers/attention/mm_encoder_attention.py
+4
-12
@@ -80,7 +80,7 @@ class MMEncoderAttention(CustomOp):
80
def enabled(cls) -> bool:
81
return True
82
83
- def maybe_reshape_qkv_to_4d(
83
+ def view_qkv_to_4d(
84
self,
85
query: torch.Tensor,
86
key: torch.Tensor,
@@ -97,11 +97,6 @@ class MMEncoderAttention(CustomOp):
97
key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
98
value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)
99
100
- if (num_repeat := self.num_queries_per_kv) > 1:
101
- # Handle MQA and GQA
102
- key = torch.repeat_interleave(key, num_repeat, dim=2)
103
- value = torch.repeat_interleave(value, num_repeat, dim=2)
104
-
100
return query, key, value
101
102
def _forward_sdpa(
@@ -119,9 +114,7 @@ class MMEncoderAttention(CustomOp):
114
kv_len = key.size(1)
115
is_reshaped = query.dim() != 4
116
122
- query, key, value = self.maybe_reshape_qkv_to_4d(
123
- query, key, value, bsz, q_len, kv_len
124
- )
117
+ query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len)
118
119
output = vit_torch_sdpa_wrapper(
120
q=query,
@@ -129,6 +122,7 @@ class MMEncoderAttention(CustomOp):
122
v=value,
123
scale=self.scale,
124
cu_seqlens=cu_seqlens,
125
+ enable_gqa=self.num_heads > self.num_kv_heads,
126
)
127
if is_reshaped:
128
output = output.reshape(bsz, q_len, -1)
@@ -154,9 +148,7 @@ class MMEncoderAttention(CustomOp):
148
kv_len = key.size(1)
149
is_reshaped = query.dim() != 4
150
157
- query, key, value = self.maybe_reshape_qkv_to_4d(
158
- query, key, value, bsz, q_len, kv_len
159
- )
151
+ query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len)
152
153
output = vit_flash_attn_wrapper(
154
q=query,
vllm/model_executor/models/molmo2.py
+1
-12
@@ -628,18 +628,6 @@ class ImagePoolingAttention(nn.Module):
628
key = key.view(bsz, kv_len, self.num_kv_heads, self.head_dim)
629
value = value.view(bsz, kv_len, self.num_kv_heads, self.head_dim)
630
631
- if self.num_heads != self.num_kv_heads:
632
- key = torch.repeat_interleave(
633
- key,
634
- self.num_heads // self.num_kv_heads,
635
- dim=2,
636
- )
637
- value = torch.repeat_interleave(
638
- value,
639
- self.num_heads // self.num_kv_heads,
640
- dim=2,
641
- )
642
-
631
query, key, value = (x.transpose(1, 2) for x in (query, key, value))
632
633
out = F.scaled_dot_product_attention(
@@ -648,6 +636,7 @@ class ImagePoolingAttention(nn.Module):
636
value,
637
attn_mask=attn_mask,
638
is_causal=False,
639
+ enable_gqa=self.num_heads > self.num_kv_heads,
640
).transpose(1, 2)
641
642
return out.reshape(bsz, q_len, -1)
vllm/v1/attention/backends/cpu_attn.py
+1
-4
@@ -398,10 +398,6 @@ class CPUAttentionBackendImpl(AttentionImpl):
398
key = key.movedim(0, key.dim() - 2)
399
value = value.movedim(0, value.dim() - 2)
400
401
- if self.num_kv_heads != self.num_heads:
402
- key = key.repeat_interleave(self.num_queries_per_kv, dim=-3)
403
- value = value.repeat_interleave(self.num_queries_per_kv, dim=-3)
404
-
401
causal_attn = attn_type == AttentionType.DECODER
402
403
sdpa_start_loc = attn_metadata.sdpa_start_loc.numpy() # type: ignore
@@ -418,6 +414,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
414
dropout_p=0.0,
415
is_causal=causal_attn and mask is None,
416
scale=self.scale,
417
+ enable_gqa=self.num_heads > self.num_kv_heads,
418
)
419
.squeeze(0)
420
.movedim(query.dim() - 2, 0)
vllm/v1/attention/ops/vit_attn_wrappers.py
+12
-4
@@ -115,13 +115,16 @@ def apply_sdpa(
115
k: torch.Tensor,
116
v: torch.Tensor,
117
scale: float | None = None,
118
+ enable_gqa: bool = False,
119
) -> torch.Tensor:
120
"""
121
Input shape:
122
(batch_size x seq_len x num_heads x head_size)
123
"""
124
q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
124
- output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, scale=scale)
125
+ output = F.scaled_dot_product_attention(
126
+ q, k, v, dropout_p=0.0, scale=scale, enable_gqa=enable_gqa
127
+ )
128
output = einops.rearrange(output, "b h s d -> b s h d ")
129
return output
130
@@ -134,6 +137,7 @@ def torch_sdpa_wrapper(
137
v: torch.Tensor,
138
scale: float | None = None,
139
cu_seqlens: torch.Tensor | None = None,
140
+ enable_gqa: bool = False,
141
) -> torch.Tensor:
142
# Never remove the contiguous logic for ROCm
143
# Without it, hallucinations occur with the backend
@@ -143,7 +147,7 @@ def torch_sdpa_wrapper(
147
v = v.contiguous()
148
149
if cu_seqlens is None:
146
- return apply_sdpa(q, k, v, scale=scale)
150
+ return apply_sdpa(q, k, v, scale=scale, enable_gqa=enable_gqa)
151
152
outputs = []
153
@@ -152,7 +156,7 @@ def torch_sdpa_wrapper(
156
k_chunks = torch.split(k, lens, dim=1)
157
v_chunks = torch.split(v, lens, dim=1)
158
for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
155
- output_i = apply_sdpa(q_i, k_i, v_i, scale=scale)
159
+ output_i = apply_sdpa(q_i, k_i, v_i, scale=scale, enable_gqa=enable_gqa)
160
outputs.append(output_i)
161
context_layer = torch.cat(outputs, dim=1)
162
return context_layer
@@ -164,6 +168,7 @@ def torch_sdpa_wrapper_fake(
168
v: torch.Tensor,
169
scale: float | None,
170
cu_seqlens: torch.Tensor | None,
171
+ enable_gqa: bool = False,
172
) -> torch.Tensor:
173
return torch.empty_like(q)
174
@@ -181,5 +186,8 @@ def vit_torch_sdpa_wrapper(
186
v: torch.Tensor,
187
scale: float | None = None,
188
cu_seqlens: torch.Tensor | None = None,
189
+ enable_gqa: bool = False,
190
) -> torch.Tensor:
185
- return torch.ops.vllm.torch_sdpa_wrapper(q, k, v, scale, cu_seqlens)
191
+ return torch.ops.vllm.torch_sdpa_wrapper(
192
+ q, k, v, scale, cu_seqlens, enable_gqa=enable_gqa
193
+ )