f37c550bf6 · tduncan/Nvidia-Optimization-Stack-VLLM

docs/serving/context_parallel_deployment.md

+9

@@ -20,6 +20,15 @@ Both approaches are under active development.

20

21 Due to the auto-regressive nature of decoding, every decoding step needs to compute a small amount of query tokens w.r.t. a large number of key/value tokens stored in the paged KV cache. The core of decode context parallel is how to shard the KV cache across GPUs.

22

23 +`engine_client.vllm_config.cache_config.effective_attention_block_size` reports the

24 +initialized full-attention block size in tokens, including DCP: a physical block of

25 +16 tokens with DCP=4 represents 64 tokens. The same value is available through

26 +gRPC `Control.GetServerInfo.effective_attention_block_size`. Physical block-size

27 +fields retain their existing meaning. The new value is unavailable (`None` in

28 +Python, absent in gRPC) if there is no common full-attention size, the scheduler

29 +does not expose it, or an older engine omits it. gRPC also omits the value when

30 +engine replicas disagree. Partial cache events carry their own actual size.

31 +

32 For a model with `H` kv-heads, a request with `T` tokens in the context needs to store `H * T` key/value tensors in the KV cache.

33

34 1. If one GPU can hold them all, and the performance is good enough, then no parallelization is needed.

rust/Cargo.toml

+1 -1

@@ -147,7 +147,7 @@ vllm-llm = { path = "src/llm" }

147 vllm-managed-engine = { path = "src/managed-engine" }

148 vllm-metrics = { path = "src/metrics" }

149 vllm-parser = { path = "src/parser" }

150 -vllm-proto = { version = "0.2.0", path = "proto" }

150 +vllm-proto = { path = "proto" }

151 vllm-server = { path = "src/server" }

152 vllm-text = { path = "src/text" }

153 vllm-tokenizer = { path = "src/tokenizer" }

rust/proto/README.md

+4 -2

@@ -2,6 +2,9 @@

2

3 This directory is the canonical source for vLLM's gRPC schema.

4

5 +See [context parallel deployment](../../docs/serving/context_parallel_deployment.md)

6 +for effective attention block-size metadata in Python and `Control.GetServerInfo`.

7 +

8 Schema updates are no longer published to the Buf Schema Registry. Rust consumers

9 should use `vllm-proto` from crates.io; consumers in other languages can generate

10 bindings from the `.proto` files in this directory. Buf still builds and lints

@@ -34,8 +37,7 @@ On pull requests and releases, `cargo-semver-checks` compares the crate with its

37 latest published version. Include any required version bump in the protocol

38 change PR. This check becomes available after the first manual publication.

39

37 -1. Update the crate version and the `vllm-proto` workspace dependency together,

38 - and update `rust/Cargo.lock`.

40 +1. Update the crate version in `rust/proto/Cargo.toml` and update `rust/Cargo.lock`.

41 2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run`

42 and the frontend gRPC tests. Record the tested vLLM releases or revisions in

43 the release notes; matching crate versions alone do not establish runtime compatibility.

rust/proto/control.proto

+2

@@ -50,6 +50,8 @@ message ServerInfo {

50 uint64 max_batched_tokens = 9;

51 uint32 max_loras = 10;

52 RlCapabilities rl_capabilities = 11;

53 + // Full-attention block size in tokens after initialization; absent if unavailable.

54 + optional uint64 effective_attention_block_size = 14;

55 }

56

57 message RlCapabilities {

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

+9

@@ -510,6 +510,15 @@ impl EngineCoreClient {

510 self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum()

511 }

512

513 + /// Return the effective attention block size if all engines report the same value.

514 + pub fn effective_attention_block_size(&self) -> Option<u64> {

515 + let size = self.ready_response().effective_attention_block_size?;

516 + self.engines

517 + .iter()

518 + .all(|engine| engine.ready_response.effective_attention_block_size == Some(size))

519 + .then_some(size)

520 + }

521 +

522 /// Return the minimum engine-reported `max_model_len` across all engines.

523 ///

524 /// This is the auto-fitted value after KV cache profiling and may differ

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

+1

@@ -72,6 +72,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {

72 weight_transfer_backend: None,

73 enable_sleep_mode: false,

74 supports_draft_weight_updates: false,

75 + effective_attention_block_size: None,

76 }

77 }

78

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

+3

@@ -104,6 +104,9 @@ pub struct EngineCoreReadyResponse {

104 /// Whether the engine has a speculative draft model that can be updated.

105 #[serde(default)]

106 pub supports_draft_weight_updates: bool,

107 + /// Full-attention block size in tokens after initialization, or unavailable.

108 + #[serde(default)]

109 + pub effective_attention_block_size: Option<u64>,

110 }

111

112 /// Frontend-owned ZMQ addresses that are sent to the engine during startup

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

+6

@@ -2846,6 +2846,11 @@ fn python_msgpack_fixtures_match_rust_encoding() {

2846

2847 let ready_response: EngineCoreReadyResponse =

2848 rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap();

2849 + let mut legacy_ready = serde_json::to_value(&ready_response).unwrap();

2850 + legacy_ready.as_object_mut().unwrap().remove("effective_attention_block_size");

2851 + let legacy_ready: EngineCoreReadyResponse =

2852 + rmp_serde::from_slice(&rmp_serde::to_vec_named(&legacy_ready).unwrap()).unwrap();

2853 + assert!(legacy_ready.effective_attention_block_size.is_none());

2854 assert!(ready_response.supports_lora);

2855 assert_eq!(ready_response.max_loras, 8);

2856 assert_eq!(

@@ -2854,6 +2859,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {

2859 );

2860 assert!(ready_response.enable_sleep_mode);

2861 assert!(ready_response.supports_draft_weight_updates);

2862 + assert_eq!(ready_response.effective_attention_block_size, Some(64));

2863 let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode");

2864 assert!(kv_events_config.enable_kv_cache_events);

2865 assert_eq!(kv_events_config.publisher, "zmq");

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

+2

@@ -423,6 +423,7 @@ class EngineCoreReadyResponse:

423 weight_transfer_backend: str | None = None

424 enable_sleep_mode: bool = False

425 supports_draft_weight_updates: bool = False

426 + effective_attention_block_size: int | None = None

427

428

429 ready_response = EngineCoreReadyResponse(

@@ -446,6 +447,7 @@ ready_response = EngineCoreReadyResponse(

447 weight_transfer_backend="nccl",

448 enable_sleep_mode=True,

449 supports_draft_weight_updates=True,

450 + effective_attention_block_size=64,

451 kv_events_config=KVEventsConfig(

452 enable_kv_cache_events=True,

453 publisher="zmq",

rust/src/server/src/grpc/control.rs

+1

@@ -195,6 +195,7 @@ impl pb::control_server::Control for ControlServiceImpl {

195 max_batched_tokens: ready.max_num_batched_tokens,

196 max_loras: ready.max_loras,

197 rl_capabilities: Some(self.rl_capabilities()),

198 + effective_attention_block_size: self.client().effective_attention_block_size(),

199 }))

200 }

201

rust/src/server/src/grpc/tests.rs

+12 -1

@@ -1911,8 +1911,16 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() {

1911 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

1912 #[serial]

1913 async fn control_reports_server_and_model_info() {

1914 + let mut ready = default_ready_response();

1915 + ready.effective_attention_block_size = Some(64);

1916 let (generate_service, control_service, engine_health, _engine_task) =

1915 - setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await;

1917 + setup_grpc_service_with_engine_script(

1918 + b"engine-grpc-info".to_vec(),

1919 + ready,

1920 + Arc::new(FakeTextBackend),

1921 + |_, _| boxed_test_future(async {}),

1922 + )

1923 + .await;

1924 let (channel, server_task) = start_grpc_test_server(

1925 generate_service,

1926 control_service,

@@ -1928,6 +1936,7 @@ async fn control_reports_server_and_model_info() {

1936 .expect("get server info")

1937 .into_inner();

1938 assert_eq!(server.engine_version, "test-vllm-version");

1939 + assert_eq!(server.effective_attention_block_size, Some(64));

1940 assert_eq!(server.api_version, "vllm");

1941 assert_eq!(server.instance_id, "test-instance");

1942 assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32);

@@ -2195,6 +2204,7 @@ async fn control_aggregates_multi_engine_capacity() {

2204 ready_0.weight_transfer_backend = Some("nccl".to_string());

2205 ready_0.enable_sleep_mode = true;

2206 ready_0.supports_draft_weight_updates = true;

2207 + ready_0.effective_attention_block_size = Some(64);

2208

2209 let mut ready_1 = default_ready_response();

2210 ready_1.max_model_len = 4_096;

@@ -2248,6 +2258,7 @@ async fn control_aggregates_multi_engine_capacity() {

2258 .into_inner();

2259 assert_eq!(server.max_model_len, 4_096);

2260 assert_eq!(server.total_kv_blocks, 30);

2261 + assert!(server.effective_attention_block_size.is_none());

2262 let rl = server.rl_capabilities.expect("RL capabilities");

2263 assert!(!rl.weight_transfer_enabled);

2264 assert!(rl.weight_transfer_backend.is_empty());

tests/config/test_config_utils.py

+3

@@ -215,6 +215,9 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():

215 base_hash = CacheConfig().compute_hash()

216 assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash

217 assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash

218 + config = CacheConfig()

219 + config.effective_attention_block_size = 64

220 + assert config.compute_hash() == base_hash

221

222

223 def test_scheduler_config_hash_includes_max_num_seqs():

tests/v1/core/test_kv_cache_utils.py

+38

@@ -279,6 +279,44 @@ def make_request(

279 )

280

281

282 +@pytest.mark.parametrize("dcp", [1, 4])

283 +def test_effective_attention_block_size_matches_events(dcp):

284 + from vllm.distributed.kv_events import BlockStored

285 + from vllm.v1.engine.core import EngineCore

286 +

287 + config = KVCacheConfig(

288 + num_blocks=32,

289 + kv_cache_tensors=[],

290 + kv_cache_groups=[

291 + KVCacheGroupSpec(["attention"], new_kv_cache_spec()),

292 + ],

293 + )

294 + manager = KVCacheManager(

295 + generate_scheduler_kv_cache_config([config]),

296 + max_model_len=256,

297 + scheduler_block_size=16 * dcp,

298 + hash_block_size=16 * dcp,

299 + dcp_world_size=dcp,

300 + enable_kv_cache_events=True,

301 + )

302 + core = EngineCore.__new__(EngineCore)

303 + core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16))

304 + core.scheduler = SimpleNamespace(kv_cache_manager=manager)

305 + core._initialize_effective_attention_block_size()

306 + block_size = core.vllm_config.cache_config.effective_attention_block_size

307 + assert block_size == 16 * dcp

308 +

309 + request = make_request(

310 + "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256

311 + )

312 + assert manager.allocate_slots(request, 64) is not None

313 + assert [

314 + event.block_size

315 + for event in manager.take_events()

316 + if isinstance(event, BlockStored)

317 + ] == [block_size]

318 +

319 +

320 def new_kv_cache_spec(

321 block_size=16,

322 num_kv_heads=2,

tests/v1/engine/test_engine_core_client.py

+28 -3

@@ -158,7 +158,10 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):

158 local_engines_only=False,

159 enable_elastic_ep=False,

160 )

161 - vllm_config = SimpleNamespace(parallel_config=parallel_config)

161 + vllm_config = SimpleNamespace(

162 + parallel_config=parallel_config,

163 + cache_config=SimpleNamespace(effective_attention_block_size=None),

164 + )

165

166 client = core_client_mod.MPClient(

167 asyncio_mode=False,

@@ -368,10 +371,15 @@ def test_dplb_finished_requests_release_inflight():

371 assert req.request_id not in client.reqs_in_flight

372

373

371 -def test_apply_ready_response_syncs_block_size():

374 +@pytest.mark.parametrize(

375 + ("effective_size", "other_size"),

376 + [(None, None), (4224, 4224), (4224, 1056), (4224, None)],

377 +)

378 +def test_apply_ready_response_syncs_block_size(effective_size, other_size):

379 import msgspec

380

381 client = object.__new__(MPClient)

382 + client._effective_attention_block_sizes = set()

383 client.vllm_config = SimpleNamespace(

384 cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),

385 model_config=SimpleNamespace(max_model_len=8192),

@@ -399,14 +407,31 @@ def test_apply_ready_response_syncs_block_size():

407 max_loras=0,

408 )

409 )

402 - client._apply_ready_response(payload)

410 + fields = msgspec.msgpack.decode(payload)

411 + if effective_size is None:

412 + del fields["effective_attention_block_size"]

413 + else:

414 + fields["effective_attention_block_size"] = effective_size

415 + client._apply_ready_response(msgspec.msgpack.encode(fields))

416 assert client.vllm_config.cache_config.block_size == 1056

417 + cache_config = client.vllm_config.cache_config

418 + assert cache_config.effective_attention_block_size == effective_size

419 +

420 + fields["effective_attention_block_size"] = other_size

421 + client._apply_ready_response(msgspec.msgpack.encode(fields))

422 + assert cache_config.effective_attention_block_size == (

423 + effective_size if effective_size == other_size else None

424 + )

425 +

426 + client._apply_ready_response(b"")

427 + assert cache_config.effective_attention_block_size is None

428

429

430 def test_apply_ready_response_syncs_mamba_block_size():

431 import msgspec

432

433 client = object.__new__(MPClient)

434 + client._effective_attention_block_sizes = set()

435 client.vllm_config = SimpleNamespace(

436 cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),

437 model_config=SimpleNamespace(max_model_len=8192),

vllm/config/cache.py

+3

@@ -211,6 +211,8 @@ class CacheConfig:

211 """The number of blocks to allocate for CPU memory."""

212

213 # Set after KV cache initialization.

214 + effective_attention_block_size: int | None = field(default=None, init=False)

215 + """Full-attention block size in tokens, including DCP, or None if unavailable."""

216 kv_cache_size_tokens: int | None = field(default=None, init=False)

217 """Per-DP-engine KV cache capacity in tokens (group-aware). Uses

218 group-aware capacity since num_gpu_blocks * block_size can be wrong

@@ -277,6 +279,7 @@ class CacheConfig:

279 # Post-init/derived counters

280 "num_gpu_blocks",

281 "num_cpu_blocks",

282 + "effective_attention_block_size",

283 "kv_cache_size_tokens",

284 "kv_cache_max_concurrency",

285 # WIP feature toggle not impacting compiled graph shape

vllm/v1/engine/__init__.py

+2

@@ -102,6 +102,8 @@ class EngineCoreReadyResponse:

102 weight_transfer_backend: str | None = None

103 enable_sleep_mode: bool = False

104 supports_draft_weight_updates: bool = False

105 + # Full-attention block size in tokens after initialization, or unavailable.

106 + effective_attention_block_size: int | None = None

107

108

109 class EngineCoreRequest(

vllm/v1/engine/core.py

+20 -1

@@ -85,7 +85,7 @@ from vllm.v1.fault_tolerance.engine_core_sentinel import (

85 EngineCoreSentinel,

86 fault_tolerant_wrapper,

87 )

88 -from vllm.v1.kv_cache_interface import KVCacheConfig

88 +from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec

89 from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats

90 from vllm.v1.outputs import ModelRunnerOutput

91 from vllm.v1.request import Request, RequestStatus

@@ -168,6 +168,7 @@ class EngineCore:

168 block_size=scheduler_block_size,

169 hash_block_size=hash_block_size,

170 )

171 + self._initialize_effective_attention_block_size()

172 self.use_spec_decode = vllm_config.speculative_config is not None

173 self.check_for_draft_tokens = (

174 self.use_spec_decode or vllm_config.model_config.is_diffusion

@@ -382,6 +383,21 @@ class EngineCore:

383 )

384 return scheduler_kv_cache_config

385

386 + def _initialize_effective_attention_block_size(self) -> None:

387 + cache_config = self.vllm_config.cache_config

388 + cache_config.effective_attention_block_size = None

389 + cache_manager = getattr(self.scheduler, "kv_cache_manager", None)

390 + if cache_manager is None:

391 + return

392 + block_sizes = {

393 + manager.block_size

394 + for manager in cache_manager.coordinator.single_type_managers

395 + if is_full_attention_spec(manager.kv_cache_spec)

396 + }

397 + cache_config.effective_attention_block_size = (

398 + block_sizes.pop() if len(block_sizes) == 1 else None

399 + )

400 +

401 def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

402 supported_tasks = self.model_executor.supported_tasks

403 self._log_pooler_config(supported_tasks)

@@ -1649,6 +1665,9 @@ class EngineCoreProc(EngineCore):

1665 num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0,

1666 block_size=self.vllm_config.cache_config.block_size,

1667 mamba_block_size=self.vllm_config.cache_config.mamba_block_size,

1668 + effective_attention_block_size=(

1669 + self.vllm_config.cache_config.effective_attention_block_size

1670 + ),

1671 dp_stats_address=self.frontend_stats_publish_address,

1672 dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."),

1673 vllm_version=VLLM_VERSION,

vllm/v1/engine/core_client.py

+15 -3

@@ -566,6 +566,7 @@ class MPClient(EngineCoreClient):

566 ):

567 self.vllm_config = vllm_config

568 self._renderer: BaseRenderer | None = renderer

569 + self._effective_attention_block_sizes: set[int | None] = set()

570

571 # ZMQ setup.

572 sync_ctx = zmq.Context(io_threads=2)

@@ -803,10 +804,21 @@ class MPClient(EngineCoreClient):

804 def _apply_ready_response(self, payload: bytes) -> None:

805 """Decode an EngineCoreReadyResponse and sync any post-initialization

806 config changes (e.g. auto-fitted max_model_len) back to the frontend."""

806 - if not payload:

807 - return

807 vllm_config = self.vllm_config

809 - response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)

808 + response = (

809 + msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse)

810 + if payload

811 + else None

812 + )

813 + self._effective_attention_block_sizes.add(

814 + response.effective_attention_block_size if response is not None else None

815 + )

816 + sizes = self._effective_attention_block_sizes

817 + vllm_config.cache_config.effective_attention_block_size = (

818 + next(iter(sizes)) if len(sizes) == 1 else None

819 + )

820 + if response is None:

821 + return

822 vllm_config.model_config.max_model_len = min(

823 vllm_config.model_config.max_model_len, response.max_model_len

824 )

[Feature][Frontend] Expose effective attention block size for DCP (#56538)