Skip to content
Open
14 changes: 14 additions & 0 deletions docs/disaggregation.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wonder if this gate can be described independently, with coordinator use as a note.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I re-wrote the comments to describe the gate independently and pull the coordinator behavior into a note. Thank you!


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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should avoid temporal language as instructed in AGENTS.md.

Would be interesting to understand if your agent "willingly" disobeyed that instruction or it just lost it in context.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vMaroon , this is an interesting point.
In this case the agent didn't work according to the instruction.
The agent states that "it didn't deliberately set the rule aside, but it also didn't consciously check it against the wording I was producing..."
We probably need to explicitly re-check/re-run against AGENTS.md

In any case, I re-wrote the comment not to contain temporal language


#### 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.
Expand Down
9 changes: 9 additions & 0 deletions pkg/epp/framework/interface/requestcontrol/plugins.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Description should be self-contained, and descriptive of what the module does.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I re-phrased the description.

Thank you!

plugin.Plugin
ShouldRejectConditionalDecode(ctx context.Context, request *fwksched.InferenceRequest, endpoint fwksched.Endpoint) bool
}
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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.

---

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand All @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
Loading
Loading