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 6, 03:15 UTC
[Bugfix] Add replacement of _compute_slot_mapping_kernel on CPU (#37987)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Mar 24, 2026 at 22:00 UTC
352b90c4a4f2af1d75cdc09fecf30a47a4ca39fb
7 files changed
+105
-9
.buildkite/hardware_tests/cpu.yaml
-7
@@ -3,7 +3,6 @@ depends_on: []
3
steps:
4
- label: CPU-Kernel Tests
5
depends_on: []
6
- soft_fail: true
6
device: intel_cpu
7
no_plugin: true
8
source_file_dependencies:
@@ -23,7 +22,6 @@ steps:
22
23
- label: CPU-Compatibility Tests
24
depends_on: []
26
- soft_fail: true
25
device: intel_cpu
26
no_plugin: true
27
source_file_dependencies:
@@ -37,7 +35,6 @@ steps:
35
36
- label: CPU-Language Generation and Pooling Model Tests
37
depends_on: []
40
- soft_fail: true
38
device: intel_cpu
39
no_plugin: true
40
source_file_dependencies:
@@ -53,7 +50,6 @@ steps:
50
51
- label: CPU-Quantization Model Tests
52
depends_on: []
56
- soft_fail: true
53
device: intel_cpu
54
no_plugin: true
55
source_file_dependencies:
@@ -73,7 +69,6 @@ steps:
69
70
- label: CPU-Distributed Tests
71
depends_on: []
76
- soft_fail: true
72
device: intel_cpu
73
no_plugin: true
74
source_file_dependencies:
@@ -92,7 +87,6 @@ steps:
87
88
- label: CPU-Multi-Modal Model Tests %N
89
depends_on: []
95
- soft_fail: true
90
device: intel_cpu
91
no_plugin: true
92
source_file_dependencies:
@@ -107,7 +101,6 @@ steps:
101
102
- label: "Arm CPU Test"
103
depends_on: []
110
- soft_fail: true
104
device: arm_cpu
105
no_plugin: true
106
commands:
csrc/cpu/torch_bindings.cpp
+12
@@ -126,6 +126,12 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
126
const torch::Tensor& topk_id, const bool skip_weighted,
127
const std::string& act, const std::string& isa);
128
129
+void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
130
+ const torch::Tensor positions,
131
+ const torch::Tensor block_table,
132
+ torch::Tensor slot_mapping,
133
+ const int64_t block_size);
134
+
135
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
136
// vLLM custom ops
137
@@ -334,6 +340,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
340
" Tensor! out, Tensor query, Tensor kv_cache,"
341
" float scale, Tensor block_tables, Tensor seq_lens) -> ()");
342
ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
343
+
344
+ ops.def(
345
+ "compute_slot_mapping_kernel_impl(Tensor query_start_loc, Tensor "
346
+ "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt "
347
+ "block_size) -> ()",
348
+ &compute_slot_mapping_kernel_impl);
349
}
350
351
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
csrc/cpu/utils.cpp
+35
@@ -189,3 +189,38 @@ ScratchPadManager* ScratchPadManager::get_scratchpad_manager() {
189
return &manager;
190
}
191
} // namespace cpu_utils
192
+
193
+void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
194
+ const torch::Tensor positions,
195
+ const torch::Tensor block_table,
196
+ torch::Tensor slot_mapping,
197
+ const int64_t block_size) {
198
+ const int32_t req_num = query_start_loc.size(0) - 1;
199
+ const int64_t block_table_stride = block_table.stride(0);
200
+
201
+ const int32_t* __restrict__ query_start_loc_ptr =
202
+ query_start_loc.data_ptr<int32_t>();
203
+ const int64_t* __restrict__ positions_ptr = positions.data_ptr<int64_t>();
204
+ const int32_t* __restrict__ blocktable_ptr = block_table.data_ptr<int32_t>();
205
+ int64_t* __restrict__ slot_mapping_ptr = slot_mapping.data_ptr<int64_t>();
206
+
207
+#pragma omp parallel for
208
+ for (int32_t req_idx = 0; req_idx < req_num; ++req_idx) {
209
+ int32_t token_start_idx = query_start_loc_ptr[req_idx];
210
+ int32_t token_end_idx = query_start_loc_ptr[req_idx + 1];
211
+ int32_t token_num = token_end_idx - token_start_idx;
212
+ const int64_t* __restrict__ curr_position_ptr =
213
+ positions_ptr + token_start_idx;
214
+ int64_t* __restrict__ curr_slot_mapping_ptr =
215
+ slot_mapping_ptr + token_start_idx;
216
+ const int32_t* __restrict__ curr_block_table_ptr =
217
+ blocktable_ptr + req_idx * block_table_stride;
218
+
219
+ for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
220
+ int64_t token_position = curr_position_ptr[token_idx];
221
+ int64_t block_id = curr_block_table_ptr[token_position / block_size];
222
+ curr_slot_mapping_ptr[token_idx] =
223
+ block_id * block_size + token_position % block_size;
224
+ }
225
+ }
226
+}
docker/Dockerfile.cpu
+1
-1
@@ -161,7 +161,7 @@ RUN ln -s /usr/bin/clangd-14 /usr/bin/clangd
161
162
# install development dependencies (for testing)
163
RUN --mount=type=cache,target=/root/.cache/uv \
164
- uv pip install -e tests/vllm_test_utils
164
+ uv pip install --no-build-isolation -e tests/vllm_test_utils
165
166
RUN --mount=type=cache,target=/root/.cache/uv \
167
--mount=type=cache,target=/root/.cache/ccache \
vllm/model_executor/layers/quantization/awq_marlin.py
+1
-1
@@ -309,7 +309,7 @@ class AWQMarlinConfig(QuantizationConfig):
309
group_size = quant_config.get("group_size")
310
zero_point = quant_config.get("zero_point")
311
312
- if not (current_platform.is_cuda_alike() or current_platform.is_cpu()):
312
+ if not current_platform.is_cuda_alike():
313
return False
314
315
if quant_method != "awq":
vllm/utils/cpu_triton_utils.py
new
+47
@@ -0,0 +1,47 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""
5
+Contains replacement functions to fallback Triton usages in CPU backend
6
+"""
7
+
8
+from collections.abc import Callable
9
+
10
+import torch
11
+
12
+
13
+class _FuncWrapper:
14
+ def __init__(self, func: Callable) -> None:
15
+ self.func = func
16
+
17
+ def __getitem__(self, *args, **kwargs) -> Callable:
18
+ return self.func
19
+
20
+
21
+# For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py
22
+def _compute_slot_mapping_kernel_impl(
23
+ num_tokens: int,
24
+ max_num_tokens: int,
25
+ query_start_loc: torch.Tensor, # [num_reqs + 1], int32
26
+ positions: torch.Tensor, # [num_tokens], int64
27
+ block_table: torch.Tensor, # [max_num_reqs, max_num_blocks_per_req], int32
28
+ block_table_stride: int, # max_num_blocks_per_req
29
+ block_size: int,
30
+ slot_mapping: torch.Tensor, # [max_num_tokens], int64
31
+ TOTAL_CP_WORLD_SIZE: int,
32
+ TOTAL_CP_RANK: int,
33
+ CP_KV_CACHE_INTERLEAVE_SIZE: int,
34
+ PAD_ID: int,
35
+ BLOCK_SIZE: int,
36
+) -> None:
37
+ assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU."
38
+ torch.ops._C.compute_slot_mapping_kernel_impl(
39
+ query_start_loc,
40
+ positions,
41
+ block_table,
42
+ slot_mapping,
43
+ block_size,
44
+ )
45
+
46
+
47
+compute_slot_mapping_kernel = _FuncWrapper(_compute_slot_mapping_kernel_impl)
vllm/v1/worker/cpu_model_runner.py
+9
@@ -6,6 +6,7 @@ from typing import Any
6
import torch
7
import torch.nn as nn
8
9
+import vllm.utils.cpu_triton_utils as cpu_tl
10
from vllm.config import VllmConfig
11
from vllm.logger import init_logger
12
from vllm.model_executor.model_loader import get_model
@@ -28,6 +29,7 @@ class CPUModelRunner(GPUModelRunner):
29
self.cascade_attn_enabled = False
30
31
self._postprocess_tensors()
32
+ self._postprocess_triton()
33
34
def _postprocess_tensors(self) -> None:
35
# Note: replace device tensors with cpu tensors
@@ -52,6 +54,13 @@ class CPUModelRunner(GPUModelRunner):
54
if isinstance(v, CpuGpuBuffer):
55
v.gpu = v.cpu
56
57
+ def _postprocess_triton(self) -> None:
58
+ import vllm.v1.worker.block_table
59
+
60
+ vllm.v1.worker.block_table._compute_slot_mapping_kernel = (
61
+ cpu_tl.compute_slot_mapping_kernel
62
+ )
63
+
64
@instrument(span_name="Loading (CPU)")
65
def load_model(self, load_dummy_weights: bool = False) -> None:
66
if load_dummy_weights: