Implement async scheduling for single card

Author: gty111

gllm/async_worker.py

@@ -169,6 +169,11 @@ async def run_server(args):
         help="Experimental feature for worker implemented by async",
         action="store_true",
     )
+    parser.add_argument(
+        "--async-scheduling",
+        help="Overlap CPU input preparation for the next batch with GPU execution of the current batch",
+        action="store_true",
+    )
     parser.add_argument(
         "--gpu-memory-util",
         type=float,
@@ -301,6 +306,7 @@ async def run_server(args):
         assigned_layers=args.assigned_layers,
         schedule_method=args.schedule_method,
         use_async_worker=args.use_async_worker,
+        async_scheduling=args.async_scheduling,
         use_thinking=not args.disable_thinking,
         disable_cuda_graph=args.disable_cuda_graph,
         max_cuda_graph_bs=args.max_cuda_graph_bs,

gllm/entrypoints/api_server.py

@@ -6,6 +6,15 @@
 class Sampler:
 
     def forward(self, logits: torch.Tensor, input_data: InputData):
+        return self.forward_gpu(logits, input_data).cpu().numpy().tolist()
+
+    def forward_gpu(self, logits: torch.Tensor, input_data: InputData) -> torch.Tensor:
+        """Same as forward() but returns a GPU tensor without D2H copy.
+
+        Used by async scheduling so the D2H transfer can be initiated with
+        non_blocking=True on a dedicated copy stream, overlapping with
+        CPU-side scheduling work for the next batch.
+        """
         # repetition_penalty
         logits /= torch.where(logits > 0, input_data.repetition_penalty, 1.0)
         logits *= torch.where(logits <= 0, 1.0, input_data.repetition_penalty)
@@ -17,8 +26,7 @@ def forward(self, logits: torch.Tensor, input_data: InputData):
 
         q = torch.empty_like(probs)
         q.exponential_()
-        return probs.div_(q).argmax(dim=1).cpu().numpy().tolist()
-        # return torch.multinomial(probs, 1).squeeze(1).cpu().numpy().tolist()
+        return probs.div_(q).argmax(dim=1)
 
     def _apply_top_k_top_p(
         self,

gllm/layers/sampler.py

@@ -6,13 +6,12 @@
 import tqdm
 from logger import logger
 
-from gllm.async_worker import AsyncWorker, run_worker_async
 from gllm.comm import IPCPackage, zmqComm
 from gllm.id_allocator import IDAllocator
-from gllm.model_runner import ModelRunner
+from gllm.model_runner import ModelRunner, AsyncModelRunner
 from gllm.sequence import Sequence
 from gllm.utils import get_model_load_pbar, init_logger, random_uuid
-from gllm.worker import Worker, run_worker
+from gllm.worker import Worker, AsyncWorker, run_worker
 
 
 class LLM:
@@ -40,6 +39,7 @@ def __init__(
         assigned_layers=None,
         schedule_method="chunked_prefill",
         use_async_worker=False,
+        async_scheduling=False,
         use_thinking=True,
         disable_cuda_graph=False,
         max_cuda_graph_bs=32,
@@ -50,7 +50,8 @@ def __init__(
         init_logger()
         self.model_path = model_path
         self.load_format = load_format
-        self.model_runner = ModelRunner(
+        model_runner_cls = AsyncModelRunner if async_scheduling else ModelRunner
+        self.model_runner = model_runner_cls(
             load_format=load_format,
             model_path=model_path,
             gpu_memory_util=gpu_memory_util,
@@ -90,8 +91,11 @@ def __init__(
         self.assigned_layers = assigned_layers
         self.schedule_method = schedule_method
         self.use_async_worker = use_async_worker
+        self.async_scheduling = async_scheduling
 
         logger.info(f"Schedule method: {schedule_method}")
+        if async_scheduling:
+            logger.info("Async scheduling enabled")
 
         # Interact with workers
         self.wait_lists: List[Sequence] = []
@@ -168,7 +172,7 @@ def init_workers(self):
             self.load_progress()
 
     def start_worker(self, local_rank, pp_rank, tp_rank):
-        worker_cls = Worker if not self.use_async_worker else AsyncWorker
+        worker_cls = Worker if not self.async_scheduling else AsyncWorker
         comm = zmqComm(
             self.host,
             self.zmq_port_base,
@@ -195,7 +199,7 @@ def start_worker(self, local_rank, pp_rank, tp_rank):
             self.schedule_method,
         )
         process = self.ctx.Process(
-            target=run_worker if not self.use_async_worker else run_worker_async,
+            target=run_worker,
             args=(worker,),
             daemon=True,
         )