Softmax VF 尾块优化（P 布局重构）方案说明

分支：flash_attn_softmax_vf_tail（基线 master 06e277852）
提交：5bf987379（实现）+ df765da7a（clang-format）
涉及文件：op_kernel/utils/attn_softmax_dn_bn1s1s2_vf.h、op_kernel/arch35/flash_attn_block_vec_dn_unmerged.h
适用：不合轴 DN 模板（templateId=2，*_dn_unmerged.h）的 softmax 向量化段

---

1. 问题

旧实现里 softmax 的列循环恒按编译期 tileCols(= s2BaseSize) 跑满：

位置	旧循环	代价
补值	for (col = validColCount; col < tileCols; ++col) 逐列 256B store	O(tileCols − actS2) 次 store
掩码 + max	for (colGroup < tileCols/4)，4 个相邻列	O(tileCols)
exp/打包	for (colGroup < tileCols/4)，4 个相隔 tileCols/4 的列	O(tileCols)，VEC 主要开销

所以尾块（actS2 < tileCols）无论实际几列都要按满宽算一遍。实测（D=128，tileCols=256，32768 个全尾块）：
旧实现的 kernel 时间几乎与 actS2 无关 —— actS2=32 时 1293us，actS2=256 时 1839us（仅 +42%）。

根因：旧 P 布局把"相隔 tileCols/4 的两列"打包进同一条 128-half store，于是"处理多少列"只能是
tileCols/4 的整数倍，补值（逐列 store）反而是最便宜的手段。

目标：把列粒度降到 2 列，让 exp/补值/max 的工作量随 actS2 走。

---

2. 约束（VF 硬限制）

约束	说明
RegTensor ≤ 32、MaskReg ≤ 8	超限会 spill / 劣化
Hardware Loop	循环变量 uint16_t、从 0、步长 1、循环体内禁 if/?:
StoreAlign<DATA_BLOCK_COPY>	一条指令搬 8 个 DataBlock(32B)，块间只有一个统一步长 dataBlockStride（单位 DataBlock）
L1 侧布局不可动	cube 的 bmm2 按 NZ [m_fractal][k_fractal] 读 P，本方案不改 cube

---

3. 方案总览

1. P(UB) 改成 32 列一组：组内 8 个槽，每槽 = 16×16 fractal(512B) + 32B pad = 544B = 17 DataBlock；组 = 8×544 = 4352B。
2. VF 组内 4 路取列 {c, c+8, c+16, c+24}（c 为组内列号 0..7），按 (c, c+16) / (c+8, c+24) 配对，
   两条 store 各写 8 块，dataBlockStride=17、repeatStride=1（等价于每次迭代列号 +1，32B）。
3. 每 32 列仍是 16 条 store（8 迭代 × 2 条），与旧实现同量，不增加指令数。
4. 尾组按 2 列收界：组数 ceil(actS2/32)，末组迭代数 k = min(actS2 % 32, 8)；
   补值范围随之从 tileCols 收到 "末组最远被碰列"（≤27 列）。
5. vec1 的 P→L1 搬运改为按 fractal 取数（每 m_fractal 一条 DataCopy），L1 布局与旧版完全一致。
6. mask_mode=0 也走本模板 VF（common 的 ProcessVec1VfDn 是旧 P 布局，与新搬运不匹配）。

---

4. P 布局精确规格

组 g（32 列） = 2 个 k_fractal(16 列) × 4 个 m_fractal(16 行) = 8 个槽
槽序:  k0.m0  k0.m1  k0.m2  k0.m3 | k1.m0  k1.m1  k1.m2  k1.m3
偏移:   +0    +544   +1088  +1632 | +2176  +2720  +3264  +3808

槽内(512B) = 16 个 s2 列 × 32B：第 i 列占 [i*32, i*32+32)
一列(64 行) = 4 个 32B 块，分别落在 m0..m3 四个槽里

地址公式（byte）：

off(g, c', m) = g*4352 + ((c'/16)*4 + m)*544 + (c'%16)*32
容量 = ceil(s2BaseSize/32) * 4352B        // = 4.25 × s2BaseSize

s2BaseSize	新容量	旧容量	差
128	17,408B	16,640B	+4.6%
224	30,464B	28,928B	+5.3%
256	34,816B	33,024B	+5.4%

4.1 为什么 pad 必须"每 fractal 一个"

一条 store 要写 (c, c+16) 这一对（每列 4 个 16 行 chunk），8 个块地址必须均匀：

pad 策略	8 块地址（每块 32B，fractal=512B）	是否可行
纯 NZ 无 pad	0, 512, 1024, 1536, 2048, 2560, 3072, 3584	❌ 步长 16 均匀，但 8 块全是 512 的倍数 → bank 相位相同，(c, c+16) 撞 bank
每组只在 16/16 交界留 1 个 pad（组 4128B）	0, 512, 1024, 1536, 2080, 2592, 3104, 3616	❌ 前 4 块 512、第 5 块起 544 → 步长不均匀，只能拆两条半空 store
每 fractal 后 +32B（槽 544B）	0, 544, 1088, 1632, 2176, 2720, 3264, 3808	✅ 步长 17 均匀，且 17 为奇数 → 相邻 fractal 起始 bank 相位错开 32B

---

5. VF 改造细节

5.1 计算（不变）

Muls(scale) → 掩码 Compare/Select → FusedExpSub → Cast<T2>。4 路仍是 4 个独立列，
寄存器预算与旧实现一致（~24 RegTensor）。

5.2 打包与写（新）

packedA = DeInterleave(cast(expValue0), cast(expValue2));  // 列 (c,   c+16) → 128 half
packedB = DeInterleave(cast(expValue1), cast(expValue3));  // 列 (c+8, c+24)
StoreAlign<T2, DATA_BLOCK_COPY>(ptrA, packedA, /*dataBlockStride=*/17, /*repeatStride=*/1, allLanesB16);
StoreAlign<T2, DATA_BLOCK_COPY>(ptrB, packedB, /*dataBlockStride=*/17, /*repeatStride=*/1, allLanesB16);

• ptrA = 组基址 + c*32（写 c 与 c+16），ptrB = ptrA + 256B（写 c+8 与 c+24）
• 一条 store 的 8 个块 = "列 X 的 4 个行块 + 列 Y 的 4 个行块"
• 实测确认：DeInterleave(A,B) = [A 的 64 lane | B 的 64 lane] 拼接

5.3 循环结构（★ 实测关键）

// 满组: 外层组(uint16) + 内层组内列(uint16), 地址只由"基址 + 迭代号*常量"构成
for (uint16_t groupIdx = 0; groupIdx < fullGroupCount; ++groupIdx) {
    __ubuf__ float* scoreGroup = scoreBuf + groupIdx * 32 * rowStride;
    __ubuf__ T2*    pOutGroup  = pOutBuf  + groupIdx * 2176;      // 4352B/2
    for (uint16_t colInGroup = 0; colInGroup < 8U; ++colInGroup) { ... }
}
// 末组: 迭代数 k = min(actS2 % 32, 8)

踩坑：第一版把"组 × 组内列"合成一个扁平循环，用 组号 = i>>3, 列号 = i&7 反推地址。
语义/精度全对，但 Compiler 不再生成 Hardware Loop：VEC 时间翻倍、AIV VEC 占比打满
（212 条语料平均 0.6503、总 0.6922）。VF 里不要用移位/取模从循环变量反推索引。

5.4 尾块裁剪与补值上界

exp 的 4 路取列会碰到末组里最远 (k-1)+24 列，而有效列只到 actS2。补值必须覆盖被 exp 到的列，
不能只覆盖 max 需要的 align4(actS2)：

• 分母：LastDivNew(..., sumUb, ...) 用 running sum 归一化 attn_out；
• LSE：lse = max + log(sum)。

lse 对 max 是不变的（max 抬 Δ，sum 乘 e^-Δ，lse 抵消），所以 LSE 出错只能来自 sum 里多了项。
实测把补值收到 align4/align16(actS2) 后（P 已全干净），LSE 立刻崩（0.0175 / 0.1254 / 1.9614），
唯一通过的是 actS2 % 32 == 31（取列集合恰好等于有效列）—— 反证成立。

结论：

fillEndCol = 32*fullGroups + min(actS2%32, 8) + 24,  再向上对齐到 4 的倍数
补值列数 ≤ 27（旧实现最多 tileCols-1 列）
掩码/max 循环收界到同一范围（4 列粒度）

5.5 寄存器预算

packLanesB16 与 allLanesB16 对 b16 是同一份"全 128 lane"掩码，合并省出 1 个；
最终 MaskReg 仍为 8/8（allLanesF32 / allLanesB16 / validRowMask / visibleMask0..3 + 常量表），
RegTensor ~24/32。

---

6. vec1 的 P → L1 搬运

L1 目的布局不变（NZ，m_fractal 步长 16*actS2Align 元素，k_fractal 连续）。新布局下每个槽
恰好是一个 16×16 fractal，于是每个 m_fractal 一条 DataCopy：

// blockCount = actS2Align/16, blockLen = 16 DataBlock(512B), srcGap = 52 DataBlock, dstGap = 0
for (uint16_t mFrac = 0; mFrac < 4U; ++mFrac) {
    DataCopy(mm2AL1Tensor[dstBase + mFrac*16*actS2Align],
             stage1CastTensor[mFrac * (544/sizeof(INPUT_T))],
             {fractalCount, 16, 52, 0});
}

（源步长 4×544B = 2176B = 512B + 1664B ✓；目的连续。）

---

7. 正确性验证

用例集	结果
P 布局精度矩阵（tile=256，32 条）	fp16 全部 PASS；2 条 bf16 与基线逐位相同
全尾块专项（21 条，S2=actS2<tileCols）	20/21 PASS（唯一 FAIL 为 bf16，与基线逐位相同 1.2174% vs 1.2179%）
s2=224 专项（10 条）	10/10 PASS
既有逐档（5 条，含 mask0 / M 尾 / FD）	5/5 PASS

• 与旧实现逐位一致：所有可比 fp16 用例的 MaxAbsErr 一致到 6 位小数 → P 列映射、vec1 搬运、容量公式均正确。
• bf16 的 2~3% fail 属既有量化口径（同形状 fp16 为 0.04~0.12%），换回旧实现数值完全相同。
• TND varlen 在当前 harness 上会 pybind11 GIL 崩溃（仓库自带用例、且不走本模板的对照用例同样崩），
  与本次改动无关，故该格未覆盖。

---

8. 性能

8.1 尾块隔离测量（S2 = actS2 < tileCols，32768 个尾块，5 runs 取 min）

actS2 / tileCols	基线 (us)	本方案 (us)	加速比
12.5%（32/256）	1292.6	1143.9	1.130
25%（64/256）	1278.0	1165.2	1.097
50%（128/256）	1289.6	1234.5	1.045
87.5%（224/256）	1828.3	1814.8	1.007
100%（256/256，满块对照）	1838.8	1876.5	0.980

• 旧实现耗时几乎与 actS2 无关，新实现随 actS2 线性 → 列裁剪确实生效；
• VEC 忙时在 32/256 档下降 ~40%（ratio 0.091→0.056，VEC 时间 ×1.67），但 softmax 只占该 shape 的 5~15%，
  所以端到端收益小于 VEC 削减幅度；
• 满块对照 −2%（全语料口径 −0.4~−0.8%）：新布局在满块上确有少量代价。

8.2 全语料（223 条，背靠背 5 runs 取 min）

分类	n	平均	中位	总耗时
全部	221	0.9986	0.9987	1.0013
满块	150	0.9922	0.9937	0.9932
尾块	71	1.0121	1.0087	1.0173

语料里尾块只占 1/nblk（最大 10%：S2=2048/tile=224 的 10 块 1 尾），所以整体接近中性；
收益集中在"尾块占比高"的形状（上表 8.1）。

8.3 测量噪声

同版本两次独立测量：多数 case ~0.1%，个别大 case 5~9%。表中 ±1% 级差异应视为噪声。

---