MLA Changes (quic#956)

fix kv head replication

---------

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

Author: quic-mamta

QEfficient/blocking/blocked_attention_forwards.py

@@ -898,25 +898,25 @@ def blocked_kv_mla_attention_forward(
             absorption = False
 
         k_heads, q_heads = compressed_kv_block.shape[1], query.shape[1]
-        num_heads_to_repeat = q_heads - k_heads
-        repeated_ckv_block = compressed_kv_block[:, 0, :, :].expand(
-            batch_size, num_heads_to_repeat, -1, module.kv_lora_rank
-        )
-        compressed_kv_block = torch.cat((compressed_kv_block, repeated_ckv_block), dim=1)
 
-        repeated_k_pe_block = k_pe_block[:, 0, :, :].expand(
-            batch_size, num_heads_to_repeat, -1, module.qk_rope_head_dim
-        )
-        k_pe_block = torch.cat((k_pe_block, repeated_k_pe_block), dim=1)
+        if k_heads > 1:
+            num_heads_to_repeat = math.ceil(q_heads / k_heads)
+            compressed_kv_block = (
+                compressed_kv_block.unsqueeze(2)
+                .expand(-1, -1, num_heads_to_repeat, -1, -1)
+                .reshape(batch_size, num_heads_to_repeat * k_heads, -1, module.config.kv_lora_rank)
+            )
+            compressed_kv_block = compressed_kv_block[:, :q_heads, :, :]
+
+            k_pe_block = (
+                k_pe_block.unsqueeze(2)
+                .expand(-1, -1, num_heads_to_repeat, -1, -1)
+                .reshape(batch_size, num_heads_to_repeat * k_heads, -1, module.config.qk_rope_head_dim)
+            )
+            k_pe_block = k_pe_block[:, :q_heads, :, :]
 
         if absorption:
             krope_nope = torch.cat((compressed_kv_block, k_pe_block), dim=-1)
-            k_heads, q_heads = krope_nope.shape[1], query.shape[1]
-            num_heads_to_repeat = q_heads - k_heads
-            repeated_k = krope_nope[:, 0, :, :].expand(
-                batch_size, num_heads_to_repeat, -1, module.qk_rope_head_dim + module.kv_lora_rank
-            )
-            krope_nope = torch.cat((krope_nope, repeated_k), dim=1)
             attn_weights_block = torch.matmul(query, krope_nope.transpose(2, 3)) * scaling
             # [1, 64, q_len, 576] X [1, 1, 576, kv_block_size] -> [1, 64, q_len, kv_block_size]
             attn_weights_block = torch.where(causal_mask_block, masked_tensor, attn_weights_block)
@@ -930,20 +930,8 @@ def blocked_kv_mla_attention_forward(
                 skip_future,
             )  # [1, 64, q_len, kv_block_size] X [1, 1, kv_block_size, 512] -> [1, 64, q_len, 512]
         else:
-            k_heads, q_heads = compressed_kv_block.shape[1], query.shape[1]
-            num_heads_to_repeat = q_heads - k_heads
-            repeated_ckv_block = compressed_kv_block[:, 0, :, :].expand(
-                batch_size, num_heads_to_repeat, -1, module.kv_lora_rank
-            )
-            compressed_kv_block = torch.cat((compressed_kv_block, repeated_ckv_block), dim=1)
             knope = torch.matmul(compressed_kv_block, per_head_k_up_normal)
-
-            repeated_k_pe_block = k_pe_block[:, 0, :, :].expand(
-                batch_size, num_heads_to_repeat, -1, module.qk_rope_head_dim
-            )
-            k_pe_block = torch.cat((k_pe_block, repeated_k_pe_block), dim=1)
-
-            krope_nope = torch.cat((knope, k_pe_block.expand(-1, num_heads, -1, -1)), dim=-1)
+            krope_nope = torch.cat((knope, k_pe_block), dim=-1)
             attn_weights_block = torch.matmul(query, krope_nope.transpose(2, 3)) * scaling
             attn_weights_block = torch.where(causal_mask_block, masked_tensor, attn_weights_block)
             current_max, current_denominator, output = update_running_softmax(

QEfficient/transformers/models/deepseek_v3/modeling_deepseek.py

@@ -451,12 +451,26 @@ def fused_forward_orig(
             k_pe = compressed_kvs.update_k_pe(k_pe, self.layer_idx, cache_kwargs)
 
         k_heads, q_heads = kva.shape[1], q_pe.shape[1]
-        num_heads_to_repeat = q_heads - k_heads
-        repeated_kva = kva[:, 0, :, :].expand(bsz, num_heads_to_repeat, -1, self.kv_lora_rank)
-        kva_expanded = torch.cat((kva, repeated_kva), dim=1)
 
-        repeated_k_pe = k_pe[:, 0, :, :].expand(bsz, num_heads_to_repeat, -1, self.qk_rope_head_dim)
-        k_pe_expanded = torch.cat((k_pe, repeated_k_pe), dim=1)
+        if k_heads > 1:
+            num_heads_to_repeat = math.ceil(q_heads / k_heads)
+
+            kva_expanded = (
+                kva.unsqueeze(2)
+                .expand(-1, -1, num_heads_to_repeat, -1, -1)
+                .reshape(bsz, num_heads_to_repeat * k_heads, -1, self.config.kv_lora_rank)
+            )
+            kva_expanded = kva_expanded[:, :q_heads, :, :]
+
+            k_pe_expanded = (
+                k_pe.unsqueeze(2)
+                .expand(-1, -1, num_heads_to_repeat, -1, -1)
+                .reshape(bsz, num_heads_to_repeat * k_heads, -1, self.config.qk_rope_head_dim)
+            )
+            k_pe_expanded = k_pe_expanded[:, :q_heads, :, :]
+        else:
+            kva_expanded = kva
+            k_pe_expanded = k_pe
 
         v_up_per_head = self.v_up.squeeze(0).view(self.kv_lora_rank, self.num_heads, self.v_head_dim).permute(1, 0, 2)
         value_states = torch.matmul(kva_expanded, v_up_per_head)