fix move_indices

Author: huangtingwei9988

python/sglang/srt/managers/cache_controller.py

@@ -495,9 +495,12 @@ def move_indices(self, host_indices, device_indices):
         if self.io_backend == "kernel":
             return host_indices.to(self.mem_pool_device.device), device_indices
         elif self.io_backend == "direct":
-            device_indices = device_indices.cpu()
-            host_indices, idx = host_indices.sort()
-            return host_indices, device_indices.index_select(0, idx)
+            if self.mem_pool_host.layout == "layer_first":
+                device_indices = device_indices.cpu()
+                host_indices, idx = host_indices.sort()
+                return host_indices, device_indices.index_select(0, idx)
+            elif self.mem_pool_host.layout == "page_first_direct":
+                return host_indices, device_indices.cpu()
         else:
             raise ValueError(f"Unsupported io backend")