10b · Step7 StateUpdate 的 Cube 主循环(2094-2289)
本章是
10-Output与StateUpdate的Cube融合.md的接续篇。!1135 的 1918-2410 这一段
0. 结论先行
锚点:chunk_gated_delta_rule.h:2094-2289(本章;2290-2410 见 10c 章)
- Step7 的 Cube 化只做了一半的乘法,另一半留给 Vector。 唯一进 Cube 的是
(K·exp(g_last-g))^T @ v_new(:2230、:2245共 4 条 Mmad);exp(g_last)·state这一项在 UB 上用一条Muls做掉(:2123), 因为它和 Cube 输出是相加关系,不需要也不应该进 L0C。 kTileM的分界是 96 而不是别的数::2104写作constexpr uint32_t kTileM = (kSpecializedDk <= 96) ? kSpecializedDk : 64;, 于是kMTileCount(:2106)在 64/80/96 桶里都是 1,只有 128 桶是 2。 可机械证明的必然而非作者断言:融合版 NZ→ND 加法(:2265-2271)要求kStateRows*16 == (kStateRows/16)*256,即 L0C 落地缓冲必须正好容纳"一整列 fractal"; 一旦 M 被切成 64 行的 tile,每个 tile 的 NZ 只有curTileM/16段 fractal,nBlock之间的步长从kStateRows*16变成curTileM*16,融合 Add 的常量repStride就不成立, 只能退回:2274-2286的逐行逐块分段加。- K^T 的转置完全发生在 UB,不进 Cube。
dav_m200的 Mmad 只接受 NZ 排布的 A/B, 没有"转置位"可设;所以:2149-2166与:2177-2190用两轮TransDataTo5HD(16×16 fractal 转置指令)把[64, dk]的 K 摆成[dk, 64]再 DMA 到 GM slot, hi/lo 两份各转一次。 - 标量
gLastExp与 Cube 结果的先后是"先衰减、后累加",且中间隔着两次全停。:2121-2123在 staging 之前就把exp(g_last)乘进stateInFp32; Cube 结果直到:2258才落到cNz、:2269/:2283才加到 state 上。:2193与:2259两道PipeBarrier<PIPE_ALL>()保证 V 侧衰减与 M 侧写回不交叠。 反过来若把Muls挪到Add之后,Cube 项会被二次缩放,语义立即错。 - A 侧拆 hi/lo、B 侧只量化一次,是本章唯一的精度妥协点。
:2147/:2171-2175走完整的"Cast→反向 Cast→Sub→再 Cast"补偿链, 而:2192的 B(v_new)用的是只 Cast 的StageHalf(:1581-1589)。chunkVFp32里的v_new本身是 Step4 的 FP32 Cube 结果(ComputeValueAndVNewCube的 L0C 读回), 所以它进 Cube 前被量化了一次;全文件只有MatmulBlockFp32Compensated(:454/:471)对 B 也拆 hi/lo。 docs/PR-1135 §2.2 把这一处描述成"无损",与快照不符。 - 两条 GM 写回出口(FP32 workspace 与 FP16
final_state转置)、transposeCapacity的贴边算术、 以及"一次算两条出口"的成本账,全部在 10c 章(2290-2410)逐行展开; 本章只到ComputeStateUpdateCube的右括号:2289为止。
1. Dispatch:一个 if constexpr 就把整条 Cube 路径编进去
锚点:chunk_gated_delta_rule.h:2094-2099
// Fixed-shape state update:
// state = exp(g_last) * state + (K * exp(g_last - g))^T @ v_new.
// The weighted K^T @ v_new product uses Cube. The 64/80/96 buckets convert the
// complete NZ result to contiguous ND and use one wide Vector Add; the 128
// bucket retains the segmented fallbac
__aicore__ inline void ComputeStateUpdateCubeDispatch(int32_t t_start, uint64_t qkHead) {
if constexpr (kSpecializedDk != 0) {
ComputeStateUpdateCube<kSpecializedDk, kSpecializedDk>(t_start, qkHead);
}
}:2094签名只带t_start与qkHead:Step7 需要的其余几何量 (chunkLen、avFp32、v_i、curV)在快路径下全部退化为模板常量,不需要传。:2095if constexpr (kSpecializedDk != 0):桶 0 是"未知形状/纯 Vector"实例, 这一行的常量折叠让桶 0 里整个 Cube 函数体都不生成代码, 这也是ComputeStateUpdateCube里敢用constexpr uint32_t数组尺寸的前提。:2096模板实参kSpecializedDk, kSpecializedDk同时喂给kStateRows与kMatmulN: 即**"状态的行 = K 的列 = 输出矩阵的列 = 桶宽"**,这是整个 Step7 固定形状的根。:2097-2098空 else(无else分支)与右括号;桶 0 直接掉回调用点的 Vector 分支。
2. 模板参数、M 拆分决策与 GM slot 预算
锚点:chunk_gated_delta_rule.h:2100-2119
template <uint32_t kStateRows, uint32_t kMatmulN>
__aicore__ inline void ComputeStateUpdateCube(int32_t t_start, uint64_t qkHead) {
// The 64/80/96 buckets fit one complete M tile in L0C; the 128 bucket keeps
// the conservative 64-row split.
constexpr uint32_t kTileM = (kSpecializedDk <= 96) ? kSpecializedDk : 64;
constexpr uint32_t kMatmulK = 64;
constexpr uint32_t kMTileCount = (kStateRows + kTileM - 1) / kTileM;
constexpr uint32_t kAElements = kStateRows * kMatmulK;
constexpr uint32_t kATileElements = kTileM * kMatmulK;
constexpr uint32_t kBElements = kMatmulK * kMatmulN;
constexpr uint32_t kStateElements = kStateRows * kMatmulN;
constexpr uint32_t kCTileElements = kTileM * kMatmulN;
__gm__ uint8_t *stageBase = GetCubeStageBase(0);
GlobalTensor<half> aStageGm;
GlobalTensor<half> bStageGm;
GlobalTensor<half> aResidualStageGm;
aStageGm.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(stageBase), kAElements);
bStageGm.SetGlobalBuffer(reinterpret_cast<__gm__ half *>(stageBase + kAElements * sizeof(half)), kBElements);
aResidualStageGm.SetGlobalBuffer(
reinterpret_cast<__gm__ half *>(stageBase + (kAElements + kBElements) * sizeof(half)), kAElements);:2100-2101template <uint32_t kStateRows, uint32_t kMatmulN>:名字与 10 章ComputeOutputCube<kStateK, kMatmulN>(:1924)不同,因为这里kStateRows是 A 的行数(dk) 而不是 K 的长度;Step7 的 K 长度恒为 64(chunkSize)。:2102-2103注释即 §0 第 2 条的原文依据:The 64/80/96 buckets fit one complete M tile in L0C; the 128 bucket keeps the conservative 64-row split."conservative" 一词说明作者当时也没有把 128 拆开的硬件必然性写死, 但:2265的kNzFractalElements = kStateRows * 16让这条注释成了代码层面的约束。:2104本章最关键的一行:kTileM = (kSpecializedDk <= 96) ? kSpecializedDk : 64。 三桶取自身(64/80/96),128 桶取 64。64 是唯一同时满足三条约束的常数: L0A/L0B 的 fractal 行对齐(16 的倍数)、TransDataTo5HDParams.repeatTimes为uint8_t、 以及 10 章已确立的"64 行 = 一个 chunk 的 token 数"。:2105kMatmulK = 64:乘法的 K 维是 chunkSize,不是 dk。:2106kMTileCount = (kStateRows + kTileM - 1) / kTileM:向上取整。四桶实测值:| -
kSpecializedDk | kTileM | kMTileCount | kStateRows 整除? |
|---|---|---|---|
| 0(纯 Vector 桶) | — | 进此函数(:2095 常量折叠掉) | |
| 64 | 64 | 1 | 是 |
| 80 | 80 | 1 | 是 |
| 96 | 96 | 1 | 是 |
| 128 | 64 | 2 | 是(128 = 2×64) |
注意 80 桶 kTileM = 80:80 + 79 = 159,159 / 80 = 1。 所以只有 128 需要拆,且拆完正好两刀,:2217 的 curTileM 永远取满 kTileM, 不存在第三块残 tile——这也是 :2250 敢用 kSpecializedDk <= 96 而不是 kMTileCount == 1 的原因。
:2107-2111五个元素数常量,逐个含义:kAElements = kStateRows*kMatmulK= 整个 A 矩阵(dk×64)元素数;kATileElements = kTileM*kMatmulK= 单个 M-tile 的 A;kBElements = kMatmulK*kMatmulN= 64×dk 的 B(v_new);kStateElements = kStateRows*kMatmulN= dk×dk 的 C/state,FP32 元素数;kCTileElements = kTileM*kMatmulN= 单 tile 的 L0C 占用。:2112GetCubeStageBase(0):与 10 章同一 slot0(:249-255),即workspace + stateWorkspace 字节数 + blockIdx*2*64KB;每个 kernel 内所有 Cube 函数共用 slot0, 靠执行顺序串行复用(详见 11 章风险清单)。:2113-2119三段 FP16 板;:2116A 占kAElements,:2117B 紧随其后偏移kAElements*sizeof(half),:2118-2119A 残差再偏移(kAElements+kBElements)*sizeof(half)。 四桶的 slot 预算(CUBE_STAGE_SLOT_BYTES = 64*1024,:39):桶 A B A_lo 合计 64 KB 余量 64 64·64·2 = 8192 64·64·2 = 8192 8192 24576 B 40960 B 80 80·64·2 = 10240 64·80·2 = 10240 10240 30720 B 34816 B 96 96·64·2 = 12288 64·96·2 = 12288 12288 36864 B 28672 B 128 128·64·2 = 16384 64·128·2 = 16384 16384 49152 B 16384 B 与 10 章
ComputeOutputCube的 128 桶合计 65536 B(贴死)不同, Step7 的 A 是dk×64而 Output 的 A 是64×dk,B 却是同一个形状—— 两者只是行列表述相反,字节数相同,故 Output 侧 A/B/A_lo = 16384+32768+16384。
3. exp(g_last) 标量与衰减:先于任何 Cube 动作
锚点:chunk_gated_delta_rule.h:2120-2130
float gLast = gCumsumFp32.GetValue(kMatmulK - 1);
float gLastExp = expGCumFp32.GetValue(kMatmulK - 1);
Muls(stateInFp32, stateInFp32, gLastExp, kStateElements);
Duplicate(deltaFp32, gLast, kMatmulK);
PipeBarrier<PIPE_V>();
Sub(deltaFp32, deltaFp32, gCumsumFp32, kMatmulK);
PipeBarrier<PIPE_V>();
Exp(deltaFp32, deltaFp32, kMatmulK);
PipeBarrier<PIPE_V>();:2120空行(函数体第一段结束于:2119)。:2121gLast = gCumsumFp32.GetValue(kMatmulK - 1):快路径下chunkLen == 64(:259), 所以索引 63 就是本块最后一个 token 的g_cumsum。用kMatmulK-1而不是chunkLen-1, 把运行期长度换成编译期常量,是桶化后的通用手法(对照 fallback:2298仍用chunkLen - 1)。:2122gLastExp = expGCumFp32.GetValue(63):不重算 exp,直接取 Step2 预计算的表 (PrepareDecayAndExp:907写好的Exp(expGCumFp32, gCumsumFp32, chunkLen))。:2123Muls(stateInFp32, stateInFp32, gLastExp, kStateElements): 一条 V 指令给整块 dk×dk 状态乘上衰减,长度用kStateElements(含 padding 列)。 这里就是 §0 第 4 条说的"先衰减"。 注意它是原地的:state 只有这一条退路,因为stateInFp32与chunkScoresFp32共用 tmpBuff 同一基址(:175-176),但 scores 在 Step7 之前已死。:2124Duplicate(deltaFp32, gLast, 64):把标量gLast广播成 64 个 float 的向量, 这就是"标量→向量"的标准落法:不是把标量直接喂给Muls,而是先在 V 上造一个等长向量。:2125PipeBarrier<PIPE_V>()::2123的Muls与:2124的Duplicate都写 UB, 下一条Sub要读deltaFp32,V 管内部必须收口。:2126Sub(deltaFp32, deltaFp32, gCumsumFp32, 64):算g_last - g_i(逐 token)。:2127又一条 V 栅栏,:2128Exp(deltaFp32, deltaFp32, 64):V 管的指数单元 (硬件是查表+多项式拟合,与:364 ScalarExp的 8 项泰勒展开不同精度路径)。:2129PipeBarrier<PIPE_V>();:2130空行。 这三行Duplicate/Sub/Exp全部在deltaFp32上原地完成, 与 10 章ComputeAttnInter的 fallback(:2307用ScalarExp逐 token 算)形成 "一次批量 vs 64 次标量"的直接对照。
4. K 板:借用死掉的输出缓冲,逐行乘衰减系数
锚点:chunk_gated_delta_rule.h:2131-2146
// Load K in [64,128] ND order, apply one decay coefficient per row, then
// transpose the FP16 high/low parts to the [128,64] A matrix.
LocalTensor<inType> keyLocal = chunkKFp32.template ReinterpretCast<inType>();
uint64_t keyOffset = (static_cast<uint64_t>(t_start) * NK_ + qkHead) * realK_;
LoadPaddedRows(chunkAttnOutFp32, keyLocal, keyGm_, keyOffset, kMatmulK, NK_ * realK_);
TQueSync<PIPE_V, PIPE_S> decaySync;
decaySync.SetFlag(0);
decaySync.WaitFlag(0);
for (uint32_t row = 0; row < kMatmulK; ++row) {
float rowScale = deltaFp32.GetValue(row);
Muls(chunkAttnOutFp32[row * kStateRows], chunkAttnOutFp32[row * kStateRows], rowScale, kStateRows);
}
PipeBarrier<PIPE_V>();
LocalTensor<half> stageLocal = stateOutQueue_.AllocTensor<half>();
LocalTensor<half> transposedLocal = kCumdecayFp32.template ReinterpretCast<half>();:2131-2132注释是本章的第二个"作者自述":先按 [64, dk] ND 载入 K、逐行乘一个衰减系数、 再把 FP16 高低两份转置成 [dk, 64] 的 A。三句话对应 §4/§5/§6 三个小节。:2133keyLocal = chunkKFp32.ReinterpretCast<inType>(): 借chunkKFp32当 FP16 staging(ReinterpretCast只改元素类型、不改地址)。 此刻chunkKFp32的身份是 10 章:1946留下的"FP16 Q 板",早已死。 容量核算:写入rows*alignK_ = 64*128 = 8192个 half = 16384 B,chunkKFp32容量cs*alignK = 8192个 float = 32768 B,只用掉一半(10 章 §4 同一结论)。:2134keyOffset = (t_start*NK_ + qkHead)*realK_:与 10 章:1947逐字符同构,只是keyGm_换queryGm_。:2135LoadPaddedRows(chunkAttnOutFp32, keyLocal, keyGm_, keyOffset, kMatmulK, NK_*realK_): dst 是chunkAttnOutFp32——刚被WriteAttnTileToGm(:1326)写空的输出板, 现在当作 [64, dk] 的 FP32 K。顺序由ProcessOneVTile保证(:1326在:1327之前)。:2136-2138TQueSync<PIPE_V, PIPE_S> decaySync+SetFlag(0)/WaitFlag(0): V 写 → S 读的跨管握手。因为:2140要用标量管从deltaFp32里GetValue。 这个方向(V→S)在全文件 12 处PIPE_V,PIPE_S中一致::401/:528/:648/:896/:937/:1156/:1188/:1207/:1390/:1517/:2136/:2398全部紧接在标量读/写 UB 之前; 而 6 处PIPE_S,PIPE_V(:330/:535/:870/:1004/:1230/:1528)全部紧跟在标量SetValue之后。 10 章 §10.3 里被 !1135 删掉的TQueSync<PIPE_S,PIPE_V> coefficientSync属于"标量只读、向量消费寄存器值", 在这个不变式里本来就不需要握手——同文件:1185-1186(GetValue后立刻Muls,无握手)是既存的旁证。:2139-2142逐行衰减:rowScale = deltaFp32.GetValue(row)(S 管),Muls(chunkAttnOutFp32[row*kStateRows], ..., rowScale, kStateRows)(V 管)。 64 次标量读 + 64 条宽 V 乘,取代 fallback 里"每个 (i,d) 一次ScalarExp"的 8192 次标量指数。 行长写kStateRows(= dk = alignK_)而不是realK_:padding 列也必须被乘, 虽然 0×系数仍是 0,但宽度对齐才能让:2141的地址步进与:2147的kAElements自洽。:2143PipeBarrier<PIPE_V>();:2144空行。:2145stageLocal = stateOutQueue_.AllocTensor<half>():stateOutQueue_在 !1135 里已经不是"状态输出队列"了,而是全kernel通用的 FP16 单槽中转队列 (声明:2423,TQue<QuePosition::VECOUT, 1>,BUFFER_NUM=1)。 队列容量在:138-143取max(stateTileBytes, chunkOutBytes): 128 桶 =max(128*128*2, 64*128*2) = 32768 B = 16384 half ≥ kAElements=8192✓。:2146transposedLocal = kCumdecayFp32.ReinterpretCast<half>(): 借 kCumdecay 当转置落地点(容量2*64*128 = 16384 half ≥ 8192✓)。 此刻kCumdecayFp32的身份是 10 章:2041的cNz(attn 的 NZ 落地板),已死。
5. 第一轮转置:hi 分量的 [64,dk] → [dk,64]
锚点:chunk_gated_delta_rule.h:2147-2166
Cast(stageLocal, chunkAttnOutFp32, RoundMode::CAST_NONE, kAElements);
PipeBarrier<PIPE_V>();
TransDataTo5HDParams transposeParams;
transposeParams.repeatTimes = static_cast<uint8_t>(kStateRows / 16);
transposeParams.srcRepStride = 1;
transposeParams.dstRepStride = kMatmulK;
for (uint32_t block = 0; block < kMatmulK / 16; ++block) {
LocalTensor<half> srcRows[16];
LocalTensor<half> dstRows[16];
for (uint32_t row = 0; row < 16; ++row) {
srcRows[row] = stageLocal[block * 16 * kStateRows + row * kStateRows];
dstRows[row] = transposedLocal[block * 16 + row * kMatmulK];
}
TransDataTo5HD<half>(dstRows, srcRows, transposeParams);
}
PipeBarrier<PIPE_V>();
event_t vectorToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE3));
SetFlag<HardEvent::V_MTE3>(vectorToMte3);
WaitFlag<HardEvent::V_MTE3>(vectorToMte3);
DataCopy(aStageGm, transposedLocal, kAElements);:2147Cast(stageLocal, chunkAttnOutFp32, RoundMode::CAST_NONE, kAElements): FP32→FP16,CAST_NONE= 工具链默认 RNE 舍入(本仓库从不指定舍入模式, 与:1569/:1573的残差 Cast 同一用法)。kAElements = dk*64 = 8192,注意元素数写的是转置后的形状,与转置前字节数相同,所以可以直接搬。:2148V 栅栏;:2149TransDataTo5HDParams transposeParams;。:2150repeatTimes = kStateRows/16:128 桶 = 8,且该字段是uint8_t(对照:1479的显式kBlockCount <= UINT8_MAX守卫),8 远未贴边。:2151srcRepStride = 1:源侧每多一个 repeat,指针前进 1 个 32B 块 = 16 个 half, 即"同一行里的下一个 16 列"。:2152dstRepStride = kMatmulK:目的侧每多一个 repeat,前进kMatmulK个块 =64*32 B = 1024 half= 正好一行 A(宽度 64)的 16 倍,即"下一个 16 行"。 这两个 stride 合起来就是"16×16 fractal 转置"的地址规则;TransDataTo5HD的跨距单位在快照里没有注释,本节按 32B 块解释, 依据是同文件CopyToGm(:49-50)、LoadPaddedRows(:289-291)对DataCopyParams的块单位用法。:2153-2161外层block循环 = A 的列块(kMatmulK/16 = 4次), 每次准备 16 个源行指针srcRows[row] = stageLocal[block*16*kStateRows + row*kStateRows](A_hi 的第 row 行、第 block 组 16 列)与 16 个目的行指针dstRows[row] = transposedLocal[block*16 + row*kMatmulK](转置矩阵的第 block 组 16 列、第 row 行)。:2160TransDataTo5HD<half>(dstRows, srcRows, transposeParams):一条指令摆 16 个 fractal。 写入上界:(block=3)*16 + (repeat=7)*16*64 + row*64 + 16 = 48 + 7168 + 960 + 16 = 8192✓ 正好kAElements。:2162PipeBarrier<PIPE_V>():转置是 V 管,后面 MTE3 要读它。:2163-2165V_MTE3事件SetFlag+WaitFlag成对: 语义是"V 已写完 → 通知 MTE3 → 原地等 MTE3 接手完成", 与StageHalfWithResidual:1564-1568的MTE3_V惯用法互为镜像。:2166DataCopy(aStageGm, transposedLocal, kAElements):MTE3 出站, 把转置好的 A_hi 写进 GM slot0 第一段。 为什么要绕 GM 一圈?dav_m200上 L1 的 NZ 装载只有DataCopy(LocalTensor<L1>, GlobalTensor, Nd2NzParams)这一条入口 (见 05 章五件套),UB→L1 没有直接通路,所以 Cube 化的每个操作数都要"UB→GM→L1"。
6. 第二轮转置:lo 残差分量与 B 板(v_new)
锚点:chunk_gated_delta_rule.h:2167-2193
event_t mte3ToVector = static_cast<event_t>(pipe_->FetchEventID(HardEvent::MTE3_V));
SetFlag<HardEvent::MTE3_V>(mte3ToVector);
WaitFlag<HardEvent::MTE3_V>(mte3ToVector);
Cast(chunkKFp32, stageLocal, RoundMode::CAST_NONE, kAElements);
PipeBarrier<PIPE_V>();
Sub(chunkKFp32, chunkAttnOutFp32, chunkKFp32, kAElements);
PipeBarrier<PIPE_V>();
Cast(stageLocal, chunkKFp32, RoundMode::CAST_NONE, kAElements);
PipeBarrier<PIPE_V>();
for (uint32_t block = 0; block < kMatmulK / 16; ++block) {
LocalTensor<half> srcRows[16];
LocalTensor<half> dstRows[16];
for (uint32_t row = 0; row < 16; ++row) {
srcRows[row] = stageLocal[block * 16 * kStateRows + row * kStateRows];
dstRows[row] = transposedLocal[block * 16 + row * kMatmulK];
}
TransDataTo5HD<half>(dstRows, srcRows, transposeParams);
}
PipeBarrier<PIPE_V>();
vectorToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE3));
SetFlag<HardEvent::V_MTE3>(vectorToMte3);
WaitFlag<HardEvent::V_MTE3>(vectorToMte3);
DataCopy(aResidualStageGm, transposedLocal, kAElements);
stateOutQueue_.FreeTensor(stageLocal);
StageHalf(chunkVFp32, bStageGm, kBElements);
PipeBarrier<PIPE_ALL>();:2168-2170MTE3_V成对事件:等 MTE3 把上一块 UB 读空。 严格讲:2171写的是chunkKFp32、:2166读的是kCumdecayFp32,二者不重叠, 但同一函数里:2184马上要覆写transposedLocal(= kCumdecayFp32), 作者把栅栏放在这里是一次保守的全停,与:1566-1568的写法逐字对应。:2171Cast(chunkKFp32, stageLocal, CAST_NONE, kAElements):FP16→FP32 反向 Cast, 写回chunkKFp32(容量8192 float= 32768 B,正好)。 这一步是"hi 分量还原成 FP32",为下一步求残差做准备。:2172V 栅栏;:2173Sub(chunkKFp32, chunkAttnOutFp32, chunkKFp32, kAElements):lo_fp32 = fp32 - fp32(hi),补偿法的残差定义。:2174-2175再栅栏 +Cast(stageLocal, chunkKFp32, CAST_NONE, kAElements):残差量化成 FP16 放回队列。:2176-2186与 §5 的:2148-2162完全同构的转置循环(block4 次 × 16 行 ×repeatTimes8), 唯一区别是源是残差版stageLocal、目的仍是transposedLocal(覆写 A_hi 的 UB 副本, A_hi 已在 GM 里)。:2187-2190V_MTE3成对 +DataCopy(aResidualStageGm, transposedLocal, kAElements): A_lo 落到 slot0 第三段。:2191stateOutQueue_.FreeTensor(stageLocal):单槽队列必须显式归还, 否则:2192的StageHalf里AllocTensor拿不到槽。:2192StageHalf(chunkVFp32, bStageGm, kBElements): B 板 =v_new(10 章离开时的身份),走:1581-1589的只量化一次版本。 这里有一个必须点明的精度事实:chunkVFp32里的v_new是 Step4 的 FP32 Cube 结果 (ComputeValueAndVNewCube的 L0C 读回),StageHalf内部Cast成 FP16 后不做 hi/lo 拆分, 所以 Step7 的第二个操作数被量化了一次。 六条主循环 Cube 函数里只有MatmulBlockFp32Compensated(:454/:471)对 B 也拆 hi/lo。 docs/PR-1135 §2.2 称 StageHalf 的 B 侧 Cast 为"无损",在本文件语境下不准确(见返回说明)。:2193PipeBarrier<PIPE_ALL>():V/MTE3 全部收尾,下面进 Cube。
7. 四级分配器与"常驻 B"
锚点:chunk_gated_delta_rule.h:2194-2214
LocalMemAllocator<Hardware::L1> l1Allocator;
LocalMemAllocator<Hardware::L0A> l0aAllocator;
LocalMemAllocator<Hardware::L0B> l0bAllocator;
LocalMemAllocator<Hardware::L0C> l0cAllocator;
LocalTensor<half> a1Local = l1Allocator.Alloc<TPosition::A1, half>(kATileElements);
LocalTensor<half> b1Local = l1Allocator.Alloc<TPosition::B1, half>(kBElements);
LocalTensor<half> a2Local = l0aAllocator.Alloc<TPosition::A2, half>(kATileElements);
LocalTensor<half> b2Local = l0bAllocator.Alloc<TPosition::B2, half>(kBElements);
LocalTensor<float> c1Local = l0cAllocator.Alloc<TPosition::CO1, float>(kCTileElements);
// B is shared by both 64-row M tiles and remains resident in L0B.
DataCopy(b1Local, bStageGm, Nd2NzParams{1, kMatmulK, kMatmulN, 0, kMatmulN, kMatmulK, 1, 0});
SetFlag<HardEvent::MTE2_MTE1>(0);
WaitFlag<HardEvent::MTE2_MTE1>(0);
for (uint32_t kBlock = 0; kBlock < kMatmulK / 16; ++kBlock) {
LoadData(b2Local[kBlock * kMatmulN * 16], b1Local[kBlock * 512 / sizeof(half)],
LoadData2DParams{0, kMatmulN / 16, kMatmulK / 16, 0, 0, true, 0});
}
SetFlag<HardEvent::MTE1_M>(0);
WaitFlag<HardEvent::MTE1_M>(0);:2194空行。:2195-2198四个LocalMemAllocator:Hardware::L1 / L0A / L0B / L0C是片上存储层级的分配器类型;:2199-2203的Alloc<TPosition::..., T>(n)才是取指针。:2199TPosition::A1(L1 的 A 区)、:2200B1、:2201A2(L0A)、:2202B2(L0B)、:2203CO1(L0C,FP32)。语义:GM --MTE2/Nd2Nz--> L1 --LoadData--> L0A/L0B --Mmad--> L0C。:2199/:2201用kATileElements(单 tile,128 桶 4096 half)而不是kAElements: L0A 只需装得下一个 M-tile,这正是kTileM拆分的第二重收益;:2200/:2202用kBElements(整个 B,128 桶 8192 half);:2203用kCTileElements。:2205注释B is shared by both 64-row M tiles and remains resident in L0B.: B 只搬一次、只装载进 L0B 一次,两个 M-tile 复用—— 这是"K^T@v_new"这类"同一 B 乘多个 A tile"的标准省流量写法。:2206DataCopy(b1Local, bStageGm, Nd2NzParams{1, kMatmulK, kMatmulN, 0, kMatmulN, kMatmulK, 1, 0}): 八字段依次为{repeatTimes, srcBlockLen, srcStride(块), srcStep, dstStride(块), dstBlockLen, ... , }, 与本函数 A 侧:2219-2221的写法对照可读性最好:A 用{1, curTileM, kMatmulK, 0, kMatmulK, curTileM, 1, 0}, 即"ND 源每重复一次跳kMatmulK个 32B 块、NZ 目的每 fractal 写curTileM/16段"。 B 的 ND 是[64, dk](K 维 64、N 维 dk),NZ 目的按 fractal 铺。:2207-2208MTE2_MTE1成对:等 L1 写完再让 MTE1(L1→L0 的搬运管)读。:2209-2212for kBlock < kMatmulK/16的LoadData:LoadData(b2Local[kBlock*kMatmulN*16], b1Local[kBlock*512/sizeof(half)], LoadData2DParams{0, kMatmulN/16, kMatmulK/16, 0, 0, true, 0})。512/sizeof(half)= 256 个 half = 一个 fractal,b1Local[kBlock*256]是第 kBlock 个 fractal 起点;LoadData2DParams末位之前的true是 B 侧的转置/连续标志(A 侧:2226同位置为false), 这正是"L0B 里 B 以列主序驻留"的硬件表达。:2213-2214MTE1_M成对:L0B 就绪,允许 M 管开算。 注意:2199的 A 板此时还是空的——A 在循环里逐 tile 搬。
8. M-tile 主循环与两次 Mmad(hi、lo 累加进同一个 L0C)
锚点:chunk_gated_delta_rule.h:2215-2248
for (uint32_t mTile = 0; mTile < kMTileCount; ++mTile) {
uint32_t curTileM = (mTile + 1) * kTileM <= kStateRows ? kTileM : kStateRows - mTile * kTileM;
uint32_t aOffset = mTile * kATileElements;
DataCopy(
a1Local, aStageGm[aOffset],
Nd2NzParams{1, static_cast<uint16_t>(curTileM), kMatmulK, 0, kMatmulK, static_cast<uint16_t>(curTileM), 1, 0});
SetFlag<HardEvent::MTE2_MTE1>(0);
WaitFlag<HardEvent::MTE2_MTE1>(0);
for (uint32_t mBlock = 0; mBlock < curTileM / 16; ++mBlock) {
LoadData(a2Local[mBlock * kMatmulK * 16], a1Local[mBlock * 512 / sizeof(half)],
LoadData2DParams{0, kMatmulK / 16, static_cast<uint16_t>(curTileM / 16), 0, 0, false, 0});
}
SetFlag<HardEvent::MTE1_M>(0);
WaitFlag<HardEvent::MTE1_M>(0);
Mmad(c1Local, a2Local, b2Local, MmadParams{static_cast<uint16_t>(curTileM), kMatmulN, kMatmulK, 0, false, true});
SetFlag<HardEvent::M_MTE1>(0);
WaitFlag<HardEvent::M_MTE1>(0);
DataCopy(
a1Local, aResidualStageGm[aOffset],
Nd2NzParams{1, static_cast<uint16_t>(curTileM), kMatmulK, 0, kMatmulK, static_cast<uint16_t>(curTileM), 1, 0});
SetFlag<HardEvent::MTE2_MTE1>(0);
WaitFlag<HardEvent::MTE2_MTE1>(0);
for (uint32_t mBlock = 0; mBlock < curTileM / 16; ++mBlock) {
LoadData(a2Local[mBlock * kMatmulK * 16], a1Local[mBlock * 512 / sizeof(half)],
LoadData2DParams{0, kMatmulK / 16, static_cast<uint16_t>(curTileM / 16), 0, 0, false, 0});
}
SetFlag<HardEvent::MTE1_M>(0);
WaitFlag<HardEvent::MTE1_M>(0);
Mmad(c1Local, a2Local, b2Local, MmadParams{static_cast<uint16_t>(curTileM), kMatmulN, kMatmulK, 0, false, false});
SetFlag<HardEvent::M_MTE1>(0);
WaitFlag<HardEvent::M_MTE1>(0);:2215空行;:2216for (mTile = 0; mTile < kMTileCount; ++mTile):128 桶跑 2 轮,其余 1 轮。:2217curTileM = (mTile+1)*kTileM <= kStateRows ? kTileM : kStateRows - mTile*kTileM: 通用残 tile 处理,但在四个桶里恒等于kTileM(§2 表已证 128=2×64、80/96 整除)。 保留这行的意义在于:一旦将来加 112 这类桶,代码不用改。:2218aOffset = mTile*kATileElements:GM 上 A 板按 tile 切。:2219-2221A_hi:DataCopy(a1Local, aStageGm[aOffset], Nd2NzParams{1, curTileM, kMatmulK, 0, kMatmulK, curTileM, 1, 0})。:2222-2223MTE2_MTE1;:2224-2227LoadData把 A 的curTileM/16个 fractal 列送进 L0A,a2Local[mBlock*kMatmulK*16]对应 A^T 的第 mBlock 组 16 行。:2228-2229MTE1_M;:2230Mmad(c1Local, a2Local, b2Local, MmadParams{curTileM, kMatmulN, kMatmulK, 0, false, true}):{m, n, k, ?, 保留位, init}。末位init=true= C 清零重写,即本轮 tile 的第一笔。:2232-2233M_MTE1:M 管读完了 L0A/L0B 才允许 MTE1 覆写——L0A 是a1Local→a2Local的目的地, 紧接着:2234就要把 A_lo 写进同一个a1Local。:2234-2245与:2219-2230逐行同构,只有两处不同: 源换成aResidualStageGm[aOffset](:2235), 以及:2245的init=false—— 同一c1Local上继续累加。hi/lo 补偿与"Cube 融合"在硬件上是同一件事:都只是一串init=true,false,...的 Mmad 打在同一块 L0C。:2247-2248M_MTE1成对:本轮 Mmad 对 L0A 的读结束(注意:这里不是M_V_MTE3, 也没有任何"等 L0C 写完"的事件——dav_m200上 L0C 的读回只能靠后面的PipeBarrier<PIPE_ALL>收口, 这正是 !1174 要修的 "vec read & cube write same address" 的根,见 docs/PR-1174 §1)。
9. L0C 落地与两条 NZ→ND 归并:融合 vs 分段
锚点:chunk_gated_delta_rule.h:2249-2289
LocalTensor<float> cNz = chunkKFp32;
if constexpr (kSpecializedDk <= 96) {
// chunkK/kCumdecay/decayMask are dead after staging A. Their combined
// contiguous UB region holds the complete 64/80/96-bucket NZ result.
cNz = tmpBuff.GetWithOffset<float>(kStateElements, 0);
}
DataCopyParams cCopyParams{static_cast<uint16_t>(kMatmulN / 16), static_cast<uint16_t>(curTileM / 16), 0, 0};
DataCopyEnhancedParams cCopyEnhanced;
cCopyEnhanced.blockMode = BlockMode::BLOCK_MODE_MATRIX;
DataCopy(cNz, c1Local, cCopyParams, cCopyEnhanced);
PipeBarrier<PIPE_ALL>();
constexpr uint16_t kNdBlockLen = 16 * sizeof(float) / BLOCK_BYTES;
if constexpr (kSpecializedDk <= 96) {
// Fuse NZ->ND addressing with state accumulation. One repeat-stride
// Add handles every row of a 16-column NZ fractal.
constexpr uint32_t kNzFractalElements = kStateRows * 16;
constexpr uint8_t kNdDstRepStride = kMatmulN * sizeof(float) / BLOCK_BYTES;
constexpr uint8_t kNzSrcRepStride = 16 * sizeof(float) / BLOCK_BYTES;
for (uint32_t nBlock = 0; nBlock < kMatmulN / 16; ++nBlock) {
Add(stateInFp32[nBlock * 16], stateInFp32[nBlock * 16], cNz[nBlock * kNzFractalElements], 16,
static_cast<uint8_t>(kStateRows),
BinaryRepeatParams{1, 1, 1, kNdDstRepStride, kNdDstRepStride, kNzSrcRepStride});
}
PipeBarrier<PIPE_V>();
} else {
uint16_t nzSrcStride = static_cast<uint16_t>((curTileM / 16 * 16 * 16 - 16) * sizeof(float) / BLOCK_BYTES);
DataCopyParams nzToNdParams{static_cast<uint16_t>(kMatmulN / 16), kNdBlockLen, nzSrcStride, 0};
uint32_t nzFractalStride = (curTileM / 16) * 16 * 16;
for (uint32_t row = 0; row < curTileM; ++row) {
uint32_t stateRow = mTile * kTileM + row;
for (uint32_t nBlock = 0; nBlock < kMatmulN / 16; ++nBlock) {
uint32_t ndOffset = stateRow * kMatmulN + nBlock * 16;
uint32_t nzOffset = row * 16 + nBlock * nzFractalStride;
Add(stateInFp32[ndOffset], stateInFp32[ndOffset], cNz[nzOffset], 16);
}
}
}
PipeBarrier<PIPE_V>();
}
}:2249LocalTensor<float> cNz = chunkKFp32;:默认借 chunkKFp32(128 桶走这条)。 容量:curTileM*kMatmulN = 64*128 = 8192 float = 32768 B=chunkKFp32全容量,贴边。 与 10 章:2041的cNz = kCumdecayFp32不同:Output 侧 kCumdecay 是 attn 落地板、chunkK 空闲, Step7 侧刚好相反(chunkK 刚被:2173的残差用完,kCumdecay 是转置板但也要留到:2190之后)。 同一批 UB 在不同函数里身份互换,是 11 章风险清单的主体。:2250-2254if constexpr (kSpecializedDk <= 96)→:2253cNz = tmpBuff.GetWithOffset<float>(kStateElements, 0): 直接从 tmpBuff 偏移 0 取一整块,等价于把chunkKFp32 | kCumdecayFp32 | decayMaskFp32三块连续区拼成一个大板(InitLocalBuffers:161-168保证它们首尾相接)。:2251-2252的注释就是这个意思:三者在 staging A 之后全死。 容量核对(kStateElements个 FP32 vs 三块之和):桶 需要 kStateElements*4三块连续区字节数 判定 64 4096·4 = 16384 16384+16384+16384 = 49152 宽松 80 6400·4 = 25600 20480+20480+16384 = 57344 宽松 96 9216·4 = 36864 24576+24576+16384 = 65536 宽松 128 16384·4 = 65536 32768+32768+16384 = 81920 不合并( :2250排除),改用单块 chunkK128 桶被排除的真正原因不是放不下,而是放不下"一个 tile"以外还需要跨 tile 连续性: 融合分支下面会用到
kStateRows*16这个常量步长。:2255-2258L0C→UB 矩阵模式搬运:DataCopyParams{repeatTimes = kMatmulN/16, blockLen = curTileM/16, 0, 0}+blockMode = BLOCK_MODE_MATRIX。 矩阵模式下两字段被重新解释为"N/16 个 fractal 列 × curTileM/16 个 cube 块", 总量(kMatmulN/16)*(curTileM/16)*256个 FP32 =kCTileElements✓(128 桶 8×4×256 = 8192 = 32768 B)。:2259PipeBarrier<PIPE_ALL>():本函数里唯一一次跨全管收口, 作用是"M 管的 L0C 写与 MTE2 的 L0C→UB 读都完成,V 管才能把它当普通 UB 读"。 没有比这更细的原语可用(无 Fixpipe、无M_V事件对 L0C 的直接支持), 代价是整条流水在这一行停顿。:2261kNdBlockLen = 16*sizeof(float)/BLOCK_BYTES = 2(块):一条 ND 段 = 16 个 FP32 = 64 B。:2262-2273融合分支(64/80/96)::2265constexpr uint32_t kNzFractalElements = kStateRows * 16;:2266kNdDstRepStride = kMatmulN*4/32(128→16 块,96→12 块,80→10 块):2267kNzSrcRepStride = 16*4/32 = 2块:2268for nBlock < kMatmulN/16→ 只有dk/16条 V 指令(64/80/96 桶分别 4/5/6 条)。:2269-2271:Add(dst=stateInFp32[nBlock*16], src0=同一个, src1=cNz[nBlock*kNzFractalElements], 16, kStateRows, BinaryRepeatParams{1,1,1, kNdDstRepStride, kNdDstRepStride, kNzSrcRepStride})读法:每个 repeat 处理 16 个元素;dst/src0 每次前进一整行 ND(kMatmulN个 float =kNdDstRepStride块), src1 每次前进 16 个 float(kNzSrcRepStride= 2 块);共kStateRows个 repeat。 为什么 src1 的常量步长能走完整块 NZ? 元素(m,n)在 L0C 的 NZ 布局里位于fractal(n/16, m/16)内偏移(m%16)*16 + (n%16), 即地址= (n/16)*numMBlocks*256 + (m/16)*256 + (m%16)*16。 本分支numMBlocks = kStateRows/16,故numMBlocks*256 = (kStateRows/16)*256 = kStateRows*16—— 正是kNzFractalElements; 而(m/16)*256 + (m%16)*16 = m*16—— 正是"每 repeat 前进 16 个 float"。 两个等式同时成立的条件就是"一个 tile 覆盖全部kStateRows行", 于是 §0 第 2 条的必然性闭环:128 桶切成 64 行 tile 后numMBlocks变成curTileM/16, 与外层nBlock的基址kStateRows*16不一致,融合写法失效,只能走分段。:2273融合分支的PipeBarrier<PIPE_V>()。:2274-2286分段分支(128 桶)::2275nzSrcStride = (curTileM/16*16*16 - 16)*4/32 = (4*256-16)/8 = 126块, 与 10 章:2049的kNzSrcStride同值(那里kMatmulM恒为 64)。 含义:一次读 2 块(16 个 float = 一个 fractal 行),再跳 126 块 = 净进 128 块 = 1024 float = 4 个 fractal = 下一个nBlock✓。:2276DataCopyParams nzToNdParams{kMatmulN/16, kNdBlockLen, nzSrcStride, 0},:2277nzFractalStride = (curTileM/16)*16*16 = 1024,:2278-2285双层循环curTileM × kMatmulN/16 = 64×8 = 512条Add(..., 16)(每个 tile 512 条,两 tile 1024 条)。 与融合分支对比:指令条数从 8 条涨到 1024 条,但省掉了 L0C→UB→ND 的二次搬运。 注意:2279stateRow = mTile*kTileM + row:分段分支必须显式补偿 M-tile 偏移, 融合分支因为覆盖全 M 而无需此项。 这里Add的 src1 已经不再是纯 NZ(:2283读的是cNz[nzOffset],nzOffset = row*16 + nBlock*1024——直接用地址算式,不再依赖 fractal 常量步长),:2275-2276的nzToNdParams实际上定义了同一套地址, 但本分支用 Add 而非 DataCopy 落位,因为要"加"而不是"搬"。:2287PipeBarrier<PIPE_V>();:2288循环右括号;:2289函数右括号。