import torch
import torch.nn as nn

class Model(nn.Module):
    """
    原始模型：RMSNorm + Residual
    """
    def __init__(self, dim, eps=1e-6):
        super(Model, self).__init__()
        self.eps = eps
        self.weight = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
        """
        原始实现：先RMSNorm，再与residual相加
        
        Args:
            x (torch.Tensor): 输入张量 shape (B, L, D)
            residual (torch.Tensor): 残差张量 shape (B, L, D)
        
        Returns:
            torch.Tensor: RMSNorm(x) + residual
        """
        dtype = x.dtype
        x = x.float()
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        rmsnorm_output = (x * self.weight).to(dtype)
        return rmsnorm_output + residual

class ModelNew(torch.nn.Module):
    """
    融合模型：直接实现RMSNorm + Residual
    """
    def __init__(self, dim, eps=1e-6):
        super(ModelNew, self).__init__()
        self.eps = eps
        self.weight = torch.nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
        """
        融合实现：在CUDA kernel中直接完成RMSNorm + Residual
        
        Args:
            x (torch.Tensor): 输入张量 shape (B, L, D)
            residual (torch.Tensor): 残差张量 shape (B, L, D)
        
        Returns:
            torch.Tensor: RMSNorm(x) + residual
        """
        # 这个将在CUDA中实现
        pass

# 测试参数
batch_size = 16
seq_len = 512
dim = 4096

def get_inputs():
    """生成测试输入"""
    x = torch.randn(batch_size, seq_len, dim)
    residual = torch.randn(batch_size, seq_len, dim)
    return [x, residual]

def get_init_inputs():
    """获取初始化参数"""
    return [dim]
