mlkem-sync/.claude/plans/keygen_plan.md

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.