feat(epp): move conditional-decode gate into prefix-based-pd-decider - #2335
feat(epp): move conditional-decode gate into prefix-based-pd-decider#2335dmitripikus wants to merge 9 commits into
Conversation
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>
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>
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>
roytman
left a comment
There was a problem hiding this comment.
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>
|
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! |
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>
What type of PR is this?
/kind feature
What this PR does / why we need it:
Moves the RFC 7240
Prefer: if-availableconditional-decode 412 gate out of the director's hard-coded check and into a configurable plugin behavior on the existingprefix-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:
PreRequesttoPrefixBasedPDDecider. When a request carriesPrefer: if-availableand the non-cached suffix is at leastnonCachedTokens(unweighted, so a RAM-cached prefix contributes its full length), the plugin returnserrcommon.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.promptTokensin the gate the same waydisaggregate()does: prompts shorter thanpromptTokensskip both remote prefill and the 412 gate, so the operator's short-prompt shortcut applies uniformly.primaryEndpointHasCachedPrefixand the inline gate block frompkg/epp/requestcontrol/director.go.nonCachedTokens: 0disables the gate (consistent with disaggregation being off).Groundwork was landed separately in #2215 (adds
errorreturn to thePreRequestinterface).Which issue(s) this PR fixes:
Fixes #1686
Release note: