Support page size > 1 (#4356)

2025-03-12 22:22:39 -07:00
parent 2f6bacee03
commit c76040e31b
23 changed files with 877 additions and 284 deletions
@@ -103,6 +103,9 @@ class TpModelWorkerClient:
            self.worker.model_runner.token_to_kv_pool_allocator,
        )

+    def get_kv_cache(self):
+        return self.worker.model_runner.token_to_kv_pool
+
    def forward_thread_func(self):
        try:
            with torch.get_device_module(self.device).stream(self.forward_stream):
@@ -203,7 +206,7 @@ class TpModelWorkerClient:
            -(self.future_token_ids_ct + 1),
            -(self.future_token_ids_ct + 1 + bs),
            -1,
-            dtype=torch.int32,
+            dtype=torch.int64,
            device=self.device,
        )
        self.future_token_ids_ct = (