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

[yashkgp]

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:

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.