Repository navigation
feat(transformer): wire split-KV decode behind TLLM_ATTN_SPLITKV (TLLM-ATTN-SPLITKV PR-C) - #13
Merged
Merged
Conversation
Connects the split-KV entry points from PR-B to the production decode path. - LayerWorkspace gains `attn_partial`, sized once for kAttnMaxSplits (32) rather than for the runtime value: the switch is parsed per call (so setenv keeps working in tests without a reset seam) while the allocation stays fixed, so a changed value cannot overflow the buffer. - TLLM_ATTN_SPLITKV = <num_splits>: unset/empty/0/1 take the single-pass entry points (bit-identical to today, no extra combine launch); 2..32 take the split-KV ones. Non-numeric or out-of-range values fail loudly (G5), with the offending value and the accepted range in the message. - Routing covers both decode branches (direct and legacy); prefill is untouched, and a test pins that. Tests: 5 new cases in test_paged_dispatch.cpp, plus runMinimalDecode extracted so the two entry-error tests stop duplicating the minimal view. Off-by-default/0/1 are bitwise identical; invalid values fail loudly; legacy and direct agree bitwise at the same num_splits; split shifts numerics only slightly versus single-pass; prefill is unaffected. --- BLOCKING FINDING, recorded because it affects whether this can ship. While developing this, the full suite failed deterministically (16 tests; the first an illegal read in attention_decode_kernel) whenever the extra device allocation was present. The read was 4 bytes past the END of the KV pool -- the kernel used a visible length (~64) far larger than the step's actual 21 -- in the *existing contiguous* path (test_paged_oracle.cpp:673). Decisive control: at the PR-B state, with no split-KV code at all, injecting one independent 8 KiB cudaMalloc/cudaFree into LayerWorkspace::allocate reproduces the identical failure. So it is a pre-existing, allocation-layout-sensitive latent bug, not something split-KV introduced. It is also build-sensitive: the failing binary failed 4/4 runs, and after a revert-and-repop rebuild the same source passes 6/6. That is the dangerous kind of latent bug -- it can flip on another toolchain, which is what CI (CUDA 11.8) versus this box (13.3) represents. Consequence: split-KV adds a device allocation and therefore perturbs the same layout. This commit is green here, but it should not be trusted until that latent bug is understood and fixed as its own task. Co-authored-by: CommandCodeBot <noreply@commandcode.ai>
2 tasks
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
1 participant
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.
TLLM-ATTN-SPLITKV PR-C：接线。stacked on #12（kernel）。把 PR-B 的
_splitkv入口接到生产 decode 路径，用开关控制。新增行为
LayerWorkspace增加attn_partial，按kAttnMaxSplits = 32上界一次性分配：开关仍每次调用重新解析（setenv在测试里即时生效、不需要 reset seam），而分配大小固定，因此运行中改值不会越界。TLLM_ATTN_SPLITKV = <num_splits>：未设置/空/0/1走单遍入口（与今天逐位相同，且不多一次 combine launch）；2..32走 split-KV；非数字或越界显式报错并带上取值与合法区间（G5）。测试
test_paged_dispatch.cpp新增 5 项：默认/0/1逐位相同；非法值显式失败（abc/8x/99/+4）；同一num_splits下 legacy 与 direct 逐位相同；split 相对单遍只有小幅数值变化（尺度相关的宽松上限，精确门禁在 kernel 级）；prefill 不受开关影响。另把最小视图抽成runMinimalDecode，去掉两处入口错误测试的 20 行重复。本机验证：全量 225 passed / 11 skipped（6 次连续一致）、
compute-sanitizer0 error、clang-format 18.1.8 clean。开发过程中，只要那次额外 device 分配存在，全量测试就确定性失败（16 项，首个是
attention_decode_kernel的非法读）。我追到根因形状：读地址在 KV pool 末尾之后 4 字节，即 kernel 用了一个远大于本步实际（visible = 21）的可见长度（约 64），而且发生在既有的连续路径（tests/test_paged_oracle.cpp:673），与 split-KV 无关。决定性对照：在 PR-B 状态（完全没有 split-KV 代码）下，只在
LayerWorkspace::allocate里注入一次独立的 8 KiBcudaMalloc/cudaFree，就复现了完全相同的失败。而且它对构建敏感：失败的二进制连续 4 次失败；revert + 重新应用后重建，同一份源码连续 6 次通过。这正是最危险的一类隐患——换个工具链就可能翻面，而 CI 用的 CUDA 11.8 与本机 13.3 正是这种差异。
结论：split-KV 会增加一次 device 分配，因而会扰动同一布局。本 PR 在本机是绿的，但在那个隐患被单独定位并修掉之前，不应视为可信。本次 CI 就是一次真实检验。
限制
num_splits最优值与"小 S 是否回归"仍待 PR-D 实测。num_splits > 1改变 fp32 求和顺序（设计包 Q3，owner 已批准）。后续
PR-D（扩
--dpa-bench测新旧三路 + 扫num_splits）。另建议把上面的隐患单独立一个 bug 任务——它比 split-KV 本身更值得优先处理。