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 13, 00:45 UTC
[NIXL] fix cpu PD after physical <> logical block_size PR (#28904)
Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Chendi.Xue committed
Nov 18, 2025 at 13:03 UTC
c3e29786209d91d3842e839b62f4d1d815902262
3 files changed
+17
-5
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+7
-2
@@ -55,7 +55,7 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
55
# Find the git repository root directory
56
GIT_ROOT=$(git rev-parse --show-toplevel)
57
58
-SMI_BIN=$(which nvidia-smi || which rocm-smi)
58
+SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
59
60
# Trap the SIGINT signal (triggered by Ctrl+C)
61
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
@@ -91,8 +91,13 @@ get_model_args() {
91
get_num_gpus() {
92
if [[ "$SMI_BIN" == *"nvidia"* ]]; then
93
echo "$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)"
94
- else
94
+ elif [[ "$SMI_BIN" == *"rocm"* ]]; then
95
echo "$($SMI_BIN -l | grep GPU | wc -l)"
96
+ else
97
+ # works for non-cuda platforms,
98
+ # assuming at least 1 device and
99
+ # let system to decide which card to use
100
+ echo "1"
101
fi
102
}
103
tools/install_nixl_from_source_ubuntu.py
+1
@@ -95,6 +95,7 @@ def install_system_dependencies():
95
"meson",
96
"libtool",
97
"libtool-bin",
98
+ "pkg-config",
99
]
100
run_command(["apt-get", "update"])
101
run_command(["apt-get", "install", "-y"] + apt_packages)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+9
-3
@@ -1161,6 +1161,14 @@ class NixlConnectorWorker:
1161
# to better exploit the memory layout (ie num_blocks is the first dim).
1162
split_k_and_v = self.kv_topo.split_k_and_v
1163
tensor_size_bytes = None
1164
+
1165
+ # TODO (NickLucche): Get kernel_block_size in a cleaner way
1166
+ # NHD default "view" for non-MLA cache
1167
+ if self.device_type == "cpu":
1168
+ block_size_position = -2
1169
+ else:
1170
+ block_size_position = -2 if self.use_mla else -3
1171
+
1172
# Enable different block lengths for different layers when MLA is used.
1173
self.block_len_per_layer = list[int]()
1174
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
@@ -1175,9 +1183,7 @@ class NixlConnectorWorker:
1183
if base_addr in seen_base_addresses:
1184
continue
1185
1178
- # TODO (NickLucche): Get kernel_block_size in a cleaner way
1179
- # NHD default "view" for non-MLA cache
1180
- kernel_block_size = cache.shape[-2] if self.use_mla else cache.shape[-3]
1186
+ kernel_block_size = cache.shape[block_size_position]
1187
1188
if self.block_size != kernel_block_size:
1189
logger.info_once(