@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 445a2a4d1a
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 5, 08:30 UTC

feat(cpu): add CPU support for draft model speculative decoding (#32662)

Signed-off-by: R <Ganesh.R@amd.com>

Ganesh R committed Apr 10, 2026 at 09:19 UTC 445a2a4d1a3a383a1a36da8acb7800bb85edabeb
7 files changed +945 -7
cmake/cpu_extension.cmake
+3
@@ -349,6 +349,7 @@ endif()
349 set(VLLM_EXT_SRC
350 "csrc/cpu/activation.cpp"
351 "csrc/cpu/utils.cpp"
352 + "csrc/cpu/spec_decode_utils.cpp"
353 "csrc/cpu/layernorm.cpp"
354 "csrc/cpu/mla_decode.cpp"
355 "csrc/cpu/pos_encoding.cpp"
@@ -383,6 +384,7 @@ if (ENABLE_X86_ISA)
384 "csrc/cpu/cpu_wna16.cpp"
385 "csrc/cpu/cpu_fused_moe.cpp"
386 "csrc/cpu/utils.cpp"
387 + "csrc/cpu/spec_decode_utils.cpp"
388 "csrc/cpu/cpu_attn.cpp"
389 "csrc/cpu/dnnl_kernels.cpp"
390 "csrc/cpu/torch_bindings.cpp"
@@ -395,6 +397,7 @@ if (ENABLE_X86_ISA)
397
398 set(VLLM_EXT_SRC_AVX2
399 "csrc/cpu/utils.cpp"
400 + "csrc/cpu/spec_decode_utils.cpp"
401 "csrc/cpu/cpu_attn.cpp"
402 "csrc/cpu/torch_bindings.cpp"
403 # TODO: Remove these files
csrc/cpu/spec_decode_utils.cpp new
+409
@@ -0,0 +1,409 @@
1 +#include "cpu_types.hpp"
2 +
3 +#include <algorithm>
4 +
5 +namespace cpu_utils {
6 +
7 +void eagle_prepare_inputs_padded_kernel_impl(
8 + const torch::Tensor& cu_num_draft_tokens,
9 + const torch::Tensor& valid_sampled_tokens_count,
10 + const torch::Tensor& query_start_loc_gpu,
11 + torch::Tensor& token_indices_to_sample,
12 + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) {
13 + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
14 + const int64_t* valid_count_ptr =
15 + valid_sampled_tokens_count.data_ptr<int64_t>();
16 + const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr<int32_t>();
17 + int32_t* indices_out_ptr = token_indices_to_sample.data_ptr<int32_t>();
18 + int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr<int64_t>();
19 +
20 +#pragma omp parallel for
21 + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
22 + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
23 + int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx;
24 + int64_t num_valid_tokens = valid_count_ptr[req_idx];
25 +
26 + int64_t num_rejected = 0;
27 + if (num_draft_tokens > 0) {
28 + num_rejected = num_draft_tokens + 1 - num_valid_tokens;
29 + }
30 +
31 + int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1;
32 + int32_t index_to_sample = q_last_tok_idx - num_rejected;
33 +
34 + indices_out_ptr[req_idx] = index_to_sample;
35 + rejected_out_ptr[req_idx] = num_rejected;
36 + }
37 +}
38 +
39 +void eagle_prepare_next_token_padded_kernel_impl(
40 + const torch::Tensor& sampled_token_ids,
41 + const torch::Tensor& discard_request_mask,
42 + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
43 + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
44 + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) {
45 + const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr<int64_t>();
46 + const bool* discard_mask_ptr = discard_request_mask.data_ptr<bool>();
47 + const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr<int64_t>();
48 + int64_t* next_ids_out_ptr = next_token_ids.data_ptr<int64_t>();
49 + int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr<int64_t>();
50 +
51 + const int64_t stride = sampled_token_ids.stride(0);
52 +
53 +#pragma omp parallel for
54 + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
55 + const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride;
56 + int64_t valid_count = 0;
57 + int64_t last_valid_token = -1;
58 +
59 + for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) {
60 + int64_t token = row_ptr[pos];
61 + if (token != -1 && token < vocab_size) {
62 + valid_count++;
63 + last_valid_token = token;
64 + }
65 + }
66 +
67 + bool discard = discard_mask_ptr[req_idx];
68 + if (discard) {
69 + next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx];
70 + valid_count_out_ptr[req_idx] = 0;
71 + } else {
72 + next_ids_out_ptr[req_idx] =
73 + (valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx];
74 + valid_count_out_ptr[req_idx] = valid_count;
75 + }
76 + }
77 +}
78 +
79 +void eagle_step_slot_mapping_metadata_kernel_impl(
80 + const torch::Tensor& positions, const torch::Tensor& block_table,
81 + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
82 + torch::Tensor& out_slot_mapping, const int64_t block_size,
83 + const int64_t max_model_len, const int64_t PAD_ID) {
84 + const int64_t batch_size = positions.size(0);
85 + const int64_t input_batch_size = out_slot_mapping.size(0);
86 +
87 + const int64_t* pos_ptr = positions.data_ptr<int64_t>();
88 + const int32_t* bt_ptr = block_table.data_ptr<int32_t>();
89 + int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>();
90 + int64_t* out_clamped_ptr = out_clamped_positions.data_ptr<int64_t>();
91 + int64_t* out_slot_ptr = out_slot_mapping.data_ptr<int64_t>();
92 +
93 + const int64_t bt_stride = block_table.stride(0);
94 + const int64_t n_blocks_per_req = block_table.size(1);
95 +
96 +#pragma omp parallel for
97 + for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) {
98 + if (req_idx >= batch_size) {
99 + out_slot_ptr[req_idx] = PAD_ID;
100 + continue;
101 + }
102 +
103 + int64_t position = pos_ptr[req_idx];
104 + int64_t new_position = position + 1;
105 + bool exceeds_max = new_position >= max_model_len;
106 + int64_t clamped_position = exceeds_max ? 0 : new_position;
107 +
108 + out_clamped_ptr[req_idx] = clamped_position;
109 +
110 + int64_t block_number = clamped_position / block_size;
111 + block_number = std::min(block_number, n_blocks_per_req - 1);
112 + int32_t block_id = bt_ptr[req_idx * bt_stride + block_number];
113 + int64_t slot_id = block_id * block_size + (clamped_position % block_size);
114 + out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id;
115 +
116 + int32_t seq_len = seq_lens_ptr[req_idx];
117 + int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1);
118 + new_seq_len = std::min(new_seq_len, static_cast<int32_t>(max_model_len));
119 + seq_lens_ptr[req_idx] = new_seq_len;
120 + }
121 +}
122 +
123 +void copy_and_expand_eagle_inputs_kernel_impl(
124 + const torch::Tensor& target_token_ids,
125 + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
126 + torch::Tensor& out_input_ids, torch::Tensor& out_positions,
127 + torch::Tensor& out_is_rejected_token_mask,
128 + torch::Tensor& out_is_masked_token_mask,
129 + torch::Tensor& out_new_token_indices,
130 + torch::Tensor& out_hidden_state_mapping,
131 + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
132 + const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
133 + const int64_t total_input_tokens,
134 + const int64_t num_padding_slots_per_request, const bool shift_input_ids) {
135 + const int64_t num_reqs = query_end_loc.size(0);
136 +
137 + const int64_t* target_ids_ptr = target_token_ids.data_ptr<int64_t>();
138 + const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
139 + const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
140 + const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
141 + const int32_t* query_end_ptr = query_end_loc.data_ptr<int32_t>();
142 +
143 + int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
144 + int64_t* out_pos_ptr = out_positions.data_ptr<int64_t>();
145 + bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr<bool>();
146 + bool* out_mask_ptr = out_is_masked_token_mask.data_ptr<bool>();
147 + int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr<int32_t>();
148 + int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr<int32_t>();
149 +
150 +#pragma omp parallel for
151 + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
152 + int32_t q_start = query_start_ptr[req_idx];
153 + int32_t next_q_start = query_start_ptr[req_idx + 1];
154 + int32_t q_end = query_end_ptr[req_idx];
155 +
156 + int64_t num_valid_tokens =
157 + shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1);
158 + int64_t input_offset = shift_input_ids ? 1 : 0;
159 +
160 + int64_t out_start = q_start + req_idx * (num_padding_slots_per_request -
161 + (shift_input_ids ? 1 : 0));
162 + int64_t num_rejected = next_q_start - q_end - 1;
163 + int64_t total_output_tokens =
164 + num_valid_tokens + num_padding_slots_per_request + num_rejected;
165 +
166 + int64_t start_pos = target_pos_ptr[q_start];
167 + int64_t bonus_token = next_ids_ptr[req_idx];
168 +
169 + for (int64_t j = 0; j < total_output_tokens; ++j) {
170 + int64_t out_idx = out_start + j;
171 + bool is_valid = j < num_valid_tokens;
172 + bool is_bonus = j == num_valid_tokens;
173 + bool is_parallel = (j > num_valid_tokens) &&
174 + (j < num_valid_tokens + num_padding_slots_per_request);
175 + bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request;
176 +
177 + int64_t in_idx =
178 + std::min(static_cast<int64_t>(q_start + input_offset + j),
179 + total_input_tokens - 1);
180 +
181 + int64_t token_id = padding_token_id;
182 + if (is_valid)
183 + token_id = target_ids_ptr[in_idx];
184 + else if (is_bonus)
185 + token_id = bonus_token;
186 + else if (is_parallel)
187 + token_id = parallel_drafting_token_id;
188 +
189 + out_ids_ptr[out_idx] = token_id;
190 + out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j);
191 + out_rej_mask_ptr[out_idx] = is_rejected;
192 + out_mask_ptr[out_idx] = is_parallel;
193 +
194 + if (is_bonus || is_parallel) {
195 + int64_t new_token_local_idx = j - num_valid_tokens;
196 + int64_t new_token_out_idx =
197 + req_idx * num_padding_slots_per_request + new_token_local_idx;
198 + out_new_idx_ptr[new_token_out_idx] = out_idx;
199 + }
200 + }
201 +
202 + if (shift_input_ids) {
203 + int64_t n_input = next_q_start - q_start;
204 + for (int64_t j = 0; j < n_input; ++j) {
205 + out_hidden_map_ptr[q_start + j] = out_start + j;
206 + }
207 + }
208 + }
209 +}
210 +
211 +void rejection_greedy_sample_kernel_impl(
212 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
213 + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
214 + const torch::Tensor& bonus_token_ids,
215 + const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len) {
216 + const int64_t batch_size = cu_num_draft_tokens.size(0);
217 +
218 + int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
219 + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
220 + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
221 + const int64_t* target_argmax_ptr = target_argmax.data_ptr<int64_t>();
222 + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
223 + const bool* greedy_ptr =
224 + is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
225 +
226 + const int64_t out_stride = output_token_ids.stride(0);
227 + const int64_t bonus_stride = bonus_token_ids.stride(0);
228 +
229 +#pragma omp parallel for
230 + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
231 + if (greedy_ptr && !greedy_ptr[req_idx]) continue;
232 +
233 + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
234 + int64_t end_idx = cu_draft_ptr[req_idx];
235 + int64_t num_draft_tokens = end_idx - start_idx;
236 +
237 + bool rejected = false;
238 + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
239 + int64_t target_id = target_argmax_ptr[start_idx + pos];
240 + out_ptr[req_idx * out_stride + pos] = target_id;
241 +
242 + if (draft_ids_ptr[start_idx + pos] != target_id) {
243 + rejected = true;
244 + break;
245 + }
246 + }
247 +
248 + if (!rejected) {
249 + out_ptr[req_idx * out_stride + num_draft_tokens] =
250 + bonus_ids_ptr[req_idx * bonus_stride];
251 + }
252 + }
253 +}
254 +
255 +void rejection_random_sample_kernel_impl(
256 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
257 + const torch::Tensor& draft_token_ids,
258 + const std::optional<torch::Tensor>& draft_probs,
259 + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
260 + const torch::Tensor& recovered_token_ids,
261 + const torch::Tensor& uniform_probs,
262 + const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
263 + const int64_t vocab_size, const bool no_draft_probs) {
264 + const int64_t batch_size = cu_num_draft_tokens.size(0);
265 +
266 + int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
267 + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
268 + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
269 + const float* draft_probs_ptr =
270 + no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
271 + const float* target_probs_ptr = target_probs.data_ptr<float>();
272 + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
273 + const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr<int64_t>();
274 + const float* uniform_probs_ptr = uniform_probs.data_ptr<float>();
275 + const bool* greedy_ptr =
276 + is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
277 +
278 + const int64_t out_stride = output_token_ids.stride(0);
279 + const int64_t bonus_stride = bonus_token_ids.stride(0);
280 + const int64_t target_stride = target_probs.stride(0);
281 + const int64_t draft_probs_stride =
282 + no_draft_probs ? 0 : draft_probs.value().stride(0);
283 +
284 +#pragma omp parallel for
285 + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
286 + if (greedy_ptr && greedy_ptr[req_idx]) continue;
287 +
288 + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
289 + int64_t end_idx = cu_draft_ptr[req_idx];
290 + int64_t num_draft_tokens = end_idx - start_idx;
291 +
292 + bool rejected = false;
293 + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
294 + int64_t token_idx = start_idx + pos;
295 + int64_t draft_id = draft_ids_ptr[token_idx];
296 +
297 + float p = target_probs_ptr[token_idx * target_stride + draft_id];
298 + float q =
299 + no_draft_probs
300 + ? 1.0f
301 + : draft_probs_ptr[token_idx * draft_probs_stride + draft_id];
302 + float uniform_p = uniform_probs_ptr[token_idx];
303 +
304 + float ratio = (q > 0.0f) ? (p / q) : 0.0f;
305 +
306 + if (ratio >= uniform_p) {
307 + out_ptr[req_idx * out_stride + pos] = draft_id;
308 + } else {
309 + out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx];
310 + rejected = true;
311 + break;
312 + }
313 + }
314 +
315 + if (!rejected) {
316 + out_ptr[req_idx * out_stride + num_draft_tokens] =
317 + bonus_ids_ptr[req_idx * bonus_stride];
318 + }
319 + }
320 +}
321 +
322 +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
323 + const torch::Tensor& cu_num_tokens,
324 + const int64_t replace_from, const int64_t replace_to) {
325 + const int64_t batch_size = cu_num_tokens.size(0);
326 + const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr<int64_t>();
327 +
328 + int64_t* out_ptr = output.data_ptr<int64_t>();
329 + const int64_t* in_ptr = input.data_ptr<int64_t>();
330 +
331 +#pragma omp parallel for
332 + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
333 + int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1];
334 + int64_t end_idx = cu_tokens_ptr[req_idx];
335 + int64_t val = in_ptr[req_idx];
336 +
337 + if (val == replace_from) {
338 + val = replace_to;
339 + }
340 +
341 + for (int64_t i = start_idx; i < end_idx; ++i) {
342 + out_ptr[i] = val;
343 + }
344 + }
345 +}
346 +
347 +void sample_recovered_tokens_kernel_impl(
348 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
349 + const torch::Tensor& draft_token_ids,
350 + const std::optional<torch::Tensor>& draft_probs,
351 + const torch::Tensor& target_probs, const torch::Tensor& inv_q,
352 + const int64_t vocab_size, const bool no_draft_probs) {
353 + const int64_t batch_size = cu_num_draft_tokens.size(0);
354 +
355 + int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
356 + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
357 + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
358 + const float* draft_probs_ptr =
359 + no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
360 + const float* target_probs_ptr = target_probs.data_ptr<float>();
361 + const float* inv_q_ptr = inv_q.data_ptr<float>();
362 +
363 + const int64_t target_stride = target_probs.stride(0);
364 + const int64_t draft_probs_stride =
365 + no_draft_probs ? 0 : draft_probs.value().stride(0);
366 + const int64_t inv_q_stride = inv_q.stride(0);
367 +
368 +#pragma omp parallel for
369 + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
370 + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
371 + int64_t end_idx = cu_draft_ptr[req_idx];
372 + int64_t num_draft_tokens = end_idx - start_idx;
373 +
374 + const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride;
375 +
376 + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
377 + int64_t token_idx = start_idx + pos;
378 + int64_t draft_id = draft_ids_ptr[token_idx];
379 +
380 + const float* token_target_probs =
381 + target_probs_ptr + token_idx * target_stride;
382 + const float* token_draft_probs =
383 + no_draft_probs ? nullptr
384 + : (draft_probs_ptr + token_idx * draft_probs_stride);
385 +
386 + int64_t best_id = 0;
387 + float best_val = -1.0f;
388 +
389 + for (int64_t v = 0; v < vocab_size; ++v) {
390 + float prob = token_target_probs[v];
391 + if (no_draft_probs) {
392 + if (v == draft_id) prob = 0.0f;
393 + } else {
394 + float diff = prob - token_draft_probs[v];
395 + prob = diff > 0.0f ? diff : 0.0f;
396 + }
397 +
398 + float val = prob * req_inv_q[v];
399 + if (val > best_val) {
400 + best_val = val;
401 + best_id = v;
402 + }
403 + }
404 + out_ptr[token_idx] = best_id;
405 + }
406 + }
407 +}
408 +
409 +} // namespace cpu_utils
csrc/cpu/torch_bindings.cpp
+119
@@ -138,6 +138,61 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
138 torch::Tensor slot_mapping,
139 const int64_t block_size);
140
141 +namespace cpu_utils {
142 +void eagle_prepare_inputs_padded_kernel_impl(
143 + const torch::Tensor& cu_num_draft_tokens,
144 + const torch::Tensor& valid_sampled_tokens_count,
145 + const torch::Tensor& query_start_loc_gpu,
146 + torch::Tensor& token_indices_to_sample,
147 + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs);
148 +void eagle_prepare_next_token_padded_kernel_impl(
149 + const torch::Tensor& sampled_token_ids,
150 + const torch::Tensor& discard_request_mask,
151 + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
152 + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
153 + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs);
154 +void eagle_step_slot_mapping_metadata_kernel_impl(
155 + const torch::Tensor& positions, const torch::Tensor& block_table,
156 + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
157 + torch::Tensor& out_slot_mapping, const int64_t block_size,
158 + const int64_t max_model_len, const int64_t PAD_ID);
159 +void copy_and_expand_eagle_inputs_kernel_impl(
160 + const torch::Tensor& target_token_ids,
161 + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
162 + torch::Tensor& out_input_ids, torch::Tensor& out_positions,
163 + torch::Tensor& out_is_rejected_token_mask,
164 + torch::Tensor& out_is_masked_token_mask,
165 + torch::Tensor& out_new_token_indices,
166 + torch::Tensor& out_hidden_state_mapping,
167 + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
168 + const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
169 + const int64_t total_input_tokens,
170 + const int64_t num_padding_slots_per_request, const bool shift_input_ids);
171 +void rejection_greedy_sample_kernel_impl(
172 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
173 + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
174 + const torch::Tensor& bonus_token_ids,
175 + const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len);
176 +void rejection_random_sample_kernel_impl(
177 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
178 + const torch::Tensor& draft_token_ids,
179 + const std::optional<torch::Tensor>& draft_probs,
180 + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
181 + const torch::Tensor& recovered_token_ids,
182 + const torch::Tensor& uniform_probs,
183 + const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
184 + const int64_t vocab_size, const bool no_draft_probs);
185 +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
186 + const torch::Tensor& cu_num_tokens,
187 + const int64_t replace_from, const int64_t replace_to);
188 +void sample_recovered_tokens_kernel_impl(
189 + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
190 + const torch::Tensor& draft_token_ids,
191 + const std::optional<torch::Tensor>& draft_probs,
192 + const torch::Tensor& target_probs, const torch::Tensor& inv_q,
193 + const int64_t vocab_size, const bool no_draft_probs);
194 +} // namespace cpu_utils
195 +
196 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
197 // vLLM custom ops
198
@@ -363,6 +418,70 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
418 "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt "
419 "block_size) -> ()",
420 &compute_slot_mapping_kernel_impl);
421 +
422 + // Speculative decoding kernels
423 + ops.def(
424 + "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
425 + "Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, "
426 + "Tensor(a3!) token_indices_to_sample, "
427 + "Tensor(a4!) num_rejected_tokens_gpu, "
428 + "SymInt num_reqs) -> ()",
429 + &cpu_utils::eagle_prepare_inputs_padded_kernel_impl);
430 + ops.def(
431 + "eagle_prepare_next_token_padded_kernel_impl("
432 + "Tensor sampled_token_ids, Tensor discard_request_mask, "
433 + "Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, "
434 + "Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, "
435 + "SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()",
436 + &cpu_utils::eagle_prepare_next_token_padded_kernel_impl);
437 + ops.def(
438 + "eagle_step_slot_mapping_metadata_kernel_impl("
439 + "Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, "
440 + "Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, "
441 + "SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()",
442 + &cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl);
443 + ops.def(
444 + "copy_and_expand_eagle_inputs_kernel_impl("
445 + "Tensor target_token_ids, Tensor target_positions, "
446 + "Tensor next_token_ids, Tensor(a3!) out_input_ids, "
447 + "Tensor(a4!) out_positions, "
448 + "Tensor(a5!) out_is_rejected_token_mask, "
449 + "Tensor(a6!) out_is_masked_token_mask, "
450 + "Tensor(a7!) out_new_token_indices, "
451 + "Tensor(a8!) out_hidden_state_mapping, "
452 + "Tensor query_start_loc, Tensor query_end_loc, "
453 + "SymInt padding_token_id, SymInt parallel_drafting_token_id, "
454 + "SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
455 + "bool shift_input_ids) -> ()",
456 + &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
457 + ops.def(
458 + "rejection_greedy_sample_kernel_impl("
459 + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
460 + "Tensor draft_token_ids, Tensor target_argmax, "
461 + "Tensor bonus_token_ids, Tensor? is_greedy, "
462 + "SymInt max_spec_len) -> ()",
463 + &cpu_utils::rejection_greedy_sample_kernel_impl);
464 + ops.def(
465 + "rejection_random_sample_kernel_impl("
466 + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
467 + "Tensor draft_token_ids, Tensor? draft_probs, "
468 + "Tensor target_probs, Tensor bonus_token_ids, "
469 + "Tensor recovered_token_ids, Tensor uniform_probs, "
470 + "Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, "
471 + "bool no_draft_probs) -> ()",
472 + &cpu_utils::rejection_random_sample_kernel_impl);
473 + ops.def(
474 + "expand_kernel_impl(Tensor(a0!) output, Tensor input, "
475 + "Tensor cu_num_tokens, SymInt replace_from, "
476 + "SymInt replace_to) -> ()",
477 + &cpu_utils::expand_kernel_impl);
478 + ops.def(
479 + "sample_recovered_tokens_kernel_impl("
480 + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
481 + "Tensor draft_token_ids, Tensor? draft_probs, "
482 + "Tensor target_probs, Tensor inv_q, SymInt vocab_size, "
483 + "bool no_draft_probs) -> ()",
484 + &cpu_utils::sample_recovered_tokens_kernel_impl);
485 }
486
487 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
vllm/utils/cpu_triton_utils.py
+274
@@ -45,3 +45,277 @@ def _compute_slot_mapping_kernel_impl(
45
46
47 compute_slot_mapping_kernel = _FuncWrapper(_compute_slot_mapping_kernel_impl)
48 +
49 +
50 +def _ensure_int64(t: torch.Tensor) -> torch.Tensor:
51 + return t if t.dtype == torch.int64 else t.to(torch.int64)
52 +
53 +
54 +def _eagle_prepare_inputs_padded_kernel_impl(
55 + cu_num_draft_tokens,
56 + valid_sampled_tokens_count,
57 + query_start_loc_gpu,
58 + token_indices_to_sample,
59 + num_rejected_tokens_gpu,
60 + num_reqs,
61 +):
62 + # C++ expects int64 for cu_num_draft_tokens, valid_sampled_tokens_count,
63 + # and num_rejected_tokens_gpu, but Python allocates them as int32.
64 + orig_rejected_dtype = num_rejected_tokens_gpu.dtype
65 + rejected_i64 = (
66 + num_rejected_tokens_gpu
67 + if orig_rejected_dtype == torch.int64
68 + else num_rejected_tokens_gpu.to(torch.int64)
69 + )
70 + torch.ops._C.eagle_prepare_inputs_padded_kernel_impl(
71 + _ensure_int64(cu_num_draft_tokens),
72 + _ensure_int64(valid_sampled_tokens_count),
73 + query_start_loc_gpu,
74 + token_indices_to_sample,
75 + rejected_i64,
76 + num_reqs,
77 + )
78 + if orig_rejected_dtype != torch.int64:
79 + num_rejected_tokens_gpu.copy_(rejected_i64.to(orig_rejected_dtype))
80 +
81 +
82 +def _eagle_prepare_next_token_padded_kernel_impl(
83 + sampled_token_ids,
84 + discard_request_mask,
85 + backup_next_token_ids,
86 + next_token_ids,
87 + valid_sampled_tokens_count,
88 + vocab_size,
89 + num_sampled_tokens_per_req,
90 + num_reqs,
91 + stride=None,
92 + BLOCK_SIZE_TOKENS=None,
93 +):
94 + # C++ reads all integer tensors as int64_t*. Output tensors are written
95 + # in-place so we create int64 copies, call C++, and copy back.
96 + orig_next_dtype = next_token_ids.dtype
97 + orig_valid_dtype = valid_sampled_tokens_count.dtype
98 + next_i64 = _ensure_int64(next_token_ids)
99 + valid_i64 = _ensure_int64(valid_sampled_tokens_count)
100 + torch.ops._C.eagle_prepare_next_token_padded_kernel_impl(
101 + _ensure_int64(sampled_token_ids),
102 + discard_request_mask,
103 + _ensure_int64(backup_next_token_ids),
104 + next_i64,
105 + valid_i64,
106 + vocab_size,
107 + num_sampled_tokens_per_req,
108 + num_reqs,
109 + )
110 + if orig_next_dtype != torch.int64:
111 + next_token_ids.copy_(next_i64.to(orig_next_dtype))
112 + if orig_valid_dtype != torch.int64:
113 + valid_sampled_tokens_count.copy_(valid_i64.to(orig_valid_dtype))
114 +
115 +
116 +def _eagle_step_slot_mapping_metadata_kernel_impl(
117 + positions,
118 + block_table,
119 + stride,
120 + seq_lens,
121 + out_clamped_positions,
122 + out_slot_mapping,
123 + block_size,
124 + max_model_len,
125 + n_blocks_per_req,
126 + PAD_ID,
127 + batch_size=None,
128 +):
129 + assert batch_size is None or batch_size == positions.shape[0], (
130 + f"batch_size mismatch: {batch_size} vs positions.shape[0]={positions.shape[0]}"
131 + )
132 + torch.ops._C.eagle_step_slot_mapping_metadata_kernel_impl(
133 + positions,
134 + block_table,
135 + seq_lens,
136 + out_clamped_positions,
137 + out_slot_mapping,
138 + block_size,
139 + max_model_len,
140 + PAD_ID,
141 + )
142 +
143 +
144 +def _copy_and_expand_eagle_inputs_kernel_impl(
145 + target_token_ids_ptr,
146 + target_positions_ptr,
147 + next_token_ids_ptr,
148 + out_input_ids_ptr,
149 + out_positions_ptr,
150 + out_is_rejected_token_mask_ptr,
151 + out_is_masked_token_mask_ptr,
152 + out_new_token_indices_ptr,
153 + out_hidden_state_mapping_ptr,
154 + query_start_loc_ptr,
155 + query_end_loc_ptr,
156 + padding_token_id,
157 + parallel_drafting_token_id,
158 + total_input_tokens,
159 + num_padding_slots_per_request,
160 + shift_input_ids,
161 + BLOCK_SIZE_TOKENS=None,
162 + BLOCK_SIZE_REQS=None,
163 +):
164 + """Adapter between Triton kernel call convention and C++ implementation.
165 +
166 + The Triton kernel uses '_ptr' suffixed parameter names and compile-time
167 + constants (BLOCK_SIZE_TOKENS, BLOCK_SIZE_REQS) which are not needed by
168 + the C++ implementation. C++ reads token id tensors as int64_t*.
169 + Output tensors that are int32 need copy-back after C++ writes int64.
170 + """
171 + orig_ids_dtype = out_input_ids_ptr.dtype
172 + orig_pos_dtype = out_positions_ptr.dtype
173 + out_ids_i64 = _ensure_int64(out_input_ids_ptr)
174 + out_pos_i64 = _ensure_int64(out_positions_ptr)
175 + torch.ops._C.copy_and_expand_eagle_inputs_kernel_impl(
176 + _ensure_int64(target_token_ids_ptr),
177 + _ensure_int64(target_positions_ptr),
178 + _ensure_int64(next_token_ids_ptr),
179 + out_ids_i64,
180 + out_pos_i64,
181 + out_is_rejected_token_mask_ptr,
182 + out_is_masked_token_mask_ptr,
183 + out_new_token_indices_ptr,
184 + out_hidden_state_mapping_ptr,
185 + query_start_loc_ptr,
186 + query_end_loc_ptr,
187 + padding_token_id,
188 + parallel_drafting_token_id,
189 + total_input_tokens,
190 + num_padding_slots_per_request,
191 + shift_input_ids,
192 + )
193 + if orig_ids_dtype != torch.int64:
194 + out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
195 + if orig_pos_dtype != torch.int64:
196 + out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype))
197 +
198 +
199 +def _rejection_greedy_sample_kernel_impl(
200 + output_token_ids,
201 + cu_num_draft_tokens,
202 + draft_token_ids,
203 + target_argmax,
204 + bonus_token_ids,
205 + is_greedy,
206 + max_spec_len,
207 +):
208 + # C++ kernel expects int64 for all integer tensors.
209 + orig_dtype = output_token_ids.dtype
210 + output_token_ids_i64 = _ensure_int64(output_token_ids)
211 + torch.ops._C.rejection_greedy_sample_kernel_impl(
212 + output_token_ids_i64,
213 + _ensure_int64(cu_num_draft_tokens),
214 + _ensure_int64(draft_token_ids),
215 + _ensure_int64(target_argmax),
216 + _ensure_int64(bonus_token_ids),
217 + is_greedy,
218 + max_spec_len,
219 + )
220 + if orig_dtype != torch.int64:
221 + output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
222 +
223 +
224 +def _rejection_random_sample_kernel_impl(
225 + output_token_ids,
226 + cu_num_draft_tokens,
227 + draft_token_ids,
228 + draft_probs,
229 + target_probs,
230 + bonus_token_ids,
231 + recovered_token_ids,
232 + uniform_probs,
233 + is_greedy,
234 + max_spec_len,
235 + vocab_size,
236 + NO_DRAFT_PROBS=False,
237 +):
238 + # C++ kernel expects int64 for all integer tensors and float32 for probs.
239 + # uniform_probs is intentionally float64 in Python to avoid exact-zero
240 + # samples; cast to float32 here for C++ compatibility.
241 + orig_dtype = output_token_ids.dtype
242 + output_token_ids_i64 = _ensure_int64(output_token_ids)
243 + torch.ops._C.rejection_random_sample_kernel_impl(
244 + output_token_ids_i64,
245 + _ensure_int64(cu_num_draft_tokens),
246 + _ensure_int64(draft_token_ids),
247 + draft_probs,
248 + target_probs,
249 + _ensure_int64(bonus_token_ids),
250 + _ensure_int64(recovered_token_ids),
251 + uniform_probs.to(torch.float32),
252 + is_greedy,
253 + max_spec_len,
254 + vocab_size,
255 + NO_DRAFT_PROBS,
256 + )
257 + if orig_dtype != torch.int64:
258 + output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
259 +
260 +
261 +def _expand_kernel_impl(
262 + output,
263 + input_val,
264 + cu_num_tokens,
265 + replace_from,
266 + replace_to,
267 + MAX_NUM_TOKENS=None,
268 +):
269 + torch.ops._C.expand_kernel_impl(
270 + _ensure_int64(output),
271 + _ensure_int64(input_val),
272 + _ensure_int64(cu_num_tokens),
273 + replace_from,
274 + replace_to,
275 + )
276 +
277 +
278 +def _sample_recovered_tokens_kernel_impl(
279 + output_token_ids,
280 + cu_num_draft_tokens,
281 + draft_token_ids,
282 + draft_probs,
283 + target_probs,
284 + inv_q,
285 + vocab_size,
286 + BLOCK_SIZE=None,
287 + NO_DRAFT_PROBS=False,
288 +):
289 + # C++ reads integer tensors as int64_t*; ensure correct dtype.
290 + orig_dtype = output_token_ids.dtype
291 + output_i64 = _ensure_int64(output_token_ids)
292 + torch.ops._C.sample_recovered_tokens_kernel_impl(
293 + output_i64,
294 + _ensure_int64(cu_num_draft_tokens),
295 + _ensure_int64(draft_token_ids),
296 + draft_probs,
297 + target_probs,
298 + inv_q,
299 + vocab_size,
300 + NO_DRAFT_PROBS,
301 + )
302 + if orig_dtype != torch.int64:
303 + output_token_ids.copy_(output_i64.to(orig_dtype))
304 +
305 +
306 +eagle_prepare_inputs_padded_kernel = _FuncWrapper(
307 + _eagle_prepare_inputs_padded_kernel_impl
308 +)
309 +eagle_prepare_next_token_padded_kernel = _FuncWrapper(
310 + _eagle_prepare_next_token_padded_kernel_impl
311 +)
312 +copy_and_expand_eagle_inputs_kernel = _FuncWrapper(
313 + _copy_and_expand_eagle_inputs_kernel_impl
314 +)
315 +eagle_step_slot_mapping_metadata_kernel = _FuncWrapper(
316 + _eagle_step_slot_mapping_metadata_kernel_impl
317 +)
318 +rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_impl)
319 +rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl)
320 +expand_kernel = _FuncWrapper(_expand_kernel_impl)
321 +sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl)
vllm/v1/spec_decode/eagle.py
+4 -5
@@ -26,7 +26,6 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
26 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
27 from vllm.multimodal import MULTIMODAL_REGISTRY
28 from vllm.platforms import current_platform
29 -from vllm.triton_utils import triton
29 from vllm.utils.platform_utils import is_pin_memory_available
30 from vllm.v1.attention.backend import CommonAttentionMetadata
31 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -48,6 +47,7 @@ from vllm.v1.spec_decode.utils import (
47 eagle_prepare_next_token_padded_kernel,
48 eagle_step_update_slot_mapping_and_metadata,
49 extend_all_queries_by_N,
50 + next_power_of_2,
51 )
52 from vllm.v1.utils import CpuGpuBuffer
53 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
@@ -689,9 +689,7 @@ class SpecDecodeBaseProposer:
689 max_num_tokens_per_request = (
690 cad.max_query_len + self.net_num_new_slots_per_request
691 )
692 - BLOCK_SIZE_TOKENS = min(
693 - 256, triton.next_power_of_2(max_num_tokens_per_request)
694 - )
692 + BLOCK_SIZE_TOKENS = min(256, next_power_of_2(max_num_tokens_per_request))
693 num_blocks = (
694 max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1
695 ) // BLOCK_SIZE_TOKENS
@@ -717,6 +715,7 @@ class SpecDecodeBaseProposer:
715 query_end_loc = cad.query_start_loc[1:] - 1
716 if num_rejected_tokens_gpu is not None:
717 query_end_loc = query_end_loc - num_rejected_tokens_gpu
718 +
719 copy_and_expand_eagle_inputs_kernel[grid](
720 # (Padded) Inputs from the target model
721 target_token_ids_ptr=target_token_ids,
@@ -899,7 +898,7 @@ class SpecDecodeBaseProposer:
898 grid = (batch_size,)
899
900 # Find the next power of 2 for block sizes
902 - BLOCK_SIZE_TOKENS = triton.next_power_of_2(num_tokens)
901 + BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)
902 eagle_prepare_next_token_padded_kernel[grid](
903 sampled_token_ids,
904 discard_request_mask,
vllm/v1/spec_decode/utils.py
+15 -1
@@ -11,6 +11,20 @@ from vllm.v1.attention.backends.utils import (
11 PADDING_SLOT_ID = -1
12
13
14 +def next_power_of_2(n: int) -> int:
15 + """Return the smallest power of 2 >= n."""
16 + if n <= 0:
17 + return 1
18 + n -= 1
19 + n |= n >> 1
20 + n |= n >> 2
21 + n |= n >> 4
22 + n |= n >> 8
23 + n |= n >> 16
24 + n |= n >> 32
25 + return n + 1
26 +
27 +
28 @triton.jit
29 def eagle_step_slot_mapping_metadata_kernel(
30 positions_ptr, # [batch_size] - current positions (1D view for M-RoPE)
@@ -102,8 +116,8 @@ def eagle_step_update_slot_mapping_and_metadata(
116 batch_size = positions_1d.shape[0]
117 if input_batch_size is None:
118 input_batch_size = batch_size
105 - n_blocks_per_req = block_table_tensor.shape[1]
119
120 + n_blocks_per_req = block_table_tensor.shape[1]
121 eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)](
122 positions_1d,
123 block_table_tensor,
vllm/v1/worker/cpu_model_runner.py
+121 -1
@@ -11,6 +11,8 @@ from vllm.config import VllmConfig
11 from vllm.logger import init_logger
12 from vllm.model_executor.model_loader import get_model
13 from vllm.tracing import instrument
14 +from vllm.v1.core.sched.output import SchedulerOutput
15 +from vllm.v1.kv_cache_interface import KVCacheConfig
16 from vllm.v1.utils import CpuGpuBuffer
17 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
18
@@ -23,7 +25,7 @@ class CPUModelRunner(GPUModelRunner):
25 super().__init__(vllm_config, device)
26
27 assert device == torch.device("cpu")
26 - assert self.speculative_config is None, "spec decode is not supported."
28 + # Note: speculative decoding is now supported on CPU with C++ native impls
29
30 self.use_cuda_graph = False
31 self.cascade_attn_enabled = False
@@ -61,6 +63,34 @@ class CPUModelRunner(GPUModelRunner):
63 cpu_tl.compute_slot_mapping_kernel
64 )
65
66 + # Speculative decoding fallbacks
67 + import vllm.v1.sample.rejection_sampler
68 + import vllm.v1.spec_decode.eagle
69 + import vllm.v1.spec_decode.utils
70 +
71 + vllm.v1.spec_decode.eagle.eagle_prepare_inputs_padded_kernel = (
72 + cpu_tl.eagle_prepare_inputs_padded_kernel
73 + )
74 + vllm.v1.spec_decode.eagle.eagle_prepare_next_token_padded_kernel = (
75 + cpu_tl.eagle_prepare_next_token_padded_kernel
76 + )
77 + vllm.v1.spec_decode.eagle.copy_and_expand_eagle_inputs_kernel = (
78 + cpu_tl.copy_and_expand_eagle_inputs_kernel
79 + )
80 + vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel = (
81 + cpu_tl.eagle_step_slot_mapping_metadata_kernel
82 + )
83 + vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = (
84 + cpu_tl.rejection_greedy_sample_kernel
85 + )
86 + vllm.v1.sample.rejection_sampler.rejection_random_sample_kernel = (
87 + cpu_tl.rejection_random_sample_kernel
88 + )
89 + vllm.v1.sample.rejection_sampler.expand_kernel = cpu_tl.expand_kernel
90 + vllm.v1.sample.rejection_sampler.sample_recovered_tokens_kernel = (
91 + cpu_tl.sample_recovered_tokens_kernel
92 + )
93 +
94 @instrument(span_name="Loading (CPU)")
95 def load_model(self, load_dummy_weights: bool = False) -> None:
96 if load_dummy_weights:
@@ -74,6 +104,10 @@ class CPUModelRunner(GPUModelRunner):
104 if self.lora_config:
105 self.model = self.load_lora_model(self.model, self.vllm_config, self.device)
106
107 + if hasattr(self, "drafter"):
108 + logger.info_once("Loading drafter model...")
109 + self.drafter.load_model(self.model)
110 +
111 def get_model(self) -> nn.Module:
112 return self.model
113
@@ -89,8 +123,29 @@ class CPUModelRunner(GPUModelRunner):
123 )
124 )
125
126 + # Warm up drafter for speculative decoding
127 + if self.speculative_config and (self.speculative_config.uses_draft_model()):
128 + from vllm.v1.spec_decode.draft_model import DraftModelProposer
129 +
130 + if isinstance(self.drafter, (DraftModelProposer)):
131 + logger.info("Warming up drafter model...")
132 + self.drafter.dummy_run(max(16, self.max_num_reqs))
133 +
134 logger.info("Warming up done.")
135
136 + def initialize_kv_cache(
137 + self,
138 + kv_cache_config: KVCacheConfig,
139 + is_profiling: bool = False,
140 + ) -> None:
141 + super().initialize_kv_cache(kv_cache_config, is_profiling)
142 +
143 + if self.speculative_config:
144 + if self.speculative_config.use_eagle():
145 + logger.info("EAGLE drafter KV cache initialized for CPU backend")
146 + elif self.speculative_config.uses_draft_model():
147 + logger.info("Draft model KV cache initialized for CPU backend")
148 +
149 def _init_device_properties(self) -> None:
150 pass
151
@@ -102,6 +157,71 @@ class CPUModelRunner(GPUModelRunner):
157 # so stale KV cache data never affects computation.
158 pass
159
160 + # =========================================================================
161 + # CPU-safe overrides for speculative decoding methods
162 + # These methods override GPU-specific implementations that use CUDA streams
163 + # =========================================================================
164 +
165 + def _copy_draft_token_ids_to_cpu(
166 + self, scheduler_output: "SchedulerOutput", zeros_only: bool = False
167 + ) -> None:
168 + """CPU-safe version: no async copy needed, tensors already on CPU."""
169 + if self.use_async_scheduling and not (
170 + scheduler_output.has_structured_output_requests
171 + or self.input_batch.sampling_metadata.output_token_ids
172 + ):
173 + return
174 + self._draft_token_req_ids = self.input_batch.req_ids.copy()
175 +
176 + draft_token_ids: torch.Tensor = self._draft_token_ids
177 + if not torch.is_tensor(draft_token_ids):
178 + return
179 +
180 + num_reqs = draft_token_ids.shape[0]
181 + if self.draft_token_ids_cpu is not None:
182 + if not zeros_only:
183 + self.draft_token_ids_cpu[:num_reqs].copy_(draft_token_ids)
184 + else:
185 + self.draft_token_ids_cpu[:num_reqs] = 0
186 +
187 + def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]:
188 + """CPU-safe version: no event synchronization needed."""
189 + if isinstance(self._draft_token_ids, list):
190 + return self._draft_token_ids, self.input_batch.req_ids
191 + req_ids = self._draft_token_req_ids
192 + if req_ids is None:
193 + return [], []
194 + if self.draft_token_ids_cpu is not None:
195 + return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids
196 + return [], []
197 +
198 + def _copy_valid_sampled_token_count(
199 + self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor
200 + ) -> None:
201 + """CPU-safe version: direct copy without CUDA streams."""
202 + if self.valid_sampled_token_count_cpu is None:
203 + return
204 +
205 + counts = valid_sampled_tokens_count
206 + counts_cpu = self.valid_sampled_token_count_cpu
207 + counts_cpu[: counts.shape[0]].copy_(counts)
208 + self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)
209 +
210 + def _get_valid_sampled_token_count(self) -> list[int]:
211 + """CPU-safe version: no event synchronization needed."""
212 + prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids
213 + if prev_sampled_token_ids is None:
214 + return []
215 +
216 + counts_cpu = self.valid_sampled_token_count_cpu
217 + if counts_cpu is None:
218 + return []
219 + return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist()
220 +
221 + def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
222 + """CPU-safe version: direct tolist() without CUDA events."""
223 + return sampled_token_ids.tolist()
224 +
225
226 @contextmanager
227 def _torch_cuda_wrapper():