[tab] miles/backends/fsdp_utils/models/qwen3_patch.py (python)
import torch
import torch.nn as nn
import logging

logger = logging.getLogger(__name__)

def fuse_qkv_for_qwen3(model: nn.Module):
    """
    Fuses q_proj, k_proj, and v_proj into a single qkv_proj for each attention layer.
    This aligns the training model with SGLang's optimized inference architecture.
    """
    from transformers.models.qwen2.modeling_qwen2 import Qwen2Attention, Qwen2FlashAttention2, Qwen2SdpaAttention
    
    # We also need to handle cases where the model might be using a custom Qwen3 class 
    # if it's not strictly using the Qwen2 implementation.
    
    fused_count = 0
    for name, module in model.named_modules():
        if isinstance(module, (Qwen2Attention, Qwen2FlashAttention2, Qwen2SdpaAttention)):
            # Found an attention layer
            q_proj = module.q_proj
            k_proj = module.k_proj
            v_proj = module.v_proj
            
            hidden_size = q_proj.in_features
            
            # Qwen2/3 uses GQA/MQA often. 
            # q_proj: [num_heads * head_dim, hidden_size]
            # k_proj: [num_kv_heads * head_dim, hidden_size]
            # v_proj: [num_kv_heads * head_dim, hidden_size]
            
            out_features = q_proj.out_features + k_proj.out_features + v_proj.out_features
            
            # Create the fused layer
            qkv_proj = nn.Linear(hidden_size, out_features, bias=q_proj.bias is not None)
            
            # Fuse weights
            with torch.no_grad():
                fused_weight = torch.cat([q_proj.weight, k_proj.weight, v_proj.weight], dim=0)
                qkv_proj.weight.copy_(fused_weight)
                if q_proj.bias is not None:
                    fused_bias = torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0)
                    qkv_proj.bias.copy_(fused_bias)
            
            # Replace the modules
            module.qkv_proj = qkv_proj
            
            # Patch the forward method
            # We need to save the original forward or implement a new one that uses qkv_proj
            # For Qwen2, the forward pass calls q_proj, k_proj, v_proj.
            # We can monkey-patch the module's forward or just replace the proj calls.
            
            def fused_forward(self, hidden_states, *args, **kwargs):
                # This is a simplified version, we might need to be more careful 
                # with the exact signature of the specific Attention class.
                # However, most Qwen2-based models use these three projections at the start of forward.
                
                # We'll implement a more robust patch by replacing the individual proj calls 
                # if the class allows it, or by overriding the whole forward.
                pass

            # Actually, a cleaner way is to keep the q_proj, k_proj, v_proj attributes 
            # but make them "views" or "proxies" to the fused layer.
            # But that's complex with FSDP.
            
            # Better: Replace the forward method of the instance.
            # Let's look at the original forward of Qwen2Attention.
            
            fused_count += 1

    logger.info(f"Fused QKV for {fused_count} attention layers.")
    return model