re-ordered the scheduling stages in the disagg handler - #2340
Conversation
github-actions
Bot
added
kind/feature
ahg-g
changed the title
github-actions
Bot
added
kind/feature
github-actions
Bot
added
size/XL
…iller is picked first, then encoder, then decoder Signed-off-by: ahg-g <ahg@google.com>
Signed-off-by: ahg-g <ahg@google.com>
|
/hold |
|
This PR changes existing semantics, which is not good. To maintain backward compatibility #2342 makes this an opt-in choice. |
elevran
left a comment
elevran
left a comment
There was a problem hiding this comment.
Reordering prefill-before-decode makes sense given #2299 (rack-affinity constraint needs the prefill pick before decode runs). Good catch fixing the nil-endpoint case in PrefixBasedPDDecider in the latest push - that was a real bug in the first version of this PR.
Two things before merge: the release note is marked NONE but this changes default routing/metrics behavior for every disagg deployment; and worth considering whether stage order should be configurable rather than hardcoded, since different deployments may want either (decode vs prefill heavy?).
| // Stages run in order: decode → encode (optional) → prefill (optional). | ||
| // Stages run in order: prefill (optional) → encode (optional) → decode. | ||
| // Returns the next profile to execute, or an empty map when all stages are done. | ||
| func (h *Handler) Pick(ctx context.Context, request *scheduling.InferenceRequest, profiles map[string]scheduling.SchedulerProfile, |
There was a problem hiding this comment.
A few comments
- IIRC, you need to run encoding to get the object's token for prefill stage,
- running prefill before decode is changed user behavior and warrants a releae note
- for backward compat, worth considering making the order configurable via DisaggProfileHandler parameters so that decode first can still be selected (might need to be the default for compatibility or at least add a deprecation note after validating the performance gain from prefill first). The parameter can be simple bool (run prefill first), or an array of steps/profiles to avoid another full reorder in the code next time the trade-off changes.
There was a problem hiding this comment.
I sent out #2342 which adds a parameter and preserves backward compatibility and adds a release note, I will close this PR
There was a problem hiding this comment.
IIRC, you need to run encoding to get the object's token for prefill stage,
This is only changing the picking order within the scheduler, not the execution order.
There was a problem hiding this comment.
what I'm not sure of is if the prefill plugins need token information that is produced as part of the encoding profile's plugins. If that's the case (might not be... could be that they all run earlier to produce request attributes), then encode profile must run before the prefill profile.
It was meant as an open question - I don't know if that's the case.
| Runs each scheduling stage in sequence and assembles the final result from all stages that ran. | ||
|
|
||
| 1. Run the decode profile (always). | ||
| 1. If a prefill decider is configured and approves the request, run the prefill profile. |
There was a problem hiding this comment.
nit: the lead-in sentence above this list (line 27) still says "decode (always), and optionally encode and prefill," but this numbered list now runs prefill -> encode -> decode. Update the lead-in to match.
2 participants
re-ordered the scheduling stages in the disagg handler such that prefiller is picked first, then encoder, then decoder
What type of PR is this?
/kind feature
What this PR does / why we need it:
With #2299 merging, picking the prefill last becomes suboptimal. The priority should be to pick the endpoint with the highest kv-cache match across domains/racks and then constrain the decoder to that rack.
Which issue(s) this PR fixes:
Related to #2315
Release note (write
NONEif no user-facing change):