Add replicatekvhead transform

Signed-off-by: Mamta Singh <mamtsing@qti.qualcomm.com>

Author: mamtsing

QEfficient/transformers/cache_utils.py

@@ -302,10 +302,10 @@ def update_ckv(self, compressed_kv, cache_kwargs):
         position_ids = cache_kwargs.get("position_ids")
         batch_index = cache_kwargs.get("batch_index", None)  # TODO: add support later
 
-        self.ckv = CtxScatterFunc3D.apply(self.ckv, position_ids, compressed_kv)
+        self.ckv = CtxScatterFunc.apply(self.ckv, position_ids, compressed_kv)
 
         ckv_out = self.ckv
-        ctx_len = ckv_out.shape[1]
+        ctx_len = ckv_out.shape[-2]
         ctx_indices = torch.arange(ctx_len)[None, ...]
         gather_limit = position_ids.max(1, keepdim=True).values
         invalid_mask = ctx_indices > gather_limit
@@ -315,7 +315,7 @@ def update_ckv(self, compressed_kv, cache_kwargs):
             invalid_idx_value = 0
         ctx_indices = torch.where(invalid_mask, invalid_idx_value, ctx_indices)
 
-        ckv_out = CtxGatherFunc3D.apply(ckv_out, ctx_indices)
+        ckv_out = CtxGatherFunc.apply(ckv_out, ctx_indices, ctx_len)
         ckv_out = torch.where(invalid_mask.unsqueeze(-1), torch.tensor(0.0, dtype=torch.float32), ckv_out)
         return ckv_out
 

QEfficient/transformers/models/deepseek_v3/modeling_deepseek_qeff.py

@@ -214,26 +214,23 @@ def __qeff_init__(
             -1, self.num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim
         ).split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
         q_up = q_up.reshape(-1, self.num_heads * self.qk_nope_head_dim).unsqueeze(0)
-        # self.register_buffer("q_up", q_up.detach().clone(), persistent=False)
+
         self.q_up = torch.nn.Parameter(q_up.detach().clone())
         q_rope = q_rope.reshape(-1, self.num_heads * self.qk_rope_head_dim).unsqueeze(0)
-        # self.register_buffer("q_rope", q_rope.detach().clone(), persistent=False)
+
         self.q_rope = torch.nn.Parameter(q_rope.detach().clone())
         k_up, v_up = self.kv_b_proj.weight.T.view(-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim).split(
             [self.qk_nope_head_dim, self.v_head_dim], dim=-1
         )
         k_up = k_up.reshape(-1, self.num_heads * self.qk_nope_head_dim).unsqueeze(0)
         v_up = v_up.reshape(-1, self.num_heads * self.v_head_dim).unsqueeze(0)
-        # self.register_buffer("k_up", k_up.detach().clone(), persistent=False)
-        # self.register_buffer("v_up", v_up.detach().clone(), persistent=False)
+
         self.k_up = torch.nn.Parameter(k_up.detach().clone())
         self.v_up = torch.nn.Parameter(v_up.detach().clone())
         per_head_q_up = self.q_up.squeeze(0).view(-1, self.num_heads, self.qk_nope_head_dim).transpose(0, 1)
         per_head_k_up = (
             self.k_up.squeeze(0).view(-1, self.num_heads, self.qk_nope_head_dim).transpose(0, 1).transpose(1, 2)
         )
-        # self.register_buffer("per_head_q_up", per_head_q_up.detach().clone(), persistent=False)
-        # self.register_buffer("per_head_k_up", per_head_k_up.detach().clone(), persistent=False)
         self.per_head_q_up = torch.nn.Parameter(per_head_q_up.detach().clone())
         self.per_head_k_up = torch.nn.Parameter(per_head_k_up.detach().clone())
 
@@ -243,12 +240,7 @@ def __qeff_init__(
             out = torch.cat((out,x), 0)
         fusedqk = out.reshape(self.num_heads, -1, self.kv_lora_rank)
 
-        #fusedqk = torch.bmm(per_head_q_up, per_head_k_up)
-        # self.register_buffer("fusedqk", fusedqk.detach().clone(), persistent=False)
         self.fusedqk = torch.nn.Parameter(fusedqk.detach().clone())
-        kv_a_proj_with_mqa_ckv, kv_a_proj_with_mqa_k_pe = self.kv_a_proj_with_mqa.weight.T.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
-        self.kv_a_proj_with_mqa_ckv = torch.nn.Parameter(kv_a_proj_with_mqa_ckv.detach().clone())
-        self.kv_a_proj_with_mqa_k_pe = torch.nn.Parameter(kv_a_proj_with_mqa_k_pe.detach().clone())
 
     def fused_forward(
         self,
@@ -267,69 +259,79 @@ def fused_forward(
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
 
-        compressed_kv = torch.matmul(hidden_states, self.kv_a_proj_with_mqa_ckv)
-        k_pe = torch.matmul(hidden_states, self.kv_a_proj_with_mqa_k_pe)
-        k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
+        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+        compressed_kv = compressed_kv.view(bsz, q_len, -1, self.kv_lora_rank+self.qk_rope_head_dim).transpose(1, 2)
+        compressed_kv, k_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
 
         q_a_proj_out = self.q_a_layernorm(self.q_a_proj(hidden_states))
         q_pe = torch.bmm(q_a_proj_out, self.q_rope)
         q_pe = q_pe.view(bsz, q_len, self.num_heads, self.qk_rope_head_dim).transpose(1, 2)
         q_nope = torch.bmm(q_a_proj_out, self.q_up)
         q_nope = q_nope.view(bsz, q_len, self.num_heads, self.qk_nope_head_dim).transpose(1, 2)
 
+        compressed_kv = self.kv_a_layernorm(compressed_kv)
         cache_kwargs = {"position_ids": position_ids, "batch_index": batch_index}
         if compressed_kvs is not None:
             compressed_kv = compressed_kvs.update_ckv(compressed_kv, self.layer_idx, cache_kwargs)
 
-        kva = self.kv_a_layernorm(compressed_kv)
-        k_nope = torch.bmm(kva, self.k_up)
-        k_nope = k_nope.view(bsz, -1, self.num_heads, self.qk_nope_head_dim).transpose(1, 2)
-        value_states = torch.bmm(kva, self.v_up)
-        value_states = value_states.view(bsz, -1, self.num_heads, self.qk_nope_head_dim).transpose(1, 2)
-
-        cos, sin = self.rotary_emb(value_states, seq_len=32 * 1024)
-        q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
-
-        if compressed_kvs is not None:
-            k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs)
+        kva = compressed_kv
 
         if mla_absorption is not None:
             enable_absorption = mla_absorption.get("enable", False)
             absorb_online = mla_absorption.get("online", False)
         else:
             enable_absorption = False
 
+        n_head_ckv = compressed_kv.shape[1]
+        p = self.num_heads//n_head_ckv
+
+        value_out = []
+        for i in range(n_head_ckv):
+          value_states_ph = torch.matmul(kva[:,i,:,:], self.v_up[:, :, i*p*self.v_head_dim: (i+1)*p*self.v_head_dim])
+          value_states_ph = value_states_ph.view(bsz, -1, p, self.qk_nope_head_dim).transpose(1, 2)
+          value_out.append(value_states_ph)
+        value_states = torch.cat(value_out, dim=1)
+
+        cos, sin = self.rotary_emb(value_states_ph, seq_len=32 * 1024)
+        q_pe, k_pe = orig_apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
+
+        if compressed_kvs is not None:
+            k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs)
+
         x = []
-        for i in range(self.num_heads):
-            if enable_absorption:
-                if absorb_online:
-                    if i==0:
-                        print("online absorption")
-                    out = torch.matmul(self.per_head_q_up[i,:,:], self.per_head_k_up[i,:,:])
-                    out = out.reshape(1, -1, self.kv_lora_rank)
-                    out2 = torch.matmul(q_a_proj_out.unsqueeze(1), out)
+        for k in range(n_head_ckv):
+            k_nope = torch.matmul(kva[:,k,:,:], self.k_up[:, :, k*p*self.qk_nope_head_dim: (k+1)*p*self.qk_nope_head_dim])
+            k_nope = k_nope.view(bsz, -1, p, self.qk_nope_head_dim).transpose(1, 2)
+
+            for i in range(k*p, (k+1)*p):
+                if enable_absorption:
+                    if absorb_online:
+                        if i==0:
+                            print("online absorption")
+                        out = torch.matmul(self.per_head_q_up[i,:,:], self.per_head_k_up[i,:,:])
+                        out = out.reshape(1, -1, self.kv_lora_rank)
+                        out2 = torch.matmul(q_a_proj_out.unsqueeze(1), out)
+                    else:
+                        if i==0:
+                            print("using fused qk")
+                        out2 = torch.matmul(q_a_proj_out.unsqueeze(1), self.fusedqk[i,:,:])
+
+                    out3 = torch.cat((out2, q_pe[:,i,:,:].unsqueeze(1)), -1)
+                    kva_kpe = torch.cat((kva[:,k,:,:],k_pe[:,k,:,:]), -1).unsqueeze(1)
+                    attn_weights = torch.matmul(out3, kva_kpe.transpose(2,3)) * self.softmax_scale
                 else:
                     if i==0:
-                        print("using fused qk")
-                    out2 = torch.matmul(q_a_proj_out.unsqueeze(1), self.fusedqk[i,:,:])
-
-                out3 = torch.cat((out2, q_pe[:,i,:,:].unsqueeze(1)), -1)
-                kva_kpe = torch.cat((kva,k_pe.squeeze(1)), -1)
-                attn_weights = torch.matmul(out3, kva_kpe.transpose(1, 2).unsqueeze(1)) * self.softmax_scale
-            else:
-                if i==0:
-                    print("no absorption")
-                query_states = torch.cat((q_nope[:,i,:,:], q_pe[:,i,:,:]), -1)
-                key_states = torch.cat((k_nope[:,i,:,:].unsqueeze(1), k_pe), -1)
-                attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
+                        print("no absorption")
+                    query_states = torch.cat((q_nope[:,i,:,:], q_pe[:,i,:,:]), -1).unsqueeze(1)
+                    key_states = torch.cat((k_nope[:,i%p,:,:], k_pe[:,k,:,:]), -1).unsqueeze(1)
+                    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale
 
-            if attention_mask is not None:  # no matter the length, we just slice it
-                attn_weights = torch.where(attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=torch.float32), attn_weights)
+                if attention_mask is not None:  # no matter the length, we just slice it
+                    attn_weights = torch.where(attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=torch.float32), attn_weights)
 
-            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q_pe.dtype)
-            attn_output = torch.matmul(attn_weights, value_states[:,i,:,:])
-
-            x.append(attn_output)
+                attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q_pe.dtype)
+                attn_output = torch.matmul(attn_weights, value_states[:,i,:,:])
+                x.append(attn_output)
 
         attn_output = torch.cat(x, dim=1)
 
@@ -455,23 +457,6 @@ def forward(self, hidden_states):
         hidden_states = hidden_states + self.shared_experts(residuals)
         return hidden_states
 
-    # def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor):
-    #     final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype)
-    #     expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts))
-    #     expert_mask = expert_mask.permute(2, 0, 1)
-
-    #     for expert_idx in range(len(self.experts)):
-    #         expert = self.experts[expert_idx]
-    #         mask = expert_mask[expert_idx]
-    #         expert_output = expert(hidden_states) * (((topk_weights * mask).sum(1))[:, None])
-    #         expert_output = torch.where(
-    #             (topk_weights * mask).sum(1).to(torch.bool)[:, None],
-    #             expert_output,
-    #             torch.tensor(0.0),
-    #         )
-    #         final_hidden_states = final_hidden_states + expert_output
-    #     return final_hidden_states.type(hidden_states.dtype)
-
 
 class QEffPrefillOnlyDeepseekV3MoE(nn.Module):
     def __qeff_init__(

QEfficient/transformers/models/modeling_auto.py

@@ -51,6 +51,7 @@
     PrefillOnlyExternalModuleMapperTransform,
     PrefillOnlyChunkedTransform,
     PrefillOnlyTransform,
+    ReplicateKVHeadTransform,
     RevertPrefillKeepAttentionTransform,
     RevertPrefillOnlyTransform,
     RevertPrefillOnlyExternalModuleMapperTransform,
@@ -2410,6 +2411,11 @@ def __init__(
         self.comp_ctx_lengths_prefill, self.comp_ctx_lengths_decode = None, None
         self.hash_params["max_seq_len_cached"] = max_seq_len_cached
 
+        if self.model.config.model_type in {"kimi_k2"}:
+            self.model, replicate_kv_transformed = ReplicateKVHeadTransform.apply(self.model, **kwargs)
+            if replicate_kv_transformed:
+                self.hash_params["config"] = model.config.to_diff_dict()
+
         # ---Sampling---
         # Note: SamplerTransform should be applied after all other transforms
         # are done. The role of the sampler is to just add nodes at the output of the
@@ -2746,11 +2752,11 @@ def export(
             output_names = [v for v in output_names if "past" not in v]
             example_inputs["compressed_kvs"] = [[] for _ in range(self.num_layers)]
             for i in range(self.num_layers):
-                ckv = torch.zeros((bs, seq_len, self.model.config.kv_lora_rank), dtype=torch.float32)
-                k_pe = torch.zeros((bs, 1, seq_len, self.model.config.qk_rope_head_dim), dtype=torch.float32)
+                ckv = torch.zeros((bs, 4, seq_len, self.model.config.kv_lora_rank), dtype=torch.float32)
+                k_pe = torch.zeros((bs, 4, seq_len, self.model.config.qk_rope_head_dim), dtype=torch.float32)
                 example_inputs["compressed_kvs"][i].append(ckv)
                 example_inputs["compressed_kvs"][i].append(k_pe)
-                dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 1: "ctx_len"}
+                dynamic_axes[f"compressed_kv.{i}"] = {0: "batch_size", 2: "ctx_len"}
                 dynamic_axes[f"k_pe.{i}"] = {0: "batch_size", 2: "ctx_len"}
                 output_names.append(f"compressed_kv.{i}_RetainedState")
                 output_names.append(f"k_pe.{i}_RetainedState")

QEfficient/transformers/models/pytorch_transforms.py

@@ -10,6 +10,7 @@
 from types import MethodType
 from typing import Callable, Optional, Tuple, Union
 
+import torch
 from torch import nn
 from transformers.models.codegen.modeling_codegen import (
     CodeGenAttention,
@@ -456,6 +457,7 @@
 from QEfficient.transformers.post_processing import build_and_attach_mlp, model_type_registry
 from QEfficient.transformers.sampler.sampler import sampler_forward
 from QEfficient.transformers.spd.spd_transform_forward import tlm_forward
+from QEfficient.utils.logging_utils import logger
 
 SPD_TARGET = "target"
 
@@ -694,6 +696,79 @@ class RevertPrefillOnlyTransform(ModuleMappingTransform):
     }
 
 
+class ReplicateKVHeadTransform:
+    """
+    Replicates KV heads in attention modules to match the number of KV heads in the target model.
+    This transform is used when the source model has fewer KV heads than required in target model.
+    """
+
+    def _duplicate_weights_for_linear_layer(
+        layer: nn.Module, orig_kv_heads: int, repeat: int, dim: int, hidden_size: int
+    ):
+        new_kv_heads = repeat #for mla
+
+        layer.weight.data = torch.repeat_interleave(
+            layer.weight.data.view(orig_kv_heads, dim, hidden_size), repeat, 0
+        ).view(new_kv_heads * dim, hidden_size)
+
+        if layer.bias is not None:
+            layer.bias.data = torch.repeat_interleave(
+                layer.bias.data.view(orig_kv_heads, dim), repeat, 0
+            ).view(new_kv_heads * dim)
+
+    def _get_text_model(model):
+        """
+        Determine and return the appropriate text_model from a given model object.
+        """
+        # Check for VLMs
+        if hasattr(model, "language_model"):
+            if hasattr(model.language_model, "model"):
+                return model.language_model.model
+            else:
+                return model.language_model
+        # Check for CausalLMs
+        if hasattr(model, "model"):
+            return model.model
+
+        raise AttributeError("No suitable text model found in the provided model.")
+
+    @classmethod
+    def apply(cls, model: nn.Module, **kwargs) -> nn.Module:
+        """
+        Replicates KV heads in attention modules based on provided multiplier.
+
+        Args:
+            model: The model to apply the transform to.
+            kwargs: Additional arguments for the transformation. Includes:
+                - num_kv_heads_repeat: The number of times to repeat the KV heads.
+        """
+        n_repeat = kwargs.pop("num_kv_heads_repeat", 1)
+        transformed = False
+        if n_repeat is not None and n_repeat > 1:
+            text_model = cls._get_text_model(model)
+
+            orig_kv_heads = 1 # for mla #text_model.config.num_key_value_heads
+            new_kv_heads = n_repeat*orig_kv_heads
+            text_model.config.orig_kv_heads = orig_kv_heads
+            text_model.config.num_key_value_heads = new_kv_heads
+
+            num_attention_heads = text_model.config.num_attention_heads
+            hidden_size = text_model.config.hidden_size
+
+            logger.warning(f"Original KV heads: {orig_kv_heads}")
+            logger.warning(f"Modified KV heads: {new_kv_heads}")
+            transformed = True
+            for block in text_model.layers:
+                attn = getattr(block, "cross_attn", getattr(block, "self_attn", None))
+                attn.num_key_value_heads = new_kv_heads
+                head_dim = attn.kv_lora_rank + attn.qk_rope_head_dim
+
+                cls._duplicate_weights_for_linear_layer(
+                    attn.kv_a_proj_with_mqa, orig_kv_heads, n_repeat, head_dim, hidden_size
+                )
+        return model, transformed
+
+
 class SpDTransform:
     """
     Apply generic QEffForCausalLM forward pass to extract `num_speculative_tokens+1` hidden states before computing logits during decode phase and extract last predicted token during prefill.