From bc1a2bdf98535a4197e865b8182fb10612091302 Mon Sep 17 00:00:00 2001 From: 50h100a Date: Mon, 16 Dec 2024 19:47:42 +0000 Subject: [PATCH 1/2] do not use cached chunks for prompt_logprobs --- aphrodite/processing/scheduler.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/aphrodite/processing/scheduler.py b/aphrodite/processing/scheduler.py index 9fecc6ebca..1452cdc75b 100644 --- a/aphrodite/processing/scheduler.py +++ b/aphrodite/processing/scheduler.py @@ -1046,9 +1046,6 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: scheduler_outputs = self._schedule() now = time.time() - if not self.cache_config.enable_prefix_caching: - common_computed_block_nums = [] - # Create input data structures. seq_group_metadata_list: List[SequenceGroupMetadata] = [] for i, scheduled_seq_group in enumerate( @@ -1079,10 +1076,14 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]: 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() From d49ead733470fdfcd74eed957bba9fb50258d5e7 Mon Sep 17 00:00:00 2001 From: 50h100a Date: Mon, 16 Dec 2024 20:00:08 +0000 Subject: [PATCH 2/2] reduce sampler peak memory usage --- aphrodite/modeling/layers/sampler.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/aphrodite/modeling/layers/sampler.py b/aphrodite/modeling/layers/sampler.py index 3c506121f4..5092dbe16e 100644 --- a/aphrodite/modeling/layers/sampler.py +++ b/aphrodite/modeling/layers/sampler.py @@ -402,9 +402,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: @@ -422,7 +423,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. sample_results, maybe_sampled_tokens_tensor = _sample( @@ -1424,7 +1427,7 @@ def _sample( # sampling_tensors) -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: @@ -1436,9 +1439,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( @@ -1538,7 +1546,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]