06 · 主循环与系数批量化(chunk_gated_delta_rule.h:686-1020)
结论先行:这 335 行里一条 Mmad 都没有(970 行只是调用点),它是 Cube 的供料线——被 Mmad 消费的每个 FP32 中间量(kCumdecayFp32、chunkKFp32、decayMaskFp32、缓存进 chunkAttnOutFp32 的 Q·scale)都在这里成型。!1135 在这一段做的不是"换算法",而是把 !1108 用标量循环堆出来的系数装载、前缀和、行缩放、decay 构造四件事改成整块 DMA + Gather + Brcb + 高维指令,发射次数从约 260 条降到约 20 条,标量 GM 访存从 128 次降到 0 次。
本段同时埋了两处后来被抓出来的雷:827 行的原地加宽 Cast(PR !1158 的一行修复点)、968/1192 行手写的快路径守卫与 257 行 IsCubeFastPath 的判定不一致(本章自查)。
0. 基本概念
本段用到的四个机器语义概念,后面反复引用:
- 高维(三阶)向量指令:
Op(dst, src0, src1, num, repeatTime, RepeatParams)一次发射描述"repeatTime 轮、每轮 num 个元素(fp32 一轮最多 64 个 = 8 个 32B 块)",各操作数按参数里的间隔在轮间/块间推进,单位是 32B 块。它把"每行发一条指令"的 N 次发射压成 1 次,是 310P Vector 管最廉价的批量手段。快照里没有 AscendC 头文件,本章对各处间隔参数一律按"净效果"解读(可由缓冲区布局反推验证),不逐字段套官方字段名。 Brcb(dst, src, repeats, BrcbRepeatParams{1, 8}):把 src 的第 i 个标量复制成 dst 第 i 个 32B 块(fp32 = 8 个元素)。它是"标量数组 → 每行一个常量块"的桥,配合高维指令的"块级步进 1"就等价于按行广播。Gather(dst, src, offsets, base, count):dst[i] = src[byte(offsets[i]) + base],偏移表与 base 都以字节为单位,于是一条指令即可从行主序板里抽走一列。PipeBarrier<PIPE_X>vs 事件/TQueSync:前者只等"本管道内已发射的指令排空",是单管自身的定序;SetFlag/WaitFlag<HardEvent::A_B>(以及封装它的TQueSync<PIPE_A,PIPE_B>)是跨管握手,必须成对出现——单向 Set 或单向 Wait 都构不成依赖。310P 的事件 ID 数量有限,跨管事件要先pipe_->FetchEventID(...)申请(本段 820/824/842/846/1063 等),TQueSync则自带槽位(本段全用 0 号槽)。
1. 686-733 · 主循环入口三兄弟(ProcessHead / ProcessHeadVTile / ProcessChunk)
锚点:chunk_gated_delta_rule.h:686-733
__aicore__ inline void ProcessHead(int32_t seq0, int32_t seq1, uint64_t head_i, uint64_t batch_i) {
uint64_t nvPerNk = (NV_ >= NK_) ? (NV_ / NK_) : 1;
uint64_t qkHead = head_i / nvPerNk;
uint64_t stateBaseOffset = (batch_i * NV_ + head_i) * realK_ * realV_;
uint64_t workspaceStateBaseOffset = (batch_i * NV_ + head_i) * realK_ * stateWorkspaceStrideV_;
uint32_t cs = chunkSize_;
int32_t totalLen = seq1 - seq0;
uint32_t numC = Ceil(static_cast<uint32_t>(totalLen), cs);
for (uint32_t c = 0; c < numC; c++) {
ProcessChunk(seq0, seq1, head_i, c, qkHead, stateBaseOffset, workspaceStateBaseOffset, c + 1 == numC);
}
}
__aicore__ inline void ProcessHeadVTile(int32_t seq0, int32_t seq1, uint64_t head_i, uint64_t batch_i,
uint32_t vTileIdx) {
uint64_t nvPerNk = (NV_ >= NK_) ? (NV_ / NK_) : 1;
uint64_t qkHead = head_i / nvPerNk;
uint64_t stateBaseOffset = (batch_i * NV_ + head_i) * realK_ * realV_;
uint64_t workspaceStateBaseOffset = (batch_i * NV_ + head_i) * realK_ * stateWorkspaceStrideV_;
uint32_t cs = chunkSize_;
int32_t totalLen = seq1 - seq0;
uint32_t numC = Ceil(static_cast<uint32_t>(totalLen), cs);
for (uint32_t c = 0; c < numC; c++) {
ProcessChunkVTile(seq0, seq1, head_i, c, qkHead, stateBaseOffset, workspaceStateBaseOffset, c + 1 == numC,
vTileIdx);
}
}
// ---- per-chunk phases (ProcessHead delegates one chunk to ProcessChunk) ----
__aicore__ inline void ProcessChunk(int32_t seq0, int32_t seq1, uint64_t head_i, uint32_t c, uint64_t qkHead,
uint64_t stateBaseOffset, uint64_t workspaceStateBaseOffset, bool isLastChunk) {
uint32_t cs = chunkSize_;
uint32_t avFp32 = vStepAligned_;
int32_t totalLen = seq1 - seq0;
int32_t t_start = seq0 + c * cs;
uint32_t chunkLen =
(static_cast<uint32_t>(totalLen) - c * cs >= cs) ? cs : static_cast<uint32_t>(totalLen) - c * cs;
LoadChunkKey(qkHead, t_start, chunkLen);
LoadChunkCoefficients(t_start, head_i, chunkLen);
ComputeKBeta(chunkLen);
PrepareDecayAndExp(chunkLen);
PrepareAttnQueryCache(qkHead, t_start, chunkLen);
ComputeAttnMatrix(qkHead, t_start, chunkLen);
ComputeKCumdecay(chunkLen);
ProcessVTiles(t_start, chunkLen, head_i, qkHead, avFp32, stateBaseOffset, workspaceStateBaseOffset, c, isLastChunk);
}686-697 与 699-712 相对 !1108 零改动,成块引用即可,要点三条:
687-688:GQA 头映射。nvPerNk = NV_/NK_(NV_<NK_时保守取 1,即不做共享),一个 K/Q head 服务nvPerNk个 V head,所以qkHead = head_i/nvPerNk是整除降维。689/690:两份 state 基址。stateBaseOffset用realK_*realV_(对外outType布局),workspaceStateBaseOffset用realK_*stateWorkspaceStrideV_(workspace 里 V 维按 8 元素对齐后的步长)——同一份状态的两个地址空间,chunk 间递推走 workspace,最后一块才写 finalState。693-696:按 chunkSize 切块并逐个委派,695末尾的c + 1 == numC就是isLastChunk旗标。700-711的ProcessHeadVTile只多一个vTileIdx,服务realV_ > vStep_的场景(217-244的Process()按 (batch, head, vTile) 三元组分派工作项)。713/715是空行,714是段落注释// ---- per-chunk phases ... ----,无代码。
716-733 ProcessChunk 是 K 侧的总谱,与 !1108 唯一差异是 726/727 两行:ComputeGCumsum(t_start, head_i, chunkLen) + ComputeKBeta(t_start, head_i, chunkLen) → LoadChunkCoefficients(t_start, head_i, chunkLen) + ComputeKBeta(chunkLen)。签名里 t_start/head_i 从 ComputeKBeta 消失,正是"标量读 GM 的责任上移到 726"的函数级证据。718-723 算出 t_start 与尾块宽度 chunkLen(不足 cs 时取余数,尾块因此天然走不进 Cube 快路径);725-731 是本章要逐段拆的 7 个阶段,732 交棒 ProcessVTiles(第 08 章)。这 8 个阶段的先后顺序 = 缓冲复用链的顺序:chunkVFp32 在 725/726/728/730 里被借了四次(见 §9 总表)。
2. 734-778 · V-tile 版委派与 LoadChunkKey 重写
锚点:chunk_gated_delta_rule.h:734-778
__aicore__ inline void ProcessChunkVTile(int32_t seq0, int32_t seq1, uint64_t head_i, uint32_t c, uint64_t qkHead,
uint64_t stateBaseOffset, uint64_t workspaceStateBaseOffset,
bool isLastChunk, uint32_t vTileIdx) {
uint32_t cs = chunkSize_;
uint32_t avFp32 = vStepAligned_;
int32_t totalLen = seq1 - seq0;
int32_t t_start = seq0 + c * cs;
uint32_t chunkLen =
(static_cast<uint32_t>(totalLen) - c * cs >= cs) ? cs : static_cast<uint32_t>(totalLen) - c * cs;
LoadChunkKey(qkHead, t_start, chunkLen);
LoadChunkCoefficients(t_start, head_i, chunkLen);
ComputeKBeta(chunkLen);
PrepareDecayAndExp(chunkLen);
PrepareAttnQueryCache(qkHead, t_start, chunkLen);
ComputeAttnMatrix(qkHead, t_start, chunkLen);
ComputeKCumdecay(chunkLen);
uint32_t v_i = vTileIdx * vStep_;
uint32_t curV = (v_i + vStep_ > realV_) ? realV_ - v_i : vStep_;
ProcessOneVTile(t_start, chunkLen, head_i, qkHead, avFp32, stateBaseOffset, workspaceStateBaseOffset, c,
isLastChunk, v_i, curV);
}
// Phase 1: load K for this chunk into chunkKFp32.
__aicore__ inline void LoadChunkKey(uint64_t qkHead, int32_t t_start, uint32_t chunkLen) {
uint64_t stagingCapacity = static_cast<uint64_t>(2) * chunkSize_ * vStepAligned_;
uint64_t keyElements = static_cast<uint64_t>(chunkLen) * alignK_;
uint64_t gmRowStride = static_cast<uint64_t>(NK_) * realK_;
bool canUseDma = chunkLen <= UINT16_MAX && keyElements <= stagingCapacity;
if (likely(canUseDma)) {
LocalTensor<inType> keyLocal = chunkVFp32.template ReinterpretCast<inType>();
uint64_t qkOff = (static_cast<uint64_t>(t_start) * NK_ + qkHead) * realK_;
LoadPaddedRows(chunkKFp32, keyLocal, keyGm_, qkOff, chunkLen, static_cast<uint32_t>(gmRowStride));
return;
}
for (uint32_t i = 0; i < chunkLen; i++) {
int32_t t = t_start + i;
uint64_t qkOff = (static_cast<uint64_t>(t) * NK_ + qkHead) * realK_;
for (uint32_t j = 0; j < realK_; j++) {
chunkKFp32.SetValue(i * alignK_ + j, static_cast<float>(keyGm_.GetValue(qkOff + j)));
}
}
}734空行;735-756 ProcessChunkVTile与 !1108:423-444 同样只差746/747两行(同一个阶段替换)。752-753计算本 V-tile 的列窗口:v_i = vTileIdx*vStep_,尾 tile 收窄到realV_ - v_i;754-755直调ProcessOneVTile而不是走ProcessVTiles的v_i循环——因为 V-tile 循环已被提升到Process()的工作项层,这是 !1135 tiling 侧"工作单元三元组化"的落点。757空行;758Phase 1 注释。759-769是本段改动最大的地方。!1108 的canUseDma有 5 个条件(含alignK_ == realK_),非对齐 dk 会整块退化到771-777的双层标量循环(最多 64×128 = 8192 次标量 GM 读)。!1135 把对齐判定下沉进LoadPaddedRows(:272),这里只剩"行数能塞进 u16"与"暂存装得下"两条(763)。761:keyElements从chunkLen*realK_改成chunkLen*alignK_——要搬的是补齐到桶宽的行,LoadPaddedRows按alignK_行距写 UB,暂存容量必须按桶宽算。765就是任务书点名的"借缓冲"首演:chunkVFp32.ReinterpretCast<inType>()拿一块 FP32 缓冲的 FP16 视图当 DMA 落地点(同一块 UB,前keyElements*2字节),767的LoadPaddedRows内部完成 strided DMA + 尾巴标量 +Cast到chunkKFp32,fp16 视图随即死亡。760的stagingCapacity = 2*chunkSize_*vStepAligned_里那个 2 就是"fp32 缓冲的 fp16 元素数翻倍"。769的return之后770是空行,771-777保留 !1108:476-482 的标量 fallback(逐字未改,含775直接写 fp32 目标以省一次 Cast),776/777/778是三层收尾花括号。
3. 779-808 · BuildCoefficientOffsets 与乒乓并行前缀和
锚点:chunk_gated_delta_rule.h:779-808
__aicore__ inline void BuildCoefficientOffsets(LocalTensor<uint32_t> offsets, uint32_t chunkLen) {
LocalTensor<int32_t> signedOffsets = offsets.template ReinterpretCast<int32_t>();
CreateVecIndex<int32_t>(signedOffsets, 0, chunkLen);
PipeBarrier<PIPE_V>();
Muls(signedOffsets, signedOffsets, static_cast<int32_t>(NV_ * sizeof(float)), chunkLen);
PipeBarrier<PIPE_V>();
}
// Inclusive prefix sum for at most one 64-row chunk. Ping-ponging between
// two UB tensors avoids the loop-carried scalar GetValue/SetValue chain.
__aicore__ inline void PrefixSumChunkG(uint32_t chunkLen) {
LocalTensor<float> src = gCumsumFp32;
LocalTensor<float> dst = deltaFp32;
bool resultInGCumsum = true;
for (uint32_t offset = 1; offset < chunkLen; offset <<= 1) {
Muls(dst, src, 1.0f, chunkLen);
PipeBarrier<PIPE_V>();
Add(dst[offset], src[offset], src, chunkLen - offset);
PipeBarrier<PIPE_V>();
LocalTensor<float> tmp = src;
src = dst;
dst = tmp;
resultInGCumsum = !resultInGCumsum;
}
if (!resultInGCumsum) {
Muls(gCumsumFp32, src, 1.0f, chunkLen);
PipeBarrier<PIPE_V>();
}
}BuildCoefficientOffsets(780-786)——Gather 的偏移向量怎么造
781的offsets.template ReinterpretCast<int32_t>()是本文件反复出现的类型技巧:Gather的偏移表形参要求LocalTensor<uint32_t>,而CreateVecIndex/Muls在本机 CANN 上只有整型(int32_t)实例可实例化。同一块 UB 用两个类型视图分别满足两个 API 的类型要求,零拷贝、无转换开销;Muls在int32_t上做整数乘,位模式重新解释成uint32_t正好就是 Gather 要的字节偏移。782CreateVecIndex(signedOffsets, 0, chunkLen)填0,1,...,chunkLen-1;783PipeBarrier<PIPE_V>;784原地乘NV_*sizeof(float)(每个 token 有NV_个 head,一个 head 步距 =NV_*4字节);785再来一条PipeBarrier<PIPE_V>。两条栅栏都是同管读写定序:782与784读写同一批地址,栅栏保证向量管按序落地(784读的是782刚写的),并保证函数返回后偏移表对后续Gather可见。786收尾。全函数 3 条发射,换来后面两次Gather复用(816/829-830/849-851 的offsetsReady就是为复用服务的)。
PrefixSumChunkG(790-808)——Hillis-Steele 并行扫描
791-793双缓冲:src = gCumsumFp32、dst = deltaFp32、resultInGCumsum记录"结果此刻在哪个缓冲"。deltaFp32的容量在InitLocalBuffers:191-194保证 ≥ cs(max(stateStrideK_, cs, avStepAligned)),所以恒装得下 64 个 fp32。794offset <<= 1→ 轮数 =floor(log2(chunkLen)):chunkLen=64 → 6 轮,32 → 5 轮,1 → 0 轮(尾块/单 token 天然退化,不进循环)。795Muls(dst, src, 1.0f, chunkLen)是向量 memcpy(×1.0 是本文件的"搬家 + 刷新"惯用法);797Add(dst[offset], src[offset], src, chunkLen-offset)是扫描本体:dst[o+i] = src[o+i] + src[i]。为什么要先整段复制——797只覆盖i ∈ [offset, chunkLen)的后段,前offset个元素必须原样带到下一轮。798定序后799-802交换 src/dst 并翻奇偶旗标;803收循环。804-807:轮数为奇数时结果落在deltaFp32,再Muls ×1.0搬回gCumsumFp32(64 → 6 轮偶数,免搬;32 → 5 轮奇数,要搬)。这样无论走几轮,函数出口处"前缀和在gCumsumFp32"这一契约恒成立,869之后的读者不必知道轮数。- 收益核算:!1108:486-492 是 64 步串行
gSum += LoadG(...)(每步一次标量 GM 读 + 一次标量 UB 写 + 一次 loop-carried 依赖),这里压成 12 条向量指令(6 轮 × 2)+ 至多 1 条搬回,64 级依赖链变 6 级,且标量管整段空闲。 - 数值代价:Hillis-Steele 改变了加法结合顺序,与标量串行累加在 ULP 级别不等价(≤64 项同号累加,量级不受影响)。这是"并行化重排求值顺序"的常规代价,值得记住的是它和 !1158 的精度回归网(tail 精度双断言)共用同一套容差。
4. 809-864 · LoadChunkCoefficients:整板 DMA + Gather,以及 !1158 的那一行
锚点:chunk_gated_delta_rule.h:809-864
// Load the complete [chunk, Hv] beta/g slabs once, then gather this head
// from UB. The same byte-offset vector is reused for beta and g.
__aicore__ inline void LoadChunkCoefficients(int32_t t_start, uint64_t head_i, uint32_t chunkLen) {
uint64_t slabElements = static_cast<uint64_t>(chunkLen) * NV_;
uint64_t slabCapacity = static_cast<uint64_t>(chunkSize_) * vStepAligned_;
LocalTensor<uint32_t> coefficientOffsets = dotProductFp32.template ReinterpretCast<uint32_t>();
bool offsetsReady = false;
bool canUseBetaDma = slabElements <= slabCapacity && slabElements * sizeof(inType) % BLOCK_BYTES == 0;
if (likely(canUseBetaDma)) {
LocalTensor<inType> betaLocal = chunkVFp32.template ReinterpretCast<inType>();
event_t vectorToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(vectorToMte2);
WaitFlag<HardEvent::V_MTE2>(vectorToMte2);
DataCopy(betaLocal, betaGm_[static_cast<uint64_t>(t_start) * NV_], slabElements);
event_t mte2ToVector = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(mte2ToVector);
WaitFlag<HardEvent::MTE2_V>(mte2ToVector);
Cast(chunkVFp32, betaLocal, RoundMode::CAST_NONE, slabElements);
PipeBarrier<PIPE_V>();
BuildCoefficientOffsets(coefficientOffsets, chunkLen);
offsetsReady = true;
Gather(betaFp32, chunkVFp32, coefficientOffsets, static_cast<uint32_t>(head_i * sizeof(float)), chunkLen);
PipeBarrier<PIPE_V>();
} else {
for (uint32_t i = 0; i < chunkLen; ++i) {
betaFp32.SetValue(i, LoadBeta(t_start + i, head_i));
}
}
float gSum = 0.0f;
bool canUseGDma = hasGamma_ != 0 && slabElements <= slabCapacity && slabElements * sizeof(float) % BLOCK_BYTES == 0;
if (likely(canUseGDma)) {
event_t vectorToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(vectorToMte2);
WaitFlag<HardEvent::V_MTE2>(vectorToMte2);
DataCopy(chunkVFp32, gGm_[static_cast<uint64_t>(t_start) * NV_], slabElements);
event_t mte2ToVector = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(mte2ToVector);
WaitFlag<HardEvent::MTE2_V>(mte2ToVector);
if (!offsetsReady) {
BuildCoefficientOffsets(coefficientOffsets, chunkLen);
}
Gather(gCumsumFp32, chunkVFp32, coefficientOffsets, static_cast<uint32_t>(head_i * sizeof(float)), chunkLen);
PipeBarrier<PIPE_V>();
PrefixSumChunkG(chunkLen);
} else if (hasGamma_ == 0) {
Duplicate(gCumsumFp32, 0.0f, chunkLen);
PipeBarrier<PIPE_V>();
} else {
for (uint32_t i = 0; i < chunkLen; ++i) {
gSum += LoadG(t_start + i, head_i);
gCumsumFp32.SetValue(i, gSum);
}
}
}810-811注释已自陈设计:"Load the complete [chunk, Hv] beta/g slabs once, then gather this head from UB. The same byte-offset vector is reused for beta and g."813-814:slabElements = chunkLen*NV_是 GM 里[T, NV]板上一个 chunk 的全部 head(beta 与 g 同布局);slabCapacity = chunkSize_*vStepAligned_是借来的chunkVFp32的 fp32 元素容量——用 fp32 容量约束 fp16 板是保守的(fp16 只占一半字节),换来的是与 819/827 两个视图共用一段判定。815:偏移表宿主是dotProductFp32的uint32_t视图。生命周期成立:本函数是chunkLen阶段序列里第一个用到dotProductFp32的人,之后它才回归DotFp32(:390) 与ComputeRecursiveAttnDiagonalBlock:525的gatherOffsets身份。816/817:offsetsReady与canUseBetaDma(装得下 + 板的字节数是 32 的倍数,DataCopy 以 32B 块为单位)。818-832beta 快路径。819是 PR !1158 的一行修复点:betaLocal借的是chunkVFp32的 fp16 视图,而827的Cast(chunkVFp32, betaLocal, CAST_NONE, slabElements)源与目的同址且是 fp16→fp32 的加宽——读2n字节、写4n字节,Vector 管按 32B 块(8 个 fp32)为一步推进,写窗口永远是读窗口的 2 倍宽且向前延伸,从第 2 个 repeat 起未读的 fp16 源已被自己的输出踩掉,只有前 8 个系数幸存。所以head_i ≥ 8(宽 GQA)时 beta 全程损坏,NV_ ≤ 8的窄 shape 完全看不出来——这正是 !1135 那轮 12P/1F 实测里没被任何用例抓到的 bug。- 它此刻"借的是哪个缓冲":
chunkVFp32,同一块 UB 既当 DMA 落地点(fp16 视图)又当加宽目的地(fp32 视图)。与765那次成功借用的唯一区别是:那里 fp16→fp32 的目的地是另一块chunkKFp32,视图内只做同宽度 DMA,读写区间不相交。修复(src/88d8b965:820,并在819补了一句注释 "Keep the FP16 source separate because an in-place widening Cast overwrites unread beta values.")把源换成chunkAttnOutFp32的 fp16 视图——它在本函数返回后才会被PrepareAttnQueryCache覆写,借用窗口夹在原主人两次使用之间。 820-826是两条必须成对的跨管事件:V_MTE2(820-822)防 WAR——上一轮 V 管对chunkVFp32的写要排空,MTE2 才准往里写;MTE2_V(824-826)防 RAW——DMA 未完成 V 管不得读。823是唯一一条DataCopy(一维长度版,落点连续、无 gap)。828PipeBarrier<PIPE_V>保证Cast落地后Gather才读;829-830造偏移表并置旗标;831Gather(betaFp32, chunkVFp32, coefficientOffsets, head_i*sizeof(float), chunkLen)一条指令抽走第head_i列(第 4 参 base 以字节为单位,故head_i*sizeof(float);净效果betaFp32[i] = slab[i*NV_ + head_i]);832定序。833-837标量 fallback:834-836就是 !1108:503 的 64 次LoadBeta(每次一次标量 GM 读),837收if/else。838空行。839float gSum = 0.0f;只有 fallback 用得到(860),快路径下它全程为 0 —— 这是 !1108ComputeGCumsum的残余声明,成本是一个标量寄存器,保留它让两条路径共用同一函数体。840-854g 快路径:840比 beta 多一个hasGamma_ != 0条件,且按sizeof(float)判 32B 对齐;842-848同样的事件对 + 一条 DMA 把 g 板直接搬进chunkVFp32——g 本身就是 fp32(gGm_是GlobalTensor<float>),所以这里是同宽度搬运,没有 Cast,!1158 的坑只存在于 beta 分支。此刻 beta 板已被831抽干,覆写合法(读者只剩betaFp32)。849-851是offsetsReady的用途:beta 走了 fallback 时偏移表还没造,这里补一次——if (!offsetsReady)保证四种路径组合下852的 Gather 都有有效偏移表。852-853Gather 出原始 g(还不是前缀和);854交给PrefixSumChunkG就地扫描。855-857hasGamma_ == 0时Duplicate(gCumsumFp32, 0.0f, chunkLen):没有 g 输入 ≡ 不衰减,与 !1108 靠LoadG恒返0.0f(:353-355) 等价,但从"64 次标量写"变成"1 条向量写"。858-863唯一保留gSum的标量累加链(861直接写gCumsumFp32,所以这条路径绕开PrefixSumChunkG,结果契约仍然一致)。864收尾。- 收益核算:!1108 = 64 次标量读 beta + 64 次标量读 g + 64 步串行累加 ≈ 128 次标量 GM 访存 + 128 次 Set/Wait;!1135 快路径 = 2 条 DMA + 1 条偏移表(3 发) + 2 条 Gather + 12 条扫描 ≈ 20 条发射,标量 GM 访存 0 次。
5. 865-899 · ComputeKBeta:Brcb 展开 + 高维 Mul 一次乘整块
锚点:chunk_gated_delta_rule.h:865-899
// Phase 2: k_beta = K * beta -> kCumdecayFp32. Brcb expands each
// per-row beta into one 32-byte block; high-dimensional Mul then broadcasts
// that block across every K block while advancing one beta per row.
__aicore__ inline void ComputeKBeta(uint32_t chunkLen) {
TQueSync<PIPE_S, PIPE_V> inputSync;
inputSync.SetFlag(0);
inputSync.WaitFlag(0);
LocalTensor<float> betaBlocks = chunkVFp32;
uint8_t brcbRepeats = static_cast<uint8_t>(Ceil(chunkLen, FP32_NUM_PER_BLOCK));
Brcb(betaBlocks, betaFp32, brcbRepeats, BrcbRepeatParams{1, 8});
PipeBarrier<PIPE_V>();
// The padded K columns must remain exact zeros because all Cube products
// consume the complete bucket width.
Duplicate(kCumdecayFp32, 0.0f, chunkLen * alignK_);
PipeBarrier<PIPE_V>();
uint8_t rowStride = static_cast<uint8_t>(alignK_ / FP32_NUM_PER_BLOCK);
BinaryRepeatParams mulParams{1, 1, 0, rowStride, rowStride, 1};
uint32_t kOffset = 0;
for (; kOffset + 64 <= realK_; kOffset += 64) {
Mul(kCumdecayFp32[kOffset], chunkKFp32[kOffset], betaBlocks, static_cast<uint64_t>(64),
static_cast<uint8_t>(chunkLen), mulParams);
}
if (kOffset < realK_) {
Mul(kCumdecayFp32[kOffset], chunkKFp32[kOffset], betaBlocks, static_cast<uint64_t>(realK_ - kOffset),
static_cast<uint8_t>(chunkLen), mulParams);
}
PipeBarrier<PIPE_V>();
TQueSync<PIPE_V, PIPE_S> betaSync;
betaSync.SetFlag(0);
betaSync.WaitFlag(0);
}866-868注释就是这段的规格书:"Brcb expands each per-row beta into one 32-byte block; high-dimensional Mul then broadcasts that block across every K block while advancing one beta per row."870-872TQueSync<PIPE_S,PIPE_V> inputSync:进方向栅栏。725/726的 fallback 分支是标量管写chunkKFp32/betaFp32的,V 管第一次读它们之前必须定序;快路径下前面已有PipeBarrier<PIPE_V>覆盖,这条 S→V 是"两条路径都正确"的兜底,与 !1108:498-500 同位同义(保留未动)。874-876:betaBlocks = chunkVFp32(第 3 次借用,此刻 beta/g 板都已死亡);brcbRepeats = ceil(chunkLen/8)——一条 Brcb 的每个 repeat 吞 8 个标量(256 bit)产出 8 个 32B 块,8 repeat 覆盖 64 个 beta。879-882:这是终态相对 a9c66ddb 新增的四行(含注释与空行)。881Duplicate(kCumdecayFp32, 0.0f, chunkLen*alignK_)把整个输出区先清 0,因为887-894的 Mul 只覆盖realK_列,桶宽的 padding 列[realK_, alignK_)必须保持精确 0——后面每个 Cube 乘积都吃完整桶宽(kAElements = 64*kSpecializedDk,且快路径下alignK_ == kSpecializedDk,恰好是881覆盖的字节数)。这是 69300ffc 泛化非对齐 Dk 时必须补的保险。884-885:rowStride = alignK_/8即"一行的 32B 块数";mulParams六个间隔值描述 (dst, src0, src1) 的两级步进,净效果是:chunkKFp32与kCumdecayFp32在行内连续推进、行末前进rowStride块,betaBlocks行内不推进、行末只前进 1 块(=每行一个常量 → 广播)。886-890:K 维以 64 列为步长切块(fp32 一次 repeat 的上限正好 64 个元素 = 8 块),891-894处理非 64 倍数的尾巴。realK_=128→ 2 条Mul;realK_=80→ 1 条整块 + 1 条 16 列;889/893的static_cast<uint8_t>(chunkLen)是 repeatTime 上限(127)之内的防御写法。对比 !1108:502-507:64 次Muls+ 64 次标量LoadBeta+ 128 次 Set/Wait → 现在 1 条 Brcb + 1 条 Duplicate + 2-3 条 Mul。895-898:PipeBarrier<PIPE_V>+TQueSync<PIPE_V,PIPE_S> betaSync(出方向):981/1202等标量读者要读由kCumdecayFp32派生的量,必须让 V 管写完成对 S 可见。899收尾。
6. 900-940 · PrepareDecayAndExp:decay 行矩阵一次构造
锚点:chunk_gated_delta_rule.h:900-940
// Phase 3: precompute exp(gCumsum) and decay_mask[i][j] = exp(g_cumsum[i] -
// g_cumsum[j]).
__aicore__ inline void PrepareDecayAndExp(uint32_t chunkLen) {
uint32_t cs = chunkSize_;
Muls(gCumsumFp32, gCumsumFp32, 1.0f, chunkLen);
PipeBarrier<PIPE_V>();
Exp(expGCumFp32, gCumsumFp32, chunkLen);
PipeBarrier<PIPE_V>();
// g is a log-decay and its prefix sum can easily be below -100 for real
// Qwen activations. Factoring exp(g_i - g_j) into exp(g_i)*exp(-g_j)
// overflows the second factor. Build all rows with high-dimensional Vector
// instructions, clear the non-causal upper triangle before Exp, then issue
// one strided Exp for the complete causal matrix.
uint8_t rowStride = static_cast<uint8_t>(cs / FP32_NUM_PER_BLOCK);
UnaryRepeatParams rowBroadcastParams{1, 1, rowStride, 0};
Muls(decayMaskFp32, gCumsumFp32, -1.0f, static_cast<uint64_t>(chunkLen), static_cast<uint8_t>(chunkLen),
rowBroadcastParams);
PipeBarrier<PIPE_V>();
LocalTensor<float> gBlocks = chunkVFp32;
uint8_t brcbRepeats = static_cast<uint8_t>(Ceil(chunkLen, FP32_NUM_PER_BLOCK));
Brcb(gBlocks, gCumsumFp32, brcbRepeats, BrcbRepeatParams{1, 8});
PipeBarrier<PIPE_V>();
BinaryRepeatParams rowAddParams{1, 1, 0, rowStride, rowStride, 1};
Add(decayMaskFp32, decayMaskFp32, gBlocks, static_cast<uint64_t>(chunkLen), static_cast<uint8_t>(chunkLen),
rowAddParams);
PipeBarrier<PIPE_V>();
for (uint32_t row = 0; row + 1 < chunkLen; ++row) {
Duplicate(decayMaskFp32[row * cs + row + 1], 0.0f, chunkLen - row - 1);
}
PipeBarrier<PIPE_V>();
UnaryRepeatParams rowExpParams{1, 1, rowStride, rowStride};
Exp(decayMaskFp32, decayMaskFp32, static_cast<uint64_t>(chunkLen), static_cast<uint8_t>(chunkLen), rowExpParams);
PipeBarrier<PIPE_V>();
TQueSync<PIPE_V, PIPE_S> decaySync;
decaySync.SetFlag(0);
decaySync.WaitFlag(0);
}900空行;901-902注释声明两个产物:expGCumFp32[i] = e^{G_i}与decayMaskFp32[i][j] = e^{G_i-G_j}。904-906Muls(gCumsum, gCumsum, 1.0f, chunkLen):×1.0 刷新(855-863两条标量路径写过gCumsumFp32,这里保证 V 管后续读到的就是最终值)。长度参数从 !1108 的cs改成chunkLen(尾块不再多算)。907-908:一次向量Exp出 64 个e^{G_i}。910-914是本 PR 最重要的数值注释:g 是对数衰减,真实 Qwen 激活里前缀和轻松低于 −100,所以exp(G_i − G_j)不能拆成exp(G_i)·exp(−G_j)——第二个因子会溢出成 inf。必须整块构造差值再整体 Exp。915-918行广播 −G_j:rowStride = cs/8(行距 8 块);rowBroadcastParams{1, 1, rowStride, 0}的净效果是"src 行间不前进、dst 行间前进 rowStride 块",于是decayMask[i][j] = -G_j对每一行都成立(同一份 −G 复制 64 行)。919空行;920-924gBlocks = chunkVFp32(第 4 次借用)+ Brcb 把+G_i展成"每行一个 32B 块"。925-928高维Add(rowAddParams与885的mulParams同形)→decayMask[i][j] = G_i − G_j,整块一次。929空行;930-932Exp 之前清非因果区:j > i处G_i − G_j > 0,Exp 出来是 inf 甚至溢出;不清零就等着后面1145/1148的整块Mul把0×inf变成 NaN、并被ComputeRecursiveAttnBlocked64的整块读取消费(分块求逆按完整 16×16/32×32 块读,上三角的垃圾值会参与乘积)。这 63 条Duplicate是本段残留的发射开销,作者选择用它换数值安全;!1108 的逐行写法天然不会生成上三角的爆炸值(写出来也没人读),所以这条陷阱是新并行化引入的,注释与代码一起把它钉住了。933-935rowExpParams{1, 1, rowStride, rowStride}+ 一条高维Exp算完整个 64×64(dst/src 行间都前进 rowStride 块)。对比 !1108:523-531:64 次标量读 + 64×(Duplicate+Sub+Exp) = 192 条指令 → 现在 6 条高维指令 + 63 条清零。936-939TQueSync<PIPE_V,PIPE_S> decaySync:出方向,981/993/1185/2140的标量GetValue全靠它定序。940收尾。
7. 941-961 · CanCacheAttnQuery 与 Q 缓存(chunkVFp32 再借一次)
锚点:chunk_gated_delta_rule.h:941-961
__aicore__ inline bool CanCacheAttnQuery(uint32_t chunkLen) {
uint64_t queryElements = static_cast<uint64_t>(chunkLen) * alignK_;
uint64_t stagingCapacity = static_cast<uint64_t>(2) * chunkSize_ * vStepAligned_;
uint64_t cacheCapacity = static_cast<uint64_t>(chunkSize_) * vStepAligned_;
return chunkLen <= UINT16_MAX && queryElements <= stagingCapacity && queryElements <= cacheCapacity;
}
// chunkAttnOutFp32 is unused until V-tile processing. When it can hold a
// full Q chunk, reuse chunkVFp32 as FP16 staging and cache Q*scale once for
// the attention QK dot products.
__aicore__ inline void PrepareAttnQueryCache(uint64_t qkHead, int32_t t_start, uint32_t chunkLen) {
if (!CanCacheAttnQuery(chunkLen)) {
return;
}
LocalTensor<inType> queryLocal = chunkVFp32.template ReinterpretCast<inType>();
uint64_t qkOff = (static_cast<uint64_t>(t_start) * NK_ + qkHead) * realK_;
LoadPaddedRows(chunkAttnOutFp32, queryLocal, queryGm_, qkOff, chunkLen, NK_ * realK_);
Muls(chunkAttnOutFp32, chunkAttnOutFp32, scale_, chunkLen * alignK_);
PipeBarrier<PIPE_V>();
}941空行。942-947 CanCacheAttnQuery相对 !1108:534-543 的两处变化:943的长度从realK_改成alignK_(桶宽),946把 5 个 DMA 对齐条件删光(下沉给LoadPaddedRows)。三条判据各管一事:946的chunkLen <= UINT16_MAX(DMA 块数字段宽度)、queryElements <= stagingCapacity(fp16 暂存放得下,944的2*chunkSize_*vStepAligned_=chunkVFp32的 fp16 元素容量)、queryElements <= cacheCapacity(945=chunkAttnOutFp32的 fp32 元素容量)。948空行;949-951注释书面交代了借缓冲的合法性:"chunkAttnOutFp32 is unused until V-tile processing"——它要到 Phase 6 Step 2 才有正式主人,此刻是空房。952-961 PrepareAttnQueryCache:!1108 的 25 行(手工算srcGapBytes、造DataCopyParams、S_MTE2/V_MTE2/MTE2_V三组事件、Cast)塌成958一行LoadPaddedRows。956就是任务书说的"chunkVFp32 被 ReinterpretCast 成 inType 当 fp16 暂存":与765、819同一手法,但这次目的地是chunkAttnOutFp32(fp32 区),源与目的天然分离,不存在 !1158 那类重叠。959的Muls(..., scale_, chunkLen*alignK_)长度从chunkLen*realK_扩到桶宽——padding 列0×scale = 0,仍然干净,而 Cube 会连这些列一起乘。960定序,961收尾。- 这个缓存的意义:Q·scale 是
ComputeAttnProductsCube第二个乘积的 A 矩阵(1113直接把它分裂出 GM),快路径下没有它就没有 QK^T。
8. 962-1014 · ComputeAttnMatrix:Cube / Fallback 分叉点
锚点:chunk_gated_delta_rule.h:962-1014
// Phase 4: attn = -((k_beta @ K^T) * decay_mask) (lower tri), recursive accumulation, identity
// diagonal; then attn_i = (Q @ K^T) * decay_mask overwrites decayMaskFp32 (lower tri + diag).
__aicore__ inline void ComputeAttnMatrix(uint64_t qkHead, int32_t t_start, uint32_t chunkLen) {
uint32_t cs = chunkSize_;
if constexpr (kSpecializedDk != 0) {
if (likely(realK_ <= kSpecializedDk && chunkLen == 64 && vStepAligned_ >= kSpecializedDk &&
CanCacheAttnQuery(chunkLen))) {
ComputeAttnProductsCube<kSpecializedDk>();
ComputeRecursiveAttn(chunkScoresFp32, chunkLen, cs);
for (uint32_t i = 0; i < chunkLen; i++) {
chunkScoresFp32.SetValue(i * cs + i, 1.0f);
}
return;
}
}
for (uint32_t i = 1; i < chunkLen; i++) {
for (uint32_t jj = 0; jj < i; jj++) {
float dot = DotFp32(kCumdecayFp32[i * alignK_], chunkKFp32[jj * alignK_], realK_);
chunkScoresFp32.SetValue(i * cs + jj, -dot * decayMaskFp32.GetValue(i * cs + jj));
}
}
ComputeRecursiveAttn(chunkScoresFp32, chunkLen, cs);
for (uint32_t i = 0; i < chunkLen; i++) {
chunkScoresFp32.SetValue(i * cs + i, 1.0f);
}
if (CanCacheAttnQuery(chunkLen)) {
for (uint32_t i = 0; i < chunkLen; i++) {
LocalTensor<float> queryRow = chunkAttnOutFp32[i * alignK_];
for (uint32_t jj = 0; jj <= i; jj++) {
float dot = DotFp32(queryRow, chunkKFp32[jj * alignK_], realK_);
float decay = decayMaskFp32.GetValue(i * cs + jj);
decayMaskFp32.SetValue(i * cs + jj, dot * decay);
}
}
} else {
for (uint32_t i = 0; i < chunkLen; i++) {
int32_t t_i = t_start + i;
uint64_t qkOff_i = (static_cast<uint64_t>(t_i) * NK_ + qkHead) * realK_;
for (uint32_t d = 0; d < realK_; d++) {
deltaFp32.SetValue(d, static_cast<float>(queryGm_.GetValue(qkOff_i + d)) * scale_);
}
TQueSync<PIPE_S, PIPE_V> querySync;
querySync.SetFlag(0);
querySync.WaitFlag(0);
for (uint32_t jj = 0; jj <= i; jj++) {
float dot = DotFp32(deltaFp32, chunkKFp32[jj * alignK_], realK_);
float decay = decayMaskFp32.GetValue(i * cs + jj);
decayMaskFp32.SetValue(i * cs + jj, dot * decay);
}
}
}
}963-964注释与 !1108:574-575 同文(Phase 4 的产物定义:A 下三角 + 递归 + 单位对角;然后QK^T ⊙ decay覆盖decayMaskFp32)。967 if constexpr (kSpecializedDk != 0):编译期分叉。kSpecializedDk == 0的通用实例里整段被折叠掉,Cube 代码不进 binary、不占寄存器;这就是"分桶常量折叠保证重构零性能损失"的机器层面依据(docs/PR-1135-d0906c3e/10§9.2 的实测结论)。725cf45d 之前这里是三段==64 / ==128 / ==80重复。968-969四个运行时守卫,逐个说为什么需要:realK_ <= kSpecializedDk(构造期90-91已令alignK_ = kSpecializedDk,K 维补齐到桶宽,Cube 的 reduction 维才等于kMatmulK);chunkLen == 64(M/N 固定 64,尾块必走慢路径,长序列最后一个 chunk 是短板);vStepAligned_ >= kSpecializedDk(1068/1113的StageHalfWithResidual借chunkVFp32当残差 scratch,要装64*桶宽个 fp32,容量式即cs*vStepAligned_ ≥ 64*kSpecializedDk);CanCacheAttnQuery(chunkLen)(第二个乘积的 A 必须是952缓存好的 Q·scale)。注意 257 行的IsCubeFastPath用的是vStepAligned_ == kSpecializedDk,与本处>=不一致,见自查。970一次ComputeAttnProductsCube<kSpecializedDk>()顶掉 4096 次DotFp32;971ComputeRecursiveAttn在629-632处再分叉到ComputeRecursiveAttnBlocked64(第 05 章);972-974标量把对角写成 1.0((I−A)⁻¹的对角本就是 1,这里显式覆盖是 !1108 的语义延续);975return,976/977两层收尾。978-1013fallback 与 !1108:578-613 逐字节相同(diff 为空):978-983A 的直接项(980每格一次DotFp32,981标量读 decay 并写-dot*decay)、984递归修正、985-987对角 1、988-996Q 已缓存时逐行逐列点积、997-1013Q 未缓存时每行现从 GM 标量读进deltaFp32(1001-1003)再点积。1004-1006的TQueSync<PIPE_S,PIPE_V> querySync是本段唯一不能省的 S→V 配对:1002用标量SetValue写deltaFp32,1008的DotFp32在 V 管读这块 UB;不配对就会拿上一轮的残留参加点积。1011/1012/1013/1014四层收尾。- 数量级对照:!1108 这一段 ≈
2016 + 2080 = 4096次DotFp32,每次约 5 条发射 → 约 2 万条;!1135 快路径 = 4 次Mmad+ 2 次 L0C→UB 转换(第 07 章)。
10. chunkVFp32 在本段的五次身份(借缓冲生命周期)
锚点:chunk_gated_delta_rule.h:759-940
| 借用点 | 视图 | 装什么 | 死亡时刻 |
|---|---|---|---|
| 765 | ReinterpretCast<inType> | K 的 fp16 DMA 暂存 | LoadPaddedRows 内 Cast 到 chunkKFp32 之后(338) |
| 819 | ReinterpretCast<inType> | beta 板 fp16 落地点 | 831 Gather 之后(但 827 与它同址 → !1158) |
| 819/827/845 | fp32 视图 | beta 板加宽结果 / g 板落地点 | 852 Gather 之后 |
| 874 | fp32 视图 | betaBlocks(Brcb 广播块) | 894 高维 Mul 之后 |
| 921 | fp32 视图 | gBlocks(Brcb 广播块) | 927 高维 Add 之后 |
五次借用互不相交的唯一硬证据是阶段顺序(§1 的 725-731):借的窗口全部夹在"上一个读者的最后一次使用"和"下一个主人的第一次写"之间。写 Cube 代码时这张表必须逐次重推——!1158 的教训正是"照抄上一处的成功答案"。