From dbfce5b45e79eeca2c21d24fcb2fc77e6db7ae41 Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 14:10:18 +0300 Subject: [PATCH 1/8] Delegate conditional-decode 412 decision to prefix-based-pd-decider Signed-off-by: Dmitri Pikus --- .../interface/requestcontrol/plugins.go | 9 ++ .../disagg/prefix_based_pd_decider.go | 89 ++++++++---- .../disagg/prefix_based_pd_decider_test.go | 70 ++++++++++ pkg/epp/requestcontrol/director.go | 61 +++------ pkg/epp/requestcontrol/director_test.go | 127 ++++++------------ .../requestcontrol/request_control_config.go | 30 +++++ 6 files changed, 238 insertions(+), 148 deletions(-) diff --git a/pkg/epp/framework/interface/requestcontrol/plugins.go b/pkg/epp/framework/interface/requestcontrol/plugins.go index 9f85c85a5f..a0ec88a1aa 100644 --- a/pkg/epp/framework/interface/requestcontrol/plugins.go +++ b/pkg/epp/framework/interface/requestcontrol/plugins.go @@ -98,3 +98,12 @@ type PreAdmitter interface { plugin.Plugin PreAdmit(ctx context.Context, request *fwksched.InferenceRequest) error } + +// ConditionalDecodeDecider answers the RFC 7240 "Prefer: if-available" gate: +// given the chosen decode endpoint, should the request be rejected with HTTP +// 412 (because the local KV cache does not cover enough of the prompt) or +// forwarded? At most one such plugin is consulted per director. +type ConditionalDecodeDecider interface { + plugin.Plugin + ShouldRejectConditionalDecode(ctx context.Context, request *fwksched.InferenceRequest, endpoint fwksched.Endpoint) bool +} diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go index 78aa904d62..8d2ebe5c77 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider.go @@ -10,6 +10,7 @@ import ( "github.com/llm-d/llm-d-router/pkg/common/observability/logging" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" + fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" ) @@ -33,8 +34,11 @@ func (p PrefixBasedPDDeciderConfig) validate() error { return nil } -// compile-time type assertion -var _ deciderPlugin = &PrefixBasedPDDecider{} +// compile-time type assertions +var ( + _ deciderPlugin = &PrefixBasedPDDecider{} + _ fwkrc.ConditionalDecodeDecider = &PrefixBasedPDDecider{} +) // PrefixBasedPDDecider is a PD decider plugin which decision is based prefix aware type PrefixBasedPDDecider struct { @@ -93,7 +97,6 @@ func (d *PrefixBasedPDDecider) WithName(name string) *PrefixBasedPDDecider { } func (d *PrefixBasedPDDecider) disaggregate(ctx context.Context, request *scheduling.InferenceRequest, endpoint scheduling.Endpoint) bool { - logger := log.FromContext(ctx) debugLogger := log.FromContext(ctx).V(logging.DEBUG) // NonCachedTokens defines the minimum number of non-cached tokens required @@ -101,50 +104,84 @@ func (d *PrefixBasedPDDecider) disaggregate(ctx context.Context, request *schedu if d.config.NonCachedTokens == 0 { return false } + nonCachedTokens, ok := d.computeNonCachedTokens(ctx, request, endpoint) + if !ok { + return false + } + if nonCachedTokens < d.config.NonCachedTokens { + debugLogger.Info("Non-cached suffix is smaller than threshold, using decode profile only") + return false // do not run prefill + } + return true +} + +// ShouldRejectConditionalDecode reports whether a conditional-decode request +// (RFC 7240 "Prefer: if-available") should be rejected with HTTP 412 because +// the chosen decode endpoint's KV cache does not cover enough of the prompt. +// +// Returns true (reject) when the non-cached suffix meets or exceeds +// NonCachedTokens, or when the prefix cache state cannot be read from the +// endpoint. Returns false when NonCachedTokens is 0 (gate disabled) or the +// non-cached suffix is below the threshold. +func (d *PrefixBasedPDDecider) ShouldRejectConditionalDecode(ctx context.Context, request *scheduling.InferenceRequest, endpoint scheduling.Endpoint) bool { + debugLogger := log.FromContext(ctx).V(logging.DEBUG) + if d.config.NonCachedTokens == 0 { + return false + } + nonCachedTokens, ok := d.computeNonCachedTokens(ctx, request, endpoint) + if !ok { + // Cannot read prefix cache state - fail closed (reject) to preserve + // the current 412 behavior when no prefix-cache producer is wired up. + return true + } + if nonCachedTokens < d.config.NonCachedTokens { + debugLogger.Info("conditional-decode: non-cached suffix below threshold, forwarding", + "nonCachedTokens", nonCachedTokens, "threshold", d.config.NonCachedTokens) + return false + } + debugLogger.Info("conditional-decode: non-cached suffix at or above threshold, rejecting", + "nonCachedTokens", nonCachedTokens, "threshold", d.config.NonCachedTokens) + return true +} + +// computeNonCachedTokens returns the length of the non-cached prompt suffix in +// tokens for the given endpoint. ok is false when the input length or the +// endpoint's PrefixCacheMatchInfo attribute could not be read. +// +// Uses the unweighted cached-block count, not the tier-weighted match score: +// a RAM-cached prefix must contribute its full token count, otherwise the +// non-cached suffix is overestimated and requests with large local-RAM hits +// are misrouted to remote prefill. +func (d *PrefixBasedPDDecider) computeNonCachedTokens(ctx context.Context, request *scheduling.InferenceRequest, endpoint scheduling.Endpoint) (int, bool) { + logger := log.FromContext(ctx) + debugLogger := logger.V(logging.DEBUG) + if endpoint == nil { logger.Error(nil, "prefix decider: endpoint is nil") - return false + return 0, false } inputTokens, err := getUserInputLenInTokens(request) if err != nil { logger.Error(err, "prefix decider: failed to get user input length in tokens") - return false - } - if inputTokens < d.config.NonCachedTokens { - debugLogger.Info("Input is shorter than the nonCachedToken, no disaggregated PD") - return false + return 0, false } - // inspect the decode endpoint to disaggregate if prefill should run or not. - // if the non-cached part is short enough - no disaggregation. prefixInfoRaw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.String()) if !ok || prefixInfoRaw == nil { logger.Error(nil, "unable to read prefix cache state") - return false + return 0, false } prefixCacheMatchInfo, ok := prefixInfoRaw.(*attrprefix.PrefixCacheMatchInfo) if !ok { logger.Error(nil, "wrong type of prefix cache match info") - return false + return 0, false } - // number of cached tokens. Use the unweighted cached-block count, not the - // tier-weighted match score: a RAM-cached prefix must contribute its full - // token count here, otherwise the non-cached suffix is overestimated and - // requests with large local-RAM hits are misrouted to remote prefill. hitPrefixTokens := prefixCacheMatchInfo.CachedBlockCount() * prefixCacheMatchInfo.BlockSizeTokens() - // length of non-cached suffix in tokens nonCachedTokens := inputTokens - hitPrefixTokens - debugLogger.Info("Computed hit percentage for prefix cache", "absolute hit prefix len (tokens)", hitPrefixTokens, "prompt length (token)", inputTokens) - - if nonCachedTokens < d.config.NonCachedTokens { - debugLogger.Info("Non-cached suffix is smaller than threshold, using decode profile only") - return false // do not run prefill - } - - return true + return nonCachedTokens, true } // getUserInputLenInTokens returns an estimated token count for the user input. diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go index 3d2b365a7a..7589a2caba 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/prefix_based_pd_decider_test.go @@ -415,6 +415,76 @@ func TestDisaggregate_UsesUnweightedCachedBlockCount(t *testing.T) { "sanity: the tier-weighted score alone undercounts cached blocks and misroutes") } +func TestShouldRejectConditionalDecode(t *testing.T) { + ctx := utils.NewTestContext(t) + + tests := []struct { + name string + nonCachedTokens int + request *scheduling.InferenceRequest + endpoint scheduling.Endpoint + expectReject bool + }{ + { + name: "threshold zero disables gate", + nonCachedTokens: 0, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpoint(0), + expectReject: false, + }, + { + name: "non-cached suffix below threshold forwards", + nonCachedTokens: 5, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpoint(8), + expectReject: false, + }, + { + name: "non-cached suffix at threshold rejects", + nonCachedTokens: 5, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpoint(5), + expectReject: true, + }, + { + name: "no cache hit at all rejects", + nonCachedTokens: 5, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpoint(0), + expectReject: true, + }, + { + name: "fully cached prompt forwards", + nonCachedTokens: 1, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpoint(10), + expectReject: false, + }, + { + name: "missing prefix info rejects (fail closed)", + nonCachedTokens: 5, + request: makeRequestWithTokens(10), + endpoint: makeTestEndpointBase(), + expectReject: true, + }, + { + name: "nil endpoint rejects", + nonCachedTokens: 5, + request: makeRequestWithTokens(10), + endpoint: nil, + expectReject: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + decider, err := NewPrefixBasedPDDecider(PrefixBasedPDDeciderConfig{NonCachedTokens: tt.nonCachedTokens}) + require.NoError(t, err) + assert.Equal(t, tt.expectReject, decider.ShouldRejectConditionalDecode(ctx, tt.request, tt.endpoint)) + }) + } +} + func TestDisaggregateNoPrefixInfo(t *testing.T) { ctx := utils.NewTestContext(t) diff --git a/pkg/epp/requestcontrol/director.go b/pkg/epp/requestcontrol/director.go index afe3ea7c6e..6ede831678 100644 --- a/pkg/epp/requestcontrol/director.go +++ b/pkg/epp/requestcontrol/director.go @@ -48,7 +48,6 @@ import ( fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling" fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling" - attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" "github.com/llm-d/llm-d-router/pkg/epp/handlers" "github.com/llm-d/llm-d-router/pkg/epp/metadata" "github.com/llm-d/llm-d-router/pkg/epp/metrics" @@ -61,49 +60,26 @@ const ( responseBodyQueueCapacity = 100 ) -// primaryEndpointHasCachedPrefix reports whether the primary profile's chosen -// endpoint has at least one matching prefix block in its KV cache, as observed -// by a precise/approximate-prefix scorer during the decode profile run. It -// returns false when the result is missing, the primary profile produced no -// endpoint, the endpoint carries no PrefixCacheMatchInfo attribute, or the -// recorded match has zero blocks. False-return reasons are logged at -// V(logutil.DEBUG) to disambiguate misconfiguration (no scorer attached) from -// a real cache miss. -func primaryEndpointHasCachedPrefix(logger logr.Logger, result *fwksched.SchedulingResult) bool { +// primaryDecodeEndpoint returns the first endpoint chosen by the primary +// profile, or nil when the result is empty or malformed. False-return reasons +// are logged at V(logutil.DEBUG) to disambiguate misconfiguration from a real +// cache miss. +func primaryDecodeEndpoint(logger logr.Logger, result *fwksched.SchedulingResult) fwksched.Endpoint { debug := logger.V(logutil.DEBUG) if result == nil { debug.Info("conditional-decode: scheduling result is nil") - return false + return nil } primary, ok := result.ProfileResults[result.PrimaryProfileName] if !ok || primary == nil { debug.Info("conditional-decode: primary profile result missing", "primary", result.PrimaryProfileName) - return false + return nil } if len(primary.TargetEndpoints) == 0 { debug.Info("conditional-decode: primary profile produced no endpoints", "primary", result.PrimaryProfileName) - return false - } - endpoint := primary.TargetEndpoints[0] - if endpoint == nil { - debug.Info("conditional-decode: primary endpoint is nil") - return false - } - raw, ok := endpoint.Get(attrprefix.PrefixCacheMatchInfoDataKey.String()) - if !ok || raw == nil { - debug.Info("conditional-decode: endpoint has no prefix-cache match attribute (no scorer attached?)") - return false - } - info, ok := raw.(*attrprefix.PrefixCacheMatchInfo) - if !ok { - debug.Info("conditional-decode: prefix-cache attribute has unexpected type", "type", fmt.Sprintf("%T", raw)) - return false - } - if info.MatchBlocks() == 0 { - debug.Info("conditional-decode: prefix-cache match has zero blocks") - return false + return nil } - return true + return primary.TargetEndpoints[0] } // Datastore defines the interface required by the Director. @@ -330,14 +306,21 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo // at encode/prefill/decode. Lives in the director (not in a profile handler) // so it fires regardless of which profile handler is configured. if routing.IsConditionalDecode(reqCtx.Request.Headers) { - if !primaryEndpointHasCachedPrefix(logger, result) { - logger.V(logutil.DEBUG).Info("conditional-decode: chosen decode worker has no cached prefix, returning 412") - return reqCtx, errcommon.Error{ - Code: errcommon.PreconditionFailed, - Msg: "no decode worker has the requested KV cache", + decider := d.requestControlPlugins.ConditionalDecodeDecider() + switch { + case decider == nil: + logger.V(logutil.DEBUG).Info("conditional-decode: no decider configured, forwarding") + default: + endpoint := primaryDecodeEndpoint(logger, result) + if endpoint == nil || decider.ShouldRejectConditionalDecode(ctx, reqCtx.SchedulingRequest, endpoint) { + logger.V(logutil.DEBUG).Info("conditional-decode: chosen decode worker has no cached prefix, returning 412") + return reqCtx, errcommon.Error{ + Code: errcommon.PreconditionFailed, + Msg: "no decode worker has the requested KV cache", + } } + logger.V(logutil.DEBUG).Info("conditional-decode: chosen decode worker has cached prefix, forwarding") } - logger.V(logutil.DEBUG).Info("conditional-decode: chosen decode worker has cached prefix, forwarding") } reqCtx.SchedulingRequest.SchedulingResult = result diff --git a/pkg/epp/requestcontrol/director_test.go b/pkg/epp/requestcontrol/director_test.go index 1bf5ecb943..11dfccf12b 100644 --- a/pkg/epp/requestcontrol/director_test.go +++ b/pkg/epp/requestcontrol/director_test.go @@ -27,7 +27,6 @@ import ( "testing" "time" - "github.com/go-logr/logr" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" "github.com/stretchr/testify/assert" @@ -56,6 +55,7 @@ import ( attrprefix "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/datalayer/attribute/prefix" "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/requesthandling/parsers/openai" sessionaffinityfilter "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/filter/sessionaffinity" + "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/profilehandler/disagg" sessionaffinityscorer "github.com/llm-d/llm-d-router/pkg/epp/framework/plugins/scheduling/scorer/sessionaffinity" "github.com/llm-d/llm-d-router/pkg/epp/handlers" "github.com/llm-d/llm-d-router/pkg/epp/metadata" @@ -1735,77 +1735,10 @@ func newResponseBodyTestRequestContext(requestID string, completionTokens int) * // ── Conditional-decode gate (Prefer: if-available) ───────────────────────── -// wrongTypeAttr is a Cloneable that is NOT *attrprefix.PrefixCacheMatchInfo, -// used to exercise the type-assertion failure branch. -type wrongTypeAttr struct{} - -func (w wrongTypeAttr) Clone() fwkdl.Cloneable { return w } - -func TestPrimaryEndpointHasCachedPrefix(t *testing.T) { - endpointWith := func(matched, total int) fwksched.Endpoint { - attrs := fwkdl.NewAttributes() - attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), - attrprefix.NewPrefixCacheMatchInfo(matched, total, 1)) - return fwksched.NewEndpoint( - &fwkdl.EndpointMetadata{NamespacedName: types.NamespacedName{Namespace: "default", Name: "p"}}, - nil, attrs, - ) - } - endpointBare := func() fwksched.Endpoint { - return fwksched.NewEndpoint( - &fwkdl.EndpointMetadata{NamespacedName: types.NamespacedName{Namespace: "default", Name: "p"}}, - nil, fwkdl.NewAttributes(), - ) - } - endpointWithWrongType := func() fwksched.Endpoint { - attrs := fwkdl.NewAttributes() - attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), wrongTypeAttr{}) - return fwksched.NewEndpoint( - &fwkdl.EndpointMetadata{NamespacedName: types.NamespacedName{Namespace: "default", Name: "p"}}, - nil, attrs, - ) - } - resultWith := func(eps ...fwksched.Endpoint) *fwksched.SchedulingResult { - return &fwksched.SchedulingResult{ - PrimaryProfileName: "decode", - ProfileResults: map[string]*fwksched.ProfileRunResult{ - "decode": {TargetEndpoints: eps}, - }, - } - } - - tests := []struct { - name string - in *fwksched.SchedulingResult - want bool - }{ - {"nil result", nil, false}, - {"empty profile results", &fwksched.SchedulingResult{PrimaryProfileName: "decode"}, false}, - {"primary profile missing", &fwksched.SchedulingResult{ - PrimaryProfileName: "decode", - ProfileResults: map[string]*fwksched.ProfileRunResult{"other": {TargetEndpoints: []fwksched.Endpoint{endpointWith(2, 4)}}}, - }, false}, - {"primary profile nil", &fwksched.SchedulingResult{ - PrimaryProfileName: "decode", - ProfileResults: map[string]*fwksched.ProfileRunResult{"decode": nil}, - }, false}, - {"primary has no endpoints", resultWith(), false}, - {"endpoint has no match info", resultWith(endpointBare()), false}, - {"wrong type in attribute", resultWith(endpointWithWrongType()), false}, - {"zero match blocks", resultWith(endpointWith(0, 4)), false}, - {"some match blocks", resultWith(endpointWith(2, 4)), true}, - {"full match", resultWith(endpointWith(4, 4)), true}, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - assert.Equal(t, tt.want, primaryEndpointHasCachedPrefix(logr.Discard(), tt.in)) - }) - } -} - // newConditionalDecodeDirector builds a minimal Director suitable for -// exercising the conditional-decode gate end-to-end. -func newConditionalDecodeDirector(t *testing.T, scheduleResult *fwksched.SchedulingResult) (*Director, context.Context) { +// exercising the conditional-decode gate end-to-end. When decider is non-nil +// it is wired into the request-control config; otherwise the gate is disabled. +func newConditionalDecodeDirector(t *testing.T, scheduleResult *fwksched.SchedulingResult, decider fwkrc.ConditionalDecodeDecider) (*Director, context.Context) { t.Helper() ctx := logutil.NewTestLoggerIntoContext(context.Background()) @@ -1837,18 +1770,30 @@ func newConditionalDecodeDirector(t *testing.T, scheduleResult *fwksched.Schedul }) mockSched := &mockScheduler{scheduleResults: scheduleResult} - cfg := NewConfig().WithAdmissionPlugins(newMockAdmissionPlugin("admit", nil)) + cfg := NewConfig(). + WithAdmissionPlugins(newMockAdmissionPlugin("admit", nil)). + WithConditionalDecodeDecider(decider) candidates := NewCachedEndpointCandidates(context.Background(), NewDatastoreEndpointCandidates(ds), time.Minute) dir := NewDirectorWithConfig(ds, mockSched, &mockAdmissionController{}, candidates, cfg) return dir, ctx } func TestDirector_HandleRequest_ConditionalDecode(t *testing.T) { - scheduleResultWith := func(matched, total int) *fwksched.SchedulingResult { + const ( + blockSize = 1 + nonCachedTokens = 5 + inputTokens = 10 + ) + + // scheduleResultWith builds a result where the chosen decode endpoint + // reports `cached` cached blocks (each `blockSize` tokens). cached == -1 + // means the endpoint carries no PrefixCacheMatchInfo attribute. + scheduleResultWith := func(cached int) *fwksched.SchedulingResult { attrs := fwkdl.NewAttributes() - if matched >= 0 { + if cached >= 0 { attrs.Put(attrprefix.PrefixCacheMatchInfoDataKey.String(), - attrprefix.NewPrefixCacheMatchInfo(matched, total, 1)) + attrprefix.NewPrefixCacheMatchInfo(cached, inputTokens, blockSize). + WithCachedBlockCount(cached)) } return &fwksched.SchedulingResult{ PrimaryProfileName: "decode", @@ -1865,23 +1810,33 @@ func TestDirector_HandleRequest_ConditionalDecode(t *testing.T) { } } + newDecider := func(t *testing.T) fwkrc.ConditionalDecodeDecider { + t.Helper() + d, err := disagg.NewPrefixBasedPDDecider(disagg.PrefixBasedPDDeciderConfig{NonCachedTokens: nonCachedTokens}) + require.NoError(t, err) + return d + } + tests := []struct { name string preferValue string // empty == no Prefer header - matched int // -1 == no PrefixCacheMatchInfo at all - total int + cached int // -1 == no PrefixCacheMatchInfo at all + decider func(t *testing.T) fwkrc.ConditionalDecodeDecider wantErrCode string }{ - {"prefer if-available + cache hit forwards", "if-available", 2, 4, ""}, - {"prefer if-available + zero match returns 412", "if-available", 0, 4, errcommon.PreconditionFailed}, - {"prefer if-available + no match info returns 412", "if-available", -1, 0, errcommon.PreconditionFailed}, - {"absent Prefer header proceeds even with no cache", "", -1, 0, ""}, - {"unrelated Prefer token proceeds even with no cache", "return=minimal", -1, 0, ""}, + // inputTokens=10, threshold=5: cached>=6 → nonCached<=4 → forward; cached<=5 → nonCached>=5 → reject. + {"prefer if-available + suffix below threshold forwards", "if-available", 8, newDecider, ""}, + {"prefer if-available + suffix at threshold returns 412", "if-available", 5, newDecider, errcommon.PreconditionFailed}, + {"prefer if-available + zero cache returns 412", "if-available", 0, newDecider, errcommon.PreconditionFailed}, + {"prefer if-available + no match info returns 412", "if-available", -1, newDecider, errcommon.PreconditionFailed}, + {"absent Prefer header proceeds regardless of cache", "", -1, newDecider, ""}, + {"unrelated Prefer token proceeds regardless of cache", "return=minimal", -1, newDecider, ""}, + {"no decider configured forwards even on cache miss", "if-available", -1, func(*testing.T) fwkrc.ConditionalDecodeDecider { return nil }, ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - dir, ctx := newConditionalDecodeDirector(t, scheduleResultWith(tt.matched, tt.total)) + dir, ctx := newConditionalDecodeDirector(t, scheduleResultWith(tt.cached), tt.decider(t)) reqCtx := &handlers.RequestContext{ Request: &handlers.Request{ @@ -1900,6 +1855,12 @@ func TestDirector_HandleRequest_ConditionalDecode(t *testing.T) { parseResult, err := openai.NewOpenAIParser().ParseRequest(ctx, reqCtx.Request.RawBody, reqCtx.Request.Headers) require.NoError(t, err) + // Inject a tokenized prompt so the decider's token-count branch + // sees inputTokens=10 instead of 0 (no tokenizer plugin runs in + // this minimal director fixture). + parseResult.Body.TokenizedPrompt = &fwkrh.TokenizedPrompt{ + PerPromptTokens: [][]uint32{make([]uint32, inputTokens)}, + } _, err = dir.HandleRequest(ctx, reqCtx, parseResult.Body) if tt.wantErrCode == "" { diff --git a/pkg/epp/requestcontrol/request_control_config.go b/pkg/epp/requestcontrol/request_control_config.go index 51697909cd..b38e217770 100644 --- a/pkg/epp/requestcontrol/request_control_config.go +++ b/pkg/epp/requestcontrol/request_control_config.go @@ -17,6 +17,8 @@ limitations under the License. package requestcontrol import ( + "sigs.k8s.io/controller-runtime/pkg/log" + "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin" fwkrc "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requestcontrol" ) @@ -41,6 +43,12 @@ type Config struct { preRequestPlugins []fwkrc.PreRequest responseReceivedPlugins []fwkrc.ResponseHeaderProcessor responseStreamingPlugins []fwkrc.ResponseBodyProcessor + + // conditionalDecodeDecider, when set, is consulted by the director on + // requests carrying RFC 7240 "Prefer: if-available" to decide whether to + // reject the request with HTTP 412. When nil, conditional-decode requests + // are forwarded unconditionally. + conditionalDecodeDecider fwkrc.ConditionalDecodeDecider } // WithPreAdmissionPlugins sets the given plugins as the PreAdmitter plugins. @@ -76,6 +84,19 @@ func (c *Config) WithDataProducerPlugins(plugins ...fwkrc.DataProducer) *Config return c } +// WithConditionalDecodeDecider sets the plugin consulted by the director to +// gate RFC 7240 "Prefer: if-available" requests. Passing nil disables the gate. +func (c *Config) WithConditionalDecodeDecider(d fwkrc.ConditionalDecodeDecider) *Config { + c.conditionalDecodeDecider = d + return c +} + +// ConditionalDecodeDecider returns the configured conditional-decode decider, +// or nil when no plugin is wired in. +func (c *Config) ConditionalDecodeDecider() fwkrc.ConditionalDecodeDecider { + return c.conditionalDecodeDecider +} + // WithAdmissionPlugins sets the given plugins as the Admit plugins. func (c *Config) WithAdmissionPlugins(plugins ...fwkrc.Admitter) *Config { c.admissionPlugins = plugins @@ -105,6 +126,15 @@ func (c *Config) AddPlugins(pluginObjects ...plugin.Plugin) { if admissionPlugin, ok := plugin.(fwkrc.Admitter); ok { c.admissionPlugins = append(c.admissionPlugins, admissionPlugin) } + if condDecider, ok := plugin.(fwkrc.ConditionalDecodeDecider); ok { + if c.conditionalDecodeDecider == nil { + c.conditionalDecodeDecider = condDecider + } else { + log.Log.Info("Multiple ConditionalDecodeDecider plugins configured; ignoring later instance", + "kept", c.conditionalDecodeDecider.TypedName().String(), + "ignored", condDecider.TypedName().String()) + } + } } } From 7bd1128388489d4607bc94b69d832e458426d87a Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 16:14:58 +0300 Subject: [PATCH 2/8] git commit -sa -m "fix(epp): flatten conditional-decode switch into if/else" Signed-off-by: Dmitri Pikus --- pkg/epp/requestcontrol/director.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pkg/epp/requestcontrol/director.go b/pkg/epp/requestcontrol/director.go index 6ede831678..c133821201 100644 --- a/pkg/epp/requestcontrol/director.go +++ b/pkg/epp/requestcontrol/director.go @@ -307,10 +307,9 @@ func (d *Director) HandleRequest(ctx context.Context, reqCtx *handlers.RequestCo // so it fires regardless of which profile handler is configured. if routing.IsConditionalDecode(reqCtx.Request.Headers) { decider := d.requestControlPlugins.ConditionalDecodeDecider() - switch { - case decider == nil: + if decider == nil { logger.V(logutil.DEBUG).Info("conditional-decode: no decider configured, forwarding") - default: + } else { endpoint := primaryDecodeEndpoint(logger, result) if endpoint == nil || decider.ShouldRejectConditionalDecode(ctx, reqCtx.SchedulingRequest, endpoint) { logger.V(logutil.DEBUG).Info("conditional-decode: chosen decode worker has no cached prefix, returning 412") From 541981e90ed3c0f0f6bb9269390a7bb286d0744f Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 17:47:34 +0300 Subject: [PATCH 3/8] docs: document conditional-decode 412 gate behavior change Signed-off-by: Dmitri Pikus --- docs/disaggregation.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/disaggregation.md b/docs/disaggregation.md index 56f9cdb398..570e23b0e6 100644 --- a/docs/disaggregation.md +++ b/docs/disaggregation.md @@ -413,6 +413,20 @@ The `prefix-based-pd-decider` plugin makes the disaggregation decision according - `nonCachedTokens`: Number of non-cached tokens that trigger disaggregation - If set to 0, disaggregation never occurs for any request +**Conditional-decode 412 gate** + +When this plugin is configured, it also drives the EPP's RFC 7240 `Prefer: if-available` (conditional-decode) gate. The coordinator uses that header to mark a speculative early-decode attempt: forward to a decode worker only if its KV cache already covers the prompt, otherwise return HTTP 412 Precondition Failed so the coordinator restarts the pipeline at encode/prefill/decode. + +The gate uses the same `nonCachedTokens` threshold as the disaggregation decision: + +- Non-cached suffix `< nonCachedTokens` → forward to the chosen decode worker. +- Non-cached suffix `≥ nonCachedTokens` → return 412. +- `nonCachedTokens: 0` → gate disabled (forwards every conditional-decode request). +- Plugin not declared in the config → gate disabled (the EPP forwards conditional-decode requests unconditionally). +- Prefix-cache state unreadable on the chosen endpoint (e.g. no approximate-prefix producer wired up) → fail closed (412). + +This behavior differs from earlier releases, where the gate was a hard-coded check inside the director that forwarded on any cache hit and rejected only on a complete cache miss. The threshold is now in tokens (not blocks), is operator-configurable, and is opt-in via the plugin. + #### Always-Disagg PD Decider The `always-disagg-pd-decider` is a simpler alternative used mainly for testing or benchmarking. It always triggers disaggregation, regardless of prefix cache state or prompt characteristics. From cd229c3ceca449e7d34f9271bd494915db5e34fb Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 18:30:03 +0300 Subject: [PATCH 4/8] docs: more clarification about conditional-decode 412 gate behavior change Signed-off-by: Dmitri Pikus --- .../profilehandler/disagg/README.md | 50 +++++++++++++++---- 1 file changed, 41 insertions(+), 9 deletions(-) diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md index ca1f53c6a5..768cd33858 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md @@ -2,6 +2,8 @@ Plugins for disaggregated inference scheduling: a profile handler that selects the active stages: EPD (no disaggregation), P/D (Prefill/Decode), E/P/D (Encode/Prefill/Decode), or E/PD (Encode/Prefill-Decode), legacy headers handlers (deprecated) kept for backward compatibility, and decider plugins that control whether each disaggregation stage runs per request. +`PrefixBasedPDDecider` additionally drives the EPP director's conditional-decode 412 gate (RFC 7240 `Prefer: if-available`) via the `ConditionalDecodeDecider` extension point — see its section for details. + ## Contents - [Profile Handlers](#profile-handlers) @@ -170,20 +172,35 @@ plugins: **Type:** `prefix-based-pd-decider` -Decides per-request whether P/D disaggregation should run, based on how much of the prompt is already cached on the selected decode pod. +Drives two decisions per request, both based on how much of the prompt is already cached on the selected decode pod: + +1. **P/D disaggregation** — whether to offload prefill to a remote prefill pod (consumed by `disagg-profile-handler` via `deciders.prefill`). +2. **Conditional-decode 412 gate** — whether a request carrying RFC 7240 `Prefer: if-available` should be rejected with HTTP 412 instead of forwarded (consumed by the EPP director through the `ConditionalDecodeDecider` extension point). + +Both decisions share the same `nonCachedTokens` threshold and the same uncached-suffix computation; they differ only in their failure semantics (see [How It Works](#how-it-works)). #### What it does -Compares the uncached portion of the request prompt against a configurable threshold, triggering P/D disaggregation only when the uncached suffix is long enough to justify the overhead. +Compares the uncached portion of the request prompt against a configurable threshold: 1. Read the prompt token count as `len(request.Body.TokenizedPrompt.TokenIDs)`. 2. Read `PrefixCacheMatchInfo` from the decode endpoint attributes. 3. Compute uncached suffix length. -4. Return true (disaggregate) if uncached tokens ≥ `nonCachedTokens`. +4. **Disaggregation:** return true (disaggregate) if uncached tokens ≥ `nonCachedTokens`. +5. **Conditional-decode gate:** return true (reject with 412) if uncached tokens ≥ `nonCachedTokens`. #### How It Works -The prompt token count is `len(request.Body.TokenizedPrompt.TokenIDs)`, populated by a `token-producer` — auto-created with the tokenizer-free `estimate` backend when none is configured. Prefix cache state is read from the `PrefixCacheMatchInfo` attribute on the decode endpoint, populated by `approx-prefix-cache-producer`. If the attribute is absent or malformed, disaggregation is skipped. Setting `nonCachedTokens: 0` disables the decider entirely (always returns false). +The prompt token count is `len(request.Body.TokenizedPrompt.TokenIDs)`, populated by a `token-producer` — auto-created with the tokenizer-free `estimate` backend when none is configured. Prefix cache state is read from the `PrefixCacheMatchInfo` attribute on the decode endpoint, populated by `approx-prefix-cache-producer`. + +Setting `nonCachedTokens: 0` disables both decisions entirely (disaggregation never runs, the conditional-decode gate always forwards). + +**Failure semantics** when prefix-cache state is unreadable (attribute missing, malformed, or producer not configured) differ between the two decisions: + +| Decision | Behavior on unreadable prefix info | Rationale | +|---|---|---| +| Disaggregation | fail open — return false (no disaggregation) | A misconfigured prefix-cache producer should not silently route every request to remote prefill. | +| Conditional-decode 412 gate | fail closed — return true (reject with 412) | The header is a "fast-fail" hint from the coordinator: if the EPP cannot prove the cache covers the prompt, the safe answer is to bounce so the coordinator falls back to the full pipeline. | #### Inputs consumed @@ -195,9 +212,9 @@ The prompt token count is `len(request.Body.TokenizedPrompt.TokenIDs)`, populate ##### Parameters | Name | Type | Required | Default | Description | |------|------|----------|---------|-------------| -| `nonCachedTokens` | `int` | No | `0` | Uncached token threshold above which P/D disaggregation is triggered. `0` disables the decider. | +| `nonCachedTokens` | `int` | No | `0` | Uncached token threshold above which P/D disaggregation is triggered, and at or above which the conditional-decode 412 gate rejects. `0` disables both behaviors. | -##### Example +##### Example — P/D disaggregation ```yaml plugins: - type: prefix-based-pd-decider @@ -209,11 +226,26 @@ plugins: prefill: prefix-based-pd-decider ``` +##### Example — conditional-decode 412 gate only (no P/D) +```yaml +plugins: + - type: prefix-based-pd-decider + parameters: + nonCachedTokens: 512 + # No disagg-profile-handler / no `prefill` profile. + # The plugin is auto-registered as the ConditionalDecodeDecider via + # AddPlugins; the director consults it whenever a request carries the + # RFC 7240 `Prefer: if-available` header. +``` + +When the plugin is *not* declared in the EPP config, the conditional-decode gate is disabled and the director forwards every `Prefer: if-available` request unconditionally. + #### Limitations -- `nonCachedTokens: 0` disables disaggregation entirely (the decider always returns false). -- A `token-producer` populates `TokenizedPrompt`; when none is configured the framework auto-creates one with the `estimate` backend, so disaggregation works without extra setup. -- Requires `PrefixCacheMatchInfo` on the decode endpoint; if absent, disaggregation is skipped with an error log. +- `nonCachedTokens: 0` disables both the disaggregation decision and the conditional-decode gate (the decider returns false for disaggregation and false for "should reject"). +- A `token-producer` populates `TokenizedPrompt`; when none is configured the framework auto-creates one with the `estimate` backend, so both behaviors work without extra setup. +- Requires `PrefixCacheMatchInfo` on the decode endpoint. If absent: disaggregation is skipped with an error log (fail open); the conditional-decode gate rejects with 412 (fail closed). +- Only one `ConditionalDecodeDecider` is consulted per director. If multiple plugins implement the interface, the first one registered through `Config.AddPlugins` wins; later instances are logged and ignored. --- From 4cf1121ed3329f064b5df834d2f0332579070673 Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 18:41:16 +0300 Subject: [PATCH 5/8] git commit -s -m "docs(epp): fix stale comment on primaryDecodeEndpoint" Signed-off-by: Dmitri Pikus --- pkg/epp/requestcontrol/director.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/epp/requestcontrol/director.go b/pkg/epp/requestcontrol/director.go index c133821201..a8f7666097 100644 --- a/pkg/epp/requestcontrol/director.go +++ b/pkg/epp/requestcontrol/director.go @@ -61,9 +61,10 @@ const ( ) // primaryDecodeEndpoint returns the first endpoint chosen by the primary -// profile, or nil when the result is empty or malformed. False-return reasons -// are logged at V(logutil.DEBUG) to disambiguate misconfiguration from a real -// cache miss. +// profile, or nil when the result is nil, the primary profile is missing, or +// it produced no endpoints. Each nil-return reason is logged at +// V(logutil.DEBUG) so an unexpected nil here can be distinguished from a +// healthy schedule with no endpoint. func primaryDecodeEndpoint(logger logr.Logger, result *fwksched.SchedulingResult) fwksched.Endpoint { debug := logger.V(logutil.DEBUG) if result == nil { From 765591d6d475f7ad06f9d7eaf7d1e3f543b2c7c5 Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 19:03:13 +0300 Subject: [PATCH 6/8] Comment is changed not to contain temporal language Signed-off-by: Dmitri Pikus --- docs/disaggregation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/disaggregation.md b/docs/disaggregation.md index 570e23b0e6..f9d5a5ddd5 100644 --- a/docs/disaggregation.md +++ b/docs/disaggregation.md @@ -425,7 +425,7 @@ The gate uses the same `nonCachedTokens` threshold as the disaggregation decisio - Plugin not declared in the config → gate disabled (the EPP forwards conditional-decode requests unconditionally). - Prefix-cache state unreadable on the chosen endpoint (e.g. no approximate-prefix producer wired up) → fail closed (412). -This behavior differs from earlier releases, where the gate was a hard-coded check inside the director that forwarded on any cache hit and rejected only on a complete cache miss. The threshold is now in tokens (not blocks), is operator-configurable, and is opt-in via the plugin. +The threshold is configured in tokens, drives both the disaggregation decision and the 412 gate, and is opt-in: omitting the plugin disables the gate. #### Always-Disagg PD Decider The `always-disagg-pd-decider` is a simpler alternative used mainly for testing or benchmarking. From 8f5e2e628f889b30cde27f4829639491fe934ae3 Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 19:11:14 +0300 Subject: [PATCH 7/8] docs: describe conditional-decode gate independently of coordinator Signed-off-by: Dmitri Pikus --- docs/disaggregation.md | 5 ++++- .../plugins/scheduling/profilehandler/disagg/README.md | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/disaggregation.md b/docs/disaggregation.md index f9d5a5ddd5..02cd96c157 100644 --- a/docs/disaggregation.md +++ b/docs/disaggregation.md @@ -415,7 +415,10 @@ The `prefix-based-pd-decider` plugin makes the disaggregation decision according **Conditional-decode 412 gate** -When this plugin is configured, it also drives the EPP's RFC 7240 `Prefer: if-available` (conditional-decode) gate. The coordinator uses that header to mark a speculative early-decode attempt: forward to a decode worker only if its KV cache already covers the prompt, otherwise return HTTP 412 Precondition Failed so the coordinator restarts the pipeline at encode/prefill/decode. +When this plugin is configured, it also gates requests carrying the RFC 7240 `Prefer: if-available` header: a request is forwarded to the chosen decode worker only when the worker's KV cache already covers the prompt; otherwise the EPP returns HTTP 412 Precondition Failed. + +> **Note** +> The llm-d coordinator uses this header to mark speculative early-decode attempts. A 412 response signals that the local cache was insufficient, prompting the coordinator to restart the pipeline at encode/prefill/decode. The gate uses the same `nonCachedTokens` threshold as the disaggregation decision: diff --git a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md index 768cd33858..0cc71f40d0 100644 --- a/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md +++ b/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/README.md @@ -200,7 +200,7 @@ Setting `nonCachedTokens: 0` disables both decisions entirely (disaggregation ne | Decision | Behavior on unreadable prefix info | Rationale | |---|---|---| | Disaggregation | fail open — return false (no disaggregation) | A misconfigured prefix-cache producer should not silently route every request to remote prefill. | -| Conditional-decode 412 gate | fail closed — return true (reject with 412) | The header is a "fast-fail" hint from the coordinator: if the EPP cannot prove the cache covers the prompt, the safe answer is to bounce so the coordinator falls back to the full pipeline. | +| Conditional-decode 412 gate | fail closed — return true (reject with 412) | The `Prefer: if-available` header is a fast-fail hint: if the EPP cannot prove the cache covers the prompt, the safe answer is to reject with 412 so the caller can fall back. | #### Inputs consumed From 4035835f2f855f57ca7f1a801994a4da9049382f Mon Sep 17 00:00:00 2001 From: Dmitri Pikus Date: Thu, 18 Jun 2026 19:25:55 +0300 Subject: [PATCH 8/8] docs(epp): make ConditionalDecodeDecider doc self-contained Signed-off-by: Dmitri Pikus --- pkg/epp/framework/interface/requestcontrol/plugins.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pkg/epp/framework/interface/requestcontrol/plugins.go b/pkg/epp/framework/interface/requestcontrol/plugins.go index a0ec88a1aa..336120d841 100644 --- a/pkg/epp/framework/interface/requestcontrol/plugins.go +++ b/pkg/epp/framework/interface/requestcontrol/plugins.go @@ -99,10 +99,9 @@ type PreAdmitter interface { PreAdmit(ctx context.Context, request *fwksched.InferenceRequest) error } -// ConditionalDecodeDecider answers the RFC 7240 "Prefer: if-available" gate: -// given the chosen decode endpoint, should the request be rejected with HTTP -// 412 (because the local KV cache does not cover enough of the prompt) or -// forwarded? At most one such plugin is consulted per director. +// ConditionalDecodeDecider decides whether a request carrying the +// RFC 7240 "Prefer: if-available" header should be rejected with HTTP 412 +// Precondition Failed or forwarded to the chosen decode endpoint. type ConditionalDecodeDecider interface { plugin.Plugin ShouldRejectConditionalDecode(ctx context.Context, request *fwksched.InferenceRequest, endpoint fwksched.Endpoint) bool