Repository: huggingface/transformers
Pull request: [`JetMoe`] Fix KV head repetition and padding free

During the vllm refactor 2 small issues have been introduced
- Jetmoe doesn't repeat the kv heads with `repeat_interleave` but uses `repeat`, i.e. different patterns
- Position ids wasn't passed down leading to padding free not working

Should hopefully align the tests again in #41377 (at least checked locally on our dgx machines that the outputs are matching again) cc @ydshieh 

Changed files:
filename: src/transformers/models/jetmoe/modeling_jetmoe.py
patch: @@ -426,7 +426,8 @@ def __init__(self, config: JetMoeConfig, layer_idx: Optional[int] = None):
                 "when creating this class."
             )
 
-        self.num_key_value_groups = config.num_experts_per_tok
+        self.num_key_value_groups = 1  # We ignore this by setting it to 1 as we have different repeat patterns
+        self.top_k = config.num_experts_per_tok
         self.attention_dropout = config.attention_dropout
         self.kv_projection_size = config.kv_channels * config.num_key_value_heads
         self.num_key_value_heads = config.num_key_value_heads
@@ -468,6 +469,11 @@ def forward(
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
+        # This is different from other models where we repeat k/v heads
+        # instead of repeat interleaving them
+        key_states = key_states.repeat(1, self.top_k, 1, 1)
+        value_states = value_states.repeat(1, self.top_k, 1, 1)
+
         attn_output, attn_weights = attention_interface(
             self,
             query_states,
@@ -479,7 +485,7 @@ def forward(
             **kwargs,
         )
 
-        attn_output = attn_output.view(*input_shape, self.num_key_value_groups, -1)
+        attn_output = attn_output.view(*input_shape, self.top_k, -1)
         attn_output = self.experts.reduce(attn_output, topo_info)
         attn_output = attn_output.view(*input_shape, -1)
         return attn_output, attn_weights, router_logits
@@ -635,6 +641,7 @@ def forward(
                 past_key_values=past_key_values,
                 use_cache=use_cache,
                 cache_position=cache_position,
+                position_ids=position_ids,
                 **kwargs,
             )
 

filename: src/transformers/models/jetmoe/modular_jetmoe.py
patch: @@ -311,7 +311,8 @@ def __init__(self, config: JetMoeConfig, layer_idx: Optional[int] = None):
                 "when creating this class."
             )
 
-        self.num_key_value_groups = config.num_experts_per_tok
+        self.num_key_value_groups = 1  # We ignore this by setting it to 1 as we have different repeat patterns
+        self.top_k = config.num_experts_per_tok
         self.attention_dropout = config.attention_dropout
         self.kv_projection_size = config.kv_channels * config.num_key_value_heads
         self.num_key_value_heads = config.num_key_value_heads
@@ -353,6 +354,11 @@ def forward(
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
+        # This is different from other models where we repeat k/v heads
+        # instead of repeat interleaving them
+        key_states = key_states.repeat(1, self.top_k, 1, 1)
+        value_states = value_states.repeat(1, self.top_k, 1, 1)
+
         attn_output, attn_weights = attention_interface(
             self,
             query_states,
@@ -364,7 +370,7 @@ def forward(
             **kwargs,
         )
 
-        attn_output = attn_output.view(*input_shape, self.num_key_value_groups, -1)
+        attn_output = attn_output.view(*input_shape, self.top_k, -1)
         attn_output = self.experts.reduce(attn_output, topo_info)
         attn_output = attn_output.view(*input_shape, -1)
         return attn_output, attn_weights, router_logits
@@ -512,6 +518,7 @@ def forward(
                 past_key_values=past_key_values,
                 use_cache=use_cache,
                 cache_position=cache_position,
+                position_ids=position_ids,
                 **kwargs,
             )