---
title: "Transformer 骨架"
description: "Phoenix 排序与召回共用的生产版 transformer：双向注意力、GQA、RoPE、QK 归一化、GELU 前馈、无 MoE，取代早期的 Grok-1 移植版"
---

> Documentation Index
> Fetch the complete documentation index at: https://x-algorithm.vergego.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# 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`）。详见 [候选隔离注意力掩码](https://x-algorithm.vergego.ai/concepts/candidate-isolation-masking/)。
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`）。

## 整体结构

```mermaid
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`）定义了骨架的核心参数：

| 字段 | 含义 | 排序旗舰取值（开源默认值） |
|------|------|------|
| `emb_size` | 隐藏维 D | 2560 |
| `num_layers` | 层数 | 8 |
| `attn_config` | 注意力配置，见下 | `RecsysAttentionConfig` |
| `ffn_config` | 前馈配置 | `widening_factor=2`、`gated_mlp=False` |
| `pre_norm` / `primer_norm` | 子层前 / 后 RMSNorm | `True` / `False` |
| `use_layer_stack` | 用 `jax.lax.scan` 堆叠层 | `True` |
| `num_candidate_only_layers` | 末尾只更新候选的层数 | 2 |
| `fp8_gemms` | 矩阵乘用 FP8 | `Fp8Gemms.ALL` |
| `remat_policy` | 反向时重算策略 | `SAVE_H100_RECSYS` |

出处：`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`）：

```python
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`）。

各配置的封顶与归一化：

| 配置 | `qk_norm` | `attn_logit_cap` | 出处 |
|------|-----------|------------------|------|
| 排序 `home_direct_packed` | `True` | -1（关） | `xrecsys.py:339-340` |
| 召回 `xrecsys_two_tower`（旗舰） | `False`（默认） | 80，`soft_sign` | `xrecsys_two_tower.py:439-513`、`:900-906` |
| 召回 `xrecsys_two_tower_combined` | `True` | -1（关） | `xrecsys_two_tower.py:274-275` |
| 召回 `xrecsys_two_tower_immersive_only` | `False` | 80 | `xrecsys_two_tower.py:378-379` |

### 注意力实现

配置项 `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`）：

```python
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 移植版的差异

| 项 | 早期 `grok.py`（已删除） | 现在 `xrex/models/transformer.py` |
|----|------|------|
| 掩码 | 因果下三角 + 候选块清零 | 非因果，按段 ID 判断 |
| 前馈 | GLU 门控 | 非门控 GELU，宽度 2× |
| logit 封顶 | `30 · tanh(x/30)` | 默认 `soft_sign`、上限 80；旗舰排序关闭，改用 QK 归一化 |
| 子层后归一化 | 总是做 | `primer_norm` 开关，排序旗舰关 |
| 末尾层 | 全序列 | 可只更新候选 |
| 规模（排序） | 演示 checkpoint：128 维 4 层 | 旗舰 2560 维 8 层 |
| MoE | 无 | 无 |

## 设计决策

| 决策 | 理由 |
|------|------|
| 非因果、按段 ID 屏蔽 | 推荐不需要自回归；用户前缀与历史互相可见，候选按规则隔离 |
| GQA | KV 头少，服务时 KV 缓存与带宽开销小 |
| QK 归一化取代 logit 封顶 | 两者都用来防止注意力 logit 过大；旗舰排序选前者。源码未写理由，此为推测 |
| 右锚定 RoPE | 最近行为与候选的相对位置恒定 |
| μP 缩放 | 宽度变化时按规则缩放学习率，小模型上调好的超参可沿用到大模型。README 只说明 512 是基准宽度（`phoenix/README.md:419-423`），超参迁移是推测 |

## 源码锚点

- `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 排序模型](https://x-algorithm.vergego.ai/concepts/phoenix-ranking/) —— 使用该骨架的排序模型
- [Phoenix 召回：双塔检索架构](https://x-algorithm.vergego.ai/concepts/phoenix-retrieval/) —— 使用该骨架的召回用户塔
- [候选隔离注意力掩码](https://x-algorithm.vergego.ai/concepts/candidate-isolation-masking/) —— 段 ID 掩码与候选专用层
- [Phoenix 训练](https://x-algorithm.vergego.ai/concepts/phoenix-training/) —— 优化器与训练步

Source: https://x-algorithm.vergego.ai/concepts/grok-transformer/index.mdx