# CTI 生成式推荐广告排序推理优化方案

> 基于 baseline 代码分析、HSTU / GRAB 论文研究、官方提交规范

---

## 0. 前置修复：接口对齐（必须首先完成）

### 0.1 发现的问题

官方 `任务提交接口说明.md` 定义了评测系统调用接口。baseline 代码与规范存在三处致命不匹配：

| 接口 | 官方要求 | Baseline 实际 | 后果 |
|------|----------|---------------|------|
| 数据集类名 | `CTRTestSeqDataset` | `CTRUserDataset` | 评测 `from infer import CTRTestSeqDataset` 失败 |
| 构造参数 | `test_logids_ordered`, `item_dict`, `user_seq`, `max_feasign_per_slot`, `max_ctx_len` | `item_dict`, `user_seq`, `max_feasign_per_slot`, `pred_logids` | 参数名和数量不匹配 |
| `load_model` | `load_model(ckpt_path: Path) -> (model, device)` | `load_model(device='cuda:0', ckpt_path=None)` | 评测调用时 Path 会被错误赋给 `device` 参数 |

评测系统会用 `python -c "from infer import CTRTestSeqDataset, load_model; ..."` 来加载你的代码。**这三处不改，任何优化都白费。**

### 0.2 修复步骤

**修复 1：重命名类并调整构造参数**

将 `CTRUserDataset` 改为 `CTRTestSeqDataset`，参数名改为 `test_logids_ordered`，增加 `max_ctx_len` 占位：

```python
class CTRTestSeqDataset(Dataset):
    """按用户组织的 CTR 测试数据集（对齐评测接口）"""

    def __init__(self, test_logids_ordered, item_dict, user_seq=None,
                 max_feasign_per_slot=None, max_ctx_len=None):
        super().__init__()
        self.item_dict = item_dict
        self.user_seq = user_seq if user_seq else {}
        self.max_feasign_per_slot = max_feasign_per_slot
        self.max_ctx_len = max_ctx_len
        self.pred_logids = set(test_logids_ordered) if test_logids_ordered else set()
        # ... 其余逻辑不变
```

**修复 2：修正 `load_model` 签名**

```python
def load_model(ckpt_path, device='cuda:0'):
    """加载模型。签名对齐评测接口：第一个参数必须是 ckpt_path。"""
    # ... 其余逻辑不变
```

### 0.3 提交验证标准

修改后必须在 AI Studio 提交一次，确认能跑通（得分 > 0），再开始做优化。这是所有优化的前提。

---

## 1. 优化起点（Baseline 数据）

| 指标 | 当前值 | 说明 |
|------|--------|------|
| 推理耗时 | 229.18s | **只计 `model(batch)` 的逐 batch 累加时间** |
| AUC | 0.759 | 阈值 ≥ 0.65 |
| PCOC | 1.110 | 阈值 [0.85, 1.15] |
| 综合得分 | 25.85 | `score_latency * 70 + score_model * 30` |

技术栈：
- Python 3.10.10 + PyTorch 2.6.0 + CUDA 12.4
- 模型：RepEncoder（28 slot / 512 维 embedding）→ 8 层 Transformer（8 头 / 512 维）→ MoE FFN（8 个 expert，Top-2 gating）
- 数据：已缓存为 9 个 shard 分片（shard_0000.pt ~ shard_0008.pt），共 2039 batch / 7774 条预测

---

## 2. 约束条件（全部来自官方规则）

| 约束 | 来源 | 说明 |
|------|------|------|
| 不能改"组网" | 赛题说明 | 模型结构、层数、头数、维度不可改 |
| 不能改参数（权重值） | 赛题说明 | 量化/稀疏化/剪枝**明确允许** |
| 不能用测试集训练 | 赛题说明 | 仅推理，不做任何训练 |
| 推理时限 ≤ 300s | 提交规范 §4.1 | 超时总分直接 0。**只计 `model(batch)` 时间，数据加载和模型加载不计** |
| build_env.sh 时限 ≤ 720s | 提交规范 §3 | 超时或非 0 退出码直接失败 |
| AUC ≥ 0.65 且 PCOC ∈ [0.85, 1.15] | 提交规范 §4.2 | 任一不满足，总分直接 0 |
| 压缩包内不能有 `dataset/` 和 `ckpt.pt` | 提交规范 §2.2 | 评测系统自行提供 |
| 压缩包后缀必须是 `.zip`/`.tar.gz`/`.tar` | 提交规范 §1 | 其他格式不识别 |
| 解压后文件必须直接在根目录 | 提交规范 §1 | 不能多一层包裹文件夹 |

---

## 3. 优化方案总览（修订版）

```
第〇步：接口对齐（不优化，先跑通）
  └── 确认能提交得分 > 0
       │
第①步：FP16 量化 ←── 明确允许，收益最大
  └── 预期：229s → ~120s
       │
第②步：Flash Attention ←── 数学等价，不改组网
  └── 预期：120s → ~90s
       │
第③步：torch.compile ←── 编译器优化，不改组网
  └── 预期：90s → ~65s
       │
第④步：数据流优化 ←── 减少 CPU→GPU 传输开销
  └── 预期：65s → ~55s
       │
第⑤步：MoE 优化 ←── "剪枝"明确允许
  └── 预期：55s → ~50s
       │
第⑥步：INT8 量化（可选） ←── 收益大但风险高
  └── 只有在①②③⑤之后仍不够时才尝试
```

每个方案完成后在 AI Studio 提交验证，确认分数提升后再进入下一步。

---

## 4. 方案一：FP16 量化

### 4.1 合规性

✅ 比赛规则明确写有"量化除外"，属于允许范围。

### 4.2 原理

当前 baseline 使用 FP32。FP16 将模型参数和激活值减半，GPU 的张量核心对 FP16 有原生加速。注意 `Embedding` 层保留 FP32（索引是整数，embedding 查表时再转换）。

### 4.3 实现

修改 `load_model` 函数中的模型加载部分：

```python
def load_model(ckpt_path, device='cuda:0'):
    # ... 模型初始化不变 ...

    dev = torch.device(device if torch.cuda.is_available() else "cpu")

    if ckpt_path is None:
        ckpt_path = Path(__file__).parent / 'ckpt.pt'
    else:
        ckpt_path = Path(ckpt_path)

    if ckpt_path.exists():
        ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
        model.load_state_dict(ckpt['model_state_dict'])

        # === FP16 优化：将模型转为半精度 ===
        model = model.half()
        # Embedding 层保留 FP32（索引是 int，不需要转）
        model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)

        print(f"[INFO] Loaded checkpoint from {ckpt_path} (epoch={ckpt.get('epoch', '?')})")
    else:
        print(f"[WARNING] Checkpoint {ckpt_path} not found, using random weights")

    model.to(dev)
    model.eval()
    print(f"[INFO] Model ready. Device: {dev}")
    return model, dev
```

`move_batch_to_device` 中同步处理输入数据精度：

```python
def move_batch_to_device(batch, device):
    if isinstance(batch, dict):
        return {k: move_batch_to_device(v, device) for k, v in batch.items()}
    elif isinstance(batch, (list, tuple)):
        return [move_batch_to_device(x, device) for x in batch]
    elif torch.is_tensor(batch):
        x = batch.to(device)
        # 浮点 tensor → FP16，整数 tensor 保持不变
        if x.dtype == torch.float32:
            x = x.half()
        return x
    else:
        return batch
```

同时修改 `main()` 中缓存 batch 的逻辑，在落地磁盘前就转 FP16，避免推理时逐 batch 转换：

```python
for batch in test_loader:
    # 预转为 FP16
    batch = move_batch_to_device(batch, torch.device('cpu'))
    all_batches.append(batch)
```

### 4.4 风险

- 极低。现代 NVIDIA GPU（V100/A100/H100 等）的 FP16 吞吐量是 FP32 的 2-8 倍
- Embedding 保留 FP32 是因为 embedding 查表操作用 int 索引，不受浮点精度影响

### 4.5 预期收益

| 指标 | 变化 |
|------|------|
| 推理时间 | 229s → ~115-150s |
| AUC | 几乎不变（差异 < 0.0001） |

---

## 5. 方案二：Flash Attention

### 5.1 合规性

✅ Flash Attention 是**数学等价**的注意力计算。它用分块（tiling）算法计算 `softmax(QK^T/√d)V`，**输出结果与标准 attention 相同**（在浮点误差范围内）。不改变模型结构、不改变参数——只改变了内存访问模式和计算顺序。这属于编译器级优化，不是组网修改。

### 5.2 原理

Baseline 的 `scaled_dot_product` 先计算完整的 `QK^T` 矩阵（O(L²) 显存），再做 softmax，再乘 V。Flash Attention 将 Q、K、V 分块加载到 SRAM，逐块计算 softmax 并累积结果，**避免完整 QK^T 驻留 HBM**。

Baseline 用了 Sequence Packing（同一用户的多个 impression 拼成一条长序列），序列越长收益越大。

### 5.3 实现

PyTorch 2.0+ 提供了 `F.scaled_dot_product_attention`，自动选择最优后端（Flash Attention / Memory Efficient Attention / 标准实现）：

```python
import torch.nn.functional as F

def scaled_dot_product(q, k, v, extension):
    """使用 PyTorch SDPA 后端（自动启用 Flash Attention）"""
    d = q.size(-1)

    if extension is not None and "mask" in extension:
        mask = extension["mask"]
        # mask 形状: [1, 1, S, S] 或 [B, 1, S, S]
        # 转换为 float mask，确保 device 和 dtype 一致
        attn_mask = mask.to(device=q.device, dtype=q.dtype)
    else:
        attn_mask = None

    return F.scaled_dot_product_attention(
        q, k, v,
        attn_mask=attn_mask,
        dropout_p=0.0,
        is_causal=False,
    )
```

验证是否启用了 Flash Attention：

```python
import torch
# 在模型加载后运行，确认后端可用
print("Flash SDP:", torch.backends.cuda.flash_sdp_enabled())
print("Mem Efficient SDP:", torch.backends.cuda.mem_efficient_sdp_enabled())
print("Math SDP:", torch.backends.cuda.math_sdp_enabled())
```

### 5.4 与 Baseline 输出的兼容性

Baseline 中 mask 是 `torch.bool` 类型，原代码用 `scores.masked_fill(mask == 0, float("-inf"))` 处理。`F.scaled_dot_product_attention` 接受 bool mask 时自动做等价处理。数学等价，不会影响 AUC/PCOC。

### 5.5 预期收益

| 序列长度 | 加速比 |
|----------|--------|
| < 512 | 1.1x - 1.3x |
| 512 - 2048 | 1.5x - 2x |
| > 2048 | 2x - 4x |

---

## 6. 方案三：torch.compile

### 6.1 合规性

✅ `torch.compile` 是纯粹的**编译器优化**（JIT 编译 + 算子融合）。它不改变模型结构、不改变权重值、不改变任何数学运算——只是把多个小 kernel 合并成一个大 kernel 以减少 GPU kernel launch 开销。

### 6.2 实现

在 `load_model` 中，`model.eval()` 之后添加一行：

```python
model = torch.compile(model, mode="reduce-overhead")
```

`mode` 选择：
| 模式 | 行为 | 适合场景 |
|------|------|----------|
| `"default"` | 平衡编译速度与运行时 | 不确定时先用这个 |
| `"reduce-overhead"` | 更激进的融合，减少 kernel launch | **推理场景首选** |
| `"max-autotune"` | 自动调优 Triton kernel（首次编译慢） | 最终提交时切换到此 |

建议先用 `"reduce-overhead"` 快速验证，最终提交时换成 `"max-autotune"`。

### 6.3 已知限制

- 首次运行会触发 JIT 编译（0.5-2 分钟），在评测环境可能计入推理时间
- 解决：在 `build_env.sh` 中用一个小输入做一次预热编译：

```bash
#!/bin/bash
set -e
pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/

# 预热 torch inductor，避免推理时编译
python -c "
import torch
@torch.compile(mode='max-autotune')
def warmup(x):
    return x * 2
x = torch.randn(100, 100, device='cuda')
warmup(x)
print('Inductor cache ready')
"

echo "build env success"
```

### 6.4 预期收益

| 模式 | 加速比 |
|------|--------|
| reduce-overhead | 1.2x - 1.4x |
| max-autotune | 1.3x - 1.5x（首次多花 1-2 分钟） |

---

## 7. 方案四：数据流优化

### 7.1 重要澄清：评测的计时逻辑

从提交规范 §4.1：
> 只统计模型前向时间（逐 batch 累加 `model(batch)` 耗时），数据加载、模型加载不计入。

所以 **DataLoader 层面的优化（num_workers、pin_memory）不影响评分**。但这不代表数据流优化没用——以下两点仍然有效：

1. **减少 `move_batch_to_device` 开销**：此函数在 `model(batch)` 之前调用，但如果它在推理循环里逐 batch 执行，仍会阻塞。解决方案：提前把所有 batch 搬到 GPU。
2. **减少 FP32→FP16 转换开销**：在缓存阶段就转好，推理循环里无需重复转换。

### 7.2 实现

在 `main()` 的数据加载部分，缓存时就完成设备搬运和类型转换：

```python
# 原代码
all_batches = [batch for batch in test_loader]

# 优化：缓存时直接搬到 GPU + 转 FP16
dev = torch.device('cuda:0')
all_batches = []
for batch in test_loader:
    batch = move_batch_to_device(batch, dev)
    all_batches.append(batch)
```

这样推理循环中不需要再调用 `move_batch_to_device`：

```python
with torch.no_grad():
    for batch in tqdm(all_batches, desc="Inference"):
        # batch 已在 GPU 上，无需 move_batch_to_device
        pred_mask = batch["pred_mask"].bool()
        t_start = time.time()
        logits, moe_loss = model(batch)
        logits = logits.squeeze(-1)
        probs = torch.sigmoid(logits)
        time_sum += time.time() - t_start
        # ...
```

### 7.3 潜在问题

全部 batch 都在 GPU 上会占用大量显存。如果数据总量超过可用显存，会 OOM。安全做法：先测试单个 shard 的显存占用，若不够则分批预加载。

如果显存不够，采用双缓冲策略（当前 batch 在 GPU 推理，下一个 batch 异步上传）。

### 7.4 预期收益

有限（因为数据加载本来就不计分），但能消除循环内的不必要开销。

---

## 8. 方案五：MoE 推理优化

### 8.1 合规性

✅ "剪枝"是规则明确允许的三种优化之一（量化/稀疏/剪枝）。Expert 剪枝属于模型剪枝的子类。

### 8.2 当前状态

```python
class SMoE(nn.Module):
    def __init__(self, d_model, dim_ff, num_experts=8, k=2):
        # 每层 8 个 FFN expert，每个 token 激活 2 个
```

8 层 × 8 expert = 64 个 FFN 模块。推理时大多数 token 只会路由到少数 expert。

### 8.3 优化方向

**a) Expert 负载统计 + 合并**

对少量数据做一次前向，统计各 expert 的被激活次数：

```python
expert_hits = torch.zeros(8, 8)  # [layers, experts]
model.eval()
with torch.no_grad():
    for batch in sample_batches[:10]:  # 只用 10 个 batch 采样
        batch = move_batch_to_device(batch, dev)
        for layer_idx, moe in enumerate(model.seq_encoder.moe):
            topk_idx, _, _ = moe.gate(x)  # x 需要从 forward 中间取
            for e in range(8):
                expert_hits[layer_idx, e] += (topk_idx == e).sum()
```

如果某个 expert 激活次数 < 1%，可将其权重合并到最相似的 expert（通过权重余弦相似度找最近邻），或直接移除（需验证 AUC 不掉）。

**b) 替换为静态 FFN（激进方案）**

如果发现某些层的 expert 负载极度不均衡（例如 95% token 走同一个 expert），可直接把该层的 SMoE 替换为单个 Expert：

```python
# 仅当某层 expert 负载极不均衡时
# 找到最常用的 expert
best = expert_hits[layer_idx].argmax().item()
# 替换 SMoE 为单个 Expert
model.seq_encoder.moe[layer_idx] = model.seq_encoder.moe[layer_idx].experts[best]
```

此方案**必须提交验证 AUC/PCOC 不跌破阈值**。

### 8.4 风险

- 负载统计需要从 Transformer forward 中间提取 gate 输出，需要修改 forward 添加 hook
- Expert 合并可能影响 AUC 0.001-0.005

### 8.5 预期收益

| 操作 | 加速比 |
|------|--------|
| 移除 1-2 个死 expert | 1.05x - 1.15x |
| 替换 2-3 层为单 FFN | 1.2x - 1.4x |

---

## 9. 方案六：INT8 量化（可选）

### 9.1 合规性

✅ "量化除外"，明确允许。

### 9.2 适用条件

只有在前 5 步做完后仍需进一步加速时才尝试。INT8 量化需要校准数据（可用测试集的一小部分做 calibration），有精度损失风险。

```python
import torch.quantization as quant

# 仅对 Transformer encoder 做 INT8 量化（Embedding 层跳过）
model.seq_encoder.qconfig = quant.get_default_qconfig('qnnpack')
model.seq_encoder = quant.quantize_dynamic(
    model.seq_encoder,
    {nn.Linear},  # 仅量化 Linear 层
    dtype=torch.qint8,
)
```

### 9.3 风险

- AUC 可能下降 0.005-0.02，需提交后验证
- 如果评测环境无 GPU（CPU-only），QNNPACK 才有用；GPU 下需用 TensorRT 做 INT8

---

## 10. CUDA Graph（为什么不做）

原方案中列了 CUDA Graph，审查后**移除**，理由如下：

1. CUDA Graph 要求**所有输入形状完全相同**。Baseline 数据经过 Sequence Packing，不同 batch 的序列长度差异很大，不满足这一前提
2. 每个 batch 的 `user_offsets` 长度不同，导致 mask 形状也不同
3. 若要强行使用，需要对 batch 做 padding 对齐，反而引入额外开销

**结论：不适用于本场景，放弃此方案。**

---

## 11. 部署要点

### 11.1 压缩包结构

```
submit.zip
├── infer.py              # 主推理脚本（实现所有必需接口）
├── requirements.txt      # Python 依赖列表
└── build_env.sh          # 环境构建脚本（可选但推荐）
```

提交前务必验证：
- [ ] 包内**没有** `dataset/` 目录
- [ ] 包内**没有** `ckpt.pt` 文件
- [ ] 包内**没有**多余的顶层文件夹
- [ ] 后缀是 `.zip`（或 `.tar.gz`）

### 11.2 build_env.sh 设计

```bash
#!/bin/bash
set -e
# 安装依赖（评测系统自动使用阿里云镜像）
pip install -r requirements.txt

# 预热 torch compile（如果方案三启用）
python -c "
import torch
@torch.compile(mode='max-autotune')
def _warmup(x): return x * 2
_warmup(torch.randn(100, 100, device='cuda'))
print('Inductor ready')
" 2>/dev/null || echo 'torch.compile not available, skipping'

echo "build env success"
```

### 11.3 requirements.txt（最小化）

```
torch==2.6.0
triton==3.2.0
numpy==2.2.6
scikit-learn==1.7.2
tqdm==4.67.3
```

- 去掉所有 `nvidia-*` 包（评测环境已预装 CUDA）
- 版本号精确锁定，避免安装时依赖冲突

---

## 12. 风险评估与底线策略

### 各方案对模型质量的影响

| 方案 | AUC 影响 | PCOC 影响 | 0 分风险 |
|------|----------|-----------|----------|
| 接口修复 | 无 | 无 | **必须做，否则直接 0 分** |
| FP16 | < 0.0001 | < 0.0001 | 极低 |
| Flash Attention | < 0.0001 | 无 | 极低 |
| torch.compile | 无 | 无 | 低（首次编译可能超时） |
| GPU 预加载 | 无 | 无 | 低（OOM 风险） |
| MoE 剪枝 | 0.001 - 0.005 | 微小 | 中（需提交验证） |
| INT8 | 0.005 - 0.02 | 可能偏移 | 高（可能跌破阈值） |

### 安全策略（必须遵守）

每完成一个优化：
1. 在 AI Studio 提交，拿到新得分
2. 记录本次得分变化
3. 如果 AUC < 0.65 或 PCOC 不在 [0.85, 1.15]，立即回退该优化
4. 如果得分上升，保留并进入下一步

### 保底策略

- 至少完成接口修复 + FP16，能稳定拿到 > 25 分
- 如果 Flash Attention / torch.compile 跑不通，回退不影响得分
- **在截止日期前一天（6 月 25 日）停止实验**，提交当前最优版本

---

## 13. 实施检查清单

### 第〇周（立即）
- [ ] 修复接口：`CTRTestSeqDataset`、`load_model(ckpt_path, ...)` 
- [ ] 在 AI Studio 提交一次，确认得分 > 0（验证接口正确）

### 第一轮
- [ ] FP16 量化：`model.half()` + embedding 保留 FP32
- [ ] 数据预加载到 GPU + 预转 FP16
- [ ] 提交验证

### 第二轮
- [ ] Flash Attention：替换 `scaled_dot_product`
- [ ] `torch.compile(mode="reduce-overhead")`
- [ ] `build_env.sh` 写预热逻辑
- [ ] 提交验证

### 第三轮（时间允许）
- [ ] MoE expert 负载分析 + 合并
- [ ] `torch.compile` 切换为 `"max-autotune"`
- [ ] INT8 量化评估（如果得分仍不满意）

---

## 14. 预期效果（修订）

| 阶段 | 预期推理时间 | 预期得分 | 主要贡献 |
|------|-------------|----------|----------|
| Baseline | 229s | 25.85 | — |
| 接口修复 | 229s | 25.85 | 确保能跑 |
| + FP16 + GPU 预加载 | ~120s | ~50 | 量化为主要贡献 |
| + Flash Attention | ~90s | ~60 | 长序列受益 |
| + torch.compile | ~65s | ~70 | 算子融合 |
| + MoE 优化 | ~50s | ~78 | 剪枝 |
| 极限（+INT8） | ~30s | ~87 | 有精度风险 |

---

## 附录 A：官方规则原文引用

来自赛题说明页面：
> 【由于是推理性能优化，组网不可进行策略性改动，不可对测试集进行训练】

来自 baseline notebook 第三单元：
> 选手不可对组网和相关参数进行修改。【违规为0分】
> **量化稀疏剪枝除外**

来自 `任务提交接口说明.md` §4-5：
> - 延迟阈值：300 秒。只统计模型前向时间（逐 batch 累加 model(batch) 耗时）
> - AUC ∈ [0.65, 1.0]、PCOC ∈ [0.85, 1.15]，任一不满足总分直接置 0
> - 压缩包自带 dataset/ 或 ckpt.pt → 评测失败
> - build_env.sh 超时 720 秒或返回码非 0 → 评测失败

## 附录 B：相关资源

| 资源 | 链接 |
|------|------|
| GRAB 论文 | arXiv 2602.01865 |
| HSTU 论文 | arXiv 2402.17152 (ICML 2024) |
| 官方 Baseline 项目 | AI Studio project/10186630 |
| 比赛主页 | aistudio.baidu.com/competition/detail/1461 |
| 提交结果页 | aistudio.baidu.com/competition/detail/1461/0/submit-result |