[HiCache] Bulk/staging reload fast-path for contiguous host runs

python/sglang/srt/mem_cache/pool_host/base.py

@@ -12,6 +12,7 @@
 import torch
 
 from sglang.srt.mem_cache.host_memory import available_host_memory_bytes
+from sglang.srt.mem_cache.pool_host.coalesce_reload import take_contiguous_run
 from sglang.srt.mem_cache.pool_host.common import (
     _cuda_host_unregister,
     get_allocator_from_storage,
@@ -512,9 +513,11 @@ def _merge_release_slots(self):
             return
 
         if len(self.free_slots) == 0 and len(self.release_slots) == 1:
-            self.free_slots = self.release_slots[0]
+            merged = self.release_slots[0]
         else:
-            self.free_slots = torch.cat([self.free_slots, *self.release_slots])
+            merged = torch.cat([self.free_slots, *self.release_slots])
+        # keep free_slots sorted so alloc() can find a contiguous run in one scan
+        self.free_slots = torch.sort(merged).values
 
         self.release_slots = []
         self.num_release_slots = 0
@@ -540,12 +543,15 @@ def alloc(self, need_size: int) -> Optional[torch.Tensor]:
         if need_size > len(self.free_slots):
             self._merge_release_slots()
 
-        select_index = self.free_slots[:need_size]
-        self.free_slots = self.free_slots[need_size:]
+        head = self.free_slots[:need_size]
+        select_index, self.free_slots = take_contiguous_run(self.free_slots, need_size)
 
-        assert not self.slot_used[select_index].any(), (
+        # take_contiguous_run can skip the head of free_slots, so check the head
+        # too, otherwise an in-use slot leaked there would never be caught
+        checked = torch.cat([head, select_index])
+        assert not self.slot_used[checked].any(), (
             f"Double-alloc detected: slots already allocated: "
-            f"{select_index[self.slot_used[select_index]].tolist()}."
+            f"{checked[self.slot_used[checked]].unique().tolist()}."
         )
         self.slot_used[select_index] = True
 

python/sglang/srt/mem_cache/pool_host/coalesce_reload.py

@@ -0,0 +1,153 @@
+# Author: adlashab <adlashab@amd.com>
+#
+# Contiguous host-slot packing and a bulk/staging reload fast-path for HiCache.
+#
+# The host KV pool hands out free slot ids first-fit and appends freed ids on
+# free, so after a few fill/evict cycles a re-admitted prefix lands on a scrambled
+# set of host offsets and the reload runs the per-index transfer kernel over
+# scattered indices. `take_contiguous_run` makes the allocator hand out a
+# contiguous run whenever one is free (the pool keeps its free list sorted), so a
+# prefix occupies a contiguous host region. `bulk_reload` then moves that region
+# in one shot: a single copy_ when the device target is also contiguous, or a bulk
+# copy into a contiguous device staging buffer followed by an on-device
+# index_copy_ when the device slots are scattered. Any non-contiguous host case,
+# dtype/shape mismatch, or non-contiguous destination returns False so the caller
+# runs the exact stock kernel, so a reload is never wrong.
+
+import json
+import os
+import threading
+
+import torch
+
+_staging = {}
+_staging_lock = threading.Lock()
+
+# Fire counters so a live run can confirm the fast-path took the contiguous path
+# rather than falling back. Inert unless SGLANG_HICACHE_BULK_STATS names a file.
+_fire = {"bulk": 0, "stage": 0, "fallback": 0, "alloc_contig": 0, "alloc_frag": 0}
+_STATS_PATH = os.environ.get("SGLANG_HICACHE_BULK_STATS")
+
+
+def bulk_stats():
+    return dict(_fire)
+
+
+def _bump(kind):
+    _fire[kind] = _fire.get(kind, 0) + 1
+    if _STATS_PATH:
+        try:
+            with open(_STATS_PATH + ".tmp", "w") as f:
+                json.dump(_fire, f)
+            os.replace(_STATS_PATH + ".tmp", _STATS_PATH)
+        except Exception:
+            pass
+
+
+def take_contiguous_run(free_slots, need_size):
+    """Pick which free slots to hand out. `free_slots` is a 1-D int64 tensor kept
+    sorted ascending by the pool. Returns (selected, remaining): `selected` is a
+    length-`need_size` contiguous unit-stride run when one is available (best fit,
+    smallest run that holds it), else the first `need_size` slots. `remaining`
+    stays sorted. Same success set as the stock first-`need_size` pick."""
+    n = int(free_slots.numel())
+    if need_size <= 0:
+        return free_slots[:0], free_slots
+    if n <= need_size:
+        return free_slots, free_slots[:0]
+    d = free_slots[1:] - free_slots[:-1]
+    brk = torch.nonzero(d != 1, as_tuple=False).flatten()
+    starts = torch.cat([free_slots.new_zeros(1), brk + 1])
+    ends = torch.cat([brk + 1, free_slots.new_full((1,), n)])
+    lens = ends - starts
+    fit = lens >= need_size
+    if bool(fit.any()):
+        cand = torch.nonzero(fit, as_tuple=False).flatten()
+        j = int(cand[torch.argmin(lens[cand])])
+        p = int(starts[j])
+        sel = free_slots[p : p + need_size]
+        rem = torch.cat([free_slots[:p], free_slots[p + need_size :]])
+        _bump("alloc_contig")
+        return sel, rem
+    _bump("alloc_frag")
+    return free_slots[:need_size], free_slots[need_size:]
+
+
+def _single_run(idx):
+    n = int(idx.numel())
+    if n == 0:
+        return None
+    start = int(idx[0].item())
+    ar = torch.arange(start, start + n, dtype=idx.dtype, device=idx.device)
+    return (start, n) if torch.equal(idx, ar) else None
+
+
+def _bulk_copy(src, dst):
+    if src.dtype != dst.dtype or src.numel() != dst.numel() or not dst.is_contiguous():
+        return False
+    dst.view(-1).copy_(src.reshape(-1), non_blocking=True)
+    return True
+
+
+def _get_staging(dst, n):
+    row = tuple(dst.shape[1:])
+    key = (str(dst.device), dst.dtype, row)
+    with _staging_lock:
+        buf = _staging.get(key)
+        if buf is None or buf.shape[0] < n:
+            cap = 1 if buf is None else buf.shape[0]
+            while cap < n:
+                cap *= 2
+            try:
+                buf = torch.empty((cap,) + row, dtype=dst.dtype, device=dst.device)
+            except Exception:
+                return None
+            _staging[key] = buf
+        return buf
+
+
+def _bulk_scatter(src_contig, dst, device_indices):
+    n = int(src_contig.shape[0])
+    if n == 0:
+        return True
+    if src_contig.dtype != dst.dtype or not dst.is_contiguous():
+        return False
+    di = device_indices
+    if str(di.device) != str(dst.device):
+        di = di.to(dst.device, non_blocking=True)
+    if di.dtype != torch.long:
+        di = di.to(torch.long)
+    if int(di.numel()) != n:
+        return False
+    staging = _get_staging(dst, n)
+    if staging is None:
+        return False
+    staging[:n].copy_(src_contig, non_blocking=True)
+    dst.index_copy_(0, di, staging[:n])
+    return True
+
+
+def bulk_reload(src_layer_buf, dst_layer_buf, host_indices, device_indices):
+    """Bulk reload one layer when the host indices are one contiguous run. Handles
+    both the MLA KV anchor (token runs) and the DSA indexer (page runs); the caller
+    passes the matching per-layer buffers. Returns True if it moved the data,
+    False to fall back to the stock per-index kernel."""
+    hr = _single_run(host_indices)
+    if hr is None:
+        _bump("fallback")
+        return False
+    h0, n = hr
+    src = src_layer_buf[h0 : h0 + n]
+    dr = _single_run(device_indices)
+    if dr is not None:
+        d0, _ = dr
+        if _bulk_copy(src, dst_layer_buf[d0 : d0 + n]):
+            _bump("bulk")
+            return True
+        _bump("fallback")
+        return False
+    if _bulk_scatter(src, dst_layer_buf, device_indices):
+        _bump("stage")
+        return True
+    _bump("fallback")
+    return False

python/sglang/srt/mem_cache/pool_host/dsa.py

@@ -374,16 +374,24 @@ def load_to_device_per_layer(
         host_page_indices, device_page_indices = self._get_indexer_page_indices(
             host_indices, device_indices
         )
+        from sglang.srt.mem_cache.pool_host.coalesce_reload import bulk_reload
+
         use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0
         if use_kernel:
             if self.layout == "layer_first":
-                transfer_kv_per_layer_mla(
-                    src=self.index_k_with_scale_buffer[host_layer_id],
-                    dst=device_pool.index_k_with_scale_buffer[device_layer_id],
-                    src_indices=host_page_indices,
-                    dst_indices=device_page_indices,
-                    item_size=self.indexer_page_stride_size,
-                )
+                if not bulk_reload(
+                    self.index_k_with_scale_buffer[host_layer_id],
+                    device_pool.index_k_with_scale_buffer[device_layer_id],
+                    host_page_indices,
+                    device_page_indices,
+                ):
+                    transfer_kv_per_layer_mla(
+                        src=self.index_k_with_scale_buffer[host_layer_id],
+                        dst=device_pool.index_k_with_scale_buffer[device_layer_id],
+                        src_indices=host_page_indices,
+                        dst_indices=device_page_indices,
+                        item_size=self.indexer_page_stride_size,
+                    )
             elif self.layout == "page_first":
                 transfer_kv_per_layer_mla_pf_lf(
                     src=self.index_k_with_scale_buffer,

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -24,6 +24,7 @@
     _WRITE_BACK_STAGING_PAGE_CHUNK,
     HostKVCache,
 )
+from sglang.srt.mem_cache.pool_host.coalesce_reload import bulk_reload
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,
     make_kernel_ptr_table,
@@ -641,7 +642,14 @@ def load_to_device_per_layer(
 
         if io_backend == "kernel":
             if self.layout == "layer_first":
-                if self.can_use_jit:
+                if bulk_reload(
+                    self.kv_buffer[host_layer_id],
+                    device_buffer,
+                    host_indices,
+                    device_indices,
+                ):
+                    pass
+                elif self.can_use_jit:
                     jit_transfer_hicache_one_layer_mla(
                         page_size=self.page_size,
                         cache_dst=device_buffer,

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -233,6 +233,16 @@ def test_double_alloc(self):
         self.assertIn("Double-alloc", msg)
         self.assertIn(f"[{int(leak[0])}]", msg)
 
+    def test_alloc_prefers_contiguous_run(self):
+        everything = self.host_pool.alloc(self.host_pool.size)
+        self.host_pool.free(everything[:2])
+        self.host_pool.free(everything[6:10])
+        # merged free list is [0, 1, 6, 7, 8, 9], first-fit would hand out the
+        # split [0, 1, 6, 7]
+        indices = self.host_pool.alloc(4)
+        self.assertEqual(indices.tolist(), everything[6:10].tolist())
+        self.assertEqual(self.host_pool.free_slots.tolist(), everything[:2].tolist())
+
     def test_double_free(self):
         indices = self.host_pool.alloc(4)
         self.assertEqual(len(indices), 4)