Repository navigation
Legacy normal kernels trap for hidden >= 8192: TMA staging buffer caps hidden with no host-side check #722
Description
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:304Any 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.cudispatch can be madehidden-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.cucombine 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.cudispatch stages a whole token (data +SourceMeta+ scales + top-k) in one
TMA, so lifting its cap needs either chunking or a larger budget (16384 -> 24576would 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.
Activity