986e217158 · tduncan/Nvidia-Optimization-Stack-VLLM

rust/proto/inference.proto

+15

@@ -56,6 +56,21 @@ message GenerateRequest {

56

57 // Apply the engine's configured watermark. Defaults to true when omitted.

58 optional bool watermarking = 16;

59 +

60 + KVHintsEnvelope kv_hints = 17;

61 +}

62 +

63 +message KVHintsEnvelope {

64 + string protocol_version = 1;

65 + string message_id = 2;

66 + repeated KVHintAction actions = 3;

67 +}

68 +

69 +message KVHintAction {

70 + string action_id = 1;

71 + string action_type = 2;

72 + string action_version = 3;

73 + google.protobuf.Struct payload = 4;

74 }

75

76 message RandomSampling {

rust/src/chat/src/lib.rs

+1

@@ -194,6 +194,7 @@ impl ChatRequestProcessor {

194 add_special_tokens: request.add_special_tokens,

195 data_parallel_rank: request.data_parallel_rank,

196 session_id: request.session_id,

197 + kv_hints: None,

198 reasoning_parser_kwargs,

199 lora_request: request.lora_request,

200 arrival_time: Some(arrival_time),

rust/src/engine-core-client/src/protocol/kv_hints.rs new

+21

@@ -0,0 +1,21 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +use std::collections::BTreeMap;

5 +

6 +use serde::{Deserialize, Serialize};

7 +

8 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

9 +pub struct KvHintAction {

10 + pub action_id: String,

11 + pub action_type: String,

12 + pub action_version: String,

13 + pub payload: BTreeMap<String, serde_json::Value>,

14 +}

15 +

16 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

17 +pub struct KvHintsEnvelope {

18 + pub protocol_version: String,

19 + pub message_id: String,

20 + pub actions: Vec<KvHintAction>,

21 +}

rust/src/engine-core-client/src/protocol/mod.rs

+1

@@ -24,6 +24,7 @@ fn is_false(v: &bool) -> bool {

24

25 pub mod dtype;

26 pub mod handshake;

27 +pub mod kv_hints;

28 pub mod logprobs;

29 pub mod lora;

30 pub mod multimodal;

rust/src/engine-core-client/src/protocol/request.rs

+16 -1

@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};

8 use serde_default::DefaultFromSerde;

9 use serde_tuple::{Deserialize_tuple, Serialize_tuple};

10

11 +use crate::protocol::kv_hints::KvHintsEnvelope;

12 use crate::protocol::multimodal::MmFeatures;

13 use crate::protocol::sampling::EngineCoreSamplingParams;

14 use crate::protocol::{OpaqueValue, lora};

@@ -127,6 +128,8 @@ pub struct EngineCoreRequest {

128 /// Stable session identity shared by related requests.

129 #[serde(default)]

130 pub session_id: Option<String>,

131 + #[serde(default)]

132 + pub kv_hints: Option<KvHintsEnvelope>,

133 }

134

135 impl EngineCoreRequest {

@@ -159,6 +162,7 @@ mod tests {

162

163 use super::*;

164 use crate::protocol::dtype::TensorDtype;

165 + use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};

166 use crate::protocol::multimodal::{

167 MmBatchedField, MmFeatureSpec, MmField, MmFieldElem, MmKwargValue, MmModality,

168 PlaceholderRange,

@@ -181,6 +185,16 @@ mod tests {

185 arrival_time: 1234.5,

186 client_index: 7,

187 session_id: Some("session-1".to_string()),

188 + kv_hints: Some(KvHintsEnvelope {

189 + protocol_version: "0.1".to_string(),

190 + message_id: "msg-1".to_string(),

191 + actions: vec![KvHintAction {

192 + action_id: "action-1".to_string(),

193 + action_type: "example.action".to_string(),

194 + action_version: "1.0".to_string(),

195 + payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]),

196 + }],

197 + }),

198 ..EngineCoreRequest::default()

199 };

200

@@ -191,13 +205,14 @@ mod tests {

205 other => panic!("expected array, got {other:?}"),

206 };

207

194 - assert_eq!(array.len(), 21);

208 + assert_eq!(array.len(), 22);

209 assert_eq!(array[0], Value::from("req-1"));

210 assert_eq!(array[2], Value::Nil);

211 assert_eq!(array[4], Value::Nil);

212 assert_eq!(array[10], Value::Nil);

213 assert_eq!(array[11], Value::from(7));

214 assert_eq!(array[20], Value::from("session-1"));

215 + assert!(matches!(&array[21], Value::Map(_)));

216 }

217

218 #[test]

rust/src/engine-core-client/src/tests/client.rs

+11

@@ -20,6 +20,7 @@ use zeromq::util::PeerIdentity;

20 use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage};

21

22 use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};

23 +use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};

24 use crate::protocol::logprobs::MaybeWireLogprobs;

25 use crate::protocol::multimodal::{

26 MmFeatureSpec, MmField, MmFieldElem, MmFlatField, MmKwargValue, MmModality, MmSlice,

@@ -165,6 +166,16 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {

166 }),

167 arrival_time: 42.5,

168 session_id: Some("session-1".to_string()),

169 + kv_hints: Some(KvHintsEnvelope {

170 + protocol_version: "0.1".to_string(),

171 + message_id: "msg-1".to_string(),

172 + actions: vec![KvHintAction {

173 + action_id: "action-1".to_string(),

174 + action_type: "example.action".to_string(),

175 + action_version: "1.0".to_string(),

176 + payload: BTreeMap::from([("key".to_string(), serde_json::json!("value"))]),

177 + }],

178 + }),

179 ..EngineCoreRequest::default()

180 }

181 }

rust/src/engine-core-client/src/tests/python_compat.py

+26

@@ -52,6 +52,19 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):

52 output_kind: RequestOutputKind = RequestOutputKind.DELTA

53

54

55 +class KvHintAction(msgspec.Struct, frozen=True):

56 + action_id: str

57 + action_type: str

58 + action_version: str

59 + payload: dict[str, object]

60 +

61 +

62 +class KvHintsEnvelope(msgspec.Struct, frozen=True):

63 + protocol_version: str

64 + message_id: str

65 + actions: list[KvHintAction]

66 +

67 +

68 class EngineCoreRequest(

69 msgspec.Struct,

70 array_like=True,

@@ -78,6 +91,7 @@ class EngineCoreRequest(

91 reasoning_parser_kwargs: dict[str, object] | None = None

92 abort_immediately: bool = False

93 session_id: str | None = None

94 + kv_hints: KvHintsEnvelope | None = None

95

96

97 class EngineCoreOutput(

@@ -153,6 +167,18 @@ request = EngineCoreRequest(

167 arrival_time=42.5,

168 client_index=0,

169 session_id="session-1",

170 + kv_hints=KvHintsEnvelope(

171 + protocol_version="0.1",

172 + message_id="msg-1",

173 + actions=[

174 + KvHintAction(

175 + action_id="action-1",

176 + action_type="example.action",

177 + action_version="1.0",

178 + payload={"key": "value"},

179 + )

180 + ],

181 + ),

182 )

183

184 # All defaults -> empty map. Regression guard for the sparse-map decode.

rust/src/llm/examples/external_engine_smoke.rs

+1

@@ -60,6 +60,7 @@ fn build_request(request_id: String, max_tokens: u32) -> GenerateRequest {

60 priority: 0,

61 data_parallel_rank: None,

62 session_id: None,

63 + kv_hints: None,

64 reasoning_parser_kwargs: None,

65 lora_request: None,

66 }

rust/src/llm/src/request.rs

+20

@@ -4,6 +4,7 @@

4 use std::collections::BTreeMap;

5

6 use uuid::Uuid;

7 +use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope;

8 use vllm_engine_core_client::protocol::lora::LoraRequest;

9 use vllm_engine_core_client::protocol::multimodal::MmFeatures;

10 use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningParserKwargs};

@@ -48,6 +49,8 @@ pub struct GenerateRequest {

49 pub data_parallel_rank: Option<u32>,

50 /// Stable session identity shared by related requests.

51 pub session_id: Option<String>,

52 + /// Optional orchestrator-originated KV hints.

53 + pub kv_hints: Option<KvHintsEnvelope>,

54 /// Optional reasoning-parser kwargs forwarded to engine-side structured

55 /// output logic.

56 pub reasoning_parser_kwargs: Option<ReasoningParserKwargs>,

@@ -79,6 +82,7 @@ impl GenerateRequest {

82 priority,

83 data_parallel_rank,

84 session_id,

85 + kv_hints,

86 reasoning_parser_kwargs,

87 lora_request,

88 } = self;

@@ -109,6 +113,7 @@ impl GenerateRequest {

113 trace_headers,

114 resumable: false,

115 session_id,

116 + kv_hints,

117 external_req_id: Some(external_request_id),

118 // Rust parser doesn't expose this information, leave it unset and let the

119 // reasoning logic in engine-sided structured output manager handle it.

@@ -134,6 +139,7 @@ impl PreparedGenerateRequest {

139 mod tests {

140 use std::collections::BTreeMap;

141

142 + use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};

143 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;

144 use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;

145

@@ -155,6 +161,16 @@ mod tests {

161 priority: 3,

162 data_parallel_rank: Some(2),

163 session_id: Some("session-1".to_string()),

164 + kv_hints: Some(KvHintsEnvelope {

165 + protocol_version: "0.1".to_string(),

166 + message_id: "msg-1".to_string(),

167 + actions: vec![KvHintAction {

168 + action_id: "action-1".to_string(),

169 + action_type: "example.action".to_string(),

170 + action_version: "1.0".to_string(),

171 + payload: BTreeMap::new(),

172 + }],

173 + }),

174 reasoning_parser_kwargs: Some(ReasoningParserKwargs {

175 chat_template_kwargs: [(

176 "chat_template_kwargs".to_string(),

@@ -183,6 +199,10 @@ mod tests {

199 assert_eq!(request.cache_salt.as_deref(), Some("salt"));

200 assert_eq!(request.data_parallel_rank, Some(2));

201 assert_eq!(request.session_id.as_deref(), Some("session-1"));

202 + assert_eq!(

203 + request.kv_hints.as_ref().map(|hints| hints.message_id.as_str()),

204 + Some("msg-1")

205 + );

206 assert_eq!(

207 request.trace_headers,

208 Some(BTreeMap::from([(

rust/src/llm/tests/generate.rs

+1

@@ -167,6 +167,7 @@ fn sample_generate_request(request_id: &str, max_tokens: u32) -> GenerateRequest

167 priority: 0,

168 data_parallel_rank: None,

169 session_id: None,

170 + kv_hints: None,

171 reasoning_parser_kwargs: None,

172 lora_request: None,

173 }

rust/src/server/src/grpc/convert.rs

+61 -1

@@ -7,6 +7,7 @@

7 use thiserror_ext::AsReport as _;

8 use tonic::Status;

9 use uuid::Uuid;

10 +use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};

11 use vllm_engine_core_client::protocol::output::StopReason;

12 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;

13 use vllm_text::{

@@ -16,6 +17,32 @@ use vllm_text::{

17

18 use super::pb;

19

20 +fn kv_hints_from_proto(hints: pb::KvHintsEnvelope) -> KvHintsEnvelope {

21 + KvHintsEnvelope {

22 + protocol_version: hints.protocol_version,

23 + message_id: hints.message_id,

24 + actions: hints

25 + .actions

26 + .into_iter()

27 + .map(|action| KvHintAction {

28 + action_id: action.action_id,

29 + action_type: action.action_type,

30 + action_version: action.action_version,

31 + payload: action

32 + .payload

33 + .map(|payload| {

34 + payload

35 + .fields

36 + .into_iter()

37 + .map(|(key, value)| (key, proto_value_to_json(&value)))

38 + .collect()

39 + })

40 + .unwrap_or_default(),

41 + })

42 + .collect(),

43 + }

44 +}

45 +

46 // ========================================================================================

47 // Request conversion

48 // ========================================================================================

@@ -56,6 +83,7 @@ pub fn to_text_request(

83 req.request_id

84 };

85 let session_id = req.session_id.filter(|s| !s.is_empty());

86 + let kv_hints = req.kv_hints.map(kv_hints_from_proto);

87

88 let sampling = req.sampling.as_ref();

89 let decoding = req.decoding.as_ref();

@@ -108,6 +136,7 @@ pub fn to_text_request(

136 add_special_tokens: true,

137 data_parallel_rank: None,

138 session_id,

139 + kv_hints,

140 reasoning_parser_kwargs: None,

141 lora_request: None,

142 arrival_time: None,

@@ -537,7 +566,9 @@ mod tests {

566 use vllm_tokenizer::test_utils::TestTokenizer;

567

568 use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason};

540 - use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request};

569 + use super::{

570 + ResponseOpts, json_to_proto_struct, pb, to_finish_info, to_sequence_output, to_text_request,

571 + };

572

573 fn base_request() -> pb::GenerateRequest {

574 pb::GenerateRequest {

@@ -578,6 +609,35 @@ mod tests {

609 assert_eq!(text.sampling_params.temperature, Some(0.7));

610 }

611

612 + #[test]

613 + fn kv_hints_propagate_from_top_level_request_field() {

614 + let req = pb::GenerateRequest {

615 + kv_hints: Some(pb::KvHintsEnvelope {

616 + protocol_version: "0.1".to_string(),

617 + message_id: "msg-1".to_string(),

618 + actions: vec![pb::KvHintAction {

619 + action_id: "action-1".to_string(),

620 + action_type: "example.action".to_string(),

621 + action_version: "1.0".to_string(),

622 + payload: json_to_proto_struct(&serde_json::json!({"source": "g2"})),

623 + }],

624 + }),

625 + ..base_request()

626 + };

627 +

628 + let text = to_text_request(req, false, &["test-model".to_string()]).expect("convert ok");

629 + let hints = text.kv_hints.expect("kv hints");

630 + assert_eq!(hints.protocol_version, "0.1");

631 + assert_eq!(hints.message_id, "msg-1");

632 + assert_eq!(hints.actions[0].action_id, "action-1");

633 + assert_eq!(hints.actions[0].action_type, "example.action");

634 + assert_eq!(hints.actions[0].action_version, "1.0");

635 + assert_eq!(

636 + hints.actions[0].payload.get("source"),

637 + Some(&serde_json::json!("g2"))

638 + );

639 + }

640 +

641 #[test]

642 fn unset_temperature_defaults_to_greedy() {

643 let text = to_text_request(base_request(), false, &["test-model".to_string()])

rust/src/server/src/routes/inference/generate/convert.rs

+1

@@ -82,6 +82,7 @@ pub(super) fn prepare_generate_request(

82 add_special_tokens: false,

83 data_parallel_rank: ctx.data_parallel_rank,

84 session_id: ctx.session_id,

85 + kv_hints: None,

86 reasoning_parser_kwargs: None,

87 lora_request: lora_resolution.lora_request.clone(),

88 arrival_time: None,

rust/src/server/src/routes/openai/completions/convert.rs

+1

@@ -155,6 +155,7 @@ pub(super) fn prepare_completion_request(

155 add_special_tokens: request.add_special_tokens,

156 data_parallel_rank: ctx.data_parallel_rank,

157 session_id,

158 + kv_hints: None,

159 reasoning_parser_kwargs: None,

160 lora_request: lora_resolution.lora_request.clone(),

161 arrival_time: None,

rust/src/server/src/routes/tokenize/types.rs

+1

@@ -54,6 +54,7 @@ impl TokenizeCompletionRequest {

54 add_special_tokens: self.add_special_tokens,

55 data_parallel_rank: None,

56 session_id: None,

57 + kv_hints: None,

58 reasoning_parser_kwargs: None,

59 lora_request: None,

60 arrival_time: None,

rust/src/text/src/lower.rs

+31

@@ -60,6 +60,7 @@ pub fn lower_text_request(

60 priority: request.priority,

61 data_parallel_rank: request.data_parallel_rank,

62 session_id: request.session_id.clone(),

63 + kv_hints: request.kv_hints.clone(),

64 reasoning_parser_kwargs: request.reasoning_parser_kwargs.clone(),

65 lora_request: request.lora_request.clone(),

66 arrival_time: request.arrival_time,

@@ -320,6 +321,7 @@ mod tests {

321 use std::collections::{BTreeSet, HashMap};

322

323 use serial_test::file_serial;

324 + use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope};

325 use vllm_engine_core_client::protocol::multimodal::{

326 MmFeatureSpec, MmModality, PlaceholderRange,

327 };

@@ -1301,6 +1303,35 @@ mod tests {

1303 assert_eq!(prepared.generate_request.request_id, "text-1");

1304 }

1305

1306 + #[test]

1307 + fn lower_text_request_passes_kv_hints_through() {

1308 + let hints = KvHintsEnvelope {

1309 + protocol_version: "0.1".to_string(),

1310 + message_id: "msg-1".to_string(),

1311 + actions: vec![KvHintAction {

1312 + action_id: "action-1".to_string(),

1313 + action_type: "example.action".to_string(),

1314 + action_version: "1.0".to_string(),

1315 + payload: Default::default(),

1316 + }],

1317 + };

1318 + let request = TextRequest {

1319 + kv_hints: Some(hints.clone()),

1320 + ..sample_request()

1321 + };

1322 +

1323 + let prepared = lower_text_request(

1324 + request,

1325 + vec![1, 2, 3],

1326 + sample_sampling_hints(),

1327 + sample_sampling_limits(),

1328 + &stub_tokenizer(),

1329 + )

1330 + .unwrap();

1331 +

1332 + assert_eq!(prepared.generate_request.kv_hints, Some(hints));

1333 + }

1334 +

1335 #[test]

1336 fn lower_text_request_passes_arrival_time_through() {

1337 let request = TextRequest {

rust/src/text/src/request.rs

+5

@@ -6,6 +6,7 @@ use std::collections::HashMap;

6 use enum_as_inner::EnumAsInner;

7 use serde::{Deserialize, Serialize};

8 use serde_json::Value;

9 +use vllm_engine_core_client::protocol::kv_hints::KvHintsEnvelope;

10 use vllm_engine_core_client::protocol::lora::LoraRequest;

11 use vllm_engine_core_client::protocol::multimodal::MmFeatures;

12 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;

@@ -211,6 +212,9 @@ pub struct TextRequest {

212 /// Stable session identity shared by related requests.

213 #[serde(default)]

214 pub session_id: Option<String>,

215 + /// Optional orchestrator-originated KV hints.

216 + #[serde(default)]

217 + pub kv_hints: Option<KvHintsEnvelope>,

218 /// Optional reasoning-parser kwargs forwarded to engine-side structured

219 /// output logic.

220 #[serde(default)]

@@ -242,6 +246,7 @@ impl TextRequest {

246 add_special_tokens: false,

247 data_parallel_rank: None,

248 session_id: None,

249 + kv_hints: None,

250 reasoning_parser_kwargs: None,

251 lora_request: None,

252 arrival_time: None,

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+18

@@ -33,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

33 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

34 OffloadingConnectorScheduler,

35 RequestOffloadState,

36 + _create_req_context,

37 get_sliding_window_size_in_chunks,

38 )

39 from vllm.v1.core.block_pool import BlockPool

@@ -60,6 +61,7 @@ from vllm.v1.kv_cache_interface import (

61 SlidingWindowSpec,

62 )

63 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

64 +from vllm.v1.kv_hints import KvHintsEnvelope

65 from vllm.v1.kv_offload.base import (

66 GPULoadStoreSpec,

67 LookupResult,

@@ -214,6 +216,7 @@ def _make_partial_tail_request(

216 request = MagicMock()

217 request.request_id = "req"

218 request.kv_transfer_params = kv_transfer_params

219 + request.kv_hints = None

220 request.num_prompt_tokens = 30

221 request.num_tokens = 30

222 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]

@@ -225,6 +228,21 @@ def _make_partial_tail_request(

228 return request

229

230

231 +def test_create_req_context_preserves_kv_hints():

232 + kv_hints = KvHintsEnvelope(

233 + protocol_version="0.1",

234 + message_id="message-1",

235 + actions=[],

236 + )

237 + request = MagicMock(

238 + request_id="req",

239 + kv_transfer_params=None,

240 + kv_hints=kv_hints,

241 + )

242 +

243 + assert _create_req_context(request).kv_hints is kv_hints

244 +

245 +

246 def _reduce_kv_connector_stats(runner):

247 reduced: dict[str, int | float] = {}

248 for payload in runner.kv_connector_stats:

tests/v1/kv_offload/tiering/test_kvcr_tier.py

+8 -4

@@ -5,12 +5,13 @@ from collections.abc import Collection, Iterable, Mapping

5 from types import SimpleNamespace

6 from typing import Any

7

8 +import msgspec

9 import numpy as np

10 import pytest

11

12 pytest.importorskip("kvcr")

13

13 -from kvcr import ROUTER_HINT_KEY, KVCRBindings

14 +from kvcr import KVCRBindings

15 from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig

16 from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy

17 from kvcr.types import (

@@ -27,6 +28,7 @@ from kvcr.types import (

28 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

29 OffloadingConnectorStats,

30 )

31 +from vllm.v1.kv_hints import KvHintsEnvelope

32 from vllm.v1.kv_offload.base import (

33 LookupResult,

34 Medium,

@@ -271,10 +273,10 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch):

273 """Check hint forwarding, key conversion, load descriptors, and cleanup."""

274 kvcr = RecordingKVCR()

275 tier = _make_tier(monkeypatch, kvcr)

274 - router_hint = {"opaque": True}

276 + kv_hint = KvHintsEnvelope(protocol_version="0.1", message_id="msg", actions=[])

277 ctx = ReqContext(

278 req_id="req",

277 - kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()},

279 + kv_hints=kv_hint,

280 )

281 key = make_offload_key((123).to_bytes(8, "big"), 0)

282 same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7)

@@ -282,7 +284,9 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch):

284

285 tier.on_new_request(ctx)

286

285 - assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})]

287 + assert kvcr.submit_hint_calls == [

288 + ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)})

289 + ]

290

291 bindings = kvcr.constructor_bindings

292 assert bindings is not None

tests/v1/test_request.py

+34

@@ -3,6 +3,7 @@

3

4 from vllm import SamplingParams

5 from vllm.v1.engine import EngineCoreRequest

6 +from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope

7 from vllm.v1.request import Request, RequestStatus

8

9

@@ -24,6 +25,7 @@ def test_request_status_fmt_str():

25

26

27 def test_request_copies_session_id_from_engine_core_request():

28 + """Test that Request preserves session ID from EngineCoreRequest."""

29 engine_request = EngineCoreRequest(

30 request_id="request-1",

31 prompt_token_ids=[1, 2, 3],

@@ -40,3 +42,35 @@ def test_request_copies_session_id_from_engine_core_request():

42 request = Request.from_engine_core_request(engine_request, block_hasher=None)

43

44 assert request.session_id == "session-1"

45 +

46 +

47 +def test_request_copies_kv_hints_from_engine_core_request():

48 + """Test that Request preserves KV hints from EngineCoreRequest."""

49 + kv_hints = KvHintsEnvelope(

50 + protocol_version="0.1",

51 + message_id="msg-1",

52 + actions=[

53 + KvHintAction(

54 + action_id="action-1",

55 + action_type="example.action",

56 + action_version="1.0",

57 + payload={},

58 + )

59 + ],

60 + )

61 + engine_request = EngineCoreRequest(

62 + request_id="request-1",

63 + prompt_token_ids=[1, 2, 3],

64 + mm_features=None,

65 + sampling_params=SamplingParams(max_tokens=1),

66 + pooling_params=None,

67 + arrival_time=0.0,

68 + lora_request=None,

69 + cache_salt=None,

70 + data_parallel_rank=None,

71 + kv_hints=kv_hints,

72 + )

73 +

74 + request = Request.from_engine_core_request(engine_request, block_hasher=None)

75 +

76 + assert request.kv_hints == kv_hints

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+1

@@ -553,6 +553,7 @@ def _create_req_context(req: Request) -> ReqContext:

553 return ReqContext(

554 req_id=req.request_id,

555 kv_transfer_params=params,

556 + kv_hints=req.kv_hints,

557 load_tier_filter=load_filter,

558 )

559

vllm/engine/protocol.py

+2

@@ -22,6 +22,7 @@ from vllm.tasks import SupportedTask

22 from vllm.v1.engine import EngineCoreRequest

23 from vllm.v1.engine.input_processor import InputProcessor

24 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult

25 +from vllm.v1.kv_hints import KvHintsEnvelope

26

27 if TYPE_CHECKING:

28 from vllm.v1.engine import PauseMode

@@ -102,6 +103,7 @@ class EngineClient(ABC):

103 priority: int = 0,

104 data_parallel_rank: int | None = None,

105 session_id: str | None = None,

106 + kv_hints: KvHintsEnvelope | None = None,

107 reasoning_ended: bool | None = None,

108 reasoning_parser_kwargs: dict[str, Any] | None = None,

109 ) -> AsyncGenerator[RequestOutput, None]:

vllm/v1/engine/__init__.py

+2

@@ -16,6 +16,7 @@ from vllm.lora.request import LoRARequest

16 from vllm.multimodal.inputs import MultiModalFeatureSpec

17 from vllm.pooling_params import PoolingParams

18 from vllm.sampling_params import SamplingParams

19 +from vllm.v1.kv_hints import KvHintsEnvelope

20 from vllm.v1.metrics.stats import (

21 PrefillStats,

22 RequestSpecDecodeMetrics,

@@ -157,6 +158,7 @@ class EngineCoreRequest(

158 abort_immediately: bool = False

159

160 session_id: str | None = None

161 + kv_hints: KvHintsEnvelope | None = None

162

163 @property

164 def params(self) -> SamplingParams | PoolingParams:

vllm/v1/engine/async_llm.py

+9

@@ -53,6 +53,7 @@ from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollec

53 from vllm.v1.engine.parallel_sampling import ParentRequest

54 from vllm.v1.executor import Executor

55 from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult

56 +from vllm.v1.kv_hints import KvHintsEnvelope

57 from vllm.v1.metrics.loggers import (

58 StatLoggerFactory,

59 StatLoggerManager,

@@ -386,6 +387,7 @@ class AsyncLLM(EngineClient):

387 prompt_text: str | None = None,

388 reasoning_ended: bool | None = None,

389 reasoning_parser_kwargs: dict[str, Any] | None = None,

390 + kv_hints: KvHintsEnvelope | None = None,

391 ) -> RequestOutputCollector:

392 """Add new request to the AsyncLLM."""

393 if self.errored:

@@ -420,6 +422,7 @@ class AsyncLLM(EngineClient):

422 priority,

423 data_parallel_rank,

424 session_id,

425 + kv_hints,

426 )

427

428 # Convert Input --> Request.

@@ -453,6 +456,7 @@ class AsyncLLM(EngineClient):

456 priority=priority,

457 data_parallel_rank=data_parallel_rank,

458 session_id=session_id,

459 + kv_hints=kv_hints,

460 )

461 else:

462 # Raw prompts require tokenization and possibly multimodal

@@ -469,6 +473,7 @@ class AsyncLLM(EngineClient):

473 priority=priority,

474 data_parallel_rank=data_parallel_rank,

475 session_id=session_id,

476 + kv_hints=kv_hints,

477 )

478 prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)

479

@@ -557,6 +562,7 @@ class AsyncLLM(EngineClient):

562 priority: int = 0,

563 data_parallel_rank: int | None = None,

564 session_id: str | None = None,

565 + kv_hints: KvHintsEnvelope | None = None,

566 ) -> RequestOutputCollector:

567 self._validate_streaming_input_sampling_params(sampling_params)

568

@@ -569,6 +575,7 @@ class AsyncLLM(EngineClient):

575 priority=priority,

576 data_parallel_rank=data_parallel_rank,

577 session_id=session_id,

578 + kv_hints=kv_hints,

579 )

580

581 if not sampling_params.skip_clone:

@@ -670,6 +677,7 @@ class AsyncLLM(EngineClient):

677 priority: int = 0,

678 data_parallel_rank: int | None = None,

679 session_id: str | None = None,

680 + kv_hints: KvHintsEnvelope | None = None,

681 reasoning_ended: bool | None = None,

682 reasoning_parser_kwargs: dict[str, Any] | None = None,

683 ) -> AsyncGenerator[RequestOutput, None]:

@@ -710,6 +718,7 @@ class AsyncLLM(EngineClient):

718 priority=priority,

719 data_parallel_rank=data_parallel_rank,

720 session_id=session_id,

721 + kv_hints=kv_hints,

722 prompt_text=prompt_text,

723 reasoning_ended=reasoning_ended,

724 reasoning_parser_kwargs=reasoning_parser_kwargs,

vllm/v1/engine/input_processor.py

+3

@@ -32,6 +32,7 @@ from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid

32 from vllm.utils.async_utils import make_async

33 from vllm.utils.jsontree import json_iter_leaves

34 from vllm.v1.engine import EngineCoreRequest

35 +from vllm.v1.kv_hints import KvHintsEnvelope

36

37 logger = init_logger(__name__)

38

@@ -318,6 +319,7 @@ class InputProcessor:

319 data_parallel_rank: int | None = None,

320 resumable: bool = False,

321 session_id: str | None = None,

322 + kv_hints: KvHintsEnvelope | None = None,

323 ) -> EngineCoreRequest:

324 self._validate_params(params, supported_tasks)

325 self._validate_lora(lora_request)

@@ -461,6 +463,7 @@ class InputProcessor:

463 trace_headers=trace_headers,

464 resumable=resumable,

465 session_id=session_id,

466 + kv_hints=kv_hints,

467 )

468

469 def _validate_prompt_len(

vllm/v1/engine/llm_engine.py

+3

@@ -34,6 +34,7 @@ from vllm.v1.engine.input_processor import InputProcessor

34 from vllm.v1.engine.output_processor import OutputProcessor

35 from vllm.v1.engine.parallel_sampling import ParentRequest

36 from vllm.v1.executor import Executor

37 +from vllm.v1.kv_hints import KvHintsEnvelope

38 from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager

39 from vllm.v1.metrics.reader import Metric, get_metrics_snapshot

40 from vllm.v1.metrics.stats import IterationStats

@@ -230,6 +231,7 @@ class LLMEngine:

231 priority: int = 0,

232 session_id: str | None = None,

233 prompt_text: str | None = None,

234 + kv_hints: KvHintsEnvelope | None = None,

235 ) -> str:

236 # Validate the request_id type.

237 if not isinstance(request_id, str):

@@ -263,6 +265,7 @@ class LLMEngine:

265 trace_headers=trace_headers,

266 priority=priority,

267 session_id=session_id,

268 + kv_hints=kv_hints,

269 )

270 prompt_text, _, _ = extract_prompt_components(self.model_config, prompt)

271

vllm/v1/kv_hints/__init__.py new

+6

@@ -0,0 +1,6 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from vllm.v1.kv_hints.protocol import KvHintAction, KvHintsEnvelope

5 +

6 +__all__ = ["KvHintAction", "KvHintsEnvelope"]

vllm/v1/kv_hints/protocol.py new

+26

@@ -0,0 +1,26 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from typing import Any

5 +

6 +import msgspec

7 +

8 +

9 +class KvHintAction(msgspec.Struct, frozen=True): # type: ignore[call-arg]

10 + """A versioned orchestrator-provided KV hint action."""

11 +

12 + action_id: str

13 + action_type: str

14 + action_version: str

15 + # `payload` must contain JSON-compatible values. Action-specific schemas are

16 + # validated by the component implementing the action.

17 + payload: dict[str, Any]

18 +

19 +

20 +class KvHintsEnvelope(msgspec.Struct, frozen=True): # type: ignore[call-arg]

21 + """Orchestrator-provided KV hints optionally attached to one inference request.

22 + A KV hint envelope is a collection of KV hint action(s)."""

23 +

24 + protocol_version: str

25 + message_id: str

26 + actions: list[KvHintAction]

vllm/v1/kv_offload/base.py

+4 -2

@@ -16,6 +16,7 @@ if TYPE_CHECKING:

16 OffloadingConnectorStats,

17 )

18

19 +from vllm.v1.kv_hints import KvHintsEnvelope

20 from vllm.v1.kv_offload.config import OffloadingConfig

21

22 # `OffloadKey` identifies an offloaded block. It combines a block hash with

@@ -89,10 +90,11 @@ TierFilter.ALL = TierFilter(matchers=(TierMatcher(),))

90 class ReqContext:

91 req_id: str

92 kv_transfer_params: dict[str, Any] | None = None

93 + kv_hints: KvHintsEnvelope | None = None

94 load_tier_filter: TierFilter = TierFilter.ALL

95 # Per-request scratch space keyed by value type, so a tier can parse

94 - # kv_transfer_params once (in on_new_request) and read the result back

95 - # on later calls for the same request.

96 + # kv_transfer_params and kv_hints once (in on_new_request) and read the

97 + # result back on later calls for the same request.

98 _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False)

99 # End-token position for each key in this request. The scheduler records

100 # these positions so managers can recover prefix order even when store

vllm/v1/kv_offload/tiering/kvcr/manager.py

+7 -10

@@ -7,15 +7,15 @@ import mmap

7 import socket

8 import time

9 import uuid

10 -from collections.abc import Collection, Iterable, Mapping

10 +from collections.abc import Collection, Iterable

11 from pathlib import Path

12 from typing import TYPE_CHECKING, Any, cast

13

14 +import msgspec

15 from kvcr import (

16 DURATION_METRIC,

17 KVCR,

18 ROUTER_HINT_CAPABILITIES,

18 - ROUTER_HINT_KEY,

19 STATE_METRIC,

20 TRANSFER_BLOCKS_METRIC,

21 TRANSFER_BYTES_METRIC,

@@ -575,14 +575,11 @@ class KVCRSecondaryTierManager(SecondaryTierManager):

575

576 @override

577 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

578 - params = getattr(req_context, "kv_transfer_params", None)

579 - if isinstance(params, Mapping):

580 - hint = params.get(ROUTER_HINT_KEY)

581 - if hint is not None:

582 - self._kvcr.submit_hint(

583 - request_id=req_context.req_id,

584 - hints=hint,

585 - )

578 + if (hint := req_context.kv_hints) is not None:

579 + self._kvcr.submit_hint(

580 + request_id=req_context.req_id,

581 + hints=msgspec.to_builtins(hint),

582 + )

583 return RequestOffloadingContext()

584

585 @override

vllm/v1/request.py

+4

@@ -20,6 +20,7 @@ from vllm.v1.engine import (

20 EngineCoreRequest,

21 FinishReason,

22 )

23 +from vllm.v1.kv_hints import KvHintsEnvelope

24 from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics

25 from vllm.v1.structured_output.request import StructuredOutputRequest

26 from vllm.v1.utils import ConstantList

@@ -78,6 +79,7 @@ class Request:

79 reasoning_ended: bool | None = None,

80 reasoning_parser_kwargs: dict[str, Any] | None = None,

81 abort_immediately: bool = False,

82 + kv_hints: KvHintsEnvelope | None = None,

83 ) -> None:

84 self.request_id = request_id

85 self.client_index = client_index

@@ -193,6 +195,7 @@ class Request:

195 # trace_headers

196 self.trace_headers = trace_headers

197 self.session_id = session_id

198 + self.kv_hints = kv_hints

199

200 # True if this request is scheduled as a non-final prefill chunk.

201 self.is_prefill_chunk = False

@@ -260,6 +263,7 @@ class Request:

263 block_hasher=block_hasher,

264 resumable=request.resumable,

265 session_id=request.session_id,

266 + kv_hints=request.kv_hints,

267 reasoning_ended=request.reasoning_ended,

268 reasoning_parser_kwargs=request.reasoning_parser_kwargs,

269 abort_immediately=request.abort_immediately,

[Feature] Add first-class KV hints request envelope for programmatic KV management (#53423)