MiniCPM-SALA：让Transformer在百万token下跑起来

┌─────────────────────────────────────────────────────┐
│  Layer 0:  Lightning Attention (线性注意力)          │  高效全局建模
├─────────────────────────────────────────────────────┤
│  Layer 1:  Lightning Attention                       │  高效全局建模
├─────────────────────────────────────────────────────┤
│  Layer 2:  Lightning Attention                       │  高效全局建模
├─────────────────────────────────────────────────────┤
│  Layer 3:  InfLLM-V2 (稀疏注意力)                    │  高保真局部精确
├─────────────────────────────────────────────────────┤
│  Layer 4:  Lightning Attention                       │  高效全局建模
├─────────────────────────────────────────────────────┤
│  ... 每4层重复一次 ...                               │
└─────────────────────────────────────────────────────┘

原始: [Attention] → [Attention] → [Attention] → [Attention] → ...
                              ↓ HALO转换
混合: [Linear] → [Linear] → [Linear] → [Sparse] → [Linear] → ...

import torch
import torch.nn as nn

class SparseAttention(nn.Module):
    """InfLLM-V2风格的稀疏注意力"""

    def __init__(self, hidden_dim, num_heads, window_size=1024, num_global_tokens=64):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.window_size = window_size
        self.num_global_tokens = num_global_tokens

        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.o_proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # 投影
        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)

        # 选择全局关键token（简化版：随机选择）
        global_indices = torch.randperm(seq_len)[:self.num_global_tokens]
        k_global = k[:, global_indices]
        v_global = v[:, global_indices]

        # 局部窗口注意力
        outputs = []
        for i in range(seq_len):
            # 窗口范围
            start = max(0, i - self.window_size // 2)
            end = min(seq_len, i + self.window_size // 2)

            # 合并局部和全局
            k_local = k[:, start:end]
            v_local = v[:, start:end]
            k_combined = torch.cat([k_local, k_global], dim=1)
            v_combined = torch.cat([v_local, v_global], dim=1)

            # 计算注意力
            attn = torch.matmul(q[:, i:i+1], k_combined.transpose(-2, -1))
            attn = torch.softmax(attn / (self.head_dim ** 0.5), dim=-1)
            out = torch.matmul(attn, v_combined)
            outputs.append(out)

        output = torch.cat(outputs, dim=1)
        return self.o_proj(output.view(batch_size, seq_len, -1))

class LightningAttention(nn.Module):
    """Lightning Attention线性注意力实现"""

    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads

        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.o_proj = nn.Linear(hidden_dim, hidden_dim)

    def feature_map(self, x):
        """特征映射函数：ELU + 1"""
        return torch.nn.functional.elu(x) + 1

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # 投影
        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)

        # 应用特征映射
        q = self.feature_map(q)
        k = self.feature_map(k)

        # 线性注意力: (K^T V) 先算，再与 Q 相乘
        # 复杂度: O(N * d^2) 而非 O(N^2 * d)

        # 计算累积和（left-to-right scan）
        # kv = cumsum(k * v)
        kv = torch.matmul(k.transpose(-2, -1), v)  # [batch, heads, d, d]

        # 应用RoPE（线性注意力层有位置编码）
        # ... RoPE实现省略 ...

        # 最终输出
        output = torch.matmul(q, kv)  # [batch, seq, heads, d]

        # 归一化
        k_sum = k.sum(dim=1, keepdim=True)  # [batch, 1, heads, d]
        output = output / (torch.matmul(q, k_sum.unsqueeze(-1)) + 1e-6)

        return self.o_proj(output.view(batch_size, seq_len, -1))

class HybridTransformerBlock(nn.Module):
    """混合Transformer块：根据层索引选择注意力类型"""

    def __init__(self, hidden_dim, num_heads, layer_idx, sparse_ratio=4):
        super().__init__()

        # 每 sparse_ratio 层使用一次稀疏注意力
        use_sparse = (layer_idx % sparse_ratio == 0)

        if use_sparse:
            self.attention = SparseAttention(hidden_dim, num_heads)
            self.use_rope = False  # 稀疏注意力不用RoPE
        else:
            self.attention = LightningAttention(hidden_dim, num_heads)
            self.use_rope = True  # 线性注意力用RoPE

        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 4),
            nn.GELU(),
            nn.Linear(hidden_dim * 4, hidden_dim)
        )

        self.norm1 = nn.LayerNorm(hidden_dim)
        self.norm2 = nn.LayerNorm(hidden_dim)

        # 输出门
        self.output_gate = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        # 注意力 + 残差
        attn_out = self.attention(self.norm1(x))
        x = x + torch.sigmoid(self.output_gate(attn_out)) * attn_out

        # FFN + 残差
        x = x + self.ffn(self.norm2(x))

        return x