[Bugfix][Model] OLMo 2: split qkv correctly for GQA and MQA (vllm-project#13687)

Signed-off-by: Louis Ulmer <[email protected]>

Author: 2015aroras

vllm/model_executor/models/olmo2.py

@@ -157,7 +157,7 @@ def forward(
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
-        q, k, v = qkv.chunk(chunks=3, dim=-1)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         q, k = self._apply_qk_norm(q, k)
         q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v, kv_cache, attn_metadata)