1inference/
 2├── model.py      # 923 行 - 核心模型定义
 3├── generate.py   # 187 行 - 推理生成逻辑
 4├── kernel.py     # 275 行 - Kernel 实现（TileLang JIT：act_quant / fp8_gemm / fp8_index）
 5├── convert.py    # 101 行 - 权重转换脚本（HF → 本仓库推理格式）
 6├── config_671B_v3.2.json  # 671B 推理 demo 的 ModelArgs 配置（与 generate.py 匹配）
 7├── requirements.txt       # 推理依赖（含 tilelang / fast_hadamard_transform）
 8└── README.md              # 推理 demo 使用说明
 9encoding/
10└── encoding_dsv32.py  # Chat/Tool 调用的 prompt 编码（与推理代码独立）

 1@dataclass
 2class ModelArgs:
 3    # 基础配置
 4    max_batch_size: int = 8           # 最大批次大小
 5    max_seq_len: int = 4096 * 4       # 最大序列长度 (16K，实际 V3.2 是 163K)
 6    dtype: Literal["bf16", "fp8"] = "bf16"  # 计算精度
 7    scale_fmt: Optional[str] = None   # FP8 激活量化 scale 的存储格式（如 ue8m0）
 8    vocab_size: int = 102400          # 词表大小 (实际 V3.2 是 129280)
 9    dim: int = 2048                   # 隐藏层维度 (实际 V3.2 是 7168)
10    inter_dim: int = 10944            # Dense MLP 的中间维度（实际 V3.2 是 18432）
11    moe_inter_dim: int = 1408         # MoE expert 的中间维度（实际 V3.2 是 2048）
12    n_layers: int = 27                # 层数 (实际 V3.2 是 61)
13    n_dense_layers: int = 1           # Dense 层数 (实际 V3.2 是 3)
14    n_heads: int = 16                 # 注意力头数 (实际 V3.2 是 128)
15
16    # MoE 配置
17    n_routed_experts: int = 64        # 路由专家数 (实际 V3.2 是 256)
18    n_shared_experts: int = 2         # 共享专家数 (实际 V3.2 是 1)
19    n_activated_experts: int = 6      # 激活专家数 (实际 V3.2 是 8)
20    n_expert_groups: int = 1          # 专家分组数 (实际 V3.2 是 8)
21    n_limited_groups: int = 1         # 选择的组数 (实际 V3.2 是 4)
22    score_func: str = "softmax"       # 评分函数 (实际 V3.2 是 sigmoid)
23    route_scale: float = 1.           # 路由缩放 (实际 V3.2 是 2.5)
24
25    # MLA 配置
26    q_lora_rank: int = 0              # Q 低秩维度 (实际 V3.2 是 1536)
27    kv_lora_rank: int = 512           # KV 低秩维度
28    qk_nope_head_dim: int = 128       # 无位置编码的 QK 维度
29    qk_rope_head_dim: int = 64        # 有位置编码的 QK 维度
30    v_head_dim: int = 128             # V 头维度
31
32    # YaRN 配置
33    original_seq_len: int = 4096      # 原始训练长度
34    rope_theta: float = 10000.0       # RoPE 基础频率
35    rope_factor: float = 40           # 长度外推因子
36    beta_fast: int = 32               # 快速衰减参数
37    beta_slow: int = 1                # 慢速衰减参数
38    mscale: float = 1.                # YaRN 额外注意力缩放（影响 softmax_scale）
39
40    # Indexer 配置
41    index_n_heads: int = 64           # Indexer 头数
42    index_head_dim: int = 128         # Indexer 头维度
43    index_topk: int = 2048            # 选择的 token 数

 1class ParallelEmbedding(nn.Module):
 2    def __init__(self, vocab_size: int, dim: int):
 3        super().__init__()
 4        self.vocab_size = vocab_size  # 完整词表大小: 129280
 5        self.dim = dim                # 嵌入维度: 7168
 6
 7        # 检查词表是否可被 GPU 数整除
 8        assert vocab_size % world_size == 0
 9
10        # 计算每个 GPU 负责的词表范围
11        self.part_vocab_size = vocab_size // world_size  # 129280/8=16160
12        self.vocab_start_idx = rank * self.part_vocab_size  # GPU0: 0, GPU1: 16160, ...
13        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
14
15        # 只存储本 GPU 负责的词表部分
16        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))

 1# model.py:13-14 (默认值)
 2world_size = 1
 3rank = 0
 4
 5# model.py:873-875 (运行时修改)
 6class Transformer(nn.Module):
 7    def __init__(self, args: ModelArgs):
 8        global world_size, rank  # 修改全局变量
 9        world_size = dist.get_world_size() if dist.is_initialized() else 1
10        rank = dist.get_rank() if dist.is_initialized() else 0

1GPU 0: 负责词 0 ~ 16159      (相当于 A-C 开头的词)
2GPU 1: 负责词 16160 ~ 32319  (相当于 D-F 开头的词)
3GPU 2: 负责词 32320 ~ 48479  (相当于 G-I 开头的词)
4...
5GPU 7: 负责词 113120 ~ 129279 (相当于 X-Z 开头的词)

1mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)

1x = x - self.vocab_start_idx  # 减去起始位置
2x[mask] = 0                   # 不是我的活就填 0（防止数组越界）

1y = F.embedding(x, self.weight)

1y[mask] = 0  # 不是我负责的位置，结果清零

1dist.all_reduce(y)  # 所有 GPU 的结果相加

1  [词100的向量,  0,            0           ]   ← GPU 0 的贡献
2+ [0,           词20000的向量, 0           ]   ← GPU 1 的贡献
3+ [0,           0,            词50000的向量]   ← GPU 3 的贡献
4+ ...                                          ← 其他 GPU 贡献 0
5─────────────────────────────────────────────────────────────────
6= [词100的向量, 词20000的向量, 词50000的向量]   ← 最终结果

 1def forward(self, x: torch.Tensor) -> torch.Tensor:
 2    if world_size > 1:
 3        # 1. 创建 mask: 标记不属于本 GPU 的 token
 4        mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
 5
 6        # 2. 转换为本地索引 (减去起始偏移)
 7        x = x - self.vocab_start_idx
 8
 9        # 3. 不属于本 GPU 的 token 索引设为 0 (防止越界)
10        x[mask] = 0
11
12    # 4. 查表得到嵌入
13    y = F.embedding(x, self.weight)
14
15    if world_size > 1:
16        # 5. 不属于本 GPU 的结果清零
17        y[mask] = 0
18
19        # 6. 所有 GPU 结果求和 (all_reduce)
20        dist.all_reduce(y)
21
22    return y

1想象你要做一道数学题：计算 [1, 2, 3] × [[a, b, c, d],
2                                      [e, f, g, h],
3                                      [i, j, k, l]]
4
5结果是 [y1, y2, y3, y4]
6
7列并行就是让 2 个 GPU 分工：
8- GPU 0: 只算 [y1, y2]  (对应矩阵的前 2 列)
9- GPU 1: 只算 [y3, y4]  (对应矩阵的后 2 列)

 1想象你要做一道数学题：计算 [1, 2, 3, 4] × [[a],
 2                                          [b],
 3                                          [c],
 4                                          [d]]
 5
 6结果是 y = 1×a + 2×b + 3×c + 4×d
 7
 8行并行就是让 2 个 GPU 分工：
 9- GPU 0: 计算 1×a + 2×b (输入的前半部分 × 权重的前半行)
10- GPU 1: 计算 3×c + 4×d (输入的后半部分 × 权重的后半行)
11- 最后：all_reduce 求和得到 y

1MLP 有 3 个矩阵：W1, W2, W3
2
3设计思路：
4- W1 和 W3 用列并行 → 输出是切片
5- W2 用行并行 → 输入正好是切片，输出是完整结果
6
7这样刚好"接龙"，只需要在 W2 后做一次 all_reduce！

 1def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None,
 2           scale_fmt: Optional[str] = None) -> torch.Tensor:
 3    assert bias is None  # 不支持 bias
 4
 5    if weight.dtype != torch.float8_e4m3fn:
 6        # BF16 权重: 使用标准 F.linear
 7        return F.linear(x, weight)
 8    else:
 9        # FP8 权重: 量化输入后使用 FP8 GEMM
10        x, scale = act_quant(x, block_size, scale_fmt)  # 量化激活
11        return fp8_gemm(x, scale, weight, weight.scale)  # FP8 矩阵乘

 1class Linear(nn.Module):
 2    dtype = torch.bfloat16      # 类级别默认精度
 3    scale_fmt: Optional[str] = None
 4
 5    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
 6        super().__init__()
 7        self.in_features = in_features
 8        self.out_features = out_features
 9
10        # 创建权重参数
11        self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype))
12
13        # FP8 模式下需要额外的 scale 参数
14        if self.weight.element_size() == 1:  # FP8 每个元素 1 字节
15            scale_out = (out_features + block_size - 1) // block_size  # 向上取整
16            scale_in = (in_features + block_size - 1) // block_size
17            # scale 按 128×128 块存储
18            self.weight.scale = self.scale = nn.Parameter(
19                torch.empty(scale_out, scale_in, dtype=torch.float32)
20            )
21        else:
22            self.register_parameter("scale", None)

1class ColumnParallelLinear(Linear):
2    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
3        assert out_features % world_size == 0
4        self.part_out_features = out_features // world_size  # 输出维度切分
5        super().__init__(in_features, self.part_out_features, bias, dtype)
6
7    def forward(self, x: torch.Tensor) -> torch.Tensor:
8        # 直接计算，输出是完整输出的一部分
9        return linear(x, self.weight, self.bias, self.scale_fmt)

 1class RowParallelLinear(Linear):
 2    def __init__(self, in_features: int, out_features: int, bias: bool = False,
 3                 reduce_output=True, dtype=None):
 4        assert in_features % world_size == 0
 5        self.part_in_features = in_features // world_size  # 输入维度切分
 6        self.reduce_output = reduce_output
 7        super().__init__(self.part_in_features, out_features, bias, dtype)
 8
 9    def forward(self, x: torch.Tensor) -> torch.Tensor:
10        y = linear(x, self.weight, None, self.scale_fmt)
11
12        if self.reduce_output and world_size > 1:
13            y = y.float()
14            dist.all_reduce(y)  # 关键: 合并所有 GPU 的结果
15
16        if self.bias is not None:
17            y += self.bias
18
19        return y.type_as(x)

 1RMSNorm 计算过程：
 2
 3Step 1: 计算平方
 4        x² = [9, 16]
 5
 6Step 2: 计算均方 (Mean of Squares)
 7        mean(x²) = (9 + 16) / 2 = 12.5
 8
 9Step 3: 计算均方根 (Root Mean Square)
10        RMS = √12.5 ≈ 3.54
11
12Step 4: 归一化
13        y = x / RMS = [3/3.54, 4/3.54] ≈ [0.85, 1.13]
14
15Step 5: 乘以可学习权重 γ
16        y = y × γ = [0.85×γ₁, 1.13×γ₂]

 1class RMSNorm(nn.Module):
 2    def __init__(self, dim: int, eps: float = 1e-6):
 3        super().__init__()
 4        self.dim = dim
 5        self.eps = eps
 6        # 注意: weight 是 FP32 精度
 7        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
 8
 9    def forward(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
10        dtype = x.dtype
11
12        if residual is None:
13            # 模式1: 普通归一化
14            x = x.float()
15            var = x.pow(2).mean(-1, keepdim=True)  # 计算方差 (实际是均方)
16            x = x * torch.rsqrt(var + self.eps)    # 归一化
17            return (self.weight * x).to(dtype)
18        else:
19            # 模式2: 融合残差的归一化 (优化显存)
20            x = residual = x.float() + residual.float()  # 先加残差
21            var = x.pow(2).mean(-1, keepdim=True)
22            x = x * torch.rsqrt(var + self.eps)
23            return (self.weight * x).to(dtype), residual.to(dtype)

 1class LayerNorm(nn.Module):
 2    def __init__(self, dim: int, eps: float = 1e-6):
 3        super().__init__()
 4        self.dim = dim
 5        self.eps = eps
 6        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
 7        self.bias = nn.Parameter(torch.zeros(dim, dtype=torch.float32))  # 比 RMSNorm 多一个 bias
 8
 9    def forward(self, x: torch.Tensor):
10        return F.layer_norm(x.float(), (self.dim,), self.weight, self.bias, self.eps).type_as(x)

1位置 0:  ↑     位置 1:  ↗     位置 2:  →     位置 3:  ↘
2         |              /              —              \

1复数 z = a + bi 可以表示为 (长度, 角度)
2复数乘法 z1 × z2 = (长度1×长度2, 角度1+角度2)

 1Step 1: 把向量看作复数
 2        x = 1 + 0i
 3
 4Step 2: 计算旋转角度
 5        角度 = m × θ = 3 × 30° = 90°
 6
 7Step 3: 计算旋转因子
 8        e^(i×90°) = cos(90°) + i×sin(90°) = 0 + 1i
 9
10Step 4: 复数乘法
11        x' = (1 + 0i) × (0 + 1i) = 0 + 1i
12
13Step 5: 转回向量
14        x' = [0, 1]  （从指向右变成指向上，旋转了 90°）

 1def precompute_freqs_cis(args: ModelArgs) -> torch.Tensor:
 2    dim = args.qk_rope_head_dim  # 64
 3    seqlen = args.max_seq_len    # 163840
 4    base = args.rope_theta       # 10000
 5    factor = args.rope_factor    # 40
 6
 7    # 计算基础频率: θ_i = 1 / (10000^(2i/64))
 8    # i ∈ [0, 2, 4, ..., 62], 共 32 个
 9    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
10    # freqs.shape = [32], 值从 1.0 衰减到 约 0.01
11
12    # 如果需要外推 (163840 > 4096)
13    if seqlen > args.original_seq_len:
14        # YaRN 调整 (详见 3.3 节)
15        low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len)
16        smooth = 1 - linear_ramp_factor(low, high, dim // 2)
17        freqs = freqs / factor * (1 - smooth) + freqs * smooth
18
19    # 计算每个位置的角度
20    t = torch.arange(seqlen)  # [0, 1, 2, ..., 163839]
21    freqs = torch.outer(t, freqs)  # [163840, 32] - 外积
22
23    # 转换为复数表示: e^(iθ) = cos(θ) + i*sin(θ)
24    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
25    # freqs_cis.shape = [163840, 32], dtype = complex64
26
27    return freqs_cis

 1def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, interleaved: bool = True) -> torch.Tensor:
 2    dtype = x.dtype
 3    shape = x.shape  # [B, S, H, D]
 4
 5    if not interleaved:
 6        # 非交错模式: [x0,x1,x2,...,x31,x32,...,x63] → [x0,x32,x1,x33,...,x31,x63]
 7        x = x.view(*shape[:-1], 2, -1).transpose(-1, -2).contiguous()
 8
 9    # 将实数对看作复数: [B,S,H,D] → [B,S,H,D/2] complex
10    x = torch.view_as_complex(x.float().view(*shape[:-1], -1, 2))
11
12    # 广播 freqs_cis 到 [1, S, 1, D/2]
13    freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
14
15    # 复数乘法实现旋转
16    y = torch.view_as_real(x * freqs_cis).flatten(3)  # [B,S,H,D]
17
18    if not interleaved:
19        # 恢复非交错布局
20        y = torch.cat([y[..., 0::2], y[..., 1::2]], dim=-1)
21
22    return y.to(dtype)

 1def find_correction_dim(num_rotations, dim, base, max_seq_len):
 2    """计算给定旋转次数对应的维度索引"""
 3    return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
 4
 5def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
 6    """计算需要调整的维度范围"""
 7    low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
 8    high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
 9    return max(low, 0), min(high, dim-1)
10
11def linear_ramp_factor(min, max, dim):
12    """生成平滑过渡的线性斜坡"""
13    linear_func = (torch.arange(dim) - min) / (max - min)
14    return torch.clamp(linear_func, 0, 1)

1if args.max_seq_len > args.original_seq_len:
2    mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
3    self.softmax_scale = self.softmax_scale * mscale * mscale

1传统 KV Cache 大小计算：
2- 序列长度: 163840 tokens
3- K 维度: 128 头 × 192 维 = 24576
4- V 维度: 128 头 × 128 维 = 16384
5- 每层缓存: 163840 × (24576 + 16384) × 2B ≈ 13.4 GB
6- 61 层总共: 13.4 × 61 ≈ 817 GB  😱

1传统：x → K (24576维) → 缓存 K
2MLA： x → c_KV (512维) → 缓存 c_KV → 展开成 K (24576维)
3
4压缩率: 512 / 24576 ≈ 2%！

1MLA KV Cache 大小：
2- 序列长度: 163840 tokens
3- 压缩后维度: 512 + 64 (RoPE) = 576
4- 每层缓存: 163840 × 576 × 2B ≈ 188 MB
5- 61 层总共: 188 × 61 ≈ 11.5 GB  ✅
6
7节省: 817 GB → 11.5 GB，减少 98.6%！

 1class MLA(nn.Module):
 2    def __init__(self, args: ModelArgs):
 3        super().__init__()
 4        self.dim = args.dim                          # 7168
 5        self.n_heads = args.n_heads                  # 128
 6        self.n_local_heads = args.n_heads // world_size  # 128/8=16 (每 GPU)
 7        self.q_lora_rank = args.q_lora_rank          # 1536
 8        self.kv_lora_rank = args.kv_lora_rank        # 512
 9        self.qk_nope_head_dim = args.qk_nope_head_dim  # 128
10        self.qk_rope_head_dim = args.qk_rope_head_dim  # 64
11        self.qk_head_dim = 128 + 64                  # 192
12        self.v_head_dim = args.v_head_dim            # 128
13
14        # === Query 路径 ===
15        self.wq_a = Linear(7168, 1536)               # 压缩
16        self.q_norm = RMSNorm(1536)
17        self.wq_b = ColumnParallelLinear(1536, 128*192)  # 展开 (列并行)
18
19        # === KV 路径 ===
20        self.wkv_a = Linear(7168, 512+64)            # 压缩 (512 for KV, 64 for RoPE)
21        self.kv_norm = RMSNorm(512)
22        self.wkv_b = ColumnParallelLinear(512, 128*(128+128))  # 展开 (k_nope + v)
23
24        # === 输出 ===
25        self.wo = RowParallelLinear(128*128, 7168)   # 行并行
26
27        # === Indexer ===
28        self.indexer = Indexer(args)
29
30        # === KV Cache ===
31        self.register_buffer("kv_cache", torch.zeros(8, 163840, 512))  # 只缓存压缩的 c_KV
32        self.register_buffer("pe_cache", torch.zeros(8, 163840, 64))   # 缓存 k_pe

 1# 在 forward() 中
 2qr = self.q_norm(self.wq_a(x))  # [B,S,1536] - 压缩 + 归一化
 3q = self.wq_b(qr)               # [B,S,16*192] - 展开 (列并行后每 GPU 16 头)
 4q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)  # [B,S,16,192]
 5
 6# 分离 nope 和 rope 部分
 7q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
 8# q_nope: [B,S,16,128], q_pe: [B,S,16,64]
 9
10q_pe = apply_rotary_emb(q_pe, freqs_cis)  # 应用位置编码

 1# 在 forward() 中
 2kv = self.wkv_a(x)  # [B,S,576] - 压缩
 3kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
 4# kv: [B,S,512], k_pe: [B,S,64]
 5
 6kv = self.kv_norm(kv)  # 归一化
 7k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)  # 位置编码
 8
 9# FP8 量化模拟 (为了与部署一致)
10kv_fp8, kv_scale = act_quant(kv, block_size, self.scale_fmt)
11kv = (kv_fp8.view(-1, block_size).float() * kv_scale.view(-1, 1)).to(kv.dtype).view_as(kv)
12
13# 存入缓存
14self.kv_cache[:bsz, start_pos:end_pos] = kv      # 只存 512 维!
15self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)

 1if mask is not None:  # Prefill
 2    q = torch.cat([q_nope, q_pe], dim=-1)  # [B,S,16,192]
 3
 4    # 展开 KV
 5    kv = self.wkv_b(kv)  # [B,S,16*256]
 6    kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
 7    k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
 8    k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)
 9
10    # 注意力计算
11    scores = torch.einsum("bshd,bthd->bsht", q, k).mul_(self.softmax_scale)
12
13    # Indexer 稀疏化
14    topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
15    index_mask = torch.full((bsz, seqlen, seqlen), float("-inf"), device=x.device)
16    index_mask = index_mask.scatter_(-1, topk_indices, 0)  # 只保留 top-k
17    scores += index_mask.unsqueeze(2) + mask
18
19    scores = scores.softmax(dim=-1)
20    x = torch.einsum("bsht,bthd->bshd", scores, v)

 1else:  # Decode
 2    # 延迟反量化 wkv_b (只做一次)
 3    if self.dequant_wkv_b is None and self.wkv_b.scale is not None:
 4        self.dequant_wkv_b = weight_dequant(self.wkv_b.weight, self.wkv_b.scale)
 5    wkv_b = self.wkv_b.weight if self.dequant_wkv_b is None else self.dequant_wkv_b
 6    wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)  # [16, 256, 512]
 7
 8    # 关键优化: Q_nope 先与 W_kv_b 的 K 部分结合
 9    q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
10    # [B,1,16,128] @ [16,128,512] → [B,1,16,512]
11
12    # 与缓存计算注意力分数
13    scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) +
14              torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale
15
16    # Indexer 稀疏化
17    topk_indices = self.indexer(x, qr, start_pos, freqs_cis, mask)
18    index_mask = torch.full((bsz, 1, end_pos), float("-inf"), device=x.device)
19    index_mask = index_mask.scatter_(-1, topk_indices, 0)
20    scores += index_mask.unsqueeze(2)
21
22    scores = scores.softmax(dim=-1)
23
24    # 计算输出
25    x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
26    x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])

1传统注意力计算量：
2- 每个新 token 要和 163840 个历史 token 计算相似度
3- 计算量 = 163840 × 头数 × 维度 ≈ 非常大！

1对于当前 token 的 query q 和历史 token 的 key k：
2
3index_score = Σ (ReLU(q·k) × weight)
4
5- q·k: 相似度分数
6- ReLU: 只保留正分数（负相关的直接忽略）
7- weight: 每个注意力头的权重
8- Σ: 对所有头求和

1原始向量: [1, 0, 0, 0, ...]   ← 很多 0，不利于量化
2变换后:   [0.1, 0.1, 0.1, 0.1, ...] ← 值更均匀，量化更准

 1class Indexer(torch.nn.Module):
 2    def __init__(self, args: ModelArgs):
 3        super().__init__()
 4        self.dim = args.dim                    # 7168
 5        self.n_heads = args.index_n_heads      # 64
 6        self.n_local_heads = 64 // world_size  # 8 (每 GPU)
 7        self.head_dim = args.index_head_dim    # 128
 8        self.rope_head_dim = args.qk_rope_head_dim  # 64
 9        self.index_topk = args.index_topk      # 2048
10        self.q_lora_rank = args.q_lora_rank    # 1536
11
12        # Indexer 自己的投影层
13        self.wq_b = Linear(1536, 64*128)       # Q 投影 (使用 MLA 的 qr)
14        self.wk = Linear(7168, 128)            # K 投影 (从原始 x)
15        self.k_norm = LayerNorm(128)           # K 归一化
16        self.weights_proj = Linear(7168, 64, dtype=torch.float32)  # 权重投影
17
18        # Indexer 的 K Cache (FP8)
19        self.register_buffer("k_cache", torch.zeros(8, 163840, 128, dtype=torch.float8_e4m3fn))
20        self.register_buffer("k_scale_cache", torch.zeros(8, 163840, 1))  # 128/128=1 个 scale

 1def forward(self, x, qr, start_pos, freqs_cis, mask):
 2    bsz, seqlen, _ = x.size()
 3    end_pos = start_pos + seqlen
 4
 5    # === Q 计算 ===
 6    q = self.wq_b(qr)  # [B,S,64*128]
 7    q = q.view(bsz, seqlen, self.n_heads, self.head_dim)  # [B,S,64,128]
 8    q_pe, q_nope = torch.split(q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
 9    # 注意: Indexer 的 RoPE 是非交错的
10    q_pe = apply_rotary_emb(q_pe, freqs_cis, False)
11    q = torch.cat([q_pe, q_nope], dim=-1)
12
13    # === K 计算 ===
14    k = self.wk(x)     # [B,S,128]
15    k = self.k_norm(k)  # LayerNorm
16    k_pe, k_nope = torch.split(k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1)
17    k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, False).squeeze(2)
18    k = torch.cat([k_pe, k_nope], dim=-1)
19
20    # === Hadamard 变换 ===
21    q = rotate_activation(q)  # 随机正交变换
22    k = rotate_activation(k)
23
24    # === FP8 量化 ===
25    q_fp8, q_scale = act_quant(q, block_size, self.scale_fmt)
26    k_fp8, k_scale = act_quant(k, block_size, self.scale_fmt)
27
28    # === 更新 K Cache ===
29    self.k_cache[:bsz, start_pos:end_pos] = k_fp8
30    self.k_scale_cache[:bsz, start_pos:end_pos] = k_scale
31
32    # === 计算 Index Score ===
33    weights = self.weights_proj(x.float()) * self.n_heads ** -0.5
34    weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
35
36    # FP8 矩阵乘 + ReLU + 求和
37    index_score = fp8_index(q_fp8.contiguous(), weights,
38                            self.k_cache[:bsz, :end_pos].contiguous(),
39                            self.k_scale_cache[:bsz, :end_pos].contiguous())
40
41    if mask is not None:
42        index_score += mask  # 加上因果 mask
43
44    # === Top-K 选择 ===
45    topk_indices = index_score.topk(min(self.index_topk, end_pos), dim=-1)[1]
46
47    # 同步验证 (确保所有 GPU 选择相同的 token)
48    topk_indices_ = topk_indices.clone()
49    dist.broadcast(topk_indices_, src=0)
50    assert torch.all(topk_indices == topk_indices_)
51
52    return topk_indices

1def rotate_activation(x: torch.Tensor) -> torch.Tensor:
2    assert x.dtype == torch.bfloat16
3    from fast_hadamard_transform import hadamard_transform
4    hidden_size = x.size(-1)
5    return hadamard_transform(x, scale=hidden_size ** -0.5)

 1class MLP(nn.Module):
 2    def __init__(self, dim: int, inter_dim: int, reduce_output: bool = True):
 3        super().__init__()
 4        self.w1 = ColumnParallelLinear(dim, inter_dim)  # gate_proj
 5        self.w2 = RowParallelLinear(inter_dim, dim, reduce_output=reduce_output)  # down_proj
 6        self.w3 = ColumnParallelLinear(dim, inter_dim)  # up_proj
 7
 8    def forward(self, x: torch.Tensor) -> torch.Tensor:
 9        # SwiGLU: SiLU(W1·x) × W3·x
10        return self.w2(
11            (F.silu(self.w1(x).float()) * self.w3(x).float()).type_as(x)
12        )

1256 个路由专家：每个专家擅长某类问题
2  - 专家 0-31: 可能擅长数学
3  - 专家 32-63: 可能擅长代码
4  - 专家 64-95: 可能擅长语言理解
5  - ... (自动学习的，无需人工指定)
6
71 个共享专家：处理通用信息，所有 token 都会经过

 1Step 1: 计算 256 个专家的分数
 2        专家 5 (数学): 0.95
 3        专家 40 (代码): 0.92
 4        专家 67 (逻辑): 0.88
 5        专家 102 (推理): 0.85
 6        ... (其他专家分数较低)
 7
 8Step 2: 分组计算组得分 (每组取 top-2 求和)
 9        组1 (0-31):   专家5=0.95 + 专家12=0.45 = 1.40
10        组2 (32-63):  专家40=0.92 + 专家55=0.50 = 1.42
11        组3 (64-95):  专家67=0.88 + 专家80=0.60 = 1.48
12        组4 (96-127): 专家102=0.85 + 专家110=0.55 = 1.40
13        ... (其他组得分较低)
14
15Step 3: 选出 top-4 组
16        ✓ 组3 (1.48)
17        ✓ 组2 (1.42)
18        ✓ 组1 (1.40)
19        ✓ 组4 (1.40)
20
21Step 4: 从这 4 组的 128 个专家中选 top-8
22        ✓ 专家 5, 40, 67, 102, 12, 55, 80, 110
23
24Step 5: 归一化权重并缩放 ×2.5
25        权重 = [0.18, 0.17, 0.16, 0.15, 0.10, 0.09, 0.08, 0.07]

1参数量计算：
2- 如果用一个大 MLP: 7168 × 18432 × 3 ≈ 4 亿参数
3- 实际用 MoE:
4  - 256 个小专家: 256 × (7168 × 2048 × 3) ≈ 113 亿参数
5  - 每次只激活 8 个: 8 × ... ≈ 3.5 亿参数激活
6
7结果：参数量增加 28 倍，但计算量几乎不变！

 1class Gate(nn.Module):
 2    def __init__(self, args: ModelArgs):
 3        super().__init__()
 4        self.dim = args.dim                      # 7168
 5        self.topk = args.n_activated_experts     # 8
 6        self.n_groups = args.n_expert_groups     # 8
 7        self.topk_groups = args.n_limited_groups # 4
 8        self.score_func = args.score_func        # "sigmoid"
 9        self.route_scale = args.route_scale      # 2.5
10
11        self.weight = nn.Parameter(torch.empty(256, 7168))  # 路由权重
12        # bias 只在 dim=7168 时启用 (即完整模型)
13        self.bias = nn.Parameter(torch.empty(256, dtype=torch.float32)) if self.dim == 7168 else None
14
15    def forward(self, x):
16        # 1. 计算原始分数
17        scores = linear(x.float(), self.weight.float())  # [B*S, 256]
18
19        # 2. Sigmoid 激活
20        scores = scores.sigmoid()
21        original_scores = scores  # 保存用于最终权重
22
23        # 3. 添加 bias (用于调整专家偏好)
24        if self.bias is not None:
25            scores = scores + self.bias
26
27        # 4. 分组路由 (256 专家分成 8 组，每组 32 个)
28        if self.n_groups > 1:
29            scores = scores.view(x.size(0), self.n_groups, -1)  # [B*S, 8, 32]
30
31            # 每组取 top-2 专家分数之和作为组分数
32            group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)  # [B*S, 8]
33
34            # 选择 top-4 组
35            indices = group_scores.topk(self.topk_groups, dim=-1)[1]  # [B*S, 4]
36
37            # mask 掉未选中的组
38            mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool)
39            mask = mask.scatter_(1, indices, False)
40            scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf"))
41            scores = scores.flatten(1)  # [B*S, 256]
42
43        # 5. 从选中的组中选 top-8 专家
44        indices = scores.topk(self.topk, dim=-1)[1]  # [B*S, 8]
45
46        # 6. 从原始分数中获取权重 (不是 mask 后的分数)
47        weights = original_scores.gather(1, indices)  # [B*S, 8]
48
49        # 7. 归一化 (sigmoid 模式下)
50        weights /= weights.sum(dim=-1, keepdim=True)
51
52        # 8. 缩放
53        weights *= self.route_scale  # × 2.5
54
55        return weights, indices

1class Expert(nn.Module):
2    def __init__(self, dim: int, inter_dim: int):
3        super().__init__()
4        self.w1 = Linear(dim, inter_dim)   # 7168 → 2048
5        self.w2 = Linear(inter_dim, dim)   # 2048 → 7168
6        self.w3 = Linear(dim, inter_dim)   # 7168 → 2048
7
8    def forward(self, x):
9        return self.w2((F.silu(self.w1(x).float()) * self.w3(x).float()).type_as(x))

 1class MoE(nn.Module):
 2    def __init__(self, args: ModelArgs):
 3        super().__init__()
 4        self.dim = args.dim
 5        self.n_routed_experts = args.n_routed_experts        # 256
 6        self.n_local_experts = args.n_routed_experts // world_size  # 32 (每 GPU)
 7        self.n_activated_experts = args.n_activated_experts  # 8
 8
 9        # 当前 GPU 负责的专家范围
10        self.experts_start_idx = rank * self.n_local_experts
11        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
12
13        self.gate = Gate(args)
14
15        # 只创建本 GPU 负责的专家
16        self.experts = nn.ModuleList([
17            Expert(args.dim, args.moe_inter_dim)
18            if self.experts_start_idx <= i < self.experts_end_idx
19            else None
20            for i in range(self.n_routed_experts)
21        ])
22
23        # 共享专家 (所有 token 都经过)
24        self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim,
25                                  reduce_output=False)
26
27    def forward(self, x):
28        shape = x.size()
29        x = x.view(-1, self.dim)  # [B*S, 7168]
30
31        # 路由
32        weights, indices = self.gate(x)  # [B*S, 8], [B*S, 8]
33
34        y = torch.zeros_like(x, dtype=torch.float32)
35
36        # 统计每个专家被选中的次数
37        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
38
39        # 只计算本 GPU 负责的专家
40        for i in range(self.experts_start_idx, self.experts_end_idx):
41            if counts[i] == 0:
42                continue
43            expert = self.experts[i]
44            # 找到选中专家 i 的 (token_idx, top_idx)
45            idx, top = torch.where(indices == i)
46            # 计算并加权累加
47            y[idx] += expert(x[idx]) * weights[idx, top, None]
48
49        # 加上共享专家输出
50        y += self.shared_experts(x)
51
52        # 跨 GPU 合并
53        if world_size > 1:
54            dist.all_reduce(y)
55
56        return y.type_as(x).view(shape)

1传统方式：x → Block → y
2残差方式：x → Block → y + x  (输出 = 处理结果 + 原始输入)

1# 传统方式 (需要存储中间结果)
2x = x + attention(norm(x))   # 第一次加法
3x = x + ffn(norm(x))         # 第二次加法
4
5# 融合方式 (节省显存)
6x, residual = norm(x, residual)  # 归一化同时处理残差
7x = attention(x)
8x, residual = norm(x, residual)  # 归一化同时处理残差
9x = ffn(x)

 1class Block(nn.Module):
 2    def __init__(self, layer_id: int, args: ModelArgs):
 3        super().__init__()
 4        self.attn = MLA(args)
 5        # 前 3 层用 Dense MLP，之后用 MoE
 6        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)
 7        self.attn_norm = RMSNorm(args.dim)
 8        self.ffn_norm = RMSNorm(args.dim)
 9
10    def forward(self, x, residual, start_pos, freqs_cis, mask):
11        # 注意力前归一化 + 残差处理
12        if residual is None:
13            x, residual = self.attn_norm(x), x
14        else:
15            x, residual = self.attn_norm(x, residual)
16
17        # 注意力
18        x = self.attn(x, start_pos, freqs_cis, mask)
19
20        # FFN 前归一化 + 残差处理
21        x, residual = self.ffn_norm(x, residual)
22
23        # FFN
24        x = self.ffn(x)
25
26        return x, residual

 1class Transformer(nn.Module):
 2    def __init__(self, args: ModelArgs):
 3        global world_size, rank
 4        world_size = dist.get_world_size() if dist.is_initialized() else 1
 5        rank = dist.get_rank() if dist.is_initialized() else 0
 6
 7        # 设置全局精度
 8        Linear.dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16
 9        Linear.scale_fmt = args.scale_fmt
10
11        super().__init__()
12        self.max_seq_len = args.max_seq_len
13        self.embed = ParallelEmbedding(args.vocab_size, args.dim)
14
15        # 61 层 Block
16        self.layers = torch.nn.ModuleList()
17        for layer_id in range(args.n_layers):
18            self.layers.append(Block(layer_id, args))
19
20        self.norm = RMSNorm(args.dim)
21        # lm_head 用 FP32 以获得更精确的 logits
22        self.head = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.float32)
23
24        # 预计算位置编码
25        self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False)
26
27    @torch.inference_mode()
28    def forward(self, tokens: torch.Tensor, start_pos: int = 0):
29        seqlen = tokens.size(1)
30
31        # 获取当前位置的频率
32        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
33
34        # 因果 mask (只在 prefill 时使用)
35        mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device).triu_(1) if seqlen > 1 else None
36
37        # 嵌入
38        h, residual = self.embed(tokens), None
39
40        # 61 层 Block
41        for layer in self.layers:
42            h, residual = layer(h, residual, start_pos, freqs_cis, mask)
43
44        # 最终归一化
45        h, _ = self.norm(h, residual)
46
47        # 输出投影 (只取最后一个 token)
48        logits = self.head(h[:, -1].float())
49
50        # 合并所有 GPU 的 logits
51        if world_size > 1:
52            all_logits = [torch.empty_like(logits) for _ in range(world_size)]
53            dist.all_gather(all_logits, logits)
54            logits = torch.cat(all_logits, dim=-1)
55
56        return logits

1def sample(logits, temperature: float = 1.0):
2    logits = logits / max(temperature, 1e-5)
3    probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
4    # Gumbel-Max trick: 等价于 multinomial 但更高效
5    return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)

1cd inference
2python convert.py \
3    --hf-ckpt-path /path/to/hf/checkpoint \
4    --save-path /path/to/converted \
5    --n-experts 256 \
6    --model-parallel 8

1cd inference
2python generate.py \
3    --ckpt-path /path/to/converted \
4    --config config_671B_v3.2.json \
5    --interactive \
6    --temperature 0.6

1cd inference
2torchrun --nproc_per_node=8 generate.py \
3    --ckpt-path /path/to/converted \
4    --config config_671B_v3.2.json \
5    --interactive \
6    --temperature 0.6