@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e33081cee
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 2, 18:15 UTC

[Attention] Make FlexAttention and FlashAttention use num-blocks first layouts (#42095)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

Lucas Wilkinson committed May 27, 2026 at 03:55 UTC 7e33081cee7bd04f9941caebff2901436e2171ee
20 files changed +212 -392
tests/v1/attention/test_attention_backends.py
+30 -19
@@ -136,7 +136,8 @@ def create_and_prepopulate_kv_cache(
136 block_table = common_attn_metadata.block_table_tensor
137 slot_mapping = common_attn_metadata.slot_mapping
138
139 - # Create KV cache
139 + # Create KV cache and populate in (2, num_blocks, ...) layout for easy
140 + # flat indexing, then transpose to (num_blocks, 2, ...) layout.
141 kv_cache = torch.zeros(
142 2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
143 )
@@ -155,6 +156,9 @@ def create_and_prepopulate_kv_cache(
156 # Stay block aligned and allocate enough blocks for the new tokens
157 start_block_idx += cdiv(int(seq_lens[i]), block_size)
158
159 + # Transpose to (num_blocks, 2, ...) layout
160 + kv_cache = kv_cache.transpose(0, 1).contiguous()
161 +
162 blocks_end = start_block_idx
163
164 # Permute the context blocks (excluding block 0 which is null)
@@ -168,7 +172,7 @@ def create_and_prepopulate_kv_cache(
172 inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
173 # Add 1 to account for starting from block 1
174 inv_perm[1:] = torch.argsort(perm) + 1
171 - kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]
175 + kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]
176
177 # Construct the right block table
178 # Start from block_id=1 since block_id=0 is considered the null block
@@ -473,28 +477,35 @@ def _test_backend_correctness(
477 # Note: flex_attention has known Triton kernel compatibility issues
478 # with test infrastructures
479 for backend_name in backend_to_test:
476 - # FlashAttentionm + FlexAttention:
477 - # [2, num_blocks, block_size, num_kv_heads, head_size]
478 - # FlashInfer + Triton:
479 - # [num_blocks, 2, block_size, num_kv_heads, head_size]
480 - # Select the appropriate KV cache format for each backend
481 - kv_cache_for_backend = kv_cache
480 reset_kv_cache_layout = False
483 - if backend_name in (
484 - AttentionBackendEnum.FLASHINFER,
485 - AttentionBackendEnum.TRITON_ATTN,
486 - ):
487 - kv_cache_for_backend = kv_cache.transpose(0, 1)
481 +
482 + # Resolve backend class for both enum and string names.
483 + actual_backend = backend_name
484 + if backend_name == "FLEX_ATTENTION_SLOW":
485 + actual_backend = AttentionBackendEnum.FLEX_ATTENTION
486 + if hasattr(actual_backend, "get_class"):
487 + backend_cls = actual_backend.get_class()
488 + else:
489 + backend_cls = None
490
491 if backend_name == AttentionBackendEnum.FLASHINFER:
490 - # For FlashInfer default to HND layout and
491 - kv_cache_for_backend = (
492 - kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
493 - )
492 set_kv_cache_layout("HND")
493 reset_kv_cache_layout = True
496 - elif backend_name == AttentionBackendEnum.TRITON_ATTN:
497 - kv_cache_for_backend = kv_cache_for_backend.contiguous()
494 +
495 + # Apply stride order like runtime does in
496 + # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
497 + # layout, make contiguous, then permute to logical layout.
498 + kv_cache_for_backend = kv_cache
499 + if backend_cls is not None:
500 + try:
501 + stride_order = backend_cls.get_kv_cache_stride_order()
502 + except (AttributeError, NotImplementedError):
503 + stride_order = tuple(range(kv_cache.ndim))
504 + if stride_order != tuple(range(kv_cache.ndim)):
505 + inv_order = [stride_order.index(i) for i in range(len(stride_order))]
506 + kv_cache_for_backend = (
507 + kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
508 + )
509
510 try:
511 backend_output = run_attention_backend(
tests/v1/attention/test_kv_head_stride_canonicalization.py
+1 -1
@@ -19,7 +19,7 @@ size-1 dimensions via torch.as_strided — zero-copy.
19
20 The degenerate stride manifests at different positions in different backends:
21 - FlashInfer: stride(-3) after kv_cache.permute() → shape [..., 1, B, D]
22 -- FlashAttention: stride(-2) after kv_cache.unbind(0) → shape [N, B, 1, D]
22 +- FlashAttention: stride(-2) after kv_cache.unbind(1) → shape [N, B, 1, D]
23 """
24
25 import torch
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+17 -77
@@ -1,14 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import defaultdict
4 -from unittest.mock import MagicMock, patch
4 +from unittest.mock import MagicMock
5
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10 from vllm.utils.torch_utils import get_dtype_size
11 -from vllm.v1.attention.backend import AttentionBackend
11 from vllm.v1.attention.backends.registry import AttentionBackendEnum
12 from vllm.v1.attention.backends.utils import set_kv_cache_layout
13 from vllm.v1.kv_cache_interface import (
@@ -85,15 +84,6 @@ def _allocate_and_reshape_kv_caches(
84 set_kv_cache_layout(None)
85
86
88 -def _make_mock_layer(backend_cls: type[AttentionBackend]):
89 - """
90 - Create a mock AttentionLayerBase whose get_attn_backend returns backend_cls.
91 - """
92 - layer = MagicMock()
93 - layer.get_attn_backend.return_value = backend_cls
94 - return layer
95 -
96 -
87 def _make_worker(kv_cache_config: KVCacheConfig):
88 """
89 Create an OffloadingConnectorWorker with mocked dependencies.
@@ -119,11 +109,7 @@ def _make_worker(kv_cache_config: KVCacheConfig):
109
110
111 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
122 -@patch(
123 - "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
124 - ".worker.get_layers_from_vllm_config"
125 -)
126 -def test_register_kv_caches(mock_get_layers, backend):
112 +def test_register_kv_caches(backend):
113 """Test register_kv_caches with multiple groups covering all layer types.
114
115 Creates one FullAttention group, one MLA group, one Mamba group, and
@@ -287,13 +273,6 @@ def test_register_kv_caches(mock_get_layers, backend):
273 device=torch.device("cuda:0"),
274 )
275
290 - mock_layers: dict[str, MagicMock] = {}
291 - for layer_name in attn_layer_names:
292 - mock_layers[layer_name] = _make_mock_layer(backend_cls)
293 - for layer_name in mla_layer_names:
294 - mock_layers[layer_name] = _make_mock_layer(DeepseekV32IndexerBackend)
295 - mock_get_layers.return_value = mock_layers
296 -
276 worker, spec = _make_worker(kv_cache_config)
277 worker.register_kv_caches(kv_caches)
278
@@ -360,11 +339,7 @@ def test_register_kv_caches(mock_get_layers, backend):
339
340
341 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
363 -@patch(
364 - "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
365 - ".worker.get_layers_from_vllm_config"
366 -)
367 -def test_register_kv_caches_uniform_type(mock_get_layers, backend):
342 +def test_register_kv_caches_uniform_type(backend):
343 """Test register_kv_caches with UniformTypeKVCacheSpecs.
344
345 Two attention layers use the same backend but different num_kv_heads,
@@ -441,64 +416,29 @@ def test_register_kv_caches_uniform_type(mock_get_layers, backend):
416 device=torch.device("cuda:0"),
417 )
418
444 - mock_get_layers.return_value = {
445 - layer_a: _make_mock_layer(backend_cls),
446 - layer_b: _make_mock_layer(backend_cls),
447 - }
448 -
419 worker, spec = _make_worker(kv_cache_config)
420 worker.register_kv_caches(kv_caches)
421
422 canonical = spec.get_handlers.call_args[0][0]
423 assert isinstance(canonical, CanonicalKVCaches)
424
455 - unbinds = backend_cls.get_name() in ("FLASH_ATTN", "FLEX_ATTENTION")
456 - tensors_per_layer = 2 if unbinds else 1
457 -
425 for block_tensor in canonical.tensors:
426 assert block_tensor.tensor.dtype == torch.int8
427
428 # Single group with refs from both layers
429 assert len(canonical.group_data_refs) == 1
430 group_refs = canonical.group_data_refs[0]
464 - assert len(group_refs) == 2 * tensors_per_layer
465 -
466 - if unbinds:
467 - half_a = spec_a.page_size_bytes // 2
468 - half_b = spec_b.page_size_bytes // 2
469 -
470 - assert len(canonical.tensors) == 4
471 - assert canonical.tensors[0].page_size_bytes == half_a
472 - assert canonical.tensors[1].page_size_bytes == half_a
473 - assert canonical.tensors[2].page_size_bytes == half_b
474 - assert canonical.tensors[3].page_size_bytes == half_b
475 - assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, half_a)
476 - assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, half_a)
477 - assert canonical.tensors[2].tensor.shape == (NUM_BLOCKS, half_b)
478 - assert canonical.tensors[3].tensor.shape == (NUM_BLOCKS, half_b)
479 -
480 - assert group_refs[0] == CanonicalKVCacheRef(
481 - tensor_idx=0, page_size_bytes=half_a
482 - )
483 - assert group_refs[1] == CanonicalKVCacheRef(
484 - tensor_idx=1, page_size_bytes=half_a
485 - )
486 - assert group_refs[2] == CanonicalKVCacheRef(
487 - tensor_idx=2, page_size_bytes=half_b
488 - )
489 - assert group_refs[3] == CanonicalKVCacheRef(
490 - tensor_idx=3, page_size_bytes=half_b
491 - )
492 - else:
493 - assert len(canonical.tensors) == 2
494 - assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
495 - assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
496 - assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
497 - assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
498 -
499 - assert group_refs[0] == CanonicalKVCacheRef(
500 - tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
501 - )
502 - assert group_refs[1] == CanonicalKVCacheRef(
503 - tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
504 - )
431 + assert len(group_refs) == 2
432 +
433 + assert len(canonical.tensors) == 2
434 + assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
435 + assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
436 + assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
437 + assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
438 +
439 + assert group_refs[0] == CanonicalKVCacheRef(
440 + tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
441 + )
442 + assert group_refs[1] == CanonicalKVCacheRef(
443 + tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
444 + )
tests/v1/kv_connector/unit/offloading_connector/utils.py
+4 -20
@@ -3,7 +3,7 @@
3 from collections.abc import Iterable, Iterator
4 from dataclasses import dataclass
5 from typing import Any
6 -from unittest.mock import MagicMock, patch
6 +from unittest.mock import MagicMock
7
8 import pytest
9 import torch
@@ -25,7 +25,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
25 )
26 from vllm.forward_context import ForwardContext
27 from vllm.utils.hashing import sha256
28 -from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
28 from vllm.v1.core.kv_cache_utils import (
29 get_request_block_hasher,
30 init_none_hash,
@@ -239,37 +238,22 @@ class RequestRunner:
238
239 # register worker kv_caches to enable OffloadingWorker creations
240 # set_current_vllm_config is needed for get_kv_cache_layout() to work
242 - # Mock get_layers_from_vllm_config so that mock layer names
243 - # resolve to layers whose get_attn_backend() returns
244 - # FlashAttentionBackend.
245 - def _mock_get_layers(_vllm_config, _layer_type, layer_names):
246 - mock_layer = MagicMock()
247 - mock_layer.get_attn_backend.return_value = FlashAttentionBackend
248 - return {name: mock_layer for name in layer_names}
249 -
241 kv_caches: dict[str, torch.Tensor] = {}
242 for group in kv_cache_groups:
243 spec = group.kv_cache_spec
244 for layer_name in group.layer_names:
245 # Shape follows FlashAttention layout:
255 - # (2, num_blocks, block_size, num_kv_heads, head_size)
246 + # Shape: (num_blocks, 2, block_size, num_kv_heads, head_size)
247 kv_caches[layer_name] = torch.empty(
257 - 2,
248 num_gpu_blocks,
249 + 2,
250 spec.block_size,
251 spec.num_kv_heads,
252 spec.head_size,
253 dtype=spec.dtype,
254 )
255
265 - with (
266 - set_current_vllm_config(vllm_config),
267 - patch(
268 - "vllm.distributed.kv_transfer.kv_connector.v1"
269 - ".offloading.worker.get_layers_from_vllm_config",
270 - side_effect=_mock_get_layers,
271 - ),
272 - ):
256 + with set_current_vllm_config(vllm_config):
257 self.worker_connector.register_kv_caches(kv_caches)
258
259 # extract connector of scheduler
tests/v1/kv_connector/unit/test_mooncake_connector.py
+65 -36
@@ -369,15 +369,31 @@ async def test_kv_producer(monkeypatch):
369 with patch.object(
370 prefill_worker, "_send_blocks", return_value=0
371 ) as mock_send_blocks:
372 + # With blocks-first layout, each block is virtually split
373 + # into K and V halves, producing non-coalesced transfers.
374 + kv_half = block_len // 2
375 +
376 + def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
377 + """Build expected (src_ptrs, dst_ptrs, lengths) for
378 + virtual-split K/V transfers."""
379 + src_ptrs, dst_ptrs, lengths = [], [], []
380 + for kv_offset in (0, kv_half):
381 + for sb, db in zip(src_blocks, dst_blocks):
382 + src_ptrs.append(src_base + sb * block_len + kv_offset)
383 + dst_ptrs.append(dst_base + db * block_len + kv_offset)
384 + lengths.append(kv_half)
385 + return src_ptrs, dst_ptrs, lengths
386 +
387 # Normal case: 2 blocks to 2 blocks
373 - # Worker processes the consumer's request
388 await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
375 - # Verify transfer parameters are correct
376 - src_ptr = 0x1000 + 10 * block_len
377 - dst_ptr = 0x2000 + 20 * block_len
378 - length = 2 * block_len
389 + src, dst, lens = expected_split_transfers(
390 + 0x1000, 0x2000, [10, 11], [20, 21]
391 + )
392 mock_send_blocks.assert_called_once_with(
380 - "consumer-host:54321", [src_ptr], [dst_ptr], [length]
393 + "consumer-host:54321",
394 + src,
395 + dst,
396 + lens,
397 )
398 mock_socket.send_multipart.assert_called_once()
399
@@ -404,11 +420,12 @@ async def test_kv_producer(monkeypatch):
420 # Worker processes the consumer's request
421 await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
422 # Verify transfer parameters are correct: 11 to 20
407 - src_ptr = 0x1000 + 11 * block_len
408 - dst_ptr = 0x2000 + 20 * block_len
409 - length = 1 * block_len
423 + src, dst, lens = expected_split_transfers(0x1000, 0x2000, [11], [20])
424 mock_send_blocks.assert_called_once_with(
411 - "consumer-host:54321", [src_ptr], [dst_ptr], [length]
425 + "consumer-host:54321",
426 + src,
427 + dst,
428 + lens,
429 )
430 mock_socket.send_multipart.assert_called_once()
431
@@ -618,18 +635,14 @@ def test_register_kv_caches():
635
636 mock_batch_register.assert_called_once()
637 registered_ptrs, registered_lens = mock_batch_register.call_args[0]
621 - expected_ptrs = {
622 - tensor.data_ptr()
623 - for kv_pair in kv_caches.values()
624 - for tensor in kv_pair
625 - }
638 + expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()}
639 assert set(registered_ptrs) == expected_ptrs
627 - assert set(registered_lens) == {tensor1[0].nbytes}
640 + assert set(registered_lens) == {tensor1.nbytes}
641
642 # Verify block_len_per_layer is set correctly.
643 assert len(worker.block_len_per_layer) == len(registered_ptrs)
644 for bl in worker.block_len_per_layer:
632 - assert bl == tensor1[0].nbytes // tensor1.shape[1]
645 + assert bl == tensor1.nbytes // tensor1.shape[0]
646
647
648 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
@@ -791,33 +804,49 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
804 # Flatten nested per-group block IDs for assertions
805 flat_local = [b for g in local_block_ids for b in g]
806 flat_remote = [b for g in remote_block_ids for b in g]
807 + num_blocks = len(flat_local)
808
795 - # Heterogeneous TP: blocks cannot be coalesced because
796 - # local and remote block_lens differ
797 - assert len(src_ptrs) == len(flat_local)
798 - assert len(dst_ptrs) == len(flat_local)
799 - assert len(lengths) == len(flat_local)
809 + # With blocks-first layout, virtual split halves block
810 + # lengths and doubles transfer regions (K + V).
811 + local_kv_block_len = local_block_len // 2
812 + remote_kv_block_len = remote_block_len // 2
813
801 - # Compute expected offsets based on TP ratio
814 + assert len(src_ptrs) == 2 * num_blocks
815 + assert len(dst_ptrs) == 2 * num_blocks
816 + assert len(lengths) == 2 * num_blocks
817 +
818 + # Compute expected offsets using kv_block_len
819 if d_tp_size <= P_TP_SIZE:
820 tp_ratio = P_TP_SIZE // d_tp_size
821 expected_src_off = 0
805 - expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
806 - expected_xfer_len = local_block_len
822 + expected_dst_off = (P_TP_RANK % tp_ratio) * local_kv_block_len
823 + expected_xfer_len = local_kv_block_len
824 else:
825 ratio_abs = d_tp_size // P_TP_SIZE
809 - expected_src_off = (d_rank % ratio_abs) * remote_block_len
826 + expected_src_off = (d_rank % ratio_abs) * remote_kv_block_len
827 expected_dst_off = 0
811 - expected_xfer_len = remote_block_len
812 -
813 - for idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
814 - assert src_ptrs[idx] == (
815 - 0x1000 + lblk * local_block_len + expected_src_off
816 - )
817 - assert dst_ptrs[idx] == (
818 - 0x2000 + rblk * remote_block_len + expected_dst_off
819 - )
820 - assert lengths[idx] == expected_xfer_len
828 + expected_xfer_len = remote_kv_block_len
829 +
830 + # First num_blocks entries are K region,
831 + # next num_blocks are V region.
832 + for region_idx in range(2):
833 + local_region_base = 0x1000 + region_idx * local_kv_block_len
834 + remote_region_base = 0x2000 + region_idx * remote_kv_block_len
835 + for blk_idx, (lblk, rblk) in enumerate(
836 + zip(flat_local, flat_remote)
837 + ):
838 + idx = region_idx * num_blocks + blk_idx
839 + assert src_ptrs[idx] == (
840 + local_region_base
841 + + lblk * local_block_len
842 + + expected_src_off
843 + )
844 + assert dst_ptrs[idx] == (
845 + remote_region_base
846 + + rblk * remote_block_len
847 + + expected_dst_off
848 + )
849 + assert lengths[idx] == expected_xfer_len
850
851 # Verify successful response sent back to consumer
852 mock_socket.send_multipart.assert_called_once()
tests/v1/kv_connector/unit/test_nixl_connector.py
+3 -2
@@ -1635,6 +1635,7 @@ def test_register_kv_caches(
1635 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1636 )
1637 is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
1638 + virtually_split = is_blocks_first and not connector.prefer_cross_layer_blocks
1639
1640 if connector.prefer_cross_layer_blocks:
1641 with set_current_vllm_config(vllm_config):
@@ -1665,7 +1666,7 @@ def test_register_kv_caches(
1666 ]
1667 expected_num_entries = 1
1668
1668 - expected_blocks_count = num_blocks * (2 if is_blocks_first else 1)
1669 + expected_blocks_count = num_blocks * (2 if virtually_split else 1)
1670
1671 kv_caches = {"all-layers": cross_layers_kv_cache}
1672 else:
@@ -1739,7 +1740,7 @@ def test_register_kv_caches(
1740 else:
1741 num_blocks = kv_cache_config.num_blocks
1742
1742 - if is_blocks_first:
1743 + if virtually_split:
1744 expected_block_len = expected_tensor_size // num_blocks // 2
1745 else:
1746 expected_block_len = expected_tensor_size // num_blocks
tests/v1/worker/test_gpu_model_runner.py
+1 -29
@@ -34,7 +34,6 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
34 from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
35 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
36 from vllm.v1.kv_cache_interface import (
37 - AttentionSpec,
37 FullAttentionSpec,
38 KVCacheConfig,
39 KVCacheGroupSpec,
@@ -44,7 +43,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
43 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
44 from vllm.v1.worker.gpu_input_batch import InputBatch
45 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
47 -from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
46 +from vllm.v1.worker.utils import select_common_block_size
47
48 BLOCK_SIZE = 16
49 NUM_BLOCKS = 10
@@ -1195,33 +1194,6 @@ def test_hybrid_attention_mamba_tensor_shapes():
1194 assert torch.equal(actual_ssm, expected_ssm)
1195
1196
1198 -def test_update_hybrid_attention_mamba_layout_with_num_block_2_rewrites_stride():
1199 - from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
1200 -
1201 - ambiguous_cache = torch.empty((2, 2, BLOCK_SIZE, 1, 8), dtype=torch.float16)
1202 - """Ambiguous, because both dims[0=kv_dim] and dims[1=num_blocks] == 2"""
1203 - hidden_size = ambiguous_cache.shape[2:].numel()
1204 - assert ambiguous_cache.stride()[:2] == (2 * hidden_size, hidden_size)
1205 -
1206 - attention_spec = AttentionSpec(
1207 - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=8, dtype=torch.float16
1208 - )
1209 - runner_stub = SimpleNamespace(
1210 - cache_config=SimpleNamespace(cache_dtype="auto"),
1211 - _kv_cache_spec_attn_group_iterator=lambda: iter(
1212 - [AttentionGroup(FlashAttentionBackend, ["attn"], attention_spec, 0)]
1213 - ),
1214 - )
1215 - GPUModelRunner._update_hybrid_attention_mamba_layout(
1216 - runner_stub, {"attn": ambiguous_cache}, [BLOCK_SIZE]
1217 - )
1218 -
1219 - assert ambiguous_cache.stride()[:2] == (hidden_size, 2 * hidden_size), """\
1220 - We expect _update_hybrid_attention_mamba_layout to re-stride the cache from:
1221 - (2, num_blocks) -> (num_blocks, 2), even when num_blocks==2,
1222 - which was ambiguous before get_kv_cache_block_dim was used"""
1223 -
1224 -
1197 def test_hybrid_block_table_initialization():
1198 """Test hybrid block table with different kernel and kvcache_manager block
1199 sizes."""
vllm/distributed/kv_transfer/kv_connector/utils.py
+10 -1
@@ -420,7 +420,7 @@ class TransferTopology:
420 head_size=1,
421 )
422 logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
423 - # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
423 + # Non-MLA backends caches have 5 dims [num_blocks, 2, H,N,D],
424 # we just mock num_blocks to 1 for the dimension check below.
425 # Hybrid SSM models assume a single blocks_first layout
426 self._is_kv_layout_blocks_first = self.is_mamba or (
@@ -484,6 +484,15 @@ class TransferTopology:
484 def cross_layers_blocks(self) -> bool:
485 return self._cross_layers_blocks
486
487 + @property
488 + def virtually_split_kv_in_blocks(self) -> bool:
489 + # Whether to logically split each block into K and V halves.
490 + # Applies when K/V are interleaved within each block (blocks-first),
491 + # but NOT when cross-layer blocks are used — cross-layer blocks have
492 + # per-layer K/V interleaving (L0_K, L0_V, L1_K, L1_V, ...) so a
493 + # simple half-split does not separate K from V.
494 + return self._is_kv_layout_blocks_first and not self._cross_layers_blocks
495 +
496 @property
497 def split_k_and_v(self) -> bool:
498 # Whether to register regions for K and V separately (when present).
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+10 -23
@@ -17,7 +17,6 @@ from vllm.logger import init_logger
17 from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
18 from vllm.utils.hashing import safe_hash
19 from vllm.v1.attention.backend import AttentionMetadata
20 -from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
20 from vllm.v1.core.sched.output import SchedulerOutput
21
22 if TYPE_CHECKING:
@@ -130,33 +129,24 @@ class ExampleConnector(KVConnectorBase_V1):
129
130 Args:
131 dst_kv_cache_layer (torch.Tensor): the destination KV cache
133 - layer. In shape [2, num_pages, page_size, xxx] if not
134 - using MLA, [num_pages, page_size, xxx] otherwise.
135 - src_kv_cache (torch.Tensor): the source KV cache. In shape
136 - [2, num_tokens, xxx] if not using MLA, [num_tokens, xxx]
137 - otherwise.
132 + layer. In shape [num_pages, page_size, xxx] for MLA,
133 + [num_pages, 2, page_size, xxx] otherwise.
134 + src_kv_cache (torch.Tensor): the source KV cache.
135 slot_mapping (torch.Tensor): the slot mapping. In shape
136 [num_tokens].
137 """
141 - dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
138 if isinstance(attn_metadata, MLACommonMetadata):
139 + dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
140 num_pages = dst_kv_cache_layer_shape[0]
141 page_size = dst_kv_cache_layer_shape[1]
142 dst_kv_cache_layer = dst_kv_cache_layer.reshape(
143 num_pages * page_size, -1
144 )
145 dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
149 - elif isinstance(attn_metadata, TritonAttentionMetadata):
146 + else:
147 block_idxs = slot_mapping // self._block_size
148 offsets = slot_mapping % self._block_size
149 dst_kv_cache_layer[block_idxs, :, offsets] = src_kv_cache
153 - else:
154 - num_pages = dst_kv_cache_layer_shape[1]
155 - page_size = dst_kv_cache_layer_shape[2]
156 - dst_kv_cache_layer = dst_kv_cache_layer.reshape(
157 - 2, num_pages * page_size, -1
158 - )
159 - dst_kv_cache_layer[:, slot_mapping, ...] = src_kv_cache
150
151 # Get the metadata
152 metadata: KVConnectorMetadata = self._get_connector_metadata()
@@ -234,18 +224,15 @@ class ExampleConnector(KVConnectorBase_V1):
224 ) -> torch.Tensor:
225 """Extract the KV cache from the layer.
226
237 - Assume the shape of the layer is (2, num_pages, page_size, xxx)
238 - if MLA is not used, and (num_pages, page_size, xxx) otherwise.
227 + Assume the shape of the layer is (num_pages, page_size, xxx)
228 + for MLA, and (num_pages, 2, page_size, xxx) otherwise.
229 """
230 if isinstance(attn_metadata, MLACommonMetadata):
231 num_pages, page_size = layer.shape[0], layer.shape[1]
232 return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
243 - elif isinstance(attn_metadata, TritonAttentionMetadata):
244 - block_idxs = slot_mapping // self._block_size
245 - offsets = slot_mapping % self._block_size
246 - return layer[block_idxs, :, offsets]
247 - num_pages, page_size = layer.shape[1], layer.shape[2]
248 - return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping, ...]
233 + block_idxs = slot_mapping // self._block_size
234 + offsets = slot_mapping % self._block_size
235 + return layer[block_idxs, :, offsets]
236
237 connector_metadata = self._get_connector_metadata()
238 assert isinstance(connector_metadata, ExampleConnectorMetadata)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+1 -1
@@ -1745,7 +1745,7 @@ class MooncakeConnectorWorker:
1745 return _expand_transfer_regions(
1746 base_addrs=base_addrs,
1747 block_lens=block_lens,
1748 - is_kv_layout_blocks_first=self.transfer_topo.is_kv_layout_blocks_first,
1748 + is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
1749 )
1750
1751 def _get_sender_transfer_plan(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+4 -4
@@ -943,7 +943,7 @@ class NixlConnectorWorker:
943 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
944 self.num_regions = len(caches_data)
945
946 - if self.transfer_topo.is_kv_layout_blocks_first:
946 + if self.transfer_topo.virtually_split_kv_in_blocks:
947 # NOTE (NickLucche) When FlashInfer is used, memory is registered
948 # with joint KV for each block. This minimizes the overhead in
949 # registerMem allowing faster descs queries. In order to be able to
@@ -1118,7 +1118,7 @@ class NixlConnectorWorker:
1118 addr = base_addr + block_offset
1119 result.append((addr, kv_block_len, self.device_id))
1120
1121 - if self.transfer_topo.is_kv_layout_blocks_first:
1121 + if self.transfer_topo.virtually_split_kv_in_blocks:
1122 # Separate and interleave K/V regions to maintain the same
1123 # descs ordering. This is needed for selecting contiguous heads
1124 # when split across TP ranks.
@@ -1167,7 +1167,7 @@ class NixlConnectorWorker:
1167 addr = base_addr + block_offset + rank_offset
1168 result.append((addr, local_block_len, nixl_agent_meta.device_id))
1169
1170 - if self.transfer_topo.is_kv_layout_blocks_first:
1170 + if self.transfer_topo.virtually_split_kv_in_blocks:
1171 # With FlashInfer index V separately to allow head splitting.
1172 second_split = self.get_backend_aware_kv_block_len(
1173 layer_idx=i, first_split=False, mamba_view=False
@@ -2416,7 +2416,7 @@ class NixlConnectorWorker:
2416 |1st_split-2nd_split| |1st_split-2nd_split |
2417 """
2418 assert self.transfer_topo is not None
2419 - if self.transfer_topo.is_kv_layout_blocks_first:
2419 + if self.transfer_topo.virtually_split_kv_in_blocks:
2420 if mamba_view:
2421 block_len = self._mamba_ssm_size[not first_split]
2422 else:
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+14 -81
@@ -5,7 +5,6 @@ from dataclasses import replace
5
6 import torch
7
8 -from vllm.config import get_layers_from_vllm_config
8 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
9 KVConnectorStats,
10 )
@@ -18,7 +17,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
17 OffloadingConnectorStats,
18 )
19 from vllm.logger import init_logger
21 -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20 from vllm.v1.attention.backend import AttentionBackend
21 from vllm.v1.kv_cache_interface import (
22 AttentionSpec,
@@ -59,25 +57,9 @@ class OffloadingConnectorWorker:
57 def register_kv_caches(
58 self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
59 ):
62 - layer_names = list(kv_caches.keys())
63 - layers = get_layers_from_vllm_config(
64 - self.spec.vllm_config,
65 - AttentionLayerBase, # type: ignore[type-abstract]
66 - layer_names,
67 - )
68 - attn_backends = {
69 - layer_name: layers[layer_name].get_attn_backend()
70 - for layer_name in layer_names
71 - if layer_name in layers
72 - }
73 -
60 num_blocks = self.spec.kv_cache_config.num_blocks
61
76 - # layer_name -> list of matching KV cache tensors
77 - # such that each tensor starts with the num_blocks dimension.
78 - # FlashAttention layers which use the (2, num_blocks, ...) layout
79 - # will possibly map to 2 tensors, one per K and one per V.
80 - # All other layers will probably map to a single tensor.
62 + # layer_name -> (num_blocks, page_size_bytes) tensor
63 tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
64 # layer_name -> size of (un-padded) page in bytes
65 unpadded_page_size_bytes: dict[str, int] = {}
@@ -99,70 +81,21 @@ class OffloadingConnectorWorker:
81 assert isinstance(layer_kv_cache, torch.Tensor)
82 assert layer_kv_cache.storage_offset() == 0
83
102 - # get the logical dimension for num_blocks
103 - test_shape = attn_backends[layer_name].get_kv_cache_shape(
104 - num_blocks=1234,
105 - block_size=16,
106 - num_kv_heads=1,
107 - head_size=256,
108 - )
109 - num_blocks_logical_dim = test_shape.index(1234)
110 -
111 - # sort the logical dimensions by stride (high to low)
112 - # to get a physical-to-logical mapping:
113 - # physical_to_logical[physical_pos] = logical_dim
114 - logical_strides = layer_kv_cache.stride()
115 - physical_to_logical = sorted(
116 - range(len(logical_strides)),
117 - key=lambda idx: logical_strides[idx],
118 - reverse=True,
84 + storage = layer_kv_cache.untyped_storage()
85 + page = layer_kv_cache_spec.page_size_bytes
86 + tensors_per_block[layer_name] = (
87 + torch.tensor(
88 + [],
89 + dtype=torch.int8,
90 + device=layer_kv_cache.device,
91 + )
92 + .set_(storage)
93 + .view(num_blocks, page),
94 )
120 -
121 - num_blocks_physical_dim = physical_to_logical.index(
122 - num_blocks_logical_dim
95 + page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
96 + unpadded_page_size_bytes[layer_name] = (
97 + layer_kv_cache_spec.real_page_size_bytes
98 )
124 - if num_blocks_physical_dim == 0:
125 - storage = layer_kv_cache.untyped_storage()
126 - page = layer_kv_cache_spec.page_size_bytes
127 - tensors_per_block[layer_name] = (
128 - torch.tensor(
129 - [],
130 - dtype=torch.int8,
131 - device=layer_kv_cache.device,
132 - )
133 - .set_(storage)
134 - .view(num_blocks, page),
135 - )
136 - page_size_bytes[layer_name] = (
137 - layer_kv_cache_spec.page_size_bytes
138 - )
139 - unpadded_page_size_bytes[layer_name] = (
140 - layer_kv_cache_spec.real_page_size_bytes
141 - )
142 - else:
143 - # Flash Attention case: (2, num_blocks, ...)
144 - assert test_shape[0] == 2
145 - assert physical_to_logical[0] == 0
146 - assert num_blocks_physical_dim == 1
147 -
148 - # unbind the tensor to separate K and V tensors
149 - half_page_size = layer_kv_cache_spec.page_size_bytes // 2
150 - storage = layer_kv_cache.untyped_storage()
151 - raw = (
152 - torch.tensor(
153 - [],
154 - dtype=torch.int8,
155 - device=layer_kv_cache.device,
156 - )
157 - .set_(storage)
158 - .view(2, num_blocks, half_page_size)
159 - )
160 - tensors_per_block[layer_name] = tuple(raw.unbind(0))
161 -
162 - page_size_bytes[layer_name] = half_page_size
163 - unpadded_page_size_bytes[layer_name] = (
164 - layer_kv_cache_spec.real_page_size_bytes // 2
165 - )
99
100 elif isinstance(layer_kv_cache_spec, MambaSpec):
101 state_tensors = kv_caches[layer_name]
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+17 -70
@@ -18,7 +18,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
18 )
19 from vllm.distributed.parallel_state import get_world_group
20 from vllm.logger import init_logger
21 -from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
21 from vllm.v1.attention.backend import AttentionMetadata
22 from vllm.v1.core.sched.output import SchedulerOutput
23
@@ -140,12 +139,9 @@ class P2pNcclConnector(KVConnectorBase_V1):
139 """
140 Inject KV cache data into a given attention layer tensor.
141
143 - This function updates `layer` in-place with values from `kv_cache`,
144 - handling different backend layouts:
145 - - MLA (Multi-Linear Attention) or FlashInfer: KV tensors are
146 - indexed along the first dimension.
147 - - FlashAttention: KV tensors are indexed along the second
148 - dimension.
142 + This function updates `layer` in-place with values from `kv_cache`.
143 + All backends (MLA, FlashAttention, FlashInfer, TritonAttention)
144 + are indexed along the first dimension (block index).
145
146 If the number of provided block IDs does not match the number of KV
147 blocks, only the overlapping portion is updated, and a warning is
@@ -160,37 +156,19 @@ class P2pNcclConnector(KVConnectorBase_V1):
156 Returns:
157 None. The function modifies `layer` in-place.
158 """
163 - if (
164 - isinstance(attn_metadata, MLACommonMetadata) or layer.shape[1] == 2
165 - ): # MLA or FlashInfer
166 - num_block = kv_cache.shape[0]
167 - self.check_tensors_except_dim(layer, kv_cache, 0)
168 - if len(block_ids) == num_block:
169 - layer[block_ids, ...] = kv_cache
170 - else:
171 - layer[block_ids[:num_block], ...] = kv_cache
172 - logger.warning(
173 - "🚧kv_cache does not match, block_ids:%d, "
174 - "num_block:%d, request_id:%s",
175 - len(block_ids),
176 - num_block,
177 - request_id,
178 - )
179 -
180 - elif layer.shape[0] == 2: # FlashAttention
181 - num_block = kv_cache.shape[1]
182 - self.check_tensors_except_dim(layer, kv_cache, 1)
183 - if len(block_ids) == num_block:
184 - layer[:, block_ids, ...] = kv_cache
185 - else:
186 - layer[:, block_ids[:num_block], ...] = kv_cache
187 - logger.warning(
188 - "🚧kv_cache does not match, block_ids:%d, "
189 - "num_block:%d, request_id:%s",
190 - len(block_ids),
191 - num_block,
192 - request_id,
193 - )
159 + num_block = kv_cache.shape[0]
160 + self.check_tensors_except_dim(layer, kv_cache, 0)
161 + if len(block_ids) == num_block:
162 + layer[block_ids, ...] = kv_cache
163 + else:
164 + layer[block_ids[:num_block], ...] = kv_cache
165 + logger.warning(
166 + "🚧kv_cache does not match, block_ids:%d, "
167 + "num_block:%d, request_id:%s",
168 + len(block_ids),
169 + num_block,
170 + request_id,
171 + )
172
173 # Get the metadata
174 metadata: KVConnectorMetadata = self._get_connector_metadata()
@@ -263,37 +241,6 @@ class P2pNcclConnector(KVConnectorBase_V1):
241
242 assert self.p2p_nccl_engine is not None
243
266 - def extract_kv_from_layer(
267 - layer: torch.Tensor,
268 - block_ids: torch.Tensor,
269 - ) -> torch.Tensor:
270 - """
271 - Extract KV cache slices from a given attention layer tensor.
272 -
273 - This function handles multiple backend layouts:
274 - - MLA (Multi-Linear Attention) or FlashInfer: KV tensors are
275 - indexed along the first dimension.
276 - - FlashAttention: KV tensors are indexed along the second
277 - dimension.
278 -
279 - Args:
280 - layer (torch.Tensor): The KV cache from the attention layer.
281 - block_ids (torch.Tensor): Indices of blocks to extract.
282 -
283 - Returns:
284 - torch.Tensor: A tensor containing the extracted KV slices.
285 - Returns None if the layout is unsupported.
286 - """
287 - if (
288 - isinstance(attn_metadata, MLACommonMetadata) or layer.shape[1] == 2
289 - ): # MLA or FlashInfer
290 - return layer[block_ids, ...]
291 -
292 - if layer.shape[0] == 2: # FlashAttention
293 - return layer[:, block_ids, ...]
294 -
295 - return None
296 -
244 connector_metadata = self._get_connector_metadata()
245 assert isinstance(connector_metadata, P2pNcclConnectorMetadata)
246 for request in connector_metadata.requests:
@@ -301,7 +248,7 @@ class P2pNcclConnector(KVConnectorBase_V1):
248 ip, port = self.parse_request_id(request_id, True)
249 remote_address = ip + ":" + str(port + self._rank)
250
304 - kv_cache = extract_kv_from_layer(kv_layer, request.block_ids)
251 + kv_cache = kv_layer[request.block_ids, ...]
252 self.p2p_nccl_engine.send_tensor(
253 request_id + "#" + layer_name, kv_cache, remote_address
254 )
vllm/platforms/cuda.py
+4 -4
@@ -530,8 +530,8 @@ class CudaPlatformBase(Platform):
530 dst_block_indices: torch.Tensor,
531 ) -> None:
532 """Copy blocks from src_cache to dst_cache on GPU."""
533 - _src_cache = src_cache[:, src_block_indices]
534 - dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
533 + _src_cache = src_cache[src_block_indices]
534 + dst_cache[dst_block_indices] = _src_cache.to(dst_cache.device)
535
536 @classmethod
537 def swap_out_blocks_to_host(
@@ -542,8 +542,8 @@ class CudaPlatformBase(Platform):
542 dst_block_indices: torch.Tensor,
543 ) -> None:
544 """Copy blocks from GPU to host (CPU)."""
545 - _src_cache = src_cache[:, src_block_indices]
546 - dst_cache[:, dst_block_indices] = _src_cache.cpu()
545 + _src_cache = src_cache[src_block_indices]
546 + dst_cache[dst_block_indices] = _src_cache.cpu()
547
548 @classmethod
549 def support_hybrid_kv_cache(cls) -> bool:
vllm/platforms/rocm.py
+4 -4
@@ -901,8 +901,8 @@ class RocmPlatform(Platform):
901 dst_block_indices: torch.Tensor,
902 ) -> None:
903 """Copy blocks from src_cache to dst_cache on GPU."""
904 - _src_cache = src_cache[:, src_block_indices]
905 - dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
904 + _src_cache = src_cache[src_block_indices]
905 + dst_cache[dst_block_indices] = _src_cache.to(dst_cache.device)
906
907 @classmethod
908 def swap_out_blocks_to_host(
@@ -913,8 +913,8 @@ class RocmPlatform(Platform):
913 dst_block_indices: torch.Tensor,
914 ) -> None:
915 """Copy blocks from GPU to host (CPU)."""
916 - _src_cache = src_cache[:, src_block_indices]
917 - dst_cache[:, dst_block_indices] = _src_cache.cpu()
916 + _src_cache = src_cache[src_block_indices]
917 + dst_cache[dst_block_indices] = _src_cache.cpu()
918
919 @classmethod
920 def support_hybrid_kv_cache(cls) -> bool:
vllm/platforms/xpu.py
+4 -4
@@ -374,8 +374,8 @@ class XPUPlatform(Platform):
374 dst_block_indices: torch.Tensor,
375 ) -> None:
376 """Copy blocks from src_cache to dst_cache on XPU."""
377 - _src_cache = src_cache[:, src_block_indices]
378 - dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
377 + _src_cache = src_cache[src_block_indices]
378 + dst_cache[dst_block_indices] = _src_cache.to(dst_cache.device)
379
380 @classmethod
381 def swap_out_blocks_to_host(
@@ -386,8 +386,8 @@ class XPUPlatform(Platform):
386 dst_block_indices: torch.Tensor,
387 ) -> None:
388 """Copy blocks from XPU to host (CPU)."""
389 - _src_cache = src_cache[:, src_block_indices]
390 - dst_cache[:, dst_block_indices] = _src_cache.cpu()
389 + _src_cache = src_cache[src_block_indices]
390 + dst_cache[dst_block_indices] = _src_cache.cpu()
391
392 @classmethod
393 def num_compute_units(cls, device_id: int = 0) -> int:
vllm/v1/attention/backends/flash_attn.py
+6 -6
@@ -146,7 +146,7 @@ class FlashAttentionBackend(AttentionBackend):
146 ) -> tuple[int, ...]:
147 if block_size % 16 != 0:
148 raise ValueError("Block size must be a multiple of 16.")
149 - return (2, num_blocks, block_size, num_kv_heads, head_size)
149 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
150
151 @staticmethod
152 def get_kv_cache_stride_order(
@@ -157,12 +157,12 @@ class FlashAttentionBackend(AttentionBackend):
157 cache_layout = get_kv_cache_layout()
158 if cache_layout == "NHD" and include_num_layers_dimension:
159 # (num_blocks, num_layers, 2, block_size, num_kv_heads, head_size)
160 - return (2, 0, 1, 3, 4, 5)
160 + return (1, 0, 2, 3, 4, 5)
161 elif cache_layout == "NHD":
162 stride_order = (0, 1, 2, 3, 4)
163 elif cache_layout == "HND" and include_num_layers_dimension:
164 # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
165 - return (2, 4, 0, 1, 3, 5)
165 + return (1, 4, 0, 2, 3, 5)
166 elif cache_layout == "HND":
167 stride_order = (0, 1, 3, 2, 4)
168 else:
@@ -683,7 +683,7 @@ class FlashAttentionImpl(AttentionImpl):
683 key: shape = [num_tokens, num_kv_heads, head_size]
684 value: shape = [num_tokens, num_kv_heads, head_size]
685 kv_cache: shape =
686 - [2, num_blocks, block_size, num_kv_heads, head_size]
686 + [num_blocks, 2, block_size, num_kv_heads, head_size]
687 attn_metadata: Metadata for attention.
688 Returns:
689 shape = [num_tokens, num_heads * head_size]
@@ -731,7 +731,7 @@ class FlashAttentionImpl(AttentionImpl):
731 )
732
733 # For decoder and cross-attention, use KV cache as before
734 - key_cache, value_cache = kv_cache.unbind(0)
734 + key_cache, value_cache = kv_cache.unbind(1)
735 # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
736 # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
737 # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
@@ -862,7 +862,7 @@ class FlashAttentionImpl(AttentionImpl):
862
863 # Scatter write into the KV cache using slot_mapping indices.
864 # No TMA kernel is invoked here, so stride canonicalization is not needed.
865 - key_cache, value_cache = kv_cache.unbind(0)
865 + key_cache, value_cache = kv_cache.unbind(1)
866
867 # Reshape the input keys and values and store them in the cache.
868 # Skip this if sharing KV cache with an earlier attention layer.
vllm/v1/attention/backends/flex_attention.py
+13 -5
@@ -125,7 +125,15 @@ class FlexAttentionBackend(AttentionBackend):
125 head_size: int,
126 cache_dtype_str: str = "auto",
127 ) -> tuple[int, ...]:
128 - return (2, num_blocks, block_size, num_kv_heads, head_size)
128 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
129 +
130 + @staticmethod
131 + def get_kv_cache_stride_order(
132 + include_num_layers_dimension: bool = False,
133 + ) -> tuple[int, ...]:
134 + if include_num_layers_dimension:
135 + return (1, 0, 3, 2, 4, 5)
136 + return (0, 2, 1, 3, 4)
137
138 @staticmethod
139 def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
@@ -1055,7 +1063,7 @@ class FlexAttentionImpl(AttentionImpl):
1063 if self.attn_type == AttentionType.ENCODER_ONLY:
1064 return
1065
1058 - key_cache, value_cache = kv_cache.unbind(0)
1066 + key_cache, value_cache = kv_cache.unbind(1)
1067 torch.ops._C_cache_ops.reshape_and_cache_flash(
1068 key,
1069 value,
@@ -1086,7 +1094,7 @@ class FlexAttentionImpl(AttentionImpl):
1094 key: shape = [num_tokens, num_kv_heads, head_size]
1095 value: shape = [num_tokens, num_kv_heads, head_size]
1096 kv_cache: shape =
1089 - [2, num_blocks, block_size, num_kv_heads, head_size]
1097 + [num_blocks, 2, block_size, num_kv_heads, head_size]
1098 attn_metadata: Metadata for attention.
1099 Returns:
1100 shape = [num_tokens, num_heads * head_size]
@@ -1160,9 +1168,9 @@ class FlexAttentionImpl(AttentionImpl):
1168
1169 else:
1170 assert self.attn_type == AttentionType.DECODER
1163 - key_cache, value_cache = kv_cache.unbind(0)
1171 + key_cache, value_cache = kv_cache.unbind(1)
1172
1165 - # View out the block_size dim
1173 + # Flatten (num_blocks, block_size) into a single token dim
1174 key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
1175 value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
1176 query, key_tensor, value_tensor = map(
vllm/v1/attention/backends/triton_attn.py
+2 -2
@@ -346,8 +346,8 @@ class TritonAttentionBackend(AttentionBackend):
346 elif cache_layout == "NHD":
347 stride_order = (0, 1, 2, 3, 4)
348 elif cache_layout == "HND" and include_num_layers_dimension:
349 - # (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
350 - return (1, 2, 4, 0, 3, 5)
349 + # (num_blocks, num_kv_heads, num_layers, 2, block_size, head_size)
350 + return (1, 4, 0, 2, 3, 5)
351 elif cache_layout == "HND":
352 stride_order = (0, 1, 3, 2, 4)
353 else:
vllm/v1/worker/gpu/attn_utils.py
+2 -3
@@ -248,9 +248,8 @@ def _reshape_kv_cache(
248 # MambaSpec handling in gpu_model_runner.py.
249 # NOTE: This assumes kv_cache_shape[0] == num_blocks
250 # (i.e. the first physical dimension is the block
251 - # index), which holds for MLA backends but NOT for
252 - # standard attention backends whose shape starts with
253 - # a K/V dimension of size 2.
251 + # index), which holds for all current backends
252 + # (MLA, FlashAttention, TritonAttention, etc.).
253 dtype_size = get_dtype_size(dtype)
254 page_stride = kv_cache_spec.page_size_bytes // dtype_size
255 strides = list(torch.empty(kv_cache_shape).stride())