Layer wise changes for kimi model

QEfficient/base/onnx_transforms.py

@@ -129,17 +129,80 @@ def apply(cls, model: ModelProto) -> bool:
         return op_applied
 
 
+class RemovePrefix(BaseOnnxTransform):
+    @classmethod
+    def apply(cls, model: ModelProto) -> bool:
+        graph = model.graph
+        renamed = False
+
+        def strip_prefix(name: str) -> str:
+            parts = name.rsplit("/", 1)
+            return parts[1] if len(parts) == 2 else parts[0]
+
+        input_names = []
+        for i, inputs in enumerate(graph.input):
+            original = inputs.name
+            new = strip_prefix(original)
+            if new != original:
+                renamed = True
+            inputs.name = new
+            graph.input[i].name = new
+            input_names.append(new)
+
+        input_name_set = set(input_names)
+        output_rename_map = {}
+
+        # Rename model graph outputs and keep mapping so producer/consumer edges can be fixed.
+        for out in graph.output:
+            original = out.name
+            new = strip_prefix(original)
+            if new != original:
+                out.name = new
+                output_rename_map[original] = new
+                renamed = True
+
+        for node in graph.node:
+            for i, out in enumerate(node.output):
+                if out in output_rename_map and output_rename_map[out] != out:
+                    node.output[i] = output_rename_map[out]
+                    renamed = True
+
+            new_inputs = []
+            for s in node.input:
+                # Keep node inputs in sync for renamed model outputs.
+                if s in output_rename_map:
+                    new_inputs.append(output_rename_map[s])
+                    continue
+
+                if s in input_name_set:
+                    new_inputs.append(s)
+                    continue
+
+                replaced = s
+                if "/" in s:
+                    tail = s.rsplit("/", 1)[1]
+                    if tail in input_name_set:
+                        replaced = tail
+                new_inputs.append(replaced)
+
+            for idx in range(len(node.input)):
+                if node.input[idx] != new_inputs[idx]:
+                    node.input[idx] = new_inputs[idx]
+                    renamed = True
+
+        return renamed
+
+
 class RenameFunctionOutputsTransform(BaseOnnxTransform):
     """Rename outputs of decoder-related functions for better clarity."""
 
     @classmethod
-    def apply(cls, model: ModelProto) -> bool:
+    def apply(cls, model: ModelProto, layer_idx=0) -> bool:
         graph = model.graph
         op_type_to_func = {f.name: f for f in model.functions}
         decoder_patterns = ["DecoderLayer", "Block", "Layer"]
         renamed = False
         model_out_map = {v.name: i for i, v in enumerate(graph.output)}
-        layer_idx = 0
 
         for node in graph.node:
             if any(p in node.name or p in node.op_type for p in decoder_patterns):
@@ -150,13 +213,16 @@ def apply(cls, model: ModelProto) -> bool:
                     if "_InternalRetainedState" in out_name:
                         renamed = True
                         orig = node.output[i]
-                        new = (
-                            f"past_key.{layer_idx}_RetainedState"
-                            if "key" in out_name
-                            else f"past_value.{layer_idx}_RetainedState"
-                            if "value" in out_name
-                            else orig
-                        )
+                        if "key" in out_name:
+                            new = f"past_key.{layer_idx}_RetainedState"
+                        elif "value" in out_name:
+                            new = f"past_value.{layer_idx}_RetainedState"
+                        elif "compressed_kv" in out_name:
+                            new = f"compressed_kv.{layer_idx}_RetainedState"
+                        elif "k_pe" in out_name:
+                            new = f"k_pe.{layer_idx}_RetainedState"
+                        else:
+                            new = orig
                         node.output[i] = new
                         if orig in model_out_map:
                             graph.output[model_out_map[orig]].name = new
@@ -275,7 +341,9 @@ def _set_external_data(tensor, file_name):
             applied[CustomOpTransform] = CustomOpTransform.apply(model)
 
         if RenameFunctionOutputsTransform in requested:
-            applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply(model)
+            applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply(
+                model, layer_idx=kwargs.get("layer_idx", 0)
+            )
 
         if AdapterWeightsToInputsTransform in requested:
             applied[AdapterWeightsToInputsTransform] = AdapterWeightsToInputsTransform.apply(model, **kwargs)

QEfficient/blocking/attention_blocking.py

@@ -9,16 +9,18 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional
 
 import torch
 from transformers.cache_utils import Cache
 
 from QEfficient.blocking.blocked_attention_forwards import (
     blocked_bhqkv_attention_forward,
     blocked_h_attention_forward,
+    blocked_h_mla_attention_forward,
     blocked_hqkv_attention_forward,
     blocked_kv_attention_forward,
+    blocked_kv_mla_attention_forward,
     blocked_q_attention_forward,
     blocked_qkv_attention_forward,
 )
@@ -57,6 +59,11 @@ def supports_blocked_kv(past_key_value: Optional[Cache]) -> bool:
     BlockingMode.BHQKV: blocked_bhqkv_attention_forward,
 }
 
+_STRATEGIES_MLA: Dict[BlockingMode, Callable] = {
+    BlockingMode.KV: blocked_kv_mla_attention_forward,
+    BlockingMode.H: blocked_h_mla_attention_forward,
+}
+
 
 # helper function needed both in generic blocked approach and in other modeling files for non-blocked approach
 def past_key_value_update(
@@ -160,3 +167,66 @@ def generic_blocked_attention_interface(
     )
 
     return attn_output, attn_weights
+
+
+def generic_blocked_mla_attention_interface(
+    module,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    mla_absorption: Dict[str, Any],
+    blocking_config: AttentionBlockingConfig,
+    query: Optional[torch.Tensor] = None,
+    q_a_proj_out: Optional[torch.Tensor] = None,
+    fusedqk: Optional[torch.Tensor] = None,
+    q_nope: Optional[torch.Tensor] = None,
+    q_pe: Optional[torch.Tensor] = None,
+    kva: Optional[torch.Tensor] = None,
+    k_pe: Optional[torch.Tensor] = None,
+    per_head_q_up: Optional[torch.Tensor] = None,
+    per_head_k_up: Optional[torch.Tensor] = None,
+    per_head_v_up: Optional[torch.Tensor] = None,
+    per_head_k_up_normal: Optional[torch.Tensor] = None,
+    layer_idx: Optional[int] = None,
+    compressed_kvs: Optional[torch.Tensor] = None,
+    comp_ctx_lengths: Optional[torch.LongTensor] = None,
+    batch_index: Optional[torch.LongTensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_seen_tokens: Optional[int] = None,
+    non_blocked_forward: Callable = None,
+    score_mod: Optional[Callable] = None,
+    position_bias: Optional[torch.Tensor] = None,
+    sinks: Optional[torch.Tensor] = None,
+    sliding_window: Optional[int] = None,
+    **kwargs,
+):
+    cache_kwargs = {"position_ids": position_ids, "batch_index": batch_index}
+    mla_blocking_strategy = _STRATEGIES_MLA.get(blocking_config.mode)
+    attn_output, attn_weights = mla_blocking_strategy(
+        module=module,
+        query=query,
+        q_a_proj_out=q_a_proj_out,
+        fusedqk=fusedqk,
+        q_nope=q_nope,
+        q_pe=q_pe,
+        kva=kva,
+        k_pe=k_pe,
+        per_head_q_up=per_head_q_up,
+        per_head_k_up=per_head_k_up,
+        per_head_v_up=per_head_v_up,
+        per_head_k_up_normal=per_head_k_up_normal,
+        attention_mask=attention_mask,
+        scaling=scaling,
+        cache_kwargs=cache_kwargs,
+        layer_idx=layer_idx,
+        compressed_kvs=compressed_kvs,
+        mla_absorption=mla_absorption,
+        num_kv_blocks=blocking_config.num_kv_blocks,
+        num_q_blocks=blocking_config.num_q_blocks,
+        head_block_size=blocking_config.head_block_size,
+        num_batch_blocks=blocking_config.num_batch_blocks,
+        score_mod=score_mod,
+        position_bias=position_bias,
+        sinks=sinks,
+    )
+
+    return attn_output, attn_weights