ops-transformer大模型算子库—Transformer模型高效加速的关键技术

引言

Transformer架构已成为现代自然语言处理和多模态模型的主流选择，从BERT、GPT到LLaMA、Qwen，各类大语言模型（LLM）层出不穷。CANN开源生态中的 ops-transformer 是专门针对Transformer类模型优化的进阶算子库，提供了注意力机制、MoE（混合专家）等核心组件的高性能实现，为大模型的训练和推理提供了强有力的加速支持。

ops-transformer算子库概述

ops-transformer是CANN生态中为Transformer类模型设计的专用算子库，包含以下核心算子类别：

核心技术特点

1. 高效注意力机制实现

注意力机制是Transformer的核心组件，ops-transformer提供了多种优化的注意力实现：

"""
ops-transformer注意力算子示例
展示高效的自注意力实现
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional


class OptimizedSelfAttention(nn.Module):
    """
    优化的自注意力实现
    使用ops-transformer提供的高效算子
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        use_flash_attention: bool = True
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.use_flash_attention = use_flash_attention

        # Q、K、V投影
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        # 输出投影
        self.out_proj = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

        # 缩放因子
        self.scale = self.head_dim ** -0.5

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        causal: bool = False
    ) -> torch.Tensor:
        """
        前向传播

        Args:
            x: 输入张量 [batch_size, seq_len, embed_dim]
            attention_mask: 注意力掩码 [batch_size, seq_len, seq_len]
            causal: 是否使用因果掩码（用于解码器）

        Returns:
            输出张量 [batch_size, seq_len, embed_dim]
        """
        batch_size, seq_len, embed_dim = x.shape

        # 计算Q、K、V
        Q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        K = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        V = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim)

        # 转置为 [batch_size, num_heads, seq_len, head_dim]
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 选择注意力计算方式
        if self.use_flash_attention:
            attn_output = self._flash_attention(Q, K, V, causal)
        else:
            attn_output = self._standard_attention(Q, K, V, attention_mask, causal)

        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, embed_dim)

        # 输出投影
        output = self.out_proj(attn_output)
        output = self.dropout(output)

        return output

    def _flash_attention(self, Q, K, V, causal):
        """
        FlashAttention实现
        ops-transformer提供的高效算子
        """
        # 实际ops-transformer中有完整的FlashAttention实现
        # 这里展示简化版本的概念

        batch_size, num_heads, seq_len, head_dim = Q.shape

        # 计算注意力分数
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale

        # 应用因果掩码（如果需要）
        if causal:
            mask = torch.tril(torch.ones(seq_len, seq_len, device=Q.device))
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        # Softmax和加权
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_output = torch.matmul(attn_weights, V)

        return attn_output

    def _standard_attention(self, Q, K, V, attention_mask, causal):
        """标准注意力实现"""
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale

        # 应用掩码
        if attention_mask is not None:
            attn_scores = attn_scores + attention_mask

        if causal:
            mask = torch.tril(torch.ones(Q.size(2), Q.size(2), device=Q.device))
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, V)

        return attn_output


class RotaryEmbedding(nn.Module):
    """
    旋转位置编码（RoPE）
    ops-transformer提供的优化实现
    """

    def __init__(self, dim: int, max_seq_len: int = 2048):
        super().__init__()
        self.dim = dim

        # 计算旋转角度
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

        # 预计算位置编码
        self._build_cache(max_seq_len)

    def _build_cache(self, max_seq_len: int):
        """构建位置编码缓存"""
        t = torch.arange(max_seq_len, device=self.inv_freq.device)
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer('cos_cached', emb.cos()[None, None, :, :])
        self.register_buffer('sin_cached', emb.sin()[None, None, :, :])

    def forward(self, x: torch.Tensor, seq_len: int) -> tuple:
        """
        应用旋转位置编码

        Args:
            x: 输入张量 [batch_size, num_heads, seq_len, head_dim]
            seq_len: 序列长度

        Returns:
            旋转后的Q和K
        """
        cos = self.cos_cached[:, :, :seq_len, :]
        sin = self.sin_cached[:, :, :seq_len, :]

        x1, x2 = x[..., :self.dim//2], x[..., self.dim//2:]
        rotated = torch.cat([
            x1 * cos - x2 * sin,
            x1 * sin + x2 * cos
        ], dim=-1)

        return rotated


class AttentionWithRoPE(nn.Module):
    """
    带旋转位置编码的注意力
    """

    def __init__(self, embed_dim: int, num_heads: int, max_seq_len: int = 2048):
        super().__init__()
        self.attention = OptimizedSelfAttention(embed_dim, num_heads)
        self.rotary_emb = RotaryEmbedding(embed_dim // num_heads, max_seq_len)

    def forward(self, x: torch.Tensor, causal: bool = False) -> torch.Tensor:
        """
        前向传播（带RoPE）
        """
        batch_size, seq_len, embed_dim = x.shape

        # 投影到Q、K、V
        Q = self.attention.q_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)
        K = self.attention.k_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)
        V = self.attention.v_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)

        # 转置
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 应用旋转位置编码
        Q = self.rotary_emb(Q, seq_len)
        K = self.rotary_emb(K, seq_len)

        # 计算注意力
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.attention.scale

        if causal:
            mask = torch.tril(torch.ones(seq_len, seq_len, device=Q.device))
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_output = torch.matmul(attn_weights, V)

        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, embed_dim)

        # 输出投影
        output = self.attention.out_proj(attn_output)

        return output


# 使用示例
def test_attention_variants():
    """测试不同注意力变体"""
    batch_size = 4
    seq_len = 512
    embed_dim = 768
    num_heads = 12

    # 准备输入
    x = torch.randn(batch_size, seq_len, embed_dim)

    # 测试标准注意力
    print("=== 测试标准注意力 ===")
    attn_standard = OptimizedSelfAttention(embed_dim, num_heads, use_flash_attention=False)
    attn_standard.eval()
    with torch.no_grad():
        output_standard = attn_standard(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output_standard.shape}")

    # 测试Flash Attention
    print("\n=== 测试Flash Attention ===")
    attn_flash = OptimizedSelfAttention(embed_dim, num_heads, use_flash_attention=True)
    attn_flash.eval()
    with torch.no_grad():
        output_flash = attn_flash(x)
    print(f"输出形状: {output_flash.shape}")

    # 测试带RoPE的注意力
    print("\n=== 测试带RoPE的注意力 ===")
    attn_rope = AttentionWithRoPE(embed_dim, num_heads)
    attn_rope.eval()
    with torch.no_grad():
        output_rope = attn_rope(x, causal=True)
    print(f"输出形状: {output_rope.shape}")


if __name__ == "__main__":
    test_attention_variants()

2. MoE（混合专家）优化算子

MoE是提升大模型容量的关键技术，ops-transformer提供了专门的MoE优化算子：

"""
ops-transformer MoE算子示例
展示混合专家模型的优化实现
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple


class MoEExpert(nn.Module):
    """
    MoE专家网络
    """

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.fc1 = nn.Linear(hidden_size, intermediate_size)
        self.fc2 = nn.Linear(intermediate_size, hidden_size)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播

        Args:
            x: [batch_size, seq_len, hidden_size]

        Returns:
            输出 [batch_size, seq_len, hidden_size]
        """
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x


class OptimizedMoELayer(nn.Module):
    """
    优化的MoE层
    使用ops-transformer的MLAPO融合算子
    """

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int,
        top_k: int = 2,
        capacity_factor: float = 1.0
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor

        # 门控网络
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)

        # 专家网络
        self.experts = nn.ModuleList([
            MoEExpert(hidden_size, intermediate_size)
            for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, dict]:
        """
        前向传播

        Args:
            x: [batch_size, seq_len, hidden_size]

        Returns:
            输出 [batch_size, seq_len, hidden_size]
            辅助信息字典
        """
        batch_size, seq_len, hidden_size = x.shape

        # 计算门控分数
        gate_logits = self.gate(x)  # [batch_size, seq_len, num_experts]
        gate_probs = F.softmax(gate_logits, dim=-1)

        # 选择top-k专家
        top_k_probs, top_k_indices = torch.topk(gate_probs, self.top_k, dim=-1)
        # 归一化top-k概率
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 计算容量因子
        capacity = int(seq_len * self.capacity_factor)

        # 初始化输出
        output = torch.zeros_like(x)

        # 辅助信息统计
        expert_stats = {i: 0 for i in range(self.num_experts)}

        # 为每个专家计算贡献
        for k in range(self.top_k):
            # 获取当前轮次的专家选择
            expert_indices = top_k_indices[:, :, k]  # [batch_size, seq_len]
            expert_weights = top_k_probs[:, :, k]    # [batch_size, seq_len]

            # 遍历每个专家
            for expert_id in range(self.num_experts):
                # 创建当前专家的掩码
                mask = (expert_indices == expert_id)
                if not mask.any():
                    continue

                # 统计使用情况
                expert_stats[expert_id] += mask.sum().item()

                # 获取该专家的输入和权重
                expert_input = x[mask]  # [num_tokens, hidden_size]
                expert_weight = expert_weights[mask]  # [num_tokens]

                # 限制容量（简单丢弃策略）
                if expert_input.shape[0] > capacity:
                    expert_input = expert_input[:capacity]
                    expert_weight = expert_weight[:capacity]

                # 专家计算
                expert_output = self.experts[expert_id](expert_input)

                # 累加到输出（需要scatter操作，这里简化）
                # 在实际ops-transformer的MLAPO算子中，这是融合的
                output_index = mask.nonzero(as_tuple=False)[:expert_input.shape[0]]
                output[output_index[:, 0], output_index[:, 1]] += (
                    expert_output * expert_weight.unsqueeze(-1)
                )

        return output, expert_stats


class MLPOMoEFusion(nn.Module):
    """
    MLAPO融合算子的概念演示
    实际ops-transformer中的实现是高度优化的
    """

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        num_experts: int
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_experts = num_experts

        # 融合的专家权重（简化表示）
        self.expert_weights = nn.Parameter(
            torch.randn(num_experts, 2, intermediate_size, hidden_size)
        )

        # 门控权重
        self.gate_weight = nn.Parameter(torch.randn(num_experts, hidden_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        MLAPO融合前向传播（概念版）

        实际ops-transformer中的实现会：
        1. 融合门控计算和专家路由
        2. 并行执行多个专家计算
        3. 融合结果聚合
        """
        batch_size, seq_len, hidden_size = x.shape

        # 融合的门控计算
        gate_logits = torch.matmul(x, self.gate_weight.t())
        gate_probs = F.softmax(gate_logits, dim=-1)

        # 融合的专家计算（简化版）
        output = torch.zeros_like(x)

        for expert_id in range(self.num_experts):
            # 门控权重
            weight = gate_probs[:, :, expert_id:expert_id+1]

            # 专家计算（这里用简化的线性层模拟）
            expert_output = torch.matmul(
                x,
                self.expert_weights[expert_id, 0].t()
            )
            expert_output = F.gelu(expert_output)
            expert_output = torch.matmul(
                expert_output,
                self.expert_weights[expert_id, 1].t()
            )

            # 加权累加
            output += expert_output * weight

        return output


class TransformerBlockWithMoE(nn.Module):
    """
    带MoE的Transformer块
    """

    def __init__(
        self,
        hidden_size: int,
        ffn_hidden_size: int,
        num_heads: int,
        num_experts: int,
        top_k: int = 2
    ):
        super().__init__()

        # 注意力层
        self.attention = OptimizedSelfAttention(hidden_size, num_heads)

        # 归一化
        self.attention_norm = nn.LayerNorm(hidden_size)
        self.moe_norm = nn.LayerNorm(hidden_size)

        # MoE层
        self.moe = OptimizedMoELayer(
            hidden_size,
            ffn_hidden_size,
            num_experts,
            top_k
        )

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, dict]:
        """
        前向传播
        """
        # 自注意力（带残差）
        residual = x
        x = self.attention_norm(x)
        x = self.attention(x)
        x = x + residual

        # MoE FFN（带残差）
        residual = x
        x = self.moe_norm(x)
        x, expert_stats = self.moe(x)
        x = x + residual

        return x, expert_stats


# 使用示例
def test_moe_layer():
    """测试MoE层"""
    batch_size = 2
    seq_len = 128
    hidden_size = 512
    ffn_hidden_size = 2048
    num_experts = 8
    num_heads = 8

    # 创建MoE层
    moe_layer = OptimizedMoELayer(
        hidden_size,
        ffn_hidden_size,
        num_experts,
        top_k=2
    )
    moe_layer.eval()

    # 准备输入
    x = torch.randn(batch_size, seq_len, hidden_size)

    # 前向传播
    with torch.no_grad():
        output, expert_stats = moe_layer(x)

    print(f"=== MoE层测试 ===")
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"专家使用统计:")
    for expert_id, count in expert_stats.items():
        print(f"  Expert {expert_id}: {count} tokens")


def test_transformer_with_moe():
    """测试带MoE的Transformer"""
    batch_size = 2
    seq_len = 128
    hidden_size = 512
    ffn_hidden_size = 2048
    num_heads = 8
    num_experts = 8

    # 创建Transformer块
    block = TransformerBlockWithMoE(
        hidden_size,
        ffn_hidden_size,
        num_heads,
        num_experts,
        top_k=2
    )
    block.eval()

    # 准备输入
    x = torch.randn(batch_size, seq_len, hidden_size)

    # 前向传播
    with torch.no_grad():
        output, expert_stats = block(x)

    print(f"\n=== Transformer块with MoE测试 ===")
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"专家使用统计:")
    for expert_id, count in expert_stats.items():
        print(f"  Expert {expert_id}: {count} tokens")


if __name__ == "__main__":
    test_moe_layer()
    test_transformer_with_moe()

3. RMSNorm归一化

RMSNorm是现代大模型广泛使用的归一化方法：

"""
ops-transformer归一化算子示例
展示RMSNorm等归一化方法
"""
import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    """
    RMSNorm归一化
    ops-transformer提供的高效实现
    """

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        RMSNorm: x / sqrt(mean(x^2) + eps) * weight

        Args:
            x: [batch_size, seq_len, dim]

        Returns:
            归一化后的张量
        """
        # 计算均方根
        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)

        # 归一化并缩放
        output = x / rms * self.weight

        return output


class SwiGLU(nn.Module):
    """
    SwiGLU激活函数
    ops-transformer提供的优化实现
    """

    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        self_gate = nn.Linear(dim, hidden_dim, bias=False)
        self.value = nn.Linear(dim, hidden_dim, bias=False)
        self.out = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        SwiGLU(x) = (Swish(x @ W_gate) * (x @ W_value)) @ W_out
        """
        gate = torch.sigmoid(self.gate(x))
        value = self.value(x)
        output = self.out(gate * value)
        return output


class TransformerFFN(nn.Module):
    """
    Transformer FFN层
    使用ops-transformer的优化算子
    """

    def __init__(
        self,
        dim: int,
        hidden_dim: int,
        use_swiglu: bool = True
    ):
        super().__init__()

        self.norm = RMSNorm(dim)

        if use_swiglu:
            self.ffn = SwiGLU(dim, hidden_dim)
        else:
            self.ffn = nn.Sequential(
                nn.Linear(dim, hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, dim)
            )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        """
        residual = x
        x = self.norm(x)
        x = self.ffn(x)
        x = x + residual
        return x


# 使用示例
def test_normalization_and_activation():
    """测试归一化和激活函数"""
    batch_size = 4
    seq_len = 128
    dim = 512
    hidden_dim = 2048

    # 准备输入
    x = torch.randn(batch_size, seq_len, dim)

    # 测试RMSNorm
    print("=== 测试RMSNorm ===")
    rms_norm = RMSNorm(dim)
    output = rms_norm(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"输入均值: {x.mean():.4f}, 标准差: {x.std():.4f}")
    print(f"输出均值: {output.mean():.4f}, 标准差: {output.std():.4f}")

    # 测试FFN with SwiGLU
    print("\n=== 测试FFN with SwiGLU ===")
    ffn = TransformerFFN(dim, hidden_dim, use_swiglu=True)
    ffn.eval()
    with torch.no_grad():
        output_ffn = ffn(x)
    print(f"输出形状: {output_ffn.shape}")


if __name__ == "__main__":
    test_normalization_and_activation()

完整示例：GPT风格Decoder层

"""
完整的GPT风格Decoder层示例
整合ops-transformer的多种算子
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple


class GPTDecoderBlock(nn.Module):
    """
    GPT风格的Decoder块
    使用ops-transformer提供的优化算子
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        ffn_hidden_size: int,
        max_seq_len: int = 2048,
        dropout: float = 0.0,
        use_swiglu: bool = True
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads

        # 注意力层（带RoPE）
        self.attention_norm = RMSNorm(hidden_size)
        self.attention = OptimizedSelfAttention(
            hidden_size,
            num_heads,
            dropout=dropout,
            use_flash_attention=True
        )
        self.rotary_emb = RotaryEmbedding(
            hidden_size // num_heads,
            max_seq_len
        )

        # FFN层
        self.ffn_norm = RMSNorm(hidden_size)

        if use_swiglu:
            self.ffn_gate = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
            self.ffn_up = nn.Linear(hidden_size, ffn_hidden_size, bias=False)
            self.ffn_down = nn.Linear(ffn_hidden_size, hidden_size, bias=False)
            self.act = nn.SiLU()
        else:
            self.ffn = nn.Sequential(
                nn.Linear(hidden_size, ffn_hidden_size),
                nn.GELU(),
                nn.Linear(ffn_hidden_size, hidden_size)
            )

        self.use_swiglu = use_swiglu
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        前向传播

        Args:
            x: [batch_size, seq_len, hidden_size]
            attention_mask: [batch_size, seq_len, seq_len]
            position_ids: [batch_size, seq_len]

        Returns:
            [batch_size, seq_len, hidden_size]
        """
        batch_size, seq_len, _ = x.shape

        # === 自注意力（带残差） ===
        residual = x
        x = self.attention_norm(x)

        # 投影到Q、K、V
        Q = self.attention.q_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)
        K = self.attention.k_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)
        V = self.attention.v_proj(x).view(batch_size, seq_len, -1, self.attention.head_dim)

        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 应用旋转位置编码
        Q = self.rotary_emb(Q, seq_len)
        K = self.rotary_emb(K, seq_len)

        # 计算注意力
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.attention.scale

        # 应用因果掩码
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
        attn_scores = attn_scores.masked_fill(causal_mask == 0, float('-inf'))

        # 应用额外掩码（如果提供）
        if attention_mask is not None:
            attn_scores = attn_scores + attention_mask

        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.attention.dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, V)

        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.hidden_size)

        # 输出投影
        attn_output = self.attention.out_proj(attn_output)
        attn_output = self.attention.dropout(attn_output)

        # 残差连接
        x = residual + attn_output

        # === FFN（带残差） ===
        residual = x
        x = self.ffn_norm(x)

        if self.use_swiglu:
            # SwiGLU激活
            gate = self.act(self.ffn_gate(x))
            up = self.ffn_up(x)
            ffn_output = self.ffn_down(gate * up)
        else:
            ffn_output = self.ffn(x)

        x = residual + self.dropout(ffn_output)

        return x


class MiniGPTModel(nn.Module):
    """
    简化的GPT模型
    展示ops-transformer算子的完整使用
    """

    def __init__(
        self,
        vocab_size: int,
        hidden_size: int,
        num_layers: int,
        num_heads: int,
        ffn_hidden_size: int,
        max_seq_len: int = 2048
    ):
        super().__init__()

        # 词嵌入
        self.embedding = nn.Embedding(vocab_size, hidden_size)

        # Transformer层
        self.layers = nn.ModuleList([
            GPTDecoderBlock(
                hidden_size,
                num_heads,
                ffn_hidden_size,
                max_seq_len
            )
            for _ in range(num_layers)
        ])

        # 最终归一化
        self.final_norm = RMSNorm(hidden_size)

        # 语言建模头
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        前向传播

        Args:
            input_ids: [batch_size, seq_len]
            attention_mask: [batch_size, seq_len]

        Returns:
            logits: [batch_size, seq_len, vocab_size]
        """
        x = self.embedding(input_ids)

        # 创建因果注意力掩码
        batch_size, seq_len = input_ids.shape
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)

        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
            attention_mask = (1.0 - attention_mask) * -10000.0
            causal_mask = causal_mask + attention_mask

        # 通过所有Transformer层
        for layer in self.layers:
            x = layer(x, causal_mask)

        x = self.final_norm(x)
        logits = self.lm_head(x)

        return logits


# 使用示例
def test_gpt_model():
    """测试GPT模型"""
    vocab_size = 50000
    hidden_size = 512
    num_layers = 6
    num_heads = 8
    ffn_hidden_size = 2048
    max_seq_len = 1024
    batch_size = 2

    # 创建模型
    model = MiniGPTModel(
        vocab_size,
        hidden_size,
        num_layers,
        num_heads,
        ffn_hidden_size,
        max_seq_len
    )
    model.eval()

    # 准备输入
    input_ids = torch.randint(0, vocab_size, (batch_size, 256))
    attention_mask = torch.ones(batch_size, 256)

    # 前向传播
    with torch.no_grad():
        logits = model(input_ids, attention_mask)

    print(f"=== GPT模型测试 ===")
    print(f"输入形状: {input_ids.shape}")
    print(f"输出形状: {logits.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")


if __name__ == "__main__":
    test_gpt_model()

应用场景

ops-transformer算子库适用于以下场景：

1. 大语言模型：GPT、LLaMA、Qwen、DeepSeek等模型的训练和推理
2. 代码生成模型：CodeLLaMA、StarCoder、CodeGeeX等
3. 多模态模型：视觉-语言模型、语音-语言模型
4. MoE模型：DeepSeek-MoE、Mixtral等混合专家模型
5. 长文本处理：长文档理解、长上下文对话

总结

ops-transformer作为CANN生态中针对Transformer模型优化的专用算子库，提供了从高效注意力、MoE优化到归一化和激活函数的完整算子集合。通过MLAPO融合算子、FlashAttention、RoPE位置编码等技术，ops-transformer能够充分释放硬件的计算能力，为大模型的训练和推理提供强有力的支撑。

相关链接

• CANN组织链接: https://atomgit.com/cann
• ops-transformer仓库链接: https://atomgit.com/cann/ops-transformer

参考资料

• CANN官方文档: https://www.hiascend.com/cann
• CANN开源项目: https://gitcode.com/cann