[HiCache & JIT Kernel] Refactoring HiCache Write-Back Kernel (#21631)

2026-06-15 19:44:27 -07:00
parent a4a8a614b1
commit b5bcd76a41
12 changed files with 1207 additions and 81 deletions
@@ -726,9 +726,15 @@ class HiCacheController:
            return

        op = CacheOperation.merge_ops(self.write_queue)
-        host_indices, device_indices = self.move_indices(
-            op.host_indices, op.device_indices
-        )
+        # For now, kernel write-back keeps host indices on CPU only for page_first.
+        # More layouts can use this path once their write-back kernels accept CPU
+        # destination indices.
+        if self.io_backend == "kernel" and self.mem_pool_host.layout == "page_first":
+            host_indices, device_indices = op.host_indices, op.device_indices
+        else:
+            host_indices, device_indices = self.move_indices(
+                op.host_indices, op.device_indices
+            )
        self.write_queue.clear()

        start_event = device_module.Event()