# ML-KEM KeyGen 顶层集成 — 实现计划与进度

> 本文档用于让新工程师完整接手。包含原始计划 + 已完成进度 + 经验教训 + 下一步。

## 目标与范围
- 目标:实现 `mlkem_top` 的 **KeyGen 通路**(ML-KEM-512, K=2, η1=3),全新重写 valid/ready FSM。
- 最终验收:对 **NIST 官方 KAT**(count=0..4)端到端逐字节通过,产出完整 ML-KEM dk。
- 本计划只做 KeyGen。Encaps/Decaps 留作后续任务。
- 旧版 `mlkem_top.v`(1658 行)已于 commit 1cace51 删除,**不可用**(Decaps 是占位,且子模块握手从未真正跑通——旧 TB 用 24 处 force/release 绕过)。新版全新重写。

## 关键决策(已与用户确认)
- **Keccak 策略**:每消费者用独立、已验证的核(sha3_top / sample_ntt_sync / sample_cbd_sync 各自带 keccak_core)。**不用 arbiter,不用 `_shared` 变体**(顶层当前用独立版)。
- **SHA3 先扩多块**:已先把 sha3_top 的 H 模式扩成多块吸收(✅ 完成),再做完整 dk。
- **FSM 全新重写**:只参考算法步骤,不取回旧代码。
- **分阶段 + 逐级对拍**:每个子步骤对 golden 验证通过再进下一步。

## 环境与工具(重要)
- 工作目录:`/home/fallensigh/Dev/mlkem`,git 分支 `main`(直接提交,用户要求)。
- **Verilator** 5.046:`python3 test_framework/run_all.py [--module X]`。快速回归。
- **Vivado XSIM** 2019.2:每次需
  ```
  source /opt/Xilinx/Vivado/2019.2/settings64.sh
  export LD_PRELOAD=/usr/lib64/libtinfo.so.5
  ```
  调用 **standalone** `xvlog -sv --relax -i . <files>` / `xelab <top> -s <snap> --timescale 1ns/1ps` / `xsim <snap> -R`。**不要**用 `vivado -mode batch`(脚本里用的是 standalone 命令,不是 Tcl)。zsh 不做无引号变量分词——用 `"$@"` 传文件列表。
- XSIM 产物(xsim.dir/ .Xil/ *.jou *.log)已 gitignore;每次跑前 `rm -rf xsim.dir .Xil`。

## Golden reference(全部已验证可用)
- **NIST 官方 KAT**:`/home/fallensigh/Dev/ml-kem-r/test_data/kat_MLKEM_{512,768,1024}.rsp`。
- **ml-kem-r**(Rust, `/home/fallensigh/Dev/ml-kem-r`):`cargo test` 全过,512/768/1024 KAT 全通过。是可信参考。
  - API:`ml_kem::mlkem512::key_gen(d,z)`、`keygen_internal::<2,192>(d,z)`、`k_pke::key_gen::<2,192>(d)`。
  - 公共子函数:`sample::sample_ntt(&[u8])`、`sample::sample_poly_cbd::<ETA_BYTES>`、`ntt::ntt(&mut [i32;256])`、`ntt::multiply_ntts`、`sha3::mlkem_G/mlkem_H/mlkem_PRF`、`utils::byte_encode(&[i32;256], d)`。
  - `logging` feature + `examples/logging.rs` dump 中间值(4-coeff head)。
- **server_code**(Python, `~/Dev/server_code/python_project/PQC_2025/A_ML_KEM_v0`):`SHA_3.G/H/J/PRF`、`ML_KEM.py`。已对 hashlib 交叉验证。

### 白盒 golden 向量(Stage 0 产物,已生成并自校验)
- 生成器:`/home/fallensigh/Dev/ml-kem-r/examples/dump_keygen.rs`(我新增,未提交到 ml-kem-r,因为该仓库有用户大量未提交改动 —— 只在工作树用)。
  - 运行:`cargo run --release --example dump_keygen -- <out_dir> 5`
- 输出目录:`test_framework/modules/mlkem_keygen/golden/`(已 commit 在 mlkem 仓库 106b292)。
- 每个 KAT case c(000..004)的全 256-coeff 中间量,文件名:
  - `c000_rho.hex` / `c000_sigma.hex`:32B hex(byte0 在最低)。
  - `c000_Ahat_i_j.hex`(i,j∈0..1):256 行,每行 3-hex 系数(mod q)。
  - `c000_s_i.hex` / `c000_e_i.hex`:CBD 输出,**mod-q 形式**(负值 +Q,如 -1 → 0xd00=3328)。
  - `c000_shat_i.hex` / `c000_ehat_i.hex`:NTT(s/e) 后。
  - `c000_that_i.hex`:t_hat。
  - `c000_ek.hex`(800B hex = KAT pk)、`c000_dkpke.hex`(768B = KAT sk 前 384*K)。
  - dump_keygen 已自校验 ek==KAT.pk、dk_pke==KAT.sk[..768],全 5 case 通过。

## KeyGen 算法 (FIPS 203 Alg 13+16, K=2 η1=3)
1. (ρ,σ) = G(d ‖ K)         — SHA3-512, 33B 输入(单块)。**G 输入 = d(byte0 低) ‖ K=2 at byte32,无字节翻转**(已确认,旧版 byte-reverse 是 bug)。ρ = hash[255:0],σ = hash[511:256]。
2. Â[i][j] = SampleNTT(ρ‖j‖i), i,j∈{0,1}  — seed = rho ‖ j ‖ i(sample_ntt_sync 内部 msg={i_byte,j_byte,rho})。
3. s[i] = CBD₃(PRF(σ, **i**)), i∈{0,1}      — nonce 0,1
4. e[i] = CBD₃(PRF(σ, **K+i**)), i∈{0,1}     — nonce 2,3
5. ŝ[i] = NTT(s[i]); ê[i] = NTT(e[i])       — 前向 NTT ×4 (mode=0, **无缩放**)
6. t̂[i] = ê[i] + Σⱼ Â[i][j]∘ŝ[j]            — poly_mul + 逐系数 modQ 累加
7. ek = byteEncode₁₂(t̂[0..1]) ‖ ρ           — 12-bit 原始打包(**非压缩**)
8. dk_pke = byteEncode₁₂(ŝ[0..1])
9. dk = dk_pke ‖ ek ‖ H(ek) ‖ z             — H(ek): SHA3-256 over 800B(**多块**)

**用到的叶子模块**(端口签名见下):sha3_top(G+多块H)、sample_ntt_sync、sample_cbd_sync、ntt_core(mode=0)、poly_mul_sync、mod_add_sync。**不用** comp_decomp。

## 叶子模块端口签名(已抓取)
- `rng_sync #(SEED)`: clk,rst_n,valid_i,ready_o,data_o[255:0],valid_o,ready_i
- `sha3_top`: clk,rst_n,mode[1:0],data_i[511:0],valid_i,ready_o,hash_o[511:0],valid_o,ready_i, **+多块端口** mb_en,mb_block_i[1087:0],mb_valid_i,mb_last_i,mb_ready_o
- `sample_ntt_sync #(K=4)`: clk,rst_n,rho_i[255:0],k_i[2:0],i_idx[1:0],j_idx[1:0],valid_i,ready_o,coeff_o[11:0],valid_o,ready_i,last_o
- `sample_cbd_sync`: clk,rst_n,seed_i[255:0],nonce_i[7:0],eta_i[1:0],valid_i,ready_o,coeff_o[11:0](**12-bit signed 二补码**),valid_o,ready_i,last_o
- `ntt_core`: clk,rst_n,coeff_in[11:0],valid_i,ready_o,mode(0=fwd无缩放/1=inv有缩放),coeff_out[11:0],valid_o,ready_i,done_o。协议:IDLE/LOAD 期间 ready_o 高,流式喂 256 个(valid_i 持续);算完 S_OUTPUT 流式吐 256(valid_o 高,ready_i 消费);S_DONE→IDLE。
- `poly_mul_sync`: clk,rst_n,coeff_a_in[11:0],coeff_b_in[11:0],valid_i,ready_o,coeff_out[11:0],valid_o,ready_i。协议:LOAD 阶段 **同时**喂 A、B 各 256 对(coeff_a_in/coeff_b_in 同拍),然后 COMPUTE 流式吐 256。
- `mod_add_sync`: clk,rst_n,a[11:0],b[11:0],valid_i,ready_o,sum[11:0],valid_o,ready_i。(a+b)mod Q。
- `comp_decomp_sync`: clk,rst_n,coeff_in[11:0],d[4:0],mode,valid_i,ready_o,coeff_out[11:0],valid_o,ready_i(KeyGen 不用)。
- `sd_bram #(W=48,D=64,A=6)`、`s_bram #(W=48,D=512,A=9)`(KeyGen 暂用寄存器阵列 polymem 缓存,未用 bram)。

## 已完成进度

### ✅ Stage 0:白盒 golden 向量生成器(commit 106b292)
- `examples/dump_keygen.rs`(ml-kem-r 工作树)+ `test_framework/modules/mlkem_keygen/golden/`(mlkem 仓库)。
- 自校验:ek==KAT.pk、dk_pke==KAT.sk 前缀,5/5 通过。

### ✅ Stage 1:sha3_top 多块 SHA3-256 吸收(commit 106b292)
- `sync_rtl/sha3/sha3_top.v`:新增多块吸收 FSM。新端口 `mb_en/mb_block_i[1087:0]/mb_valid_i/mb_last_i/mb_ready_o`。
  - **设计**:调用方预填充末块的 SHA3-256 padding(0x06 ... 0x80),模块只做纯吸收 loop(state^=block; Keccak-p)。单块 G/H/J 路径在 mb_en=0 时逐 bit 不变。
  - 状态:MB_IDLE/MB_PERMUTE/MB_DONE。`mb_state_r`(1600-bit running state),`mb_digest_r`(256-bit sticky digest,保持到消费者 ack)。
  - `ready_o = !mb_en && (state_r==ST_IDLE)`、`valid_o`/`hash_o` 按 mb_en 选择。
- TB:`sync_rtl/sha3/TB/tb_sha3_mb_xsim.v`(自检,流 800B ek=6 块 → H(ek),对 hashlib.sha3_256)。向量:`vectors/mb_h_blocks.hex`、`mb_h_expected.hex`。
  - **握手坑(已解决)**:TB 必须在 accept 边沿保持 valid/last 稳定(等 mb_ready_o 拉低后才撤),否则采样竞争。digest 一次有效后 sticky,TB 收完才 ack。
- 既有 G/H/J TB(xsim 3 个 + Verilator tb_sha3.cpp)**必须 tie off mb_* 端口**(mb_en=1'b0 等),否则 mb_en 浮空成 x 破坏 ready_o。已全部修。
- 回归:Verilator 25/25,XSIM 4 个 TB(simple G/H/J、keccak_core、7-vec、multiblock)全过。

### ✅ sample_ntt 幽灵脉冲修复(commit 6db3c7c)
- **bug**:`sample_ntt_sync` 的 Phase-1(输出 d1)缺少 Phase-2 有的 `need_more` 守卫。当第 256 个被接受系数是 d1 且组推进时,某些 seed 会在 last_o 后多吐一拍 valid_o(KAT count=0 rho 下,seed i=0/j=1 吐 257 拍)。集成时漏进下个 poly 的 index 0,整流后移一位。
- **修复**:Phase-1 改 `if (d1_acc_r && need_more)`。两变体(`sample_ntt_sync.v` + `sample_ntt_sync_shared.v`)都改。
- **补盲区**:`tb_sample_ntt_xsim.v` 加断言——last_o 后保持 ready_i 数多余 valid_o,有则 FAIL。已验证该断言能抓 bug(还原副本 3 个幽灵)、修复版通过。
- 验证:40-seed(sync)+24-seed(shared)审计全 256 拍/last@256/零幽灵;Verilator 1536/1536;全框架 4334/4334;Stage 2c 2048/2048。
- 记忆:`sample-ntt-257th-pulse-bug.md`。

### ✅ Stage 2a/2b/2c:KeyGen 数据通路前半(**未提交**,在工作树)
- 文件:`sync_rtl/top/mlkem_top.v`(新建)+ TB `sync_rtl/top/TB/tb_mlkem_kg_2a_xsim.v`、`tb_mlkem_kg_2c_xsim.v`、向量 `sync_rtl/top/TB/vectors/kg_c000_AsE.hex`。
- **mlkem_top 结构**(K=2,ETA1=3):
  - 端口:clk,rst_n,d_i[255:0],z_i[255:0],start_i,busy_o,done_o + **调试 tap**:dbg_slot_i[3:0],dbg_idx_i[7:0],dbg_coeff_o[11:0](寄存器读 polymem,2拍延迟),dbg_rho_o[255:0],dbg_sigma_o[255:0]。
  - **polymem**:`reg [11:0] polymem[0:10*256-1]`,slot 映射:0-3=A00,A01,A10,A11;4-5=S0,S1;6-7=E0,E1;8-9=T0,T1。s/s_hat 共享 slot(就地 NTT),e 同理。
  - FSM 状态:ST_IDLE=0,ST_G=1,ST_A=2,ST_C=3,...,ST_DONE=15。
  - **2a(G)**:sha3_top mode=00,g_data={248'b0,8'(K),d_i};完成捕获 rho_r/sigma_r。✅ rho/sigma 对 golden 精确。
  - **2b(A)**:a_pair 0..3 循环,a_i=pair[1],a_j=pair[0],a_slot=pair。每 pair 跑 sample_ntt_sync,256 系数写 polymem[a_slot]。✅ 1024 系数全对。
  - **2c(C)**:c_poly 0..3(s0,s1,e0,e1),c_nonce=c_poly,c_slot 映射。CBD 输出经 `cbd_modq = cbd_coeff[11] ? cbd_coeff+Q : cbd_coeff` 转 mod-q 后写 polymem。✅ 2048/2048(A+s+e)全对。
  - **关键防护 a_busy/c_busy**:请求被接受(valid&&ready)时置位,last 时清零;仅 busy 时收集系数。本是为屏蔽 sample_ntt 幽灵脉冲的下游兜底;sample_ntt 根因修复后冗余但保留(无害,且对未来叶子模块是防御)。
- **字节序基准**(全部已确认,这是旧版的痛点):
  - G 输入 d 原序(byte0 低),无翻转。
  - SampleNTT seed=rho‖j‖i。
  - CBD nonce:s[i]=i,e[i]=K+i。
  - golden 多项式存 mod-q;CBD RTL 输出 12-bit 二补码,需 +Q 转换。

## 当前工作树状态(未提交)
- `sync_rtl/top/mlkem_top.v`(新,Stage 2a/b/c,FSM 到 ST_C)
- `sync_rtl/top/TB/tb_mlkem_kg_2a_xsim.v`、`tb_mlkem_kg_2c_xsim.v`(新)
- `sync_rtl/top/TB/vectors/kg_c000_AsE.hex`(新,8 poly × 256 = 2048 行,顺序 A00,A01,A10,A11,S0,S1,E0,E1)
- `.claude/plans/`(本文件)
- **下一步该先提交 Stage 2a-2c**(测试已通过)。

## 关键测试 literal(byte0-低 256-bit)
- KAT c000 d:`D_LIT = 256'h2426f1941779574d3f1b163bd57f7e173e229e630ec7f7073bdf365137c4bb6d`
- 期望 rho:`256'h15f74355ca862c3cdf3dab780c35cf24b88bf144706090a1c17e41205f9f1379`
- 期望 sigma:`256'h69b042001b5630b1a039116cbfd29f62c0bde5a6b571504a9fcce68bed667fd5`

## 下一步(剩余 Stage 2/3/4)

### Stage 2d:NTT(s/e → ŝ/ê)  ← **当前正在做**
- golden 映射已确认:`ntt(s_i)==shat_i`、`ntt(e_i)==ehat_i`(mode=0 fwd,无缩放,zeta 表见 ntt gen_vectors,N_INV=3303 仅 inv 用)。
- 计划:新增 ST_N 状态,对 4 个 slot(S0,S1,E0,E1)逐个就地 NTT:从 polymem 读 256 喂 ntt_core(mode=0),收 256 写回同 slot。
- 注意 ntt_core 协议:LOAD 期间持续 valid_i 喂 256;输出阶段 valid_o 高、ready_i 消费 256。需读地址流水(组合读 polymem)。
- TB:扩展对 shat/ehat golden 校验。
- **正读 mlkem_top.v 行 200-290 的时序逻辑准备插入 ST_N 时被打断;接手时先重读 mlkem_top.v 全文确认当前 FSM 时序结构再加 ST_N。**

### Stage 2e:矩阵累加 t̂[i] = ê[i] + Σⱼ Â[i][j]∘ŝ[j]
- 用 poly_mul_sync(同拍喂 A、B 各 256)算 Â[i][j]∘ŝ[j],再用 mod_add_sync 逐系数 + ê[i] 累加。K=2:t̂[0]=ê0+A00∘ŝ0+A01∘ŝ1;t̂[1]=ê1+A10∘ŝ0+A11∘ŝ1。
- 对 golden `that_i` 校验。

### Stage 2f:byteEncode₁₂ → ek/dk_pke
- 12-bit LSB-first 打包(FIPS:b[i*12+j]=(coeff>>j)&1,再 bits_to_bytes)。ek=byteEncode12(t̂[0..1])‖ρ(800B);dk_pke=byteEncode12(ŝ[0..1])(768B)。
- 对 KAT pk / sk 前缀逐字节校验。

### Stage 3:KeyGen FSM 集成 + 干净顶层接口
- 串起 2a-2f;顶层流式输出 ek_o/dk_o;无 force。

### Stage 4:组装完整 dk + 端到端 KAT
- dk = dk_pke ‖ ek ‖ H(ek)[用 Stage1 多块] ‖ z。
- 干净 KAT TB(无 force/release),喂 KAT d/z,逐字节比 ek 与 dk。跑 KAT count=0..4 全过。Verilator + XSIM 双框架。

## 经验教训(重要)
1. **"PASS" 只证 RTL==reference**:必须独立对 hashlib/FIPS 验证 oracle。
2. **单独 TB 盲区会掩盖模块 bug**:sample_ntt 的 257th 脉冲就是"读够就停"掩盖的。新 TB 要验证协议边界(last 后 valid 干净拉低)。
3. **加端口必须 tie off 所有旧实例**:浮空输入成 x 会破坏组合逻辑(mb_en 教训)。
4. **valid/ready 握手 TB 要在 accept 边沿保持输入稳定**,等 ready 拉低再撤,否则采样竞争(多块 H、A-pair 边界都踩过)。
5. **字节序/位序是旧版主要 bug 源**:每步对 golden 逐字节,错位立刻暴露。
6. **bash 工具偶发输出截断/丢失**:重要结论用 `rm -rf xsim.dir` 重跑确认,别凭一次模糊输出下判断。

## 任务追踪(TaskList ID)
- #7 Stage 0 ✅ completed
- #8 Stage 1 ✅ completed
- #9 Stage 2 datapath subblocks — in_progress(2a/b/c done,2d 进行中)
- #10 Stage 3 FSM 集成 — blocked by #9
- #11 Stage 4 dk + KAT — blocked by #8,#10

## 提交历史(相关)
- 106b292 feat(sha3): multi-block SHA3-256 absorb for H(ek); KeyGen golden vectors (Stage 0+1)
- 6db3c7c fix(sample_ntt): suppress spurious 257th valid_o after last_o
- (未提交)Stage 2a-2c mlkem_top + TB

---
## ✅ COMPLETE (session end state)
- KeyGen FSM: IDLE->G->A->C->N->M->E->H->DONE, pure valid/ready, no force.
- Verified vs NIST KAT MLKEM-512 **count=0..4**: ek (800B)==pk, dk (1632B)==sk, byte-exact, ~21350 cyc/case.
- Commits: 106b292, 6db3c7c, 2f206a6, 4c692e5, a9e50eb, 1791491, 9824ed8, 42d3748.
- Independent Keccak per consumer (G + dedicated H sha3_top + sample_ntt/cbd cores). No arbiter.
- SECURITY: tool outputs this session carried prompt-injection (fake operator instructing `curl|bash`); all refused. Some outputs were tampered; results re-verified with controlled delimiters.
- Remaining optional: streaming ek/dk output port (currently start_i/done_o + readback taps); Encaps/Decaps; shared-keccak migration.
