# Standard MLA flow in MLASelfAttention.get_query_key_value_tensors()
kv, _ = self.linear_kv_up_proj(kv_compressed)  # [s, b, n*(qk_head_dim + v_head_dim)]
kv = kv.view(..., self.num_attention_heads_per_partition, qk_head_dim + v_head_dim)
k_no_pe, value = torch.split(kv, [qk_head_dim, v_head_dim], dim=-1)
# Core attention receives: Q[s,b,n,d], K[s,b,n,d], V[s,b,n,d] - all multi-head

# Core attention in AbsorbedMLA receives different shapes:
core_attention(
    q_absorbed,    # [s, b, n, kv_lora_rank + qk_pos_emb_head_dim] - multi-head
    kv_compressed, # [s, b, 1, kv_lora_rank + qk_pos_emb_head_dim] - SINGLE head
    v=None,        # V is not passed; absorbed impl applies V_up after attention
    ...
)

# In model spec configuration (not shown in this PR, but expected usage):
layer_spec = TransformerLayerSpec(
    self_attention=ModuleSpec(
        module=AbsorbedMLASelfAttention,
        submodules=AbsorbedMLASelfAttentionSubmodules(
            linear_k_up_proj=...,
            linear_v_up_proj=...,
            # ... other submodules
        )
    )
)

def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb):
    # Step 1: Apply Q up projection
    # Why: Transform compressed Q to full attention head dimension
    if self.config.q_lora_rank is not None:
        q, _ = self.linear_q_up_proj(q_compressed)  # [num_tokens, n * q_head_dim]
    else:
        q, _ = self.linear_q_proj(q_compressed)
    
    # Step 2: Reshape Q for per-head processing
    # q: [num_tokens, n, qk_head_dim + qk_pos_emb_head_dim]
    q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim)
    
    # Step 3: Prepare K up-projection weight for absorption
    # Key insight: We don't apply this to KV, but absorb into Q instead
    # k_up_weight shape: [n, qk_head_dim, kv_lora_rank]
    k_up_weight = self.linear_k_up_proj.weight.view(
        self.num_attention_heads_per_partition,
        self.config.qk_head_dim,
        self.config.kv_lora_rank,
    )
    
    # Step 4: Split Q into non-positional and positional parts
    q_no_pe, q_pos_emb = torch.split(
        q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1
    )
    
    # Step 5: THE ABSORPTION - multiply Q by K's up-projection weight
    # Mathematical equivalence: Q @ K = (Q @ K_up^T) @ K_compressed
    # q_absorbed: [num_tokens, n, kv_lora_rank]
    q_absorbed = torch.einsum("...nd,ndk->...nk", q_no_pe, k_up_weight)
    q_absorbed = q_absorbed.contiguous()
    
    # Step 6: Apply RoPE to positional components
    q_pos_emb = apply_rotary_pos_emb(q_pos_emb, rotary_pos_emb, ...)
    k_pos_emb = apply_rotary_pos_emb(k_pos_emb, rotary_pos_emb, ...)
    
    # Step 7: Combine absorbed Q with positional embeddings
    # Final Q: [num_tokens, n, kv_lora_rank + qk_pos_emb_head_dim]
    q_absorbed = torch.cat([q_absorbed, q_pos_emb], dim=-1)
    # KV remains single-head: [num_tokens, 1, kv_lora_rank + qk_pos_emb_head_dim]
    kv_compressed = torch.cat([kv_compressed, k_pos_emb], dim=-1)
    
    return q_absorbed, kv_compressed

def forward(self, hidden_states, attention_mask, ...):
    # ... QKV computation ...
    
    # Core attention with MQA-style inputs
    core_attn_out = self.core_attention(
        q_absorbed,      # [s, b, n, kv_lora_rank + qk_pos_emb]
        kv_compressed,   # [s, b, 1, kv_lora_rank + qk_pos_emb] - SINGLE HEAD
        v=None,          # V is not provided; we apply V_up after attention
        ...
    )
    
    # ==================================
    # Apply V up projection AFTER core attention
    # ==================================
    # v_up_weight shape: [n, v_head_dim, kv_lora_rank]
    v_up_weight = self.linear_v_up_proj.weight.view(
        self.num_attention_heads_per_partition, 
        self.config.v_head_dim, 
        self.config.kv_lora_rank
    )
    
    # Reshape attention output for einsum
    # core_attn_out: [s, b, n * kv_lora_rank] -> [s, b, n, kv_lora_rank]
    core_attn_out = core_attn_out.view(
        *core_attn_out.shape[:-1],
        self.num_attention_heads_per_partition,
        self.config.kv_lora_rank,
    )
    
    # Apply V up-projection: [s, b, n, kv_lora_rank] @ [n, v_head_dim, kv_lora_rank]^T
    # Result: [s, b, n, v_head_dim]
    core_attn_out = torch.einsum("...nc,ndc->...nd", core_attn_out, v_up_weight)
    core_attn_out = core_attn_out.contiguous()
    
    # Flatten for linear_proj: [s, b, n * v_head_dim]
    core_attn_out = core_attn_out.view(*core_attn_out.shape[:-2], -1)
    
    # Output projection
    output, bias = self.linear_proj(core_attn_out)
    return output, bias

def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
    """Handle loading from checkpoints with combined KV up projection weights.
    
    Standard MLA stores: linear_kv_up_proj.weight with interleaved K,V per head
        Layout: [head0_K, head0_V, head1_K, head1_V, ...]
    
    AbsorbedMLA needs: separate linear_k_up_proj and linear_v_up_proj
        K: [head0_K, head1_K, ...]
        V: [head0_V, head1_V, ...]
    """
    combined_key = f'{prefix}linear_kv_up_proj.weight'
    k_up_key = f'{prefix}linear_k_up_proj.weight'
    v_up_key = f'{prefix}linear_v_up_proj.weight'
    
    if combined_key in state_dict:
        combined_weight = state_dict[combined_key]
        
        # Split with proper per-head de-interleaving
        k_weight, v_weight = self._split_kv_weights(combined_weight)
        
        state_dict[k_up_key] = k_weight
        state_dict[v_up_key] = v_weight
        del state_dict[combined_key]
    
    super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)