Fix legacy normal kernels trapping when hidden exceeds the TMA staging buffer - #723
Fix legacy normal kernels trapping when hidden exceeds the TMA staging buffer#723yashkgp wants to merge 2 commits into
Conversation
…g buffer
The legacy high-throughput dispatch/combine kernels stage every token through a
fixed-size per-warp TMA buffer, but the staged payload grows with `hidden`. The
only guard is an `EP_DEVICE_ASSERT`, which `trap`s inside the kernel, so an
oversized hidden shows up as an unrelated asynchronous CUDA error instead of a
diagnosable failure. With BF16 and `num_topk=8` the effective caps were:
- intranode dispatch : hidden <= 8184
- internode dispatch fwd/recv: hidden <= 8144
- internode combine sender : hidden <= 8160 (also hit by FP8 dispatch,
since combine is always BF16)
So `hidden=8192` could not run on the normal path at all, even though the
low-latency kernels explicitly instantiate it in `SWITCH_HIDDEN`.
- intranode dispatch: stage each token in as many TMA chunks as the buffer can
hold, instead of hard-coding two halves. Every hidden that fitted in two
chunks before keeps the exact same split, so behaviour and shared memory
usage are unchanged there; larger hidden sizes now simply use more chunks.
- internode combine: the forwarder warps already size this kernel's dynamic
shared memory, so the senders were leaving ~11KB per warp unused. Derive the
sender budget from the forwarder one, which raises its cap to 13840 while
keeping the total dynamic shared memory byte-for-byte identical (221952 B).
- internode dispatch: keep the budget as is (raising it would grow dynamic
shared memory from 128KB to 217KB and shrink L1), but check the limit on the
host so it fails immediately with an actionable message.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`format.sh` reformats every file a change touches, so this normalizes the pre-existing lines of the two files above with the pinned clang-format 15.0.7. No functional change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
| EP_HOST_ASSERT(get_num_bytes_per_token(hidden_int4, num_scales, num_topk, num_topk) + sizeof(uint64_t) <= kNumTMABytesPerWarp and | ||
| "Hidden size is too large for the internode dispatch TMA buffer"); |
There was a problem hiding this comment.
🟡 warning: 将容量检查移到 CUDA 流切换之前: 当 allocate_on_comm_stream=True 且 hidden 超限时，调用方已在 csrc/legacy/buffer.hpp:992 将当前流切换为通信流；这里抛出的异常会跳过 buffer.hpp:1228 的恢复逻辑，导致调用者捕获异常后，后续计算仍被静默提交到通信流。combine 新增的检查也有相同问题；应在切换流之前校验，或使用 RAII 保证异常时恢复。
🤖 v6
| #ifndef DISABLE_SM90_FEATURES | ||
| #pragma unroll | ||
| for (int i = 0; i < 2; ++i) { | ||
| for (int offset_int4 = 0; offset_int4 < hidden_int4; offset_int4 += tma_chunk_int4) { |
There was a problem hiding this comment.
🔵 suggestion: 接收端拷贝循环从 '#pragma unroll for (i < 2)' 常量边界改为按 tma_chunk_int4 动态步进后，编译器不再能完全展开该循环，这是本 PR 唯一的潜在性能疑点。合入前建议由有硬件的 reviewer 在典型配置（如 hidden=7168、BF16）下跑一次 intranode dispatch 吞吐对比，确认无可测量回退；同时执行 'python tests/legacy/test_intranode.py --hidden 8192' 验证原本 trap 的用例现在通过。
🤖 v5
| auto half_hidden_bytes = half_hidden_int4 * static_cast<int>(sizeof(int4)); | ||
| // Stage each token through the per-warp TMA buffer, using as few chunks as the buffer can hold. | ||
| // NOTES: the m-barrier lives right after the staged chunk, so a chunk may only use | ||
| // `kNumTMABytesPerWarp - sizeof(uint64_t)` bytes. We always use at least 2 chunks, so that the |
There was a problem hiding this comment.
🔵 suggestion: PR 描述与提交信息中的旧上限数字有一处出入，建议修正以免误导后续读者：描述称旧 intranode dispatch 上限为 hidden ≤ 8184，但旧代码同时断言 hidden_int4 % 2 == 0，故实际旧上限是 8176（hidden_int4=1022）；8184（hidden_int4=1023，奇数）在旧代码下会触发断言，是本改动移除偶数约束后才首次可用。这不影响代码正确性，只需更新文字表述。
🤖 v5
🤖 ds-review-bot Code Reviewv6核心 TMA 分块和共享内存预算调整看起来合理，但新增的可诊断异常在公开的通信流分配模式下会泄漏当前 CUDA 流状态，影响异常恢复后的后续计算。 v5改动与描述一致，审查通过（LGTM，附带两条 suggestion）。提交 4e38668 修复了 legacy normal dispatch/combine 内核在 hidden 超出每 warp TMA 暂存缓冲时以 EP_DEVICE_ASSERT trap（表现为异步 'illegal instruction' CUDA 错误）的问题。逐项核验结论：(1) intranode.cu dispatch 的动态分块公式已解析证明对所有 hidden_int4 ≥ 1 都满足 kNumTMABytesPerWarp=8192 约束（kNumMaxTMAChunkInt4=511；h≤1022 时 chunk=⌈h/2⌉≤511，h>1022 时 n=⌈h/511⌉ ⇒ ⌈h/n⌉≤511），比描述中穷举 [1,200000] 更强；原先所有可运行的配置（偶数 hidden_int4≤1022）得到 num_tma_chunks==2、chunk==hidden_int4/2 的完全相同切分，mbarrier 偏移与旧布局一致，行为与共享内存用量位相同；尾块 min(...) 处理及 hidden_int4=1 边界均安全；grep 确认 half_hidden 无残留引用。(2) internode.cu combine 的 sender 预算 align_down(924824/8,16)=27744B（本身 16 对齐），smem_size=max(277448, 924824)=221952B 与原值逐字节相同；EP_STATIC_ASSERT(>=16384) 有效防止预算回退；按 num_bytes_per_token=align_up(hidden_bytes+8(SourceMeta)+num_topk4,16) 核算，BF16+num_topk=8 下新上限恰为 13840，与描述一致；host 检查显式以 nv_bfloat16 换算正确（combine 始终以 nv_bfloat16 实例化），与内核 sender 侧 EP_DEVICE_ASSERT 逐项对应。(3) internode.cu dispatch 的 host 侧 EP_HOST_ASSERT 公式与内核第 527/581 行完全一致，能在 launch 前以可诊断异常替代 trap；保持 16384 预算（避免动态共享内存 128KB→217KB 挤压 L1）是合理的保守选择；'cond and "msg"' 写法与仓库既有用法（internode_ll.cu:510、nccl.cu:104）一致。(4) 提交 ba6a33b（Run format.sh）逐 hunk 检查确认为纯换行/缩进重排（internode.cu 长表达式折行、intranode.cu BARRIER_LAUNCH_CASE 缩进），无语义变化且未触碰新增代码。验证限制：本环境无 GPU 且 shell 受限，无法编译、无法运行 tests/legacy/test_intranode.py --hidden 8192 / test_internode.py，也无法用 clang-format 15.0.7 复核格式；需有硬件的 reviewer 按描述执行运行时检查。 v4p该 MR 修复 legacy normal 路径下 hidden 尺寸超过 TMA staging buffer 时内核 trap 的问题：intranode dispatch 改为按 TMA buffer 容量动态分块，internode combine 提高 sender warp 的 TMA 预算，并对 internode dispatch/combine 增加主机侧 EP_HOST_ASSERT。整体实现与既有模式一致，未发现功能性问题。 Files reviewed: 2 |
2 participants
Fixes #722.
The legacy normal
dispatch/combinekernels stage each token through a fixed-size per-warp TMAbuffer, but the staged payload grows with
hidden. The only guard isEP_DEVICE_ASSERT, whichtraps inside the kernel, so an oversizedhiddensurfaces as an asynchronousCUDA error: an illegal instruction was encounteredrather than a diagnosable failure. With BF16 andnum_topk=8the effective caps werehidden <= 8184(intranode dispatch),<= 8144(internodedispatch) and
<= 8160(internode combine sender), sohidden = 8192could not run on the normalpath at all — while the low-latency kernels explicitly instantiate
case 8192.Changes
intranode.cudispatch — lift the cap. Stage each token in as many TMA chunks as the buffer canhold instead of hard-coding two halves:
Any
hiddenthat fitted in two chunks before yieldsnum_tma_chunks == 2andtma_chunk_int4 == hidden_int4 / 2, i.e. the identical split, so behaviour and shared memory usageare unchanged for every configuration that works today. Larger
hiddensimply uses more chunks. Thisalso drops the
hidden_int4 % 2 == 0requirement, sinceceil_divhandles the remainder.internode.cucombine — free headroom. The forwarder warps already set this kernel's dynamicshared memory to
9248 * 24 = 221952 B, while the senders claimed only16384 * 8 = 131072 B.Deriving the sender budget from the forwarder one gives
27744 Bper sender warp and raises the capfrom 8160 to 13840, with
smem_sizestaying byte-for-byte identical at221952 B(
max(27744 * 8, 9248 * 24) == max(16384 * 8, 9248 * 24)). AnEP_STATIC_ASSERTpins the budget soit can never shrink below the previous 16384.
internode.cudispatch — host-side check only. This one stages a whole token (data +SourceMeta+ scales + top-k) in a single TMA, so lifting its cap needs chunking or a bigger budget;16384 -> 24576would grow dynamic shared memory from 128KB to 217KB and shrink L1, which wants abenchmark I cannot run. Left as is, but the limit is now checked on the host so it raises an
immediate, actionable
EP_HOST_ASSERTinstead of trapping in the kernel. Same check added for thecombine sender.
The second commit is just
bash format.sh(clang-format 15.0.7, as pinned inrequirements-lint.txt) normalizing the pre-existing lines of the two touched files, sinceformat.shreformats whole files. It contains no functional change and touches none of the new code.Verification — please read
I could not build or run this: no NVIDIA GPU available. What I did verify, by deriving both the host
and kernel formulas from the source and checking them exhaustively:
kNumTMABytesPerWarpfor everyhidden_int4in[1, 200000];bit-identical for
hidden <= 8184;smem_sizeis unchanged at221952 B.Runtime check for a reviewer with hardware — this traps on
mainand should pass here:python tests/legacy/test_intranode.py --hidden 8192 python tests/legacy/test_internode.py --hidden 8192 # still expected to raise EP_HOST_ASSERT (internode dispatch cap)A throughput sanity check at
hidden = 7168would also be worth doing, to confirm the intranodedispatch loop becoming dynamically bounded costs nothing measurable.
🤖 Generated with Claude Code