05 · K 侧准备阶段:Phase 1~5 与四层编排
本篇覆盖 src/00e39ab5/op_kernel/chunk_gated_delta_rule.h 第 374-648 行。结论先行:
- 这 275 行是编排层(4 个
Process*)+ 准备层(7 个 Phase 函数)。编排层没有任何向量指令,准备层才是全部 K 侧算法。 - Phase 1~5 与 V 维无关,产出三样跨 V-tile 复用的东西:
chunkScoresFp32(逆矩阵 A)、chunkKFp32(被覆盖成 K_cd)、decayMaskFp32(被覆盖成 QK 注意力矩阵)。"每 chunk 一次"与"每 tile 一次"的分界就在ProcessChunk的最后一行(:420)。 - 本区间共有 9 处
TQueSync对象 + 3 组事件三段式。它们不是同一种东西:有 3 处是真正必需的"标量写 UB → 向量读 UB",有 3 处是块级保守收口,有 3 处只是给标量侧的行系数"发指令前打点"。逐处给理由见 §13 的总表。
先给层级图(箭头 = 调用,括号内是行号):
Process(209) ─┬─ ProcessHead(374) ──── ProcessChunk(404) ─┬─ Phase1 LoadChunkKey(447)
│ (chunk 循环 382) Phase1b ComputeGCumsum(486)
│ Phase2 ComputeKBeta(495)
│ Phase3 PrepareDecayAndExp(514)
│ Phase3.5 PrepareAttnQueryCache(548)
│ Phase4 ComputeAttnMatrix(576) ── DotFp32(289)
│ │ └─ ComputeRecursiveAttn(317)
│ Phase5 ComputeKCumdecay(617)
│ Phase6 ProcessVTiles(649) → ProcessOneVTile(659) ×tile
└─ ProcessHeadVTile(387) ─ ProcessChunkVTile(423) ── 重复 Phase1~5 + 单个 ProcessOneVTile1. ProcessHead:一个 value-head 的全部 chunk
锚点:chunk_gated_delta_rule.h:374-386
__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);
}
}逐行:
- :375
nvPerNk:GQA 式分组。NV_ < NK_时取 1(不做除法,避免head_i/0之外的整数截断意外,且此时语义是"每个 V-head 用自己的 QK-head")。约定NK_整除NV_,否则该映射会退化,kernel 不校验(tiling 侧也不校验,属隐式前置条件)。 - :376
qkHead:Q/K 的一切 GM 偏移只用 qkHead;beta/g/state 的偏移只用 head_i。这个二分贯穿全 kernel(:457/:478/:553/:600/:818/:863 用 qkHead,:255/:259/:503/:706 用 head_i),混用会造成静默的跨 head 数据污染。 - :377/:378 两个 state 基址:前者是公共接口
[B, Nv, dv, dk](元素数dk*dv,乘积顺序不影响基址但影响行内寻址,见 06 篇 §3),后者是 kernel 自定义 workspace[B, Nv, dk, align32(dv)]。同一 head 的起始元素偏移都是(b·NV_+h)·dk·宽度,差别只在宽度取realV_还是stateWorkspaceStrideV_。 - :381
numC = Ceil(totalLen, cs):按本 batch 真实长度算块数,不读 tiling 的numChunks_。所以numChunks_(:953)与T_(:945)在本 commit 里是只赋值不读取的死成员(实测:T_只出现在 :77/:945,numChunks_只出现在 :84/:953)。它们的存在是为了与 host tiling 结构体字段一一镜像,便于CGDR_DEBUG_LOG对账。 - :383
c + 1 == numC就是isLastChunk,决定 state 写 FP32 workspace 还是 cast 成 FP16 写finalState(06 篇 §9)。
chunk 之间严格串行(
2. ProcessHeadVTile:同一件事,多带一个 tile 下标
锚点:chunk_gated_delta_rule.h:387-403
__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) ----:389-395 与 :375-381 逐字符相同,:397-398 只是多传一个 vTileIdx。作者选择复制而不是抽公共函数,理由是这些函数全部 __aicore__ inline:任何抽象都会变成设备侧的真实函数调用或模板参数膨胀,而复制只多 7 行标量代码(无向量指令,体积与时间都可忽略)。代价是双份维护:改 GQA 映射或 state 基址公式必须同时改两处——这是后续 PR 最容易漏改的地方之一(PR-1135 的 Cube 化确实重排了这两处)。
3. ProcessChunk:八段 phase 的唯一编排点
锚点:chunk_gated_delta_rule.h:404-422
__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);
ComputeGCumsum(t_start, head_i, chunkLen);
ComputeKBeta(t_start, head_i, 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);
}- :410-411
chunkLen:在 uint32 域里比较totalLen - c*cs >= cs,尾块取余数。这个式子保证1 ≤ chunkLen ≤ cs(因为c < numC⇒c*cs < totalLen)。所有三角循环的上界都是chunkLen,所有行跨距/缓冲宽度都是chunkSize_(ld)——这一对区分是本 kernel 处理变长尾块不越界的全部依据(尾块之后未写到的矩阵元素保持上一 chunk 的残留,但因为读的时候同样只用< chunkLen,永远读不到)。 - :407
avFp32 = vStepAligned_:V 侧所有行宽都用它。 - :413-419 的顺序是硬约束,逐条给出理由:
LoadChunkKey→ComputeKBeta:后者 :506 读chunkKFp32。ComputeGCumsum→PrepareDecayAndExp:后者 :516-527 读gCumsumFp32。PrepareDecayAndExp→ComputeAttnMatrix::581/:593 读decayMaskFp32、:621(Phase 5)读expGCumFp32。PrepareAttnQueryCache→ComputeAttnMatrix:后者 :590 读缓存。且它必须排在 Phase 1/Phase 3 之后——它借chunkVFp32当 FP16 暂存(:552),而 Phase 1 也借同一块(:456):两次借用之间必须有一次把前一次的数据用完(Phase 1 的Cast在 :471 完成消费)。所以 Phase 3.5 不能提前到 Phase 1 之前。ComputeAttnMatrix→ComputeKCumdecay:后者 :635 读 A。ComputeKCumdecay→ProcessVTiles:Step 4 用 K_cd,且 K_cd 覆盖在chunkKFp32上——覆盖之后原始 K 在 UB 里就没了,所以任何还需要 UB 里原始 K 的阶段都必须排在前面(Phase 4 正是最后一个使用者)。PrepareAttnQueryCache还必须排在ProcessVTiles之前:它把 Q 缓存放进chunkAttnOutFp32,而 Step 1(:687)会把同一块借去当 V 的 FP16 暂存。Q 缓存在 Phase 4 之后立刻变成死数据,这正是 Step 5 只能从 GM 现读 Q(:821)的原因——一条完整的因果链(06 篇 §6)。
4. ProcessChunkVTile:为并行度付出的计算冗余
锚点:chunk_gated_delta_rule.h:423-446
__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);
ComputeGCumsum(t_start, head_i, chunkLen);
ComputeKBeta(t_start, head_i, 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.:433-439 与 :413-419 完全相同 ⇒ tile 模式下 K 侧准备被执行 vTileCount 次(每个负责一个 V-tile 的核各做一遍)。这不是疏漏而是设计:Phase 1~5 的产出是每核私有 UB 里的数据,跨核无法共享(本 commit 没有任何跨核通信原语);而该模式的启用条件(ShouldSplitVTiles)恰好是"核有空闲",此时冗余计算的时间是白送的(并行度收益 > 冗余成本)。当 vTileCount 大到核不再空闲时,host 侧就不给这么大的 blockDim 了(blockDim = min(hv·vTileCount, aivNum))。
:440-441 的 v_i/curV 与 ProcessVTiles(:652-653)里的循环体是同一对公式的两种写法。二者必须永远等价,否则 tile 模式下最后一个 tile 会少算/多算 dv % vStep 列。
5. LoadChunkKey:2D strided DMA 快路径与非对齐回退
锚点:chunk_gated_delta_rule.h:447-485
__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) * realK_;
uint64_t gmRowStride = static_cast<uint64_t>(NK_) * realK_;
uint64_t srcGapBytes = (gmRowStride - realK_) * sizeof(inType);
bool canUseDma = chunkLen <= UINT16_MAX && keyElements <= stagingCapacity && alignK_ == realK_ &&
(realK_ * sizeof(inType)) % BLOCK_BYTES == 0 && srcGapBytes % BLOCK_BYTES == 0 &&
srcGapBytes / BLOCK_BYTES <= UINT16_MAX;
if (likely(canUseDma)) {
LocalTensor<inType> keyLocal = chunkVFp32.template ReinterpretCast<inType>();
uint64_t qkOff = (static_cast<uint64_t>(t_start) * NK_ + qkHead) * realK_;
DataCopyParams keyParams{static_cast<uint16_t>(chunkLen),
static_cast<uint16_t>(realK_ * sizeof(inType) / BLOCK_BYTES),
static_cast<uint16_t>(srcGapBytes / BLOCK_BYTES), 0};
event_t scalarToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE2));
SetFlag<HardEvent::S_MTE2>(scalarToMte2);
WaitFlag<HardEvent::S_MTE2>(scalarToMte2);
event_t vectorToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(vectorToMte2);
WaitFlag<HardEvent::V_MTE2>(vectorToMte2);
DataCopy(keyLocal, keyGm_[qkOff], keyParams);
event_t mte2ToVector = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(mte2ToVector);
WaitFlag<HardEvent::MTE2_V>(mte2ToVector);
Cast(chunkKFp32, keyLocal, RoundMode::CAST_NONE, chunkLen * realK_);
PipeBarrier<PIPE_V>();
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)));
}
}
}
// Phase 1b: cumulative sum of g into gCumsumFp32.为什么需要暂存 + Cast 两步而不是 DMA 直接进 FP32:GM 里 K 是 FP16,UB 目标 chunkKFp32 是 FP32,DataCopy 不做类型转换,所以必须先按 FP16 搬进一块 FP32 空间(省不到第二块缓冲),再用向量 Cast 无损扩展。
- :448
stagingCapacity = 2·cs·vStepAligned_:借的是chunkVFp32,容量cs·vStepA个 FP32 =cs·vStepA·4字节;按 FP16 元素数就是2·cs·vStepA。因子 2 正是 FP32→FP16 的元素数翻倍,不是保守估计。 - :449
keyElements = chunkLen·realK_:待搬的 FP16 元素数(K 无 padding,紧凑)。 - :450-451
gmRowStride = NK_·realK_是 GM 的行跨距(一次t+1要跳过的元素数),srcGapBytes = (NK_-1)·realK_·2即"跳过同一 token 其余 K-head"。 - :452-454 六个条件逐个:
chunkLen ≤ UINT16_MAX(blockCount 是 u16,cs=64 恒成立,防御)、keyElements ≤ stagingCapacity(借的暂存装得下)、alignK_ == realK_(关键:见下)、realK_·2 % 32 == 0(FP16 行宽整块)、srcGapBytes % 32 == 0(gap 整块)、srcGapBytes/32 ≤ 65535(gap 塞得进 u16)。
alignK_ == realK_ 为什么同时挡住两件事::471 的 Cast(chunkKFp32, keyLocal, ..., chunkLen*realK_) 是扁平一次转,源和目标都按连续元素排。目标 chunkKFp32 的行宽是 alignK_(:156 用 cs*ak 分配),只有 alignK_ == realK_ 时"连续第 m 个元素"才恰好等于"第 m/realK_ 行第 m%realK_ 列"。否则需要逐行 Cast 并在目标侧留 gap——Cast 这类逐元素向量指令没有 dstGap 参数,只能写成 64 次调用。所以这一个条件同时保证了"FP16 暂存紧凑"和"FP32 目标无 padding",快路径才有可能只用一条 DMA + 一条 Cast。
- :456
chunkVFp32.template ReinterpretCast<inType>():只改类型视图不改地址,此刻chunkVFp32的原主人(V 数据)还没被加载(Phase 6 Step 1),借用合法。 - :458-460
DataCopyParams{blockCount, blockLen, srcGap, dstGap}={chunkLen 行, 每行 realK_·2/32 块, 行末跳 (NK_-1)·realK_·2/32 块, UB 侧不跳}。DataCopyParams的四字段单位都是 32B block(blockCount 除外,它是行数),这正是三个% BLOCK_BYTES判定的由来。 - :461-466 事件三段式的前两段:MTE2 即将覆写 UB 暂存区,所以必须先确认(a)标量管道此前对该区间的写已结束(
S_MTE2)、(b)向量管道此前对该区间的读已结束(V_MTE2)。SetFlag与WaitFlag成对:前者由"源侧"打点,后者由"目的侧"等待;中间没有别的指令,所以它实际是"就地阻塞直到上一批动作完成"。 为什么用pipe_->FetchEventID(HardEvent::X_Y)而不是TQueSync<X,Y>:TQueSync是队列级的方向性同步(S/V 之间用它最顺手),而 MTE 与 S/V 的配对在 API 上暴露为硬件事件 ID,需要先从TPipe要一个空闲 ID(FetchEventID),再SetFlag/WaitFlag<HardEvent::X_Y>。事件 ID 是每核有限的资源,所以要FetchEventID而不是写死数字。 - :471
chunkLen * realK_用真实元素数(不含 tail padding):暂存区可能有剩余空间是脏的,但只转 chunkLen 行 × realK_ 列 ⇒chunkKFp32的行 padding 不被本路径写(后续所有读取都用realK_为长度,不触及 padding,见 §7 末尾)。 - :472
PipeBarrier<PIPE_V>:Cast之后立刻返回,函数出口处必须保证向量写完成——因为ComputeKBeta的第一个动作是TQueSync<PIPE_S,PIPE_V>(:498),那是 S→V 方向,挡不住向量写对向量读的可见性。 - :473
return:快路径结束,绝不落入慢路径。 - :474 之后的
for(:476-482)是慢路径:chunkLen × realK_次标量 GM 读 + 标量 UB 写(最坏 8192 次访存/块,比快路径慢一到两个量级)。它直接写 FP32 目标且用alignK_作行宽(:480),所以alignK_ ≠ realK_时行 padding 仍然存在(但不被读)。static_cast<float>(keyGm_.GetValue(...))把类型转换搬到标量侧,省掉Cast与暂存区——这就是慢路径不需要chunkVFp32也能工作的原因。
DataCopyParams 与 DataCopyExtParams 的分工(本区间只用前者,后者见 04 篇 :366):
| 结构 | 字段 | 单位 | 能力 |
|---|---|---|---|
DataCopyParams | {blockCount, blockLen, srcGap, dstGap} | 块数(32B) | 快,但行宽与两侧 gap 必须 32B 整除 |
DataCopyExtParams | {blockCount, blockLen, dstStride, srcStride, srcL1Size} | blockLen 是字节,stride 是字节 | 允许非整块搬运;blockCount 可为 1,配 CopyToGm(:41-57) 做"整块直搬 / 非整块逐块 + PipeBarrier<PIPE_MTE3>"回退 |
CopyToGm 的回退逻辑就是"非对齐回退"的正例:alignElem == elem 时转成 DataCopyParams;否则 blockCount=1 逐行搬、每行一次 MTE3 栅栏。test_non_multiple_reduce_width(docs/PR-406 §4.4)覆盖的正是 dk/dv 非 32B 倍数时 canUseDma=false 的整条慢路径。
6. ComputeGCumsum:块内前缀和(不是全局前缀和)
锚点:chunk_gated_delta_rule.h:486-494
__aicore__ inline void ComputeGCumsum(int32_t t_start, uint64_t head_i, uint32_t chunkLen) {
float gSum = 0.0f;
for (uint32_t i = 0; i < chunkLen; i++) {
gSum += LoadG(t_start + i, head_i);
gCumsumFp32.SetValue(i, gSum);
}
}
// Phase 2: k_beta = K * beta -> kCumdecayFp32.gSum每 chunk 从 0 起 ⇒G_i是块内前缀和。这是 chunkwise 算法的数学要求:所有用到 G 的地方(e^{G_i}、e^{G_i-G_j}、e^{G_{L-1}})都是块内相对量,块间耦合全部由 state 矩阵承担。后果是结果对 chunkSize 不敏感的前提(配合 FP32 workspace 传 state,见 01 篇 §6);若这里是全局前缀和,换 chunkSize 就会改变衰减因子的数值组合顺序。- 纯标量 64 次:加法链不可向量化(真依赖),且长度远小于一条向量指令的 64 元素宽度。
SetValue走标量管道写 UB ⇒ 下游 Phase 2/3 的 S→V 同步与Muls(×1.0)刷新(:516)都是为它服务。 LoadG返回 0 时本函数变成"写 64 个 0",hasGamma_ == 0的整个成本就是 64 次标量写。
7. ComputeKBeta:三处同步各有职责
锚点:chunk_gated_delta_rule.h:495-513
__aicore__ inline void ComputeKBeta(int32_t t_start, uint64_t head_i, uint32_t chunkLen) {
// LoadChunkKey and ComputeGCumsum populate UB through the scalar pipe.
// Synchronize before the first vector read of chunkKFp32.
TQueSync<PIPE_S, PIPE_V> inputSync;
inputSync.SetFlag(0);
inputSync.WaitFlag(0);
TQueSync<PIPE_S, PIPE_V> betaValueSync;
for (uint32_t i = 0; i < chunkLen; i++) {
float beta_val = LoadBeta(t_start + i, head_i);
betaValueSync.SetFlag(0);
betaValueSync.WaitFlag(0);
Muls(kCumdecayFp32[i * alignK_], chunkKFp32[i * alignK_], beta_val, realK_);
}
TQueSync<PIPE_V, PIPE_S> betaSync;
betaSync.SetFlag(0);
betaSync.WaitFlag(0);
}
// Phase 3: precompute exp(gCumsum) and decay_mask[i][j] = exp(g_cumsum[i] - g_cumsum[j]).- :498-500
inputSync(S→V,块级):真正必需。上一函数(LoadChunkKey慢路径 /ComputeGCumsum)用标量管道写了 UB,本函数第一条Muls用向量管道读它。两条管道各有独立的访存队列,标量的 store 可能在向量指令发射后才落地 ⇒ 必须打点。注释(:496-497)写的就是这个理由。 - :501 把
betaValueSync声明在循环外、:504-505 在循环内 Set/Wait:对象只需一次构造(拿同步资源),每行复用。这是本 kernel 的惯用写法;:604 的querySync声明在循环内是反例(§11)。 - :504-505 严格说偏保守:
Muls的第三参数是标量立即数,在指令发射时就被捕获进指令流,不存在"向量去 UB 读 beta"的动作。保留它换来的是"每行一次 S/V 对齐点"的统一形式,与 Phase 5/Step 1 的同构代码一致。真正的代价是每 chunk 64 对事件操作。 - :506 目的/源用
i * alignK_跨距、长度realK_:行跨距与有效长度分离。kCumdecayFp32的行 padding(alignK_ - realK_列,本 commit 常见情形下为 0)从不被写也从不被读;下游全部按realK_长度消费(:580/:624/:638)。 - :508-510
betaSync(V→S,块级收口):本函数自己的产物(kCumdecayFp32)是给向量用的(:580 的Mul),按方向并不需要一个 V→S 点;它的作用是"在返回调用方之前把向量管道排空并立一个可见性点",让随后PrepareDecayAndExp的标量侧与更后面的标量读有一个确定的边界。同类的还有 :626、:642,三处形态完全一致 ⇒ 归类为"保守但对称"的出口同步。
8. PrepareDecayAndExp:exp 预计算与 decay 行的构造
锚点:chunk_gated_delta_rule.h:514-533
__aicore__ inline void PrepareDecayAndExp(uint32_t chunkLen) {
uint32_t cs = chunkSize_;
Muls(gCumsumFp32, gCumsumFp32, 1.0f, cs);
PipeBarrier<PIPE_V>();
Exp(expGCumFp32, gCumsumFp32, cs);
PipeBarrier<PIPE_V>();
TQueSync<PIPE_V, PIPE_S> expSync;
expSync.SetFlag(0);
expSync.WaitFlag(0);
for (uint32_t i = 0; i < chunkLen; i++) {
float gi = gCumsumFp32.GetValue(i);
Duplicate(deltaFp32, gi, cs);
PipeBarrier<PIPE_V>();
Sub(deltaFp32, deltaFp32, gCumsumFp32, cs);
PipeBarrier<PIPE_V>();
Exp(decayMaskFp32[i * cs], deltaFp32, cs);
PipeBarrier<PIPE_V>();
}
}- :516
Muls(×1.0f, cs):gCumsumFp32由标量管道写(:490),紧接着 :518 的Exp由向量管道读 ⇒ 需要"向量管道整体过一遍 +PipeBarrier"。这与 04 篇 :345、06 篇 :898 是同一惯用法。注意长度是cs(64)不是chunkLen:尾块时[chunkLen, 64)区间是上一 chunk 的残留,这里把它们一并过一遍向量管道,保证Exp读的 64 个 lane 都是"向量管道写过的 UB",不残留标量管道未落地状态。 - :518
Exp(expGCumFp32, gCumsumFp32, cs):一条指令算 64 个(硬件 exp 查表 + 插值)。写入 expGCumFp32,长度cs⇒ 同样覆盖到残留区(见下"溢出隐患")。 - :520-522
expSync(V→S):这是全 kernel 关键的一条。此后所有标量侧读gCumsumFp32(:524、:864)、expGCumFp32(:621、:819、:853)、decayMaskFp32(:581 起)之前,都会先越过这一个 V→S 点(或被 :300-302DotFp32内部的同类点覆盖)。删掉它,Phase 5 的gExp与 Step 7 的gLastExp都可能读到旧值。 - :523-531 每行四步:标量取
→ Duplicate广播成[G_i]×64→Sub得deltaFp32[j] = G_i - G_j→Exp写decayMaskFp32第 i 行 64 列。三条PipeBarrier<PIPE_V>是同一条管道上的 WAR/WAW 保序:deltaFp32被连续读写三次(写、读改写、读),下一行的Duplicate也不能早于上一行的Exp读它。 - 每行
Exp的宽度是整cs=64,包括(上三角)。 时 单调不增 ⇒ 上三角的 ⇒ ,即 decay 缓冲里"合法但有害"的值真实存在。 - lower-triangular 到底在哪里构造:
decayMaskFp32本身是完整 64×64 的,并不是下三角;因果性(三角性)完全由四个消费点的循环上界保证::578-579 jj < i(A 的严格下三角)、:591/:607jj <= i(QK 覆盖只写下三角+对角)、:837jj <= i(AccumOutput只累加历史)。所以下三角是"读出来的",不是"造出来的"。这个分工换来一次Duplicate+Sub+Exp覆盖整行(向量效率最高),代价是必须靠纪律(三处上界一致)而非数据布局来守因果性——若上界写错,的未来信息会以 的形式静默参与计算,不会 NaN、不会越界,只是结果错。这是本 kernel 最危险的一处"约定式正确"。 真正被写成下三角的只有两个缓冲: chunkScoresFp32(:581 只写jj<i,:586 写对角 1,上三角保持 0 ——ComputeRecursiveAttn的注释 :316 声明的正是这个不变量)和覆盖后的decayMaskFp32下三角部分。 - exp 的溢出/下溢隐患,四条:
Exp的输入本应 (g 是 log 域衰减)。kernel 与 tiling 都不校验 g 的符号:若 g>0(误用/上游 bug), expGCumFp32可达;FP32 上限约 ,超出即 inf,随后Muls(..., gExp, ...)(:624)把 inf 带进 K_cd,再经Axpy污染 S 与输出,且 inf×0 会产生 NaN 并沿"×1.0 刷新"一路写进 GM。- 反向(很负)时
下溢为 0,数学上正确(完全遗忘),但 decayMaskFp32的上三角会同时出现inf(很正),虽然不读, Exp的硬件实现仍会写inf进 UB —— 排查"矩阵里有 inf 但结果对"时不要误判。 length = cs而非chunkLen:尾块残留区(上一 chunk 的值)也会被 exp。残留 一定比当前 chunk 的更负(单调性沿 token 累积),所以最坏是更严重的下溢 ⇒ 0,安全;但如果上游给的 g 变号,这条保护就不存在。 Sub(deltaFp32, deltaFp32, gCumsumFp32, cs)要求deltaFp32 ≥ cs元素:deltaSize = max(dotProductSize, vStepA) ≥ cs(:186-187,dotProductSize = max(stateStrideK_, cs))恒成立 ⇒ 缓冲复用是安全的,这个不等式是它唯一的守卫。
可优化点(读后续 PR 时对照)::523-531 的 64 次"广播-减-exp"完全可以展平成一次 4096 元素的 Sub + 一次 4096 元素 Exp(decayMaskFp32 自身容量就是 4096)。作者选按行是为了不占用第二块 16KB UB(deltaFp32 只有 128 元素)。
9. CanCacheAttnQuery:借地盘的资格判定
锚点:chunk_gated_delta_rule.h:534-547
__aicore__ inline bool CanCacheAttnQuery(uint32_t chunkLen) {
uint64_t queryElements = static_cast<uint64_t>(chunkLen) * realK_;
uint64_t stagingCapacity = static_cast<uint64_t>(2) * chunkSize_ * vStepAligned_;
uint64_t cacheCapacity = static_cast<uint64_t>(chunkSize_) * vStepAligned_;
uint64_t gmRowStride = static_cast<uint64_t>(NK_) * realK_;
uint64_t srcGapBytes = (gmRowStride - realK_) * sizeof(inType);
return chunkLen <= UINT16_MAX && alignK_ == realK_ && queryElements <= stagingCapacity &&
queryElements <= cacheCapacity && (realK_ * sizeof(inType)) % BLOCK_BYTES == 0 &&
srcGapBytes % BLOCK_BYTES == 0 && srcGapBytes / BLOCK_BYTES <= UINT16_MAX;
}(544 行为空行;545-547 是 PrepareAttnQueryCache 的文档注释。)
- 与
canUseDma(:452-454)的差别只有一项:多了queryElements <= cacheCapacity。因为 Q 缓存要同时满足两个容量:FP16 暂存(借chunkVFp32,2·cs·vStepA个 FP16)与 FP32 落点(借chunkAttnOutFp32,cs·vStepA个 FP32)。后者才是真紧的约束:cacheCapacity恰是stagingCapacity的一半,所以stagingCapacity那条在这里其实恒真(被cacheCapacity蕴含)——保留它是为了与LoadChunkKey的判定保持同形。 cacheCapacity = chunkSize_ · vStepAligned_:Q 缓存能否住进输出缓冲,取决于realK_ ≤ vStepAligned_(chunkLen ≤ cs 时)。典型 dk=128、vStep=128 ⇒ 成立;dk=256 而 vStep=64 ⇒ 不成立,Phase 4 走慢路径(每行现读 GM)。所以这个谓词就是"快路径 QK 点积"和"慢路径 QK 点积"的唯一分岔。- 该函数是纯函数(只读成员与
chunkLen),所以在 :549(做决定)与 :588(重看决定)两次调用必然同值。代价:它内部有 5 次 64 位乘除,被调用两次;写成"算一次、存个 bool"更省,但要多一个成员变量。作者选了无状态方案,避免成员膨胀(见 07 篇标量表:确实没有对应的qCached_成员)。 NK_ - 1在 :539 通过gmRowStride - realK_表达,而 :554 直接写NK_ - 1。二者等价的前提是NK_ ≥ 1;若NK_ == 0,:554 的 uint32 减法会下溢成巨大值而 :539 的 uint64 减法得到巨大正数后* sizeof也错——两处都不可靠,但NK_==0在合法 shape 下不可能(infershape 保证 head 数 ≥ 1)。
10. PrepareAttnQueryCache:借 chunkVFp32 做 FP16 暂存
锚点:chunk_gated_delta_rule.h:548-575
// 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_;
uint64_t srcGapBytes = static_cast<uint64_t>(NK_ - 1) * realK_ * sizeof(inType);
DataCopyParams queryParams{static_cast<uint16_t>(chunkLen),
static_cast<uint16_t>(realK_ * sizeof(inType) / BLOCK_BYTES),
static_cast<uint16_t>(srcGapBytes / BLOCK_BYTES), 0};
event_t scalarToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE2));
SetFlag<HardEvent::S_MTE2>(scalarToMte2);
WaitFlag<HardEvent::S_MTE2>(scalarToMte2);
event_t vectorToMte2 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2));
SetFlag<HardEvent::V_MTE2>(vectorToMte2);
WaitFlag<HardEvent::V_MTE2>(vectorToMte2);
DataCopy(queryLocal, queryGm_[qkOff], queryParams);
event_t mte2ToVector = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(mte2ToVector);
WaitFlag<HardEvent::MTE2_V>(mte2ToVector);
Cast(chunkAttnOutFp32, queryLocal, RoundMode::CAST_NONE, chunkLen * realK_);
PipeBarrier<PIPE_V>();
Muls(chunkAttnOutFp32, chunkAttnOutFp32, scale_, chunkLen * realK_);
PipeBarrier<PIPE_V>();
}(573 空行;574-575 为 Phase 4 的文档注释。)
- 与
LoadChunkKey快路径逐行同构(借chunkVFp32、同样四元组、同样三段事件、同样扁平Cast),差别只有两处:Cast的目的地是chunkAttnOutFp32(:568)而不是chunkKFp32;- 多了一步
Muls(..., scale_, ...)(:570)。
- :552 本 PR 最著名的借用:Q 的 FP16 暂存放进
chunkVFp32。合法性链条:chunkVFp32的 FP32 内容上一个主人是 Phase 1 的 K 暂存(同一段地址),Phase 1 的Cast(:471)已经把数据取出并PipeBarrier(:472);chunkVFp32真正的 FP32 身份(v_beta)要到 Phase 6 Step 1(:676)才被写。所以 Phase 3.5 落在这两个使用之间的空档里。 scale_为什么在这里折进缓存:Phase 4 的 QK 点积每 chunk 要做约chunkLen²/2次,每次的系数是。 scale是全局常量,折进 Q 缓存后点积侧完全不出现 scale;对比慢路径 :602 是在标量侧逐元素乘scale_(每行 realK_ 次标量乘)。同一个 scale 在两条路径里的施加位置不同:向量批量 vs 标量逐元素,这正好是"缓存换标量算力"的定价示例。- :570 的
Muls长度是chunkLen * realK_(真实元素,扁平),不是alignK_行宽——再次依赖CanCacheAttnQuery里的alignK_ == realK_。 - :571 的
PipeBarrier<PIPE_V>之后紧接 Phase 4 的DotFp32(向量读chunkAttnOutFp32)——同管道,barrier 足够,不需要跨管事件。
11. ComputeAttnMatrix:两次 O(n²) 点积 + 一次递归 + 一次原地覆盖
锚点:chunk_gated_delta_rule.h:576-616
// 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_;
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);
}
}
}
}四段职责:
- :578-583 直接项
。 ,其中 就是 kCumdecayFp32第 i 行(Phase 2 的产物,身份β_i k_i)。取负在标量侧做(-dot * decay)。行 0 不生成(i从 1 起),因为第 0 行没有,其值将只由 :586 置 1。 - :584 递归修正(04 篇 §7):把严格下三角
变成 。这是"用递推替代求逆"的唯一实现点。 - :585-587 单位对角:
A[i][i]=1,与 :584 之前的 0 对角合起来构成的单位下三角结构。注意它排在递归之后:递归过程要求对角保持 0( Axpy的前缀长度k天然不碰下标 ≥ k,但 :320 的系数缓存会读buf[i*ld+k],若对角已是 1 就会把当系数用)——顺序不可调换。 - :588-613 QK 注意力并原地改写
decayMaskFp32:(快路径里 已含 scale;慢路径 :602 显式乘 scale_)。此后 :838 读到的decayMaskFp32身份已是 QK 矩阵。读改写都在标量侧(:593-594、:609-610),同一线程内顺序确定,无竞争。
其它要点:
- :581 与 :593-594 都读
decayMaskFp32,但用途不同:前者读的是纯 decay,后者读后立刻覆盖成 QK 矩阵。同一份 decay 被用两次(A 的直接项 + QK 的权重),所以覆盖必须排在 :578-583 之后 —— 这就是"用完即毁"的准确表述:毁点在 Phase 4 第四段,而不是 Phase 3 之后。 - :580/:592 传的是
LocalTensor行视图(kCumdecayFp32[i*alignK_],operator[]按元素偏移返回同类型视图),零拷贝;:590 亦然。这三处是LocalTensor作为"带基址的视图值类型"的典型用法。 - 慢路径 :601-603 把 GM 的 Q 行逐元素乘 scale 后写进
deltaFp32(容量deltaSize ≥ stateStrideK_ ≥ realK_✓),:604-606 的TQueSync<PIPE_S,PIPE_V>是真正必需的(标量写 UB → :608Mul向量读 UB)。但对象声明在 i 循环内部:每行重新构造/析构一个同步对象,比 §7 的循环外声明多 64 次开销,语义完全相同。两处写法不一致,属可统一的小瑕疵。 - :608 的
DotFp32(deltaFp32, ...):lhs参数是deltaFp32,而DotFp32内部 :290 会把它和rhs的积写进dotProductFp32(另一块缓冲),所以deltaFp32不会被归约覆盖 ⇒ 一行 Q 可以连续和 i+1 个 K 行做点积。这是deltaFp32/dotProductFp32必须分成两块的原因(若在同一个缓冲里做 in-place Mul,第 2 个点积起就把第 1 个的输入踩掉了)。 - 快/慢路径互斥且都在函数内完成同一数学,所以 :614 之后无需任何"哪条路径"的状态传递。
12. ComputeKCumdecay:K_cd = A · diag(e^G) · K_beta
锚点:chunk_gated_delta_rule.h:617-648
// Phase 5: k_cumdecay = attn @ (k_beta * exp(g_cumsum)); result stored back into chunkKFp32.
__aicore__ inline void ComputeKCumdecay(uint32_t chunkLen) {
uint32_t cs = chunkSize_;
TQueSync<PIPE_S, PIPE_V> expValueSync;
for (uint32_t i = 0; i < chunkLen; i++) {
float gExp = expGCumFp32.GetValue(i);
expValueSync.SetFlag(0);
expValueSync.WaitFlag(0);
Muls(kCumdecayFp32[i * alignK_], kCumdecayFp32[i * alignK_], gExp, realK_);
}
TQueSync<PIPE_V, PIPE_S> decaySync;
decaySync.SetFlag(0);
decaySync.WaitFlag(0);
TQueSync<PIPE_S, PIPE_V> coefficientSync;
for (uint32_t i = 0; i < chunkLen; i++) {
LocalTensor<float> outRow = chunkKFp32[i * alignK_];
Duplicate(outRow, 0.0f, realK_);
PipeBarrier<PIPE_V>();
for (uint32_t k = 0; k <= i; k++) {
float score = chunkScoresFp32.GetValue(i * cs + k);
coefficientSync.SetFlag(0);
coefficientSync.WaitFlag(0);
Axpy(outRow, kCumdecayFp32[k * alignK_], score, realK_);
PipeBarrier<PIPE_V>();
}
}
TQueSync<PIPE_V, PIPE_S> outputSync;
outputSync.SetFlag(0);
outputSync.WaitFlag(0);
}- 前半 :620-625:
kCumdecayFp32 ← kCumdecayFp32 · diag(e^{G_i}),即每行乘自己的 ,原地(dst 与 src 同一行)。原地在此安全: Muls是"读完一行再写同一行"的向量指令,同管道内自序,且PipeBarrier被块级decaySync(:626-628)覆盖。gExp从expGCumFp32标量读——它的可见性由 Phase 3 的 :520-522 V→S 点保证,而 :622-623 的expValueSync是 S→V 方向,不能替代它。这是一对常被读混的相邻同步:一进一出,方向不同、保护对象不同。 - :626-628
decaySync(V→S):块级收口(与 :508-510 同形)。 - :629-641 后半是标准软件 GEMV:外层输出行
i,内层累加项k ≤ i(下三角含对角),系数A[i][k]标量取,源行K'_beta[k]用Axpy累加。这是本 kernel 里矩阵乘的规范形态,06 篇的 Step 2/4/5/6/7 全是它的变体。k ≤ i的上界:是单位下三角, 恒 0,跳过即省一半。 - :632
Duplicate(outRow, 0.0f, realK_):必须先清零,因为outRow是chunkKFp32的行,里面还躺着原始 K,而Axpy是"加"不是"写"。清零长度realK_⇒ 行 padding 保留原 K 的 padding(永远不被读,见 §5)。 - :633
PipeBarrier<PIPE_V>:Duplicate 与随后第一条 Axpy 写同一行,同管道必须排空(否则 Axpy 可能读到未清零的行)。 - :634-640:每次
Axpy前一次标量GetValue、后一次PipeBarrier<PIPE_V>⇒条 Axpy、同样多的同步。与 ComputeRecursiveAttn合计,K 侧每 chunk 约 4000 条 Axpy + 4096 次点积,这就是"向量化基线"的延迟底账(docs/PR-406 §7 ①)。
- :642-644
outputSync(V→S):这条不是保守装饰。Phase 6 Step 4(:796)要用标量GetValue读chunkKFp32(此时身份已是 K_cd)。中间隔着PrepareAttnQueryCache? 不——顺序是 Phase 5 →ProcessVTiles→ProcessOneVTile→ Step1/2/3 → Step4,中间 Step 2 有 S→V 与 Step 3 的 DMA 事件,但没有任何 V→S 点,所以 :642-644 是 :796 那次标量读的唯一可见性来源。删掉它,K_cd 可能读到旧 K ⇒ 结果错但不崩。这是"必须存在"的一类跨函数同步,值得在改动时重点守。 chunkSize_在 :618 被取出成cs,仅用于 :635 的行跨距(的 ld恒为chunkSize_,与chunkLen无关)。
13. 本区间的同步点总表(回答"每处配对为什么需要")
| 行 | 类型 | 方向 | 保护的具体数据流 | 严格必要性 |
|---|---|---|---|---|
| :461-466 | 事件 S_MTE2+V_MTE2 | S/V→MTE2 | MTE2 即将覆写 chunkVFp32 暂存区 | 必需(写-写/读-写竞争) |
| :468-470 | 事件 MTE2_V | MTE2→V | Cast 要读刚搬进来的 K | 必需 |
:498-500 inputSync | TQueSync | S→V | 标量写的 chunkKFp32(:480) 给 Muls(:506) 读 | 必需(慢路径下) |
:504-505 betaValueSync | TQueSync | S→V | 标量 beta → Muls 立即数 | 保守(发射时捕获) |
:508-510 betaSync | TQueSync | V→S | Phase 2 出口收口 | 保守 |
| :516-517 | Muls×1.0+barrier | V 内 | 标量写的 gCumsumFp32 落进向量管道视图 | 必需(与 :520 合起来) |
| :519 | barrier | V 内 | Exp 完成才可读 expGCumFp32 | 必需 |
:520-522 expSync | TQueSync | V→S | 向量写的 expGCumFp32/gCumsumFp32 给后续所有标量读(:524/:621/:819/:853/:864) | 必需,跨函数 |
| :526/:528/:530 | barrier ×3 | V 内 | deltaFp32 同管道连续读写 | 必需 |
| :558-563 / :565-567 | 事件三段式 | 同上 | Q 的 DMA | 必需 |
| :571 | barrier | V 内 | Q 缓存写完给 Phase 4 向量读 | 必需 |
:604-606 querySync | TQueSync | S→V | 标量写的 deltaFp32(Q 行) 给 DotFp32 的 Mul 读 | 必需(慢路径) |
:622-623 expValueSync | TQueSync | S→V | 标量 gExp → Muls | 保守 |
:626-628 decaySync | TQueSync | V→S | Phase 5 前半出口收口 | 保守 |
:636-637 coefficientSync | TQueSync | S→V | 标量 score(A[i][k]) → Axpy | 保守(每 Axpy 一次,2080 次/chunk) |
| :633/:639 | barrier ×2 | V 内 | 行清零/行累加 | 必需 |
:642-644 outputSync | TQueSync | V→S | K_cd 给 :796 的标量读 | 必需,跨函数 |
规律:所有真正必需的同步都发生在"数据跨过管道边界"的那一刻(标量↔UB、UB↔DMA);所有"每条 Axpy 前一次 S→V"都是把同一件事做 N 次的保守版。合并后者是零风险优化,合并前者会直接错——这条界线是评估后续 PR 改动安全性的尺子。