Add device support (#1607)

2024-10-11 17:05:58 +08:00
parent 5476ccad8f
commit 8275049ce3
5 changed files with 96 additions and 52 deletions
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -423,6 +423,9 @@ class ScheduleBatch:
    # Stream
    has_stream: bool = False

+    # device
+    device: str = "cuda"
+
    # Has regex
    has_regex: bool = False

@@ -439,6 +442,7 @@ class ScheduleBatch:
            tree_cache=tree_cache,
            return_logprob=return_logprob,
            has_stream=has_stream,
+            device=req_to_token_pool.device,
            has_regex=has_regex,
        )