deepseek-ai / deepseek-ai/DeepEP

Legacy normal kernels trap for hidden >= 8192: TMA staging buffer caps hidden with no host-side check

Open
#722 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.