fix(r3): balance replay padding across expert ranks

[EazyReal]

Summary

Router replay needs valid expert indices for every model token, including rows without a captured route. SkyRL currently fills every uncaptured or padded route with the same [0, topk) expert IDs. Megatron's dropless MoE dispatcher still sends those tokens through expert permutation and EP all-to-all, so the constant fill concentrates padding work on the first experts and their contiguous EP shard.

This change repairs only masked replay rows after token metadata alignment and TP slicing. It enumerates padding assignments in Megatron's sequence-major model order, then maps each assignment rank-first across EP shards and local experts:

assignment = padding_row_ordinal * topk + column
expert_rank = assignment % expert_parallel_size
local_expert = (assignment // expert_parallel_size) % num_local_experts
expert_id = expert_rank * num_local_experts + local_expert

Every padding row retains distinct, in-range expert IDs. Every assignment prefix is balanced across contiguous EP ownership, while captured rollout routes remain unchanged.

The aligned local-layer replay tensor is widened to int32 before these IDs are written. This is required when observed rollout routes fit in compact uint8 storage but the model has more than 256 experts. The per-layer split also follows Megatron's unpacked [sequence, batch] flattening order; packed replay remains unchanged because its aligned batch dimension is one.

Verification

• 24 passed across the focused replay and token-metadata suites
• Regression coverage for uint8 observed routes with 384 experts and EP=8, including installed expert IDs 288 and 336
• Regression coverage for captured-route preservation, per-row distinctness, expert/EP distribution over one to five padding rows, unpacked sequence-major order, packed B=1 order, and compact/wide dtype equivalence
• Black, Ruff, and git diff --check

No fresh GPU Megatron/EP end-to-end run was performed.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to distribute padding routes across experts in model token order within the replay utility functions, modifying how replay indices are split and handled during tensor parallel (TP) splitting. A review comment highlights a potential compatibility and correctness issue where calling cumsum directly on a boolean tensor can fail or behave unexpectedly depending on the PyTorch version; casting the tensor to torch.long beforehand is recommended.

[gemini-code-assist]

Calling cumsum directly on a boolean tensor (flat_padding_mask) can lead to issues depending on the PyTorch version and backend. In older PyTorch versions or on certain devices, it will raise a RuntimeError (e.g., cumsum_cpu/cumsum_cuda not implemented for 'Bool'). In other versions, it may perform boolean accumulation (logical OR) and return a boolean tensor, which would cause all padding_ordinals after the first True to be 0 (since True - 1 = 0), leading to silent correctness bugs where all padding rows are assigned to the same expert.

To ensure compatibility and correctness across all PyTorch versions and devices, explicitly cast the boolean tensor to torch.long before calling cumsum. This also aligns with the existing pattern in token_metadata.py (line 197) where attention_mask is cast to torch.long before cumsum.

Suggested change

	padding_ordinals = flat_padding_mask.cumsum(0)[flat_padding_mask] - 1
	padding_ordinals = flat_padding_mask.to(torch.long).cumsum(0)[flat_padding_mask] - 1