fix(replay): draw -1 padding fillers from per-row complement set

[DavidBellamy]

Problem

miles/utils/replay_base.py::BaseReplayManager.get_topk_fn._get_replay_result fills -1 slots in top_indices (rollout-engine truncated/aborted token routings) by walking a flat arange mod num_experts:

padding_mask = top_indices == -1
if padding_mask.any():
    top_indices[padding_mask] = (
        torch.arange(padding_mask.sum(), device=top_indices.device, dtype=top_indices.dtype)
        % scores.shape[1]
    )

This ignores the [num_tokens, topk] row structure. Within a single token's row, the cyclic filler can land on an expert id already present in that row's existing topk picks, producing a within-row duplicate.

Downstream the router builds routing_map via scatter(top_indices, True) into a [num_tokens, num_experts] boolean matrix. Duplicates within a row collapse silently. As a result routing_map.sum() < num_tokens * topk. The MoE token dispatcher then computes input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(axis=1) — i.e. sum(input_splits) = routing_map.sum() < num_tokens * topk. But num_out_tokens for the permuted buffer is routing_map.size(0) * topk (static, dropless path). At dispatch time, all_to_all_single is called with permuted_tokens.shape[0] = num_tokens * topk and input_splits summing to less, so it raises:

RuntimeError: Split sizes doesn't match total dim 0 size

This is the canonical R3 crash described in radixark#1002 — and reproduced repeatedly by Phase-3 RL360 jobs.

Why it shows up intermittently

The probability of collision per row depends on:

1. The truncation pattern at the rollout engine (which rows have any -1s, and how many).
2. The topk / num_experts ratio. For GLM-4.7-Flash (topk=4, num_experts=64) collisions are relatively rare; for higher topk:experts ratios (or edge configs) they are common.

This explains why the same code path can run many MoE forwards cleanly on one config and crash on another.

Fix

Replace the cyclic filler with a per-row complement-set draw: for each row, fill its -1 slots with the highest-scoring experts NOT already used in that row's existing topk picks. By construction, no within-row duplicate is ever produced.

n_experts = scores.shape[1]
non_pad = ~padding_mask
# [n_tokens, n_experts] used-mask, OR-aggregated from non-pad picks
used_mask = (
    F.one_hot(top_indices.clamp(min=0).long(), num_classes=n_experts)
    .mul(non_pad.long().unsqueeze(-1))
    .sum(dim=1)
    .bool()
)
# rank free experts per row by score; used experts sink to the tail via -inf
masked_scores = scores.masked_fill(used_mask, float("-inf"))
_, sorted_free = masked_scores.sort(dim=1, descending=True)
# k-th -1 in each row gets sorted_free[row, k]
pad_cumsum = torch.cumsum(padding_mask.long(), dim=1) - 1
fill_values = torch.gather(sorted_free, 1, pad_cumsum.clamp(min=0)).to(top_indices.dtype)
top_indices = torch.where(padding_mask, fill_values, top_indices)

Notes:

• top_indices is int32 per Sample.rollout_routed_experts; F.one_hot requires int64, hence .long().
• pad_cumsum.clamp(min=0) is safe — for non-pad positions torch.where masks the gathered value.
• Score-rank ordering picks "most plausible" fillers (the experts the model was next-most-likely to route to anyway), keeping replay close to a counterfactual on-policy roll.

Validation

End-to-end iter-scale RL run on M2 (job 1581417, GLM-4.7-Flash, topk=4, num_experts=64, expert-model-parallel-size=8, real terminus-2 agent, R3 on):

Signal	Result
dup_rows_after_pad (any > 0?)	0 across 620 R3_DIAG prints
Dispatcher gap = num_tokens*topk − routing_map.sum()	0 across all dispatcher prints
Split sizes doesn't match total dim 0 size	0
Timer train_step cycles completed	6 in 1h6min
Step-record log lines	144

The job eventually exited from an unrelated NCCL ALLREDUCE timeout in log_rollout_data → gather_log_data (wandb gather), well after R3 had repeatedly run cleanly through the dispatch path.

The diagnostic that confirmed root cause is preserved as an R3_DIAG env-gated print in the same function on a separate working branch; it was not included here to keep this PR scoped to the fix.

Closes radixark#1002.