diff --git a/aphrodite/modeling/layers/sampler.py b/aphrodite/modeling/layers/sampler.py index 0bc03c22a6..328aec8aed 100644 --- a/aphrodite/modeling/layers/sampler.py +++ b/aphrodite/modeling/layers/sampler.py @@ -526,9 +526,10 @@ def forward( sampling_tensors.xtc_probabilities) - # We use float32 for probabilities and log probabilities. # Compute the probabilities. - probs = torch.softmax(logits, dim=-1, dtype=torch.float) + # No cast required, as softmaxes internally accumulate with >= fp32. + # See https://github.com/pytorch/pytorch/blob/v2.4.0/aten/src/ATen/native/cuda/PersistentSoftmax.cuh#L62 + probs = torch.softmax(logits, dim=-1) # skew needs to be applied post-softmax if do_skew: @@ -546,7 +547,9 @@ def forward( logits = torch.log(probs) # Compute the log probabilities. - logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) + logprobs = torch.log_softmax(logits, dim=-1) + # Logits unused past this point. They are HUGE, for prompt_logprobs. + del logits # Sample the next tokens. maybe_deferred_sample_results, maybe_sampled_tokens_tensor = _sample( @@ -1575,7 +1578,7 @@ def _sample( ) -def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: +def _get_ranks(logprobs: torch.Tensor, query_indices: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: """ This function calculates the ranks of the chosen tokens in a logprob tensor. Args: @@ -1587,9 +1590,14 @@ def _get_ranks(x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: Each element in the returned tensor represents the rank of the chosen token in the input logprob tensor. """ - vals = x[torch.arange(0, len(x), device=x.device, dtype=indices.dtype), - indices] - return (x > vals[:, None]).long().sum(1).add_(1) + expanded_indices = torch.zeros(logprobs.shape[0], dtype=indices.dtype, device=indices.device) + expanded_indices.scatter_(-1, query_indices, indices) + vals = torch.gather(logprobs, -1, expanded_indices.unsqueeze(dim=1)) + # doing this in a single pass is not practical for prompt_logprobs. + outs = torch.zeros_like(indices) + for i in range(0, vals.shape[0], 10): + outs[i:i+10] = (logprobs[i:i+10] > vals[i:i+10]).sum(dim=-1, dtype=indices.dtype) + return outs[query_indices].add_(1) def get_logprobs( @@ -1689,7 +1697,8 @@ def get_logprobs( next_token_ids_gpu, ]] ranks = _get_ranks( - logprobs[query_indices_gpu], + logprobs, + query_indices_gpu, next_token_ids_gpu, ) assert selected_logprobs.shape[0] == ranks.shape[0] diff --git a/aphrodite/processing/scheduler.py b/aphrodite/processing/scheduler.py index d1d1842e81..cd5155af11 100644 --- a/aphrodite/processing/scheduler.py +++ b/aphrodite/processing/scheduler.py @@ -1239,10 +1239,14 @@ def schedule( block_tables[seq_id] = self.block_manager.get_block_table(seq) self.block_manager.access_all_blocks_in_seq(seq, now) - if self.cache_config.enable_prefix_caching: + + if (self.cache_config.enable_prefix_caching and + not seq_group.sampling_params.prompt_logprobs): common_computed_block_nums = ( self.block_manager.get_common_computed_block_ids( seq_group.get_seqs(status=SequenceStatus.RUNNING))) + else: + common_computed_block_nums = [] do_sample = True is_prompt = seq_group.is_prefill()