# SPDX-License-Identifier: Apache-2.0 from dataclasses import replace import torch from vllm.distributed.parallel_state import get_tp_group from vllm.logger import logger from vllm.triton_utils import HAS_TRITON from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.ops.bad_words import apply_bad_words_with_drafts from vllm.v1.sample.rejection_sampler import ( GREEDY_TEMPERATURE, MAX_SPEC_LEN, PLACEHOLDER_TOKEN_ID, RejectionSampler, generate_uniform_probs, ) from vllm.v1.sample.sampler import Sampler from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ops.triton.reject_sample import ( cal_grid_and_block_size, expand_triton, rejection_greedy_sample_with_triton, rejection_random_sample_block_verify_kernel, rejection_random_sample_kernel, sample_recovered_tokens_kernel, ) from vllm_ascend.sample.penalties import apply_all_penalties from vllm_ascend.sample.sampler import apply_top_k_top_p from vllm_ascend.utils import is_310p class AscendRejectionSampler(RejectionSampler): """Ascend-optimized rejection sampler for speculative decoding. This class overrides key methods from the base RejectionSampler to provide Ascend-specific optimizations: - Optimized greedy sampling with reduced communication - Distributed top-k/top-p sampling - Efficient batch expansion operations """ @staticmethod def apply_penalties( logits: torch.Tensor, sampling_metadata: SamplingMetadata, metadata: SpecDecodeMetadata, repeat_indices: torch.Tensor, output_token_ids: list[list[int]], ) -> torch.Tensor: if sampling_metadata.no_penalties: return logits """Use Triton-Ascend penalties on NPU when Triton is available; else vLLM default.""" if not HAS_TRITON: logger.warning_once( "[sample/rejection_sampler] Triton not available, falling back to vLLM default " "penalty implementation in rejection sampler. Rejection sampling performance " "may be degraded on NPU. " ) return Sampler.apply_penalties(logits, sampling_metadata, output_token_ids) assert sampling_metadata.prompt_token_ids is not None prompt_token_ids = sampling_metadata.prompt_token_ids[repeat_indices] presence_penalties = sampling_metadata.presence_penalties[repeat_indices] frequency_penalties = sampling_metadata.frequency_penalties[repeat_indices] repetition_penalties = sampling_metadata.repetition_penalties[repeat_indices] return apply_all_penalties( logits, prompt_token_ids, presence_penalties, frequency_penalties, repetition_penalties, output_token_ids, ) def prepare_sampling(self, top_k): if top_k is not None: self.top_k = top_k else: self.top_k = None def __init__(self, sampler): super().__init__(sampler) # Store Ascend-specific optimizations self._ascend_optimizations_enabled = True self.top_k = None logger.debug( "[sample/rejection_sampler] AscendRejectionSampler initialized. " "ascend_optimizations_enabled=%s, triton_available=%s, " "reduce_sample=%s", self._ascend_optimizations_enabled, HAS_TRITON, get_ascend_config().enable_reduce_sample, ) def apply_logits_processors( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, metadata: SpecDecodeMetadata, ) -> torch.Tensor: has_penalties = not sampling_metadata.no_penalties any_penalties_or_bad_words = sampling_metadata.bad_words_token_ids or has_penalties output_token_ids = sampling_metadata.output_token_ids if any_penalties_or_bad_words: output_token_ids = self._combine_outputs_with_spec_tokens( output_token_ids, sampling_metadata.spec_token_ids, ) # Calculate indices of target logits. if sampling_metadata.allowed_token_ids_mask is not None or has_penalties: num_requests = len(metadata.num_draft_tokens) # TODO: The apply_logits_processors function originally reused the base class from the # upper-level vLLM module. However, the current vLLM implementation introduces synchronous # host-to-device (H2D) copy operations. This function will be removed once PR # https://github.com/vllm-project/vllm/pull/46323 is merged into the upstream vLLM repository. original_indices = torch.arange(num_requests, device=logits.device, dtype=torch.long) repeat_indices = expand_batch_to_tokens( original_indices, metadata.cu_num_draft_tokens, logits.shape[0], ) logits = self.apply_penalties(logits, sampling_metadata, metadata, repeat_indices, output_token_ids) # Apply allowed token ids. if sampling_metadata.allowed_token_ids_mask is not None: token_mask = sampling_metadata.allowed_token_ids_mask[repeat_indices] logits.masked_fill_(token_mask, float("-inf")) # Apply bad words exclusion. if bad_words_token_ids := sampling_metadata.bad_words_token_ids: apply_bad_words_with_drafts(logits, bad_words_token_ids, output_token_ids, metadata.num_draft_tokens) for processor in sampling_metadata.logitsprocs.non_argmax_invariant: if isinstance(processor, MinTokensLogitsProcessor): logits = processor.apply_with_spec_decode(logits, metadata.num_draft_tokens) return logits def forward( self, metadata: SpecDecodeMetadata, # [num_tokens, vocab_size] draft_probs: torch.Tensor | None, # [num_tokens + batch_size, vocab_size] logits: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> SamplerOutput: """ Args: metadata: Metadata for spec decoding. draft_probs (Optional[torch.Tensor]): Probability distribution for the draft tokens. Shape is [num_tokens, vocab_size]. Can be None if probabilities are not provided, which is the case for ngram spec decode. logits (torch.Tensor): Target model's logits probability distribution. Shape is [num_tokens + batch_size, vocab_size]. Here, probabilities from different requests are flattened into a single tensor because this is the shape of the output logits. NOTE: `logits` can be updated in place to save memory. sampling_metadata (vllm.v1.sample.metadata.SamplingMetadata): Additional metadata needed for sampling, such as temperature, top-k/top-p parameters, or other relevant information. Returns: SamplerOutput: Contains the final output token IDs and their logprobs if requested. """ assert metadata.max_spec_len <= MAX_SPEC_LEN bonus_logits_indices = metadata.bonus_logits_indices target_logits_indices = metadata.target_logits_indices # When indexing with a tensor (bonus_logits_indices), PyTorch # creates a new tensor with separate storage from the original # logits tensor. This means any in-place operations on bonus_logits # won't affect the original logits tensor. assert logits is not None bonus_logits = logits[bonus_logits_indices] bonus_sampler_output = self.sampler( logits=bonus_logits, sampling_metadata=replace( sampling_metadata, max_num_logprobs=-1, ), predict_bonus_token=True, # Override the logprobs mode to return logits because they are # needed later to compute the accepted token logprobs. logprobs_mode_override="processed_logits" if self.is_processed_logprobs_mode else "raw_logits", ) bonus_token_ids = bonus_sampler_output.sampled_token_ids # Just like `bonus_logits`, `target_logits` is a new tensor with # separate storage from the original `logits` tensor. Therefore, # it is safe to update `target_logits` in place. raw_target_logits = logits[target_logits_indices] # Use float32 for the target_logits. raw_target_logits = raw_target_logits.to(torch.float32) target_logits = raw_target_logits if not self.is_processed_logprobs_mode: # Clone raw_target_logits before applying processors to preserve # the original raw logits for logprobs computation, since # apply_logits_processors modifies the tensor in-place. target_logits = target_logits.clone() # Clean NaN/inf without introducing CPU sync. # torch.nan_to_num is a pure element-wise op that replaces # NaN/±inf in a single pass, avoiding the .any()/.item() sync # triggered by `if tensor.any()` branches. No-op when clean. if not is_310p(): info = torch.finfo(target_logits.dtype) target_logits = torch.nan_to_num( target_logits, nan=0.0, posinf=info.max, neginf=info.min, ) target_logits = self.apply_logits_processors(target_logits, sampling_metadata, metadata) # [num_tokens, vocab_size] # NOTE(woosuk): `target_logits` can be updated in place inside the # `apply_sampling_constraints` function. target_logits = apply_sampling_constraints( target_logits, metadata.cu_num_draft_tokens, sampling_metadata, self.top_k ) output_token_ids = rejection_sample( metadata.draft_token_ids, metadata.num_draft_tokens, metadata.max_spec_len, metadata.cu_num_draft_tokens, draft_probs, target_logits, bonus_token_ids, sampling_metadata, ori_target_logits=raw_target_logits, ) logprobs_tensors = None if sampling_metadata.max_num_logprobs is not None: logprobs_tensors = self._get_logprobs_tensors( sampling_metadata.max_num_logprobs, metadata, logits, target_logits if self.is_processed_logprobs_mode else raw_target_logits, bonus_sampler_output.logprobs_tensors.logprobs, output_token_ids, ) return SamplerOutput( sampled_token_ids=output_token_ids, logprobs_tensors=logprobs_tensors, ) def greedy_sample(logits: torch.Tensor) -> torch.Tensor: tp_group = get_tp_group() B, V_local = logits.shape rank = tp_group.rank_in_group local_max_logits, local_max_indices = logits.max(dim=-1) local_global_idx = local_max_indices + rank * V_local # [B] # [B, world_size] gathered_logits = tp_group.all_gather(local_max_logits.unsqueeze(-1), dim=-1) gathered_global_idx = tp_group.all_gather(local_global_idx.unsqueeze(-1), dim=-1) # [B, world_size] global_max_rank = gathered_logits.argmax(dim=-1) # [B] target_argmax = gathered_global_idx.gather(dim=-1, index=global_max_rank.unsqueeze(-1)).squeeze(-1) # [B] return target_argmax def apply_sampling_constraints( logits: torch.Tensor, # [num_tokens, vocab_size//tp_size] cu_num_draft_tokens: torch.Tensor, # [batch_size] sampling_metadata: SamplingMetadata, top_k, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Process logits based on sampling metadata for distributed scenario. This function applies temperature scaling to the logits, then top-k, allgather, and top-p. For greedy decoding, it returns the original logits. Args: logits: Input logits tensor to be processed (local vocab partition). cu_num_draft_tokens: Cumulative number of draft tokens. sampling_metadata: Metadata containing sampling parameters such as temperature and whether greedy sampling is used. Returns: tuple[torch.Tensor, torch.Tensor | None]: - Processed logits of shape [num_tokens, top_k*tp_size] or [num_tokens, vocab_size//tp_size] for greedy - Indices tensor of shape [num_tokens, top_k*tp_size] or None for greedy """ assert logits.ndim == 2 assert cu_num_draft_tokens.ndim == 1 if sampling_metadata.all_greedy: # return logits return logits, None num_tokens = logits.shape[0] temperature = expand_batch_to_tokens( sampling_metadata.temperature, cu_num_draft_tokens, num_tokens, replace_from=GREEDY_TEMPERATURE, replace_to=1, ) # NOTE(woosuk): Update `logits` in place to avoid allocating a new tensor. logits.div_(temperature.unsqueeze(-1)) # Get expanded top_k and top_p tensors. k = None if sampling_metadata.top_k is not None: k = expand_batch_to_tokens( sampling_metadata.top_k, cu_num_draft_tokens, num_tokens, ) p = None if sampling_metadata.top_p is not None: p = expand_batch_to_tokens( sampling_metadata.top_p, cu_num_draft_tokens, num_tokens, ) # New flow: top_k -> allgather -> top_p # Returns processed logits and indices if get_ascend_config().enable_reduce_sample: logger.debug_once( "[sample/rejection_sampler] Using reduce-sample path for " "apply_sampling_constraints. top-k/top-p with TP all-gather.", ) return apply_top_k_top_p(logits, k, p, top_k) else: return apply_top_k_top_p(logits, k, p) def 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//tp_size] or tuple of (logits, indices) # For greedy: Tensor [num_tokens, vocab_size//tp_size] # For random: tuple of (logits [num_tokens, top_k*tp_size], indices [num_tokens, top_k*tp_size]) target_logits_or_tuple: torch.Tensor | tuple[torch.Tensor, torch.Tensor | None], # [batch_size, 1] bonus_token_ids: torch.Tensor, sampling_metadata: SamplingMetadata, synthetic_mode: bool = False, synthetic_conditional_rates: torch.Tensor | None = None, ori_target_logits: torch.Tensor | None = None, ) -> torch.Tensor: """ Rejection sampling for speculative decoding in distributed setting. Args: draft_token_ids: Draft token IDs [num_tokens] num_draft_tokens: Number of draft tokens per request max_spec_len: Maximum speculative length cu_num_draft_tokens: Cumulative draft tokens [batch_size] draft_probs: Draft probabilities [num_tokens, vocab_size] or None for ngram target_logits_or_tuple: Target logits or tuple of (logits, indices) - For greedy: Tensor [num_tokens, vocab_size//tp_size] - For random: tuple of (selected_logits, indices) where - selected_logits: [num_tokens, top_k*tp_size] - indices: [num_tokens, top_k*tp_size] global vocabulary indices or None bonus_token_ids: Bonus token IDs [batch_size, 1] sampling_metadata: Sampling metadata Returns: output_token_ids: [batch_size, max_spec_len + 1] """ # Unpack target_logits_or_tuple if isinstance(target_logits_or_tuple, tuple): target_logits, target_indices = target_logits_or_tuple else: target_logits = target_logits_or_tuple target_indices = None 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_logits.ndim == 2 batch_size = len(num_draft_tokens) num_tokens = draft_token_ids.shape[0] device = target_logits.device assert draft_token_ids.is_contiguous() assert draft_probs is None or draft_probs.is_contiguous() assert target_logits.is_contiguous() assert bonus_token_ids.is_contiguous() assert target_logits.shape[0] == num_tokens # Block verify requires enable_block_verify config and max_spec_len >= 3. using_block_verify = max_spec_len >= 3 and bool(get_ascend_config().rejection_sampler_config.enable_block_verify) using_entropy_verify = bool(get_ascend_config().rejection_sampler_config.enable_entropy_verify) posterior_threshold = float(get_ascend_config().rejection_sampler_config.posterior_threshold) posterior_alpha = float(get_ascend_config().rejection_sampler_config.posterior_alpha) logger.debug_once( "[sample/rejection_sampler] Rejection sampling path: " "block_verify=%s, entropy_verify=%s, all_greedy=%s, all_random=%s, " "reduce_sample=%s, triton=%s", using_block_verify, using_entropy_verify, sampling_metadata.all_greedy, sampling_metadata.all_random, get_ascend_config().enable_reduce_sample, HAS_TRITON, ) if using_entropy_verify and ori_target_logits is not None: ori_target_probs = ori_target_logits.softmax(dim=-1, dtype=torch.float32) else: ori_target_probs = None # Create output buffer. output_token_ids = torch.empty( (batch_size, max_spec_len + 1), dtype=torch.int32, # Consistent with SamplerOutput.sampled_token_ids. device=device, ) output_token_ids.fill_(PLACEHOLDER_TOKEN_ID) if sampling_metadata.all_greedy: is_greedy = None else: is_greedy = sampling_metadata.temperature == GREEDY_TEMPERATURE if HAS_TRITON: grid, block_size = cal_grid_and_block_size(batch_size) if using_block_verify or using_entropy_verify: logger.info_once( "RejectionSampler config: block_verify=%s, entropy_verify=%s, " "posterior_threshold=%s, posterior_alpha=%s, reduce_sample=%s, " "has_triton=%s, all_greedy=%s, all_random=%s", using_block_verify, using_entropy_verify, posterior_threshold, posterior_alpha, target_indices is not None, HAS_TRITON, sampling_metadata.all_greedy, sampling_metadata.all_random, ) # For greedy sampling, we need to do allgather first to get global argmax if not sampling_metadata.all_random: if get_ascend_config().enable_reduce_sample: target_argmax = greedy_sample(target_logits) else: target_argmax = target_logits.argmax(dim=-1).view(-1) if HAS_TRITON: rejection_greedy_sample_with_triton( output_token_ids, num_draft_tokens, cu_num_draft_tokens, draft_token_ids, target_argmax, bonus_token_ids, is_greedy, max_spec_len, grid, block_size, ) else: if min(num_draft_tokens) == 1 and max(num_draft_tokens) == 1 and sampling_metadata.all_greedy: rejection_greedy_sample_spec_len_1_pytorch( output_token_ids, draft_token_ids, target_argmax, bonus_token_ids, ) else: rejection_greedy_sample_pytorch( output_token_ids, cu_num_draft_tokens, draft_token_ids, target_argmax, bonus_token_ids, num_draft_tokens, max_spec_len, is_greedy, ) if sampling_metadata.all_greedy: return output_token_ids # For random sampling with selected logits # target_logits is [num_tokens, top_k*tp_size] with indices [num_tokens, top_k*tp_size] if target_indices is not None: # Enable reduce_sampling: logits are [num_tokens, top_k*tp_size] # We need to handle rejection sampling with selected vocab selected_vocab_size = target_logits.shape[-1] global_vocab_size = draft_probs.shape[-1] if draft_probs is not None else selected_vocab_size # Compute probability distribution from target logits target_probs = target_logits.softmax(dim=-1, dtype=torch.float32) assert target_probs.is_contiguous() # Generate uniform probabilities for rejection sampling uniform_probs = generate_uniform_probs( num_tokens, num_draft_tokens, sampling_metadata.generators, device, ) # Sample recovered tokens for each position 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, target_indices=target_indices, global_vocab_size=global_vocab_size, enable_reduce_sampling=True, ) if not using_block_verify: # Rejection sampling for random sampling requests with selected logits if HAS_TRITON: rejection_random_sample_kernel[(grid,)]( output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, target_indices, bonus_token_ids, recovered_token_ids, uniform_probs.to(torch.float32), is_greedy, max_spec_len, selected_vocab_size, global_vocab_size, batch_size, ori_target_probs, NO_ORI_TARGET_PROBS=ori_target_probs is None, NO_DRAFT_PROBS=draft_probs is None, ENABLE_REDUCE_SAMPLING=True, ENTROPY_VERIFY=using_entropy_verify, BLOCK_SIZE=block_size, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, SUB_BLOCK=4 * 1024, EPSILON=1e-10, ) else: rejection_random_sample_pytorch( 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, selected_vocab_size, IS_NGRAM=draft_probs is None, target_indices=target_indices, enable_reduce_sampling=True, ENTROPY_VERIFY=using_entropy_verify, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, EPSILON=1e-10, ori_target_probs=ori_target_probs, ) else: # MagicMTP: Improving acceptance rate with Block Verify. # Entropy_verify: Improving acceptance rate with entropy Verify. if HAS_TRITON: rejection_random_sample_block_verify_kernel[(grid,)]( output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, target_indices, bonus_token_ids, recovered_token_ids, uniform_probs.to(torch.float32), is_greedy, max_spec_len, selected_vocab_size, global_vocab_size, batch_size, ori_target_probs, NO_ORI_TARGET_PROBS=ori_target_probs is None, NO_DRAFT_PROBS=draft_probs is None, ENABLE_REDUCE_SAMPLING=True, ENTROPY_VERIFY=using_entropy_verify, BLOCK_SIZE=block_size, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, SUB_BLOCK=4 * 1024, EPSILON=1e-10, ) else: rejection_random_sample_block_verify_pytorch( 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, selected_vocab_size, IS_NGRAM=draft_probs is None, target_indices=target_indices, enable_reduce_sampling=True, ENTROPY_VERIFY=using_entropy_verify, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, EPSILON=1e-10, ori_target_probs=ori_target_probs, ) else: # Fallback to original mode # This path should not be used in the new distributed flow logger.warning_once( "[sample/rejection_sampler] Using fallback (non-reduce-sample) path in " "rejection_sample. This path should not be used in the new distributed flow. " "enable_reduce_sample=%s, has_target_indices=%s", get_ascend_config().enable_reduce_sample, target_indices is not None, ) vocab_size = target_logits.shape[-1] global_vocab_size = draft_probs.shape[-1] if draft_probs is not None else vocab_size # Compute probability distribution from target logits target_probs = target_logits.softmax(dim=-1, dtype=torch.float32) assert target_probs.is_contiguous() # Generate uniform probabilities for rejection sampling uniform_probs = generate_uniform_probs( num_tokens, num_draft_tokens, sampling_metadata.generators, device, ) # Sample recovered tokens for each position 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, target_indices=None, global_vocab_size=vocab_size, enable_reduce_sampling=False, ) if not using_block_verify: if HAS_TRITON: rejection_random_sample_kernel[(grid,)]( output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, None, # target_indices bonus_token_ids, recovered_token_ids, uniform_probs.to(torch.float32), is_greedy, max_spec_len, vocab_size, global_vocab_size, # global_vocab_size batch_size, ori_target_probs, NO_ORI_TARGET_PROBS=ori_target_probs is None, NO_DRAFT_PROBS=draft_probs is None, ENABLE_REDUCE_SAMPLING=False, ENTROPY_VERIFY=using_entropy_verify, BLOCK_SIZE=block_size, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, SUB_BLOCK=4 * 1024, EPSILON=1e-10, ) else: rejection_random_sample_pytorch( 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, IS_NGRAM=draft_probs is None, target_indices=None, enable_reduce_sampling=False, ENTROPY_VERIFY=using_entropy_verify, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, EPSILON=1e-10, ori_target_probs=ori_target_probs, ) else: if HAS_TRITON: rejection_random_sample_block_verify_kernel[(grid,)]( output_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, None, # target_indices bonus_token_ids, recovered_token_ids, uniform_probs.to(torch.float32), is_greedy, max_spec_len, vocab_size, global_vocab_size, # global_vocab_size batch_size, ori_target_probs, NO_ORI_TARGET_PROBS=ori_target_probs is None, NO_DRAFT_PROBS=draft_probs is None, ENABLE_REDUCE_SAMPLING=False, ENTROPY_VERIFY=using_entropy_verify, BLOCK_SIZE=block_size, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, SUB_BLOCK=4 * 1024, EPSILON=1e-10, ) else: rejection_random_sample_block_verify_pytorch( 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, IS_NGRAM=draft_probs is None, target_indices=None, enable_reduce_sampling=False, ENTROPY_VERIFY=using_entropy_verify, POSTERIOR_THRESHOLD=posterior_threshold, POSTERIOR_ALPHA=posterior_alpha, EPSILON=1e-10, ori_target_probs=ori_target_probs, ) return output_token_ids def 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 expanded_x = x.new_empty(num_tokens) if HAS_TRITON: expand_triton(batch_size, expanded_x, x, cu_num_tokens, replace_from, replace_to, max_num_tokens=MAX_SPEC_LEN) else: expand_pytorch( expanded_x, x, cu_num_tokens, replace_from, replace_to, MAX_NUM_TOKENS=MAX_SPEC_LEN, # To avoid recompilation. ) return expanded_x def sample_recovered_tokens( max_spec_len: int, num_draft_tokens: list[int], cu_num_draft_tokens: torch.Tensor, draft_token_ids: torch.Tensor, draft_probs: torch.Tensor | None, target_probs: torch.Tensor, sampling_metadata: SamplingMetadata, device: torch.device, use_block_verify: bool = False, target_indices: torch.Tensor | None = None, global_vocab_size: int | None = None, enable_reduce_sampling: bool = False, ) -> torch.Tensor: batch_size = len(num_draft_tokens) vocab_size = target_probs.shape[-1] q = torch.empty( (batch_size, vocab_size), dtype=torch.float32, device=device, ) q.exponential_() num_draft_tensor = torch.tensor(num_draft_tokens, pin_memory=True).to(device, non_blocking=True) has_draft_mask = num_draft_tensor > 0 for i, generator in sampling_metadata.generators.items(): temp_q = torch.empty_like(q[i]) temp_q.exponential_(generator=generator) q[i] = torch.where(has_draft_mask[i], temp_q, q[i]) recovered_token_ids = torch.empty_like(draft_token_ids) if HAS_TRITON: sample_recovered_tokens_kernel[(batch_size, max_spec_len)]( recovered_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, target_indices, # None for normal mode q, vocab_size, global_vocab_size if global_vocab_size is not None else vocab_size, NO_DRAFT_PROBS=draft_probs is None, ENABLE_REDUCE_SAMPLING=enable_reduce_sampling, VOCAB_BLOCK_SIZE=512, SUB_BLOCK=4 * 1024, # TODO: enable multibuffer when accuracy problem is solved. multibuffer=False, ) elif use_block_verify: sample_recovered_tokens_blockwise_pytorch( recovered_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, q, vocab_size, IS_NGRAM=draft_probs is None, target_indices=target_indices, enable_reduce_sampling=enable_reduce_sampling, ) else: sample_recovered_tokens_pytorch( recovered_token_ids, cu_num_draft_tokens, draft_token_ids, draft_probs, target_probs, q, vocab_size, IS_NGRAM=draft_probs is None, target_indices=target_indices, enable_reduce_sampling=enable_reduce_sampling, ) return recovered_token_ids def rejection_greedy_sample_spec_len_1_pytorch( output_token_ids, # [batch_size, 2] draft_token_ids, # [num_tokens] target_argmax, # [num_tokens] bonus_token_ids, # [batch_size] ): batch_size = output_token_ids.size(0) num_tokens = draft_token_ids.size(0) assert batch_size == num_tokens accept_req_mask = draft_token_ids == target_argmax output_token_ids[:, 0] = target_argmax bonus_token_ids = bonus_token_ids.squeeze(1) output_token_ids[:, 1] = torch.where(accept_req_mask, bonus_token_ids, output_token_ids[:, 1]) def rejection_greedy_sample_pytorch( output_token_ids, # [batch_size, max_spec_len + 1] cu_num_draft_tokens, # [batch_size] draft_token_ids, # [num_tokens] target_argmax, # [num_tokens] bonus_token_ids, # [batch_size] draft_tokens_per_req, # [batch_size], list max_spec_len, is_greedy=None, # [batch_size] or None ): batch_size = output_token_ids.size(0) num_tokens = draft_token_ids.size(0) device = output_token_ids.device draft_tokens_per_req = torch.tensor(draft_tokens_per_req).to(device, non_blocking=True) if is_greedy is None: is_greedy = torch.ones(batch_size, dtype=torch.bool, device=device) start_indices = cu_num_draft_tokens - draft_tokens_per_req req_ids = torch.arange(batch_size, device=device) token_req_ids = torch.repeat_interleave(req_ids, draft_tokens_per_req) token_positions = torch.arange(num_tokens, device=device) - start_indices[token_req_ids] # Find the first mismatch position of each request. mismatch_global = draft_token_ids != target_argmax if max_spec_len == 0: first_mismatch_pos_per_req = torch.zeros(batch_size, dtype=torch.long, device=device) else: # [bs, max_spec_len] pos_matrix = torch.full((batch_size, max_spec_len), -1, dtype=torch.long, device=device) pos_matrix[token_req_ids, token_positions] = token_positions mismatch_matrix = torch.full((batch_size, max_spec_len), False, dtype=torch.bool, device=device) mismatch_matrix[token_req_ids, token_positions] = mismatch_global mismatch_positions = torch.where(mismatch_matrix, pos_matrix, max_spec_len * 2) first_mismatch_pos_per_req, _ = torch.min(mismatch_positions, dim=1) no_mismatch_mask = first_mismatch_pos_per_req == max_spec_len * 2 first_mismatch_pos_per_req[no_mismatch_mask] = draft_tokens_per_req[no_mismatch_mask] # Copy matched target tokens into output. copy_len = torch.minimum(first_mismatch_pos_per_req + 1, draft_tokens_per_req) copy_indices = torch.arange(max_spec_len + 1, device=device).expand(batch_size, -1) copy_mask = copy_indices < copy_len.unsqueeze(1) greedy_mask = is_greedy.unsqueeze(1) final_copy_mask = copy_mask & greedy_mask global_idx = start_indices.unsqueeze(1) + copy_indices output_token_ids[final_copy_mask] = target_argmax[global_idx[final_copy_mask]].to(output_token_ids.dtype) # Fill bonus token. needs_bonus = is_greedy & (first_mismatch_pos_per_req >= draft_tokens_per_req) if torch.any(needs_bonus): bonus_rows = torch.where(needs_bonus)[0] bonus_cols = draft_tokens_per_req[bonus_rows] bonus_token_ids = bonus_token_ids.squeeze(1) output_token_ids[bonus_rows, bonus_cols] = bonus_token_ids[bonus_rows] def rejection_random_sample_pytorch( output_token_ids, # [batch_size, max_spec_len + 1] cu_num_draft_tokens, # [batch_size] draft_token_ids, # [num_tokens] draft_probs, # [num_tokens, vocab_size] or None target_probs, # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] bonus_token_ids, # [batch_size] recovered_token_ids, # [num_tokens] uniform_probs, # [num_tokens] is_greedy, # [batch_size] max_spec_len, vocab_size, IS_NGRAM=False, target_indices=None, # [num_tokens, selected_vocab_size] global vocab indices enable_reduce_sampling=False, ENTROPY_VERIFY=False, POSTERIOR_THRESHOLD=0.95, POSTERIOR_ALPHA=0.4, EPSILON=1e-10, ori_target_probs=None, ): """ This function implements the Speculative Decoding rejection sampling step. Instead of looping through each request and each token (which causes high overhead), it uses a fully vectorized approach: 1. **Index Mapping**: Converts the flattened 1D token arrays into a 2D [batch_size, max_draft_len] grid using 'cu_num_draft_tokens' to handle variable-length sequences in the batch. 2. **Parallel Validation**: Calculates the acceptance condition (target_prob / draft_prob >= uniform_sample) for ALL draft tokens simultaneously across the entire batch. 3. **Short-circuit Simulation**: In the loop version, once a token is rejected, subsequent tokens are ignored. Here, we simulate this by finding the 'first_reject_pos' using argmax on the rejection mask and creating a 'should_skip' mask for all indices after the first failure. 4. **Token Selection**: Uses 'torch.where' to select: - Draft tokens (if accepted) - Recovered tokens (at the point of first rejection) - Bonus tokens (if all tokens in a sequence were accepted) 5. **Masking**: Ensures operations only apply to non-greedy requests and within valid sequence lengths. """ batch_size = output_token_ids.shape[0] device = output_token_ids.device zero_cpu = torch.tensor([0], pin_memory=True) zero_device = zero_cpu.to(device, non_blocking=True) cu_start = torch.cat([zero_device, cu_num_draft_tokens[:-1]]) cu_end = cu_num_draft_tokens num_draft_per_batch = cu_end - cu_start max_draft_len = max_spec_len pos_indices_cpu = torch.arange(max_draft_len, pin_memory=True) pos_indices = pos_indices_cpu.to(device, non_blocking=True)[None, :] valid_mask = pos_indices < num_draft_per_batch[:, None] global_token_indices = cu_start[:, None] + pos_indices global_token_indices = global_token_indices.clamp(0, draft_token_ids.shape[0] - 1) draft_tokens = draft_token_ids[global_token_indices] # [batch_size, max_draft_len] placeholder_mask = draft_tokens == PLACEHOLDER_TOKEN_ID safe_draft_tokens = draft_tokens.masked_fill(placeholder_mask, 0) if IS_NGRAM: ones_cpu = torch.ones(1, pin_memory=True, dtype=torch.float32) draft_token_probs = ones_cpu.to(device, non_blocking=True).expand_as(draft_tokens) else: flat_indices = global_token_indices.flatten() flat_draft_tokens = safe_draft_tokens.flatten() flat_draft_probs = draft_probs[flat_indices, flat_draft_tokens] draft_token_probs = flat_draft_probs.view(batch_size, max_draft_len) # Get target token probs if enable_reduce_sampling: # When enable_reduce_sampling, need to search for draft token in candidates flat_global_indices = global_token_indices.flatten() flat_draft_tokens = draft_tokens.flatten() flat_target_indices = target_indices[flat_global_indices] flat_target_probs = target_probs[flat_global_indices] # Check if draft token is in candidates draft_expanded = flat_draft_tokens.unsqueeze(1) is_in_candidates = flat_target_indices == draft_expanded # Get the probability of draft token from target (if present) target_token_probs_flat = torch.where( is_in_candidates, flat_target_probs, torch.tensor(0.0, device=device) ).sum(dim=1) target_token_probs = target_token_probs_flat.view(batch_size, max_draft_len) else: flat_indices = global_token_indices.flatten() flat_draft_tokens = safe_draft_tokens.flatten() flat_target_probs = target_probs[flat_indices, flat_draft_tokens] target_token_probs = flat_target_probs.view(batch_size, max_draft_len) uniform_token_probs = uniform_probs[global_token_indices] recovered_tokens = recovered_token_ids[global_token_indices] zero_threshold_cpu = torch.tensor([0.0], pin_memory=True, dtype=torch.float32) zero_threshold = zero_threshold_cpu.to(device, non_blocking=True) if ENTROPY_VERIFY: entropy_probs = ori_target_probs if ori_target_probs is not None else target_probs all_target_dist = entropy_probs[global_token_indices] entropy = -(all_target_dist * torch.log(all_target_dist + EPSILON)).sum(dim=-1) exp_neg_entropy = torch.exp(-entropy * POSTERIOR_ALPHA) posterior_threshold_device = torch.tensor(POSTERIOR_THRESHOLD, device=device, dtype=torch.float32) threshold = torch.minimum(exp_neg_entropy, posterior_threshold_device) modified_uniform_token_probs = threshold * uniform_token_probs acceptance_condition = (draft_token_probs > zero_threshold) & ( target_token_probs / draft_token_probs >= modified_uniform_token_probs ) else: acceptance_condition = (draft_token_probs > zero_threshold) & ( target_token_probs / draft_token_probs >= uniform_token_probs ) acceptance_condition = acceptance_condition & (~placeholder_mask) first_rejection = (~acceptance_condition) & valid_mask default_pos_cpu = torch.full([batch_size, 1], max_draft_len, pin_memory=True) default_pos = default_pos_cpu.to(device, non_blocking=True) first_reject_pos = torch.where( first_rejection.any(dim=1, keepdim=True), first_rejection.float().argmax(dim=1, keepdim=True), default_pos ) pos_mask = pos_indices >= first_reject_pos should_skip = pos_mask & valid_mask final_acceptance = acceptance_condition & (~should_skip) non_greedy_mask = ~is_greedy update_mask = non_greedy_mask[:, None] & valid_mask & (~should_skip) first_reject_mask = (pos_indices == first_reject_pos) & valid_mask & non_greedy_mask[:, None] final_update_mask = update_mask | first_reject_mask final_tokens = torch.where( first_reject_mask, recovered_tokens, torch.where(final_acceptance, draft_tokens, output_token_ids[:, :max_draft_len]), ) output_token_ids[:, :max_draft_len] = torch.where( final_update_mask, final_tokens, output_token_ids[:, :max_draft_len] ) no_rejection = first_reject_pos.squeeze(1) >= num_draft_per_batch should_add_bonus = non_greedy_mask & no_rejection bonus_positions = num_draft_per_batch # [batch_size] seq_len = output_token_ids.shape[1] all_positions_cpu = torch.arange(seq_len, pin_memory=True) all_positions = all_positions_cpu.to(device, non_blocking=True)[None, :] # [1, seq_len] batch_bonus_positions = bonus_positions[:, None] # [batch_size, 1] max_spec_len_cpu = torch.tensor([max_spec_len], pin_memory=True) max_spec_len_device = max_spec_len_cpu.to(device, non_blocking=True) valid_bonus_pos = bonus_positions < (max_spec_len_device + 1) final_bonus_mask = should_add_bonus & valid_bonus_pos bonus_pos_match = all_positions == batch_bonus_positions bonus_pos_mask = bonus_pos_match & final_bonus_mask[:, None] bonus_values_expanded = bonus_token_ids.view(-1, 1).expand(-1, seq_len) output_token_ids[:] = torch.where(bonus_pos_mask, bonus_values_expanded, output_token_ids) def expand_pytorch( output_ptr, # [num_tokens] input_ptr, # [batch_size] cu_num_tokens_ptr, # [batch_size] replace_from, replace_to, MAX_NUM_TOKENS, ): """ This function broadcasts batch-level values (input_ptr) to token-level positions (output_ptr) based on cumulative token offsets. It acts like a "scatter" or "repeat_interleave" operation but with custom logic: 1. **Range Broadcasting**: It creates a boolean matrix 'in_range' of size [num_tokens, batch_size] that identifies which batch index each token belongs to by checking if the token index falls between cu_start and cu_end. 2. **Conditional Replacement**: Before expansion, it replaces specific values (e.g., padding or special markers) in the input to prepare the data. 3. **Matrix-based Mapping**: It uses 'torch.einsum' to perform a weighted sum that effectively "picks" the correct batch value for every token position simultaneously, avoiding a Python loop over the batch. """ device = cu_num_tokens_ptr.device batch_size = input_ptr.shape[0] num_tokens = output_ptr.shape[0] if batch_size == 0 or num_tokens == 0: return cu_start = torch.cat([torch.tensor([0], pin_memory=True).to(device, non_blocking=True), cu_num_tokens_ptr[:-1]]) cu_end = cu_num_tokens_ptr token_indices = torch.arange(num_tokens, device=device)[:, None] # [num_tokens, 1] cu_start_exp = cu_start[None, :] # [1, batch_size] cu_end_exp = cu_end[None, :] # [1, batch_size] in_range = (token_indices >= cu_start_exp) & (token_indices < cu_end_exp) replaced_input = torch.where(input_ptr == replace_from, replace_to, input_ptr).float() token_values = torch.einsum("tb,b->t", in_range.float(), replaced_input) needs_update = in_range.any(dim=1) output_ptr[:] = torch.where(needs_update, token_values, output_ptr) def sample_recovered_tokens_pytorch( output_token_ids, # [num_tokens] cu_num_draft_tokens, # [batch_size] draft_token_ids, # [num_tokens] draft_probs, # [num_tokens, vocab_size] or None target_probs, # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] q, # [batch_size, vocab_size] or [batch_size, selected_vocab_size] vocab_size, IS_NGRAM=False, target_indices=None, # [num_tokens, selected_vocab_size] global vocab indices enable_reduce_sampling=False, ): """ When a draft token is rejected, we must sample a "recovered" token from a modified distribution. This function calculates that distribution across the entire flattened batch. 1. **Token-to-Batch Mapping**: Using the cumulative draft token counts, it determines which request in the batch each token belongs to. This is necessary because 'q' (normalization factor) is stored per-request. 2. **Probability Adjustment**: - If N-GRAM: It zeroes out the draft token's probability in the target. - If Probabilistic: It calculates max(0, target_probs - draft_probs) as per the standard speculative decoding algorithm. 3. **Normalization & Sampling**: It divides the adjusted probabilities by the normalization distribution 'q'. To remain vectorized, it broadcasts 'q' from [batch_size, vocab] to [num_tokens, vocab]. 4. **Argmax Selection**: It selects the best recovery token for every position in one pass using torch.argmax. """ device = output_token_ids.device num_tokens = output_token_ids.shape[0] if num_tokens == 0: return cu_start = torch.cat( [ torch.tensor([0], pin_memory=True).to(device, non_blocking=True), cu_num_draft_tokens[:-1], ] ) cu_end = cu_num_draft_tokens token_indices = torch.arange(num_tokens, device=device) # [num_tokens] token_indices_expanded = token_indices[:, None] # [num_tokens, 1] cu_start_expanded = cu_start[None, :] # [1, batch_size] cu_end_expanded = cu_end[None, :] # [1, batch_size] in_range_mask = (token_indices_expanded >= cu_start_expanded) & (token_indices_expanded < cu_end_expanded) token_to_batch = torch.argmax(in_range_mask.int(), dim=1) has_match = in_range_mask.any(dim=1) token_to_batch = torch.where(has_match, token_to_batch, 0) if enable_reduce_sampling: # enable reduce_sampling: target_probs is [num_tokens, selected_vocab_size] # target_indices maps compressed indices to global vocab indices if IS_NGRAM: # Zero out the draft token in target_probs prob = target_probs.clone() for i in range(num_tokens): draft_id = draft_token_ids[i] if draft_id != PLACEHOLDER_TOKEN_ID: mask = target_indices[i] == draft_id prob[i, mask] = 0 else: # Gather draft probs at candidate indices flat_indices = target_indices.flatten() token_offsets = torch.arange(num_tokens, device=device)[:, None] * draft_probs.shape[1] flat_token_offsets = token_offsets.expand_as(target_indices).flatten() draft_probs_flat = draft_probs.flatten() valid_mask = flat_indices < draft_probs.shape[1] flat_draft_probs_at_indices = torch.where( valid_mask, draft_probs_flat[flat_token_offsets + flat_indices], torch.tensor(0.0, device=device) ) draft_probs_at_indices = flat_draft_probs_at_indices.view(num_tokens, vocab_size) prob = torch.maximum( target_probs - draft_probs_at_indices, torch.tensor(0.0, device=device), ) else: # normal mode if IS_NGRAM: token_indices = torch.arange(num_tokens, device=device) modified_target_probs = target_probs.clone() valid_draft_mask = draft_token_ids != PLACEHOLDER_TOKEN_ID modified_target_probs[ token_indices[valid_draft_mask], draft_token_ids[valid_draft_mask], ] = 0 prob = modified_target_probs else: prob = torch.maximum( target_probs - draft_probs, torch.tensor(0.0, pin_memory=True).to(device, non_blocking=True), ) q_values = q[token_to_batch] # [num_tokens, vocab_size] epsilon = 1e-10 q_values_safe = torch.where(q_values == 0, epsilon, q_values) q_values_safe = torch.where(torch.isinf(q_values), epsilon, q_values_safe) prob_over_q = prob / q_values_safe prob_over_q = torch.where((q_values == 0) | torch.isinf(q_values), -1e10, prob_over_q) if enable_reduce_sampling: # Get the index in selected vocab indices = torch.argmax(prob_over_q, dim=1) # Convert to global vocabulary indices output_token_ids[:] = target_indices[torch.arange(num_tokens, device=device), indices] else: recovered_ids = torch.argmax(prob_over_q, dim=1) output_token_ids[:] = recovered_ids def rejection_random_sample_block_verify_pytorch( output_token_ids, # [batch_size, max_spec_len + 1] cu_num_draft_tokens, # [batch_size] draft_token_ids, # [num_tokens] draft_probs, # [num_tokens, vocab_size] or None target_probs, # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] bonus_token_ids, # [batch_size] recovered_token_ids, # [num_tokens] uniform_probs, # [num_tokens] is_greedy, # [batch_size] max_spec_len, vocab_size, IS_NGRAM=False, target_indices=None, # [num_tokens, selected_vocab_size] global vocab indices enable_reduce_sampling=False, ENTROPY_VERIFY=False, POSTERIOR_THRESHOLD=0.95, POSTERIOR_ALPHA=0.4, EPSILON=1e-10, ori_target_probs=None, ): batch_size = output_token_ids.shape[0] device = output_token_ids.device zero_cpu = torch.tensor([0], pin_memory=True) zero_device = zero_cpu.to(device, non_blocking=True) cu_start = torch.cat([zero_device, cu_num_draft_tokens[:-1]]) cu_end = cu_num_draft_tokens num_draft_per_batch = (cu_end - cu_start)[:, None] pos_indices_cpu = torch.arange(max_spec_len, pin_memory=True) pos_indices = pos_indices_cpu.to(device, non_blocking=True)[None, :] valid_mask = pos_indices < num_draft_per_batch global_token_indices = cu_start[:, None] + pos_indices global_token_indices = global_token_indices.clamp(0, draft_token_ids.shape[0] - 1) draft_tokens = draft_token_ids[global_token_indices] placeholder_mask = draft_tokens == PLACEHOLDER_TOKEN_ID safe_draft_tokens = draft_tokens.masked_fill(placeholder_mask, 0) if IS_NGRAM: ones_cpu = torch.ones(1, pin_memory=True, dtype=torch.float32) draft_token_probs = ones_cpu.to(device, non_blocking=True).expand_as(draft_tokens) else: flat_indices = global_token_indices.flatten() flat_draft_tokens = safe_draft_tokens.flatten() flat_draft_probs = draft_probs[flat_indices, flat_draft_tokens] draft_token_probs = flat_draft_probs.view(batch_size, max_spec_len) # Get target token probs if enable_reduce_sampling: # When enable_reduce_sampling, need to search for draft token in candidates flat_global_indices = global_token_indices.flatten() flat_draft_tokens = draft_tokens.flatten() flat_target_indices = target_indices[flat_global_indices] flat_target_probs = target_probs[flat_global_indices] # Check if draft token is in candidates draft_expanded = flat_draft_tokens.unsqueeze(1) is_in_candidates = flat_target_indices == draft_expanded # Get the probability of draft token from target (if present) target_token_probs_flat = torch.where( is_in_candidates, flat_target_probs, torch.tensor(0.0, device=device) ).sum(dim=1) target_token_probs = target_token_probs_flat.view(batch_size, max_spec_len) else: flat_indices = global_token_indices.flatten() flat_draft_tokens = safe_draft_tokens.flatten() flat_target_probs = target_probs[flat_indices, flat_draft_tokens] target_token_probs = flat_target_probs.view(batch_size, max_spec_len) uniform_token_probs = uniform_probs[global_token_indices] recovered_tokens = recovered_token_ids[global_token_indices] pi = target_token_probs / draft_token_probs pi = pi.clamp(max=1.0) pi = torch.cumprod(pi, dim=-1) cum_uniform_token_probs = torch.cumprod(uniform_token_probs, dim=-1) if ENTROPY_VERIFY: entropy_probs = ori_target_probs if ori_target_probs is not None else target_probs all_target_dist = entropy_probs[global_token_indices] entropy = -(all_target_dist * torch.log(all_target_dist + EPSILON)).sum(dim=-1) exp_neg_entropy = torch.exp(-entropy * POSTERIOR_ALPHA) posterior_threshold_device = torch.tensor(POSTERIOR_THRESHOLD, device=device, dtype=torch.float32) threshold = torch.minimum(exp_neg_entropy, posterior_threshold_device) modified_cum_uniform_token_probs = threshold * cum_uniform_token_probs legal_mask = (draft_token_probs > 0) & (pi >= modified_cum_uniform_token_probs) else: legal_mask = (draft_token_probs > 0) & (pi >= cum_uniform_token_probs) legal_mask = legal_mask & valid_mask & (~placeholder_mask) last_accept_pos = torch.where( legal_mask.any(dim=-1, keepdim=True), (max_spec_len - legal_mask.flip(dims=[-1]).float().argmax(dim=-1, keepdim=True) - 1), -1, ) non_greedy_mask = (~is_greedy)[:, None] accept_mask = (pos_indices <= last_accept_pos) & valid_mask & non_greedy_mask output_token_ids[:, :max_spec_len] = torch.where(accept_mask, draft_tokens, output_token_ids[:, :max_spec_len]) reject_mask = (pos_indices == last_accept_pos + 1) & valid_mask & non_greedy_mask output_token_ids[:, :max_spec_len] = torch.where(reject_mask, recovered_tokens, output_token_ids[:, :max_spec_len]) bonus_mask = (last_accept_pos + 1 >= num_draft_per_batch) & non_greedy_mask all_positions_cpu = torch.arange(max_spec_len + 1, pin_memory=True) all_positions = all_positions_cpu.to(device, non_blocking=True)[None, :] bonus_pos_match = all_positions == num_draft_per_batch bonus_mask = bonus_mask & bonus_pos_match bonus_values_expanded = bonus_token_ids.view(-1, 1).expand(-1, max_spec_len + 1) output_token_ids[:] = torch.where(bonus_mask, bonus_values_expanded, output_token_ids) def sample_recovered_tokens_blockwise_pytorch( output_token_ids, # [num_tokens] cu_num_draft_tokens, # [batch_size] draft_token_ids, # [num_tokens] draft_probs, # [num_tokens, vocab_size] or None target_probs, # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] q, # [batch_size, vocab_size] or [batch_size, selected_vocab_size] vocab_size, IS_NGRAM=False, target_indices=None, # [num_tokens, selected_vocab_size] global vocab indices enable_reduce_sampling=False, ): _ = vocab_size device = output_token_ids.device num_tokens = output_token_ids.shape[0] batch_size = cu_num_draft_tokens.shape[0] if num_tokens == 0: return cu_start = torch.cat( [ torch.tensor([0], pin_memory=True).to(device, non_blocking=True), cu_num_draft_tokens[:-1], ] ) cu_end = cu_num_draft_tokens token_indices = torch.arange(num_tokens, device=device) in_range_mask = (token_indices[:, None] >= cu_start[None, :]) & (token_indices[:, None] < cu_end[None, :]) token_to_batch = torch.argmax(in_range_mask.int(), dim=1) token_to_batch = torch.where(in_range_mask.any(dim=1), token_to_batch, torch.zeros_like(token_to_batch)) pos_in_seq = token_indices - cu_start[token_to_batch] max_spec_len = int((cu_end - cu_start).max().item()) if IS_NGRAM: draft_token_scalar_probs = torch.ones(num_tokens, device=device, dtype=torch.float32) else: valid_draft_mask = draft_token_ids != PLACEHOLDER_TOKEN_ID safe_draft_token_ids = draft_token_ids.masked_fill(~valid_draft_mask, 0) draft_token_scalar_probs = draft_probs[token_indices, safe_draft_token_ids] draft_token_scalar_probs = torch.where( valid_draft_mask, draft_token_scalar_probs, torch.zeros_like(draft_token_scalar_probs), ) # Get target probability for each draft token if enable_reduce_sampling: # When enable_reduce_sampling, target_probs is [num_tokens, selected_vocab_size] # and target_indices maps selected positions to global vocab indices. # We need to search for draft_token_id in the selected candidates. draft_expanded = draft_token_ids[:, None] # [num_tokens, 1] is_in_candidates = target_indices == draft_expanded # [num_tokens, selected_vocab_size] target_token_scalar_probs = torch.where( is_in_candidates, target_probs, torch.tensor(0.0, device=device), ).sum(dim=1) # [num_tokens] else: valid_draft_mask = draft_token_ids != PLACEHOLDER_TOKEN_ID safe_draft_token_ids = draft_token_ids.masked_fill(~valid_draft_mask, 0) target_token_scalar_probs = target_probs[token_indices, safe_draft_token_ids] target_token_scalar_probs = torch.where( valid_draft_mask, target_token_scalar_probs, torch.zeros_like(target_token_scalar_probs), ) per_token_ratio = torch.where( draft_token_scalar_probs > 0, target_token_scalar_probs / draft_token_scalar_probs.clamp(min=1e-10), torch.zeros_like(target_token_scalar_probs), ) ratio_grid = torch.ones(batch_size, max_spec_len, device=device, dtype=torch.float32) ratio_grid[token_to_batch, pos_in_seq] = per_token_ratio p_prefix = torch.ones(batch_size, max_spec_len + 1, device=device, dtype=torch.float32) for k in range(max_spec_len): p_prefix[:, k + 1] = (p_prefix[:, k] * ratio_grid[:, k]).clamp(max=1.0) p_i = p_prefix[token_to_batch, pos_in_seq] p_i_expanded = p_i[:, None] if enable_reduce_sampling: # enable reduce_sampling: residual computation with selected vocab if IS_NGRAM: # Zero out the draft token in target_probs prob = target_probs.clone() for i in range(num_tokens): draft_id = draft_token_ids[i] if draft_id != PLACEHOLDER_TOKEN_ID: mask = target_indices[i] == draft_id prob[i, mask] = 0 residual = torch.clamp(p_i_expanded * prob, min=0.0) else: # Gather draft probs at candidate indices (same as sample_recovered_tokens_pytorch) flat_indices = target_indices.flatten() token_offsets = torch.arange(num_tokens, device=device)[:, None] * draft_probs.shape[1] flat_token_offsets = token_offsets.expand_as(target_indices).flatten() draft_probs_flat = draft_probs.flatten() valid_mask = flat_indices < draft_probs.shape[1] flat_draft_probs_at_indices = torch.where( valid_mask, draft_probs_flat[flat_token_offsets + flat_indices], torch.tensor(0.0, device=device) ) draft_probs_at_indices = flat_draft_probs_at_indices.view(num_tokens, -1) residual = torch.clamp(p_i_expanded * target_probs - draft_probs_at_indices, min=0.0) else: # normal mode if IS_NGRAM: modified_target = target_probs.clone() valid_draft_mask = draft_token_ids != PLACEHOLDER_TOKEN_ID modified_target[ token_indices[valid_draft_mask], draft_token_ids[valid_draft_mask], ] = 0.0 residual = torch.clamp(p_i_expanded * modified_target, min=0.0) else: residual = torch.clamp(p_i_expanded * target_probs - draft_probs, min=0.0) q_values = q[token_to_batch] epsilon = 1e-10 q_values_safe = torch.where(q_values == 0, epsilon, q_values) q_values_safe = torch.where(torch.isinf(q_values), epsilon, q_values_safe) prob_over_q = torch.where( (q_values == 0) | torch.isinf(q_values), torch.full_like(residual, -1e10), residual / q_values_safe, ) if enable_reduce_sampling: # Get the index in selected vocab, then convert to global vocab indices indices = torch.argmax(prob_over_q, dim=1) output_token_ids[:] = target_indices[torch.arange(num_tokens, device=device), indices] else: output_token_ids[:] = torch.argmax(prob_over_q, dim=1)