358 lines
15 KiB
Python
358 lines
15 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from vllm.config import VllmConfig, replace
|
||
|
|
from vllm.triton_utils import tl, triton
|
||
|
|
from vllm.v1.attention.backends.utils import (
|
||
|
|
CommonAttentionMetadata,
|
||
|
|
)
|
||
|
|
|
||
|
|
PADDING_SLOT_ID = -1
|
||
|
|
|
||
|
|
|
||
|
|
@triton.jit
|
||
|
|
def eagle_prepare_inputs_padded_kernel(
|
||
|
|
cu_num_draft_tokens_ptr, # [num_reqs]
|
||
|
|
valid_sampled_tokens_count_ptr, # [num_reqs]
|
||
|
|
query_start_loc_gpu_ptr, # [num_reqs + 1]
|
||
|
|
token_indices_to_sample_ptr, # [num_reqs] (output)
|
||
|
|
num_rejected_tokens_gpu_ptr, # [num_reqs] (output)
|
||
|
|
num_reqs, # tl.int32
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Fused kernel for Eagle prepare_input_padded. This kernel computes the
|
||
|
|
token index to sample for each request, taking into account the number
|
||
|
|
of draft tokens and the number of valid sampled tokens (which is one more than
|
||
|
|
the number of accepted tokens).
|
||
|
|
"""
|
||
|
|
req_idx = tl.program_id(axis=0)
|
||
|
|
if req_idx >= num_reqs:
|
||
|
|
return
|
||
|
|
|
||
|
|
# Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive
|
||
|
|
# cumulative sum (first entry is the first value, not zero).
|
||
|
|
cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx)
|
||
|
|
|
||
|
|
num_draft_tokens = 0
|
||
|
|
if req_idx == 0:
|
||
|
|
num_draft_tokens = cu_draft_curr
|
||
|
|
else:
|
||
|
|
cu_draft_prev = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
|
||
|
|
num_draft_tokens = cu_draft_curr - cu_draft_prev
|
||
|
|
|
||
|
|
valid_count = tl.load(valid_sampled_tokens_count_ptr + req_idx)
|
||
|
|
num_rejected_tokens = num_draft_tokens + 1 - valid_count
|
||
|
|
num_rejected_tokens = tl.where(num_draft_tokens > 0, num_rejected_tokens, 0)
|
||
|
|
|
||
|
|
# query_start_loc[req_idx + 1] is the start position of the next request,
|
||
|
|
# which is one past the last token of this request.
|
||
|
|
q_last_tok_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1) - 1
|
||
|
|
|
||
|
|
index_to_sample = q_last_tok_idx - num_rejected_tokens
|
||
|
|
tl.store(token_indices_to_sample_ptr + req_idx, index_to_sample)
|
||
|
|
tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens)
|
||
|
|
|
||
|
|
|
||
|
|
@triton.jit
|
||
|
|
def eagle_prepare_next_token_padded_kernel(
|
||
|
|
sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req]
|
||
|
|
discard_request_mask_ptr, # [num_reqs]
|
||
|
|
backup_next_token_ids_ptr, # [num_reqs]
|
||
|
|
next_token_ids_ptr, # [num_reqs] (output)
|
||
|
|
valid_sampled_tokens_count_ptr, # [num_reqs] (output)
|
||
|
|
vocab_size, # tl.int32
|
||
|
|
num_sampled_tokens_per_req, # tl.int32 (num_spec_tokens + 1)
|
||
|
|
num_reqs, # tl.int32
|
||
|
|
stride_sampled_token_ids, # tl.int32 (stride for dim 0)
|
||
|
|
BLOCK_SIZE_TOKENS: tl.constexpr, # Power-of-2 >= num_sampled_tokens_per_req
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the
|
||
|
|
number of valid (1 + accepted) tokens for each request, and the corresponding
|
||
|
|
"next" token id to sample from during speculative decoding. This is the
|
||
|
|
"last accepted token" from the sampled tokens, or the backup token if no
|
||
|
|
tokens were accepted or if the request is marked as discarded.
|
||
|
|
"""
|
||
|
|
req_idx = tl.program_id(axis=0)
|
||
|
|
if req_idx >= num_reqs:
|
||
|
|
return
|
||
|
|
|
||
|
|
# Check if this request is discarded.
|
||
|
|
is_discarded = tl.load(discard_request_mask_ptr + req_idx)
|
||
|
|
|
||
|
|
if is_discarded:
|
||
|
|
backup_token = tl.load(backup_next_token_ids_ptr + req_idx)
|
||
|
|
valid_count = tl.full((), 0, dtype=tl.uint32)
|
||
|
|
tl.store(next_token_ids_ptr + req_idx, backup_token)
|
||
|
|
tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
|
||
|
|
else:
|
||
|
|
# Count the number of valid tokens among the sampled tokens.
|
||
|
|
token_offs = tl.arange(0, BLOCK_SIZE_TOKENS)
|
||
|
|
token_mask = token_offs < num_sampled_tokens_per_req
|
||
|
|
|
||
|
|
row_ptr = sampled_token_ids_ptr + req_idx * stride_sampled_token_ids
|
||
|
|
token_ids = tl.load(row_ptr + token_offs, mask=token_mask, other=-1)
|
||
|
|
|
||
|
|
# Rejected tokens are -1, valid tokens are in [0, vocab_size)
|
||
|
|
is_valid_mask = (token_ids != -1) & (token_ids < vocab_size) & token_mask
|
||
|
|
valid_count = tl.sum(is_valid_mask)
|
||
|
|
|
||
|
|
if valid_count > 0:
|
||
|
|
# Guaranteed to be well-defined since
|
||
|
|
# valid_count > 0 implies is_valid_mask is not empty
|
||
|
|
last_valid_index = tl.max(tl.where(is_valid_mask, token_offs, -1))
|
||
|
|
|
||
|
|
# Select the token at that index, using a sum trick since
|
||
|
|
# we don't want to load again to access token_ids[last_valid_index].
|
||
|
|
last_valid_token = tl.sum(
|
||
|
|
tl.where(token_offs == last_valid_index, token_ids, 0)
|
||
|
|
)
|
||
|
|
tl.store(next_token_ids_ptr + req_idx, last_valid_token)
|
||
|
|
else:
|
||
|
|
# No valid tokens found, use backup token
|
||
|
|
backup_token = tl.load(backup_next_token_ids_ptr + req_idx)
|
||
|
|
tl.store(next_token_ids_ptr + req_idx, backup_token)
|
||
|
|
|
||
|
|
tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count)
|
||
|
|
|
||
|
|
|
||
|
|
def compute_new_slot_mapping(
|
||
|
|
cad: CommonAttentionMetadata,
|
||
|
|
new_positions: torch.Tensor,
|
||
|
|
is_rejected_token_mask: torch.Tensor,
|
||
|
|
block_size: int,
|
||
|
|
num_new_tokens: int,
|
||
|
|
max_model_len: int,
|
||
|
|
):
|
||
|
|
batch_size, n_blocks_per_req = cad.block_table_tensor.shape
|
||
|
|
req_indices = torch.arange(batch_size, device=cad.query_start_loc.device)
|
||
|
|
req_indices = torch.repeat_interleave(
|
||
|
|
req_indices,
|
||
|
|
cad.naive_query_lens() + num_new_tokens,
|
||
|
|
output_size=len(new_positions),
|
||
|
|
)
|
||
|
|
# Clamp the positions to prevent an out-of-bounds error when indexing
|
||
|
|
# into block_table_tensor.
|
||
|
|
clamped_positions = torch.clamp(new_positions, max=max_model_len - 1)
|
||
|
|
block_table_indices = (
|
||
|
|
req_indices * n_blocks_per_req + clamped_positions // block_size
|
||
|
|
)
|
||
|
|
block_nums = cad.block_table_tensor.view(-1)[block_table_indices]
|
||
|
|
block_offsets = clamped_positions % block_size
|
||
|
|
new_slot_mapping = block_nums * block_size + block_offsets
|
||
|
|
# Mask out the position ids that exceed the max model length.
|
||
|
|
exceeds_max_model_len = new_positions >= max_model_len
|
||
|
|
new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID)
|
||
|
|
# Mask out rejected tokens to prevent saves to the KV cache.
|
||
|
|
new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID)
|
||
|
|
return new_slot_mapping
|
||
|
|
|
||
|
|
|
||
|
|
def create_vllm_config_for_draft_model(
|
||
|
|
target_model_vllm_config: VllmConfig,
|
||
|
|
) -> VllmConfig:
|
||
|
|
"""The vllm_config is configured for the target model, e.g.
|
||
|
|
its quant_config and parallel_config. But the draft model is potentially
|
||
|
|
quantized differently, and has potentially different tensor_parallel_size.
|
||
|
|
This function creates a new vllm_config configured for the drafter.
|
||
|
|
The vllm_config is useful when loading the draft model with get_model().
|
||
|
|
"""
|
||
|
|
old = target_model_vllm_config
|
||
|
|
assert old.speculative_config is not None, "speculative_config is not set"
|
||
|
|
old_spec_config = old.speculative_config
|
||
|
|
new_parallel_config = replace(
|
||
|
|
old_spec_config.draft_parallel_config, rank=old.parallel_config.rank
|
||
|
|
)
|
||
|
|
new: VllmConfig = replace(
|
||
|
|
old,
|
||
|
|
quant_config=None,
|
||
|
|
parallel_config=new_parallel_config,
|
||
|
|
model_config=old_spec_config.draft_model_config,
|
||
|
|
)
|
||
|
|
return new
|
||
|
|
|
||
|
|
|
||
|
|
def extend_all_queries_by_N(
|
||
|
|
common_attn_metadata: CommonAttentionMetadata,
|
||
|
|
N: int,
|
||
|
|
arange: torch.Tensor,
|
||
|
|
new_slot_mapping: torch.Tensor,
|
||
|
|
) -> CommonAttentionMetadata:
|
||
|
|
"""
|
||
|
|
Creates a new CommonAttentionMetadata with all query lengths increased by N.
|
||
|
|
Also all seq lens are increased by N.
|
||
|
|
This is useful e.g. in speculative decoding with parallel drafting, where we
|
||
|
|
extend each sequence by N tokens and predict all tokens in one pass.
|
||
|
|
The slot mapping is computed externally, as it requires more information.
|
||
|
|
"""
|
||
|
|
cad = common_attn_metadata
|
||
|
|
# query start loc must be increased by [+0, +N, +2N, ..., +batch_size * N]
|
||
|
|
new_query_start_loc = cad.query_start_loc + N * arange[: len(cad.query_start_loc)]
|
||
|
|
new_query_start_loc_cpu = cad.query_start_loc_cpu + N * torch.arange(
|
||
|
|
len(cad.query_start_loc_cpu), dtype=torch.int32
|
||
|
|
)
|
||
|
|
new_cad = cad.replace(
|
||
|
|
query_start_loc=new_query_start_loc,
|
||
|
|
query_start_loc_cpu=new_query_start_loc_cpu,
|
||
|
|
seq_lens=cad.seq_lens + N,
|
||
|
|
# each request is extended by N tokens -> batch_size * N tokens are added
|
||
|
|
num_actual_tokens=cad.num_actual_tokens + cad.batch_size() * N,
|
||
|
|
# All query lens increase by N, so max query len increases by N
|
||
|
|
max_query_len=cad.max_query_len + N,
|
||
|
|
max_seq_len=cad.max_seq_len + N,
|
||
|
|
slot_mapping=new_slot_mapping,
|
||
|
|
)
|
||
|
|
return new_cad
|
||
|
|
|
||
|
|
|
||
|
|
# Unified copy/expand kernel
|
||
|
|
@triton.jit
|
||
|
|
def copy_and_expand_eagle_inputs_kernel(
|
||
|
|
# (Padded) Inputs from the target model
|
||
|
|
target_token_ids_ptr, # [total_tokens_in_batch]
|
||
|
|
target_positions_ptr, # [total_tokens_in_batch]
|
||
|
|
next_token_ids_ptr, # [num_reqs]
|
||
|
|
# Outputs to the drafting buffers
|
||
|
|
out_input_ids_ptr, # [total_draft_tokens_in_batch] (output)
|
||
|
|
out_positions_ptr, # [total_draft_tokens_in_batch] (output)
|
||
|
|
out_is_rejected_token_mask_ptr, # [total_draft_tokens_in_batch] (output)
|
||
|
|
out_is_masked_token_mask_ptr, # [total_draft_tokens_in_batch] (output)
|
||
|
|
out_new_token_indices_ptr, # [num_padding_slots_per_request * num_reqs] (output)
|
||
|
|
out_hidden_state_mapping_ptr, # [total_tokens_in_batch]
|
||
|
|
# Input metadata
|
||
|
|
query_start_loc_ptr, # [num_reqs + 1], last value is the total num input tokens
|
||
|
|
query_end_loc_ptr, # [num_reqs]
|
||
|
|
padding_token_id, # tl.int32
|
||
|
|
parallel_drafting_token_id, # tl.int32
|
||
|
|
# Sizing info
|
||
|
|
total_input_tokens, # tl.int32
|
||
|
|
num_padding_slots_per_request, # tl.int32
|
||
|
|
shift_input_ids, # tl.bool
|
||
|
|
BLOCK_SIZE_TOKENS: tl.constexpr, # Blocks along token dim to handle prefills
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Copy and expand inputs from the target model to the drafting buffers for Eagle
|
||
|
|
speculative decoding. This kernel handles padding slots and parallel drafting
|
||
|
|
tokens, if enabled.
|
||
|
|
"""
|
||
|
|
request_idx = tl.program_id(axis=0)
|
||
|
|
token_batch_idx = tl.program_id(axis=1)
|
||
|
|
|
||
|
|
# Load query locations
|
||
|
|
query_start_loc = tl.load(query_start_loc_ptr + request_idx)
|
||
|
|
next_query_start_loc = tl.load(query_start_loc_ptr + request_idx + 1)
|
||
|
|
query_end_loc = tl.load(query_end_loc_ptr + request_idx)
|
||
|
|
|
||
|
|
# Calculate number of valid tokens to copy and input offset
|
||
|
|
# With shift_input_ids=True, we skip the first token
|
||
|
|
# Output layout: each request gets (input_len + num_padding_slots_per_request) slots
|
||
|
|
# But with shift, we lose one token per request
|
||
|
|
if shift_input_ids:
|
||
|
|
num_valid_tokens = query_end_loc - query_start_loc
|
||
|
|
input_offset = 1
|
||
|
|
output_start = query_start_loc + request_idx * (
|
||
|
|
num_padding_slots_per_request - 1
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
num_valid_tokens = query_end_loc - query_start_loc + 1
|
||
|
|
input_offset = 0
|
||
|
|
output_start = query_start_loc + request_idx * num_padding_slots_per_request
|
||
|
|
|
||
|
|
# Number of rejected tokens from previous speculation
|
||
|
|
num_rejected = next_query_start_loc - query_end_loc - 1
|
||
|
|
|
||
|
|
# Total output tokens for this request
|
||
|
|
total_output_tokens = (
|
||
|
|
num_valid_tokens + num_padding_slots_per_request + num_rejected
|
||
|
|
)
|
||
|
|
|
||
|
|
# Process tokens in this block
|
||
|
|
j = token_batch_idx * BLOCK_SIZE_TOKENS + tl.arange(0, BLOCK_SIZE_TOKENS)
|
||
|
|
|
||
|
|
# Compute masks for different output regions:
|
||
|
|
# [0, num_valid_tokens): valid tokens copied from input
|
||
|
|
# [num_valid_tokens]: bonus token from next_token_ids
|
||
|
|
# (num_valid_tokens, num_valid_tokens + num_padding_slots_per_request):
|
||
|
|
# parallel drafting slots
|
||
|
|
# [num_valid_tokens + num_padding_slots_per_request, total_output_tokens):
|
||
|
|
# rejected slots
|
||
|
|
in_bounds = j < total_output_tokens
|
||
|
|
is_valid_region = j < num_valid_tokens
|
||
|
|
is_bonus_region = j == num_valid_tokens
|
||
|
|
is_parallel_draft_region = (j > num_valid_tokens) & (
|
||
|
|
j < num_valid_tokens + num_padding_slots_per_request
|
||
|
|
)
|
||
|
|
is_rejected_region = j >= num_valid_tokens + num_padding_slots_per_request
|
||
|
|
|
||
|
|
# Compute output indices
|
||
|
|
out_idx = output_start + j
|
||
|
|
|
||
|
|
# For valid tokens, compute input index
|
||
|
|
in_idx = query_start_loc + input_offset + j
|
||
|
|
# Clamp to avoid out-of-bounds access (masked loads still need valid addresses)
|
||
|
|
in_idx_clamped = tl.minimum(in_idx, total_input_tokens - 1)
|
||
|
|
|
||
|
|
# Load input tokens (masked to valid region)
|
||
|
|
token_ids = tl.load(
|
||
|
|
target_token_ids_ptr + in_idx_clamped, mask=is_valid_region & in_bounds, other=0
|
||
|
|
)
|
||
|
|
|
||
|
|
# Load the starting position for this request (first position in the sequence)
|
||
|
|
start_pos = tl.load(target_positions_ptr + query_start_loc)
|
||
|
|
|
||
|
|
# Load bonus token for this request
|
||
|
|
bonus_token = tl.load(next_token_ids_ptr + request_idx)
|
||
|
|
|
||
|
|
# Build final token_ids based on region
|
||
|
|
token_ids = tl.where(is_bonus_region, bonus_token, token_ids)
|
||
|
|
token_ids = tl.where(
|
||
|
|
is_parallel_draft_region, parallel_drafting_token_id, token_ids
|
||
|
|
)
|
||
|
|
token_ids = tl.where(is_rejected_region, padding_token_id, token_ids)
|
||
|
|
|
||
|
|
# Build final positions:
|
||
|
|
# Positions are NOT shifted - they start from the first input position and increment
|
||
|
|
# Output position j gets start_pos + j
|
||
|
|
# (e.g., input positions [5,6,7] -> output [5,6,7,8,9,...])
|
||
|
|
positions = start_pos + j
|
||
|
|
# Rejected positions are don't-care, set to 0
|
||
|
|
positions = tl.where(is_rejected_region, 0, positions)
|
||
|
|
|
||
|
|
# Compute output masks
|
||
|
|
is_rejected_out = is_rejected_region & in_bounds
|
||
|
|
is_masked_out = is_parallel_draft_region & in_bounds
|
||
|
|
|
||
|
|
# Compute indices of new tokens (bonus + parallel drafting) for sampling
|
||
|
|
# New tokens are at positions
|
||
|
|
# [num_valid_tokens, num_valid_tokens + num_padding_slots_per_request)
|
||
|
|
is_new_token_region = (j >= num_valid_tokens) & (
|
||
|
|
j < num_valid_tokens + num_padding_slots_per_request
|
||
|
|
)
|
||
|
|
new_token_local_idx = (
|
||
|
|
j - num_valid_tokens
|
||
|
|
) # 0 for bonus, 1, 2, ... for parallel drafting
|
||
|
|
new_token_out_idx = (
|
||
|
|
request_idx * num_padding_slots_per_request + new_token_local_idx
|
||
|
|
)
|
||
|
|
|
||
|
|
# Compute hidden state mapping (source index -> destination index)
|
||
|
|
# This maps each input position to its corresponding output position
|
||
|
|
# Hidden states don't get shifted, so we map all input tokens (including rejected)
|
||
|
|
if shift_input_ids:
|
||
|
|
num_input_tokens_this_request = next_query_start_loc - query_start_loc
|
||
|
|
is_input_region = j < num_input_tokens_this_request
|
||
|
|
src_idx = query_start_loc + j
|
||
|
|
tl.store(out_hidden_state_mapping_ptr + src_idx, out_idx, mask=is_input_region)
|
||
|
|
|
||
|
|
# Store outputs
|
||
|
|
tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds)
|
||
|
|
tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds)
|
||
|
|
tl.store(out_is_rejected_token_mask_ptr + out_idx, is_rejected_out, mask=in_bounds)
|
||
|
|
tl.store(out_is_masked_token_mask_ptr + out_idx, is_masked_out, mask=in_bounds)
|
||
|
|
tl.store(
|
||
|
|
out_new_token_indices_ptr + new_token_out_idx,
|
||
|
|
out_idx,
|
||
|
|
mask=is_new_token_region & in_bounds,
|
||
|
|
)
|