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 15, 22:15 UTC
[Bugfix][Model] OLMo 2: split qkv correctly for GQA and MQA (#13687)
Shane A committed
Feb 21, 2025 at 22:07 UTC
9a1f1da5d1f2e94adb49e0d82b464dc3c1318cc7
1 file changed
+1
-1
vllm/model_executor/models/olmo2.py
+1
-1
@@ -157,7 +157,7 @@ class Olmo2Attention(nn.Module):
157
attn_metadata: AttentionMetadata,
158
) -> torch.Tensor:
159
qkv, _ = self.qkv_proj(hidden_states)
160
- q, k, v = qkv.chunk(chunks=3, dim=-1)
160
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
161
q, k = self._apply_qk_norm(q, k)
162
q, k = self.rotary_emb(positions, q, k)
163
attn_output = self.attn(q, k, v, kv_cache, attn_metadata)