Skip to content

docs(design): split-KV decode attention design record (TLLM-ATTN-SPLITKV) - #11

Merged
holtwood merged 2 commits into
masterfrom
tllm-attn-splitkv-design
Sep 15, 2026
Merged

holtwood merged 2 commits into
masterfrom
tllm-attn-splitkv-design

Conversation

@holtwood

Copy link
Copy Markdown
Member

设计记录(PR-A),不含任何代码。 按 L3_L4_DESIGN_REVIEW_PACKAGES.md §3 模板填写,为后续 hot-path kernel 改动取得批准载体。stacked on #10。

为什么(诊断,非猜测)

PR-5 的收尾建议是"修第 4 步的 V 访问"。profiler 数据推翻了它。 ncu(2026-09-14,visible=2048 / block_size=16):

指标 值 读数
grid / block (14,1,1) / 128 一个 block 一个 head,仅此而已
sm__warps_active(achieved occupancy) 8.33% 约 70 个 SM 里 56+ 个完全空闲
registers / smem 39 / 928 B 都不是限制因素
DRAM / L2 / L1 throughput 0.44% / 0.96% / 4.0% 没有任何资源饱和
long-scoreboard / barrier stall 45.8% / 38.6% 延迟无法隐藏 + 每个 tile 的 __syncthreads()

结论:并行度不足导致的延迟受限。第 4 步的低效是"只有 4 个 warp"的症状,不是原因——只重写它会治标不治本(grid 仍是 14 个 block)。

设计

  • partial kernel:grid = (Hq, num_splits),每块按 device 端 visible_tokens 现场算自己的逻辑区间(chunk = ceil(visible/splits)),产出局部 (m_i, l_i, acc_i);combine kernel:grid = (Hq,1) 按 flash-decoding 公式合并。
  • 保住 CUDA Graph 可捕获:num_splits 是 host 参数(grid 捕获时固定),可见长度仍走 device int,partial 缓冲预分配、捕获期零分配零 D2H。
  • 不复制归约:decode_online_softmax → decode_online_softmax_range<Rows, kFinalize>,ContiguousRows/PagedRows、冻结的地址公式与零行语义全部不动。
  • API 只增不改:新增两个 _splitkv 入口;既有 attention_decode/attention_decode_paged 保持 num_splits = 1 的等价语义,因此 ~18 个既有调用点不用改,单遍路径同时充当回归锚点与回滚点。C ABI 不变,与 paged-serving 无强制顺序。

需要 reviewer 明确决议的点

  • Q3(唯一有真实爆炸半径的数值决定):num_splits > 1 会故意打破 §10.2 的「contiguous 路径逐位不变」指纹(求和顺序变了,仍在 oracle 2e-3 容差内);direct vs legacy 的逐位相等仍然成立(同一份循环)。
    缓解是 §4.4 的逐位锚点:num_splits = 1 时 exp(0)==1.0f、1/(l+1e-9) 不变 ⇒ 与今天逐位相同,从而把"重构风险"与"数值改动"干净分开。建议接受 Q3。
  • Q1 combine 用独立 kernel(建议)还是单 kernel 信号量归约;Q2 num_splits 走 host 参数;Q4 partial 放 LayerWorkspace 新缓冲(不复用 direct 路径冗余的 k_scratch);Q5 小 S 是否允许调用方回落到 num_splits = 1(不做内核内分支)。

验证

限制

  • 本包由实现者本人编写,不满足"实现 Agent 不应是唯一 reviewer";§12 已记录该独立性缺陷。改动落在两条生产 attention 路径共用的热路径 kernel 上,与 PR-1 属同一类爆炸半径 ⇒ 合并前需第二方复核。
  • 不声称任何端到端收益;kernel 级数字不得写成 TTFT/TPOT(G7)。
  • num_splits 最优值、小 S 是否回归,均列为 §2 的 Unknown,由 PR-D 实测关闭。

后续

批准后:PR-B(kernel + 测试)→ PR-C(LayerWorkspace + 开关 + 路由)→ PR-D(benchmark + 结果报告)。默认值只在 PR-D 通过后才讨论打开。

holtwood and others added 2 commits September 14, 2026 22:20
…SPLITKV)

Profiling (ncu, 2026-09-14, visible=2048/block_size=16) shows the decode
attention kernel launches as grid=(14,1,1) x 128 threads and runs at 8.33%
achieved occupancy with nothing saturated: DRAM 0.44%, L2 0.96%, L1 4.0%,
39 registers/thread, 928 B dynamic smem. Issue slots are lost to
long-scoreboard 45.8% and barrier 38.6%. The bottleneck is parallelism, not the
step-4 V access pattern that PR-5's closing note suggested.

Design: split the visible KV window across blocks (grid = (Hq, num_splits)),
each block producing a local (m, l, acc) triple, merged by a second, small
combine kernel. Split ranges are derived in-kernel from the device
visible_tokens and num_splits is a host parameter, so the grid stays fixed per
launch and CUDA Graph capture is preserved. The reduction loop is not copied:
decode_online_softmax becomes decode_online_softmax_range<Rows, kFinalize>, so
ContiguousRows and PagedRows -- and therefore the frozen address formula and
zero-row semantics -- keep working unchanged.

The strongest gate is that num_splits = 1 is bitwise identical to today
(exp(0) == 1.0f and the same 1/(l + 1e-9)), which isolates the refactor from
the numerical change. num_splits > 1 deliberately breaks the 10.2
"contiguous path bitwise-unchanged" fingerprint; that is recorded as a decision
that needs explicit approval (Q3) rather than being slipped in.

Design record only, no code: per this repo's rules a hot-path change of this
blast radius needs an approved design and a second reviewer before PR-B.

Co-authored-by: CommandCodeBot <noreply@commandcode.ai>
Owner authorised proceeding through PR-A..PR-D, so the decision field moves off
"changes_requested". This records what actually happened rather than leaving the
record contradicting the work that follows.

Q3 (accepting that num_splits > 1 breaks the 10.2 "contiguous bitwise-unchanged"
fingerprint) is marked approved: the split-KV option the owner selected already
stated that the output would no longer be bitwise identical but would stay
inside the oracle tolerance. The num_splits = 1 bitwise anchor remains a hard
merge gate.

The independence defect stands and is restated: this package was written by the
implementer, so PR-B/PR-C still need a second-party review of the diff and the
before/after performance data.

Co-authored-by: CommandCodeBot <noreply@commandcode.ai>
@holtwood
holtwood changed the base branch from tllm-dpa-pr5-benchmark to master September 15, 2026 04:21
@holtwood
holtwood merged commit efa5f2b into master Sep 15, 2026
2 checks passed
@holtwood
holtwood deleted the tllm-attn-splitkv-design branch September 16, 2026 06:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant