perf: use native enable_gqa — matches PR openai#114 attention path, 12ms/step faster

Author: turazashvili

records/track_10min_16mb/2026-03-20_StackedV1/train_gpt.py

@@ -621,12 +621,6 @@ def __init__(
         self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
         self.rotary = Rotary(self.head_dim, base=rope_base)
 
-    def _repeat_kv(self, k: Tensor, v: Tensor, bsz: int, seqlen: int) -> tuple[Tensor, Tensor]:
-        reps = self.num_heads // self.num_kv_heads
-        k = k[:, :, None, :, :].expand(-1, -1, reps, -1, -1).reshape(bsz, self.num_heads, seqlen, self.head_dim)
-        v = v[:, :, None, :, :].expand(-1, -1, reps, -1, -1).reshape(bsz, self.num_heads, seqlen, self.head_dim)
-        return k, v
-
     def forward(self, x: Tensor) -> Tensor:
         bsz, seqlen, dim = x.shape
         q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
@@ -638,10 +632,9 @@ def forward(self, x: Tensor) -> Tensor:
         q = apply_rotary_emb(q, cos, sin)
         k = apply_rotary_emb(k, cos, sin)
         q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None]
-        if self.num_kv_heads != self.num_heads:
-            k, v = self._repeat_kv(k, v, bsz, seqlen)
         y = F.scaled_dot_product_attention(
             q, k, v, attn_mask=None, is_causal=True,
+            enable_gqa=(self.num_kv_heads != self.num_heads),
         ).transpose(1, 2).contiguous().reshape(bsz, seqlen, dim)
         return self.proj(y)