Transformer 骨架

这一页回答什么

Phoenix 排序模型与召回用户塔共用同一套生产版 transformer 架构。本页解析该架构的内部结构：网络单层构成、注意力计算流程、位置编码与归一化方法、关键配置开关，以及该架构与早期 Grok-1 移植版本的差异。

页面路径沿用 grok-transformer，内容已全面更新为生产代码。早期开源版 phoenix/grok.py（移植自 Grok-1 的演示 transformer）已于 2026-08-13 移除；官方文档说明指出“本次发布即为生产实现本身”（phoenix/README.md:9-21）。当前骨架代码位于 phoenix/xrex/models/transformer.py，注意力算子位于 phoenix/xrex/models/layers.py、attention.py 与 recsys_attention.py。

核心结论

1. 单一骨架服务双模型：排序模型 RecsysAggregatedModel 与召回模型 RecsysTwoTowerModel 均通过 TransformerConfig.make() 构建 Transformer 实例（phoenix/xrex/models/transformer.py:72-158、:682）。
2. 非因果注意力机制：两个模型均将注意力配置设置为 causal=False；Token 之间的可见性完全由段 ID 决定（phoenix/xrex/configs/xrecsys.py:1038、phoenix/xrex/configs/xrecsys_two_tower.py:908）。详见 候选隔离注意力掩码。
3. 标准两层 GELU 前馈网络取代门控结构：推荐系统配置硬编码 gated_mlp=False 与 widening_factor=2（xrecsys.py:1043-1046、xrecsys_two_tower.py:913-916）。早期文档描述的“GLU 门控 FFN”已废弃。当前架构未包含 MoE 专家路由组件。
4. 分组查询注意力（GQA）：排序旗舰配置分配 20 个查询头与 4 个 KV 头，召回旗舰配置分配 16 个查询头与 4 个 KV 头，两者的 key_size 均为 128（xrecsys.py:323-324、:1026-1028）。
5. 右锚定 RoPE 位置编码：历史互动序列的最新记录固定锚定于相同位置索引，所有候选帖子统一分配“历史序列末尾”位置编码（phoenix/xrex/models/recsys_model.py:260-281）。
6. 注意力 Logit 封顶机制变更：默认采用 soft_sign 方式，截断上限为 80，即 x / (1 + |x|/80)；旗舰排序与 combined 召回配置关闭封顶，改为采用 QK 归一化（phoenix/xrex/models/attention.py:229-237、xrecsys.py:339-340）。早期文档描述的 30·tanh(x/30) 已废弃。
7. 双重归一化支持：网络在每个子层前应用 RMSNorm；当 primer_norm=True 时，在子层输出加回残差流之前追加一次 RMSNorm。旗舰排序配置关闭后置归一化（primer_norm=False），召回配置默认开启（transformer.py:495-547、:578-582，xrecsys.py:341、xrecsys_two_tower.py:921）。
8. 末尾层支持候选专用更新：通过 num_candidate_only_layers 参数配置，旗舰排序将其设为 2（transformer.py:86、:789-945）。

整体结构

flowchart TB
    IN["输入 token B × T × D"] --> STK["层栈 共 num_layers 层"]
    STK --> FN["final RMSNorm"]
    FN --> OUT["输出 B × T × D"]
    subgraph DL["每个 DecoderLayer"]
        H0["残差流 h"] --> N1["RMSNorm"]
        N1 --> ATT["多头注意力<br/>GQA + RoPE + QK 归一化"]
        ATT --> PN1["可选 RMSNorm<br/>primer_norm"]
        PN1 --> A1["h = h + attn"]
        A1 --> N2["RMSNorm"]
        N2 --> FFN["前馈 up → GELU → down"]
        FFN --> PN2["可选 RMSNorm<br/>primer_norm"]
        PN2 --> A2["h = h + ffn"]
    end

最终的 final_layer_norm 由调用方排序模型在 transformer 输出后执行（recsys_model.py:4369-4373）。

TransformerConfig：配置入口

TransformerConfig（transformer.py:72-103）定义了骨架的核心参数：

出处：xrecsys.py:532-577、:1024-1062。

注意力

配置

AttentionConfig（attention.py:17-38）与推荐专用配置 RecsysAttentionConfig（phoenix/xrex/models/recsys_attention.py:11-15，扩展包含 history_seq_len、candidate_seq_len 与 num_user_prefix_tokens）。排序模型采用如下初始化配置（xrecsys.py:1025-1042）：

RecsysAttentionConfig(
    key_size=128,
    num_q_heads=mparams["query_heads"], num_kv_heads=mparams["kv_heads"],
    attn_logit_cap=mparams.get("attn_logit_cap", 80.0),
    attn_logit_cap_method=mparams.get("attn_logit_cap_method", "soft_sign"),
    qk_norm=mparams.get("qk_norm", False),
    qkv_merge=True, rotary=True, causal=False, ...
)

计算步骤

MultiHeadAttention（layers.py:102-577）按如下步骤执行计算：

1. Q/K/V 合并投影：当 qkv_merge=True 时，算子执行单次矩阵乘法生成 Q、K、V 张量，随后按注意力头切分（layers.py:350-381）。
2. QK 归一化与 RoPE：当 qk_norm=True 时，算子对 Q 和 K 分别执行带可学习增益的 RMSNorm，随后施加 RoPE；在支持的硬件上两项计算融合为单算子执行（layers.py:394-439、:728-770）。
3. RoPE 位置编码：采用标准一维旋转位置编码，底数设为 rope_base_exponent=10000，位置取自调用方传入的 positions[..., 0]（layers.py:45-61、:382-393，attention.py:31-32）。
4. Logit 缩放与封顶：Logit 乘以系数 sqrt(128) / key_size（当 key_size=128 时等价于 1/sqrt(128)），随后按配置执行截断（phoenix/xrex/models/scaling.py:46-47、attention.py:229-237）：
   • tanh：计算 cap · tanh(x / cap)；
   • soft_sign：计算 x / (1 + |x| / cap)（推荐配置默认模式，cap=80）；
   • cap <= 0：关闭封顶截断。
5. 段掩码、Softmax 与加权求和：算子依据段 ID 屏蔽无效位置，执行 softmax 归一化后计算加权求和，最终执行线性输出投影（layers.py:511-545）。

各配置的封顶与归一化：

注意力实现

配置项 attn_impl 负责选取具体计算内核（layers.py:440-510）：可选内核包括 jax_attn（纯 JAX 参考实现）、Pallas 编写的推荐专用算子（pallas_ranker_attn / pallas_ranker_attn_infer / pallas_ranker_varlen_attn），以及 CuTe DSL 编写的 FlashAttention-4 风格算子（cutedsl_ranker_attn / cutedsl_ranker_varlen_attn，代码基于 Dao-AILab 的 quack 项目衍生开发，见 phoenix/THIRD_PARTY_NOTICES.md）。其中 varlen 版本用于支持训练阶段的序列打包机制。旗舰排序配置在 H100 GPU 上采用 cutedsl_ranker_varlen_attn（xrecsys.py:395）。

位置编码：右锚定 RoPE

右锚定位置编码通过对齐历史末尾消除序列长度差异产生的位置偏移。函数 right_anchored_rope_positions（recsys_model.py:260-281）按如下规则分配位置索引：

• 用户前缀位置顺次分配 0, 1, …。
• 历史序列采用右对齐：无论历史实际长度如何，最后一条历史记录始终对齐至 history_end - 1。当用户历史较短时，空缺位置位于左侧。
• 所有候选帖子的位置索引统一设为 history_end。

该机制确保最新历史记录与候选帖子之间的相对距离保持恒定，避免模型针对不同历史长度学习不同的相对位置分布。旗舰排序与全部召回配置均启用该机制（xrecsys.py:333、xrecsys_two_tower.py:473）。

前馈层

前馈网络采用非门控两层 GELU 结构。前馈模块位于 DenseBlock（transformer.py:248-388）。推荐配置设置 gated_mlp=False，执行非门控计算路径（:350-381）：

h_up = Linear(ffn_size)(inputs)           # up_proj
h_up = gelu(h_up, approximate=True)
h_dense = Linear(model_size)(h_up)        # down_proj

前馈中间维度计算公式为 ffn_size = widening_factor × emb_size，计算结果向上对齐至 32 的整数倍（phoenix/xrex/utils/utils.py:35-50）。在排序旗舰配置中，该维度为 2 × 2560 = 5120。代码库仍包含 gated_mlp=True 时的门控分支代码（计算 up × gelu(gate)，transformer.py:289-348），推荐配置不激活此分支。

训练稳定性与规模化

• μP 风格参数缩放：ScaleConfig 以 512 为基准隐藏宽度，依据实际宽度按规则缩放各层学习率。例如隐藏层投影缩放系数为 512 / fan_in，输入层乘以 sqrt(emb_size)，输出层乘以 sqrt(512 / emb_size) × 2（scaling.py:10-47）。官方文档说明指出 emb_size=512 为基准宽度，在 nano 配置下主干网络随宽度变化的倍率恰好为 1（phoenix/README.md:419-423）。
• 激活 L2 范数监控：网络逐层统计注意力与前馈输出的均方值；最末层输出的均方值计为 act-l2-loss，召回损失函数通过权重 act_l2_weight 引入该惩罚项（transformer.py:519-526、:984-991，phoenix/xrex/models/recsys_two_tower_model.py:2809-2811）。
• FP8 计算与重算机制：参数 fp8_gemms 启用 FP8 精度执行矩阵乘法；remat_policy 控制反向传播阶段需要保留与重算的中间激活（transformer.py:90-94）。

与早期 Grok-1 移植版的差异

设计决策

源码锚点

• phoenix/xrex/models/transformer.py:72-158 —— TransformerConfig
• phoenix/xrex/models/transformer.py:444-618 —— DecoderLayer
• phoenix/xrex/models/transformer.py:248-388 —— DenseBlock 前馈
• phoenix/xrex/models/transformer.py:682-991 —— Transformer.__call__ 与候选专用层
• phoenix/xrex/models/layers.py:102-577 —— MultiHeadAttention
• phoenix/xrex/models/attention.py:17-38、:229-237 —— 注意力配置与 logit 封顶
• phoenix/xrex/models/recsys_model.py:260-281 —— 右锚定 RoPE 位置
• phoenix/xrex/models/scaling.py:10-47 —— μP 缩放

相关页面

• Phoenix 排序模型 —— 使用该骨架的排序模型
• Phoenix 召回：双塔检索架构 —— 使用该骨架的召回用户塔
• 候选隔离注意力掩码 —— 段 ID 掩码与候选专用层
• Phoenix 训练 —— 优化器与训练步