Skip to content

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

cpp
  __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

cpp
  __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 空行;758 Phase 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

cpp
  __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 要的字节偏移。
  • 782 CreateVecIndex(signedOffsets, 0, chunkLen) 填 0,1,...,chunkLen-1;783 PipeBarrier<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。
  • 794 offset <<= 1 → 轮数 = floor(log2(chunkLen)):chunkLen=64 → 6 轮,32 → 5 轮,1 → 0 轮(尾块/单 token 天然退化,不进循环)。
  • 795 Muls(dst, src, 1.0f, chunkLen) 是向量 memcpy(×1.0 是本文件的"搬家 + 刷新"惯用法);797 Add(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

cpp
  // 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-832 beta 快路径。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)。
  • 828 PipeBarrier<PIPE_V> 保证 Cast 落地后 Gather 才读;829-830 造偏移表并置旗标;831 Gather(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 空行。
  • 839 float gSum = 0.0f; 只有 fallback 用得到(860),快路径下它全程为 0 —— 这是 !1108 ComputeGCumsum 的残余声明,成本是一个标量寄存器,保留它让两条路径共用同一函数体。
  • 840-854 g 快路径: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-853 Gather 出原始 g(还不是前缀和);854 交给 PrefixSumChunkG 就地扫描。
  • 855-857 hasGamma_ == 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

cpp

  // 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-872 TQueSync<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 新增的四行(含注释与空行)。881 Duplicate(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

cpp
  // 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-906 Muls(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-924 gBlocks = chunkVFp32(第 4 次借用)+ Brcb 把 +G_i 展成"每行一个 32B 块"。
  • 925-928 高维 Add(rowAddParams 与 885 的 mulParams 同形)→ decayMask[i][j] = G_i − G_j,整块一次。
  • 929 空行;930-932 Exp 之前清非因果区:j > i 处 G_i − G_j > 0,Exp 出来是 inf 甚至溢出;不清零就等着后面 1145/1148 的整块 Mul 把 0×inf 变成 NaN、并被 ComputeRecursiveAttnBlocked64 的整块读取消费(分块求逆按完整 16×16/32×32 块读,上三角的垃圾值会参与乘积)。这 63 条 Duplicate 是本段残留的发射开销,作者选择用它换数值安全;!1108 的逐行写法天然不会生成上三角的爆炸值(写出来也没人读),所以这条陷阱是新并行化引入的,注释与代码一起把它钉住了。
  • 933-935 rowExpParams{1, 1, rowStride, rowStride} + 一条高维 Exp 算完整个 64×64(dst/src 行间都前进 rowStride 块)。对比 !1108:523-531:64 次标量读 + 64×(Duplicate+Sub+Exp) = 192 条指令 → 现在 6 条高维指令 + 63 条清零。
  • 936-939 TQueSync<PIPE_V,PIPE_S> decaySync:出方向,981/993/1185/2140 的标量 GetValue 全靠它定序。940 收尾。

7. 941-961 · CanCacheAttnQuery 与 Q 缓存(chunkVFp32 再借一次) ​

锚点:chunk_gated_delta_rule.h:941-961

cpp
  __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

cpp
  // 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;971 ComputeRecursiveAttn 在 629-632 处再分叉到 ComputeRecursiveAttnBlocked64(第 05 章);972-974 标量把对角写成 1.0((I−A)⁻¹ 的对角本就是 1,这里显式覆盖是 !1108 的语义延续);975 return,976/977 两层收尾。
  • 978-1013 fallback 与 !1108:578-613 逐字节相同(diff 为空):978-983 A 的直接项(980 每格一次 DotFp32,981 标量读 decay 并写 -dot*decay)、984 递归修正、985-987 对角 1、988-996 Q 已缓存时逐行逐列点积、997-1013 Q 未缓存时每行现从 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

借用点视图装什么死亡时刻
765ReinterpretCast<inType>K 的 fp16 DMA 暂存LoadPaddedRows 内 Cast 到 chunkKFp32 之后(338)
819ReinterpretCast<inType>beta 板 fp16 落地点831 Gather 之后(但 827 与它同址 → !1158)
819/827/845fp32 视图beta 板加宽结果 / g 板落地点852 Gather 之后
874fp32 视图betaBlocks(Brcb 广播块)894 高维 Mul 之后
921fp32 视图gBlocks(Brcb 广播块)927 高维 Add 之后

五次借用互不相交的唯一硬证据是阶段顺序(§1 的 725-731):借的窗口全部夹在"上一个读者的最后一次使用"和"下一个主人的第一次写"之间。写 Cube 代码时这张表必须逐次重推——!1158 的教训正是"照抄上一处的成功答案"。