| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Contains replacement functions to fallback Triton usages in CPU backend.""" |
| 5 | |
| 6 | import ctypes |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | |
| 11 | # For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py |
| 12 | def _compute_slot_mapping_kernel_impl( |
| 13 | num_tokens: int, |
| 14 | max_num_tokens: int, |
| 15 | query_start_loc: torch.Tensor, # [num_reqs + 1], int32 |
| 16 | positions: torch.Tensor, # [num_tokens], int64 |
| 17 | block_table: torch.Tensor, # [max_num_reqs, max_num_blocks_per_req], int32 |
| 18 | block_table_stride: int, # max_num_blocks_per_req |
| 19 | block_size: int, |
| 20 | slot_mapping: torch.Tensor, # [max_num_tokens], int64 |
| 21 | KV_CACHE_BLOCK_SIZE: int | None = None, |
| 22 | BLOCKS_PER_KV_BLOCK: int = 1, |
| 23 | TOTAL_CP_WORLD_SIZE: int = 1, |
| 24 | TOTAL_CP_RANK: int = 0, |
| 25 | CP_KV_CACHE_INTERLEAVE_SIZE: int = 1, |
| 26 | PAD_ID: int = -1, |
| 27 | BLOCK_SIZE: int = 1024, |
| 28 | ) -> None: |
| 29 | assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU." |
| 30 | if BLOCKS_PER_KV_BLOCK != 1: |
| 31 | assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE |
| 32 | torch.ops._C.compute_slot_mapping_kernel_impl( |
| 33 | query_start_loc, |
| 34 | positions, |
| 35 | block_table, |
| 36 | slot_mapping, |
| 37 | block_size, |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | def _ensure_int64(t: torch.Tensor) -> torch.Tensor: |
| 42 | return t if t.dtype == torch.int64 else t.to(torch.int64) |
| 43 | |
| 44 | |
| 45 | def _eagle_prepare_inputs_padded_kernel_impl( |
| 46 | cu_num_draft_tokens, |
| 47 | valid_sampled_tokens_count, |
| 48 | query_start_loc_gpu, |
| 49 | token_indices_to_sample, |
| 50 | num_rejected_tokens_gpu, |
| 51 | num_reqs, |
| 52 | ): |
| 53 | # C++ expects int64 for cu_num_draft_tokens, valid_sampled_tokens_count, |
| 54 | # and num_rejected_tokens_gpu, but Python allocates them as int32. |
| 55 | orig_rejected_dtype = num_rejected_tokens_gpu.dtype |
| 56 | rejected_i64 = ( |
| 57 | num_rejected_tokens_gpu |
| 58 | if orig_rejected_dtype == torch.int64 |
| 59 | else num_rejected_tokens_gpu.to(torch.int64) |
| 60 | ) |
| 61 | torch.ops._C.eagle_prepare_inputs_padded_kernel_impl( |
| 62 | _ensure_int64(cu_num_draft_tokens), |
| 63 | _ensure_int64(valid_sampled_tokens_count), |
| 64 | query_start_loc_gpu, |
| 65 | token_indices_to_sample, |
| 66 | rejected_i64, |
| 67 | num_reqs, |
| 68 | ) |
| 69 | if orig_rejected_dtype != torch.int64: |
| 70 | num_rejected_tokens_gpu.copy_(rejected_i64.to(orig_rejected_dtype)) |
| 71 | |
| 72 | |
| 73 | def _eagle_prepare_next_token_padded_kernel_impl( |
| 74 | sampled_token_ids, |
| 75 | discard_request_mask, |
| 76 | backup_next_token_ids, |
| 77 | next_token_ids, |
| 78 | valid_sampled_tokens_count, |
| 79 | vocab_size, |
| 80 | num_sampled_tokens_per_req, |
| 81 | num_reqs, |
| 82 | stride=None, |
| 83 | BLOCK_SIZE_TOKENS=None, |
| 84 | ): |
| 85 | # C++ reads all integer tensors as int64_t*. Output tensors are written |
| 86 | # in-place so we create int64 copies, call C++, and copy back. |
| 87 | orig_next_dtype = next_token_ids.dtype |
| 88 | orig_valid_dtype = valid_sampled_tokens_count.dtype |
| 89 | next_i64 = _ensure_int64(next_token_ids) |
| 90 | valid_i64 = _ensure_int64(valid_sampled_tokens_count) |
| 91 | torch.ops._C.eagle_prepare_next_token_padded_kernel_impl( |
| 92 | _ensure_int64(sampled_token_ids), |
| 93 | discard_request_mask, |
| 94 | _ensure_int64(backup_next_token_ids), |
| 95 | next_i64, |
| 96 | valid_i64, |
| 97 | vocab_size, |
| 98 | num_sampled_tokens_per_req, |
| 99 | num_reqs, |
| 100 | ) |
| 101 | if orig_next_dtype != torch.int64: |
| 102 | next_token_ids.copy_(next_i64.to(orig_next_dtype)) |
| 103 | if orig_valid_dtype != torch.int64: |
| 104 | valid_sampled_tokens_count.copy_(valid_i64.to(orig_valid_dtype)) |
| 105 | |
| 106 | |
| 107 | def _eagle_step_slot_mapping_metadata_kernel_impl( |
| 108 | positions, |
| 109 | block_table, |
| 110 | stride, |
| 111 | seq_lens, |
| 112 | out_clamped_positions, |
| 113 | out_slot_mapping, |
| 114 | block_size, |
| 115 | max_model_len, |
| 116 | n_blocks_per_req, |
| 117 | PAD_ID, |
| 118 | batch_size=None, |
| 119 | ): |
| 120 | assert batch_size is None or batch_size == positions.shape[0], ( |
| 121 | f"batch_size mismatch: {batch_size} vs positions.shape[0]={positions.shape[0]}" |
| 122 | ) |
| 123 | torch.ops._C.eagle_step_slot_mapping_metadata_kernel_impl( |
| 124 | positions, |
| 125 | block_table, |
| 126 | seq_lens, |
| 127 | out_clamped_positions, |
| 128 | out_slot_mapping, |
| 129 | block_size, |
| 130 | max_model_len, |
| 131 | PAD_ID, |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | def _copy_and_expand_eagle_inputs_kernel_impl( |
| 136 | target_token_ids_ptr, |
| 137 | target_positions_ptr, |
| 138 | next_token_ids_ptr, |
| 139 | out_input_ids_ptr, |
| 140 | out_positions_ptr, |
| 141 | out_is_rejected_token_mask_ptr, |
| 142 | out_is_masked_token_mask_ptr, |
| 143 | out_new_token_indices_ptr, |
| 144 | out_hidden_state_mapping_ptr, |
| 145 | query_start_loc_ptr, |
| 146 | query_end_loc_ptr, |
| 147 | padding_token_id, |
| 148 | parallel_drafting_token_id, |
| 149 | total_input_tokens, |
| 150 | num_padding_slots_per_request, |
| 151 | shift_input_ids, |
| 152 | BLOCK_SIZE_TOKENS=None, |
| 153 | BLOCK_SIZE_REQS=None, |
| 154 | ): |
| 155 | """Adapter between Triton kernel call convention and C++ implementation. |
| 156 | |
| 157 | The Triton kernel uses '_ptr' suffixed parameter names and compile-time |
| 158 | constants (BLOCK_SIZE_TOKENS, BLOCK_SIZE_REQS) which are not needed by |
| 159 | the C++ implementation. C++ reads token id tensors as int64_t*. |
| 160 | Output tensors that are int32 need copy-back after C++ writes int64. |
| 161 | """ |
| 162 | orig_ids_dtype = out_input_ids_ptr.dtype |
| 163 | orig_pos_dtype = out_positions_ptr.dtype |
| 164 | out_ids_i64 = _ensure_int64(out_input_ids_ptr) |
| 165 | out_pos_i64 = _ensure_int64(out_positions_ptr) |
| 166 | torch.ops._C.copy_and_expand_eagle_inputs_kernel_impl( |
| 167 | _ensure_int64(target_token_ids_ptr), |
| 168 | _ensure_int64(target_positions_ptr), |
| 169 | _ensure_int64(next_token_ids_ptr), |
| 170 | out_ids_i64, |
| 171 | out_pos_i64, |
| 172 | out_is_rejected_token_mask_ptr, |
| 173 | out_is_masked_token_mask_ptr, |
| 174 | out_new_token_indices_ptr, |
| 175 | out_hidden_state_mapping_ptr, |
| 176 | query_start_loc_ptr, |
| 177 | query_end_loc_ptr, |
| 178 | padding_token_id, |
| 179 | parallel_drafting_token_id, |
| 180 | total_input_tokens, |
| 181 | num_padding_slots_per_request, |
| 182 | shift_input_ids, |
| 183 | ) |
| 184 | if orig_ids_dtype != torch.int64: |
| 185 | out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype)) |
| 186 | if orig_pos_dtype != torch.int64: |
| 187 | out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype)) |
| 188 | |
| 189 | |
| 190 | def _copy_and_expand_dflash_inputs_kernel_impl( |
| 191 | next_token_ids_ptr, |
| 192 | target_positions_ptr, |
| 193 | out_input_ids_ptr, |
| 194 | out_context_positions_ptr, |
| 195 | out_query_positions_ptr, |
| 196 | out_context_slot_mapping_ptr, |
| 197 | out_query_slot_mapping_ptr, |
| 198 | out_token_indices_ptr, |
| 199 | block_table_ptr, |
| 200 | block_table_stride, |
| 201 | query_start_loc_ptr, |
| 202 | num_rejected_tokens_ptr, |
| 203 | parallel_drafting_token_id, |
| 204 | block_size, |
| 205 | num_query_per_req, |
| 206 | num_speculative_tokens, |
| 207 | total_input_tokens, |
| 208 | BLOCK_SIZE=None, |
| 209 | HAS_NUM_REJECTED=False, |
| 210 | ): |
| 211 | """Adapter between the DFlash Triton launch and the C++ CPU op.""" |
| 212 | assert block_table_stride == block_table_ptr.stride(0), ( |
| 213 | "block_table_stride mismatch: " |
| 214 | f"{block_table_stride} vs {block_table_ptr.stride(0)}" |
| 215 | ) |
| 216 | |
| 217 | orig_ids_dtype = out_input_ids_ptr.dtype |
| 218 | orig_context_positions_dtype = out_context_positions_ptr.dtype |
| 219 | orig_query_positions_dtype = out_query_positions_ptr.dtype |
| 220 | orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype |
| 221 | orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype |
| 222 | out_ids_i64 = _ensure_int64(out_input_ids_ptr) |
| 223 | out_context_positions_i64 = _ensure_int64(out_context_positions_ptr) |
| 224 | out_query_positions_i64 = _ensure_int64(out_query_positions_ptr) |
| 225 | out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr) |
| 226 | out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr) |
| 227 | rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None |
| 228 | |
| 229 | if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"): |
| 230 | torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl( |
| 231 | _ensure_int64(next_token_ids_ptr), |
| 232 | _ensure_int64(target_positions_ptr), |
| 233 | out_ids_i64, |
| 234 | out_context_positions_i64, |
| 235 | out_query_positions_i64, |
| 236 | out_context_slot_mapping_i64, |
| 237 | out_query_slot_mapping_i64, |
| 238 | out_token_indices_ptr, |
| 239 | block_table_ptr, |
| 240 | query_start_loc_ptr, |
| 241 | rejected_i64, |
| 242 | parallel_drafting_token_id, |
| 243 | block_size, |
| 244 | num_query_per_req, |
| 245 | num_speculative_tokens, |
| 246 | total_input_tokens, |
| 247 | HAS_NUM_REJECTED, |
| 248 | ) |
| 249 | else: |
| 250 | next_ids_i64 = _ensure_int64(next_token_ids_ptr) |
| 251 | target_positions_i64 = _ensure_int64(target_positions_ptr) |
| 252 | block_table_stride = block_table_ptr.stride(0) |
| 253 | num_reqs = query_start_loc_ptr.shape[0] - 1 |
| 254 | |
| 255 | for req_idx in range(num_reqs): |
| 256 | ctx_start = int(query_start_loc_ptr[req_idx].item()) |
| 257 | ctx_end = int(query_start_loc_ptr[req_idx + 1].item()) |
| 258 | num_ctx = ctx_end - ctx_start |
| 259 | valid_ctx_end = ctx_end |
| 260 | if rejected_i64 is not None: |
| 261 | valid_ctx_end -= int(rejected_i64[req_idx].item()) |
| 262 | # Guard against out-of-bounds: ensure valid_ctx_end > ctx_start. |
| 263 | valid_ctx_end = max(valid_ctx_end, ctx_start + 1) |
| 264 | |
| 265 | last_pos = int(target_positions_i64[valid_ctx_end - 1].item()) |
| 266 | |
| 267 | for j in range(num_ctx): |
| 268 | ctx_idx = ctx_start + j |
| 269 | ctx_pos_idx = min(ctx_idx, total_input_tokens - 1) |
| 270 | position = int(target_positions_i64[ctx_pos_idx].item()) |
| 271 | block_num = min(position // block_size, block_table_stride - 1) |
| 272 | block_id = int(block_table_ptr[req_idx, block_num].item()) |
| 273 | slot = block_id * block_size + (position % block_size) |
| 274 | |
| 275 | out_context_positions_i64[ctx_idx] = position |
| 276 | out_context_slot_mapping_i64[ctx_idx] = slot |
| 277 | |
| 278 | for query_off in range(num_query_per_req): |
| 279 | query_out = req_idx * num_query_per_req + query_off |
| 280 | position = last_pos + 1 + query_off |
| 281 | block_num = min(position // block_size, block_table_stride - 1) |
| 282 | block_id = int(block_table_ptr[req_idx, block_num].item()) |
| 283 | slot = block_id * block_size + (position % block_size) |
| 284 | |
| 285 | out_query_positions_i64[query_out] = position |
| 286 | out_query_slot_mapping_i64[query_out] = slot |
| 287 | out_ids_i64[query_out] = ( |
| 288 | int(next_ids_i64[req_idx].item()) |
| 289 | if query_off == 0 |
| 290 | else parallel_drafting_token_id |
| 291 | ) |
| 292 | |
| 293 | if query_off > 0: |
| 294 | sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1) |
| 295 | out_token_indices_ptr[sample_out_idx] = query_out |
| 296 | |
| 297 | if orig_ids_dtype != torch.int64: |
| 298 | out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype)) |
| 299 | if orig_context_positions_dtype != torch.int64: |
| 300 | out_context_positions_ptr.copy_( |
| 301 | out_context_positions_i64.to(orig_context_positions_dtype) |
| 302 | ) |
| 303 | if orig_query_positions_dtype != torch.int64: |
| 304 | out_query_positions_ptr.copy_( |
| 305 | out_query_positions_i64.to(orig_query_positions_dtype) |
| 306 | ) |
| 307 | if orig_context_slot_mapping_dtype != torch.int64: |
| 308 | out_context_slot_mapping_ptr.copy_( |
| 309 | out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype) |
| 310 | ) |
| 311 | if orig_query_slot_mapping_dtype != torch.int64: |
| 312 | out_query_slot_mapping_ptr.copy_( |
| 313 | out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype) |
| 314 | ) |
| 315 | |
| 316 | |
| 317 | def _rejection_greedy_sample_kernel_impl( |
| 318 | output_token_ids, |
| 319 | cu_num_draft_tokens, |
| 320 | draft_token_ids, |
| 321 | target_argmax, |
| 322 | bonus_token_ids, |
| 323 | is_greedy, |
| 324 | max_spec_len, |
| 325 | uniform_probs=None, |
| 326 | synthetic_conditional_rates=None, |
| 327 | SYNTHETIC_MODE=False, |
| 328 | ): |
| 329 | # C++ kernel expects int64 for all integer tensors. |
| 330 | # Note: uniform_probs, synthetic_conditional_rates, and SYNTHETIC_MODE are |
| 331 | # passed by the rejection sampler for synthetic mode support, but are not |
| 332 | # yet implemented in the C++ CPU kernel. We accept them here to maintain |
| 333 | # compatibility with the kernel calling convention. |
| 334 | assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling" |
| 335 | orig_dtype = output_token_ids.dtype |
| 336 | output_token_ids_i64 = _ensure_int64(output_token_ids) |
| 337 | torch.ops._C.rejection_greedy_sample_kernel_impl( |
| 338 | output_token_ids_i64, |
| 339 | _ensure_int64(cu_num_draft_tokens), |
| 340 | _ensure_int64(draft_token_ids), |
| 341 | _ensure_int64(target_argmax), |
| 342 | _ensure_int64(bonus_token_ids), |
| 343 | is_greedy, |
| 344 | max_spec_len, |
| 345 | ) |
| 346 | if orig_dtype != torch.int64: |
| 347 | output_token_ids.copy_(output_token_ids_i64.to(orig_dtype)) |
| 348 | |
| 349 | |
| 350 | def _rejection_random_sample_kernel_impl( |
| 351 | output_token_ids, |
| 352 | cu_num_draft_tokens, |
| 353 | draft_token_ids, |
| 354 | draft_probs, |
| 355 | target_probs, |
| 356 | bonus_token_ids, |
| 357 | recovered_token_ids, |
| 358 | uniform_probs, |
| 359 | is_greedy, |
| 360 | max_spec_len, |
| 361 | vocab_size, |
| 362 | synthetic_conditional_rates=None, |
| 363 | NO_DRAFT_PROBS=False, |
| 364 | SYNTHETIC_MODE=False, |
| 365 | ): |
| 366 | # C++ kernel expects int64 for all integer tensors and float32 for probs. |
| 367 | # uniform_probs is intentionally float64 in Python to avoid exact-zero |
| 368 | # samples; cast to float32 here for C++ compatibility. |
| 369 | # Note: synthetic_conditional_rates and SYNTHETIC_MODE are passed by the |
| 370 | # rejection sampler for synthetic mode support, but are not yet implemented |
| 371 | # in the C++ CPU kernel. We accept them here to maintain compatibility with |
| 372 | # the kernel calling convention. |
| 373 | assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling" |
| 374 | orig_dtype = output_token_ids.dtype |
| 375 | output_token_ids_i64 = _ensure_int64(output_token_ids) |
| 376 | torch.ops._C.rejection_random_sample_kernel_impl( |
| 377 | output_token_ids_i64, |
| 378 | _ensure_int64(cu_num_draft_tokens), |
| 379 | _ensure_int64(draft_token_ids), |
| 380 | draft_probs, |
| 381 | target_probs, |
| 382 | _ensure_int64(bonus_token_ids), |
| 383 | _ensure_int64(recovered_token_ids), |
| 384 | uniform_probs.to(torch.float32), |
| 385 | is_greedy, |
| 386 | max_spec_len, |
| 387 | vocab_size, |
| 388 | NO_DRAFT_PROBS, |
| 389 | ) |
| 390 | if orig_dtype != torch.int64: |
| 391 | output_token_ids.copy_(output_token_ids_i64.to(orig_dtype)) |
| 392 | |
| 393 | |
| 394 | def _expand_kernel_impl( |
| 395 | output, |
| 396 | input_val, |
| 397 | cu_num_tokens, |
| 398 | replace_from, |
| 399 | replace_to, |
| 400 | MAX_NUM_TOKENS=None, |
| 401 | ): |
| 402 | torch.ops._C.expand_kernel_impl( |
| 403 | _ensure_int64(output), |
| 404 | _ensure_int64(input_val), |
| 405 | _ensure_int64(cu_num_tokens), |
| 406 | replace_from, |
| 407 | replace_to, |
| 408 | ) |
| 409 | |
| 410 | |
| 411 | def _sample_recovered_tokens_kernel_impl( |
| 412 | output_token_ids, |
| 413 | cu_num_draft_tokens, |
| 414 | draft_token_ids, |
| 415 | draft_probs, |
| 416 | target_probs, |
| 417 | inv_q, |
| 418 | vocab_size, |
| 419 | BLOCK_SIZE=None, |
| 420 | NO_DRAFT_PROBS=False, |
| 421 | USE_FP64_GUMBEL=False, |
| 422 | ): |
| 423 | # USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller |
| 424 | # has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes |
| 425 | # `inv_q` directly, so the flag is accepted for interface parity and the |
| 426 | # value is read at its existing dtype. |
| 427 | # C++ reads integer tensors as int64_t*; ensure correct dtype. |
| 428 | orig_dtype = output_token_ids.dtype |
| 429 | output_i64 = _ensure_int64(output_token_ids) |
| 430 | torch.ops._C.sample_recovered_tokens_kernel_impl( |
| 431 | output_i64, |
| 432 | _ensure_int64(cu_num_draft_tokens), |
| 433 | _ensure_int64(draft_token_ids), |
| 434 | draft_probs, |
| 435 | target_probs, |
| 436 | # C++ kernel reads inv_q as float32. |
| 437 | inv_q.to(torch.float32), |
| 438 | vocab_size, |
| 439 | NO_DRAFT_PROBS, |
| 440 | ) |
| 441 | if orig_dtype != torch.int64: |
| 442 | output_token_ids.copy_(output_i64.to(orig_dtype)) |
| 443 | |
| 444 | |
| 445 | def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None): |
| 446 | # BLOCK_SIZE is unused; kept for signature parity with the Triton kernel. |
| 447 | for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()): |
| 448 | ctypes.memmove(dst, src, size) |