第2章：新型架构探索 (New Architectures)#

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoELayer(nn.Module):
    def __init__(self, hidden_dim, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # 1. 门控网络 (Router)
        self.gate = nn.Linear(hidden_dim, num_experts)

        # 2. 专家网络 (这里用简单的 Linear 模拟 MLP)
        self.experts = nn.ModuleList([
            nn.Linear(hidden_dim, hidden_dim) for _ in range(num_experts)
        ])

    def forward(self, x):
        # x: [batch, seq, dim]
        batch, seq, dim = x.shape
        x_flat = x.view(-1, dim)

        # 1. 计算路由概率
        logits = self.gate(x_flat)  # [total_tokens, num_experts]
        probs = F.softmax(logits, dim=-1)

        # 2. 选出 Top-K 专家
        # indices: [total_tokens, top_k] (专家的 ID)
        # weights: [total_tokens, top_k] (路由权重)
        weights, indices = torch.topk(probs, self.top_k, dim=-1)

        # 归一化权重 (让选中的 k 个权重之和为 1)
        weights = weights / weights.sum(dim=-1, keepdim=True)

        # 3. 专家计算 (这里用循环模拟，实际用 CUDA Kernel 优化)
        final_output = torch.zeros_like(x_flat)
        for i in range(self.num_experts):
            # 找到分配给专家 i 的 token
            # mask: [total_tokens, top_k]
            mask = (indices == i)
            # batch_idx, k_idx: 哪些 token 的第几个选择是专家 i
            batch_idx, k_idx = torch.where(mask)

            if len(batch_idx) == 0:
                continue

            # 提取输入
            expert_input = x_flat[batch_idx]

            # 专家前向传播
            expert_output = self.experts[i](expert_input)

            # 加权累加回输出
            # weights[batch_idx, k_idx]: 对应的路由权重
            scale = weights[batch_idx, k_idx].unsqueeze(-1)
            final_output.index_add_(0, batch_idx, expert_output * scale)

        return final_output.view(batch, seq, dim)

# 测试
moe = MoELayer(hidden_dim=128, num_experts=8, top_k=2)
x = torch.randn(2, 10, 128)
out = moe(x)
print(f"Input: {x.shape}, Output: {out.shape}")

import torch
import torch.nn as nn

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16):
        super().__init__()
        # 1. 线性投影
        self.in_proj = nn.Linear(d_model, d_model * 2)
        self.x_proj = nn.Linear(d_model, d_model + d_state * 2) # 生成 delta, B, C

        # 2. SSM 参数 (A 是固定的/缓慢变化的)
        self.A_log = nn.Parameter(torch.log(torch.arange(1, d_state + 1, dtype=torch.float32)))

    def forward(self, x):
        # 简化版前向传播 (非并行优化的 Scan 实现)
        batch, seq, dim = x.shape
        states = torch.zeros(batch, dim, self.d_state).to(x.device)
        outputs = []

        for t in range(seq):
            xt = x[:, t, :]

            # 动态计算参数 (Selective Scan)
            # 这里省略了复杂的离散化 (Discretization) 步骤
            # h_t = A * h_{t-1} + B * x_t
            # y_t = C * h_t

            # ... (完整实现需 CUDA Kernel)
            pass
        return x # Placeholder