Delegate conditional-decode 412 decision to prefix-based-pd-decider

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
+}

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,58 +97,91 @@ 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
 	// to trigger disaggregated PD. A value of 0 disables disaggregation.
 	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.

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)
 

pkg/epp/requestcontrol/director.go

@@ -48,7 +48,6 @@
 	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 @@
 	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 @@
 	// 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