Skip the query/key head repeat for grouped value heads in gated DeltaNet (#49141)

* Skip the query/key head repeat for grouped value heads in gated DeltaNet

FLA's chunked and fused recurrent gated delta rule kernels accept fewer
query/key heads than value heads since 0.5.0, and the hub kernel does too.
The torch fallbacks now repeat query/key themselves, so the modeling code
only repeats when an older FLA is installed. This removes a query/key copy
and its backward reduction in every linear-attention layer of Qwen3-Next,
Qwen3.5, Qwen3.5-MoE, Qwen4-Exp and OLMo-Hybrid.

* Pass g and beta by keyword in test_gdn_chunked_matches_recurrent

FLA's fused recurrent kernel takes gk and gv between g and beta, so the
positional call fed beta into gk whenever FLA was installed. FLA's chunked
kernel also runs its dot products in TF32, so the tolerance follows the
backend.

* Name the FLA version check fla_needs_equal_qk_heads

* Always pass grouped query/key heads to the gated delta rule kernels

Drop the FLA version check and the layer-level repeat; the torch fallbacks
repeat query/key themselves when there are fewer of them.

* Add a helper for the gated delta rule test inputs

* Move the gated delta rule input helper into the model tester

Author: manav-simile

src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py

@@ -398,11 +398,6 @@ def forward(
         k = k.view(batch_size, seq_len, -1, self.head_k_dim)
         v = v.view(batch_size, seq_len, -1, self.head_v_dim)
 
-        if self.num_v_heads > self.num_k_heads:
-            expand_ratio = self.num_v_heads // self.num_k_heads
-            q = q.repeat_interleave(expand_ratio, dim=2)
-            k = k.repeat_interleave(expand_ratio, dim=2)
-
         beta = self.b_proj(hidden_states).sigmoid()
         if self.allow_neg_eigval:
             beta = beta * 2.0

src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py

@@ -297,9 +297,6 @@ def forward(
         beta = b.sigmoid()
         # If the model is loaded in fp16, without the .float() here, A might be -inf
         g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
-        if self.num_v_heads // self.num_k_heads > 1:
-            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
-            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
 
         if recurrent_state is not None and seq_len == 1:
             core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(

src/transformers/models/qwen3_5/modeling_qwen3_5.py

@@ -240,12 +240,17 @@ def torch_chunk_gated_delta_rule(
         - Either None or the new recurrent state tensor of shape [batch_size, num_v_heads, k_head_dim, v_head_dim]
     """
     initial_dtype = query.dtype
-    batch_size, sequence_length, _, k_head_dim = key.shape
+    batch_size, sequence_length, num_k_heads, k_head_dim = key.shape
     num_v_heads, v_head_dim = value.shape[-2:]
     recurrent_state_shape = (batch_size, num_v_heads, k_head_dim, v_head_dim)
     padded_output_shape = (batch_size, num_v_heads, -1, v_head_dim)  # -1 is the padded sequence length
     decay = g  # rename for clarity: argument name must stay "g" to match flash_linear_attention's API
 
+    # Grouped value attention: each query/key head serves num_v_heads // num_k_heads consecutive value heads
+    if num_k_heads != num_v_heads:
+        query = query.repeat_interleave(num_v_heads // num_k_heads, dim=2)
+        key = key.repeat_interleave(num_v_heads // num_k_heads, dim=2)
+
     # Make sure all tensors are fp32 and reshape them to [batch_size, num_*_heads, seqlen, ...]
     query, key, value, beta, decay = [
         x.transpose(1, 2).to(torch.float32, memory_format=torch.contiguous_format)
@@ -362,10 +367,15 @@ def torch_recurrent_gated_delta_rule(
     Same args and return value as torch_chunk_gated_delta_rule, except for `chunk_size` because the sequence dim is not
     chunked."""
     initial_dtype = query.dtype
-    batch_size, sequence_length, _, k_head_dim = key.shape
+    batch_size, sequence_length, num_k_heads, k_head_dim = key.shape
     num_v_heads, v_head_dim = value.shape[-2:]
     decay = g  # rename for clarity: argument name must stay "g" to match flash_linear_attention's API
 
+    # Grouped value attention: each query/key head serves num_v_heads // num_k_heads consecutive value heads
+    if num_k_heads != num_v_heads:
+        query = query.repeat_interleave(num_v_heads // num_k_heads, dim=2)
+        key = key.repeat_interleave(num_v_heads // num_k_heads, dim=2)
+
     # Make sure all tensors are fp32 and reshape them to [batch_size, num_*_heads, seqlen, ...]
     query, key, value, beta, decay = [
         x.transpose(1, 2).to(torch.float32, memory_format=torch.contiguous_format)
@@ -558,9 +568,6 @@ def forward(
         beta = b.sigmoid()
         # If the model is loaded in fp16, without the .float() here, A might be -inf
         g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
-        if self.num_v_heads // self.num_k_heads > 1:
-            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
-            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
 
         if recurrent_state is not None and seq_len == 1:
             core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule(