feat: Add selective KV transfer policy

[albertoperdomo2]

What type of PR is this?

/kind feature

What this PR does / why we need it:

This PR adds the Alpha selective-kv-policy request-control plugin. It lets the EPP control KV loading and offloading independently for each request.

The load policy can preserve the client setting, disable external loading, or load only when the selected endpoint reports enough reusable KV outside GPU memory. Threshold mode can also veto loading when an endpoint-local EWMA of vllm:num_requests_waiting reaches a configured limit. The offload policy can preserve offloading or set max_offload_tokens: 0.

Disabling a load sets kv_transfer_params.max_load_tokens: 0. Allowing a load removes that field to restore the vLLM default. The plugin leaves kv_load_tiers unchanged. This requires vLLM PR #55885 (merged), or a build with the same max_load_tokens API.

The plugin is disabled unless EPP starts with --allow-experimental-plugins. Unsupported request shapes fail open and remain unchanged. The coordinator prefill path is not supported because it replaces kv_transfer_params after this hook runs.

Performance evaluation

The experiment asks whether the proposed 1024 + 8 per-request selective policy can outperform the existing behavior of loading every reusable external prefix while avoiding the cost of always recomputing reusable KV.

Deployment

The deployment used llm-d Precise Prefix Caching with Qwen3-32B. It used the custom vLLM image (based on the aforementioned PR) to deploy 4 replicas of TP2 in one 8xH100 node.

Deployment profiles

Profile	Role	Load behavior
multi-tier-offloading-nvme	Primary baseline	Existing behavior without selective loading. External KV loading is uncapped.
rhoai-distributed-default-selective-loading-disabled	Forced-recompute control	External loading is disabled for every request. Local GPU prefix reuse remains enabled.
rhoai-distributed-default-selective-loading-combined	Feature configuration	Applies the 1,024-token threshold and vetoes loading when the endpoint-local waiting-queue EWMA reaches 8 requests. The gate reopens at 4 (50%).

Tiered offloading included a 256 GiB CPU tier and a local NVMe filesystem tier with 64 read threads and 64 write threads:

--kv-transfer-config={"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"spec_name":"TieringOffloadingSpec","cpu_bytes_to_use":274877906944,"secondary_tiers":[{"type":"fs","root_dir":"/mnt/nvme-kv-cache","n_read_threads":64,"n_write_threads":64}]}}

The EPP config for the selective loading looked like this:

        - type: selective-kv-policy
          name: selective-kv-policy
          parameters:
            loadPolicy: threshold
            minExternalReusableTokens: 1024
            maxWaitingRequests: 8

The forced-load profile represents the behavior currently present in deployments without this feature. Forced recompute is a diagnostic control that establishes the other static extreme. The feature arm is the exact proposed configuration: minExternalReusableTokens: 1024 and maxWaitingRequests: 8.

Workloads

Each GuideLLM profile used a 2,048-request warmup followed by a 300-second concurrent phase. Requests contained a reusable prefix, 256 uncached prompt tokens, and 128 requested output tokens. The prefix corpus contained 2,048 entries, each run used one turn, and the random seed was fixed at 20260909.

The five workload cells target regions where earlier crossover measurements showed different preferred static actions:

Benchmark profile	Reusable prefix	Concurrency
guidellm-selective-loading-calibration-p1024-c64	1,024 tokens	64
guidellm-selective-loading-calibration-p2048-c64	2,048 tokens	64
guidellm-selective-loading-calibration-p4096-c32	4,096 tokens	32
guidellm-selective-loading-calibration-p8192-c32	8,192 tokens	32
guidellm-selective-loading-calibration-p8192-c64	8,192 tokens	64

This evaluation includes pressure-heavy workloads designed to stress external KV restoration: reusable prefixes up to 8,192 tokens at concurrency 64, using CPU and NVMe offloading. Under uncapped loading, restore-path contention caused p99 TTFT to reach 57.5 seconds at 8K/C64 and limited output throughput to 962.5 tokens/s. The proposed 1024 + 8 policy reduced p99 TTFT to 5.48 seconds and increased output throughput to 1,565.6 tokens/s. These results demonstrate the policy’s behavior when external KV loading becomes a bottleneck. The collected telemetry does not establish whether NVMe, CPU staging, CPU-to-GPU transfer, or their combined queueing was the individually saturated component.

Please, keep in mind that the figures added to this PR description come from another (larger) report, therefore the image indices are not correct.

Output-token throughput

Output-token throughput is the rate of generated tokens, excluding input tokens. Because every request requested 128 output tokens, it follows successful request throughput while remaining directly relevant to serving capacity.

Workload	Forced load output tok/s	Forced recompute output tok/s	Proposed 1024 + 8 output tok/s
1K / C64	3,276.8	3,065.4	3,367.2
2K / C64	2,350.8	2,545.6	2,744.4
4K / C32	1,495.9	1,352.1	1,487.1
8K / C32	873.1	995.0	1,162.2
8K / C64	962.5	1,319.1	1,565.6

The proposed policy exceeded the current forced-load baseline in four of five cells. Across the matrix, its geometric-mean output-token throughput was 20.6% higher than forced load and 12.6% higher than forced recompute. At 8K/C64, output throughput increased from 962.5 to 1,565.6 tokens/s, a 62.7% improvement over current behavior.

Figure 3: Output-token throughput from GuideLLM. The proposed 1024 + 8 policy improves serving throughput over uncapped loading in four cells and avoids the forced-load collapse at 8K/C32 and 8K/C64.

Time to first token

TTFT captures lookup, restore, recomputation, scheduling, and prefill delay before generation starts. It is the latency metric most directly affected by the load-versus-recompute decision.

Workload	Forced load mean / p99 TTFT	Forced recompute mean / p99 TTFT	Proposed 1024 + 8 mean / p99 TTFT
1K / C64	196 / 2,507 ms	126 / 425 ms	113 / 213 ms
2K / C64	833 / 16,180 ms	178 / 814 ms	193 / 1,256 ms
4K / C32	233 / 989 ms	232 / 578 ms	199 / 447 ms
8K / C32	1,316 / 29,864 ms	438 / 1,701 ms	365 / 1,834 ms
8K / C64	2,584 / 57,523 ms	541 / 3,441 ms	632 / 5,476 ms

Relative to current forced loading, the proposed policy reduced geometric-mean mean TTFT by 62.2% and p99 TTFT by 88.8%. At 8K/C64, mean TTFT fell from 2,584 ms to 632 ms and p99 fell from 57.5 seconds to 5.48 seconds. Forced recompute retained the best TTFT at some high-pressure points, but it completed fewer requests and produced fewer output tokens than the proposed policy.

Figure 5 and Figure 6: Mean and p99 TTFT from GuideLLM. Uncapped loading develops severe tails when the restore path is under pressure; selective loading avoids that collapse while retaining external reuse.

Inter-token latency

ITL measures spacing between streamed output tokens after the first token. It primarily reflects decode behavior and interference after prefill, so it exposes a different tradeoff from TTFT.

Workload	Forced load mean / p99 ITL	Forced recompute mean / p99 ITL	Proposed 1024 + 8 mean / p99 ITL
1K / C64	18.08 / 24.03 ms	20.06 / 23.27 ms	18.27 / 21.91 ms
2K / C64	20.75 / 31.87 ms	23.94 / 29.28 ms	22.00 / 28.07 ms
4K / C32	19.72 / 27.09 ms	22.03 / 29.88 ms	20.12 / 27.10 ms
8K / C32	23.86 / 47.23 ms	28.95 / 43.65 ms	24.86 / 39.77 ms
8K / C64	31.83 / 69.39 ms	44.65 / 68.63 ms	36.27 / 59.12 ms

Forced loading had the lowest mean ITL in all cells, but it paid for that result with substantially worse TTFT and completed throughput under restore pressure. The proposed policy's geometric-mean ITL was 5.4% higher than forced load and 11.8% lower than forced recompute. Its p99 ITL was 10.4% lower than forced load and 8.5% lower than forced recompute. The policy therefore improves admission and prefill behavior without introducing a large decode-latency penalty.

Figure 7: Mean ITL from GuideLLM. The proposed policy remains between the two static controls, showing that the throughput and TTFT gains do not come from severely degraded steady-state generation.

Which issue(s) this PR fixes:

Fixes #TODO

Test plan:

• Verify load preserve, disable, and reusable-token threshold behavior.
• Verify waiting-queue EWMA close, hysteresis, and missing-metric behavior.
• Verify independent load and offload decisions and preservation of client tier filters.
• Verify unsupported request shapes fail open.
• Run go test ./pkg/epp/framework/plugins/requestcontrol/selectivekv.

AI assistance from Codex was used to inspect the implementation, update tests and documentation.

Release note:

Adds an Alpha EPP request-control plugin for per-request selective KV loading and offloading.

[vMaroon]

Awesome work, will review asap

cc @nilig @ahg-g @kaushikmitr

[albertoperdomo2]

I also validated the proposed minExternalReusableTokens: 1024 plus maxWaitingRequests: 8 policy in low-pressure conditions. I tested an 8K reusable prefix at concurrency 4 to verify that the policy does not penalize loading when the system is healthy. Compared with current unconditional loading, the policy produced equivalent results: +1.34% output throughput, −3.80% mean TTFT, −2.98% p99 TTFT, and −1.02% mean ITL. External KV transfers remained active and the sampled waiting queue stayed below the configured threshold.

[Copilot]

🟡 Changes recommended

Unresolved payload-handling and queue-gating correctness issues remain.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds an Alpha EPP plugin for selective per-request KV loading and offloading.

Changes:

• Implements preserve, disable, and threshold-based policies.
• Adds queue-pressure EWMA gating, tests, and documentation.
• Registers the plugin with the EPP.

File summaries

File	Description
pkg/epp/framework/plugins/requestcontrol/selectivekv/README.md	Documents configuration and behavior.
pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go	Implements selective KV policy logic.
pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin_test.go	Tests configuration and policy behavior.
cmd/epp/runner/runner.go	Registers the Alpha plugin.

Review details

Suppressed comments (6)

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:289

• Once a high sample closes the gate, this code only updates state.ewma when a strictly newer UpdateTime is observed. During a metrics scrape outage, requests keep reusing the old timestamp and closed state, so external loading can remain disabled indefinitely instead of failing open. Apply a freshness bound to metrics.UpdateTime and reset or skip the veto for stale samples, consistent with the existing metrics staleness handling.

	if !found || sampleTime.Before(state.sampleTime) {
		state = waitingQueueState{ewma: sample, sampleTime: sampleTime}
	} else if sampleTime.After(state.sampleTime) {
		elapsed := sampleTime.Sub(state.sampleTime)
		weight := 1 - math.Exp(-math.Ln2*
			elapsed.Seconds()/waitingQueueEWMAHalfLife.Seconds())
		state.ewma += weight * (sample - state.ewma)
		state.sampleTime = sampleTime
	}

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:334

• The documented fail-open behavior for invalid tier evidence is not implemented here: negative counts are accepted. For example, a negative GPU count increases maxExternalBlocks-gpuBlocks and can enable loading for an invalid sample, while a negative external count can be converted into a valid zero. Reject any negative tier count, including the GPU and speculative entries, before calculating the token count and return (0, false).

	gpuBlocks := byTier[gpuTierKey]
	maxExternalBlocks := 0
	for tier, blocks := range byTier {
		if tier == gpuTierKey || tier == attrprefix.SpeculativeTierKey {
			continue
		}
		maxExternalBlocks = max(maxExternalBlocks, blocks)

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:225

• The closure can make no payload change when threshold loading is allowed and offloading is preserved, especially for requests without kv_transfer_params, but MutatePayloadMap still sets Mutated=true. Director.repackage then marshals every such request instead of forwarding the original body, adding avoidable per-request serialization and potentially changing the wire representation. Track whether a controlled field actually changed and mark the body mutated only for real edits.

	request.Body.MutatePayloadMap(func(payload requesthandling.PayloadMap) {
		params, _ := payload[commonrequest.FieldKVTransferParams].(map[string]any)
		if params == nil {
			if load != loadDisable && p.offloadPolicy != PolicyDisable {
				return
			}
			params = map[string]any{}
			payload[commonrequest.FieldKVTransferParams] = params
		}
		switch load {
		case loadDisable:
			params[fieldMaxLoadTokens] = json.Number("0")
		case loadEnable:
			delete(params, fieldMaxLoadTokens)
		}
		if p.offloadPolicy == PolicyDisable {
			params[fieldMaxOffloadTokens] = json.Number("0")
		}
	})

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:334

• The per-tier counts are summed across prompts before this function runs (preciseprefixcache/producer.go:addPodMatch). Taking the max after that sum is not equivalent to summing each prompt's longest external tier: {gpu:10,cpu:20} plus {gpu:0,storage:20} becomes 10 external blocks here, although the request has 30 blocks outside GPU. Batched requests can therefore be forced to recompute despite meeting the threshold; publish an aggregate external count or retain per-prompt tier data.

	for tier, blocks := range byTier {
		if tier == gpuTierKey || tier == attrprefix.SpeculativeTierKey {
			continue
		}
		maxExternalBlocks = max(maxExternalBlocks, blocks)

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:245

• The decision ignores the client's kv_load_tiers, even though vLLM uses that field to filter secondary tiers while CPU remains a direct source. A request that excludes storage can therefore pass this threshold from a large storage-only match and remove max_load_tokens, despite having too few permitted external tokens to justify loading. Evaluate the threshold against the tiers permitted by the request, or explicitly fail open when the filter cannot be interpreted.

		tokens, ok := p.externalReusableTokens(endpoint)
		if !ok {
			logger.Info("no tier evidence; preserving KV load")
			return loadPreserve
		}

pkg/epp/framework/plugins/requestcontrol/selectivekv/plugin.go:250

• When tokens < minExternalReusableTokens, this branch skips waitingQueueVeto, so those requests never contribute samples to the endpoint EWMA. The first qualifying long-prefix request can therefore bypass the configured smoothing and hysteresis after pressure was observed only during shorter requests. Update the EWMA for every valid endpoint sample before combining it with the token-threshold decision.

		action := loadEnable
		if tokens < p.minExternalReusableTokens {
			action = loadDisable
		} else if p.waitingQueueVeto(endpoint) {
			action = loadDisable

• Files reviewed: 4/4 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[albertoperdomo2]

This is a wider issue that will be addressed in another PR (See #2890)

[ahg-g]

This is very nice! What are the next steps? is this ready to review?

[albertoperdomo2]

Tested the latest EPP changes with:

• EPP: quay.io/rh-ee-aperdomo/llm-d-router-endpoint-picker:dev-1b096e59
• vLLM: quay.io/rh-ee-aperdomo/vllm:v0.29.0-max-load-tokens, based on vLLM v0.29.0 with max_load_tokens support

Ready for review! @ahg-g

[ahg-g]

Nice work, and the benchmark write-up is thorough. Build, vet and go test -race pass locally after merging main (92.8% coverage), and the earlier Copilot findings on RawMessage handling and per-endpoint state cleanup are addressed.

1. gpuTierKey = "gpu" is the third local copy of that tier name (kvevents.defaultEventSourceDeviceTier, kvcache/backend.go). Could it use a shared exported constant instead?
2. In the README, please state the minimum vLLM release containing #55885, and what happens on older engines (is max_load_tokens ignored, making the policy a silent no-op?).
3. Nit: the YAML in the description has maxWaitingReques. Worth correcting so the tested config is unambiguous.

[albertoperdomo2]

thanks @ahg-g, I just addressed the comments 👍🏽