[hicache] Optimization for DMA copy (#8245)

2025-08-09 17:16:07 -07:00
parent ef48d5547e
commit 5c31b35db2
2 changed files with 26 additions and 6 deletions
--- a/python/sglang/srt/managers/cache_controller.py
+++ b/python/sglang/srt/managers/cache_controller.py
@@ -433,7 +433,9 @@ class HiCacheController:
        if self.io_backend == "kernel":
            return host_indices.to(self.mem_pool_device.device), device_indices
        elif self.io_backend == "direct":
-            return host_indices, device_indices.cpu()
+            device_indices = device_indices.cpu()
+            host_indices, idx = host_indices.sort()
+            return host_indices, device_indices.index_select(0, idx)
        else:
            raise ValueError(f"Unsupported io backend")