Fix weight sync selector for frozen speculative drafts

[XinyuJiangCMU]

Co-authored-with: @JessicaJiang-123

The problem

In raw conversion mode, a trainer can enable speculative decoding without training MTP layers. In that configuration, the target model changes during training, but the speculative draft is frozen and should not receive weight updates.

Before this change, Miles did not tell SGLang which model runners a weight update covered. Both the update session and the weight payloads therefore used SGLang's default selector, "all", which includes the target and the draft.

That default is wrong for a frozen draft. On every training step, SGLang still includes the draft in the weight update. The draft loads none of the incoming tensors because its loader filters out every name, but SGLang still runs the post-update processing on it.

The failure is silent. The job does not crash, no assertion fires, and the training loss can still look normal because the target model is updated correctly. The visible symptom is on the rollout side: spec_accept_length drops from about 2.71 to 1.00, so speculative decoding stops accepting useful draft tokens.

How the draft gets corrupted

Before this change, every weight update also processed the frozen draft, even though the draft did not receive any new weights.

The target receives fresh, unshuffled weights, so applying the AITER shuffle is correct. The draft keeps its old weights, which have already been shuffled.

SGLang therefore applies the shuffle to the draft a second time. Since the shuffle is not idempotent, this produces a different weight permutation, and the inference kernel then reads the weights in the wrong layout. As a result, the draft generates incorrect predictions without raising an error.

A minimal reproduction is shown below:

from aiter.ops.shuffle import shuffle_weight
import torch

w = torch.arange(16 * 32, dtype=torch.float32, device="cuda").reshape(16, 32)
s1 = shuffle_weight(w, (16, 16))
s2 = shuffle_weight(s1, (16, 16))

row 0, first 8 cells
w                   :  0  1  2  3    4    5    6    7
shuffle(w)          :  0  1  2  3   32   33   34   35
shuffle(shuffle(w)) :  0  1  2  3  256  257  258  259

Changes

• Add weight_update_selector(args) as the single place that decides which SGLang runners a weight sync covers.
• Return "target" only when the trainer can prove that the draft is frozen:
  • speculative decoding is enabled;
  • mtp_num_layers is missing or zero;
  • the conversion path is not bridge mode.
• Pass the same selector to both begin_weight_update and all weight payloads, so the update session and the actual weight loading cover the same runners.

Verification

The end-to-end test used 4 nodes, DeepSeek-V4-Flash FP8 on gfx950, colocated rollout, and a 16k response length.

The two weight-sync runs used the same configuration and differed only by this fix. A run with the draft kept frozen was used as the reference.

Configuration	spec_accept_length
Frozen-draft reference	2.7056
Weight sync without this fix	1.00
Weight sync with this fix	2.7490

Without the selector, the update session corrupts the draft and accepted length collapses to 1.00. With the selector, accepted length returns to the frozen-draft reference range.

The fixed run also recovers the expected rollout benefit: rollout time is about 320 seconds with MTP off and about 210 seconds with MTP on, roughly 34 percent lower.

[gemini-code-assist]

Caution

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

[guapisolo]

LGTM