[ch-wan]

This PR is part of a stack (oldest at bottom):

• [Refactor] Let the remaining placement consumers read the parallel context #42348
• [Fix] Skip a draft's decode recapture when it owns no decode graph #42347
• -> [Fix] Load a draft's tensor weight update from its deployment TP rank #42346

Motivation

With EAGLE3 and attention data parallelism, a draft worker owns attention and is built inside the draft TP scope, where the context answers its narrowed attention-TP rank. WeightUpdater records that rank at construction. But update_weights_from_tensor receives a LocalSerializedTensor holding one entry per deployment TP rank, so indexing it with the narrowed rank opens another rank's CUDA IPC handle: the draft loads another rank's tensor whenever the entries differ, and creates a context on that rank's device.

Modifications

Index the payload with the deployment TP rank, read at update time. The scheduler runs weight updates outside any draft placement scope, so get_parallel().tp_rank there is the deployment rank.

init_weights_update_group keeps the recorded rank: only the target joins that group.

Accuracy Tests

H200, Qwen3-8B with the AngelSlim/Qwen3-8B_eagle3 draft, --tp-size 2, GSM8K 300 questions:

config	base	this PR
plain	0.953 / 0.947 (two runs)	0.950
--attn-dp-size 2	0.950 / 0.950 (two runs)	0.947

Greedy text (4 prompts × 48 tokens) is byte-identical across all three plain runs and all three attention-DP runs. Accept length and throughput sit inside the base spread.

The premise was confirmed on a live server: under --attn-dp-size 2, world rank 1 reports the draft's recorded rank as 0 while its deployment rank is 1, so the pre-fix index would have opened rank 0's handle. Under plain --tp 2 both are 1, matching the claim that the bug needs attention DP.

test_weight_update_tensor_rank.py fails with the one line reverted and passes with it.

unit/spec and unit/model_executor: same results as the parent commit.

Speed Tests and Profiling

Not applicable.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #37156465078
Latest PR Test (Extra): ❌ Run #37156464945
Latest PR Test (AMD ROCm 10): ❌ Run #37156465035