feat: support router replay r2

[SolenoidWGT]

PR README: feat: support router replay r2

Summary

Adds router replay r2 support to SamplingClient.compute_logprobs, allowing MoE (Mixture of Experts) models to retrieve forward expert routing data when computing logprobs for use in subsequent forward_backward training.

Change Summary

Item	Details
Commit	de974c56a799b6f5e3ea27a8fea52ab42b44a705
Files Changed	weaver/sampling_client.py
Diff	+28 lines, -3 lines

Main Changes

1. New parameter router_replay_r2

   • Adds optional parameter router_replay_r2: bool = False to compute_logprobs
   • When True, the request body includes router_replay_r2: True so the server returns expert routing data

2. Return value behavior

   • router_replay_r2=False (default): Returns List[float | None] (per-token logprobs), same as before
   • router_replay_r2=True: Returns a Dict with:
     • "logprobs": per-token logprobs
     • "fwd_expert_replay_data": forward expert routing data (or None for non-MoE or when unavailable)

3. Purpose

   • fwd_expert_replay_data supports MoE router replay by reusing forward expert selections during backward, keeping training consistent

Usage Example

# Default: only logprobs
logprobs = client.compute_logprobs(prompt=prompt)

# With router replay r2: logprobs and expert routing data
result = client.compute_logprobs(prompt=prompt, router_replay_r2=True)
logprobs = result["logprobs"]
fwd_expert_replay_data = result["fwd_expert_replay_data"]  # for forward_backward

Compatibility

• Default router_replay_r2=False preserves previous behavior and remains backward compatible
• Return type and request content change only when router_replay_r2=True is passed

Testing Checklist

• With router_replay_r2=False, return type and content match the previous implementation
• With router_replay_r2=True, returned dict includes logprobs and fwd_expert_replay_data
• On MoE models, fwd_expert_replay_data works correctly with forward_backward

[void-main]

LGTM