Adjust overlap event loop (#11507)

python/sglang/srt/managers/tp_worker.py

@@ -168,6 +168,7 @@ class TpModelWorker:
         )[0]
         set_random_seed(self.random_seed)
 
+        self.enable_overlap = not server_args.disable_overlap_schedule
         self.hicache_layer_transfer_counter = None
 
     def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter):
@@ -266,9 +267,18 @@ class TpModelWorker:
                 # Skip sampling and return logits for target forward
                 return batch_result
 
-            if model_worker_batch.delay_sample_launch:
-                batch_result.delay_sample_launch = True
-                batch_result.forward_batch = forward_batch
+            if (
+                self.enable_overlap
+                and model_worker_batch.sampling_info.grammars is not None
+            ):
+
+                def sample_batch_func():
+                    batch_result.next_token_ids = self.model_runner.sample(
+                        logits_output, forward_batch
+                    )
+                    return batch_result
+
+                batch_result.delay_sample_func = sample_batch_func
                 return batch_result
 
             if model_worker_batch.is_prefill_only: