deepseek-ai / deepseek-ai/DeepEP
Legacy normal kernels trap for hidden >= 8192: TMA staging buffer caps hidden with no host-side check
- Dominant language
- Cuda
- Stars
- 10.1k
- Forks
- 1.4k
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 2
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 `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.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the staging-buffer checks in csrc/kernels/legacy/intranode.cu:304 and csrc/kernels/legacy/internode.cu:578,1816, then read deep_ep/include/deep_ep/common/exception.cuh:41 and the launch cases in csrc/kernels/legacy/launch.cuh. Use the listed tests at hidden 7168 and 8192 to reproduce where hardware is available. Done means oversized hidden values receive an actionable host-side failure or no longer trap in the affected normal kernels.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- distributed-systems, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100