Simplify logits penalizer (#2086)

python/sglang/srt/managers/schedule_batch.py

@@ -1019,7 +1019,7 @@ class ScheduleBatch:
             extend_prefix_lens = self.prefix_lens
             extend_logprob_start_lens = self.extend_logprob_start_lens
 
-        if self.sampling_info is not None:
+        if self.sampling_info:
             if self.has_grammar:
                 self.sampling_info.grammars = [req.grammar for req in self.reqs]
             else:
@@ -1063,6 +1063,7 @@ class ScheduleBatch:
             out_cache_loc=self.out_cache_loc,
             return_logprob=self.return_logprob,
             decoding_reqs=self.decoding_reqs,
+            sampling_info=dataclasses.replace(self.sampling_info),
         )
 
     def __str__(self):
@@ -1122,20 +1123,6 @@ class ModelWorkerBatch:
     # Sampling info
     sampling_info: SamplingBatchInfo
 
-    def copy(self):
-        return dataclasses.replace(self, sampling_info=self.sampling_info.copy())
-
-    def to(self, device: str):
-        self.input_ids = self.input_ids.to(device, non_blocking=True)
-        self.req_pool_indices = self.req_pool_indices.to(device, non_blocking=True)
-        self.seq_lens = self.seq_lens.to(device, non_blocking=True)
-        self.out_cache_loc = self.out_cache_loc.to(device, non_blocking=True)
-        self.req_to_token_pool_records = [
-            (x, y.to(device, non_blocking=True))
-            for x, y in self.req_to_token_pool_records
-        ]
-        self.sampling_info.to(device)
-
 
 @triton.jit
 def write_req_to_token_pool_triton(

python/sglang/srt/managers/scheduler.py

@@ -931,14 +931,14 @@ class Scheduler:
 
             # Check finish conditions
             logprob_pt = 0
-            for i, req in enumerate(batch.reqs):
+            for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
                 if req.is_retracted:
                     continue
 
                 if req.is_being_chunked <= 0:
                     # Inflight reqs' prefill is not finished
                     req.completion_tokens_wo_jump_forward += 1
-                    req.output_ids.append(next_token_ids[i])
+                    req.output_ids.append(next_token_id)
                     req.check_finished()
 
                     if req.finished():
@@ -947,7 +947,7 @@ class Scheduler:
                         self.tree_cache.cache_unfinished_req(req)
 
                     if req.grammar is not None:
-                        req.grammar.accept_token(next_token_ids[i])
+                        req.grammar.accept_token(next_token_id)
 
                     if req.return_logprob:
                         logprob_pt += self.add_logprob_return_values(

python/sglang/srt/managers/tp_worker.py

@@ -16,6 +16,7 @@ limitations under the License.
 """A tensor parallel worker."""
 
 import logging
+import threading
 from typing import Optional
 
 from sglang.srt.configs.model_config import ModelConfig
@@ -138,9 +139,15 @@ class TpModelWorker:
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
         self.model_runner.forward(forward_batch)
 
-    def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
+    def forward_batch_generation(
+        self,
+        model_worker_batch: ModelWorkerBatch,
+        launch_event: Optional[threading.Event] = None,
+    ):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
         logits_output = self.model_runner.forward(forward_batch)
+        if launch_event:
+            launch_event.set()
         next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
         return logits_output, next_token_ids
 

python/sglang/srt/managers/tp_worker_overlap_thread.py

@@ -15,6 +15,7 @@ limitations under the License.
 
 """A tensor parallel worker."""
 
+import dataclasses
 import logging
 import threading
 import time
@@ -107,7 +108,7 @@ class TpModelWorkerClient:
 
             # Run forward
             logits_output, next_token_ids = self.worker.forward_batch_generation(
-                model_worker_batch
+                model_worker_batch, self.launch_event
             )
 
             # Update the future token ids map
@@ -134,7 +135,6 @@ class TpModelWorkerClient:
             next_token_ids = next_token_ids.to("cpu", non_blocking=True)
             copy_event.record()
 
-            self.launch_event.set()
             self.output_queue.put((copy_event, logits_output, next_token_ids))
 
     def resolve_batch_result(self, bid: int):
@@ -159,7 +159,10 @@ class TpModelWorkerClient:
 
     def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
         # Push a new batch to the queue
-        self.input_queue.put((model_worker_batch.copy(), self.future_token_ids_ct))
+        model_worker_batch.sampling_info = dataclasses.replace(
+            model_worker_batch.sampling_info
+        )
+        self.input_queue.put((model_worker_batch, self.future_token_ids_ct))
 
         # Allocate output future objects
         bs = len(model_worker_batch.seq_lens)