fix(dp-attn): DP-attention buffer sizing + EAGLE padding for require_attn_tp_gather=True

[ashutoshuiuc]

Summary

Fixes 7 bugs triggered when running DP-attention with require_mlp_tp_gather=False and require_attn_tp_gather=True (e.g. moe_dense_tp_size=1, ep_size=32, dp_size=4, attn_tp_size=8). Discovered while running GLM-4.5-355B-A32B training with EAGLE speculative decoding on 32 nodes (128 GPUs).

Fixes

1. global_num_tokens_gpu buffer sizing (scheduler_dp_attn_mixin.py, cuda_graph_runner.py, eagle_draft_cuda_graph_runner.py, eagle_draft_extend_cuda_graph_runner.py, model_runner.py, forward_batch_info.py)

Buffer allocation and _update_gather_batch only checked require_mlp_tp_gather to decide between dp_size vs 1 element. With our config (mlp=False, attn=True), get_dp_local_info() indexed out-of-bounds by dp_rank, producing garbage that crashed flashinfer with q.shape[0] does not match qo_indptr[-1].

Fix: change all guards to if require_mlp_tp_gather or require_attn_tp_gather.

2. EAGLE draft extend padding for DP-attention (flashinfer_backend.py, eagle_info.py)

EAGLE draft extend has num_tokens = batch_size * num_draft_tokens, not necessarily divisible by attn_tp_size. The DP-attention all-gather requires alignment but flashinfer metadata (qo_indptr, kv_indptr, kv_indices) wasn't padded to match.

Fix: append dummy padding request to flashinfer metadata covering the extra positions, same pattern the non-speculative extend path already uses for piecewise CUDA graph padding.

3. Truncate padded logits in EAGLE verify (eagle_info.py)

After DP-attention-padded forward, logits have ceil_align(n, attn_tp_size) rows but verify expects exactly bs * draft_token_num.

Fix: truncate next_token_logits before verification.

4. Reset heartbeat before declaring server ready (http_server.py)

CUDA graph capture takes ~53s for 355B. TokenizerManager.last_receive_tstamp was set at object creation, so /health_generate returned 503 immediately after startup. With DP-attention, one "unhealthy" shard causes NCCL hangs on all others.

Fix: reset last_receive_tstamp right before "fired up and ready to roll" log.

Config that triggers these bugs

--sglang-tp-size 32
--sglang-dp-size 4
--sglang-ep-size 32
--sglang-moe-dense-tp-size 1
--sglang-enable-dp-attention
--sglang-enable-dp-lm-head
--sglang-speculative-algorithm EAGLE

This gives attn_tp_size = tp_size / dp_size = 8, require_mlp_tp_gather = False (because moe_dense_tp_size == 1), require_attn_tp_gather = True.

Test plan

• Tested with GLM-4.5-355B-A32B on 32x4-GPU nodes (128 GPUs total)
• Without these fixes: crashes immediately on first decode batch (fixes 1-2) or first EAGLE draft extend (fixes 3-6)
• With these fixes: rollout generation runs successfully through full batch generation
• Fix 7 tested by observing /health_generate returning 200 immediately after server ready instead of 503

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.