Chuoer47's Blog

Appearance

04 · 手工双缓冲：状态预取状态机与输出回收 ​

横切视角。台账 also: 行：…rule.h:382-403、:212-252（primary 属 02）、 …rule.h:587-611、:526-556、:474-493（primary 属 02b）、 …_tiling.cpp:59-107（primary 属 01）、:222-251（primary 属 01b）。 本篇不重贴上述区间已有的整段引用块，只做时序重排与新增推导。

1. 结论先行 ​

  1. 这个 kernel 里没有教科书意义的双缓冲。 输入队列全部是 TQue<QuePosition::VECIN, 1>（:629-634），深度 1 写进类型系统 （02b §7 末），InitBuffer 只能减不能加。所谓"双缓冲"实际是 一次 DMA 在飞 + 一次计算的错位窗口，靠 stateInQueue_ 的 1 个槽与 stateInUb 这 1 份 UB 副本共同构成——严格说是"单缓冲 + 提前发射下一笔读"。
  2. 窗口长度等于一次 ProcessSequenceChunk 的耗时（:606 → :608）， 所以 seqLen 是流水收益的唯一放大因子：seqLen=1 时窗口只有一次 Compute， seqLen=8（MAX_MTP 上限）时窗口 8 倍。MTP 不只增加算量，还顺带养活了指令流水。
  3. 输出侧才有真正的 2 槽能力（:635-636 + :225-226）， 但 QueueAttnOutput/QueueStateOutput 的 pending 生命周期只在 单个 (head, v 分片) 内（:560-563 是函数局部），跨分片不保留。 所以深度 2 的收益上限是 seqLen-1 次重叠。
  4. host 侧的候选搜索几乎总是选深度 1：01b 篇 §5.1 已用 TQue 槽位预算证明 {2,2} 不可达、且向量 gamaK 时退化为 {1,1}。本篇 §6 给出为什么这不亏： 深度 2 换来的重叠窗口只有"一次 MTE3 搬运"，而它要多付 alignK_*vStep_*2 + vStep_*2 字节 UB。
  5. 三处同步 API 并存（SetWaitFlag 自由函数、WaitVectorToScalar 成员、 TQueSync<PIPE_S, PIPE_V> 局部对象），是流水设计逐步补丁化的化石层，见 §7。

2. 状态机的完整状态表 ​

把 ProcessHead（:587-611）与它调用的两个函数合成一台状态机。 记 S = stateInQueue_ 的槽内容、U = stateInUb 的内容、Q = S 的队列状态：

步锚点动作SUQ
0:597PrefetchState(片0) → AllocTensor + DMA + EnQue片0(half)陈旧满
1:600 → :399DeQue 取走片0空未变空
2:400CastInputToFp32(U, 片0, alignK_*alignedV)空片0(fp32)空
3:401SetWaitFlag<V_MTE2>空片0空
4:402FreeTensor(片0) → 槽归还空片0可用
5:606PrefetchState(片1) → 槽被片1 占据片1(half)片0满
6:608ProcessSequenceChunk 用 U 算片0 的所有 token片1片0' (原地更新)满
7下一轮 :600DeQue 片1 → CastInputToFp32 覆写 U空片1空

关键：步骤 5 与 6 之间是唯一的重叠窗口——MTE2 在搬片1，Vector 在算片0。 窗口能成立的前提是"片0 有 UB 副本、片1 有队列槽"两份存储同时存在， 即 stateInUb（tmpBuff 里，cubeSize 字节，:237-238）+ stateInQueue_ 槽 （alignK_*vStep_*sizeof(inType) 字节，:220）。 这就是"双缓冲没有铺开"的确切代价：多占一个 half 队列槽， 而它换到的只有一笔 DMA 的重叠。

:596/:604 两次算 nextStateOffset 是因为步骤 0 的片号是 0（偏移无 v_i 项）， 步骤 5 起才有 + nextVOffset * realK_。把两者合并成 base + v_i * realK_ 就能只留一行——代码选择写两遍， 于是 stateOffset * NV_ + head_i 这个式子在 :596 与 :604 各出现一次， 而写侧 :570 用的是第三份同构式子（槽号不同）。三份重复是本文件最容易改错的地方。

PrefetchState（:382-395）的快慢两路（realK_ == alignK_ 与否） 在 02a 篇 §6 已述：慢路不做真正的 pad，因此 stateInUb 列尾残留会进 DotFp32。从流水视角补一句： 残留来自队列槽的上一代内容（步骤 7 覆写前槽里是片0 的 half）， 所以 dk 非 16 倍数时"片 v 的列尾"读到的是"片 v-1 的列尾"—— 污染是有结构的，不是随机噪声。这使 02a 篇 §6 的风险在多次迭代下 可能表现为系统性偏差而非乱码。真机验证方法见 00 篇「未做」第 5 项。

3. 为什么衰减系数是成员而不是入参 ​

ProcessSequenceChunk 的 :572-574 把 gama_/beta_/gamaK_ 写进成员， Compute 再从成员读（:438、:442、:454）。 这让 Compute 不是纯函数，但换来的是 Compute 的签名只有 3 个参数 （:435：curSingleV、curQKOffset、curVOffset）。

从流水角度这个取舍有实质影响：三个标量走寄存器/常量通路而非 UB， :438 的 Muls 直接用标量立即数式的乘子，不需要为它腾 UB 区。 对比 gamaKInUb（:642）——它是向量衰减，必须占 UB， 于是它成为唯一"跨函数持久、来自队列而非 tmpBuff"的 LocalTensor（02b §7 第 3 点）。 同一物理量（逐维衰减）的两种表达（标量成员 / UB 向量）走了两条完全不同的存储通路， 这正是 gamaKScalar_ 这个开关的代价：它不是"性能选项"而是"存储布局选项"。 01b 篇 §4.2 记录 host 用元素数探测来决定它， 01b 篇 §5.1 记录它额外吃掉 fixedQueueBufferNum 的一个队列槽（:233 的 5U + ((hasGamaK == 1U && gamaKScalar == 0U) ? 1U : 0U)，在 :235 的槽位预算守卫里消费）—— 标量版每 vStep 省 2*MAX_MTP*alignK_ 字节 UB 与 1 个队列槽。

4. CopyInQKV 与预取的次序选择 ​

:590（CopyInQKV）在 :597（首个 PrefetchState）之前。两种可选顺序的比较：

顺序首个 token 可开算的时刻代价
QKV 先（现状）QKV DMA + Cast 完成后state 的 DMA 晚发射，与 QKV 的 DMA 串行
state 先state DMA 完成 + Cast 后QKV 可以与 state 的 Cast 重叠

现状把 q/k/v/beta/gk 全部落地后才开始搬 state。 state 是最大的一块（alignK_ * vStep_ 个 half，dk=128、vStep=16 时 4KB/片）， 把它排在最后发射意味着 MTE2 队列里最长的一笔没有与前几笔重叠。 但反过来：CopyInQKV 内部末尾有 SetWaitFlag<V_MTE2>（:376）， 它已经强制 MTE2 与 V 之间同步一次，若先发射 state 再算 QKV， stateInQueue_ 的槽会与 qLocal 等争用 UB 峰值—— InitBuffer 的分配是一次性的（:217-227），所以争用只体现在峰值占用而非时序， 而峰值由 host 的 ubRestBytes 固定（01b 篇 §5.3）。 结论：顺序选择对 UB 峰值无影响，纯影响 MTE2 流水重叠；作者选了"先小后大"， 牺牲了一次重叠换来"QKV 落地后 Compute 需要的 kInUb/vInUb 一定就绪"的简单推理。 这条判断无法从快照内证实动机（无 commit message 可查），记为报告推断。

5. 输出回收：pending 机制的确切收益 ​

QueueAttnOutput/QueueStateOutput（:532-556，全录在 02b §5）的分支： 深度 1 → 立刻 CopyOut；深度 2 → 记下、下次再 CopyOut、函数退出时刷尾（:579-584）。

在深度 1 下，每 token 的节奏是：

Compute(:575) → 两个 EnQue(:468,:471) → CopyOutAttn(:535→:474) → SetWaitFlag<MTE3_V>(:480) → FreeTensor
              → CopyOutState(:548→:484) → SetWaitFlag<MTE3_V>(:491) → FreeTensor

:480 的 SetWaitFlag<HardEvent::MTE3_V> 是屏障（02 篇 §4：Set 后立刻 Wait）， 所以 Vector 侧必须等 MTE3 搬完才能继续。它在每个 token 上出现两次 （attn 一次、state 一次），即 2 × seqLen × (v 分片数) 次全同步。

深度 2 下：第 n 个 token 只 EnQue，第 n+1 个 token 的 Queue*Output 才执行 CopyOut(pendingOffset, ...)。于是 Compute(n+1) 与 MTE3 搬运(n) 可以重叠—— 但重叠窗口被 CopyOut 之后的 :480 屏障截断（它等的是上一次的 MTE3）。 净收益 ≈ 一次 Compute 的时间。而 CopyOutState 是 curSingleV 次 DataCopy（:486-490）， curSingleV 可到 vStep_，这笔搬运不短。

所以深度 2 是划算的（省掉一次几乎全等的 Compute 等待）， 但 host 侧预算常不允许（01b 篇 §5.1）。这个矛盾是 RGDR 手写流水的真实瓶颈： 不是设计不想重叠，是 310P 的 UB 装不下第二个 alignK_ × vStep_ 输出缓冲。:225 的单槽成本 alignK_*vStep_*sizeof(outType) = 128×16×2 = 4KB（dk=128）， stateOutBufferNum 从 1 到 2 就是多 4KB——在 01b 篇 §5.3 的 ubRestBytes 余量下，vStep 大时给不出这 4KB。

CopyOutAttn（:474-482）与 CopyOutState（:484-493）的 DeQue → DataCopy → barrier → FreeTensor 四步同构。不同点只有一处：attn 一次 DataCopy（:477），state 逐 v 行循环（:486-490）。 差别来源：attn 的输出是 curSingleV 个连续元素（[T, NV, DV] 里固定 (seq,head) 的一段）， state 的输出是 curSingleV 条行（每行 realK_，行距 realK_，UB 侧行距 alignK_）—— UB/GM 行距不等时无法一次搬完（02 篇 §5.1 的 gmStepPerRow 问题在写侧的再现）。

6. 与 host 候选搜索的闭环 ​

01b 篇 §5 从算术侧推出 {2,2} 不可达；本篇从代码侧补上为什么搜索空间只设计到 2：

  1. TQue<..., MAX_OUT_BUFFER_NUM>（:635-636）的模板第二参是 2， 写 3 会越界；MAX_OUT_BUFFER_NUM = 2 定义在 :25。 类型层把上限钉死，host 的候选表 {1,1},{1,2},{2,2} 正是它的子集枚举。
  2. :173-174 的再钳制（02 篇 §8.3）保证即使 host 写错也不会突破 2。
  3. 输入侧没有候选：BUFFER_NUM = 1（:24）被 :217-226 的 6 个输入队列共用， 且 TQue<VECIN, 1> 的模板参数也是 1。 所以 host 的搜索只针对输出，不针对输入——01b 篇 §5.1 的候选表只有 (stateOutBufferNum, attnOutBufferNum) 两个维度，与此对应。

EvaluateBufferProfile（tiling :73-97，primary 属 01）里 vStep 的两步求解 （先按系数取整、再 CeilDiv(dv, repeatTime) 回填）与本篇的窗口分析直接相关： repeatTime = CeilDiv(dv, vStep) 就是 v 分片数， 分片数越多 → 每片越小 → Compute 越快但 PrefetchState 次数越多。 vStep 必须 16 对齐（01b §5.2）的第二个理由在本篇 §2 步骤 2： CastInputToFp32 的长度是 alignK_ * alignedV， alignedV = Ceil(curSingleV,16)*16（:398），若 curSingleV 本身是 16 的倍数则 alignedV == curSingleV，Cast 无浪费。非对齐时 Cast 多算 alignK_ * (alignedV - curSingleV) 个元素——纯浪费的 Vector 周期。 host 选 16 对齐的 vStep 同时省了 UB 尾部和 Vector 空转， 但最后一篇（dv % vStep != 0 的尾片）无法对齐，浪费不可避免。

7. 三种同步 API 的化石层 ​

API定义事件/管道调用点数用途
SetWaitFlag<HardEvent>（自由函数模板）:34-39任意18（实测 grep -c）通用屏障
WaitVectorToScalar()（成员）:254-258固定 V_S2（:414、:523）Vector 写 → Scalar 读
TQueSync<PIPE_S, PIPE_V>（局部对象）无定义（AscendC 提供）S→V1（:463-465）Scalar 写 → Vector 读

三条方向互不重复：任意、V→S、S→V。所以并存不是冗余， 而是 TQueSync 只支持"队列式同步对"、SetFlag/WaitFlag 需要手取 event id 的能力差异。 但 SetWaitFlag 已经能表达 V_S（02 篇 §4）， :254-258 的成员版确实是重复实现——它的唯一附加价值是用 pipe_ 成员而非 GetTPipePtr() 全局。两处用法在 Init 之前/之后的可用性上没有差别 （pipe_ 在 :190 赋值，:34-39 的调用点全在 Init 之后）。 判定：可收敛为一处，属未清理的实现重复，非缺陷。

06 篇会核对：TQueSync 那 3 行（:463-465）是否为两快照之间新增 ——若新增，则化石层是"补丁逐层累加"的直接证据。

8. 量化：一次 Process 的指令与搬运次数 ​

以 b=1、NV=32、NK=8、dk=128、dv=128、seqLen=1（decode 常态）、 vStep=16（host 大概率给值，01b 篇 §5.2）为例，逐式给出计数（不是估算）：

  • taskUnits = b*NV = 32（:176），workBlockDim_ = min(32, coreNum)；310P 的 AIV 数 在本快照内查不到（GetCoreNumAiv 是 runtime 接口，01b 篇 §3.1）， 所以按 coreNum ≥ 32 与 < 32 两档讨论。
  • 每核 head 数：NV / workBlockDim_（整除时）。
  • 每 head：v 分片数 = ceil(128/16) = 8（tiling.cpp:86 的 repeatTime 同值）。
  • 每分片：CopyInQKV 0 次（它在分片循环之外，:590）+ PrefetchState 1 次 + LoadPrefetchedState 1 次 + ProcessSequenceChunk 1 次。
  • 每 chunk（seqLen=1）：Compute 1 次 = Muls 至多 2 次 + curSingleV(=16) 次 DotFp32（每次 1 Mul + 1 WholeReduceSum + 2 标量加， dk=128 时 fullRepeats=2，02a 篇 §7 第 3 点）+ 16 次 Muls+Add 对 + 16 次 DotFp32 + 2 次 CastFp32ToOutput。 合计每分片 32 次 DotFp32（memory 一遍、out 一遍）。
  • 每 head 的 GM 写入：8 × (1 + 128/16 × 8?)…——写侧 attn 每分片 1 次 DataCopy， state 每分片 curSingleV=16 次 DataCopy（:486）⇒ 每 head 8×16 = 128 次 state DMA
    • 8 次 attn DMA。每核每 head 136 笔 MTE3，其中 128 笔的 blockLen 只有 realK_*2 = 256 字节（:487）。

这张账说明两件事：（a）CopyOutState 的逐行循环是 MTE3 事务数的主导者 （128:8 = 16:1）；（b）dk 非 16 倍数时那 128 笔每笔还要被 AlignUp 到 32 的倍数 （02 篇 §5.2），事务数不变但每笔多拷。 若把 :486-490 的循环换成 DataCopyExtParams 带 dstStride 的批量写， 事务数可从 128 降到 8——这是 00 篇「未做」第 5 项之外唯一能纯代码内指出的优化方向， 本篇只登记、不实施（目标 §11：不顺手优化）。