sglang 0.1.21__py3-none-any.whl → 0.1.22__py3-none-any.whl

sglang/bench_latency.py

@@ -156,14 +156,14 @@ def extend(reqs, model_runner):
     )
     batch.prepare_for_extend(model_runner.model_config.vocab_size, None)
     output = model_runner.forward(batch, ForwardMode.EXTEND)
-    next_token_ids, _ = batch.sample(output.next_token_logits)
+    next_token_ids = batch.sample(output.next_token_logits)
     return next_token_ids, output.next_token_logits, batch
 
 
 def decode(input_token_ids, batch, model_runner):
     batch.prepare_for_decode(input_token_ids.cpu().numpy())
     output = model_runner.forward(batch, ForwardMode.DECODE)
-    next_token_ids, _ = batch.sample(output.next_token_logits)
+    next_token_ids = batch.sample(output.next_token_logits)
     return next_token_ids, output.next_token_logits
 
 
@@ -284,23 +284,26 @@ def main(server_args, bench_args):
     else:
         work_func = latency_test
 
-    workers = []
-    for tp_rank in range(server_args.tp_size):
-        proc = multiprocessing.Process(
-            target=work_func,
-            args=(
-                server_args,
-                bench_args,
-                tp_rank,
-            ),
-        )
-        proc.start()
-        workers.append(proc)
-
-    for proc in workers:
-        proc.join()
-
-    proc.terminate()
+    if server_args.tp_size == 1:
+        work_func(server_args, bench_args, 0)
+    else:
+        workers = []
+        for tp_rank in range(server_args.tp_size):
+            proc = multiprocessing.Process(
+                target=work_func,
+                args=(
+                    server_args,
+                    bench_args,
+                    tp_rank,
+                ),
+            )
+            proc.start()
+            workers.append(proc)
+
+        for proc in workers:
+            proc.join()
+
+        proc.terminate()
 
 
 if __name__ == "__main__":