EngineX-Cambricon/enginex-mlu590-vllm
10
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[Model] Support DeepSeek-V4

Browse Source
This commit is contained in:
chenxb002
2026-04-24 09:50:34 +08:00
commit b9925203b8
172 changed files with 44780 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

946
vllm_mlu/v1/sample/rejection_sampler.py Normal file
View File

@@ -0,0 +1,946 @@
# 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)

118
vllm_mlu/v1/sample/sampler.py Normal file
View File

@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-MLU project
# SPDX-License-Identifier: Apache-2.0
import torch
from vllm.config.model import LogprobsMode
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.sample.sampler import Sampler, _SAMPLING_EPS
from vllm_mlu._mlu_utils import *
from vllm_mlu import _mlu_ops as mlu_ops
"""
@brief: use tmo random_sample
"""
def mlu_random_sample(
probs: torch.Tensor,
generators: dict[int, torch.Generator],
) -> torch.Tensor:
is_gumbel_max = True
return mlu_ops.random_sample(probs, is_gumbel_max, generators).view(-1)
class MluSampler(Sampler):
def sample(
self,
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
logprobs_mode_override: LogprobsMode | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Sample logits based on sampling metadata.
The various logits processing functions called in this method
may update the logits tensor in-place.
"""
logprobs_mode = logprobs_mode_override or self.logprobs_mode
assert not (sampling_metadata.all_greedy and sampling_metadata.all_random)
if sampling_metadata.all_random:
greedy_sampled = None
else:
greedy_sampled = self.greedy_sample(logits)
if sampling_metadata.all_greedy:
processed_logprobs = None
if sampling_metadata.max_num_logprobs is not None:
if logprobs_mode == "processed_logits":
processed_logprobs = logits
elif logprobs_mode == "processed_logprobs":
processed_logprobs = self.compute_logprobs(logits)
return greedy_sampled, processed_logprobs
assert sampling_metadata.temperature is not None
"""
=============================
Modify by vllm_mlu
=============================
@brief: use tmo topk_topp_sampler to sample.
"""
use_tmo = (sampling_metadata.top_k is not None) or (sampling_metadata.top_p is not None)
if use_tmo:
batch_size, vocab_size = logits.shape
index_in = torch.arange(vocab_size, dtype=torch.int32, device=logits.device)
(
logits_out,
sorted_logits_out,
index_out,
true_select_len,
) = mlu_ops.apply_topkp_v2(
logits,
index_in,
sampling_metadata.temperature,
None,
sampling_metadata.top_k.to(torch.int32) if sampling_metadata.top_k is not None else None,
sampling_metadata.top_p,
)
processed_logprobs = None
if logprobs_mode == "processed_logits":
processed_logprobs = logits
elif logprobs_mode == "processed_logprobs":
processed_logprobs = self.compute_logprobs(logits)
probs = logits_out.softmax(dim=-1, dtype=torch.float32)
random_sampled = mlu_random_sample(probs, sampling_metadata.generators)
else:
# Apply temperature.
logits = self.apply_temperature(
logits, sampling_metadata.temperature, sampling_metadata.all_random
)
# Apply logits processors that only apply to random sampling
# (argmax invariant)
for processor in sampling_metadata.logitsprocs.argmax_invariant:
logits = processor.apply(logits)
# Apply top_k and/or top_p.
random_sampled, processed_logprobs = self.topk_topp_sampler(
logits,
sampling_metadata.generators,
sampling_metadata.top_k,
sampling_metadata.top_p,
)
"""
=================
End of MLU Hijack
=================
"""
if greedy_sampled is None:
return random_sampled, processed_logprobs
sampled = torch.where(
sampling_metadata.temperature < _SAMPLING_EPS,
greedy_sampled,
random_sampled,
out=greedy_sampled, # Reuse tensor
)
return sampled, processed_logprobs