[PR #723] Fix legacy normal kernels trapping when hidden exceeds the TMA staging buffer #731

New issue

Open

opened 2026-10-27 21:05:42 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:05:42 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/723
Author: @yashkgp
Created: 8/12/2026
Status: 🔄 Open

Base: main ← Head: fix/legacy-normal-tma-hidden-limit

📝 Commits (2)

4e38668 Fix legacy normal kernels trapping when hidden exceeds the TMA staging buffer
ba6a33b Run format.sh

📊 Changes

2 files changed (+61 additions, -28 deletions)

View changed files

📝 csrc/kernels/legacy/internode.cu (+42 -17)
📝 csrc/kernels/legacy/intranode.cu (+19 -11)

📄 Description

Fixes #722.

The legacy normal dispatch/combine kernels stage each token through a fixed-size per-warp TMA
buffer, but the staged payload grows with hidden. The only guard is EP_DEVICE_ASSERT, which
traps inside the kernel, so an oversized hidden surfaces as an asynchronous
CUDA error: an illegal instruction was encountered rather than a diagnosable failure. With BF16 and
num_topk=8 the effective caps were hidden <= 8184 (intranode dispatch), <= 8144 (internode
dispatch) and <= 8160 (internode combine sender), so hidden = 8192 could not run on the normal
path at all — while the low-latency kernels explicitly instantiate case 8192.

Changes

intranode.cu dispatch — lift the cap. Stage each token in as many TMA chunks as the buffer can
hold instead of hard-coding two halves:

num_tma_chunks = max(2, ceil_div(hidden_int4, (kNumTMABytesPerWarp - sizeof(uint64_t)) / sizeof(int4)))
tma_chunk_int4 = ceil_div(hidden_int4, num_tma_chunks)

Any hidden that fitted in two chunks before yields num_tma_chunks == 2 and
tma_chunk_int4 == hidden_int4 / 2, i.e. the identical split, so behaviour and shared memory usage
are unchanged for every configuration that works today. Larger hidden simply uses more chunks. This
also drops the hidden_int4 % 2 == 0 requirement, since ceil_div handles the remainder.

internode.cu combine — free headroom. The forwarder warps already set this kernel's dynamic
shared memory to 9248 * 24 = 221952 B, while the senders claimed only 16384 * 8 = 131072 B.
Deriving the sender budget from the forwarder one gives 27744 B per sender warp and raises the cap
from 8160 to 13840, with smem_size staying byte-for-byte identical at 221952 B
(max(27744 * 8, 9248 * 24) == max(16384 * 8, 9248 * 24)). An EP_STATIC_ASSERT pins the budget so
it can never shrink below the previous 16384.

internode.cu dispatch — 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 -> 24576 would grow dynamic shared memory from 128KB to 217KB and shrink L1, which wants a
benchmark I cannot run. Left as is, but the limit is now checked on the host so it raises an
immediate, actionable EP_HOST_ASSERT instead of trapping in the kernel. Same check added for the
combine sender.

The second commit is just bash format.sh (clang-format 15.0.7, as pinned in
requirements-lint.txt) normalizing the pre-existing lines of the two touched files, since
format.sh reformats 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:

the new chunking fits kNumTMABytesPerWarp for every hidden_int4 in [1, 200000];
zero previously-working sizes change their chunk split, so the intranode path is
bit-identical for hidden <= 8184;
the internode combine smem_size is unchanged at 221952 B.

Runtime check for a reviewer with hardware — this traps on main and 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 = 7168 would also be worth doing, to confirm the intranode
dispatch loop becoming dynamically bounded costs nothing measurable.

🤖 Generated with Claude Code

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/723 **Author:** [@yashkgp](https://github.com/yashkgp) **Created:** 8/12/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/legacy-normal-tma-hidden-limit` --- ### 📝 Commits (2) - [`4e38668`](https://github.com/deepseek-ai/DeepEP/commit/4e38668b79720cd0f0caca8441ff9bf63ff82b08) Fix legacy normal kernels trapping when hidden exceeds the TMA staging buffer - [`ba6a33b`](https://github.com/deepseek-ai/DeepEP/commit/ba6a33b7fe938749208d8ee8b033afdc17814dbe) Run format.sh ### 📊 Changes **2 files changed** (+61 additions, -28 deletions) <details> <summary>View changed files</summary> 📝 `csrc/kernels/legacy/internode.cu` (+42 -17) 📝 `csrc/kernels/legacy/intranode.cu` (+19 -11) </details> ### 📄 Description Fixes #722. The legacy normal `dispatch`/`combine` kernels stage each token through a fixed-size per-warp TMA buffer, but the staged payload grows with `hidden`. The only guard is `EP_DEVICE_ASSERT`, which `trap`s inside the kernel, so an oversized `hidden` surfaces as an asynchronous `CUDA error: an illegal instruction was encountered` rather than a diagnosable failure. With BF16 and `num_topk=8` the effective caps were `hidden <= 8184` (intranode dispatch), `<= 8144` (internode dispatch) and `<= 8160` (internode combine sender), so `hidden = 8192` could not run on the normal path at all — while the low-latency kernels explicitly instantiate `case 8192`. ### Changes **`intranode.cu` dispatch — lift the cap.** Stage each token in as many TMA chunks as the buffer can hold instead of hard-coding two halves: ``` num_tma_chunks = max(2, ceil_div(hidden_int4, (kNumTMABytesPerWarp - sizeof(uint64_t)) / sizeof(int4))) tma_chunk_int4 = ceil_div(hidden_int4, num_tma_chunks) ``` Any `hidden` that fitted in two chunks before yields `num_tma_chunks == 2` and `tma_chunk_int4 == hidden_int4 / 2`, i.e. the identical split, so behaviour and shared memory usage are unchanged for every configuration that works today. Larger `hidden` simply uses more chunks. This also drops the `hidden_int4 % 2 == 0` requirement, since `ceil_div` handles the remainder. **`internode.cu` combine — free headroom.** The forwarder warps already set this kernel's dynamic shared memory to `9248 * 24 = 221952 B`, while the senders claimed only `16384 * 8 = 131072 B`. Deriving the sender budget from the forwarder one gives `27744 B` per sender warp and raises the cap from 8160 to 13840, with `smem_size` staying **byte-for-byte identical** at `221952 B` (`max(27744 * 8, 9248 * 24) == max(16384 * 8, 9248 * 24)`). An `EP_STATIC_ASSERT` pins the budget so it can never shrink below the previous 16384. **`internode.cu` dispatch — 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 -> 24576` would grow dynamic shared memory from 128KB to 217KB and shrink L1, which wants a benchmark I cannot run. Left as is, but the limit is now checked on the host so it raises an immediate, actionable `EP_HOST_ASSERT` instead of trapping in the kernel. Same check added for the combine sender. The second commit is just `bash format.sh` (clang-format 15.0.7, as pinned in `requirements-lint.txt`) normalizing the pre-existing lines of the two touched files, since `format.sh` reformats 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: - the new chunking fits `kNumTMABytesPerWarp` for **every** `hidden_int4` in `[1, 200000]`; - **zero** previously-working sizes change their chunk split, so the intranode path is bit-identical for `hidden <= 8184`; - the internode combine `smem_size` is unchanged at `221952 B`. Runtime check for a reviewer with hardware — this traps on `main` and should pass here: ```bash 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 = 7168` would also be worth doing, to confirm the intranode dispatch loop becoming dynamically bounded costs nothing measurable. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:05:42 +00:00

bsakaguchi referenced this issue

2026-10-27 21:05:45 +00:00

[PR #731] [CLOSED] feat(low_latency): support top-16 routing and hidden-dim 3584 #736