[GH-ISSUE #722] Legacy normal kernels trap for hidden >= 8192: TMA staging buffer caps hidden with no host-side check #409

New issue

Open

opened 2026-10-26 10:00:14 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 10:00:14 +00:00

Owner

Originally created by @yashkgp on GitHub (Aug 12, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/722

Summary

The legacy "normal" (high-throughput) dispatch/combine kernels stage each token through a
fixed-size per-warp TMA buffer in shared memory. The buffer size is a hard-coded constexpr, but the
staged payload grows with hidden, so hidden is silently capped. Exceeding the cap is caught only
by a device-side assertion, which traps inside the kernel.

As a result hidden = 8192 cannot run on the normal path at all, even though the low-latency
kernels explicitly instantiate it (case 8192 in SWITCH_HIDDEN, csrc/kernels/legacy/launch.cuh).

Affected sites

site	runtime constraint	max BF16 `hidden` (`num_topk=8`)
`csrc/kernels/legacy/intranode.cu:304` (dispatch)	`hidden_int4 / 2 * 16 + 8 <= 8192`	8184
`csrc/kernels/legacy/internode.cu:578` (dispatch forwarder/receiver)	`num_bytes_per_token + 8 <= 16384`	8144
`csrc/kernels/legacy/internode.cu:1816` (combine sender)	`num_bytes_per_token + 8 <= 16384`	8160

The combine site is also hit by FP8 dispatch, because combine is always BF16.

intranode combine and internode combine's forwarder path are not affected: their budgets are
checked with EP_STATIC_ASSERT against hidden-independent quantities.

Why it fails badly

EP_DEVICE_ASSERT (deep_ep/include/deep_ep/common/exception.cuh:41) is printf + asm("trap;")
with no NDEBUG escape, and there is no host-side check for these quantities anywhere. So the
failure is not a Python-level error; it is an in-kernel trap that surfaces as an asynchronous
CUDA error: an illegal instruction was encountered, frequently reported against an unrelated later
call. The only clue is a printf from the device.

Reproduction

# 7168 (the default) sits just under every cap, which is why CI/tests do not catch this
python tests/legacy/test_intranode.py --hidden 7168   # ok
python tests/legacy/test_intranode.py --hidden 8192   # Assertion failed: .../intranode.cu:304

Any BF16 dispatch with hidden >= 8192 reproduces it, including the backward pass of a model with
hidden = 8192 (Llama-3-70B, Qwen2-72B), where dispatch of gradients is BF16 even if the forward
dispatch is FP8.

Suggested fix

hidden should not be bounded by an internal staging buffer. Concretely:

intranode.cu dispatch can be made hidden-agnostic cheaply: it already splits each token
into 2 TMA chunks, so it only needs to split into as many chunks as the buffer can hold.
internode.cu combine sender has free headroom: the forwarder warps already size this
kernel's dynamic shared memory (9248 * 24 = 221952 B), while the senders only claim
16384 * 8 = 131072 B. Deriving the sender budget from the forwarder budget raises its cap with
the total dynamic shared memory unchanged.
internode.cu dispatch stages a whole token (data + SourceMeta + scales + top-k) in one
TMA, so lifting its cap needs either chunking or a larger budget (16384 -> 24576 would grow
dynamic shared memory from 128KB to 217KB and shrink L1, so it wants a benchmark).

At minimum, these limits should be checked on the host so they raise an actionable exception
instead of trapping in the kernel.

PR with 1, 2 and host-side checks for 3: yashkgp/DeepEP#fix/legacy-normal-tma-hidden-limit

Note on verification

The caps above were derived from the source and checked exhaustively with a script, not on hardware
— I do not currently have an NVIDIA GPU to run the repro on. Please treat the exact numbers as
"derived from the formulas in the two files" rather than measured.

Originally created by @yashkgp on GitHub (Aug 12, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/722 ### Summary The legacy "normal" (high-throughput) `dispatch`/`combine` kernels stage each token through a fixed-size per-warp TMA buffer in shared memory. The buffer size is a hard-coded `constexpr`, but the staged payload grows with `hidden`, so `hidden` is silently capped. Exceeding the cap is caught only by a **device-side** assertion, which `trap`s inside the kernel. As a result **`hidden = 8192` cannot run on the normal path at all**, even though the low-latency kernels explicitly instantiate it (`case 8192` in `SWITCH_HIDDEN`, `csrc/kernels/legacy/launch.cuh`). ### Affected sites | site | runtime constraint | max BF16 `hidden` (`num_topk=8`) | |---|---|---| | `csrc/kernels/legacy/intranode.cu:304` (dispatch) | `hidden_int4 / 2 * 16 + 8 <= 8192` | **8184** | | `csrc/kernels/legacy/internode.cu:578` (dispatch forwarder/receiver) | `num_bytes_per_token + 8 <= 16384` | **8144** | | `csrc/kernels/legacy/internode.cu:1816` (combine sender) | `num_bytes_per_token + 8 <= 16384` | **8160** | The combine site is also hit by **FP8** dispatch, because combine is always BF16. `intranode` `combine` and `internode` `combine`'s forwarder path are not affected: their budgets are checked with `EP_STATIC_ASSERT` against hidden-independent quantities. ### Why it fails badly `EP_DEVICE_ASSERT` (`deep_ep/include/deep_ep/common/exception.cuh:41`) is `printf` + `asm("trap;")` with no `NDEBUG` escape, and there is no host-side check for these quantities anywhere. So the failure is not a Python-level error; it is an in-kernel trap that surfaces as an asynchronous `CUDA error: an illegal instruction was encountered`, frequently reported against an unrelated later call. The only clue is a `printf` from the device. ### Reproduction ```bash # 7168 (the default) sits just under every cap, which is why CI/tests do not catch this python tests/legacy/test_intranode.py --hidden 7168 # ok python tests/legacy/test_intranode.py --hidden 8192 # Assertion failed: .../intranode.cu:304 ``` Any BF16 `dispatch` with `hidden >= 8192` reproduces it, including the backward pass of a model with `hidden = 8192` (Llama-3-70B, Qwen2-72B), where dispatch of gradients is BF16 even if the forward dispatch is FP8. ### Suggested fix `hidden` should not be bounded by an internal staging buffer. Concretely: 1. **`intranode.cu` dispatch** can be made `hidden`-agnostic cheaply: it already splits each token into 2 TMA chunks, so it only needs to split into as many chunks as the buffer can hold. 2. **`internode.cu` combine sender** has free headroom: the *forwarder* warps already size this kernel's dynamic shared memory (`9248 * 24 = 221952 B`), while the senders only claim `16384 * 8 = 131072 B`. Deriving the sender budget from the forwarder budget raises its cap with the total dynamic shared memory unchanged. 3. **`internode.cu` dispatch** stages a whole token (data + `SourceMeta` + scales + top-k) in one TMA, so lifting its cap needs either chunking or a larger budget (`16384 -> 24576` would grow dynamic shared memory from 128KB to 217KB and shrink L1, so it wants a benchmark). At minimum, these limits should be checked on the **host** so they raise an actionable exception instead of trapping in the kernel. PR with 1, 2 and host-side checks for 3: yashkgp/DeepEP#fix/legacy-normal-tma-hidden-limit ### Note on verification The caps above were derived from the source and checked exhaustively with a script, not on hardware — I do not currently have an NVIDIA GPU to run the repro on. Please treat the exact numbers as "derived from the formulas in the two files" rather than measured.

bsakaguchi referenced this issue

2026-10-27 21:03:58 +00:00

[PR #409] [CLOSED] atomic_clean_flag in low latency combine, seems useless ? #554