wehub-resource-sync
487 lines
17 KiB
Python
487 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Contains replacement functions to fallback Triton usages in CPU backend
|
|
"""
|
|
|
|
import ctypes
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
|
|
|
|
class _FuncWrapper:
|
|
def __init__(self, func: Callable) -> None:
|
|
self.func = func
|
|
|
|
def __getitem__(self, *args, **kwargs) -> Callable:
|
|
return self.func
|
|
|
|
|
|
# For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py
|
|
def _compute_slot_mapping_kernel_impl(
|
|
num_tokens: int,
|
|
max_num_tokens: int,
|
|
query_start_loc: torch.Tensor, # [num_reqs + 1], int32
|
|
positions: torch.Tensor, # [num_tokens], int64
|
|
block_table: torch.Tensor, # [max_num_reqs, max_num_blocks_per_req], int32
|
|
block_table_stride: int, # max_num_blocks_per_req
|
|
block_size: int,
|
|
slot_mapping: torch.Tensor, # [max_num_tokens], int64
|
|
KV_CACHE_BLOCK_SIZE: int | None = None,
|
|
BLOCKS_PER_KV_BLOCK: int = 1,
|
|
TOTAL_CP_WORLD_SIZE: int = 1,
|
|
TOTAL_CP_RANK: int = 0,
|
|
CP_KV_CACHE_INTERLEAVE_SIZE: int = 1,
|
|
PAD_ID: int = -1,
|
|
BLOCK_SIZE: int = 1024,
|
|
) -> None:
|
|
assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU."
|
|
if BLOCKS_PER_KV_BLOCK != 1:
|
|
assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE
|
|
torch.ops._C.compute_slot_mapping_kernel_impl(
|
|
query_start_loc,
|
|
positions,
|
|
block_table,
|
|
slot_mapping,
|
|
block_size,
|
|
)
|
|
|
|
|
|
compute_slot_mapping_kernel = _FuncWrapper(_compute_slot_mapping_kernel_impl)
|
|
|
|
|
|
def _ensure_int64(t: torch.Tensor) -> torch.Tensor:
|
|
return t if t.dtype == torch.int64 else t.to(torch.int64)
|
|
|
|
|
|
def _eagle_prepare_inputs_padded_kernel_impl(
|
|
cu_num_draft_tokens,
|
|
valid_sampled_tokens_count,
|
|
query_start_loc_gpu,
|
|
token_indices_to_sample,
|
|
num_rejected_tokens_gpu,
|
|
num_reqs,
|
|
):
|
|
# C++ expects int64 for cu_num_draft_tokens, valid_sampled_tokens_count,
|
|
# and num_rejected_tokens_gpu, but Python allocates them as int32.
|
|
orig_rejected_dtype = num_rejected_tokens_gpu.dtype
|
|
rejected_i64 = (
|
|
num_rejected_tokens_gpu
|
|
if orig_rejected_dtype == torch.int64
|
|
else num_rejected_tokens_gpu.to(torch.int64)
|
|
)
|
|
torch.ops._C.eagle_prepare_inputs_padded_kernel_impl(
|
|
_ensure_int64(cu_num_draft_tokens),
|
|
_ensure_int64(valid_sampled_tokens_count),
|
|
query_start_loc_gpu,
|
|
token_indices_to_sample,
|
|
rejected_i64,
|
|
num_reqs,
|
|
)
|
|
if orig_rejected_dtype != torch.int64:
|
|
num_rejected_tokens_gpu.copy_(rejected_i64.to(orig_rejected_dtype))
|
|
|
|
|
|
def _eagle_prepare_next_token_padded_kernel_impl(
|
|
sampled_token_ids,
|
|
discard_request_mask,
|
|
backup_next_token_ids,
|
|
next_token_ids,
|
|
valid_sampled_tokens_count,
|
|
vocab_size,
|
|
num_sampled_tokens_per_req,
|
|
num_reqs,
|
|
stride=None,
|
|
BLOCK_SIZE_TOKENS=None,
|
|
):
|
|
# C++ reads all integer tensors as int64_t*. Output tensors are written
|
|
# in-place so we create int64 copies, call C++, and copy back.
|
|
orig_next_dtype = next_token_ids.dtype
|
|
orig_valid_dtype = valid_sampled_tokens_count.dtype
|
|
next_i64 = _ensure_int64(next_token_ids)
|
|
valid_i64 = _ensure_int64(valid_sampled_tokens_count)
|
|
torch.ops._C.eagle_prepare_next_token_padded_kernel_impl(
|
|
_ensure_int64(sampled_token_ids),
|
|
discard_request_mask,
|
|
_ensure_int64(backup_next_token_ids),
|
|
next_i64,
|
|
valid_i64,
|
|
vocab_size,
|
|
num_sampled_tokens_per_req,
|
|
num_reqs,
|
|
)
|
|
if orig_next_dtype != torch.int64:
|
|
next_token_ids.copy_(next_i64.to(orig_next_dtype))
|
|
if orig_valid_dtype != torch.int64:
|
|
valid_sampled_tokens_count.copy_(valid_i64.to(orig_valid_dtype))
|
|
|
|
|
|
def _eagle_step_slot_mapping_metadata_kernel_impl(
|
|
positions,
|
|
block_table,
|
|
stride,
|
|
seq_lens,
|
|
out_clamped_positions,
|
|
out_slot_mapping,
|
|
block_size,
|
|
max_model_len,
|
|
n_blocks_per_req,
|
|
PAD_ID,
|
|
batch_size=None,
|
|
):
|
|
assert batch_size is None or batch_size == positions.shape[0], (
|
|
f"batch_size mismatch: {batch_size} vs positions.shape[0]={positions.shape[0]}"
|
|
)
|
|
torch.ops._C.eagle_step_slot_mapping_metadata_kernel_impl(
|
|
positions,
|
|
block_table,
|
|
seq_lens,
|
|
out_clamped_positions,
|
|
out_slot_mapping,
|
|
block_size,
|
|
max_model_len,
|
|
PAD_ID,
|
|
)
|
|
|
|
|
|
def _copy_and_expand_eagle_inputs_kernel_impl(
|
|
target_token_ids_ptr,
|
|
target_positions_ptr,
|
|
next_token_ids_ptr,
|
|
out_input_ids_ptr,
|
|
out_positions_ptr,
|
|
out_is_rejected_token_mask_ptr,
|
|
out_is_masked_token_mask_ptr,
|
|
out_new_token_indices_ptr,
|
|
out_hidden_state_mapping_ptr,
|
|
query_start_loc_ptr,
|
|
query_end_loc_ptr,
|
|
padding_token_id,
|
|
parallel_drafting_token_id,
|
|
total_input_tokens,
|
|
num_padding_slots_per_request,
|
|
shift_input_ids,
|
|
BLOCK_SIZE_TOKENS=None,
|
|
BLOCK_SIZE_REQS=None,
|
|
):
|
|
"""Adapter between Triton kernel call convention and C++ implementation.
|
|
|
|
The Triton kernel uses '_ptr' suffixed parameter names and compile-time
|
|
constants (BLOCK_SIZE_TOKENS, BLOCK_SIZE_REQS) which are not needed by
|
|
the C++ implementation. C++ reads token id tensors as int64_t*.
|
|
Output tensors that are int32 need copy-back after C++ writes int64.
|
|
"""
|
|
orig_ids_dtype = out_input_ids_ptr.dtype
|
|
orig_pos_dtype = out_positions_ptr.dtype
|
|
out_ids_i64 = _ensure_int64(out_input_ids_ptr)
|
|
out_pos_i64 = _ensure_int64(out_positions_ptr)
|
|
torch.ops._C.copy_and_expand_eagle_inputs_kernel_impl(
|
|
_ensure_int64(target_token_ids_ptr),
|
|
_ensure_int64(target_positions_ptr),
|
|
_ensure_int64(next_token_ids_ptr),
|
|
out_ids_i64,
|
|
out_pos_i64,
|
|
out_is_rejected_token_mask_ptr,
|
|
out_is_masked_token_mask_ptr,
|
|
out_new_token_indices_ptr,
|
|
out_hidden_state_mapping_ptr,
|
|
query_start_loc_ptr,
|
|
query_end_loc_ptr,
|
|
padding_token_id,
|
|
parallel_drafting_token_id,
|
|
total_input_tokens,
|
|
num_padding_slots_per_request,
|
|
shift_input_ids,
|
|
)
|
|
if orig_ids_dtype != torch.int64:
|
|
out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
|
|
if orig_pos_dtype != torch.int64:
|
|
out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype))
|
|
|
|
|
|
def _copy_and_expand_dflash_inputs_kernel_impl(
|
|
next_token_ids_ptr,
|
|
target_positions_ptr,
|
|
out_input_ids_ptr,
|
|
out_context_positions_ptr,
|
|
out_query_positions_ptr,
|
|
out_context_slot_mapping_ptr,
|
|
out_query_slot_mapping_ptr,
|
|
out_token_indices_ptr,
|
|
block_table_ptr,
|
|
block_table_stride,
|
|
query_start_loc_ptr,
|
|
num_rejected_tokens_ptr,
|
|
parallel_drafting_token_id,
|
|
block_size,
|
|
num_query_per_req,
|
|
num_speculative_tokens,
|
|
total_input_tokens,
|
|
BLOCK_SIZE=None,
|
|
HAS_NUM_REJECTED=False,
|
|
):
|
|
"""Adapter between the DFlash Triton launch and the C++ CPU op."""
|
|
assert block_table_stride == block_table_ptr.stride(0), (
|
|
"block_table_stride mismatch: "
|
|
f"{block_table_stride} vs {block_table_ptr.stride(0)}"
|
|
)
|
|
|
|
orig_ids_dtype = out_input_ids_ptr.dtype
|
|
orig_context_positions_dtype = out_context_positions_ptr.dtype
|
|
orig_query_positions_dtype = out_query_positions_ptr.dtype
|
|
orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype
|
|
orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype
|
|
out_ids_i64 = _ensure_int64(out_input_ids_ptr)
|
|
out_context_positions_i64 = _ensure_int64(out_context_positions_ptr)
|
|
out_query_positions_i64 = _ensure_int64(out_query_positions_ptr)
|
|
out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr)
|
|
out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr)
|
|
rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None
|
|
|
|
if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"):
|
|
torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl(
|
|
_ensure_int64(next_token_ids_ptr),
|
|
_ensure_int64(target_positions_ptr),
|
|
out_ids_i64,
|
|
out_context_positions_i64,
|
|
out_query_positions_i64,
|
|
out_context_slot_mapping_i64,
|
|
out_query_slot_mapping_i64,
|
|
out_token_indices_ptr,
|
|
block_table_ptr,
|
|
query_start_loc_ptr,
|
|
rejected_i64,
|
|
parallel_drafting_token_id,
|
|
block_size,
|
|
num_query_per_req,
|
|
num_speculative_tokens,
|
|
total_input_tokens,
|
|
HAS_NUM_REJECTED,
|
|
)
|
|
else:
|
|
next_ids_i64 = _ensure_int64(next_token_ids_ptr)
|
|
target_positions_i64 = _ensure_int64(target_positions_ptr)
|
|
block_table_stride = block_table_ptr.stride(0)
|
|
num_reqs = query_start_loc_ptr.shape[0] - 1
|
|
|
|
for req_idx in range(num_reqs):
|
|
ctx_start = int(query_start_loc_ptr[req_idx].item())
|
|
ctx_end = int(query_start_loc_ptr[req_idx + 1].item())
|
|
num_ctx = ctx_end - ctx_start
|
|
valid_ctx_end = ctx_end
|
|
if rejected_i64 is not None:
|
|
valid_ctx_end -= int(rejected_i64[req_idx].item())
|
|
# Guard against out-of-bounds: ensure valid_ctx_end > ctx_start.
|
|
valid_ctx_end = max(valid_ctx_end, ctx_start + 1)
|
|
|
|
last_pos = int(target_positions_i64[valid_ctx_end - 1].item())
|
|
|
|
for j in range(num_ctx):
|
|
ctx_idx = ctx_start + j
|
|
ctx_pos_idx = min(ctx_idx, total_input_tokens - 1)
|
|
position = int(target_positions_i64[ctx_pos_idx].item())
|
|
block_num = min(position // block_size, block_table_stride - 1)
|
|
block_id = int(block_table_ptr[req_idx, block_num].item())
|
|
slot = block_id * block_size + (position % block_size)
|
|
|
|
out_context_positions_i64[ctx_idx] = position
|
|
out_context_slot_mapping_i64[ctx_idx] = slot
|
|
|
|
for query_off in range(num_query_per_req):
|
|
query_out = req_idx * num_query_per_req + query_off
|
|
position = last_pos + 1 + query_off
|
|
block_num = min(position // block_size, block_table_stride - 1)
|
|
block_id = int(block_table_ptr[req_idx, block_num].item())
|
|
slot = block_id * block_size + (position % block_size)
|
|
|
|
out_query_positions_i64[query_out] = position
|
|
out_query_slot_mapping_i64[query_out] = slot
|
|
out_ids_i64[query_out] = (
|
|
int(next_ids_i64[req_idx].item())
|
|
if query_off == 0
|
|
else parallel_drafting_token_id
|
|
)
|
|
|
|
if query_off > 0:
|
|
sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1)
|
|
out_token_indices_ptr[sample_out_idx] = query_out
|
|
|
|
if orig_ids_dtype != torch.int64:
|
|
out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
|
|
if orig_context_positions_dtype != torch.int64:
|
|
out_context_positions_ptr.copy_(
|
|
out_context_positions_i64.to(orig_context_positions_dtype)
|
|
)
|
|
if orig_query_positions_dtype != torch.int64:
|
|
out_query_positions_ptr.copy_(
|
|
out_query_positions_i64.to(orig_query_positions_dtype)
|
|
)
|
|
if orig_context_slot_mapping_dtype != torch.int64:
|
|
out_context_slot_mapping_ptr.copy_(
|
|
out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype)
|
|
)
|
|
if orig_query_slot_mapping_dtype != torch.int64:
|
|
out_query_slot_mapping_ptr.copy_(
|
|
out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype)
|
|
)
|
|
|
|
|
|
def _rejection_greedy_sample_kernel_impl(
|
|
output_token_ids,
|
|
cu_num_draft_tokens,
|
|
draft_token_ids,
|
|
target_argmax,
|
|
bonus_token_ids,
|
|
is_greedy,
|
|
max_spec_len,
|
|
uniform_probs=None,
|
|
synthetic_conditional_rates=None,
|
|
SYNTHETIC_MODE=False,
|
|
):
|
|
# C++ kernel expects int64 for all integer tensors.
|
|
# Note: uniform_probs, synthetic_conditional_rates, and SYNTHETIC_MODE are
|
|
# passed by the rejection sampler for synthetic mode support, but are not
|
|
# yet implemented in the C++ CPU kernel. We accept them here to maintain
|
|
# compatibility with the kernel calling convention.
|
|
assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
|
|
orig_dtype = output_token_ids.dtype
|
|
output_token_ids_i64 = _ensure_int64(output_token_ids)
|
|
torch.ops._C.rejection_greedy_sample_kernel_impl(
|
|
output_token_ids_i64,
|
|
_ensure_int64(cu_num_draft_tokens),
|
|
_ensure_int64(draft_token_ids),
|
|
_ensure_int64(target_argmax),
|
|
_ensure_int64(bonus_token_ids),
|
|
is_greedy,
|
|
max_spec_len,
|
|
)
|
|
if orig_dtype != torch.int64:
|
|
output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
|
|
|
|
|
|
def _rejection_random_sample_kernel_impl(
|
|
output_token_ids,
|
|
cu_num_draft_tokens,
|
|
draft_token_ids,
|
|
draft_probs,
|
|
target_probs,
|
|
bonus_token_ids,
|
|
recovered_token_ids,
|
|
uniform_probs,
|
|
is_greedy,
|
|
max_spec_len,
|
|
vocab_size,
|
|
synthetic_conditional_rates=None,
|
|
NO_DRAFT_PROBS=False,
|
|
SYNTHETIC_MODE=False,
|
|
):
|
|
# C++ kernel expects int64 for all integer tensors and float32 for probs.
|
|
# uniform_probs is intentionally float64 in Python to avoid exact-zero
|
|
# samples; cast to float32 here for C++ compatibility.
|
|
# Note: synthetic_conditional_rates and SYNTHETIC_MODE are passed by the
|
|
# rejection sampler for synthetic mode support, but are not yet implemented
|
|
# in the C++ CPU kernel. We accept them here to maintain compatibility with
|
|
# the kernel calling convention.
|
|
assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
|
|
orig_dtype = output_token_ids.dtype
|
|
output_token_ids_i64 = _ensure_int64(output_token_ids)
|
|
torch.ops._C.rejection_random_sample_kernel_impl(
|
|
output_token_ids_i64,
|
|
_ensure_int64(cu_num_draft_tokens),
|
|
_ensure_int64(draft_token_ids),
|
|
draft_probs,
|
|
target_probs,
|
|
_ensure_int64(bonus_token_ids),
|
|
_ensure_int64(recovered_token_ids),
|
|
uniform_probs.to(torch.float32),
|
|
is_greedy,
|
|
max_spec_len,
|
|
vocab_size,
|
|
NO_DRAFT_PROBS,
|
|
)
|
|
if orig_dtype != torch.int64:
|
|
output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
|
|
|
|
|
|
def _expand_kernel_impl(
|
|
output,
|
|
input_val,
|
|
cu_num_tokens,
|
|
replace_from,
|
|
replace_to,
|
|
MAX_NUM_TOKENS=None,
|
|
):
|
|
torch.ops._C.expand_kernel_impl(
|
|
_ensure_int64(output),
|
|
_ensure_int64(input_val),
|
|
_ensure_int64(cu_num_tokens),
|
|
replace_from,
|
|
replace_to,
|
|
)
|
|
|
|
|
|
def _sample_recovered_tokens_kernel_impl(
|
|
output_token_ids,
|
|
cu_num_draft_tokens,
|
|
draft_token_ids,
|
|
draft_probs,
|
|
target_probs,
|
|
inv_q,
|
|
vocab_size,
|
|
BLOCK_SIZE=None,
|
|
NO_DRAFT_PROBS=False,
|
|
USE_FP64_GUMBEL=False,
|
|
):
|
|
# USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller
|
|
# has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes
|
|
# `inv_q` directly, so the flag is accepted for interface parity and the
|
|
# value is read at its existing dtype.
|
|
# C++ reads integer tensors as int64_t*; ensure correct dtype.
|
|
orig_dtype = output_token_ids.dtype
|
|
output_i64 = _ensure_int64(output_token_ids)
|
|
torch.ops._C.sample_recovered_tokens_kernel_impl(
|
|
output_i64,
|
|
_ensure_int64(cu_num_draft_tokens),
|
|
_ensure_int64(draft_token_ids),
|
|
draft_probs,
|
|
target_probs,
|
|
# C++ kernel reads inv_q as float32.
|
|
inv_q.to(torch.float32),
|
|
vocab_size,
|
|
NO_DRAFT_PROBS,
|
|
)
|
|
if orig_dtype != torch.int64:
|
|
output_token_ids.copy_(output_i64.to(orig_dtype))
|
|
|
|
|
|
eagle_prepare_inputs_padded_kernel = _FuncWrapper(
|
|
_eagle_prepare_inputs_padded_kernel_impl
|
|
)
|
|
eagle_prepare_next_token_padded_kernel = _FuncWrapper(
|
|
_eagle_prepare_next_token_padded_kernel_impl
|
|
)
|
|
copy_and_expand_eagle_inputs_kernel = _FuncWrapper(
|
|
_copy_and_expand_eagle_inputs_kernel_impl
|
|
)
|
|
copy_and_expand_dflash_inputs_kernel = _FuncWrapper(
|
|
_copy_and_expand_dflash_inputs_kernel_impl
|
|
)
|
|
eagle_step_slot_mapping_metadata_kernel = _FuncWrapper(
|
|
_eagle_step_slot_mapping_metadata_kernel_impl
|
|
)
|
|
rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_impl)
|
|
rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl)
|
|
expand_kernel = _FuncWrapper(_expand_kernel_impl)
|
|
sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl)
|
|
|
|
|
|
def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None):
|
|
# BLOCK_SIZE is unused; kept for signature parity with the Triton kernel.
|
|
for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()):
|
|
ctypes.memmove(dst, src, size)
|
|
|
|
|
|
batch_memcpy_kernel = _FuncWrapper(_batch_memcpy_impl)
|