Fix DeepSeek-V4/DeepSeek-V4-Pro DP-attention gather semantics

[mikekg]

Motivation

Fixes #31699.

DeepSeek-V4 DP-attention produces numerically invalid output when
moe_a2a_backend=none, data_parallel_size>1, and attn_tp_size>1.

By the time the MoE gather runs, self_attn has already reduced its output
across the attention-TP group. The hidden states are therefore replicated
across attention-TP ranks.

The existing dp_gather_partial calls treat those tensors as unreduced partial
contributions. Their reduce-scatter sums the replicated values, multiplying the
hidden-state magnitude by attn_tp_size at every MoE layer.

The same replicated-versus-partial distinction applies to the NextN input-ID
gather.

Modifications

• Use dp_gather_replicate for the normal post-attention hidden-state gather.
• Use dp_gather_replicate for the NextN input-ID gather.
• Clone input IDs before both main-model and NextN replicate gathers. The
  MAX_LEN implementation can zero its local input on non-leader attention-TP
  ranks, and input_ids[:, None] otherwise aliases the caller-owned tensor.

The hidden-state inputs are not cloned because they are dead after their
gathers; cloning those large activations once per layer would add unnecessary
overhead.

Reproduction

Environment

The captured A/B used:

• Two hosts with 8 NVIDIA H200 GPUs each: 16 GPUs total
• nvidia/DeepSeek-V4-Pro-NVFP4
• SGLang revision a9cf5e68e
• Container lmsysorg/sglang:nightly-dev-cu13-20260715-50d1edaa
• Global tensor parallel size 16
• Data parallel size 4, producing attention TP4 x DP4 within those 16 ranks
• moe_a2a_backend omitted, so it defaults to none
• SGLANG_SHARED_EXPERT_TP1=1
• --enable-dp-lm-head, to exclude the LM-head path as a confounding variable
• No explicit --quantization; checkpoint quantization is auto-detected

The same failure was independently reproduced with attention TP8 x DP2 and on
two hosts with 8 H100 GPUs each.

The PR diff applies cleanly to the tested a9cf5e68e revision. Use the
unmodified revision for the failing run and the same revision with this PR diff
applied for the passing run. Keep the checkpoint and every launch argument
identical.

Stage the checkpoint on a login or transfer host, not in the GPU allocation:

python3 -m venv /shared/venvs/huggingface
/shared/venvs/huggingface/bin/python -m pip install --upgrade huggingface_hub
/shared/venvs/huggingface/bin/hf download \
  nvidia/DeepSeek-V4-Pro-NVFP4 \
  --local-dir /shared/models/DeepSeek-V4-Pro-NVFP4

The model and source checkout must be visible at the same paths on both GPU
hosts.

Run one sglang.launch_server process on each host. Both processes use global
--tensor-parallel-size 16; this is 16 total ranks across the two hosts, not 16
GPUs per host.

On host 0:

export NODE_RANK=0
export HEAD_NODE='<host-0 hostname or IP>'

On host 1:

export NODE_RANK=1
export HEAD_NODE='<host-0 hostname or IP>'

Then run this command inside the SGLang container/environment on both hosts:

export SGLANG_SHARED_EXPERT_TP1=1

python3 -m sglang.launch_server \
  --model-path /shared/models/DeepSeek-V4-Pro-NVFP4 \
  --served-model-name deepseek-ai/DeepSeek-V4-Pro \
  --tensor-parallel-size 16 \
  --data-parallel-size 4 \
  --enable-dp-attention \
  --enable-dp-lm-head \
  --nnodes 2 \
  --node-rank "$NODE_RANK" \
  --dist-init-addr "$HEAD_NODE:29500" \
  --kv-cache-dtype fp8_e4m3 \
  --moe-runner-backend marlin \
  --fp4-gemm-backend marlin \
  --chunked-prefill-size 4096 \
  --disable-flashinfer-autotune \
  --mem-fraction-static 0.85 \
  --context-length 2200 \
  --cuda-graph-max-bs 256 \
  --max-running-requests 256 \
  --disable-radix-cache \
  --host 0.0.0.0 \
  --port 8000

Do not pass --quantization or --moe-a2a-backend.

Once host 0 reports ready, send:

curl -s "http://$HEAD_NODE:8000/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "deepseek-ai/DeepSeek-V4-Pro",
    "messages": [{"role": "user", "content": "What is the capital of France?"}],
    "temperature": 0,
    "max_tokens": 128
  }'

For attention TP8 x DP2 on the same 16 GPUs, change only
--data-parallel-size 4 to --data-parallel-size 2.

A single 8-GPU H200 host can run flat TP8, but flat TP does not exercise this
DP-attention gather and is not the demonstrated reproduction.

Accuracy Tests

Before: unpatched source fails

The unpatched server reports healthy and returns successful, well-formed JSON,
but the generated content is numerically corrupted.

For What is the capital of France? at temperature=0, captured pre-fix
content included:

" +1 dep voquest m LH(-- _mi(zip bell {,PPTkinsagastruct C bell financially
BOuye em A_->TR,s. gland Amar [-]-bBuff ia Jol ke quer occurring ,holt PB
sakcka..."

For What is 4 times 5 plus 2? Answer with just the number. the pre-fix server
returned:

"ffbs, bottles--> basisryl YusDL prime Scale Giber late pus , reckpp Cochrane
elem hydro Hugtered verte Jez-"

Neither response answers its prompt. The failure occurred with no server
exception and no request error.

After: patched source passes

With the patch applied, the same deterministic checks returned:

What is the capital of France?
=> The capital of France is Paris.

What is 4 times 5 plus 2?
=> 4 * 5 + 2 = 20 + 2 = 22

A store had 15 apples, sold 6, and received 20 more. How many now?
=> 29

Results:

• 3/3 prompts produced the same correct answer as the coherent flat-TP16
  baseline.
• The checks were repeated across seven independent patched server launches.
• Both TP4 x DP4 and TP8 x DP2 produced coherent output after the fix.
• No patched launch regressed to the pre-fix token-soup behavior.
• Generated-text samples were inspected throughout the 1K and 8K sweeps.
• The complete 8K sweep finished all ten concurrency levels, 1 through 512,
  with zero request errors and coherent generated text.

The synchronous path was exercised directly. The NextN call site was not enabled\nin these runtime tests; it was changed because it has the same replicated-input\ncollective semantics.

Local checks against current main:

• Ruff 0.15.1: passed
• Black 26.1.0: passed
• isort 7.0.0: passed
• py_compile: passed
• git diff --check: passed

Speed Tests and Profiling

The patched TP4 x DP4 8K-input/1K-output run completed its full sweep:

Concurrency	Completed	Total token throughput	Throughput/GPU
1	10/10	281.34 tok/s	17.584 tok/s/GPU
8	80/80	1,811.15 tok/s	113.197 tok/s/GPU
64	640/640	6,936.50 tok/s	433.531 tok/s/GPU
256	2,560/2,560	10,468.91 tok/s	654.307 tok/s/GPU
512	5,120/5,120	12,274.31 tok/s	767.144 tok/s/GPU

A pre-fix performance comparison is not meaningful because every pre-fix
generated response was numerically invalid. Those jobs were stopped after the
correctness failure was confirmed.

Checklist

• Format the code according to the pre-commit configuration.
• The affected path requires distributed attention-TP and
  DP collectives; it was validated on two-node, 16-GPU systems.
• No user-facing documentation change is required.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code-style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #30963574545
Latest PR Test (Extra): ❌ Run #30963574465

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[ch-wan]

/rerun-test test_moe_ep.py

[github-actions]

Results for /rerun-test test_moe_ep.py:

🚀 2-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/moe/test_moe_ep.py

[ch-wan]

/rerun-test test_moe_ep.py test_moe_ep_extra.py

[github-actions]

Results for /rerun-test test_moe_ep.py test_moe_ep_extra.py:

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/moe/test_moe_ep.py
cd test/ && python3 registered/moe/test_moe_ep_extra.py

[ch-wan]

/rerun-test test_deepep_small.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit cdff33d.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[ch-wan]

/rerun-test test_deepep_small.py

[github-actions]

Results for /rerun-test test_deepep_small.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/ep/test_deepep_small.py

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

All NV pipelines have passed. @ch-wan could you help to review and merge? Thanks!