Repository navigation
docs(design): split-KV decode attention design record (TLLM-ATTN-SPLITKV) - #11
Merged
Merged
Conversation
…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>
This was referenced Sep 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
设计记录(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):sm__warps_active(achieved occupancy)__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 公式合并。num_splits是 host 参数(grid 捕获时固定),可见长度仍走 device int,partial 缓冲预分配、捕获期零分配零 D2H。decode_online_softmax→decode_online_softmax_range<Rows, kFinalize>,ContiguousRows/PagedRows、冻结的地址公式与零行语义全部不动。_splitkv入口;既有attention_decode/attention_decode_paged保持num_splits = 1的等价语义,因此 ~18 个既有调用点不用改,单遍路径同时充当回归锚点与回滚点。C ABI 不变,与paged-serving无强制顺序。需要 reviewer 明确决议的点
num_splits > 1会故意打破 §10.2 的「contiguous 路径逐位不变」指纹(求和顺序变了,仍在 oracle2e-3容差内);directvslegacy的逐位相等仍然成立(同一份循环)。缓解是 §4.4 的逐位锚点:
num_splits = 1时exp(0)==1.0f、1/(l+1e-9)不变 ⇒ 与今天逐位相同,从而把"重构风险"与"数值改动"干净分开。建议接受 Q3。num_splits走 host 参数;Q4 partial 放LayerWorkspace新缓冲(不复用 direct 路径冗余的k_scratch);Q5 小 S 是否允许调用方回落到num_splits = 1(不做内核内分支)。验证
docs npm run buildexit 0(修掉了一处指向 PR docs(design): TLLM-DPA direct paged decode attention design package (TLLM-P0-004 设计门禁) #5 未合入文档的死链)。限制
num_splits最优值、小 S 是否回归,均列为 §2 的 Unknown,由 PR-D 实测关闭。后续
批准后:PR-B(kernel + 测试)→ PR-C(LayerWorkspace + 开关 + 路由)→ PR-D(benchmark + 结果报告)。默认值只在 PR-D 通过后才讨论打开。