Skip to content

feat(epp): move conditional-decode gate into prefix-based-pd-decider - #2335

Open
dmitripikus wants to merge 9 commits into
llm-d:mainfrom
dmitripikus:cond-decode-config-new
Open

feat(epp): move conditional-decode gate into prefix-based-pd-decider#2335
dmitripikus wants to merge 9 commits into
llm-d:mainfrom
dmitripikus:cond-decode-config-new

Conversation

@dmitripikus

@dmitripikus dmitripikus commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?

/kind feature

What this PR does / why we need it:

Moves the RFC 7240 Prefer: if-available conditional-decode 412 gate out of the director's hard-coded check and into a configurable plugin behavior on the existing prefix-based-pd-decider.

Today the director hard-codes "reject when the primary decode endpoint has zero matched prefix blocks", which is expressed in tier-weighted blocks and cannot be tuned. Operators have no way to reject speculative decodes based on the actual token count that would need to be prefilled locally.

This PR:

  • Adds PreRequest to PrefixBasedPDDecider. When a request carries Prefer: if-available and the non-cached suffix is at least nonCachedTokens (unweighted, so a RAM-cached prefix contributes its full length), the plugin returns errcommon.Error{Code: PreconditionFailed} — the director's runner preserves the typed code end-to-end so the 412 response is byte-identical to what the director used to emit.
  • Honors promptTokens in the gate the same way disaggregate() does: prompts shorter than promptTokens skip both remote prefill and the 412 gate, so the operator's short-prompt shortcut applies uniformly.
  • Removes primaryEndpointHasCachedPrefix and the inline gate block from pkg/epp/requestcontrol/director.go.
  • Makes the gate opt-in: deployments without the plugin no longer emit 412 and always forward speculative decode requests.
  • nonCachedTokens: 0 disables the gate (consistent with disaggregation being off).

Groundwork was landed separately in #2215 (adds error return to the PreRequest interface).

Which issue(s) this PR fixes:

Fixes #1686

Release note:

The conditional-decode 412 gate (RFC 7240 `Prefer: if-available`) is now enforced by the `prefix-based-pd-decider` plugin instead of a hard-coded director check. Deployments that declare the plugin enforce the gate at the configured `nonCachedTokens` threshold (unweighted cached-block count) and honor its `promptTokens` short-prompt bypass; deployments without the plugin no longer emit 412 and always forward speculative decode requests. To restore the previous "any weighted cache hit lets it through" behavior, set `nonCachedTokens: 1`.

The RFC 7240 Prefer: if-available cache-adequacy gate is expressed via the
existing prefix-based-pd-decider plugin instead of a hard-coded director
check. The gate becomes opt-in (deployments without the plugin forward
speculative decode requests unconditionally) and its threshold is expressed
in non-cached tokens, using the plugin's existing nonCachedTokens knob and
the unweighted cached-block count so RAM-cached prefixes are counted at
their full length.

The director's runner already aggregates PreRequest errors and preserves a
typed errcommon.Error{Code: PreconditionFailed} when it is the sole
failure, so the 412 response the coordinator restart step (pkg/coordinator/
steps/decode_proxy) sees is byte-for-byte identical to what the director
used to emit.

Closes llm-d#1686

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
@dmitripikus
dmitripikus requested a review from a team as a code owner August 9, 2026 11:40
@dmitripikus
dmitripikus requested review from ahg-g and elevran August 9, 2026 11:40
@github-actions github-actions Bot added size/XL Denotes a PR that changes 500-999 lines, ignoring generated files. kind/feature Categorizes issue or PR as related to a new feature. and removed kind/feature Categorizes issue or PR as related to a new feature. labels Aug 9, 2026
Comment thread docs/communication.md Outdated
Comment thread docs/coordinator_architecture.md Outdated
Comment thread docs/disaggregation.md Outdated
Align the PreRequest gate with disaggregate(): a prompt shorter than
promptTokens skips remote prefill in both paths, so the operator's short-
prompt shortcut is not contradicted by a 412 that forces the coordinator
to run remote prefill anyway.

Collapse the 412-gate write-up into a single canonical explanation in
docs/disaggregation.md; communication.md and coordinator_architecture.md
become one-line pointers.

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
@github-actions github-actions Bot added kind/feature Categorizes issue or PR as related to a new feature. and removed kind/feature Categorizes issue or PR as related to a new feature. labels Aug 10, 2026
Extract needsRemotePrefill, the private helper that answers "does this
request's non-cached suffix meet NonCachedTokens?", and rewrite both
disaggregate and PreRequest as thin wrappers over it. The two callers
keep their distinct policies for read-failure — disaggregate logs and
falls back to no-disagg (fail-soft, internal routing); PreRequest maps
it to a 412 (fail-closed, speculative-decode gate).

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
@dmitripikus
dmitripikus requested a review from roytman August 10, 2026 12:08
Comment thread docs/disaggregation.md Outdated
Comment thread docs/disaggregation.md

@roytman roytman left a 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.

There is another inherited issue:
The gate fails open, not closed, when the tokenizer hasn't run — getUserInputLenInTokens returns (0, nil) for a nil TokenizedPrompt (not an error). That makes nonCachedTokens = 0 - hitPrefixTokens, always below threshold, so the request is forwarded — the opposite of the "fails closed" claim in the PreRequest doc comment. This is reachable in production: director.go explicitly swallows DataProducer (tokenizer) failures, so a tokenizer backend hiccup silently disables the whole gate. The new test "no TokenizedPrompt: zero-token prompt trivially covered" (wantReject: false) locks this behavior in rather than catching it.

We should validate if we want to leave this behaviour or now.
The possible fix is:

func getUserInputLenInTokens(request *scheduling.InferenceRequest) (int, error) {
	if request == nil || request.Body == nil {
		return 0, errors.New("request or request body is nil")
	}
	if request.Body.TokenizedPrompt == nil {
		return 0, errors.New("prompt not tokenized")
	}
	return request.Body.TokenizedPrompt.TokenCount(), nil
}

One test needs to flip with it - prefix_based_pd_decider_test.go, "no TokenizedPrompt: zero-token prompt trivially covered":

{
    name:            "no TokenizedPrompt: cache state unknown, fails closed",
    nonCachedTokens: 5,
    headers:         preferIfAvailable,
    result:          resultWithEndpoint(makeTestEndpoint(0)),
    request:         completionsRequestWithPrompt(fwkrh.Prompt{}),
    wantReject:      true,
},

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
disaggregate() and PreRequest() ran the same computation twice per
conditional-decode request against the same decode endpoint. Cache the
outcome on the request via the plugin-local attribute store so PreRequest
reuses the decision disaggregate computed during scheduling.

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
Introduce the plugin's two roles in the section intro and add a
gate-only EndpointPickerConfig example so the standalone usage is
discoverable next to the 412 gate description.

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
A nil TokenizedPrompt on the request read as zero input tokens, which
tripped the short-input shortcut and silently forwarded the request.
Director swallows DataProducer errors, so a tokenizer-backend hiccup
disabled the whole gate in production. Treat missing tokenization as an
error in getUserInputLenInTokens: disaggregate still soft-fails to
decode-only via its existing Error-log branch, and PreRequest returns
412 via its existing typed-error branch.

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
@dmitripikus

Copy link
Copy Markdown
Contributor Author

Regarding the issue in getUserInputLenInTokens(): I fixed it, so that getUserInputLenInTokens now returns an error when TokenizedPrompt == nil. disaggregate still soft-fails to decode-only via its existing Error-log branch, and PreRequest returns 412 via its existing typed-error branch. Test is updated. Thanks!

@dmitripikus
dmitripikus requested a review from roytman August 11, 2026 09:16
The prefix-based-pd-decider memoizes its decision on the request's
attribute store, so tests that reused one *InferenceRequest across
multiple table cases or sequence steps saw the first call's memo on
later calls. Rebuild the request per subtest and per series step to
match production, where each HTTP request is its own *InferenceRequest.

Signed-off-by: Dmitri Pikus <DPIKUS@il.ibm.com>
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/XL Denotes a PR that changes 500-999 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Conditional-decode 412 gate is hard-coded; make it configurable via a plugin

2 participants