Scan epilogue tokens for NaN and Inf behind EP_DEBUG_CHECKS by hershys-aws · Pull Request #33 · amazon-contributing/DeepEP

hershys-aws · 2026-10-26T16:54:58Z

Description of changes:

EP_DEBUG_CHECKS=1 compiles a token check into the kernels. The flag is off by default. With the flag off, the kernels do not change. With the flag on, the dispatch epilogue scans each token before it stores to recv_x, and the combine epilogue scans each reduced token before it stores to combined_x. A hit prints the rank, the token, the source rank and token, the expert, and the first bad element. A hit in dispatch means the value arrived bad. A hit in combine after a clean dispatch means a partial arrived bad, or the reduce made it. The scan handles BF16 and FP8 E4M3. It tests all elements of a 32-bit word at once and has no branches on the clean path.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Xuan-1998 · 2026-10-26T22:17:41Z

Can we add some tests under tests/ to exercise this path?

hershys-aws · 2026-10-27T17:09:45Z

Can we add some tests under tests/ to exercise this path?

Added!

vladimiraerov · 2026-10-29T12:03:02Z

The change is quite invasive and big, make the divergence from the mainstream code much bigger and harder to maintain. I think we should keep it as a side branch, and if we need to use it ASAP we can quickly rebase the side branch and provide this to who ever need it.

Xuan-1998 · 2026-10-29T21:15:46Z

hange is quite invasive and big, make the divergence from the mainstream code much bigger and harder to maintain. I think we should keep it as a side branch, and if we need to use it ASAP we can quickly rebase the side branch and provide this to who ever need it.

as discussed, let's keep it on main behind a debug flag

whn09 · 2026-11-03T08:36:29Z

Tested on 2× p6-b300 (main 2f4ce10 + this PR, NCCL 2.31.2). test_debug_checks.py passes on both nodes (needs EP_NIC_NAME set to a device that exists, since it doesn't pass --num-sms). With the flag off, perf is flat and the full matrix + pressure pass. With EP_DEBUG_CHECKS=1, copy/reduce epilogues cost ~+1 µs (≤1.7%). The JIT cache key includes the flags, so toggling it is safe.

The token scan is the `tokens` check of `EP_DEBUG_CHECKS`, bit 1 in `compiled.cuh`. The host JIT generator of each kernel with the check passes the mask as the template argument `kDebugChecks`, so the value is part of the kernel cache key. The check tests its bit with `if constexpr`, so with the bit clear it adds no instructions. `diag.cuh` has the first check: a warp scan that counts the non-finite elements of one token and finds the first one. A BF16 or FP8 E4M3 element is non-finite when its exponent bits are all ones. The scan tests all elements of a 32-bit word with one add and two masks. The fast pass reads the token once, with no branches, and asks the warp "did any lane see a hit". A clean token pays only this cost. The slow pass counts and locates the bad elements. It runs only after a hit. Reports go through `EP_DIAG_REPORT`. Each kernel variant prints 32 reports in each window of about 1 s, and then the number of reports it did not print. A window that starts during a burst of reports can print up to 64. One bad buffer can hold thousands of bad tokens, and device printf is slow: 32 lines add about 110 us to a launch. The window is in time, not in launches, so the reports also start again under CUDA graph replay.

The copy epilogue stages each received token in shared memory before the TMA store to `recv_x`. When `EP_DEBUG_CHECKS` has the `tokens` check, the epilogue scans the token there. The scan adds no global memory traffic, no launch and no sync. Dispatch is a copy, so a hit means the bad value was already in the receive buffer. The report names the rank, the received token, the source rank and token, the local expert, the number of bad elements, and the index and bits of the first one. The bits separate Inf (an overflow upstream) from NaN. `get_debug_check_mask` reads the mask of `EP_DEBUG_CHECKS` once per process, and `launch_dispatch_copy_epilogue` passes it to the kernel as the template argument `kDebugChecks`. Dispatch takes BF16, or FP8 E4M3 with scale factors, so the kernel gets the element type from `kNumSFPacks`, with no other template argument. Measured on 2 x p5en (H200), hidden 3328, top-6, 8192 tokens: with the flag off, the kernel has the same instructions as before. With the flag on, the copy epilogue time is within noise for BF16 and FP8 (144 test cases, 0 false reports), and a planted NaN is reported with the correct source rank, token and element.

The reduce epilogue sums the partials of each token in shared memory before the TMA store to `combined_x`. When `EP_DEBUG_CHECKS` has the `tokens` check, the epilogue scans the sum there, with the same scan as the dispatch epilogue. The output is always BF16. The report names the rank, the token, the number of sources, the number of bad elements, and the index and bits of the first one. The host passes `EP_DEBUG_CHECKS` to the kernel in the same way as for dispatch. Together, the two scans bracket a bad value. A hit in dispatch means the value arrived bad. A hit in combine after a clean dispatch means a partial arrived bad, or the reduce made the bad value. Measured on 2 x p5en (H200), hidden 3328, top-6, 8192 tokens: with the flag off, the kernel has the same instructions as before. With the flag on, the reduce epilogue takes 3 us more in the test (70 to 73 us) and the same time in isolation under Nsight Compute. A planted NaN is reported with the correct token and element. The README documents the `tokens` check under `EP_DEBUG_CHECKS`.

`tests/elastic/test_debug_checks.py` compiles the kernels with `EP_DEBUG_CHECKS=tokens`. A clean dispatch and combine must give no report. One NaN in the dispatch input, BF16 and FP8, must give one report on each rank that receives the token, with the source rank, the source token and the element. One NaN in the combine input must give one report on the rank that owns the token. A NaN in every token must give no more reports than the rate limit allows, and the next window must print how many reports were not printed. The reports come from device printf, so the test captures file descriptor 1 around each step.

`test_corruption` compares the token check with a bitwise check against the clean data. A second stream overwrites a few tokens of the dispatch or combine input after a random delay, while the kernels read them. The check must report exactly the changed tokens that hold a NaN or Inf. The test counts the changed tokens that hold only valid numbers, which the check cannot see. A second part flips one bit in many tokens, for each bit of BF16 and FP8. The check must report exactly the flips that make a NaN or Inf. The test keeps the reports under the rate limit, so a missing report is a real miss. On 2 x p5en (H200) and 2 x p6-b300, 48 iterations of each race: every copy with a NaN or Inf had a report, a third to a half of the changed tokens held only valid numbers, and in BF16 only a flip of bit 14 made a NaN or Inf. The step adds about 12 s to the test.

hershys-aws marked this pull request as ready for review October 26, 2026 16:56

hershys-aws force-pushed the feat/epilogue-nonfinite-check branch 2 times, most recently from ba44a72 to b3ddfd9 Compare October 27, 2026 16:41

hershys-aws self-assigned this Oct 27, 2026

Xuan-1998 previously approved these changes Oct 27, 2026

View reviewed changes

Comment thread csrc/elastic/buffer.hpp Resolved

hershys-aws dismissed Xuan-1998’s stale review via 0af4589 October 27, 2026 18:35

hershys-aws force-pushed the feat/epilogue-nonfinite-check branch from b3ddfd9 to 0af4589 Compare October 27, 2026 18:35

hershys-aws requested review from Xuan-1998 and vladimiraerov October 27, 2026 18:35

vladimiraerov reviewed Oct 29, 2026

View reviewed changes

Comment thread csrc/jit/compiler.hpp Resolved

vladimiraerov mentioned this pull request Oct 29, 2026

Check the GIN signals at kernel entry behind EP_DEBUG_CHECKS #34

Open

hershys-aws force-pushed the feat/epilogue-nonfinite-check branch 2 times, most recently from 3e019e4 to 552e47a Compare October 29, 2026 19:51

Xuan-1998 previously approved these changes Oct 29, 2026

View reviewed changes

whn09 mentioned this pull request Nov 3, 2026

fix: adding logging/timeout protection #27

Merged