CTI-Inference-Opt/docs/superpowers/plans/2026-06-14-cti-auc-recovery.md

CTI 推理优化冲击 80+ 实现计划

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: 在不改模型结构、不训练测试集的前提下，先找回当前推理丢失的 AUC，再做结构性延迟重写，把榜上分数从 58.86 推向 80+。

Architecture: 在 AI Studio notebook（A800 + dataset + ckpt.pt）里，先建一个带同步计时和配置开关的测量闭环 bench.py；阶段 A 用消融实验定位并找回 AUC（30 分桶）；阶段 B 用数值等价的内核重写压低延迟（块对角注意力 / MoE 向量化 / embedding 融合）。每步过本地关卡，再用有限的提交确认验证集。

Tech Stack: Python 3.10, PyTorch 2.6.0 (CUDA 12.4), NVIDIA A800 (SM80), sklearn (AUC), AI Studio notebook。

---

执行环境约定

• 所有运行都在 AI Studio notebook 内（本地 Windows 只装了 numpy+tqdm，跑不了 torch）。
• 提交文件只有 infer.py / requirements.txt / build_env.sh 会被打包；bench.py、tests/ 绝不进提交包。
• 每个改 infer.py 的任务，最后都要确认 bench.py 默认配置仍能复现「当前最优」，避免污染提交版本。
• 数据路径（notebook 内）：代码/code/dataset/（软链）、代码/code/ckpt.pt、本地标签 dataset/label_data.txt。

文件结构

文件	职责	是否提交
代码/code/infer.py	提交主脚本。引入模块级 CONFIG 开关；load_model/RepEncoder/SMoE/注意力按 CONFIG 行为，默认值=当前最优	✅
代码/code/bench.py	测量闭环。设置 infer.CONFIG，跑本地推理，同步计时，打印 AUC/PCOC/延迟/总分；支持配置扫描	❌
代码/code/tests/test_equiv.py	阶段 B 重写的数值等价测试（新实现 vs 原实现 allclose）	❌
代码/code/EXPERIMENTS.md	实验记录表（配置 → AUC/PCOC/延迟/本地分/提交分）	❌（可入 git，不入提交包）

---

阶段 0：测量闭环

Task 1: 给 infer.py 加 CONFIG 开关板

Files:

• Modify: 代码/code/infer.py（顶部新增 CONFIG；改 load_model、RepEncoder.forward）

• Step 1: 在 import 之后、数据加载层之前插入模块级 CONFIG

# ============================================================
# 实验配置开关（提交时保持默认 = 当前最优行为）
# bench.py 会在 import 后覆盖这些值；评测系统不碰它，用默认值。
# ============================================================
CONFIG = {
    "fp16": True,            # True=半精度；False=FP32 参考
    "keep_fp32_modules": (), # 在 fp16 下仍保留 FP32 的子模块名前缀，如 ("rep_encoder.emb",)
    "expert_merge": True,    # 是否做 expert 相似度合并
    "merge_threshold": 0.90, # 合并余弦阈值
    "signid_mode": "clamp",  # "clamp" 或 "modulo"，处理超界 sign id
    "sync_timing": False,    # bench 里设 True，做 torch.cuda.synchronize 真实计时
}

• Step 2: 改 RepEncoder.forward，按 CONFIG 处理 sign id

把 代码/code/infer.py 中 RepEncoder.forward 的这一行：

            values = values.clamp(0, max_idx)  # 超出 vocab_size 的 sign id 截断，避免越界

替换为：

            if CONFIG["signid_mode"] == "modulo":
                values = values % self.emb.num_embeddings
            else:
                values = values.clamp(0, max_idx)

• Step 3: 改 load_model，按 CONFIG 控制 fp16 / 保留 FP32 模块 / expert 合并

把 load_model 中从 model = model.half() 到 _merge_experts(...) 这一段：

        # === FP16 量化：模型参数转半精度，Embedding 保留 FP32 ===
        model = model.half()
        model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
        print("[INFO] Model converted to FP16 (embedding kept in FP32)")

        # === 按 Expert 权重相似度合并冗余 expert ===
        _merge_experts(model, sim_threshold=0.90)

替换为：

        if CONFIG["fp16"]:
            model = model.half()
            # embedding 始终保留 FP32（int 索引查表）
            model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
            # 额外保留 FP32 的模块（精度敏感层）
            for name, module in model.named_modules():
                if any(name.startswith(p) for p in CONFIG["keep_fp32_modules"]):
                    module.to(torch.float32)
            print(f"[INFO] FP16 on; FP32-kept: {('rep_encoder.emb',) + CONFIG['keep_fp32_modules']}")
        else:
            model = model.float()
            print("[INFO] FP32 reference (no half)")

        if CONFIG["expert_merge"]:
            _merge_experts(model, sim_threshold=CONFIG["merge_threshold"])
        else:
            print("[INFO] expert_merge off")

注意：keep_fp32_modules 里若含某层（如 seq_encoder.norm1），其输入需在该层处转回 FP32。先只用整体 fp16/fp32 与 emb，敏感层在 Task 5 单独处理；本任务只接好开关。

• Step 4: 在 notebook 跑一遍默认配置，确认行为未变

Run（notebook cell）：

%cd /home/aistudio/code
!python infer.py

Expected：打印 FP16 on、expert 合并日志，AUC ≈ 0.759、PCOC ≈ 1.05~1.11（与改动前一致，证明开关默认值没改变行为）。

• Step 5: Commit

git add 代码/code/infer.py
git commit -m "feat: infer.py 增加 CONFIG 实验开关（默认=当前最优行为）"

Task 2: 建 bench.py 测量闭环

Files:

• Create: 代码/code/bench.py

• Step 1: 写 bench.py

"""本地测量闭环：设置 infer.CONFIG，跑推理，同步计时，打印指标。不进提交包。"""
import sys, time, io
from pathlib import Path
import torch
from torch.utils.data import DataLoader

import infer  # 同目录


def run_once(config_override: dict, batch_size: int = 50, max_batches: int | None = None):
    infer.CONFIG.update(config_override)
    infer.CONFIG["sync_timing"] = True

    cur = Path(__file__).parent
    ref = cur / "dataset"
    history = ref / "history"
    test_csv = ref / "test.csv"
    label_file = ref / "label_data.txt"

    files = (sorted(history.glob("*.csv")) if history.exists() else []) + [test_csv]
    item_dict, user_seq = infer.load_sample_files(files)
    test_logids = infer.load_logids_from_file(test_csv)
    ds = infer.CTRTestSeqDataset(
        test_logids_ordered=list(test_logids), item_dict=item_dict,
        user_seq=user_seq, max_feasign_per_slot={1: 2}, max_ctx_len=None,
    )
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0,
                        collate_fn=infer.make_collate_fn(ds.max_slot_id))
    batches = []
    for b in loader:
        batches.append(infer.move_batch_to_device(b, torch.device("cpu")))
        if max_batches and len(batches) >= max_batches:
            break

    model, dev = infer.load_model(ckpt_path=None)
    logid2p, t_sum = {}, 0.0
    with torch.inference_mode():
        for b in batches:
            b = infer.move_batch_to_device(b, dev)
            pm = b["pred_mask"].bool()
            torch.cuda.synchronize()
            t0 = time.time()
            logits, _ = model(b)
            probs = torch.sigmoid(logits.squeeze(-1))
            torch.cuda.synchronize()
            t_sum += time.time() - t0
            for lid, p in zip(b["logid"][pm].cpu().tolist(), probs[pm].cpu().tolist()):
                logid2p[lid] = p

    # 按 test.csv 顺序写 predict 并打分
    order = [int(l.split(",")[0]) for l in open(test_csv) if l.strip()]
    pred_path = cur / "predict.txt"
    with open(pred_path, "w") as f:
        for lid in order:
            f.write(f"{logid2p[lid]}\n")
    res = infer._cal_score(pred_path, label_file, default_latency=t_sum)
    print(f"[BENCH] cfg={config_override} bs={batch_size} -> "
          f"AUC={res['auc']:.5f} PCOC={res['pcoc']:.4f} "
          f"lat={res['latency']:.2f}s score={res['score_all']:.2f}")
    return res


if __name__ == "__main__":
    run_once({})  # 默认配置基准

• Step 2: 跑默认配置，建立本地基准

Run：

%cd /home/aistudio/code
!python bench.py

Expected：打印 [BENCH] 一行，记录 AUC/PCOC/同步后真实延迟/本地分。这是后续所有对比的锚点。

• Step 3: 建实验记录表并记录第一行

Create 代码/code/EXPERIMENTS.md，写入表头与默认配置那一行（数值用 Step 2 实测填）：

| 配置 | AUC | PCOC | 延迟(同步) | 本地分 | 提交分 |
|------|-----|------|-----------|--------|--------|
| 默认(当前最优) | <实测> | <实测> | <实测> | <实测> | 58.86 |

• Step 4: Commit

git add 代码/code/bench.py 代码/code/EXPERIMENTS.md
git commit -m "feat: 新增 bench.py 测量闭环 + 实验记录表"

---

阶段 A：找回 AUC（30 分桶，最高优先）

Task 3: FP32 参考跑 —— 确立 AUC 天花板（核心前提验证）

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 跑纯 FP32、不合并 expert、clamp

Run（notebook）：

import bench
bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "clamp"})

Expected：打印一行 AUC/PCOC/延迟。记录这个 AUC —— 它是当前代码路径下模型的真实可达上限。

• Step 2: 判定核心前提

把结果记入 EXPERIMENTS.md。判定：

• 若 FP32 AUC 明显 > 默认配置 AUC（如 ≥ +0.01）→ 说明 fp16/合并在掉精度，Task 4/5 有收益。

• 若 FP32 AUC 仍 ≈ 0.759（验证集对应 ~0.7526）→ 当前数据路径触不到更高 AUC；缺口可能在 sign-id/特征/上下文（Task 3.5/6），或「80 目标」前提存疑，需暂停并与队友/官方答疑核对（见 spec §10）。

• Step 3: Commit

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: FP32 参考跑，记录 AUC 天花板"

Task 4: Sign-ID 取模 vs clamp

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 先查 max_sign_id 是否超 5M 词表

Run（notebook）：

import infer
from pathlib import Path
files = sorted(Path("dataset/history").glob("*.csv")) + [Path("dataset/test.csv")]
item_dict, user_seq = infer.load_sample_files(files)
mx = max(int(s) for r in item_dict.values() for s in r["signs"].tolist())
print("max_sign_id =", mx, "vocab =", 5000000, "超界比例可观?", mx >= 5000000)

Expected：打印最大 sign id。若 mx >= 5_000_000，clamp 会把大量 id 压到同一行 —— 头号嫌疑成立。

• Step 2: FP32 下对比 clamp vs modulo

Run：

import bench
bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "clamp"})
bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "modulo"})

Expected：两行 AUC。

• Step 3: 判定 + 记录

• modulo 的 AUC 明显更高 → 训练用的就是取模哈希，保留 modulo（合规：只是正确还原模型输入，不改结构/权重）。

• 两者相近或 modulo 更差 → 训练用 clamp/或 id 不超界，保留 clamp。 记入 EXPERIMENTS.md。

• Step 4: Commit

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: sign-id clamp vs modulo 对比"

Task 5: 精度摆放（混合精度找回 AUC）

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 逐步把敏感层保留 FP32，对比 AUC

用上一步定下的 signid_mode（记为 SM），依次跑：

import bench
bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM,
                "keep_fp32_modules": ()})                       # 纯 fp16
bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM,
                "keep_fp32_modules": ("linear",)})              # 保留最终输出头
bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM,
                "keep_fp32_modules": ("linear", "rep_encoder.input_norm",
                                       "rep_encoder.linear")})  # +RepEncoder 头

Expected：三行 AUC + 延迟。

• Step 2: 选「AUC 最接近 FP32 且延迟可接受」的组合

记 KEEP = 选中的 keep_fp32_modules。判定标准：相对 FP32 参考，AUC 损失 ≤ 0.001 优先；若纯 fp16 已无损，则 KEEP=()。记入 EXPERIMENTS.md。

• Step 3: Commit

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: 混合精度摆放，确定 keep_fp32_modules"

Task 6: Expert 合并的 AUC 代价

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 在选定精度下对比 expert_merge 开/关

import bench
bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP,
                "expert_merge": False})
bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP,
                "expert_merge": True, "merge_threshold": 0.90})

Expected：两行，含 AUC 与延迟。

• Step 2: 判定

• 合并掉 AUC（> 0.0005）但只省一点延迟 → 关掉合并（延迟从阶段 B 补，那里不损精度）。

• 合并不掉 AUC → 保留。记 MERGE = 最终决定。记入 EXPERIMENTS.md。

• Step 3: Commit

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: 量化 expert 合并的 AUC 代价并决定开关"

Task 7: 特征与上下文完整性核查

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 核查 max_feasign_per_slot 截断的影响

import bench
bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP,
                "expert_merge": MERGE})  # 当前 dataset 用 {1:2}

然后改 bench.run_once 里 max_feasign_per_slot={1: 2} 为 None（临时编辑 bench.py 或加参数），再跑一次，对比 AUC。 Expected：两行。若去掉截断 AUC 升高，说明截断在丢信息。

注意：评测系统构造 CTRTestSeqDataset 时传哪些 max_feasign_per_slot/max_ctx_len 由评测端决定，我们不一定能控制。本步先确认「完整特征是否更好」，若是，则在 CTRTestSeqDataset.__init__ 里对截断做更保守的默认（仅在确证合规、不属"序列截断"违规的前提下）。

• Step 2: 核查每条测试样本是否 attend 到完整用户历史

import infer
from pathlib import Path
files = sorted(Path("dataset/history").glob("*.csv")) + [Path("dataset/test.csv")]
item_dict, user_seq = infer.load_sample_files(files)
test_uids = {item_dict[l]["userid"] for l in infer.load_logids_from_file(Path("dataset/test.csv"))}
have_hist = sum(1 for u in test_uids if len(user_seq.get(u, [])) > 1)
print(f"测试用户 {len(test_uids)}，其中有历史序列(>1)的 {have_hist} "
      f"({have_hist/len(test_uids):.1%})；序列长度分布：")
import numpy as np
lens = np.array([len(user_seq.get(u, [])) for u in test_uids])
print("min/median/max =", lens.min(), int(np.median(lens)), lens.max())

Expected：绝大多数测试用户应有较长历史序列。若大量用户只有长度 1（无历史），说明历史没正确挂上 —— 这会严重压低生成式模型 AUC，需排查 load_sample_files 的 userid 关联与排序。

• Step 3: 记录结论 + Commit

把两步结论记入 EXPERIMENTS.md。

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: 特征截断与上下文完整性核查"

Task 8: 锁定阶段 A 最优配置并设为 infer.py 默认 + 提交验证

Files:

• Modify: 代码/code/infer.py（把 CONFIG 默认值改为阶段 A 选定组合）

• Step 1: 更新 infer.py 的 CONFIG 默认值

把 CONFIG 默认值改成 Task 4~7 选定的 signid_mode=SM、keep_fp32_modules=KEEP、expert_merge=MERGE、merge_threshold 等（sync_timing 保持 False）。

• Step 2: 跑默认配置确认达到阶段 A 最优本地分

%cd /home/aistudio/code
!python bench.py

Expected：AUC ≥ 默认基准，本地分高于先前。

• Step 3: 打包并提交一次（消耗 1 次/天额度）

cd /home/aistudio/code
rm -f predict.txt
zip -y ../eval.zip infer.py requirements.txt build_env.sh
# 确认包内无 dataset/、无 ckpt.pt、无 bench.py/tests/
unzip -l ../eval.zip

然后在 AI Studio 提交页提交 eval.zip。

• Step 4: 记录验证集分数 + Commit

把提交得到的验证集 AUC/PCOC/延迟/分数记入 EXPERIMENTS.md。

git add 代码/code/infer.py 代码/code/EXPERIMENTS.md
git commit -m "feat: 锁定阶段A最优配置为默认 + 验证集提交结果"

---

阶段 B：结构性延迟重写（数值等价，不动 AUC）

每个重写任务都先写「新实现 vs 原实现 allclose」等价测试，再替换，最后用 bench 确认 AUC 不变、延迟下降。

Task 9: 块对角因果注意力（FlexAttention）

Files:

• Create: 代码/code/tests/test_equiv.py

• Modify: 代码/code/infer.py（scaled_dot_product / CTRModel.forward mask 路径）

• Step 1: 写等价测试（先失败）

Create 代码/code/tests/test_equiv.py：

import torch, torch.nn.functional as F
import sys; sys.path.insert(0, "..")
import infer

def _dense_attn(q, k, v, mask):
    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask.to(q.dtype).bool())

def test_flex_matches_dense():
    torch.manual_seed(0)
    B, H, S, Dh = 1, 8, 37, 64
    q, k, v = [torch.randn(B, H, S, Dh, device="cuda") for _ in range(3)]
    # 构造 3 个用户的 user_offsets：长度 10/12/15
    offsets = torch.tensor([0, 10, 22, 37], device="cuda")
    m = infer.CTRModel.get_sequence_causal_mask.__get__(object())(offsets)  # 见下
    dense = _dense_attn(q, k, v, m.unsqueeze(0).unsqueeze(0))
    flex = infer.flex_block_causal_attn(q, k, v, offsets)
    assert torch.allclose(dense, flex, atol=1e-3, rtol=1e-3), (dense - flex).abs().max()

说明：get_sequence_causal_mask 是实例方法，测试里改成直接调用一个等价的独立函数 infer._build_dense_causal_mask(offsets)（Step 3 会把现有逻辑抽成模块级函数，便于测试与复用）。把上面 m = ... 那行改为 m = infer._build_dense_causal_mask(offsets)。

• Step 2: 跑测试确认失败

Run：

%cd /home/aistudio/code/tests
!python -m pytest test_equiv.py::test_flex_matches_dense -v

Expected：FAIL（infer.flex_block_causal_attn / _build_dense_causal_mask 未定义）。

• Step 3: 在 infer.py 实现 FlexAttention 路径

把 CTRModel.get_sequence_causal_mask 的逻辑抽为模块级函数，并新增 flex 实现：

from torch.nn.attention.flex_attention import flex_attention, create_block_mask

def _build_dense_causal_mask(user_offsets):
    lengths = user_offsets[1:] - user_offsets[:-1]
    idx = torch.repeat_interleave(
        torch.arange(lengths.numel(), device=user_offsets.device), lengths)
    same = idx.view(1, -1) == idx.view(-1, 1)
    causal = torch.tril(torch.ones_like(same, dtype=torch.bool))
    return same & causal

def flex_block_causal_attn(q, k, v, user_offsets):
    S = q.size(-2)
    lengths = user_offsets[1:] - user_offsets[:-1]
    doc_id = torch.repeat_interleave(
        torch.arange(lengths.numel(), device=q.device), lengths)
    def mask_mod(b, h, qi, ki):
        return (qi >= ki) & (doc_id[qi] == doc_id[ki])
    block_mask = create_block_mask(mask_mod, B=None, H=None, Q_LEN=S, KV_LEN=S, device=q.device)
    return flex_attention(q, k, v, block_mask=block_mask)

然后改 CTRModel.forward：mask 不再现造稠密矩阵传给 SDPA，而是把 user_offsets 透传，调用 flex_block_causal_attn。把 scaled_dot_product 改为接收 extension={"user_offsets": ...} 并走 flex；get_sequence_causal_mask 保留供测试/回退。

兼容性：FlexAttention 要求 q/k/v 为 [B,H,S,Dh]（现有 forward 已是该布局）。FP16 下 atol 放宽到 2e-2 重测。

• Step 4: 跑测试确认通过

Run：

!python -m pytest test_equiv.py::test_flex_matches_dense -v

Expected：PASS。

• Step 5: bench 确认 AUC 不变、延迟下降

import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench)
bench.run_once({})

Expected：AUC 与 Task 8 一致（±0.0005），延迟较 Task 8 下降。记入 EXPERIMENTS.md。

• Step 6: Commit

git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md
git commit -m "perf: 块对角因果注意力改用 FlexAttention（数值等价，提速）"

Task 10: MoE 向量化（消除 Python 循环与同步）

Files:

• Modify: 代码/code/infer.py（SMoE.__init__ 预堆叠权重；SMoE.forward 稠密批量计算）

• Modify: 代码/code/tests/test_equiv.py（加 MoE 等价测试）

• Step 1: 写 MoE 等价测试（先失败）

在 test_equiv.py 追加：

def test_smoe_vectorized_matches_loop():
    torch.manual_seed(0)
    m = infer.SMoE(d_model=512, dim_ff=1024, num_experts=8, k=2).cuda().eval()
    x = torch.randn(1, 50, 512, device="cuda")
    with torch.no_grad():
        ref, _ = infer._smoe_forward_loop(m, x)   # 原实现（保留为参考函数）
        new, _ = m(x)                              # 新向量化实现
    assert torch.allclose(ref, new, atol=1e-4, rtol=1e-4), (ref - new).abs().max()

• Step 2: 跑测试确认失败

Run：!python -m pytest test_equiv.py::test_smoe_vectorized_matches_loop -v Expected：FAIL（_smoe_forward_loop 未定义 / 新旧不一致）。

• Step 3: 实现向量化 SMoE

把现有 SMoE.forward 的循环体抽成模块级 _smoe_forward_loop(moe, x)（保留作参考/回退），新 forward 改为稠密批量（8 个小 FFN 全算，再按 top-k 选取加权 —— 数学等价，GPU 上无 gather/同步更快）：

class SMoE(nn.Module):
    def __init__(self, d_model, dim_ff, num_experts, k=2):
        super().__init__()
        self.num_experts = num_experts
        self.k = k
        self.experts = nn.ModuleList([Expert(d_model, dim_ff) for _ in range(num_experts)])
        self.gate = TopKGate(d_model, num_experts, k=k)
        self._stacked = False

    def _stack_weights(self):
        self.register_buffer("W1", torch.stack([e.fc1.weight for e in self.experts]))  # [E,F,D]
        self.register_buffer("b1", torch.stack([e.fc1.bias   for e in self.experts]))  # [E,F]
        self.register_buffer("W2", torch.stack([e.fc2.weight for e in self.experts]))  # [E,D,F]
        self.register_buffer("b2", torch.stack([e.fc2.bias   for e in self.experts]))  # [E,D]
        self._stacked = True

    def forward(self, x):
        if not self._stacked:
            self._stack_weights()
        B, S, D = x.shape
        topk_idx, topk_score, probs = self.gate(x)
        xf = x.reshape(-1, D)                                   # [N,D]
        h = torch.einsum("nd,efd->enf", xf, self.W1) + self.b1[:, None, :]   # [E,N,F]
        h = F.relu(h)
        o = torch.einsum("enf,eDf->enD", h, self.W2) + self.b2[:, None, :]   # [E,N,D]
        o = o.permute(1, 0, 2)                                  # [N,E,D]
        idx = topk_idx.reshape(-1, self.k)                      # [N,k]
        sc = topk_score.reshape(-1, self.k)                     # [N,k]
        sel = torch.gather(o, 1, idx.unsqueeze(-1).expand(-1, -1, D))        # [N,k,D]
        out = (sel * sc.unsqueeze(-1)).sum(1).reshape(B, S, D)
        moe_loss = probs.sum(dim=(0, 1)).std() / (probs.sum(dim=(0, 1)).mean() + 1e-6)
        return out, moe_loss

注意：合并 expert（Task 6 若开启）会改变 num_experts 和权重 —— _stack_weights 必须在合并之后、首次 forward 时调用（上面 lazy 实现已满足）。dtype 要与 x 一致（fp16 时 stack 出来即 fp16）。

• Step 4: 跑测试确认通过

Run：!python -m pytest test_equiv.py::test_smoe_vectorized_matches_loop -v Expected：PASS。

• Step 5: bench 确认 AUC 不变、延迟下降

import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench)
bench.run_once({})

Expected：AUC 一致，延迟较 Task 9 下降。记入 EXPERIMENTS.md。

• Step 6: Commit

git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md
git commit -m "perf: SMoE 稠密向量化（数值等价，消除循环/同步）"

Task 11: Embedding 池化融合（28 次 segment_reduce → 1 次）

Files:

• Modify: 代码/code/infer.py（RepEncoder.forward）

• Modify: 代码/code/tests/test_equiv.py

• Step 1: 写等价测试（先失败）

在 test_equiv.py 追加，对比融合实现与逐 slot 实现在同一输入上的输出 allclose（构造一个 28-slot 的小 batch dict，调用 infer._rep_forward_perslot(enc, batch) 参考实现 vs enc(batch)）。

def test_rep_fused_matches_perslot():
    torch.manual_seed(0)
    enc = infer.RepEncoder(vocab_size=1000, emb_dim=512, slot_num=28, d_model=512).cuda().eval()
    batch = {}
    for s in range(1, 29):
        n = torch.randint(1, 5, (10,))                      # 每样本 1~4 个 sign
        vals = torch.randint(0, 1000, (int(n.sum()),))
        offs = torch.cat([torch.zeros(1, dtype=torch.long), n.cumsum(0)])
        batch[s] = (vals.cuda(), offs.cuda())
    with torch.no_grad():
        ref = infer._rep_forward_perslot(enc, batch)
        new = enc(batch)
    assert torch.allclose(ref, new, atol=1e-4), (ref - new).abs().max()

• Step 2: 跑测试确认失败

Run：!python -m pytest test_equiv.py::test_rep_fused_matches_perslot -v Expected：FAIL（_rep_forward_perslot 未定义）。

• Step 3: 实现融合

把现有逐 slot 循环抽为 _rep_forward_perslot(enc, batch)（参考/回退）。新 RepEncoder.forward 把 28 个 slot 的 values 拼成一条，offsets 平移拼接成覆盖 28*N 段的单一 offsets，一次 segment_reduce，再 reshape [28, N, emb] → permute/cat 成 [N, 28*emb]：

def forward(self, batch):
    max_idx = self.emb.num_embeddings - 1
    target_dtype = self.input_norm.weight.dtype
    N = batch[1][1].numel() - 1                      # 样本数 = offsets 段数
    all_vals, seg_offsets, base = [], [0], 0
    for s in range(1, self.slot_num + 1):
        vals, offs = batch[s]
        if CONFIG["signid_mode"] == "modulo":
            vals = vals % self.emb.num_embeddings
        else:
            vals = vals.clamp(0, max_idx)
        all_vals.append(vals)
        seg_offsets.extend((offs[1:] + base).tolist())
        base += vals.numel()
    cat_vals = torch.cat(all_vals)
    seg = torch.tensor(seg_offsets, device=cat_vals.device, dtype=torch.long)
    emb = self.emb(cat_vals).to(target_dtype)
    pooled = torch.segment_reduce(emb, reduce="sum", offsets=seg, initial=0)  # [28*N, emb]
    pooled = pooled.view(self.slot_num, N, self.emb_dim).permute(1, 0, 2).reshape(N, -1)
    return self.linear(self.input_norm(pooled))

验证点：seg_offsets 构造正确性强依赖每个 slot 的 offsets 含开头的 0 —— 测试里务必覆盖「某样本某 slot 为空」的情况（offsets 出现连续相等）。FP16 下放宽 atol。

• Step 4: 跑测试确认通过

Run：!python -m pytest test_equiv.py::test_rep_fused_matches_perslot -v Expected：PASS。

• Step 5: bench 确认 AUC 不变、延迟下降 + Commit

import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench)
bench.run_once({})

Expected：AUC 一致，延迟下降。记入 EXPERIMENTS.md。

git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md
git commit -m "perf: RepEncoder 融合 28 次 segment_reduce 为单次"

Task 12: 确认 batch_size 控制权并（若可）扫描最优

Files:

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 判断评测端是否固定 batch_size

查 代码/任务提交接口说明.md 与 baseline notebook：评测端自建 DataLoader 时 batch_size 是否由其设定。若由评测端固定 → 我们无法在评测改 batch（跳过本任务，只在本地扫描了解趋势）。若 infer.py 的 main() 才建 loader 而评测复用我们的某入口 → 记录可控。

• Step 2: 本地扫描 batch_size 的延迟趋势

import bench
for bs in [50, 100, 200, 400]:
    bench.run_once({}, batch_size=bs)

Expected：延迟随 bs 变化曲线（注意显存）。记入 EXPERIMENTS.md，作为「若可控则用」的参考。

• Step 3: Commit

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: batch_size 控制权确认与延迟扫描"

Task 13: 重估 torch.compile / CUDA Graph（图理干净后）

Files:

• Modify: 代码/code/infer.py、代码/code/build_env.sh

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 对干净后的模型试 torch.compile

在 load_model 末尾（model.eval() 后）加可开关的：

if CONFIG.get("compile", False):
    model = torch.compile(model, mode="max-autotune", dynamic=True)

build_env.sh 加预热（按 spec §11 模板）。bench 对比开/关。

FlexAttention 与 torch.compile 通常配合良好（flex 本就鼓励 compile）；这次重估可能与上次（失败）结果不同。

• Step 2: bench 对比 + 判定

import bench
bench.run_once({"compile": False})
bench.run_once({"compile": True})

若 compile 提速且 AUC 不变 → 保留并把 compile 默认设 True；否则关掉。CUDA Graph 仅在序列长度分桶后另行评估，本任务不强求。记入 EXPERIMENTS.md。

• Step 3: Commit

git add 代码/code/infer.py 代码/code/build_env.sh 代码/code/EXPERIMENTS.md
git commit -m "exp: 图清理后重估 torch.compile"

---

阶段 C：收尾

Task 14: PCOC 校准（可选，免费零头）

Files:

• Modify: 代码/code/infer.py（输出处单调缩放）

• Modify: 代码/code/EXPERIMENTS.md

• Step 1: 在历史数据上估校准系数

用带标签的历史数据估一个对 logit 的温度/偏移 (a, b)，使 mean(sigmoid(a*logit+b)) ≈ mean(label)（只在历史上拟合，不碰测试集）。把系数写入 CONFIG（如 "calib": (a, b)），在 CTRModel.forward 输出前应用：pred_logits = a * pred_logits + b（单调，不改 AUC）。

• Step 2: bench 确认 PCOC 趋近 1、AUC 不变

import bench
bench.run_once({})

Expected：PCOC 更接近 1.0，AUC 不变。记入 EXPERIMENTS.md。

• Step 3: Commit

git add 代码/code/infer.py 代码/code/EXPERIMENTS.md
git commit -m "feat: 历史数据 PCOC 单调校准（不改 AUC）"

Task 15: 最终提交 + 保底

Files:

• 无代码改动（打包提交）

• Step 1: 全测试 + bench 总确认

%cd /home/aistudio/code/tests
!python -m pytest -v
%cd /home/aistudio/code
!python bench.py

Expected：所有等价测试 PASS；本地分为历史最高。

• Step 2: 打包并校验包内容

cd /home/aistudio/code
rm -f predict.txt
zip -y ../eval.zip infer.py requirements.txt build_env.sh
unzip -l ../eval.zip   # 确认无 dataset/、ckpt.pt、bench.py、tests/

• Step 3: 提交并记录；保留保底版本

提交 eval.zip，把验证集分数记入 EXPERIMENTS.md。若新版翻车，立即回退到已知保底（当前 58.86 对应的 commit）。

git add 代码/code/EXPERIMENTS.md
git commit -m "exp: 最终版本提交结果"
git tag best-$(date +%m%d)   # 标记当前最优，便于回退

---

自检（计划 vs spec）

• spec §4 测量闭环 → Task 1–2 ✅
• spec §5 阶段 A（sign-id/精度/expert合并/特征/上下文）→ Task 3–8 ✅
• spec §6 阶段 B（注意力/MoE/embedding/batch/compile）→ Task 9–13 ✅
• spec §7 PCOC 校准 → Task 14 ✅
• spec §8 合规与提交纪律（10次/天、保底、包校验）→ Task 8/15 ✅
• spec §9 成功标准（FP32 天花板、≥0.01 AUC 杠杆、延迟≤25s、PCOC∈[0.95,1.05]）→ Task 3/4-5/9-13/14 的关卡 ✅
• spec §10 前提验证（验证集 AUC 是否 > 0.7526）→ Task 3 Step 2 判定门 ✅

已知风险/未决（继承自 spec §10）：

• 评测端是否固定 batch_size、传哪些截断参数 —— Task 7/12 先确认，控制权不在我方则相应任务降级为「仅本地参考」。
• 核心前提（验证集 AUC 有上行空间）若被 Task 3 证伪，暂停阶段 B，回到与队友/官方答疑核对目标。