diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py
index de5087f31..97754f362 100644
--- a/python/sglang/srt/models/olmo2.py
+++ b/python/sglang/srt/models/olmo2.py
@@ -101,6 +101,7 @@ class Olmo2Attention(nn.Module):
             self.hidden_size,
             self.head_dim,
             self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
             bias=config.attention_bias,
             quant_config=quant_config,
             prefix=add_prefix("qkv_proj", prefix),