[PR #739] [CLOSED] fix: invalidate non-sync dispatch padding metadata #744

New issue

Closed

opened 2026-10-27 21:05:48 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:05:48 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/739
Author: @waizuichougou
Created: 8/26/2026
Status: ❌ Closed

Base: main ← Head: fix/elastic-dispatch-padding-metadata

📝 Commits (2)

69db086 fix: invalidate non-sync dispatch padding metadata
b6a8911 test: clarify dispatch padding allocator poisoning

📊 Changes

2 files changed (+36 additions, -2 deletions)

View changed files

📝 deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh (+20 -2)
📝 tests/elastic/test_ep.py (+16 -0)

📄 Description

Summary

Initialize the unused suffix of recv_topk_idx to -1 for non-expand
dispatches that allocate worst-case output buffers without a CPU
synchronization.

This fixes a correctness issue encountered by vLLM's DeepEP V2
decode/CUDA Graph path, where stale padding entries can be interpreted as
valid expert IDs and corrupt routing data. vLLM currently works around the
issue by sanitizing the metadata downstream:

github.com/vllm-project/vllm@04c2a8deac

The initialization is performed by the existing dispatch copy epilogue, so it
does not add another kernel launch or require a host synchronization.

Problem

When ElasticBuffer.dispatch runs with do_cpu_sync=False and
do_expand=False, its output tensors are allocated for the worst-case number
of received tokens. The dispatch copy epilogue reads the actual received token
count from the GPU prefix sum, but it only writes the valid prefix of each
output tensor.

Because recv_topk_idx is allocated with torch::empty, rows after the actual
received count may contain stale allocator data. A consumer that processes the
complete fixed-shape CUDA Graph buffer can interpret a stale, non-negative
value as a valid local expert ID. This can add padding rows to expert routing
lists and corrupt the results of valid tokens.

The linked vLLM fix sanitizes the same uninitialized DeepEP V2 metadata in its
downstream decode/CUDA Graph integration.

DeepEP's legacy dispatch kernels also use -1 for unused recv_topk_idx
entries, so the fix makes the elastic path follow the existing invalid-expert
sentinel convention.

Fix

Before replacing the worst-case num_recv_tokens argument with the
GPU-resident actual count, record whether the epilogue received a worst-case
output allocation. For non-expand dispatch, use the existing epilogue
kernel/grid to write -1 to recv_topk_idx[actual_count:capacity].

This change:

covers both normal and cached non-sync dispatch;
adds no kernel launch and no host synchronization;
is compiled out for expand dispatch;
does not change the observable result of CPU-synchronized dispatch; and
supports both 32-bit and 64-bit topk_idx_t.

Tests

Added allocator-poisoning assertions for normal and cached non-expand,
non-sync dispatch.
Built the full extension for SM90 with CUDA 12.8 and NCCL 2.30.7.
Compiled the changed CUDA template for SM90 with both 32-bit and 64-bit
topk_idx_t.
Ran a standalone CUDA suffix-sanitization probe on both available NVIDIA A100
GPUs for actual counts 0, 1, 2, capacity - 1, and capacity.
Ran the standalone probe under compute-sanitizer --tool memcheck; result:
ERROR SUMMARY: 0 errors.
Ran Python syntax checks and git diff --check.

CI note

The repository's current format.sh reports pre-existing B023 findings in
tests/elastic/test_ep.py and YAPF formatting differences against the current
main version of that file. This change does not introduce a new Ruff finding;
the focused syntax and whitespace checks listed above pass.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/739 **Author:** [@waizuichougou](https://github.com/waizuichougou) **Created:** 8/26/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/elastic-dispatch-padding-metadata` --- ### 📝 Commits (2) - [`69db086`](https://github.com/deepseek-ai/DeepEP/commit/69db086b6ab741081ab6ae53e4642b9317af9ab5) fix: invalidate non-sync dispatch padding metadata - [`b6a8911`](https://github.com/deepseek-ai/DeepEP/commit/b6a8911210b011582392b3d1a280d57fea6cf3ae) test: clarify dispatch padding allocator poisoning ### 📊 Changes **2 files changed** (+36 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh` (+20 -2) 📝 `tests/elastic/test_ep.py` (+16 -0) </details> ### 📄 Description ## Summary Initialize the unused suffix of `recv_topk_idx` to `-1` for non-expand dispatches that allocate worst-case output buffers without a CPU synchronization. This fixes a correctness issue encountered by vLLM's DeepEP V2 decode/CUDA Graph path, where stale padding entries can be interpreted as valid expert IDs and corrupt routing data. vLLM currently works around the issue by sanitizing the metadata downstream: https://github.com/vllm-project/vllm/commit/04c2a8deac44fdb1ca3e2b5ec3e6bf16f3f6a914 The initialization is performed by the existing dispatch copy epilogue, so it does not add another kernel launch or require a host synchronization. ## Problem When `ElasticBuffer.dispatch` runs with `do_cpu_sync=False` and `do_expand=False`, its output tensors are allocated for the worst-case number of received tokens. The dispatch copy epilogue reads the actual received token count from the GPU prefix sum, but it only writes the valid prefix of each output tensor. Because `recv_topk_idx` is allocated with `torch::empty`, rows after the actual received count may contain stale allocator data. A consumer that processes the complete fixed-shape CUDA Graph buffer can interpret a stale, non-negative value as a valid local expert ID. This can add padding rows to expert routing lists and corrupt the results of valid tokens. The linked vLLM fix sanitizes the same uninitialized DeepEP V2 metadata in its downstream decode/CUDA Graph integration. DeepEP's legacy dispatch kernels also use `-1` for unused `recv_topk_idx` entries, so the fix makes the elastic path follow the existing invalid-expert sentinel convention. ## Fix Before replacing the worst-case `num_recv_tokens` argument with the GPU-resident actual count, record whether the epilogue received a worst-case output allocation. For non-expand dispatch, use the existing epilogue kernel/grid to write `-1` to `recv_topk_idx[actual_count:capacity]`. This change: - covers both normal and cached non-sync dispatch; - adds no kernel launch and no host synchronization; - is compiled out for expand dispatch; - does not change the observable result of CPU-synchronized dispatch; and - supports both 32-bit and 64-bit `topk_idx_t`. ## Tests - Added allocator-poisoning assertions for normal and cached non-expand, non-sync dispatch. - Built the full extension for SM90 with CUDA 12.8 and NCCL 2.30.7. - Compiled the changed CUDA template for SM90 with both 32-bit and 64-bit `topk_idx_t`. - Ran a standalone CUDA suffix-sanitization probe on both available NVIDIA A100 GPUs for actual counts `0`, `1`, `2`, `capacity - 1`, and `capacity`. - Ran the standalone probe under `compute-sanitizer --tool memcheck`; result: `ERROR SUMMARY: 0 errors`. - Ran Python syntax checks and `git diff --check`. ## CI note The repository's current `format.sh` reports pre-existing `B023` findings in `tests/elastic/test_ep.py` and YAPF formatting differences against the current `main` version of that file. This change does not introduce a new Ruff finding; the focused syntax and whitespace checks listed above pass. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-27 21:05:48 +00:00