[faresobeid]

Purpose

P/D router replay fails every request on models with more than 256 experts (Qwen3.5-397B, Kimi-K2, Nemotron Super/Ultra). Since PrimeIntellect-ai/prime-rl#3965, prime-rl sends vLLM's per-model routing dtype, which is uint16 for those models. The inline P/D merge assumed 1 byte per id and dropped dtype on re-encode.

This PR also changes how the merge cuts the decode rows, so prime-rl no longer has to zero-pad decode responses to the prompt start (PrimeIntellect-ai/prime-rl#3830).

Changes (routed_experts_merge.rs)

• dtype is required (uint8 or uint16). Row bytes = itemsize(dtype) * layers * topk. dtype is kept on re-encode.
• Prefill and decode must have the same dtype (always true since prime-rl#3965); otherwise the merge errors.
• Decode rows are cut by position: decode.suffix_rows((prompt.start + prompt.seq_len).saturating_sub(decode.start)). Prefill rows are followed by the decode rows from the prefill end on; a gap is an error.
  • A decode response whose rows start at its first forward (prompt_len - 1) works without padding.
  • A decode response padded to the prompt start (current prime-rl) gives the same cut as before.

Release

Ships as router v0.2.3. The version bump is on branch release/v0.2.3 (this PR + release: v0.2.3). After this PR merges, rebase that branch onto main and open the release/v0.2.3 PR; merging it tags v0.2.3 (.github/workflows/release.yml). prime-rl then bumps its router pin.

Test Plan

• Unit tests: same-dtype merge of padded decode rows; uint16 merge of decode rows that start at the first forward.
• cargo test --lib routed_experts, cargo fmt --check.

GPU validation with router v0.2.3 (2026-10-14)

1x 8xH200, vLLM 0.31.0, Qwen3-30B-A3B TP1, VLLM_BATCH_INVARIANT=1, router replay + sampling masks. prime-rl = PrimeIntellect-ai/prime-rl#3830 (978146cf5) merged with PrimeIntellect-ai/prime-rl#3900 (20a322525); router = release/v0.2.3 (ff818da, this PR + the version bump), wheel built locally. 1P1D (NIXL) through the router vs a colocated server. Identical means: same tokens, same routing rows for prompt and completion (100%, rows = prompt + completion - 1, merged start = 0), same sampling masks, same logprobs.

• Concurrency 1 (prompts 64-6000 tokens, 128 completion tokens): 10/10 identical.
• Concurrency 10, fresh cache_salt on the colocated side: 20/20 identical.
• uint16: no knob exists (vLLM picks uint8 for <=256 experts), so a test-only sitecustomize forced vLLM's capturer to uint16 on P, D and the colocated server. The responses carry dtype: uint16. Concurrency 1: 10/10 identical; concurrency 10: 20/20 identical.
• Forced decode preemption (D with num_gpu_blocks_override=1300, 1000-4000 token prompts, up to 1024 tokens, concurrency 10): 6 preemptions on D, 20/20 identical.
• Router logs: no errors or warnings during the runs (only the shutdown warning). Engines: all Using V2 Model Runner, no BlockObjectStoreError.
• Negative control, router v0.2.2 with the same P/D pair: 8/10 requests fail (HTTP 500 decode routed_experts has 128 rows, expected at least <prompt_len>). The 2 prompts of 64 tokens return 128 rows instead of 191, and the completion rows are wrong.

Note

Honor routed_experts dtype in inline merge row slicing and concat

• Adds a required dtype field (uint8 or uint16) to routed-experts payloads, used to compute row byte widths in row_bytes, suffix_rows, and concat_rows in routed_experts_merge.rs
• Merge handler now removes only decode rows that overlap the prefill range, instead of a fixed prefill sequence length, so payloads with nonzero row start merge correctly
• Adds checked arithmetic for row ranges and shape byte sizes; overflow returns an error instead of a wrong result
• Risk: decoding now rejects routed-experts payloads without a valid dtype or with a data length not matching itemsize × shape, so existing uint8-only payloads must include dtype

^{Macroscope summarized 09aebbc.}

---

Note

Medium Risk
Changes P/D response stitching for routed_experts with stricter dtype and row-range validation; wrong overlap logic would corrupt routing metadata for replay.

Overview
Fixes inline routed_experts merge for P/D replay when expert IDs use uint16 (models with >256 experts). Payloads now require a dtype (uint8 or uint16); byte length checks and row slicing use itemsize × layers × topk, and merged JSON preserves dtype.

Merge no longer always drops the first prefill.seq_len decode rows. It strips only rows already covered by prefill using prompt.end() - decode.start, and concat_rows errors on dtype mismatch or a gap between prefill end and decode start (so decode can start at first-forward without zero-padding).

Tests cover padded uint8 merge, uint16 merge with non-zero start, and overflow/invalid metadata rejection.

^{Reviewed by Cursor Bugbot for commit 09aebbc. Bugbot is set up for automated code reviews on this repo. Configure here.}

[macroscopeapp]

Approvability

Verdict: Approved at 09aebbc

Macroscope's review found this PR approvable — This is a focused, one-file fix to the existing routed-experts P/D merge, adding dtype-aware byte handling and position-aware row stitching with validation and targeted tests. Its effects remain contained to the established merge path, with no new workflow or infrastructure introduced.

^{You can add or adjust custom eligibility rules. Learn more.}