diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fce74e237..42e841fa4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e5ba805aa..93a5ccce0 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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,