[Rust Frontend][RL] add routed expert prompt offset (#52703)
Signed-off-by: Biswa Panda <biswa.panda@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Biswa Panda committed
Aug 18, 2026 at 13:04 UTC
6a391a931a92250aa45aa711aaa3f099f5a8d05f
4 files changed
+15
rust/src/engine-core-client/src/protocol/sampling.rs
+4
@@ -145,6 +145,9 @@ pub struct EngineCoreSamplingParams {
145
pub skip_reading_prefix_cache: Option<bool>,
146
/// Additional request parameters for custom extensions (from `vllm_xargs`).
147
pub extra_args: Option<HashMap<String, serde_json::Value>>,
148
+ /// Number of prompt tokens to skip from returned routed-expert data.
149
+ /// A value of zero returns routing data for the entire prompt.
150
+ pub routed_experts_prompt_start: u32,
151
}
152
153
impl EngineCoreSamplingParams {
@@ -175,6 +178,7 @@ impl EngineCoreSamplingParams {
178
logprob_token_ids: None,
179
skip_reading_prefix_cache: None,
180
extra_args: None,
181
+ routed_experts_prompt_start: 0,
182
}
183
}
184
}
rust/src/engine-core-client/src/tests/client.rs
+2
@@ -159,6 +159,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
159
stop_token_ids: vec![151643],
160
eos_token_id: Some(151645),
161
all_stop_token_ids: BTreeSet::from([151643, 151645]),
162
+ routed_experts_prompt_start: 1,
163
..EngineCoreSamplingParams::for_test()
164
}),
165
arrival_time: 42.5,
@@ -2680,6 +2681,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2681
logprob_token_ids: None,
2682
skip_reading_prefix_cache: None,
2683
extra_args: None,
2684
+ routed_experts_prompt_start: 0,
2685
},
2686
);
2687
rust/src/engine-core-client/src/tests/python_compat.py
+2
@@ -47,6 +47,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
47
stop_token_ids: list[int] = []
48
_eos_token_id: int | None = None
49
_all_stop_token_ids: set[int] = set()
50
+ routed_experts_prompt_start: int = 0
51
output_kind: RequestOutputKind = RequestOutputKind.DELTA
52
53
@@ -135,6 +136,7 @@ request = EngineCoreRequest(
136
stop_token_ids=[151643],
137
_eos_token_id=151645,
138
_all_stop_token_ids={151643, 151645},
139
+ routed_experts_prompt_start=1,
140
output_kind=RequestOutputKind.FINAL_ONLY,
141
),
142
pooling_params=None,
rust/src/text/src/lower.rs
+7
@@ -188,6 +188,7 @@ pub fn lower_sampling_params(
188
logprob_token_ids,
189
skip_reading_prefix_cache,
190
extra_args: vllm_xargs,
191
+ routed_experts_prompt_start: 0,
192
};
193
validate_resolved_sampling_params(¶ms)?;
194
validate_vocab_range(¶ms, &sampling_limits)?;
@@ -645,6 +646,7 @@ mod tests {
646
logprob_token_ids: None,
647
skip_reading_prefix_cache: None,
648
extra_args: None,
649
+ routed_experts_prompt_start: 0,
650
}
651
"#]]
652
.assert_debug_eq(¶ms);
@@ -694,6 +696,7 @@ mod tests {
696
logprob_token_ids: None,
697
skip_reading_prefix_cache: None,
698
extra_args: None,
699
+ routed_experts_prompt_start: 0,
700
}
701
"#]]
702
.assert_debug_eq(¶ms);
@@ -859,6 +862,7 @@ mod tests {
862
logprob_token_ids: None,
863
skip_reading_prefix_cache: None,
864
extra_args: None,
865
+ routed_experts_prompt_start: 0,
866
}
867
"#]]
868
.assert_debug_eq(¶ms);
@@ -926,6 +930,7 @@ mod tests {
930
logprob_token_ids: None,
931
skip_reading_prefix_cache: None,
932
extra_args: None,
933
+ routed_experts_prompt_start: 0,
934
}
935
"#]]
936
.assert_debug_eq(¶ms);
@@ -986,6 +991,7 @@ mod tests {
991
logprob_token_ids: None,
992
skip_reading_prefix_cache: None,
993
extra_args: None,
994
+ routed_experts_prompt_start: 0,
995
}
996
"#]]
997
.assert_debug_eq(¶ms);
@@ -1235,6 +1241,7 @@ mod tests {
1241
logprob_token_ids: None,
1242
skip_reading_prefix_cache: None,
1243
extra_args: None,
1244
+ routed_experts_prompt_start: 0,
1245
}
1246
"#]]
1247
.assert_debug_eq(¶ms);