Skip to content

[Flow Control] Stage-Aware Saturation Gating for disaggregated P/D deployments - #2221

Open
loicmarchal wants to merge 4 commits into
llm-d:mainfrom
loicmarchal:feature/pd-stage-aware-saturation
Open

[Flow Control] Stage-Aware Saturation Gating for disaggregated P/D deployments#2221
loicmarchal wants to merge 4 commits into
llm-d:mainfrom
loicmarchal:feature/pd-stage-aware-saturation

Conversation

@loicmarchal

@loicmarchal loicmarchal commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?

/kind feature

What this PR does / why we need it:

The flow control system currently computes a single saturation value across all endpoints. In disaggregated prefill/decode (P/D) deployments, this masks stage-specific bottlenecks: a prefill surge can overwhelm
prefill workers while the global average stays low because idle decode workers dilute the signal, so backpressure never kicks in when it should.

This PR makes the dispatch cycle stage-aware by:

  1. Adding a partitionEndpoints helper that classifies endpoints into prefill, decode, and interleaved buckets using the llm-d.ai/role pod label (reusing role constants from bylabel/roles.go). Endpoints without a
    role label, without metadata, or with an unrecognized role value default to the decode bucket, matching the existing NewDecodeRole(allowsNoLabel: true) convention for monolithic deployment safety.
  2. Computing saturation independently per non-empty stage, then using max(S_prefill, S_decode, S_interleaved) as the global saturation fed into priority ceilings and HoL blocking. Empty partitions are skipped
    (not treated as saturated). When all partitions are empty, the detector is called on the full pool to preserve existing empty-pool behavior.
  3. Adding a "stage" label to both flow_control_pool_saturation Prometheus gauges, emitting per-stage ("prefill", "decode", "interleaved") and "global" values each dispatch cycle.

The existing gating logic (priority ceilings, HoL blocking) is unchanged -- it consumes the effective saturation value transparently. Monolithic deployments (no role labels) behave identically to before: all endpoints land in the decode bucket, producing a single-stage saturation equal to the previous flat-pool value

Which issue(s) this PR fixes:

Fixes #1186

Release note:

Flow control saturation gating is now stage-aware in disaggregated prefill/decode deployments. Endpoints are partitioned by the `llm-d.ai/role` pod label, saturation is computed per stage, and backpressure triggers on the worst stage (`max(S_prefill, S_decode, S_interleaved)`). A new `stage` label on the `flow_control_pool_saturation` metric enables per-stage observability. Monolithic deployments are unaffected.

@loicmarchal
loicmarchal requested review from a team, LukeAVanDrie and shmuelk as code owners July 29, 2026 03:59
@loicmarchal
loicmarchal requested review from ahg-g and elevran July 29, 2026 03:59
@github-actions github-actions Bot added kind/feature Categorizes issue or PR as related to a new feature. size/L Denotes a PR that changes 100-499 lines, ignoring generated files. and removed kind/feature Categorizes issue or PR as related to a new feature. labels Jul 29, 2026
@loicmarchal
loicmarchal force-pushed the feature/pd-stage-aware-saturation branch 2 times, most recently from fa4b5e6 to d866463 Compare July 29, 2026 19:02
}
stageSat := p.saturationDetector.Saturation(ctx, part.endpoints)
metrics.RecordFlowControlPoolSaturation(p.poolName, part.name, stageSat)
if stageSat > saturation {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is currently implementing the simple approach to calculate saturation = max(S_prefill, S_decode, S_interleaved).

I'm looking into a solution to count inflight prefill request into the decode saturation to correctly count for projected decode load and not gate too late (as discussed in #1186 (comment)).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we land simple first and defer this. I am also interested in investigating some approaches here if you're open to collaborating.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LukeAVanDrie so you suggest we go with this solution first and look into the projected decode saturation as a followup PR? That's ok with me. Sure, of course I'll be happy to collaborate on that one.

Signed-off-by: Loic Marchal <lmarchal@redhat.com>
Signed-off-by: Loic Marchal <lmarchal@redhat.com>
Signed-off-by: Loic Marchal <lmarchal@redhat.com>
Signed-off-by: Loic Marchal <lmarchal@redhat.com>
@loicmarchal
loicmarchal force-pushed the feature/pd-stage-aware-saturation branch from d866463 to 4b3faf7 Compare July 30, 2026 03:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/feature Categorizes issue or PR as related to a new feature. size/L Denotes a PR that changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Flow Control / Disaggregation] Stage-Aware Saturation Gating for Prefill/Decode Pools

2 participants