@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bdf13965ab
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

[V1] Support cross-layer KV sharing (#18212)

Signed-off-by: Yong Hoon Shin <yhshin@meta.com>

Yong Hoon Shin committed Jun 3, 2025 at 13:33 UTC bdf13965ab4a528d30cb82854487910189865d9d
31 files changed +733 -73
tests/v1/tpu/worker/test_tpu_model_runner.py
+226 -1
@@ -4,8 +4,13 @@ import unittest.mock as mock
4
5 import pytest
6
7 -from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
7 +from vllm.attention.layer import Attention
8 +from vllm.config import (CacheConfig, ModelConfig, SchedulerConfig, VllmConfig,
9 + set_current_vllm_config)
10 from vllm.sampling_params import SamplingParams
11 +from vllm.utils import GiB_bytes
12 +from vllm.v1.core.kv_cache_utils import (estimate_max_model_len,
13 + get_kv_cache_config)
14 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
15 SchedulerOutput)
16 from vllm.v1.worker.tpu_model_runner import (
@@ -363,3 +368,223 @@ def test_get_req_paddings():
368 assert _get_req_paddings(1, 32) == [8, 16, 32]
369 assert _get_req_paddings(8, 32) == [8, 16, 32]
370 assert _get_req_paddings(8, 36) == [8, 16, 32, 36]
371 +
372 +
373 +def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
374 + layer_0 = "model.layers.0.self_attn.attn"
375 + layer_1 = "model.layers.1.self_attn.attn"
376 + error_msg = f"{layer_1} must come before the current layer"
377 + with pytest.raises(ValueError, match=error_msg):
378 + fwd_context = {
379 + # initialization below will fail because target layer is invalid;
380 + # the target layer needs to come before layer 1
381 + layer_0:
382 + Attention(
383 + num_heads=8,
384 + head_size=64,
385 + scale=1.0,
386 + prefix=layer_0,
387 + kv_sharing_target_layer_name=layer_1,
388 + ),
389 + layer_1:
390 + Attention(
391 + num_heads=8,
392 + head_size=64,
393 + scale=1.0,
394 + prefix=layer_1,
395 + )
396 + }
397 + # suppress var not used error
398 + assert fwd_context is not None
399 +
400 +
401 +def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
402 + layer_0 = "model.layers.0.self_attn.attn"
403 + layer_1 = "model.layers.1.self_attn.attn"
404 + invalid_layer = "model.layers.0.cross_attn.attn"
405 + error_msg = f"{invalid_layer} is not a valid Attention layer in the model"
406 + with pytest.raises(ValueError, match=error_msg):
407 + fwd_context = {
408 + layer_0:
409 + Attention(
410 + num_heads=8,
411 + head_size=64,
412 + scale=1.0,
413 + prefix=layer_0,
414 + ),
415 + layer_1:
416 + Attention(
417 + num_heads=8,
418 + head_size=64,
419 + scale=1.0,
420 + prefix=layer_1,
421 + # invalid layer: cross_attn.atn doesn't exist!
422 + kv_sharing_target_layer_name=invalid_layer,
423 + )
424 + }
425 + # suppress var not used error
426 + assert fwd_context is not None
427 +
428 +
429 +def test_init_kv_cache_with_kv_sharing_target_same_as_current():
430 + layer_0 = "model.layers.0.self_attn.attn"
431 + layer_1 = "model.layers.1.self_attn.attn"
432 + error_msg = f"{layer_1} cannot be the same as the current layer"
433 + with pytest.raises(ValueError, match=error_msg):
434 + fwd_context = {
435 + # initialization below will fail because target layer is invalid;
436 + # the target layer needs to come before layer 1
437 + layer_0:
438 + Attention(
439 + num_heads=8,
440 + head_size=64,
441 + scale=1.0,
442 + prefix=layer_0,
443 + ),
444 + layer_1:
445 + Attention(
446 + num_heads=8,
447 + head_size=64,
448 + scale=1.0,
449 + prefix=layer_1,
450 + kv_sharing_target_layer_name=layer_1,
451 + )
452 + }
453 + # suppress var not used error
454 + assert fwd_context is not None
455 +
456 +
457 +def test_init_kv_cache_without_kv_sharing(model_runner):
458 + layer_0 = "model.layers.0.self_attn.attn"
459 + layer_1 = "model.layers.1.self_attn.attn"
460 + vllm_config = model_runner.vllm_config
461 + with set_current_vllm_config(vllm_config):
462 + fwd_context = {
463 + layer_0:
464 + Attention(
465 + num_heads=8,
466 + head_size=64,
467 + scale=1.0,
468 + prefix=layer_0,
469 + ),
470 + layer_1:
471 + Attention(
472 + num_heads=8,
473 + head_size=64,
474 + scale=1.0,
475 + prefix=layer_1,
476 + )
477 + }
478 + # suppress var not used error
479 + assert fwd_context is not None
480 + # Set high context length to test max context length estimation
481 + vllm_config.model_config.max_model_len = 3_000_000
482 + vllm_ctx = vllm_config.compilation_config.static_forward_context
483 + kv_cache_spec = model_runner.get_kv_cache_spec()
484 + assert len(kv_cache_spec) == 2
485 + assert len(model_runner.shared_kv_cache_layers) == 0
486 +
487 + available_memory = 20 * GiB_bytes
488 + # page size for layer 0's kv_cache_spec is 32KB
489 + num_expected_blocks = 327680 # 20GB / 32KB / 2 (num layers)
490 + kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
491 + available_memory)
492 + assert kv_cache_config.num_blocks == num_expected_blocks
493 + assert len(kv_cache_config.tensors) == 2
494 + assert kv_cache_config.tensors[layer_0].size == available_memory // 2
495 + assert kv_cache_config.tensors[layer_1].size == available_memory // 2
496 +
497 + max_context_len =\
498 + estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
499 + # max context len with KV sharing should be 2x as large as without
500 + assert max_context_len == 1310720
501 +
502 + # important: override tensor size to prevent large mem alloc during test
503 + # this will only allocate 2 block worth of memory (2 * 32kb)
504 + kv_cache_config.num_blocks = 1
505 + for layer in kv_cache_config.tensors:
506 + kv_cache_config.tensors[layer].size =\
507 + kv_cache_spec[layer].page_size_bytes
508 +
509 + model_runner.initialize_kv_cache(kv_cache_config)
510 +
511 + layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
512 + layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
513 + # check layer 1 kv cache does NOT share memory with layer 0
514 + assert id(layer_1_kv) != id(layer_0_kv)
515 +
516 + # check layer 1 added to kv cache group's layer names
517 + assert len(kv_cache_config.kv_cache_groups) == 1
518 + assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
519 + assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
520 + assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
521 +
522 +
523 +def test_init_kv_cache_with_kv_sharing_valid(model_runner):
524 + layer_0 = "model.layers.0.self_attn.attn"
525 + layer_1 = "model.layers.1.self_attn.attn"
526 + vllm_config = model_runner.vllm_config
527 + with set_current_vllm_config(vllm_config):
528 + fwd_context = {
529 + layer_0:
530 + Attention(
531 + num_heads=8,
532 + head_size=64,
533 + scale=1.0,
534 + prefix=layer_0,
535 + ),
536 + layer_1:
537 + Attention(
538 + num_heads=8,
539 + head_size=64,
540 + scale=1.0,
541 + prefix=layer_1,
542 + kv_sharing_target_layer_name="model.layers.0.self_attn.attn",
543 + )
544 + }
545 + # suppress var not used error
546 + assert fwd_context is not None
547 + # Set high context length to test max context length estimation
548 + vllm_config.model_config.max_model_len = 3_000_000
549 + vllm_ctx = vllm_config.compilation_config.static_forward_context
550 + kv_cache_spec = model_runner.get_kv_cache_spec()
551 + assert len(kv_cache_spec) == 1
552 + assert layer_0 in kv_cache_spec
553 + assert model_runner.shared_kv_cache_layers[layer_1] == layer_0
554 +
555 + available_memory = 20 * GiB_bytes
556 + # page size for layer 0's kv_cache_spec is 32KB
557 + # with KV sharing, we can allocate (available_mem//page_size//1) blocks
558 + # which is twice as many as without KV sharing
559 + num_expected_blocks = 655360 # 20GB / 32KB
560 + kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
561 + available_memory)
562 + assert kv_cache_config.num_blocks == num_expected_blocks
563 + assert len(kv_cache_config.tensors) == 1
564 + # Each layer now has twice the available memory for KV cache
565 + # compared to no KV sharing
566 + assert kv_cache_config.tensors[layer_0].size == available_memory
567 +
568 + max_context_len =\
569 + estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
570 + # max context len with KV sharing should be 2x as large as without
571 + assert max_context_len == 2 * 1310720
572 +
573 + # important: override tensor size to prevent large mem alloc during test
574 + # this will only allocate 1 block worth of memory (32kb)
575 + kv_cache_config.num_blocks = 1
576 + kv_cache_config.tensors[layer_0].size =\
577 + kv_cache_spec[layer_0].page_size_bytes
578 +
579 + model_runner.initialize_kv_cache(kv_cache_config)
580 +
581 + layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
582 + layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
583 + # check layer 1 kv cache shares memory with layer 0
584 + assert id(layer_1_kv) == id(layer_0_kv)
585 +
586 + # check layer 1 added to kv cache group's layer names
587 + assert len(kv_cache_config.kv_cache_groups) == 1
588 + assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
589 + assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
590 + assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
tests/v1/worker/test_gpu_model_runner.py
+237 -7
@@ -7,8 +7,11 @@ import pytest
7
8 from vllm.attention import Attention
9 from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
10 - SchedulerConfig, VllmConfig)
10 + SchedulerConfig, VllmConfig, set_current_vllm_config)
11 from vllm.sampling_params import SamplingParams
12 +from vllm.utils import GiB_bytes
13 +from vllm.v1.core.kv_cache_utils import (estimate_max_model_len,
14 + get_kv_cache_config)
15 from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
16 SchedulerOutput)
17 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -19,6 +22,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner
22
23 BLOCK_SIZE = 16
24 NUM_BLOCKS = 10
25 +DEVICE = "cuda"
26
27
28 def initialize_kv_cache(runner: GPUModelRunner):
@@ -55,8 +59,7 @@ def initialize_kv_cache(runner: GPUModelRunner):
59 runner.initialize_attn_backend(kv_cache_config)
60
61
58 -@pytest.fixture
59 -def model_runner():
62 +def get_vllm_config():
63 scheduler_config = SchedulerConfig(
64 max_num_seqs=10,
65 max_num_batched_tokens=512,
@@ -84,13 +87,18 @@ def model_runner():
87 scheduler_config=scheduler_config,
88 parallel_config=parallel_config,
89 )
87 - num_heads = model_config.get_num_kv_heads(parallel_config)
90 + return vllm_config
91 +
92 +
93 +@pytest.fixture
94 +def model_runner():
95 + vllm_config = get_vllm_config()
96 + model_config = vllm_config.model_config
97 + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)
98 head_size = model_config.get_head_size()
99 vllm_config.compilation_config.static_forward_context[
100 "layer.0"] = Attention(num_heads, head_size, 0.1)
91 -
92 - device = "cuda"
93 - runner = GPUModelRunner(vllm_config, device)
101 + runner = GPUModelRunner(vllm_config, DEVICE)
102 initialize_kv_cache(runner)
103 return runner
104
@@ -385,3 +393,225 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
393 model_runner_2.load_model() # Load real weights inplace
394 assert str(model_runner.get_model().state_dict()) == str(
395 model_runner_2.get_model().state_dict())
396 +
397 +
398 +def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
399 + layer_0 = "model.layers.0.self_attn.attn"
400 + layer_1 = "model.layers.1.self_attn.attn"
401 + error_msg = f"{layer_1} must come before the current layer"
402 + with pytest.raises(ValueError, match=error_msg):
403 + fwd_context = {
404 + # initialization below will fail because target layer is invalid;
405 + # the target layer needs to come before layer 1
406 + layer_0:
407 + Attention(
408 + num_heads=8,
409 + head_size=64,
410 + scale=1.0,
411 + prefix=layer_0,
412 + kv_sharing_target_layer_name=layer_1,
413 + ),
414 + layer_1:
415 + Attention(
416 + num_heads=8,
417 + head_size=64,
418 + scale=1.0,
419 + prefix=layer_1,
420 + )
421 + }
422 + # suppress var not used error
423 + assert fwd_context is not None
424 +
425 +
426 +def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():
427 + layer_0 = "model.layers.0.self_attn.attn"
428 + layer_1 = "model.layers.1.self_attn.attn"
429 + invalid_layer = "model.layers.0.cross_attn.attn"
430 + error_msg = f"{invalid_layer} is not a valid Attention layer in the model"
431 + with pytest.raises(ValueError, match=error_msg):
432 + fwd_context = {
433 + layer_0:
434 + Attention(
435 + num_heads=8,
436 + head_size=64,
437 + scale=1.0,
438 + prefix=layer_0,
439 + ),
440 + layer_1:
441 + Attention(
442 + num_heads=8,
443 + head_size=64,
444 + scale=1.0,
445 + prefix=layer_1,
446 + # invalid layer: cross_attn.atn doesn't exist!
447 + kv_sharing_target_layer_name=invalid_layer,
448 + )
449 + }
450 + # suppress var not used error
451 + assert fwd_context is not None
452 +
453 +
454 +def test_init_kv_cache_with_kv_sharing_target_same_as_current():
455 + layer_0 = "model.layers.0.self_attn.attn"
456 + layer_1 = "model.layers.1.self_attn.attn"
457 + error_msg = f"{layer_1} cannot be the same as the current layer"
458 + with pytest.raises(ValueError, match=error_msg):
459 + fwd_context = {
460 + # initialization below will fail because target layer is invalid;
461 + # the target layer needs to come before layer 1
462 + layer_0:
463 + Attention(
464 + num_heads=8,
465 + head_size=64,
466 + scale=1.0,
467 + prefix=layer_0,
468 + ),
469 + layer_1:
470 + Attention(
471 + num_heads=8,
472 + head_size=64,
473 + scale=1.0,
474 + prefix=layer_1,
475 + kv_sharing_target_layer_name=layer_1,
476 + )
477 + }
478 + # suppress var not used error
479 + assert fwd_context is not None
480 +
481 +
482 +def test_init_kv_cache_without_kv_sharing():
483 + layer_0 = "model.layers.0.self_attn.attn"
484 + layer_1 = "model.layers.1.self_attn.attn"
485 + vllm_config = get_vllm_config()
486 + with set_current_vllm_config(vllm_config):
487 + fwd_context = {
488 + layer_0:
489 + Attention(
490 + num_heads=8,
491 + head_size=64,
492 + scale=1.0,
493 + prefix=layer_0,
494 + ),
495 + layer_1:
496 + Attention(
497 + num_heads=8,
498 + head_size=64,
499 + scale=1.0,
500 + prefix=layer_1,
501 + )
502 + }
503 + # suppress var not used error
504 + assert fwd_context is not None
505 + # Set high context length to test max context length estimation
506 + vllm_config.model_config.max_model_len = 3_000_000
507 + vllm_ctx = vllm_config.compilation_config.static_forward_context
508 + runner = GPUModelRunner(vllm_config, DEVICE)
509 + kv_cache_spec = runner.get_kv_cache_spec()
510 + assert len(kv_cache_spec) == 2
511 + assert len(runner.shared_kv_cache_layers) == 0
512 +
513 + available_memory = 20 * GiB_bytes
514 + # page size for layer 0's kv_cache_spec is 32KB
515 + num_expected_blocks = 327680 # 20GB / 32KB / 2 (num layers)
516 + kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
517 + available_memory)
518 + assert kv_cache_config.num_blocks == num_expected_blocks
519 + assert len(kv_cache_config.tensors) == 2
520 + assert kv_cache_config.tensors[layer_0].size == available_memory // 2
521 + assert kv_cache_config.tensors[layer_1].size == available_memory // 2
522 +
523 + max_context_len =\
524 + estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
525 + # max context len with KV sharing should be 2x as large as without
526 + assert max_context_len == 1310720
527 +
528 + # important: override tensor size to prevent large mem alloc during test
529 + # this will only allocate 2 block worth of memory (2 * 32kb)
530 + kv_cache_config.num_blocks = 1
531 + for layer in kv_cache_config.tensors:
532 + kv_cache_config.tensors[layer].size =\
533 + kv_cache_spec[layer].page_size_bytes
534 +
535 + runner.initialize_kv_cache(kv_cache_config)
536 +
537 + layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
538 + layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
539 + # check layer 1 kv cache does NOT share memory with layer 0
540 + assert id(layer_1_kv) != id(layer_0_kv)
541 +
542 + # check layer 1 added to kv cache group's layer names
543 + assert len(kv_cache_config.kv_cache_groups) == 1
544 + assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
545 + assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
546 + assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
547 +
548 +
549 +def test_init_kv_cache_with_kv_sharing_valid():
550 + layer_0 = "model.layers.0.self_attn.attn"
551 + layer_1 = "model.layers.1.self_attn.attn"
552 + vllm_config = get_vllm_config()
553 + with set_current_vllm_config(vllm_config):
554 + fwd_context = {
555 + layer_0:
556 + Attention(
557 + num_heads=8,
558 + head_size=64,
559 + scale=1.0,
560 + prefix=layer_0,
561 + ),
562 + layer_1:
563 + Attention(
564 + num_heads=8,
565 + head_size=64,
566 + scale=1.0,
567 + prefix=layer_1,
568 + kv_sharing_target_layer_name="model.layers.0.self_attn.attn",
569 + )
570 + }
571 + # suppress var not used error
572 + assert fwd_context is not None
573 + # Set high context length to test max context length estimation
574 + vllm_config.model_config.max_model_len = 3_000_000
575 + vllm_ctx = vllm_config.compilation_config.static_forward_context
576 + runner = GPUModelRunner(vllm_config, DEVICE)
577 + kv_cache_spec = runner.get_kv_cache_spec()
578 + assert len(kv_cache_spec) == 1
579 + assert layer_0 in kv_cache_spec
580 + assert runner.shared_kv_cache_layers[layer_1] == layer_0
581 +
582 + available_memory = 20 * GiB_bytes
583 + # page size for layer 0's kv_cache_spec is 32KB
584 + # with KV sharing, we can allocate (available_mem//page_size//1) blocks
585 + # which is twice as many as without KV sharing
586 + num_expected_blocks = 655360 # 20GB / 32KB
587 + kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
588 + available_memory)
589 + assert kv_cache_config.num_blocks == num_expected_blocks
590 + assert len(kv_cache_config.tensors) == 1
591 + # Each layer now has twice the available memory for KV cache
592 + # compared to no KV sharing
593 + assert kv_cache_config.tensors[layer_0].size == available_memory
594 +
595 + max_context_len =\
596 + estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
597 + # max context len with KV sharing should be 2x as large as without
598 + assert max_context_len == 2 * 1310720
599 +
600 + # important: override tensor size to prevent large mem alloc during test
601 + # this will only allocate 1 block worth of memory (32kb)
602 + kv_cache_config.num_blocks = 1
603 + kv_cache_config.tensors[layer_0].size =\
604 + kv_cache_spec[layer_0].page_size_bytes
605 +
606 + runner.initialize_kv_cache(kv_cache_config)
607 +
608 + layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
609 + layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
610 + # check layer 1 kv cache shares memory with layer 0
611 + assert id(layer_1_kv) == id(layer_0_kv)
612 +
613 + # check layer 1 added to kv cache group's layer names
614 + assert len(kv_cache_config.kv_cache_groups) == 1
615 + assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
616 + assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
617 + assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
vllm/attention/backends/abstract.py
+1
@@ -270,6 +270,7 @@ class AttentionImpl(ABC, Generic[T]):
270 blocksparse_params: Optional[Dict[str, Any]] = None,
271 logits_soft_cap: Optional[float] = None,
272 attn_type: str = AttentionType.DECODER,
273 + kv_sharing_target_layer_name: Optional[str] = None,
274 ) -> None:
275 raise NotImplementedError
276
vllm/attention/backends/blocksparse_attn.py
+3
@@ -306,7 +306,10 @@ class BlocksparseFlashAttentionImpl(AttentionImpl):
306 blocksparse_params: Optional[Dict[str, Any]] = None,
307 logits_soft_cap: Optional[float] = None,
308 attn_type: str = AttentionType.DECODER,
309 + kv_sharing_target_layer_name: Optional[str] = None,
310 ) -> None:
311 + if kv_sharing_target_layer_name is not None:
312 + raise NotImplementedError("KV sharing is not supported in V0.")
313 assert blocksparse_params is not None
314 assert alibi_slopes is None, ValueError(
315 "Alibi not support for blocksparse flash attention.")
vllm/attention/backends/cpu_mla.py
+2 -1
@@ -206,12 +206,13 @@ class CPUMLAImpl(MLACommonImpl[CPUMLAMetadata]):
206 blocksparse_params: Optional[Dict[str, Any]],
207 logits_soft_cap: Optional[float],
208 attn_type: str,
209 + kv_sharing_target_layer_name: Optional[str],
210 # MLA Specific Arguments
211 **mla_args) -> None:
212 super().__init__(num_heads, head_size, scale, num_kv_heads,
213 alibi_slopes, sliding_window, kv_cache_dtype,
214 blocksparse_params, logits_soft_cap, attn_type,
214 - **mla_args)
215 + kv_sharing_target_layer_name, **mla_args)
216
217 unsupported_features = [
218 alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
vllm/attention/backends/dual_chunk_flash_attn.py
+3
@@ -290,9 +290,12 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
290 blocksparse_params: Optional[Dict[str, Any]] = None,
291 logits_soft_cap: Optional[float] = None,
292 attn_type: str = AttentionType.DECODER,
293 + kv_sharing_target_layer_name: Optional[str] = None,
294 layer_idx: int = -1,
295 dual_chunk_attention_config: Optional[Dict[str, Any]] = None,
296 ) -> None:
297 + if kv_sharing_target_layer_name is not None:
298 + raise NotImplementedError("KV sharing is not supported in V0.")
299 self.num_heads = num_heads
300 self.head_size = head_size
301 self.scale = float(scale)
vllm/attention/backends/flash_attn.py
+3
@@ -618,8 +618,11 @@ class FlashAttentionImpl(AttentionImpl):
618 blocksparse_params: Optional[Dict[str, Any]] = None,
619 logits_soft_cap: Optional[float] = None,
620 attn_type: str = AttentionType.DECODER,
621 + kv_sharing_target_layer_name: Optional[str] = None,
622 use_irope: bool = False,
623 ) -> None:
624 + if kv_sharing_target_layer_name is not None:
625 + raise NotImplementedError("KV sharing is not supported in V0.")
626 if blocksparse_params is not None:
627 raise ValueError(
628 "FlashAttention does not support block-sparse attention.")
vllm/attention/backends/flashinfer.py
+3
@@ -936,8 +936,11 @@ class FlashInferImpl(AttentionImpl):
936 blocksparse_params: Optional[Dict[str, Any]] = None,
937 logits_soft_cap: Optional[float] = None,
938 attn_type: str = AttentionType.DECODER,
939 + kv_sharing_target_layer_name: Optional[str] = None,
940 use_irope: bool = False,
941 ) -> None:
942 + if kv_sharing_target_layer_name is not None:
943 + raise NotImplementedError("KV sharing is not supported in V0.")
944 if use_irope:
945 logger.warning_once(
946 "Using irope in FlashInfer is not supported yet, it will fall"
vllm/attention/backends/flashmla.py
+2 -1
@@ -184,12 +184,13 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
184 blocksparse_params: Optional[Dict[str, Any]],
185 logits_soft_cap: Optional[float],
186 attn_type: str,
187 + kv_sharing_target_layer_name: Optional[str] = None,
188 # MLA Specific Arguments
189 **mla_args) -> None:
190 super().__init__(num_heads, head_size, scale, num_kv_heads,
191 alibi_slopes, sliding_window, kv_cache_dtype,
192 blocksparse_params, logits_soft_cap, attn_type,
192 - **mla_args)
193 + kv_sharing_target_layer_name, **mla_args)
194
195 assert is_flashmla_supported(), \
196 "FlashMLA is not supported on this device"
vllm/attention/backends/hpu_attn.py
+3
@@ -110,9 +110,12 @@ class HPUAttentionImpl(AttentionImpl, torch.nn.Module):
110 blocksparse_params: Optional[Dict[str, Any]] = None,
111 max_seq_len: int = 4096,
112 attn_type: str = AttentionType.DECODER,
113 + kv_sharing_target_layer_name: Optional[str] = None,
114 use_irope: bool = False,
115 ) -> None:
116 super(AttentionImpl, self).__init__()
117 + if kv_sharing_target_layer_name is not None:
118 + raise NotImplementedError("KV sharing is not supported in V0.")
119 if use_irope:
120 logger.warning_once(
121 "Using irope in HPU is not supported yet, it will fall back "
vllm/attention/backends/ipex_attn.py
+3
@@ -123,8 +123,11 @@ class IpexAttnBackendImpl(AttentionImpl[IpexAttnMetadata]):
123 blocksparse_params: Optional[Dict[str, Any]] = None,
124 logits_soft_cap: Optional[float] = None,
125 attn_type: str = AttentionType.DECODER,
126 + kv_sharing_target_layer_name: Optional[str] = None,
127 use_irope: bool = False,
128 ) -> None:
129 + if kv_sharing_target_layer_name is not None:
130 + raise NotImplementedError("KV sharing is not supported in V0.")
131 if use_irope:
132 logger.warning_once(
133 "Using irope in Ipex is not supported yet, it will fall"
vllm/attention/backends/mla/common.py
+3
@@ -1000,6 +1000,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1000 blocksparse_params: Optional[Dict[str, Any]],
1001 logits_soft_cap: Optional[float],
1002 attn_type: str,
1003 + kv_sharing_target_layer_name: Optional[str],
1004 # MLA Specific Arguments
1005 q_lora_rank: Optional[int],
1006 kv_lora_rank: int,
@@ -1009,6 +1010,8 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1010 v_head_dim: int,
1011 kv_b_proj: ColumnParallelLinear,
1012 ) -> None:
1013 + if kv_sharing_target_layer_name is not None:
1014 + raise NotImplementedError("KV sharing not supported in V0.")
1015 self.num_heads = num_heads
1016 self.head_size = head_size
1017 self.scale = float(scale)
vllm/attention/backends/pallas.py
+3
@@ -109,8 +109,11 @@ class PallasAttentionBackendImpl(AttentionImpl):
109 blocksparse_params: Optional[Dict[str, Any]] = None,
110 logits_soft_cap: Optional[float] = None,
111 attn_type: str = AttentionType.DECODER,
112 + kv_sharing_target_layer_name: Optional[str] = None,
113 use_irope: bool = False,
114 ) -> None:
115 + if kv_sharing_target_layer_name is not None:
116 + raise NotImplementedError("KV sharing is not supported in V0.")
117 if use_irope:
118 logger.warning_once(
119 "Using irope in Pallas is not supported yet, it will fall back "
vllm/attention/backends/rocm_aiter_mla.py
+2 -1
@@ -370,12 +370,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
370 blocksparse_params: Optional[dict[str, Any]],
371 logits_soft_cap: Optional[float],
372 attn_type: str,
373 + kv_sharing_target_layer_name: Optional[str],
374 # MLA Specific Arguments
375 **mla_args) -> None:
376 super().__init__(num_heads, head_size, scale, num_kv_heads,
377 alibi_slopes, sliding_window, kv_cache_dtype,
378 blocksparse_params, logits_soft_cap, attn_type,
378 - **mla_args)
379 + kv_sharing_target_layer_name, **mla_args)
380
381 unsupported_features = [
382 alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
vllm/attention/backends/rocm_flash_attn.py
+3
@@ -494,8 +494,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
494 blocksparse_params: Optional[Dict[str, Any]] = None,
495 logits_soft_cap: Optional[float] = None,
496 attn_type: str = AttentionType.DECODER,
497 + kv_sharing_target_layer_name: Optional[str] = None,
498 use_irope: bool = False,
499 ) -> None:
500 + if kv_sharing_target_layer_name is not None:
501 + raise NotImplementedError("KV sharing is not supported in V0.")
502 if use_irope:
503 logger.warning_once(
504 "Using irope in ROCm Flash Attention is not supported yet, it "
vllm/attention/backends/torch_sdpa.py
+3
@@ -405,8 +405,11 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
405 blocksparse_params: Optional[Dict[str, Any]] = None,
406 logits_soft_cap: Optional[float] = None,
407 attn_type: str = AttentionType.DECODER,
408 + kv_sharing_target_layer_name: Optional[str] = None,
409 use_irope: bool = False,
410 ) -> None:
411 + if kv_sharing_target_layer_name is not None:
412 + raise NotImplementedError("KV sharing is not supported in V0.")
413 if blocksparse_params is not None:
414 raise ValueError(
415 "Torch SPDA does not support block-sparse attention.")
vllm/attention/backends/triton_mla.py
+2 -1
@@ -38,12 +38,13 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
38 blocksparse_params: Optional[Dict[str, Any]],
39 logits_soft_cap: Optional[float],
40 attn_type: str,
41 + kv_sharing_target_layer_name: Optional[str],
42 # MLA Specific Arguments
43 **mla_args) -> None:
44 super().__init__(num_heads, head_size, scale, num_kv_heads,
45 alibi_slopes, sliding_window, kv_cache_dtype,
46 blocksparse_params, logits_soft_cap, attn_type,
46 - **mla_args)
47 + kv_sharing_target_layer_name, **mla_args)
48
49 unsupported_features = [
50 alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
vllm/attention/backends/xformers.py
+3
@@ -390,8 +390,11 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
390 blocksparse_params: Optional[Dict[str, Any]] = None,
391 logits_soft_cap: Optional[float] = None,
392 attn_type: str = AttentionType.DECODER,
393 + kv_sharing_target_layer_name: Optional[str] = None,
394 use_irope: bool = False,
395 ) -> None:
396 + if kv_sharing_target_layer_name is not None:
397 + raise NotImplementedError("KV sharing is not supported in V0.")
398 if blocksparse_params is not None:
399 raise ValueError(
400 "XFormers does not support block-sparse attention.")
vllm/attention/layer.py
+16 -1
@@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.base_config import (
21 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
22 from vllm.platforms import _Backend, current_platform
23 from vllm.utils import direct_register_custom_op
24 +from vllm.v1.attention.backends.utils import validate_kv_sharing_target
25
26
27 class Attention(nn.Module):
@@ -50,6 +51,7 @@ class Attention(nn.Module):
51 use_mla: bool = False,
52 prefix: str = "",
53 attn_type: str = AttentionType.DECODER,
54 + kv_sharing_target_layer_name: Optional[str] = None,
55 **extra_impl_args,
56 ) -> None:
57 """
@@ -135,7 +137,7 @@ class Attention(nn.Module):
137 self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
138 alibi_slopes, sliding_window, kv_cache_dtype,
139 blocksparse_params, logits_soft_cap, attn_type,
138 - **extra_impl_args)
140 + kv_sharing_target_layer_name, **extra_impl_args)
141 self.backend = backend_name_to_enum(attn_backend.get_name())
142 self.dtype = dtype
143
@@ -153,6 +155,19 @@ class Attention(nn.Module):
155 compilation_config.static_forward_context[prefix] = self
156 self.layer_name = prefix
157 self.attn_type = attn_type
158 +
159 + if kv_sharing_target_layer_name is not None:
160 + if not envs.VLLM_USE_V1:
161 + raise NotImplementedError(
162 + "Cross-layer KV sharing is not supported in V0.")
163 +
164 + validate_kv_sharing_target(
165 + prefix,
166 + kv_sharing_target_layer_name,
167 + compilation_config.static_forward_context,
168 + )
169 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
170 +
171 # use a placeholder kv cache tensor during init, which will be replaced
172 # by bind_kv_cache
173 # this variable will not be accessed if use_direct_call is True
vllm/v1/attention/backends/flash_attn.py
+21 -15
@@ -485,6 +485,7 @@ class FlashAttentionImpl(AttentionImpl):
485 blocksparse_params: Optional[dict[str, Any]] = None,
486 logits_soft_cap: Optional[float] = None,
487 attn_type: AttentionType = AttentionType.DECODER,
488 + kv_sharing_target_layer_name: Optional[str] = None,
489 use_irope: bool = False,
490 ) -> None:
491 if blocksparse_params is not None:
@@ -506,6 +507,7 @@ class FlashAttentionImpl(AttentionImpl):
507 # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
508 logits_soft_cap = 0
509 self.logits_soft_cap = logits_soft_cap
510 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
511
512 assert self.num_heads % self.num_kv_heads == 0
513 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -569,22 +571,26 @@ class FlashAttentionImpl(AttentionImpl):
571 # performance to make sure it does not introduce any overhead.
572
573 num_actual_tokens = attn_metadata.num_actual_tokens
572 - # Reshape the input keys and values and store them in the cache.
573 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is
574 - # not padded. However, we don't need to do key[:num_actual_tokens] and
575 - # value[:num_actual_tokens] because the reshape_and_cache_flash op uses
576 - # the slot_mapping's shape to determine the number of actual tokens.
574 key_cache, value_cache = kv_cache.unbind(0)
578 - torch.ops._C_cache_ops.reshape_and_cache_flash(
579 - key,
580 - value,
581 - key_cache,
582 - value_cache,
583 - attn_metadata.slot_mapping,
584 - self.kv_cache_dtype,
585 - layer._k_scale,
586 - layer._v_scale,
587 - )
575 +
576 + if self.kv_sharing_target_layer_name is None:
577 + # Reshape the input keys and values and store them in the cache.
578 + # Skip this if sharing KV cache with an earlier attention layer.
579 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is
580 + # not padded. However, we don't need to do key[:num_actual_tokens]
581 + # and value[:num_actual_tokens] because the reshape_and_cache_flash
582 + # op uses the slot_mapping's shape to determine the number of
583 + # actual tokens.
584 + torch.ops._C_cache_ops.reshape_and_cache_flash(
585 + key,
586 + value,
587 + key_cache,
588 + value_cache,
589 + attn_metadata.slot_mapping,
590 + self.kv_cache_dtype,
591 + layer._k_scale,
592 + layer._v_scale,
593 + )
594
595 if self.kv_cache_dtype.startswith("fp8"):
596 key_cache = key_cache.view(torch.float8_e4m3fn)
vllm/v1/attention/backends/flashinfer.py
+21 -15
@@ -507,6 +507,7 @@ class FlashInferImpl(AttentionImpl):
507 blocksparse_params: Optional[dict[str, Any]] = None,
508 logits_soft_cap: Optional[float] = None,
509 attn_type: AttentionType = AttentionType.DECODER,
510 + kv_sharing_target_layer_name: Optional[int] = None,
511 ) -> None:
512 self.num_heads = num_heads
513 self.head_size = head_size
@@ -521,6 +522,7 @@ class FlashInferImpl(AttentionImpl):
522 self.sliding_window = (sliding_window - 1, 0)
523 self.kv_cache_dtype = kv_cache_dtype
524 self.logits_soft_cap = logits_soft_cap
525 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
526
527 assert self.num_heads % self.num_kv_heads == 0
528 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -568,21 +570,25 @@ class FlashInferImpl(AttentionImpl):
570 # performance to make sure it does not introduce any overhead.
571
572 num_actual_tokens = attn_metadata.num_actual_tokens
571 - # Reshape the input keys and values and store them in the cache.
572 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is
573 - # not padded. However, we don't need to do key[:num_actual_tokens] and
574 - # value[:num_actual_tokens] because the reshape_and_cache_flash op uses
575 - # the slot_mapping's shape to determine the number of actual tokens.
576 - torch.ops._C_cache_ops.reshape_and_cache_flash(
577 - key,
578 - value,
579 - kv_cache[:, 0],
580 - kv_cache[:, 1],
581 - attn_metadata.slot_mapping,
582 - self.kv_cache_dtype,
583 - layer._k_scale,
584 - layer._v_scale,
585 - )
573 +
574 + if self.kv_sharing_target_layer_name is None:
575 + # Reshape the input keys and values and store them in the cache.
576 + # Skip this if sharing KV cache with an earlier attention layer.
577 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is
578 + # not padded. However, we don't need to do key[:num_actual_tokens]
579 + # and value[:num_actual_tokens] because the reshape_and_cache_flash
580 + # op uses the slot_mapping's shape to determine the number of
581 + # actual tokens.
582 + torch.ops._C_cache_ops.reshape_and_cache_flash(
583 + key,
584 + value,
585 + kv_cache[:, 0],
586 + kv_cache[:, 1],
587 + attn_metadata.slot_mapping,
588 + self.kv_cache_dtype,
589 + layer._k_scale,
590 + layer._v_scale,
591 + )
592
593 window_left = (self.sliding_window[0]
594 if self.sliding_window is not None else -1)
vllm/v1/attention/backends/mla/common.py
+4
@@ -586,6 +586,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
586 blocksparse_params: Optional[dict[str, Any]],
587 logits_soft_cap: Optional[float],
588 attn_type: str,
589 + kv_sharing_target_layer_name: Optional[str],
590 # MLA Specific Arguments
591 q_lora_rank: Optional[int],
592 kv_lora_rank: int,
@@ -595,6 +596,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
596 v_head_dim: int,
597 kv_b_proj: ColumnParallelLinear,
598 ) -> None:
599 + if kv_sharing_target_layer_name is not None:
600 + raise NotImplementedError("KV sharing is not supported for MLA")
601 +
602 self.num_heads = num_heads
603 self.head_size = head_size
604 self.scale = float(scale)
vllm/v1/attention/backends/mla/flashmla.py
+2 -1
@@ -93,12 +93,13 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
93 blocksparse_params: Optional[dict[str, Any]],
94 logits_soft_cap: Optional[float],
95 attn_type: str,
96 + kv_sharing_target_layer_name: Optional[str],
97 # MLA Specific Arguments
98 **mla_args) -> None:
99 super().__init__(num_heads, head_size, scale, num_kv_heads,
100 alibi_slopes, sliding_window, kv_cache_dtype,
101 blocksparse_params, logits_soft_cap, attn_type,
101 - **mla_args)
102 + kv_sharing_target_layer_name, **mla_args)
103
104 assert is_flashmla_supported(), \
105 "FlashMLA is not supported on this device"
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+2 -1
@@ -139,12 +139,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
139 blocksparse_params: Optional[dict[str, Any]],
140 logits_soft_cap: Optional[float],
141 attn_type: str,
142 + kv_sharing_target_layer_name: Optional[str],
143 # MLA Specific Arguments
144 **mla_args) -> None:
145 super().__init__(num_heads, head_size, scale, num_kv_heads,
146 alibi_slopes, sliding_window, kv_cache_dtype,
147 blocksparse_params, logits_soft_cap, attn_type,
147 - **mla_args)
148 + kv_sharing_target_layer_name, **mla_args)
149 assert (num_heads == 16 or num_heads == 128), (
150 f"Aiter MLA only supports 16 or 128 number of heads.\n"
151 f"Provided {num_heads} number of heads.\n"
vllm/v1/attention/backends/mla/triton_mla.py
+2 -1
@@ -41,12 +41,13 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
41 blocksparse_params: Optional[dict[str, Any]],
42 logits_soft_cap: Optional[float],
43 attn_type: str,
44 + kv_sharing_target_layer_name: Optional[str],
45 # MLA Specific Arguments
46 **mla_args) -> None:
47 super().__init__(num_heads, head_size, scale, num_kv_heads,
48 alibi_slopes, sliding_window, kv_cache_dtype,
49 blocksparse_params, logits_soft_cap, attn_type,
49 - **mla_args)
50 + kv_sharing_target_layer_name, **mla_args)
51
52 unsupported_features = [
53 alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
vllm/v1/attention/backends/pallas.py
+5 -1
@@ -113,6 +113,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
113 blocksparse_params: Optional[dict[str, Any]] = None,
114 logits_soft_cap: Optional[float] = None,
115 attn_type: str = AttentionType.DECODER,
116 + kv_sharing_target_layer_name: Optional[int] = None,
117 use_irope: bool = False,
118 ) -> None:
119 if use_irope:
@@ -128,6 +129,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
129 self.num_kv_heads = num_kv_heads
130 self.sliding_window = sliding_window
131 self.logits_soft_cap = logits_soft_cap
132 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
133
134 assert self.num_heads % self.num_kv_heads == 0
135 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -181,7 +183,9 @@ class PallasAttentionBackendImpl(AttentionImpl):
183 num_tokens, hidden_size = query.shape
184 query = query.view(num_tokens, self.num_heads, self.head_size)
185
184 - if kv_cache.numel() > 0:
186 + if self.kv_sharing_target_layer_name is None and kv_cache.numel() > 0:
187 + # Write input keys and values to the KV cache.
188 + # Skip this if sharing KV cache with an earlier attention layer.
189 slot_mapping = attn_metadata.slot_mapping
190 write_to_kv_cache(key, value, kv_cache, slot_mapping)
191
vllm/v1/attention/backends/triton_attn.py
+28 -23
@@ -88,6 +88,7 @@ class TritonAttentionImpl(AttentionImpl):
88 blocksparse_params: Optional[dict[str, Any]] = None,
89 logits_soft_cap: Optional[float] = None,
90 attn_type: AttentionType = AttentionType.DECODER,
91 + kv_sharing_target_layer_name: Optional[int] = None,
92 use_irope: bool = False,
93 ) -> None:
94 if blocksparse_params is not None:
@@ -109,6 +110,7 @@ class TritonAttentionImpl(AttentionImpl):
110 # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
111 logits_soft_cap = 0
112 self.logits_soft_cap = logits_soft_cap
113 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
114
115 self.use_irope = use_irope
116
@@ -178,31 +180,34 @@ class TritonAttentionImpl(AttentionImpl):
180 if use_prefill_decode_attn:
181 key_cache, value_cache = PagedAttention.split_kv_cache(
182 kv_cache, self.num_kv_heads, self.head_size)
181 -
182 - # Reshape the input keys and values and store them in the cache.
183 - PagedAttention.write_to_paged_cache(
184 - key,
185 - value,
186 - key_cache,
187 - value_cache,
188 - attn_metadata.slot_mapping,
189 - self.kv_cache_dtype,
190 - layer._k_scale,
191 - layer._v_scale,
192 - )
193 -
183 else:
184 key_cache, value_cache = kv_cache.unbind(0)
196 - torch.ops._C_cache_ops.reshape_and_cache_flash(
197 - key,
198 - value,
199 - key_cache,
200 - value_cache,
201 - attn_metadata.slot_mapping,
202 - self.kv_cache_dtype,
203 - layer._k_scale,
204 - layer._v_scale,
205 - )
185 +
186 + if self.kv_sharing_target_layer_name is None:
187 + # Reshape the input keys and values and store them in the cache.
188 + # Skip this if sharing KV cache with an earlier attention layer.
189 + if use_prefill_decode_attn:
190 + PagedAttention.write_to_paged_cache(
191 + key,
192 + value,
193 + key_cache,
194 + value_cache,
195 + attn_metadata.slot_mapping,
196 + self.kv_cache_dtype,
197 + layer._k_scale,
198 + layer._v_scale,
199 + )
200 + else:
201 + torch.ops._C_cache_ops.reshape_and_cache_flash(
202 + key,
203 + value,
204 + key_cache,
205 + value_cache,
206 + attn_metadata.slot_mapping,
207 + self.kv_cache_dtype,
208 + layer._k_scale,
209 + layer._v_scale,
210 + )
211
212 if self.kv_cache_dtype.startswith("fp8"):
213 key_cache = key_cache.view(self.fp8_dtype)
vllm/v1/attention/backends/utils.py
+33
@@ -17,3 +17,36 @@ class CommonAttentionMetadata:
17 seq_lens: torch.Tensor
18 """(batch_size,), the length of each request including both computed tokens
19 and newly scheduled tokens"""
20 +
21 +
22 +def validate_kv_sharing_target(current_layer_name, target_layer_name,
23 + static_forward_context):
24 + error_msg = (f"Specified KV sharing target layer for {current_layer_name} "
25 + f"is not valid: target layer {target_layer_name} ")
26 +
27 + if current_layer_name == target_layer_name:
28 + raise ValueError(error_msg +
29 + "cannot be the same as the current layer.")
30 +
31 + if target_layer_name not in static_forward_context:
32 + from vllm.model_executor.models.utils import extract_layer_index
33 +
34 + # If target layer name is not in the static fwd context, it means either
35 + # a) the target layer does not come BEFORE the current layer, or
36 + # b) the target layer is not an Attention layer that exists in the model
37 + current_layer_idx = extract_layer_index(current_layer_name)
38 + target_layer_idx = extract_layer_index(target_layer_name)
39 + if current_layer_idx <= target_layer_idx:
40 + raise ValueError(error_msg + "must come before the current layer.")
41 + else:
42 + raise ValueError(error_msg +
43 + "is not a valid Attention layer in the model.")
44 +
45 + # Currently KV sharing is only supported between layers of the same type
46 + target_layer_attn_type = static_forward_context[
47 + target_layer_name].attn_type
48 + expected = static_forward_context[current_layer_name].attn_type
49 + if target_layer_attn_type != expected:
50 + raise ValueError(
51 + error_msg +
52 + f"must be the same type as the current layer ({expected}).")
vllm/v1/worker/gpu_model_runner.py
+29 -2
@@ -59,8 +59,8 @@ from vllm.v1.worker.block_table import BlockTable
59 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
60 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
61
62 -from .utils import (gather_mm_placeholders, sanity_check_mm_encoder_outputs,
63 - scatter_mm_placeholders)
62 +from .utils import (gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
63 + sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
64
65 if TYPE_CHECKING:
66 import xgrammar as xgr
@@ -276,6 +276,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
276 pin_memory=self.pin_memory)
277 self.seq_lens_np = self.seq_lens_cpu.numpy()
278
279 + # Layer pairings for cross-layer KV sharing.
280 + # If an Attention layer `layer_name` is in the keys of this dict, it
281 + # means this layer will perform attention using the keys and values
282 + # from the KV cache of `shared_kv_cache_layers[layer_name]`.
283 + self.shared_kv_cache_layers: dict[str, str] = {}
284 +
285 def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> bool:
286 """
287 Update the order of requests in the batch based on the attention
@@ -2097,6 +2103,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2103 # KV cache specs.
2104 raise ValueError("Unknown KV cache spec type.")
2105
2106 + # Setup `kv_cache_config` and `kv_caches` for models
2107 + # with cross-layer KV sharing
2108 + if self.shared_kv_cache_layers:
2109 + initialize_kv_cache_for_kv_sharing(
2110 + self.shared_kv_cache_layers,
2111 + kv_cache_config.kv_cache_groups,
2112 + kv_caches,
2113 + )
2114 +
2115 if self.speculative_config and self.speculative_config.use_eagle():
2116 assert isinstance(self.drafter, EagleProposer)
2117 # validate all draft model layers belong to the same kv cache
@@ -2125,6 +2140,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2140 use_mla = self.vllm_config.model_config.use_mla
2141 kv_cache_spec: dict[str, KVCacheSpec] = {}
2142 for layer_name, attn_module in layers.items():
2143 + if (kv_tgt_layer :=
2144 + attn_module.kv_sharing_target_layer_name) is not None:
2145 + # The layer doesn't need its own KV cache and will use that of
2146 + # the target layer. We skip creating a KVCacheSpec for it, so
2147 + # that KV cache management logic will act as this layer does
2148 + # not exist, and doesn't allocate KV cache for the layer. This
2149 + # enables the memory saving of cross-layer kv sharing, allowing
2150 + # a given amount of memory to accommodate longer context lengths
2151 + # or enable more requests to be processed simultaneously.
2152 + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
2153 + continue
2154 +
2155 # TODO: Support other attention modules, e.g., cross-attention
2156 if attn_module.attn_type == AttentionType.DECODER:
2157 if attn_module.sliding_window is not None:
vllm/v1/worker/tpu_model_runner.py
+29 -1
@@ -44,7 +44,8 @@ from vllm.v1.utils import bind_kv_cache
44 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
45 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
46
47 -from .utils import sanity_check_mm_encoder_outputs
47 +from .utils import (initialize_kv_cache_for_kv_sharing,
48 + sanity_check_mm_encoder_outputs)
49
50 if TYPE_CHECKING:
51 from vllm.v1.core.sched.output import SchedulerOutput
@@ -238,6 +239,12 @@ class TPUModelRunner(LoRAModelRunnerMixin):
239 self.num_reqs_paddings = _get_req_paddings(
240 min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs)
241
242 + # Layer pairings for cross-layer KV sharing.
243 + # If an Attention layer `layer_name` is in the keys of this dict, it
244 + # means this layer will perform attention using the keys and values
245 + # from the KV cache of `shared_kv_cache_layers[layer_name]`.
246 + self.shared_kv_cache_layers: dict[str, str] = {}
247 +
248 # tensors for structured decoding
249 self.grammar_bitmask_cpu = torch.zeros(
250 (self.max_num_reqs, cdiv(self.vocab_size, 32)),
@@ -455,6 +462,18 @@ class TPUModelRunner(LoRAModelRunnerMixin):
462 block_size = self.vllm_config.cache_config.block_size
463 kv_cache_spec: dict[str, KVCacheSpec] = {}
464 for layer_name, attn_module in layers.items():
465 + if (kv_tgt_layer :=
466 + attn_module.kv_sharing_target_layer_name) is not None:
467 + # The layer doesn't need its own KV cache and will use that of
468 + # the target layer. We skip creating a KVCacheSpec for it, so
469 + # that KV cache management logic will act as this layer does
470 + # not exist, and doesn't allocate KV cache for the layer. This
471 + # enables the memory saving of cross-layer kv sharing, allowing
472 + # a given amount of memory to accommodate longer context lengths
473 + # or enable more requests to be processed simultaneously.
474 + self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
475 + continue
476 +
477 if attn_module.attn_type == AttentionType.DECODER:
478 if attn_module.sliding_window is not None:
479 kv_cache_spec[layer_name] = SlidingWindowSpec(
@@ -1376,6 +1395,15 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1395 else:
1396 raise NotImplementedError
1397
1398 + # Setup `kv_cache_config` and `kv_caches` for models
1399 + # with cross-layer KV sharing
1400 + if self.shared_kv_cache_layers:
1401 + initialize_kv_cache_for_kv_sharing(
1402 + self.shared_kv_cache_layers,
1403 + kv_cache_config.kv_cache_groups,
1404 + kv_caches,
1405 + )
1406 +
1407 bind_kv_cache(
1408 kv_caches,
1409 self.vllm_config.compilation_config.static_forward_context,
vllm/v1/worker/utils.py
+36
@@ -4,6 +4,8 @@ from typing import Optional
4
5 import torch
6
7 +from vllm.v1.kv_cache_interface import KVCacheGroupSpec
8 +
9
10 def sanity_check_mm_encoder_outputs(
11 mm_embeddings: object,
@@ -73,3 +75,37 @@ def gather_mm_placeholders(
75 return placeholders
76
77 return placeholders[is_embed]
78 +
79 +
80 +def initialize_kv_cache_for_kv_sharing(
81 + shared_kv_cache_layers: dict[str, str],
82 + kv_cache_groups: list[KVCacheGroupSpec],
83 + kv_caches: dict[str, torch.Tensor],
84 +) -> None:
85 + """
86 + Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
87 + for layers that do not allocate its own KV cache, based on the mapping in
88 + `shared_kv_cache_layers`. Adds these layers to the corresponding KV cache
89 + group, which is needed to ensure that attention metadata is assigned later.
90 +
91 + Args:
92 + shared_kv_cache_layers: Layer pairings for cross-layer KV sharing.
93 + If an Attention layer `layer_name` is in the keys of this dict, it
94 + means this layer will perform attention using the keys and values
95 + from the KV cache of `shared_kv_cache_layers[layer_name]`.
96 + kv_cache_groups: The KV cache groups of the model.
97 + kv_caches: The allocated kv_caches with layer names as keys.
98 + Note that layers in shared_kv_cache_layers.keys() are not
99 + originally included as it only contains layers which have its own
100 + KV cache allocation.
101 + """
102 + # Record index of KV cache group for each layer that allocates a KV cache.
103 + layer_to_kv_cache_group_idx: dict[str, int] = {}
104 + for i, kv_cache_group in enumerate(kv_cache_groups):
105 + for layer_name in kv_cache_group.layer_names:
106 + layer_to_kv_cache_group_idx[layer_name] = i
107 +
108 + for layer_name, target_layer_name in shared_kv_cache_layers.items():
109 + kv_caches[layer_name] = kv_caches[target_layer_name]
110 + group_idx = layer_to_kv_cache_group_idx[target_layer_name]
111 + kv_cache_groups[group_idx].layer_names.append(layer_name)