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, 06:00 UTC
[V1][Kernel] Flashinfer HND KV cache layout (#19280)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 17, 2025 at 15:09 UTC
4c8f64faa74132580ecb5f5ed5184821b5fbb013
6 files changed
+64
-20
vllm/attention/backends/flashinfer.py
+1
-3
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import dataclasses
5
-import os
5
from collections import defaultdict
6
from contextlib import contextmanager
7
from dataclasses import dataclass
@@ -50,8 +49,7 @@ if TYPE_CHECKING:
49
from vllm.worker.model_runner import (ModelInputForGPUBuilder,
50
ModelInputForGPUWithSamplingMetadata)
51
53
-FLASHINFER_KV_CACHE_LAYOUT: str = os.getenv("FLASHINFER_KV_CACHE_LAYOUT",
54
- "NHD").upper()
52
+FLASHINFER_KV_CACHE_LAYOUT: str = envs.VLLM_KV_CACHE_LAYOUT or "NHD"
53
54
55
class FlashInferBackend(AttentionBackend):
vllm/distributed/kv_transfer/kv_connector/utils.py
+5
-4
@@ -3,7 +3,6 @@
3
"""
4
KV cache helper for store.
5
"""
6
-
6
import torch
7
8
import vllm.envs as envs
@@ -94,15 +93,17 @@ class model_aware_kv_ops_helper:
93
94
95
def get_kv_connector_cache_layout():
96
+ # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
97
+ # used for faster transfer.
98
vllm_config = get_current_vllm_config()
99
kv_config = vllm_config.kv_transfer_config
99
- if vllm_config.model_config is None:
100
- logger.warning("Unable to detect current VLLM config. " \
100
+ if vllm_config.model_config is None or kv_config is None:
101
+ logger.warning_once("Unable to detect current VLLM config. " \
102
"Defaulting to NHD kv cache layout.")
103
else:
104
use_mla = vllm_config.model_config.use_mla
105
if not use_mla and kv_config.kv_connector == "NixlConnector":
105
- logger.info("NixlConnector detected. Setting KV cache " \
106
+ logger.info_once("NixlConnector detected. Setting KV cache " \
107
"layout to HND for better xfer performance.")
108
return "HND"
109
return "NHD"
vllm/envs.py
+11
@@ -128,6 +128,7 @@ if TYPE_CHECKING:
128
VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
129
VLLM_SLEEP_WHEN_IDLE: bool = False
130
VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
131
+ VLLM_KV_CACHE_LAYOUT: Optional[str] = None
132
133
134
def get_default_cache_root():
@@ -879,6 +880,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
880
# processes via zmq.
881
"VLLM_MQ_MAX_CHUNK_BYTES_MB":
882
lambda: int(os.getenv("VLLM_MQ_MAX_CHUNK_BYTES_MB", "16")),
883
+
884
+ # KV Cache layout used throughout vllm.
885
+ # Some common values are:
886
+ # - NHD
887
+ # - HND
888
+ # Where N=num_blocks, H=num_heads and D=head_size. The default value will
889
+ # leave the layout choice to the backend. Mind that backends may only
890
+ # implement and support a subset of all possible layouts.
891
+ "VLLM_KV_CACHE_LAYOUT":
892
+ lambda: os.getenv("VLLM_KV_CACHE_LAYOUT", None)
893
}
894
895
# --8<-- [end:env-vars-definition]
vllm/v1/attention/backends/flash_attn.py
+5
-7
@@ -16,13 +16,12 @@ from vllm.attention.ops.merge_attn_states import merge_attn_states
16
from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
17
get_flash_attn_version)
18
from vllm.config import VllmConfig, get_layers_from_vllm_config
19
-from vllm.distributed.kv_transfer.kv_connector.utils import (
20
- get_kv_connector_cache_layout)
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
from vllm.utils import cdiv
22
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
25
- CommonAttentionMetadata)
23
+ CommonAttentionMetadata,
24
+ get_kv_cache_layout)
25
from vllm.v1.kv_cache_interface import AttentionSpec
26
from vllm.v1.worker.block_table import BlockTable
27
@@ -73,16 +72,15 @@ class FlashAttentionBackend(AttentionBackend):
72
73
@staticmethod
74
def get_kv_cache_stride_order() -> tuple[int, ...]:
76
- # NOTE When running disaggregated PD with NIXL, HND layout is used for
77
- # faster transfer. `stride_order` indicates the permutation that gets
75
+ # `stride_order` indicates the permutation that gets
76
# us from `get_kv_cache_shape` to the actual memory layout we want.
79
- cache_layout = get_kv_connector_cache_layout()
77
+ cache_layout = get_kv_cache_layout()
78
if cache_layout == "NHD":
79
stride_order = (0, 1, 2, 3, 4)
80
elif cache_layout == "HND":
81
stride_order = (0, 1, 3, 2, 4)
82
else:
85
- raise ValueError("Unknown cache layout format %s.", cache_layout)
83
+ raise ValueError(f"Unknown cache layout format {cache_layout}.")
84
return stride_order
85
86
vllm/v1/attention/backends/flashinfer.py
+21
-6
@@ -19,7 +19,8 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
19
from vllm.logger import init_logger
20
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
21
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
22
- CommonAttentionMetadata)
22
+ CommonAttentionMetadata,
23
+ get_kv_cache_layout)
24
from vllm.v1.kv_cache_interface import AttentionSpec
25
from vllm.v1.worker.block_table import BlockTable
26
@@ -66,6 +67,19 @@ class FlashInferBackend(AttentionBackend):
67
) -> tuple[int, ...]:
68
return (num_blocks, 2, block_size, num_kv_heads, head_size)
69
70
+ @staticmethod
71
+ def get_kv_cache_stride_order() -> tuple[int, ...]:
72
+ # `stride_order` indicates the permutation that gets us from
73
+ # `get_kv_cache_shape` to the actual memory layout we want.
74
+ cache_layout = get_kv_cache_layout()
75
+ if cache_layout == "NHD":
76
+ stride_order = (0, 1, 2, 3, 4)
77
+ elif cache_layout == "HND":
78
+ stride_order = (0, 1, 3, 2, 4)
79
+ else:
80
+ raise ValueError(f"Unknown cache layout format {cache_layout}.")
81
+ return stride_order
82
+
83
84
@dataclass
85
class PerLayerParameters:
@@ -290,7 +304,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
304
def _get_prefill_wrapper(self):
305
if self._prefill_wrapper is None:
306
self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
293
- self._get_workspace_buffer(), "NHD")
307
+ self._get_workspace_buffer(), get_kv_cache_layout())
308
return self._prefill_wrapper
309
310
def _get_decode_wrapper(self):
@@ -303,14 +317,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
317
num_qo_heads // num_kv_heads > 4)
318
self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
319
self._get_workspace_buffer(),
306
- "NHD",
320
+ get_kv_cache_layout(),
321
use_tensor_cores=use_tensor_cores)
322
return self._decode_wrapper
323
324
def _get_cascade_wrapper(self):
325
if self._cascade_wrapper is None:
326
self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
313
- 2, self._get_workspace_buffer(), "NHD")
327
+ 2, self._get_workspace_buffer(), get_kv_cache_layout())
328
return self._cascade_wrapper
329
330
def _plan(self, attn_metadata: FlashInferMetadata):
@@ -620,6 +634,7 @@ class FlashInferImpl(AttentionImpl):
634
num_decode_tokens = attn_metadata.num_decode_tokens
635
num_prefill_tokens = attn_metadata.num_prefill_tokens
636
637
+ stride_order = FlashInferBackend.get_kv_cache_stride_order()
638
# Regular attention (common case).
639
# Decodes are at the front and prefills are at the back,
640
# according to reorder_batch()
@@ -634,7 +649,7 @@ class FlashInferImpl(AttentionImpl):
649
assert prefill_wrapper._sm_scale == self.scale
650
prefill_wrapper.run(
651
prefill_query,
637
- kv_cache,
652
+ kv_cache.permute(*stride_order),
653
k_scale=layer._k_scale_float,
654
v_scale=layer._v_scale_float,
655
out=output[num_decode_tokens:],
@@ -650,7 +665,7 @@ class FlashInferImpl(AttentionImpl):
665
assert decode_wrapper._sm_scale == self.scale
666
decode_wrapper.run(
667
decode_query,
653
- kv_cache,
668
+ kv_cache.permute(*stride_order),
669
k_scale=layer._k_scale_float,
670
v_scale=layer._v_scale_float,
671
out=output[:num_decode_tokens],
vllm/v1/attention/backends/utils.py
+21
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import abc
4
+import functools
5
from abc import abstractmethod
6
from dataclasses import dataclass
7
from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
@@ -12,6 +13,13 @@ if TYPE_CHECKING:
13
from vllm.v1.core.sched.output import SchedulerOutput
14
from vllm.v1.worker.gpu_input_batch import InputBatch
15
16
+import vllm.envs as envs
17
+from vllm.distributed.kv_transfer.kv_connector.utils import (
18
+ get_kv_connector_cache_layout)
19
+from vllm.logger import init_logger
20
+
21
+logger = init_logger(__name__)
22
+
23
24
@dataclass
25
class CommonAttentionMetadata:
@@ -119,3 +127,16 @@ def validate_kv_sharing_target(current_layer_name, target_layer_name,
127
raise ValueError(
128
error_msg +
129
f"must be the same type as the current layer ({expected}).")
130
+
131
+
132
+@functools.lru_cache
133
+def get_kv_cache_layout():
134
+ # Override with format specified by the user.
135
+ cache_layout = envs.VLLM_KV_CACHE_LAYOUT
136
+ if cache_layout is None:
137
+ cache_layout = get_kv_connector_cache_layout()
138
+ else:
139
+ logger.info_once("`FLASHINFER_KV_CACHE_LAYOUT` environment variable " \
140
+ "detected. Setting KV cache layout to %s.", cache_layout)
141
+
142
+ return cache_layout