This commit is contained in:
@@ -2634,6 +2634,16 @@ class Scheduler(
|
||||
self.future_map.store_to_map(batch_result.future_indices, batch_result)
|
||||
batch_result.copy_to_cpu(return_logprob=self.cur_batch.return_logprob)
|
||||
|
||||
# Release the closure and large GPU tensors that are no longer needed.
|
||||
# The delay_sample_func closure captures forward_batch (which holds
|
||||
# sampling_info with vocab_mask) and logits_output (which holds
|
||||
# next_token_logits). Without clearing these, they stay alive via
|
||||
# batch_result in result_queue and batch_record_buf until the next
|
||||
# iteration, causing a steady VRAM leak with structured output.
|
||||
batch_result.delay_sample_func = None
|
||||
if batch_result.logits_output is not None:
|
||||
batch_result.logits_output.next_token_logits = None
|
||||
|
||||
def process_batch_result(
|
||||
self,
|
||||
batch: ScheduleBatch,
|
||||
|
||||
@@ -2727,6 +2727,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
sampling_info.update_regex_vocab_mask()
|
||||
sampling_info.apply_logits_bias(logits_output.next_token_logits)
|
||||
|
||||
# Release the vocab_mask GPU tensor immediately after it has been applied
|
||||
# to the logits. In overlap scheduling, the sampling_info (and its
|
||||
# vocab_mask) can be kept alive by the delay_sample_func closure and
|
||||
# batch_record_buf until the next iteration, causing a steady VRAM leak
|
||||
# when structured output (grammar) is used.
|
||||
sampling_info.vocab_mask = None
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits_output: LogitsProcessorOutput,
|
||||
|
||||
Reference in New Issue
Block a user