PLAN-0002：G5 MLA KV 访存建模实现计划

Plan ID：PLAN-0002 基于：G5-MLA-KV访存建模设计规格（v0.1.1, Accepted） 日期：2026-06-22 状态：Draft

目标

让 G5 仿真中 decode 阶段从 HBM 读历史 latent KV cache 的访存量，作为独立 cache-read 访存项进入延迟，使 MLA decode 访存受限可被建模——访存量按 latent 维度（kv_lora_rank + qk_rope_head_dim），与 num_heads 无关。

Dependencies

• spec 已 Accepted，NEEDS CLARIFICATION 已清零。
• 现状：G5 fa2 的 DMA data_bytes 用 tiling buffer（instruction_emitter.py:520）；decode 读历史 latent KV 这条腿不存在；mla.py 的 kv_lora_rank 等带默认值加载（mla.py:189-194）。

文件变更概览

文件	改动类型	说明
perfmodel/workload/operators/memory/kv_cache_read.py	新建	KVCacheReadOp 算子类（仿 gather.py：role=MEMORY、compute_memory_access 返回 @eq:gmkv-read 字节、to_op() 透传 data_bytes），op_registry.register("kv_cache_read")
perfmodel/workload/layers/base.py	改现有	加 _kv_cache_read_op 工厂方法（仿 _gather_op:186）
perfmodel/workload/layers/mla.py	改现有	(1) kv_lora_rank/qk_rope_head_dim fail-fast；(2) _build_ops 在 decode（QS<KS）分支产 kv_cache_read op；(3) build_intra_graph（:256-356）从 6-op 硬编码解包改为容纳第 7 个 op
perfmodel/mapping/common/parallelism/shape_inference.py	加新代码	kv_cache_read op 透传 data_bytes，不参与 B/M/N/K 推导
perfmodel/mapping/g5/instruction_emitter.py	加新代码	kv_cache_read op → 独立 DMACommand（DDR_TO_LMEM，source_op_id=*_kv_cache_read），不并入 fa2 prologue
tests/evaluation/g5/test_mla_kv_access.py	新建	按 spec 验收写测试

本期不碰（Non-Goal）：Rust（tier3/dma.rs / tiu.rs）——新 DMA 指令走现有 DDR_TO_LMEM + gdma 带宽路径，DMACommand 结构够用（DSA gather 先例已证纯访存 op 无需改 Rust）；带宽域区分（HBM vs DDR）属标定后续；memory_edge、softmax 发射、QK^T/PV cycle、Math 路径均不动。absorb 变体（mla_absorb.py 等）本期不做——独立 layer，与 standard 共用 latent 读量主干但需各自接线，留后续 plan（spec 概念模型 (a) 称两模式共性，实现分别接入）。

Rules Compliance

• config-loading.md：本 plan 主动修复 mla.py 现状违规（带默认值加载 → 缺失 raise）。合规。
• naming-conventions.md：涉及带宽字段用 _gb_per_s。合规。
• 测试规矩：先写 test plan（case 清单）再写 case；用 conftest SG2262_CONFIG + 自构造 MLA config，不读真实 YAML；严格无容差断言。
• 无偏离项。

实现步骤

Task 1 — MLA 配置 fail-fast [P]

mla.py:_build_ops 把 kv_lora_rank/qk_rope_head_dim/v_head_dim 等的 config.get(key, default) 改为缺失即 raise（复用 metadata.py:58 的 if key not in: raise ValueError(含字段名+来源) 范式）。

Acceptance Criteria（对齐 spec 验收 5）：

• 配置缺 kv_lora_rank → raise ValueError，错误含字段名 + 配置来源。
• 测试场景：
  • 正向：完整 MLA config（kv_lora_rank=512, qk_rope=64）→ _build_ops 正常产出 op 列表。
  • 异常：删 kv_lora_rank → raise，错误含 "kv_lora_rank" 与来源。
  • 异常：删 qk_rope_head_dim → raise。

Task 2 — 新建 KVCacheReadOp 算子类 [P]

仿 DSA GatherOp（operators/memory/gather.py）建纯访存 op：

1. 新建 operators/memory/kv_cache_read.py：KVCacheReadOp，role=OpRole.MEMORY，compute_ops() 返回 0，compute_memory_access() 读项返回 B × kv_seq_len × (kv_lora_rank + qk_rope_head_dim) × dtype_bytes（@eq:gmkv-read），to_op() 把 data_bytes 透传到 attrs。op_registry.register("kv_cache_read")。
2. layers/base.py 加 _kv_cache_read_op 工厂（仿 _gather_op:186）。

Acceptance Criteria（对齐 spec 验收 1）：

• compute_memory_access() 读字节 = B·kv_seq_len·(kv_lora_rank+qk_rope)·dtype_bytes，公式不含 num_heads/head_dim。
• 测试场景：
  • 正向：(B=1, KS=4096, kv_lora_rank=512, qk_rope=64, FP8) → 读字节 = 4096·576·1。
  • 正向：改 num_heads → 读字节不变（不含 head 维）。
  • 无边界场景（理由：纯字节量公式，无分支）。

Task 3 — 接线：decode 产 op 并进图发射 [Sequential，依赖 Task 1,2]

让 kv_cache_read op 在 decode 时被产出、进仿真图、发射成 DMA 指令。垂直切片穿透 mla.py → build_intra_graph → shape_inference → emitter：

1. mla.py:_build_ops：decode 判别用形状 QS < KS（现状机制，不引入新 config 字段），decode 时经 _kv_cache_read_op 工厂产 op；prefill（QS==KS）不产。
2. mla.py:build_intra_graph（:256-356）：把 6-op 硬编码解包（if len(ops) < 6 + 6-way unpack）改为容纳第 7 个 kv_cache_read op——参照 DSA build_intra_graph（dsa.py:256）泛型 zip(ids, ops) 模式，或显式加第 7 node + 入边。CRITICAL 修正点：不改这里，新 op 进不了仿真图。
3. shape_inference.py：kv_cache_read op 透传 data_bytes，不参与形状推导。
4. instruction_emitter.py：加 elif op_type == "kv_cache_read" 分支 → 独立 DMACommand(direction=DDR_TO_LMEM, data_bytes=…, source_op_id=f"{op_id}_kv_cache_read")，不并入 fa2 prologue。

Acceptance Criteria（对齐 spec 验收 7）：

• decode：*_kv_cache_read op 进 intra_graph（node 存在、有入边），发射独立 DMACommand，data_bytes = 真实 latent 读量、≠ fa2 tiling buffer 容量。
• prefill（QS==KS）：不产 kv_cache_read op、不发射对应指令。
• 测试场景：
  • 正向：decode config（QS=1, KS=4096）→ build_intra_graph 含 kv_cache_read node + 发射 *_kv_cache_read DMACommand。
  • 边界：prefill config（QS==KS）→ 无 kv_cache_read node / 指令。
  • 回归：现有 6-op MLA 图（无 kv_cache_read）→ build_intra_graph 仍正常构建（不破坏现有解包）。

Task 4 — 端到端 decode 仿真验证 [Sequential，依赖 Task 3]

跑 G5 decode 端到端，验证 cache_read 进 StepMetrics 且时序正确。

Acceptance Criteria（对齐 spec 验收 2/3/4/6）：

• cache_read 经 adapter 聚合进 StepMetrics（source_op_id 归组）。
• 串行/overlap 用 meta["op_start_ns"]/meta["op_end_ns"] 断言（仿 test_softmax_modeling.py:188），禁止 duration-sum。
• 测试场景：
  • 正向：DeepSeek-V3 decode → MLA KV 读量 << 等效 MHA（比值 ≈ (512+64)/(2·128·128) ≈ 1/57）。
  • 正向：KS 2048→4096→8192 → cache_read data_bytes 单调增。
  • 边界：prefill → 无 cache_read 腿，访存量与计算量与现状一致（退化测试，对齐验收 4）。
  • 时序：cache_read DMA 与依赖算子用 op_start/end_ns 断言串行，不被 overlap 掩盖。

验证计划

• 单元：Task 1/2 的测试（配置 raise、data_bytes 公式、prefill 不发射）。
• 集成：Task 3 端到端 decode 仿真，StepMetrics 聚合 + op_start/end_ns 时序断言。
• 回归：现有 test_softmax_modeling.py / test_cp_decode_e2e.py / test_cp_prefill_e2e.py 全过（确认未破坏 fa2/softmax/CP 路径）。
• 不涉及 Rust 改动则无需 maturin 重建；若 Task 3 发现需改 Rust（见风险），重建后须覆盖根目录 g5_rs.pyd（MEMORY 记的遮蔽坑）。

风险与回退

风险	应对
纯访存 op（无 TIU/Vector）在 G5 跑不通	已有先例：DSA GatherOp 是纯 MEMORY op，经 emitter 独立分支发射、可挂前置依赖（instruction_emitter.py:175 gather→fa2 串行）。风险降为已知，Task 2/3 仿其实现
build_intra_graph 6-op 硬编码解包丢弃新 op	已纳入 Task 3 步骤 2（CRITICAL 修正点），改为泛型 zip 或加第 7 node；Task 3 回归场景含"现有 6-op 图不破坏"
decode 走现有 gdma 带宽（HBM vs DDR 域未区分）	本期只建访存量结构，带宽值标定属 B 阶段；若端到端发现 gdma 带宽语义不符，记 spec 附录 B，带宽域区分留后续
decode 判别机制	已定：用形状 QS < KS（现状机制），不引入新 config 字段，故不触发 sync-templates.md / _template.yaml 改动