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

docs/disaggregation.md

@@ -413,6 +413,23 @@ 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 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:
+
+- 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).
+
+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.
 It always triggers disaggregation, regardless of prefix cache state or prompt characteristics.

pkg/epp/framework/interface/requestcontrol/plugins.go

@@ -98,3 +98,11 @@ type PreAdmitter interface {
 	plugin.Plugin
 	PreAdmit(ctx context.Context, request *fwksched.InferenceRequest) error
 }
+
+// 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
+}

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 `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
 
@@ -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.
 
 ---
 

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)