947 lines
39 KiB
Python
947 lines
39 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-MLU project
|
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
from typing import Optional
|
|
import math
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
import vllm
|
|
from vllm.v1.sample.metadata import SamplingMetadata
|
|
from vllm.v1.sample import rejection_sampler
|
|
from vllm.v1.sample.rejection_sampler import sample_recovered_tokens
|
|
|
|
from vllm_mlu.mlu_hijack_utils import MluHijackObject
|
|
from vllm_mlu._mlu_utils import *
|
|
from vllm_mlu import _mlu_ops as mlu_ops
|
|
|
|
PLACEHOLDER_TOKEN_ID: tl.constexpr = -1
|
|
GREEDY_TEMPERATURE: tl.constexpr = 0
|
|
# Maximum number of speculative draft tokens allowed per request in a single
|
|
# step. This value is chosen to be large enough to handle typical use cases.
|
|
MAX_SPEC_LEN = 128
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief:
|
|
- Limit maximum batch size due to NRAM memory constraints
|
|
- Add generate_recovered_uniform_probs function for tmo rejection sampler
|
|
'''
|
|
MAX_BATCH_SIZE = 65536
|
|
|
|
def generate_recovered_uniform_probs(
|
|
num_tokens: int,
|
|
vocab_size: int,
|
|
num_draft_tokens: list[int],
|
|
sampling_metadata: SamplingMetadata,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
q = torch.empty(
|
|
(num_tokens, vocab_size),
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
q.exponential_()
|
|
for i, generator in sampling_metadata.generators.items():
|
|
# Do not generate random numbers for requests with no draft tokens.
|
|
# This can be important for reproducibility.
|
|
if num_draft_tokens[i] > 0:
|
|
q[i].exponential_(generator=generator)
|
|
return q
|
|
|
|
'''
|
|
=============================
|
|
End of MLU Hijack
|
|
=============================
|
|
'''
|
|
|
|
def vllm__v1__sample__rejection_sampler__expand_batch_to_tokens(
|
|
x: torch.Tensor, # [batch_size]
|
|
cu_num_tokens: torch.Tensor, # [batch_size]
|
|
num_tokens: int,
|
|
replace_from: int = 0,
|
|
replace_to: int = 0,
|
|
) -> torch.Tensor:
|
|
"""Expand [batch_size] tensor to [num_tokens] tensor based on the number of
|
|
tokens per batch in cu_num_tokens.
|
|
|
|
For example, if x = [a, b, c] and cu_num_tokens = [2, 5, 6], then
|
|
num_tokens = 6, and expanded_x = [a, a, b, b, b, c].
|
|
|
|
Args:
|
|
x: [batch_size] tensor to expand.
|
|
cu_num_tokens: [batch_size] tensor containing the cumulative number of
|
|
tokens per batch. Each element represents the total number of
|
|
tokens up to and including that batch.
|
|
num_tokens: Total number of tokens.
|
|
replace_from: int = 0
|
|
Value to be replaced if it is found in x.
|
|
replace_to: int = 0
|
|
Value to replace with when replace_from is found.
|
|
Returns:
|
|
expanded_x: [num_tokens] tensor.
|
|
"""
|
|
batch_size = x.shape[0]
|
|
assert cu_num_tokens.shape[0] == batch_size
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
'''
|
|
if batch_size > MAX_BATCH_SIZE:
|
|
raise ValueError(f"Rejection Sampler Not Supported: "
|
|
f"Batch size exceeds the maximum allowed value of {MAX_BATCH_SIZE}")
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
expanded_x = x.new_empty(num_tokens)
|
|
vllm__v1__sample__rejection_sampler__expand_kernel[(batch_size, )](
|
|
expanded_x,
|
|
x,
|
|
cu_num_tokens,
|
|
replace_from,
|
|
replace_to,
|
|
MAX_NUM_TOKENS=MAX_SPEC_LEN, # To avoid recompilation.
|
|
)
|
|
return expanded_x
|
|
|
|
# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
|
|
@triton.jit(do_not_specialize=["replace_from", "replace_to"])
|
|
def vllm__v1__sample__rejection_sampler__expand_kernel(
|
|
output_ptr, # [num_tokens]
|
|
input_ptr, # [batch_size]
|
|
cu_num_tokens_ptr, # [batch_size]
|
|
replace_from,
|
|
replace_to,
|
|
MAX_NUM_TOKENS: tl.constexpr,
|
|
):
|
|
req_idx = tl.program_id(0)
|
|
if req_idx == 0: # noqa: SIM108
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
'''
|
|
# Ensure data types are consistent
|
|
start_idx = tl.full((), 0, tl.int64)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
else:
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
'''
|
|
start_idx = tl.load(cu_num_tokens_ptr + req_idx - 1).to(tl.int64)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
end_idx = tl.load(cu_num_tokens_ptr + req_idx)
|
|
num_tokens = end_idx - start_idx
|
|
|
|
src_val = tl.load(input_ptr + req_idx)
|
|
src_val = tl.where(src_val == replace_from, replace_to, src_val)
|
|
offset = tl.arange(0, MAX_NUM_TOKENS)
|
|
tl.store(output_ptr + start_idx + offset, src_val, mask=offset < num_tokens)
|
|
|
|
|
|
@triton.jit
|
|
def vllm__v1__sample__rejection_sampler__sample_recovered_tokens_kernel(
|
|
output_token_ids_ptr, # [num_tokens]
|
|
cu_num_draft_tokens_ptr, # [batch_size]
|
|
draft_token_ids_ptr, # [num_tokens]
|
|
draft_probs_ptr, # [num_tokens, vocab_size] or None
|
|
target_probs_ptr, # [num_tokens, vocab_size]
|
|
q_ptr, # [batch_size, vocab_size]
|
|
vocab_size,
|
|
PADDED_VOCAB_SIZE: tl.constexpr,
|
|
NO_DRAFT_PROBS: tl.constexpr,
|
|
BLOCK_VOCAB: tl.constexpr = 2048,
|
|
):
|
|
req_idx = tl.program_id(0)
|
|
start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
|
|
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
|
|
num_draft_tokens = end_idx - start_idx
|
|
|
|
# Early exit for out-of-range positions.
|
|
pos = tl.program_id(1)
|
|
if pos >= num_draft_tokens:
|
|
return
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
'''
|
|
max_score = -float("inf")
|
|
max_index = 0
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
|
|
if NO_DRAFT_PROBS:
|
|
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
|
|
orig_prob = tl.load(target_probs_ptr + (start_idx + pos) * vocab_size +
|
|
draft_token_id)
|
|
# Temporarily zero out the probability of the draft token.
|
|
# This is essentially the same as target_prob - draft_prob, except that
|
|
# n-gram does not have draft_prob. We regard it as 1.
|
|
tl.store(
|
|
target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
|
|
0)
|
|
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: Replace with block loop due to ngram limitations
|
|
'''
|
|
num_blocks = tl.cdiv(PADDED_VOCAB_SIZE, BLOCK_VOCAB)
|
|
|
|
for i in tl.range(0, num_blocks):
|
|
offset = i * BLOCK_VOCAB + tl.arange(0, BLOCK_VOCAB)
|
|
mask = offset < vocab_size
|
|
|
|
if NO_DRAFT_PROBS:
|
|
prob = tl.load(
|
|
target_probs_ptr + (start_idx + pos) * vocab_size + offset,
|
|
mask=mask,
|
|
other=0
|
|
)
|
|
else:
|
|
draft_prob = tl.load(
|
|
draft_probs_ptr + (start_idx + pos) * vocab_size + offset,
|
|
mask=mask,
|
|
other=0
|
|
)
|
|
target_prob = tl.load(
|
|
target_probs_ptr + (start_idx + pos) * vocab_size + offset,
|
|
mask=mask,
|
|
other=0
|
|
)
|
|
prob = tl.maximum(target_prob - draft_prob, 0)
|
|
# NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
|
|
# `tl.argmax` will select the maximum value.
|
|
q = tl.load(q_ptr + req_idx * vocab_size + offset,
|
|
mask=mask,
|
|
other=float("-inf"))
|
|
score = prob / q # Broadcasting elementwise
|
|
cur_max = tl.argmax(score, axis=0)
|
|
|
|
cur_score = score[cur_max]
|
|
cur_index = offset[cur_max]
|
|
|
|
# Manually maintain argmax.
|
|
if cur_score > max_score:
|
|
max_score = cur_score
|
|
max_index = cur_index
|
|
|
|
tl.store(output_token_ids_ptr + start_idx + pos, max_index)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
|
|
if NO_DRAFT_PROBS:
|
|
# Restore the original probability.
|
|
tl.store(
|
|
target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id,
|
|
orig_prob)
|
|
|
|
"""
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
"""
|
|
def filter_with_acceptance_rate(output_token_ids, # [batch_size, max_spec_len + 1]
|
|
fixed_acceptance_rate):
|
|
"""
|
|
Filter speculative tokens based on a fixed acceptance rate using batch-level accept/reject decisions.
|
|
|
|
This function implements an adaptive acceptance rate control mechanism that maintains a target
|
|
acceptance rate over time through error compensation and PID-style adjustments.
|
|
|
|
Args:
|
|
output_token_ids (torch.Tensor): Input tensor of shape [batch_size, max_spec_len + 1]
|
|
where the first column contains base tokens and remaining columns contain speculative tokens
|
|
fixed_acceptance_rate (float or None): Target acceptance rate between 0.0 and 1.0
|
|
If None, returns input tensor unchanged
|
|
|
|
Returns:
|
|
torch.Tensor: Modified tensor where rejected batches have all speculative tokens
|
|
(columns 1 to max_spec_len) set to PLACEHOLDER_TOKEN_ID
|
|
|
|
Algorithm Flow:
|
|
1. **Initialization Phase**:
|
|
- Extract batch dimensions and device information
|
|
- Initialize static variables for tracking acceptance statistics:
|
|
* cumulative_error: Long-term error accumulation
|
|
* total_batches/accepted_batches: Global acceptance tracking
|
|
* acceptance_history: Sliding window for recent performance
|
|
* precision_adjustment: PID controller adjustment factor
|
|
* recent_adjustments: Error history for PID calculation
|
|
|
|
2. **Statistics Calculation**:
|
|
- Calculate global acceptance rate from all historical data
|
|
- Calculate sliding window acceptance rate from recent batches
|
|
- Compute combined error using weighted average of global and window errors
|
|
- Weight transitions from global-focused (early) to window-focused (later)
|
|
|
|
3. **PID Controller Adjustment** (after 50+ batches):
|
|
- Proportional term: Current error magnitude
|
|
- Integral term: Accumulated error over recent history
|
|
- Derivative term: Rate of error change
|
|
- Combines P, I, D terms to compute precision adjustment factor
|
|
- Limits adjustment range to prevent over-correction
|
|
|
|
4. **Error Correction**:
|
|
- Applies smooth nonlinear correction based on combined error magnitude
|
|
- Uses exponential decay mapping for gradual adjustment strength
|
|
- Handles boundary cases (0.0, 1.0, very low rates) specially
|
|
|
|
5. **Gap-based Adjustment**:
|
|
- Calculates difference between target and actual accepted batches
|
|
- Applies adaptive threshold-based corrections
|
|
- Uses exponential smoothing for adjustment strength
|
|
- Adjustment strength decreases as total batch count increases
|
|
|
|
6. **Random Perturbation** (after 100+ batches):
|
|
- Adds small random noise to prevent local minima
|
|
- Noise amplitude decreases over time for stability
|
|
|
|
7. **Batch Decision**:
|
|
- Generates random value and compares with adjusted acceptance rate
|
|
- Makes binary accept/reject decision for entire batch
|
|
|
|
8. **Token Modification**:
|
|
- If accepted: Leave all tokens unchanged
|
|
- If rejected: Set all speculative tokens (columns 1:) to PLACEHOLDER_TOKEN_ID
|
|
- This ensures token-level acceptance rate matches batch-level rate
|
|
|
|
9. **State Updates**:
|
|
- Update acceptance counters and history
|
|
- Update cumulative error using exponential moving average
|
|
- Prepare state for next function call
|
|
|
|
Key Features:
|
|
- **Batch-level consistency**: All samples in a batch share the same accept/reject fate
|
|
- **Adaptive control**: Uses multiple feedback mechanisms (global, windowed, PID)
|
|
- **Error compensation**: Corrects for deviations from target rate over time
|
|
- **Stability mechanisms**: Includes smoothing, limits, and perturbation for robustness
|
|
- **Token-level alignment**: Ensures token acceptance rate matches batch acceptance rate
|
|
|
|
Note: This function maintains internal state across calls through static variables,
|
|
so it will converge to the target acceptance rate over multiple invocations.
|
|
"""
|
|
if fixed_acceptance_rate is None:
|
|
return output_token_ids
|
|
else:
|
|
# Apply accept/reject decisions for the entire batch based on fixed_acceptance_rate
|
|
batch_size = output_token_ids.shape[0]
|
|
max_spec_len = output_token_ids.shape[1] - 1 # Get max_spec_len
|
|
device = output_token_ids.device
|
|
|
|
assert fixed_acceptance_rate >= 0 and fixed_acceptance_rate <= 1
|
|
|
|
# Use error compensation method to track global acceptance rate
|
|
# These are static variables that persist between calls
|
|
if not hasattr(filter_with_acceptance_rate, "cumulative_error"):
|
|
filter_with_acceptance_rate.cumulative_error = 0.0
|
|
if not hasattr(filter_with_acceptance_rate, "total_batches"):
|
|
filter_with_acceptance_rate.total_batches = 0
|
|
if not hasattr(filter_with_acceptance_rate, "accepted_batches"):
|
|
filter_with_acceptance_rate.accepted_batches = 0
|
|
if not hasattr(filter_with_acceptance_rate, "window_size"):
|
|
filter_with_acceptance_rate.window_size = 1000 # Sliding window size
|
|
if not hasattr(filter_with_acceptance_rate, "acceptance_history"):
|
|
filter_with_acceptance_rate.acceptance_history = [] # Track recent accept/reject history
|
|
if not hasattr(filter_with_acceptance_rate, "precision_adjustment"):
|
|
filter_with_acceptance_rate.precision_adjustment = 0.0 # Precision adjustment factor
|
|
if not hasattr(filter_with_acceptance_rate, "recent_adjustments"):
|
|
filter_with_acceptance_rate.recent_adjustments = [] # Recent adjustment history
|
|
if not hasattr(filter_with_acceptance_rate, "target_rate"):
|
|
filter_with_acceptance_rate.target_rate = fixed_acceptance_rate # Record target acceptance rate
|
|
else:
|
|
# If target acceptance rate changes, reset adjustment state
|
|
if filter_with_acceptance_rate.target_rate != fixed_acceptance_rate:
|
|
filter_with_acceptance_rate.precision_adjustment = 0.0
|
|
filter_with_acceptance_rate.recent_adjustments = []
|
|
filter_with_acceptance_rate.target_rate = fixed_acceptance_rate
|
|
|
|
# Update batch count
|
|
filter_with_acceptance_rate.total_batches += 1
|
|
|
|
# Calculate current global acceptance rate
|
|
global_rate = (filter_with_acceptance_rate.accepted_batches /
|
|
filter_with_acceptance_rate.total_batches if
|
|
filter_with_acceptance_rate.total_batches > 0 else 0.0)
|
|
|
|
# Calculate sliding window acceptance rate (focusing on recent performance)
|
|
filter_with_acceptance_rate.acceptance_history.append(0) # Default to reject
|
|
if len(filter_with_acceptance_rate.acceptance_history) > filter_with_acceptance_rate.window_size:
|
|
filter_with_acceptance_rate.acceptance_history.pop(0) # Remove oldest record
|
|
|
|
window_rate = sum(filter_with_acceptance_rate.acceptance_history) / len(filter_with_acceptance_rate.acceptance_history)
|
|
|
|
# Enhance precision for small batches - use smoother weight function
|
|
batch_weight_factor = 1.0 - math.exp(-filter_with_acceptance_rate.total_batches / 30.0) # Exponential smooth transition
|
|
|
|
# Dynamically adjust error weights: rely more on global error for fewer batches,
|
|
# more on sliding window error as batch count increases
|
|
window_size = len(filter_with_acceptance_rate.acceptance_history)
|
|
window_significance = min(window_size / 100.0, 0.9) # Window significance depends on historical data volume
|
|
window_weight = window_significance * batch_weight_factor
|
|
global_weight = 1.0 - window_weight
|
|
|
|
# Consider both global error and window error
|
|
combined_error = (global_weight * (global_rate - fixed_acceptance_rate) +
|
|
window_weight * (window_rate - fixed_acceptance_rate))
|
|
|
|
# Update precision adjustment factor - use PID controller style adjustment
|
|
if filter_with_acceptance_rate.total_batches > 50:
|
|
# Only perform precision adjustment when there's enough data
|
|
current_error = global_rate - fixed_acceptance_rate
|
|
|
|
# Save recent adjustment history
|
|
filter_with_acceptance_rate.recent_adjustments.append(current_error)
|
|
if len(filter_with_acceptance_rate.recent_adjustments) > 20: # Keep recent 20 errors
|
|
filter_with_acceptance_rate.recent_adjustments.pop(0)
|
|
|
|
# PID controller parameters
|
|
kp = 0.05 # Proportional coefficient
|
|
ki = 0.001 # Integral coefficient
|
|
kd = 0.01 # Derivative coefficient
|
|
|
|
# Proportional term - current error
|
|
p_term = current_error
|
|
|
|
# Integral term - accumulated error
|
|
i_term = sum(filter_with_acceptance_rate.recent_adjustments)
|
|
|
|
# Derivative term - error change rate
|
|
d_term = 0.0
|
|
if len(filter_with_acceptance_rate.recent_adjustments) >= 2:
|
|
d_term = filter_with_acceptance_rate.recent_adjustments[-1] - filter_with_acceptance_rate.recent_adjustments[-2]
|
|
|
|
# Calculate PID adjustment
|
|
pid_adjustment = kp * p_term + ki * i_term + kd * d_term
|
|
|
|
# Update precision adjustment factor
|
|
filter_with_acceptance_rate.precision_adjustment = pid_adjustment
|
|
|
|
# Limit adjustment factor range to prevent over-adjustment
|
|
max_adjustment = 0.02 + 0.03 * (1.0 - math.exp(-filter_with_acceptance_rate.total_batches / 500.0))
|
|
filter_with_acceptance_rate.precision_adjustment = max(-max_adjustment, min(max_adjustment, filter_with_acceptance_rate.precision_adjustment))
|
|
|
|
# Calculate acceptance rate correction factor
|
|
error_magnitude = abs(combined_error)
|
|
correction_factor = 1.0
|
|
|
|
# Use more refined error correction logic - use smooth nonlinear correction function
|
|
if error_magnitude > 0.0005: # Correct even smaller errors
|
|
# Use smooth correction function instead of piecewise function
|
|
base_strength = 2.0
|
|
error_scale = 1.0 - math.exp(-error_magnitude * 50.0) # Exponential decay mapping to [0,1]
|
|
correction_strength = base_strength + error_scale * 1.5 # Range from 2.0 to 3.5
|
|
|
|
# Smooth correction
|
|
sign = 1 if combined_error > 0 else -1
|
|
correction_factor = 1.0 + (correction_strength * error_magnitude * sign)
|
|
|
|
# Handle boundary cases to avoid division by zero
|
|
if correction_factor == 0.0:
|
|
correction_factor = 1.0
|
|
|
|
# Apply correction factor
|
|
adjusted_rate = max(0.0, min(1.0, fixed_acceptance_rate * (1.0 / correction_factor)))
|
|
|
|
# Apply precision adjustment factor
|
|
adjusted_rate = max(0.0, min(1.0, adjusted_rate - filter_with_acceptance_rate.precision_adjustment))
|
|
|
|
# More precise boundary case handling
|
|
if fixed_acceptance_rate > 0 and fixed_acceptance_rate < 0.05:
|
|
if filter_with_acceptance_rate.total_batches % int(1/fixed_acceptance_rate) == 0:
|
|
adjusted_rate = 1.0 # Periodically force accept to ensure accuracy in low acceptance rate scenarios
|
|
# If fixed_acceptance_rate is 0, directly reject
|
|
elif fixed_acceptance_rate == 0.0:
|
|
adjusted_rate = 0.0
|
|
# If fixed_acceptance_rate is 1, directly accept
|
|
elif fixed_acceptance_rate == 1.0:
|
|
adjusted_rate = 1.0
|
|
|
|
# Make precise adjustments for cases with large remaining errors
|
|
target_accepted = int(filter_with_acceptance_rate.total_batches * fixed_acceptance_rate + 0.5) # Round to nearest
|
|
actual_accepted = filter_with_acceptance_rate.accepted_batches
|
|
acceptance_gap = target_accepted - actual_accepted
|
|
|
|
# More aggressive gap adjustment strategy - use adaptive threshold and smooth adjustment
|
|
gap_relative = abs(acceptance_gap) / max(1, filter_with_acceptance_rate.total_batches)
|
|
gap_threshold = max(1, int(filter_with_acceptance_rate.total_batches * 0.005)) # Smaller dynamic threshold, at least 1
|
|
|
|
# Dynamically adjust acceptance rate based on the gap
|
|
if abs(acceptance_gap) >= gap_threshold: # Use dynamic threshold
|
|
# Use smooth adjustment strategy
|
|
if acceptance_gap > 0: # Need to accept more
|
|
# Use exponential function for smooth adjustment
|
|
gap_importance = 1.0 - math.exp(-gap_relative * 50.0) # Map to [0,1]
|
|
# Adjustment strength decreases as total batch count increases
|
|
strength_factor = math.exp(-filter_with_acceptance_rate.total_batches / 1000.0)
|
|
boost_factor = gap_importance * (0.2 + 0.8 * strength_factor) # Range from 0 to 1, decreases with total batch count
|
|
adjusted_rate = min(1.0, adjusted_rate + (1.0 - adjusted_rate) * boost_factor)
|
|
else: # Accepted too many, need to reject
|
|
# Use exponential function for smooth adjustment
|
|
gap_importance = 1.0 - math.exp(-gap_relative * 50.0) # Map to [0,1]
|
|
# Adjustment strength decreases as total batch count increases
|
|
strength_factor = math.exp(-filter_with_acceptance_rate.total_batches / 1000.0)
|
|
reduction_factor = gap_importance * (0.2 + 0.8 * strength_factor) # Range from 0 to 1, decreases with total batch count
|
|
adjusted_rate = max(0.0, adjusted_rate * (1.0 - reduction_factor))
|
|
|
|
# Add small random perturbation in fixed intervals to enhance convergence
|
|
if 0.01 < adjusted_rate < 0.99 and filter_with_acceptance_rate.total_batches > 100:
|
|
# Random perturbation amplitude decreases as batch count increases
|
|
noise_amplitude = 0.01 * math.exp(-filter_with_acceptance_rate.total_batches / 500.0)
|
|
noise = (torch.rand(1, device=device).item() * 2 - 1) * noise_amplitude
|
|
adjusted_rate = max(0.0, min(1.0, adjusted_rate + noise))
|
|
|
|
# Generate a random number to decide whether to accept the current batch
|
|
random_value = torch.rand(1, device=device).item()
|
|
accept_batch = random_value < adjusted_rate
|
|
|
|
# Set some tokens to PLACEHOLDER_TOKEN_ID to achieve specified acceptance rate
|
|
# Support max_spec_len > 1 cases
|
|
if accept_batch:
|
|
# Accept batch - don't modify token_ids
|
|
filter_with_acceptance_rate.accepted_batches += 1
|
|
filter_with_acceptance_rate.acceptance_history[-1] = 1 # Update the most recent acceptance status
|
|
else:
|
|
# Reject batch - set all speculative tokens (except first column) to PLACEHOLDER_TOKEN_ID
|
|
# This ensures token-level acceptance rate matches batch-level acceptance rate
|
|
output_token_ids[:, 1:] = PLACEHOLDER_TOKEN_ID
|
|
|
|
# Note: acceptance rate calculation is still based on entire batch accept/reject, no modification needed
|
|
# But we can add a comment explaining how actual token-level acceptance rate is calculated
|
|
# Actual token-level acceptance rate = 1 - (number of PLACEHOLDER_TOKEN_ID in output_token_ids / max_spec_len)
|
|
|
|
# Update cumulative error - use exponential moving average for smoother error adjustment
|
|
actual_rate = filter_with_acceptance_rate.accepted_batches / filter_with_acceptance_rate.total_batches
|
|
# Use EMA to smooth error updates - use adaptive EMA coefficient
|
|
alpha = 0.05 * math.exp(-filter_with_acceptance_rate.total_batches / 200.0) + 0.01 # EMA coefficient gradually decreases over time
|
|
filter_with_acceptance_rate.cumulative_error = (alpha * (actual_rate - fixed_acceptance_rate) +
|
|
(1 - alpha) * filter_with_acceptance_rate.cumulative_error)
|
|
|
|
return output_token_ids
|
|
"""
|
|
=============================
|
|
End of MLU Hijack
|
|
=============================
|
|
"""
|
|
|
|
def vllm__v1__sample__rejection_sampler__rejection_sample(
|
|
# [num_tokens]
|
|
draft_token_ids: torch.Tensor,
|
|
# [batch_size]
|
|
num_draft_tokens: list[int],
|
|
max_spec_len: int,
|
|
# [batch_size]
|
|
cu_num_draft_tokens: torch.Tensor,
|
|
# [num_tokens, vocab_size]
|
|
draft_probs: torch.Tensor | None,
|
|
# [num_tokens, vocab_size]
|
|
target_probs: torch.Tensor,
|
|
# [batch_size, 1]
|
|
bonus_token_ids: torch.Tensor,
|
|
sampling_metadata: SamplingMetadata,
|
|
) -> torch.Tensor:
|
|
assert draft_token_ids.ndim == 1
|
|
assert draft_probs is None or draft_probs.ndim == 2
|
|
assert cu_num_draft_tokens.ndim == 1
|
|
assert target_probs.ndim == 2
|
|
|
|
batch_size = len(num_draft_tokens)
|
|
num_tokens = draft_token_ids.shape[0]
|
|
vocab_size = target_probs.shape[-1]
|
|
device = target_probs.device
|
|
assert draft_token_ids.is_contiguous()
|
|
assert draft_probs is None or draft_probs.is_contiguous()
|
|
assert target_probs.is_contiguous()
|
|
assert bonus_token_ids.is_contiguous()
|
|
assert target_probs.shape == (num_tokens, vocab_size)
|
|
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: use tmo rejection_sample for all random sampling requests
|
|
'''
|
|
fixed_acceptance_rate = VLLM_MTP_FIXED_ACCEPTANCE_RATE
|
|
use_fusion_kernel = (sampling_metadata.all_random
|
|
and max_spec_len == 1
|
|
and (num_draft_tokens is not None
|
|
and 0 not in num_draft_tokens))
|
|
if use_fusion_kernel:
|
|
# All data is random, use tmo rejection_sample
|
|
# Generate uniform probabilities for rejection sampling.
|
|
# [num_tokens]
|
|
uniform_rand = vllm__v1__sample__rejection_sampler__generate_uniform_probs(
|
|
num_tokens,
|
|
num_draft_tokens,
|
|
sampling_metadata.generators,
|
|
device,
|
|
)
|
|
# generate random probs for recovered tokens
|
|
uniform_probs = generate_recovered_uniform_probs(
|
|
num_tokens,
|
|
vocab_size,
|
|
num_draft_tokens,
|
|
sampling_metadata,
|
|
device,
|
|
)
|
|
# num_draft_tokens need to be a tensor
|
|
num_draft_tokens_tensor = torch.tensor(num_draft_tokens, dtype=torch.int32, device=device)
|
|
# tmo rejection_sample dtype need to be int32
|
|
bonus_token_ids = bonus_token_ids.to(torch.int32)
|
|
draft_token_ids = draft_token_ids.to(torch.int32)
|
|
# use tmo rejection_sample
|
|
output_token_ids = mlu_ops.rejection_sample(
|
|
draft_token_ids,
|
|
num_draft_tokens_tensor,
|
|
cu_num_draft_tokens,
|
|
draft_probs,
|
|
target_probs,
|
|
bonus_token_ids,
|
|
uniform_rand,
|
|
uniform_probs,
|
|
max_spec_len,
|
|
high_acc=True # for now, only support high_acc
|
|
).view(batch_size, max_spec_len + 1)
|
|
if fixed_acceptance_rate is not None:
|
|
# set all speculative tokens to placeholder token
|
|
output_token_ids[:, 1:] = 0
|
|
output_token_ids = filter_with_acceptance_rate(output_token_ids, fixed_acceptance_rate)
|
|
return output_token_ids
|
|
'''
|
|
=============================
|
|
End of MLU Hijack
|
|
=============================
|
|
'''
|
|
|
|
# Create output buffer.
|
|
output_token_ids = torch.full(
|
|
(batch_size, max_spec_len + 1),
|
|
PLACEHOLDER_TOKEN_ID,
|
|
dtype=torch.int32, # Consistent with SamplerOutput.sampled_token_ids.
|
|
device=device,
|
|
)
|
|
|
|
if sampling_metadata.all_greedy:
|
|
is_greedy = None
|
|
else:
|
|
is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE
|
|
if not sampling_metadata.all_random:
|
|
# Rejection sampling for greedy sampling requests.
|
|
target_argmax = target_probs.argmax(dim=-1)
|
|
vllm__v1__sample__rejection_sampler__rejection_greedy_sample_kernel[(batch_size, )](
|
|
output_token_ids,
|
|
cu_num_draft_tokens,
|
|
draft_token_ids,
|
|
target_argmax,
|
|
bonus_token_ids,
|
|
is_greedy,
|
|
max_spec_len,
|
|
has_acceptance_rate=fixed_acceptance_rate is not None,
|
|
)
|
|
if sampling_metadata.all_greedy:
|
|
output_token_ids = filter_with_acceptance_rate(output_token_ids, fixed_acceptance_rate)
|
|
return output_token_ids
|
|
|
|
# Generate uniform probabilities for rejection sampling.
|
|
# [num_tokens]
|
|
uniform_probs = vllm__v1__sample__rejection_sampler__generate_uniform_probs(
|
|
num_tokens,
|
|
num_draft_tokens,
|
|
sampling_metadata.generators,
|
|
device,
|
|
)
|
|
|
|
# Sample recovered tokens for each position.
|
|
# [num_tokens]
|
|
recovered_token_ids = sample_recovered_tokens(
|
|
max_spec_len,
|
|
num_draft_tokens,
|
|
cu_num_draft_tokens,
|
|
draft_token_ids,
|
|
draft_probs,
|
|
target_probs,
|
|
sampling_metadata,
|
|
device,
|
|
)
|
|
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: Add fixed acceptance rate check
|
|
'''
|
|
# Rejection sampling for random sampling requests.
|
|
vllm__v1__sample__rejection_sampler__rejection_random_sample_kernel[(batch_size, )](
|
|
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,
|
|
NO_DRAFT_PROBS=draft_probs is None,
|
|
has_acceptance_rate=fixed_acceptance_rate is not None,
|
|
)
|
|
output_token_ids = filter_with_acceptance_rate(output_token_ids, fixed_acceptance_rate)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
return output_token_ids
|
|
|
|
# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
|
|
@triton.jit(do_not_specialize=["max_spec_len"])
|
|
def vllm__v1__sample__rejection_sampler__rejection_random_sample_kernel(
|
|
output_token_ids_ptr, # [batch_size, max_spec_len + 1]
|
|
cu_num_draft_tokens_ptr, # [batch_size]
|
|
draft_token_ids_ptr, # [num_tokens]
|
|
draft_probs_ptr, # [num_tokens, vocab_size] or None
|
|
target_probs_ptr, # [num_tokens, vocab_size]
|
|
bonus_token_ids_ptr, # [batch_size]
|
|
recovered_token_ids_ptr, # [num_tokens]
|
|
uniform_probs_ptr, # [num_tokens]
|
|
is_greedy_ptr, # [batch_size]
|
|
max_spec_len,
|
|
vocab_size,
|
|
NO_DRAFT_PROBS: tl.constexpr,
|
|
has_acceptance_rate: tl.constexpr,
|
|
):
|
|
req_idx = tl.program_id(0)
|
|
is_greedy = tl.load(is_greedy_ptr + req_idx)
|
|
if is_greedy:
|
|
# Early exit for greedy sampling requests.
|
|
return
|
|
|
|
start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
|
|
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
|
|
num_draft_tokens = end_idx - start_idx
|
|
|
|
rejected = False
|
|
for pos in range(num_draft_tokens):
|
|
if not rejected:
|
|
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
|
|
if NO_DRAFT_PROBS:
|
|
draft_prob = 1
|
|
else:
|
|
draft_prob = tl.load(
|
|
draft_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
|
|
)
|
|
target_prob = tl.load(
|
|
target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
|
|
)
|
|
uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: add accept rate check, always accept if has_acceptance_rate is True
|
|
'''
|
|
# NOTE(woosuk): While the draft probability should never be 0,
|
|
# we check it to avoid NaNs. If it happens to be 0, we reject.
|
|
if draft_prob > 0 and target_prob / draft_prob >= uniform_prob or has_acceptance_rate:
|
|
# Accept.
|
|
token_id = draft_token_id
|
|
else:
|
|
# Reject. Use recovered token.
|
|
rejected = True
|
|
token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)
|
|
'''
|
|
=============================
|
|
End of MLU Hijack
|
|
=============================
|
|
'''
|
|
tl.store(
|
|
output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id
|
|
)
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: Check whether to accept bonus token through acceptance_rate_ptr
|
|
'''
|
|
# If has acceptance rate, all tokens are accepted
|
|
if has_acceptance_rate:
|
|
rejected = False
|
|
if not rejected:
|
|
# If all tokens are accepted, append the bonus token.
|
|
bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
|
|
tl.store(
|
|
output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens,
|
|
bonus_token_id,
|
|
)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
|
|
|
|
# NOTE(woosuk): Avoid specialization to prevent unnecessary recompilation.
|
|
@triton.jit(do_not_specialize=["max_spec_len"])
|
|
def vllm__v1__sample__rejection_sampler__rejection_greedy_sample_kernel(
|
|
output_token_ids_ptr, # [batch_size, max_spec_len + 1]
|
|
cu_num_draft_tokens_ptr, # [batch_size]
|
|
draft_token_ids_ptr, # [num_tokens]
|
|
target_argmax_ptr, # [num_tokens]
|
|
bonus_token_ids_ptr, # [batch_size]
|
|
is_greedy_ptr, # [batch_size] or None
|
|
max_spec_len,
|
|
has_acceptance_rate: tl.constexpr,
|
|
):
|
|
req_idx = tl.program_id(0)
|
|
# FIXME(woosuk): Because is_greedy_ptr is not None at profiling run,
|
|
# re-compilation may happen during runtime when is_greedy_ptr is None.
|
|
is_greedy = True if is_greedy_ptr is None else tl.load(is_greedy_ptr + req_idx)
|
|
if not is_greedy:
|
|
# Early exit for non-greedy sampling requests.
|
|
return
|
|
|
|
start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
|
|
end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
|
|
num_draft_tokens = end_idx - start_idx
|
|
|
|
rejected = False
|
|
for pos in range(num_draft_tokens):
|
|
if not rejected:
|
|
draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
|
|
target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos)
|
|
tl.store(
|
|
output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
|
|
target_argmax_id,
|
|
)
|
|
if draft_token_id != target_argmax_id:
|
|
# Reject.
|
|
rejected = True
|
|
if has_acceptance_rate:
|
|
rejected = False
|
|
if not rejected:
|
|
# If all tokens are accepted, append the bonus token.
|
|
bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
|
|
tl.store(
|
|
output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens,
|
|
bonus_token_id,
|
|
)
|
|
|
|
|
|
def vllm__v1__sample__rejection_sampler__generate_uniform_probs(
|
|
num_tokens: int,
|
|
num_draft_tokens: list[int],
|
|
generators: dict[int, torch.Generator],
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Generates a batch of uniform random samples, with optional seeding
|
|
if available.
|
|
|
|
This method creates a tensor of shape `(num_tokens, )` filled
|
|
with uniform random values in the range [0, 1). If `generators` is provided,
|
|
the requests with their own seeds will use the provided `torch.Generator`
|
|
for reproducibility. The samples for the other requests will be generated
|
|
without a seed.
|
|
|
|
Args:
|
|
num_tokens: int
|
|
Total number of tokens.
|
|
num_draft_tokens: List[List[int]]
|
|
Number of draft tokens per request.
|
|
generators: Optional[Dict[int, torch.Generator]]
|
|
A dictionary mapping indices in the batch to
|
|
`torch.Generator` objects.
|
|
device: torch.device
|
|
The device on which to allocate the tensor.
|
|
Returns:
|
|
uniform_rand: torch.Tensor
|
|
A tensor of shape `(num_tokens, )` containing uniform
|
|
random values in the range [0, 1).
|
|
"""
|
|
# NOTE(woosuk): We deliberately use float64 instead of float32 here
|
|
# because when using float32, there's a non-negligible chance that
|
|
# uniform_prob is sampled to be exact 0.0 as reported in
|
|
# https://github.com/pytorch/pytorch/issues/16706. Using float64
|
|
# mitigates the issue.
|
|
'''
|
|
=============================
|
|
Modify by vllm_mlu
|
|
=============================
|
|
@brief: Changed torch.float64 to torch.float32
|
|
'''
|
|
uniform_probs = torch.rand(
|
|
(num_tokens,),
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
'''
|
|
==================
|
|
End of MLU Hijack
|
|
==================
|
|
'''
|
|
start_idx = 0
|
|
for req_idx, n in enumerate(num_draft_tokens):
|
|
# Do not generate random numbers for requests with no draft tokens.
|
|
# This can be important for reproducibility.
|
|
if n == 0:
|
|
continue
|
|
end_idx = start_idx + n
|
|
generator = generators.get(req_idx)
|
|
if generator is not None:
|
|
uniform_probs[start_idx:end_idx].uniform_(generator=generator)
|
|
start_idx = end_idx
|
|
return uniform_probs
|
|
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.generate_uniform_probs,
|
|
vllm__v1__sample__rejection_sampler__generate_uniform_probs)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.expand_batch_to_tokens,
|
|
vllm__v1__sample__rejection_sampler__expand_batch_to_tokens)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.expand_kernel,
|
|
vllm__v1__sample__rejection_sampler__expand_kernel)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.sample_recovered_tokens_kernel,
|
|
vllm__v1__sample__rejection_sampler__sample_recovered_tokens_kernel)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.rejection_sample,
|
|
vllm__v1__sample__rejection_sampler__rejection_sample)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.rejection_random_sample_kernel,
|
|
vllm__v1__sample__rejection_sampler__rejection_random_sample_kernel)
|
|
|
|
MluHijackObject.apply_hijack(rejection_sampler,
|
|
rejection_sampler.rejection_greedy_sample_kernel,
|
|
vllm__v1__sample__rejection_sampler__rejection_greedy_sample_kernel)
|