[Session R3] Add routed_experts_start_len for absolute routing slice control

[ByronHsu]

Motivation

In multi-turn RL rollouts with MoE models (e.g. Kimi-K2, Qwen3-30B-A3B), return_routed_experts returns the full conversation-length routing data on every request — including the prefix-cached range. As conversations grow, the host gather + ZMQ payload scales as O(seqlen), producing ~1s ITL spikes on long-context requests with high cache-hit ratios. This compounds with DP attention: since all DP ranks synchronize on every decode step, a single rank stalled on a long routed-experts gather blocks the entire batch across all ranks. This makes it the dominant decode bottleneck in our production RL training loop.

The problem

Turn 1:  prompt=[P1]  → generate → routings for [P1 + O1]        # 100 rows, OK
Turn 2:  prompt=[P1+O1+P2] → generate → routings for [P1+O1+P2+O2]  # 500 rows, but only ~100 are new
Turn N:  prompt=[...8K tokens...] → generate → routings for [8K+ON]   # 8K rows, ~100 new — 1s stall

The caller (RL rollout client) only needs the new routing rows, not the full prefix it already accumulated from prior turns.

The solution

Add routed_experts_start_len: Optional[int] that lets callers specify an absolute start position. The server returns routings covering [start_len, seqlen - 1) instead of the full [0, seqlen - 1). The caller sets start_len = len(accumulated_prompt) and gets back only the output-token routings.

Turn 2 with start_len=len(P1+O1):
  → routings for [P2+O2] only  # ~100 rows instead of 500

Experiment results

Kimi-K2-Instruct, pure TP=8 (8× H200), --enable-return-routed-experts --load-format dummy --moe-runner-backend triton --disable-piecewise-cuda-graph --cuda-graph-max-bs 4. 95% prefix-cache hit, 100 output tokens. Timer placed after GPU sync.

Headline numbers

 in_len   cached     full_ms  sliced_ms   speedup
   2048     2047        7.78      0.65    11.96×
   4096     4095       11.37      0.73    15.59×
   8192     8191       14.16      0.86    16.54×
  16384    16383       49.80      1.17    42.46×
  32768    32767      121.23      1.90    63.77×

Per-component breakdown

in_len	mode	elapsed	routed_experts	release_kv	other (stream_output etc)
2048	full	7.78 ms	2.19 ms	0.20 ms	5.39 ms
2048	sliced	0.65 ms	0.10 ms	0.16 ms	0.39 ms
8192	full	14.16 ms	2.98 ms	0.40 ms	10.78 ms
8192	sliced	0.86 ms	0.09 ms	0.39 ms	0.38 ms
32768	full	121.23 ms	30.12 ms	1.51 ms	89.60 ms
32768	sliced	1.90 ms	0.10 ms	1.43 ms	0.37 ms

Key findings:

• routed_experts host gather (aten::index): grows linearly in full mode (2.2 ms at 2k → 30 ms at 32k); flat ~0.1 ms with start_len at all sizes. This PR eliminates this cost.
• stream_output (ZMQ/pickle serialization): 5 ms → 90 ms in full mode (proportional to the 60 MB payload); flat ~0.4 ms with start_len (tiny payload).
• release_kv: grows linearly with seqlen in both modes (0.16 ms → 1.43 ms). This is the radix tree walk — unrelated to routed experts and the sole source of small residual growth.

Changes

Request lifecycle plumbing — routed_experts_start_len: Optional[int] = None added to:

• GenerateReqInput, TokenizedGenerateReqInput (io_struct.py)
• CompletionRequest, ChatCompletionRequest (protocol.py)
• Engine.generate, Engine.async_generate (engine.py)
• Req (schedule_batch.py)
• Pass-through in serving_chat, serving_completions, tokenizer_manager, scheduler, session_controller, encode_receiver

Server-side logic:

• BaseTopkCapturer.get_topk() gains a start_len parameter with defensive clamping
• maybe_collect_routed_experts() honors req.routed_experts_start_len, early-returns when return_routed_experts is False, and logs row-count mismatches as soft warnings
• Scheduler validates start_len <= prompt_tokens and aborts with a clear error otherwise

Tests:

• TestRoutedExpertsStartLen: 4 test cases covering default behavior, row-count correctness, bounds checking (abort on too-large start_len), and cache-hit interaction (radix prefix extends past start_len)

Test plan

• Existing TestReturnRoutedExperts passes (no regression to full-sequence return)
• New TestRoutedExpertsStartLen passes with TP=2 on Qwen3-30B-A3B
• start_len=None (default) produces identical output to omitting the field
• start_len=N returns exactly seqlen - 1 - N rows matching the tail of full return
• start_len > prompt_tokens aborts with clear error message
• Cache-hit case: radix prefix allowed to extend past start_len

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ByronHsu]

cc @zyzshishui to take over

[ispobock]

/tag-and-rerun-ci

[zyzshishui]

Maybe add if recv_req.return_routed_experts and

[Qiaolin-Yu]

better to add some doc

[zyzshishui]

better to add some doc

added, plz check again