Skip to content

feat(pd): stitch routed_expert_weights payload segments like routed_experts - #62

Draft
faresobeid wants to merge 4 commits into
feat/pd-merge-payload-segmentsfrom
feat/pd-merge-payload-fields
Draft

faresobeid wants to merge 4 commits into
feat/pd-merge-payload-segmentsfrom
feat/pd-merge-payload-fields

Conversation

@faresobeid

@faresobeid faresobeid commented Oct 4, 2026 •

Copy link
Copy Markdown

Stacked on #60. Draft, parked. No prime-rl producer yet: a future TRR fields PR would write routed_expert_weights segments (PrimeIntellect-ai/prime-rl#3904 is closed).

STITCHED_FIELDS = ["routed_experts", "routed_expert_weights"] lists the payload fields that the forward pass produces for every position. Each is stitched separately: the prefill segments come first, then the decode segments clipped at that field's own prefill coverage end. Row bytes are dtype-aware: uint8/uint16 ids, float32 weights (float32 is added to dtype_itemsize here). Every other field (sampling_mask, sampling_mask_logprobs) comes from decode only and passes through. A decode stitched field with no prefill counterpart is an error, as before.

Checks: cargo fmt --check, cargo clippy --lib --tests -D warnings, cargo test --lib. I extended the existing merge test with weights and logprobs.

Note

Stitch routed_expert_weights segments alongside routed_experts in HTTP router

  • Adds a STITCHED_FIELDS set with routed_experts and routed_expert_weights in routed_experts_merge.rs, so payload segments for both fields are collected and stitched
  • merge_routed_payload_in_json now tracks prefill coverage per field in a HashMap and clips each decode segment against its own field's prefill endpoint. Non-stitched fields like sampling masks pass through unchanged
  • A decode segment for a stitched field with no matching prefill coverage now returns an error naming that field
  • Behavioral Change: decode-only routed_expert_weights payloads that previously passed through now fail with an error instead of being silently unmerged

Macroscope summarized 4c87216. (Automatic summaries will resume when PR exits draft mode or review begins).

…perts

STITCHED_FIELDS lists the by-handle fields the forward pass produces for every
position (routed_experts and Total Router Recall's routed_expert_weights). Each
is stitched from prefill + decode with its own prefill cut; sampling masks and
their logprobs stay decode-only and pass through.
@faresobeid
faresobeid force-pushed the feat/pd-merge-payload-fields branch from 4c87216 to 0dc7ac3 Compare October 10, 2026 19:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant