[Bugfix] Fix DeepSeek-V4 non-EP TBO for attention TP > 1

[mikekg]

Motivation

Depends on #31700.

DeepSeek-V4 Two-Batch Overlap (TBO) with DP attention,
moe_a2a_backend=none, and attention tensor parallelism greater than one
applies full-TP variable-length collectives to tensors replicated within each
attention-TP group. Treating those replicas as independent shards corrupts the
gather and combine layouts and produces invalid model output.

This change gives each full-TP rank one shard of its replicated DP-local tensor
and reconstructs the DP-local result across the attention-TP group after the
combine.

Modifications

• Derive and validate DP-local token counts and full-TP shard counts for each
  TBO child batch.
• Shard replicated hidden states and input IDs across attention-TP ranks before
  the full-TP variable-length gather.
• Reduce-scatter the MoE output into one shard per full-TP rank, then
  all-gather within the attention-TP group to reconstruct the replicated
  DP-local output.
• Compute the replicated TP1 shared expert on the DP-local tensor and add it
  after the reducing combine.
• Select the non-EP TBO strategy for CP1, multi-rank DP, and SUM_LEN padding.
• Exercise attention TP in the registered DeepSeek-V4 TBO accuracy test and
  add focused tests for collective-size mapping and validation.

Accuracy Tests

H100 end-to-end setup

• Model: nvidia/DeepSeek-V4-Pro-NVFP4
• Hardware: 2 hosts × 8 NVIDIA H100 80GB
• Container: lmsysorg/sglang:nightly-dev-cu13-20260801-e4c4faf8
• Topology: TP16 / DP2 / attention-TP8
• MoE A2A backend: none
• TBO enabled
• Infrastructure: NemoRL Sandbox Infrastructure Report

On host 0:

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

On host 1:

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

Before: failure reproduction

The failure was reproduced both on pristine upstream 574ead753 and on the
parent of this change, 68d70aa8f. Both reached the real two-chunk TBO path
and returned HTTP 200 responses with corrupted generated text.

The following server command was run on both hosts:

unset SGLANG_OPT_DEEPGEMM_HC_PRENORM
unset SGLANG_OPT_USE_TILELANG_MHC_PRE
unset SGLANG_OPT_USE_TILELANG_INDEXER
export SGLANG_SHARED_EXPERT_TP1=1
export SGLANG_DP_USE_GATHERV=1
export SGLANG_DP_USE_REDUCE_SCATTER=1
export SGLANG_TBO_DEBUG=1

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

This exact probe and assertion were run after the server became ready:

curl --connect-timeout 2 --max-time 180 -sS \
  "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 7 + 5? Answer with just the number."}],"temperature":0,"max_tokens":32}' \
  > single.json
jq -r '.choices[0].message.content' single.json | tee single.txt
grep -q 12 single.txt

Pristine upstream produced:

我这里特 connKadaghanJF by,]_  . coding sni (* t LSD hob abcd[_ Witt vizunal 01 discern Amp Loy Vincent eb Wikisource  bur

grep exited 1: FAIL.

The parent of this change produced:

رشف3223232323232323232323a3a3a3a

grep exited 1: FAIL.

The parent run was Slurm job 5616864 and also recorded 48 rank-level
is_enable_two_chunk=True events.

The pristine-upstream run recorded 48 rank-level
is_enable_two_chunk=True events, representing three global two-chunk
forwards. Slurm job 5618929 completed 0:0; completion means the harness
successfully captured the expected failure.

After: patched result

Source 6c9e7a54a was run with the same model, container, hosts, and topology.
The following server command was run on both hosts:

unset SGLANG_OPT_DEEPGEMM_HC_PRENORM
unset SGLANG_OPT_USE_TILELANG_MHC_PRE
unset SGLANG_OPT_USE_TILELANG_INDEXER
export SGLANG_SHARED_EXPERT_TP1=1
export SGLANG_DP_USE_GATHERV=1
export SGLANG_DP_USE_REDUCE_SCATTER=1
export SGLANG_TBO_DEBUG=1

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

This exact probe and assertion were then run:

curl --connect-timeout 2 --max-time 180 -sS \
  "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 7 + 5? Answer with just the number."}],"temperature":0,"max_tokens":32}' \
  > single.json
jq -r '.choices[0].message.content' single.json | tee single.txt
grep -q 12 single.txt

Patched output:

12

grep exited 0: PASS.

The following multi-answer check was run twice so both DP routes served it:

curl --connect-timeout 2 --max-time 180 -sS \
  "http://$HEAD_NODE:8000/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d '{"model":"deepseek-ai/DeepSeek-V4-Pro","messages":[{"role":"user","content":"Answer each item concisely. 1. What is the capital of France? 2. What is the capital of Japan? 3. What is 4 times 5 plus 2?"}],"temperature":0,"max_tokens":192}' \
  > facts.json
jq -r '.choices[0].message.content' facts.json | tee facts.txt
grep -qi Paris facts.txt &&
grep -qi Tokyo facts.txt &&
grep -q 22 facts.txt

Both runs produced:

1. Paris
2. Tokyo
3. 22

All three assertions exited 0: PASS.

The sky/seasons check also passed on both DP routes: its assertions found
blue, scatter, season, and tilt.

The full GSM8K sweep used:

OPENAI_API_KEY=EMPTY HF_HOME="$HF_HOME" "$EVAL_PY" -m lm_eval \
  --model local-chat-completions \
  --apply_chat_template \
  --include_path "$EVAL_TASK_DIR" \
  --tasks gsm8k \
  --output_path "$RUN_DIR/gsm8k" \
  --log_samples \
  --model_args "model=deepseek-ai/DeepSeek-V4-Pro,base_url=http://$HEAD_NODE:8000/v1/chat/completions,api_key=EMPTY,eos_string=</s>,max_retries=5,num_concurrent=16,timeout=1800,tokenized_requests=False,max_length=2200" \
  --gen_kwargs 'max_tokens=512,temperature=0,top_p=1'

Requests	Flexible exact match	Strict exact match
1,319 / 1,319	0.9598 ± 0.0054	0.9598 ± 0.0054

Runtime coverage and stability:

• 9,864 is_enable_two_chunk=True rank records
• Maximum logged SWA token usage: 0.80
• Request errors: 0
• Retractions: 0
• Workload CUDA/OOM/server errors: 0
• Slurm job 5616859: COMPLETED 0:0

Focused collective-size tests passed 3/3 in the matching container (Slurm job
5618770, COMPLETED 0:0).

Speed Tests and Profiling

The 1,319-request GSM8K API phase completed in 9m01s at concurrency 16
(2.44 completed requests/s). The complete job, including checkpoint staging,
weight loading, JIT compilation, CUDA-graph capture, deterministic checks, and
GSM8K, completed in 26m03s.

The control configurations produce invalid output, so their generated-token
throughput is not a correctness-preserving performance baseline.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Documentation applicability reviewed; no user-facing documentation change is needed.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #34951900568
Latest PR Test (Extra): ❌ Run #34951900219
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[nvpohanh]

[by Codex] Review comments are inline.

[nvpohanh]

[by Codex] Severity: style | Confidence: High

Please avoid both dynamic getattr probes in this modified expression. DeepseekV4DecoderLayer.mlp is a DeepseekV2MoE, so expose typed state—for example, initialize its optional shared_experts member to None and read self.mlp.shared_experts and self.mlp._shared_expert_tp1 directly.

[nvpohanh]

[by Codex] Severity: style | Confidence: High

Black reformats this expression, and the PR's lint check currently fails. Please apply the repository formatter and commit the resulting flattening of this boolean expression.

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci