Fix VRAM leak in overlap scheduling with structured output (#20640) (#20697)

This commit is contained in:
Cishoon
2026-03-23 08:07:39 +08:00
committed by GitHub
parent 343998865a
commit 999bad5aba
2 changed files with 17 additions and 0 deletions

View File

@@ -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,

View File

@@ -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,