| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Physical KV cache layout descriptor. |
| 4 | |
| 5 | A leaf module so ``vllm.config`` can import the enum without pulling in the |
| 6 | full KV cache interface. |
| 7 | """ |
| 8 | |
| 9 | from enum import Enum |
| 10 | |
| 11 | # Logical dim indices in the 5D stride permutation [L, B, H, N, C] (see: RFC #42082). |
| 12 | _DIM_L, _DIM_B, _DIM_H, _DIM_N, _DIM_C = 0, 1, 2, 3, 4 |
| 13 | |
| 14 | |
| 15 | class KVCacheLayout(Enum): |
| 16 | """Physical layout descriptor for a KV cache group. |
| 17 | |
| 18 | The logical shape is always [L, B, H, N, <content>] (RFC #42082). |
| 19 | Each member's value is a stride permutation that maps logical axes |
| 20 | to physical (memory) order. |
| 21 | """ |
| 22 | |
| 23 | LBHNC = (0, 1, 2, 3, 4) # [L, B, H, N, C] (identity) |
| 24 | LBNHC = (0, 1, 3, 2, 4) # [L, B, N, H, C] |
| 25 | LHBNC = (0, 2, 1, 3, 4) # [L, H, B, N, C] |
| 26 | BLHNC = (1, 0, 2, 3, 4) # [B, L, H, N, C] |
| 27 | BLNHC = (1, 0, 3, 2, 4) # [B, L, N, H, C] |
| 28 | BHLNC = (1, 2, 0, 3, 4) # [B, H, L, N, C] |
| 29 | |
| 30 | @property |
| 31 | def stride_order(self) -> tuple[int, ...]: |
| 32 | return self.value |
| 33 | |
| 34 | @property |
| 35 | def layer_view_order(self) -> tuple[int, ...]: |
| 36 | """Physical axis order of a logical 4D per-layer cache view.""" |
| 37 | return tuple(i - 1 for i in self.value if i != _DIM_L) |
| 38 | |
| 39 | @property |
| 40 | def is_layer_compact(self) -> bool: |
| 41 | """True when the layer is compact; i.e. the L dimension is outermost.""" |
| 42 | return self.value[_DIM_L] == 0 |
| 43 | |
| 44 | @property |
| 45 | def is_block_contiguous(self) -> bool: |
| 46 | """True when [H, N, C] is contiguous within a block.""" |
| 47 | return self.value[-3:] == (_DIM_H, _DIM_N, _DIM_C) |
| 48 | |
| 49 | @property |
| 50 | def is_block_compact(self) -> bool: |
| 51 | """True when each page's [H, N, C] bytes form one contiguous run; i.e. |
| 52 | the L and B dimensions are outermost.""" |
| 53 | return set(self.value[:2]) == {_DIM_L, _DIM_B} |
| 54 | |
| 55 | @property |
| 56 | def is_block_outermost(self) -> bool: |
| 57 | """True when B is the outermost physical dimension.""" |
| 58 | return self.value[0] == _DIM_B |