OpenDLSS-NR 窗口注意力拆解：余弦归一化、exp位运算与固定序Softmax

OpenDLSS-NR 是一个把 NVIDIA DLSS 5 神经渲染（Neural Rendering）网络用 Vulkan 重新实现的开源项目，与原始实现逐位（bit-exact）一致。窗口注意力是这条 71 块 Swin/ViT 网络里最关键的算子之一：每一块都在 8×8 的 token 窗口内做一遍注意力。本文拆解它的三大核心机制——余弦归一化、exp 位运算技巧和固定顺序 Softmax，帮你看懂这套"逐位复刻"到底复刻了什么。

8×8 窗口：注意力的舞台 🧩

• 窗口大小固定为 8×8：64 个 token（一个 token = 一个像素），每个注意力头对它们做完整的 QKᵀ·softmax·PV。
• 窗口原点会平移。每个 block 从四个相位里选一个，x、y 各为 0 或 4：

相位每块推进一次，解码器阶段接着编码器留下的计数继续——x 边界与 y 边界由不同 block 各自跨越，而不是常见的两相位交替。

• 窗口越界不裁剪：越界 token 是零向量，它们的 key/value 贡献为零，但分数（= 学习先验 + 0）的指数仍然进入 Softmax 分母。把它们 mask 掉会得到一个不同的函数。

• 两种 token 顺序：查询用自然顺序（行优先）；键/值用物理顺序（4×4 平铺，p = tile*16 + within）。物理顺序让 4×4 块内相邻像素在内存中连续，是共享内存加载与矩阵乘分块的基础；先验、Softmax 求和树、PV 矩阵乘的 K 维度共用同一套索引。

余弦归一化：先"摆正"再相乘 📐

QKV 投影产出 q、k、v 之后，q 和 k 不做原始点积，而是先归一化（实现见 shaders/fused_block32.comp 的normalizeInPlace，独立 kernel 版在 shaders/window_normalize.comp）：

1. 逐 32 维求平方和：16 对分量打包做 FMA ——r[c] = fma(v[c], v[c], f16(v[c+16]²))（高 16 维的平方先单独舍入到半精度，再并入低 16 维）；
2. 树状归约：stride 8/4/2/1，共 15 次加法，每一步都是半精度舍入；
3. inversesqrt在 f32 中计算，再舍回半精度；
4. q 额外乘一个学习的每头 f32 温度 scale，k 只归一化，v 只做量化；
5. 结果发布为E4M3（8 位 FP8）。

两个容易漏掉的细节：

• 零行会产生 NaN：窗口边界的零向量归一化是0 × inf，NaN 会毒化后续矩阵乘。因此本项目的发布规则是NaN 发布为 +0，而-0的符号必须保留——差一个符号位，字节比对就不一致。
• 为什么用余弦而不是点积：余弦相似度把分数限制在 [−1, 1]，再乘学习温度；logit 不再随向量范数增长，FP16 累加器不会失控、3 位尾数的权重网格才可用，Softmax 也因此不需要减最大值（max-subtraction）——这正是下一节 bit trick 能成立的前提。

exp 位运算技巧：不用超越函数算指数 ⚡

Softmax 需要 exp(s)。窗口注意力没有调用任何超越函数，而是用 FP16 的位模式做了一次仿射 + 一次位操作（shaders/window_attend.comp 的expWeight，fused 版配对写法见 shaders/attention_common.glsl）：

x = clamp(0.044921875 * s + 1.30078125, 1.03125, 1.5693359375); bits = f16Bits(x); result = f16FromBits((bits << 5u) + 0x8000u); // == (bits << 5) ^ 0x8000

为什么这一步"数学上成立"：

• 对一个落在 [1, 2) 的半精度数，左移 5 位再+ 0x8000等价于指数域 +5，即数值乘以 2⁵（尾数不动、不引入新舍入）；
• 于是结果 =2^(1.4375·s − 5.375)：指数部分是 s 的线性函数，尾数则携带一个仿射近似——这正是exp(s)乘一个常数倍，而常数倍会被 Softmax 的分母除掉，归一化结果不受影响；
• 有效范围约s ∈ [−6, +6]，足够覆盖余弦 logit + 先验的全部取值。

clamp是真正的安全绳：它把每个指数压进 [2⁵×1.03125, 2⁵×1.5693359375] ≈[33, 50.2]，64 个键求和最多约 3200，稳稳落在半精度（最大 65504）之内——这就是"没有 max-subtraction 也能用 f16 分母"的原因。

固定序 Softmax：顺序本身就是语义 🔁

求和不是"随便加"，而是严格按规范（docs/numerics.md 的 "The softmax" 一节，GLSL 参考实现见 shaders/window_attend.comp）：

e[k] = expWeight(S[k]) # 64 个键，物理顺序 b[g][j] = e[g + 16j] + e[g + 16j + 8] # 先配对 (k, k+8) t[g] = ((b0 + b1) + b2) + b3 total = (((t0 + t2) + t4) + t6) + (((t1 + t3) + t5) + t7) # 偶/奇半分别求和 w[k] = E4M3( e[k] * f16(1 / total) )

• 每一个+都是半精度加法，63 次加法的顺序固定。f16 下任何重新结合（比如a+(b+c)与(a+b)+c）都可能改变位模式——顺序不是性能选择，而是语义的一部分。
• 倒数单独成步：1/total先在全精度算、舍入到 f16，再乘进每个 e[k]，最后舍入为 E4M3 发布。
• E4M3 权重：1 符号 + 4 指数 + 3 尾数，最大 448、无无穷。3 位尾数的相对步长约 1/8，单看很粗，但 64 个权重归一化后误差互相抵消，输出与原生一致。
• 在 MMA 累加器里零开销完成：16×16 累加器的 lane 布局使(k, k+8)配对与偶/奇分割都落在单个 lane 内，整行只需两次 shuffle——shaders/fused_block32.comp 里整段 Softmax 都在寄存器中执行。

一个 workgroup 装下一个窗口：融合 kernel 全景 🔩

fused_block32.comp 把"一个 8×8 窗口的完整注意力"装进一个 workgroup：128 线程 = 4 个 subgroup，每个 subgroup 负责 16 个 query；FFN → QKV → 注意力 → 投影整条链尽量留在寄存器里：

1. QKV 投影产出 q、k、v 三个累加器；余弦归一化直接在累加器寄存器内完成，K/V 行以物理顺序发布到共享内存；
2. S = QKᵀ + prior：16 个 E4M3 协同矩阵乘法（16×16×32），学习的 64×64 先验矩阵直接作为 MMA 的累加器初值——偏置加法与 QKᵀ 在同一条硬件指令里完成；
3. Softmax 全程在寄存器：exp 位技巧、树状求和、倒数、E4M3 权重；
4. O = PV：16 个 MMA 直接消费寄存器里的 E4M3 权重，输出再 E4M3 发布。

顺带一提：ViT 全局注意力差在哪 🔭

网络最底部的 ViT 块对全局 token 做注意力（shaders/global_attend.comp），指数位技巧相同、常数不同，但有四处分叉很容易猜错：

ViT 的指数常数是0.08953947·s + 1.70936143、clamp 到 [1.439453125, 1.9775390625]、左移 4 位后+ 0x4000，等价2^(1.4326·s − 3.6502)，有效范围约[−3, +3]。

Bit-exact：这些"怪癖"为什么值得复刻 🔒

• dlss5vk parity 工具在 512×512 下逐位核对75 个边界、共 57,704,448 个 E4M3 字节，判定标准严苛到+0与-0都算不一致；
• GLSL 参考路径、PTX 快速路径、浏览器 WebGPU 移植（ports/browser-webgpu/src/window/）三套实现互相印证；
• CPU 参考实现（src/reference.cpp）配合verify工具可把分歧定位到单个 kernel。

复刻的从来不是"一个数值函数"，而是一整条舍入序列：归约顺序、clamp 常数、位运算、E4M3 发布点——缺一不可。

相关代码与文档入口 📚

小结

窗口注意力 =8×8 窗口+余弦温度 logit+学习先验+位技巧 exp+固定顺序 Softmax+E4M3 发布。OpenDLSS-NR 把这些顺序、常数与舍入点完整复刻到 Vulkan 上，让每一帧都与原始 DLSS-NR 字节级一致——理解了这三块机制，你就读懂了这条网络"精确"的底层逻辑。