5080deploy: merge upstream #500, single-launch prefill buffer invalidation without hidden sync

python/freetoken/kernel/triton/moe/__init__.py

@@ -0,0 +1,5 @@
+"""Triton kernels for the MoE offload cache (``moe/offload_cache.py``)."""
+
+from .invalidate import invalidate_prefill_slots
+
+__all__ = ["invalidate_prefill_slots"]

python/freetoken/kernel/triton/moe/invalidate.py

@@ -0,0 +1,75 @@
+"""Single-launch invalidation of the MoE prefill overlap buffers.
+
+``OffloadMoeCache._invalidate_prefill_buffer`` runs twice per prefill chunk (once per
+double-buffer slot) and used to clear the buffer's slot map with a boolean-mask index:
+
+    slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
+
+A boolean index produces a data-dependent shape, so every call hid a device-to-host
+synchronization. With two buffer reuses per chunk x 48 layers, the host waited -- per
+layer -- for ALL enqueued GPU work (the cached-context attention included) before
+prefetching the next layer, which serialized the prefill-overlap pipeline. On long
+cached contexts that turned a ~3 s chunk into a 20-150 s turn (py-spy: 92% of the
+scheduler inside this invalidation + ``synchronize``, GPU 0-3% utilized).
+
+This kernel is fixed-shape: one launch, no synchronization, identical result. Zeroing
+``usage`` here (instead of a separate ``zero_``) keeps the buffer's slots as the
+argmin(usage) eviction victims, exactly like the code it replaces.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _invalidate_slots_kernel(
+    id_of_slot_ptr,   # (num_slots,) int32 -- slot -> expert id, -1 if empty
+    slot_for_id_ptr,  # (num_layers * num_experts) int32 -- expert id -> slot, -1 if not resident
+    usage_ptr,        # (num_slots,) int64 -- last-used step; zeroed so these slots are
+                      # the argmin(usage) victims in ensure_experts' eviction
+    slot_start,
+    n,
+    BLOCK: tl.constexpr,
+):
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    m = i < n
+    e = tl.load(id_of_slot_ptr + slot_start + i, mask=m, other=-1)
+    # clear the expert's residency flag; empty slots (-1) are masked out
+    tl.store(slot_for_id_ptr + e, -1, mask=m & (e >= 0))
+    tl.store(id_of_slot_ptr + slot_start + i, -1, mask=m)
+    tl.store(usage_ptr + slot_start + i, 0, mask=m)
+
+
+def invalidate_prefill_slots(
+    id_of_slot: torch.Tensor,
+    slot_for_id: torch.Tensor,
+    usage: torch.Tensor,
+    slot_start: int,
+    num_experts: int,
+) -> None:
+    """Free the ``num_experts`` slots starting at ``slot_start`` (fixed shape, no sync).
+
+    ``slot_for_id`` may be the 2-D ``(num_layers, num_experts)`` map: ids are flat
+    (``layer * num_experts + expert``) and the tensor is contiguous, so the kernel's
+    flat addressing matches the ``view(-1)`` the eager path used.
+    """
+    # The raw-pointer kernel cannot clamp like the eager slices did; refuse a buffer
+    # that does not fit instead of reading past the cache.
+    if slot_start + num_experts > id_of_slot.numel():
+        raise ValueError(
+            f"prefill buffer [{slot_start}, {slot_start + num_experts}) exceeds "
+            f"id_of_slot size {id_of_slot.numel()}"
+        )
+    # The kernel needs a CUDA context; CPU test caches keep the old eager path, where
+    # the hidden sync is harmless (no enqueued GPU work to drain).
+    if id_of_slot.device.type != "cuda":
+        old_ids = id_of_slot[slot_start : slot_start + num_experts]
+        slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
+        old_ids.fill_(-1)
+        usage[slot_start : slot_start + num_experts].zero_()
+        return
+    grid = ((num_experts + 255) // 256,)
+    _invalidate_slots_kernel[grid](
+        id_of_slot, slot_for_id, usage, slot_start, num_experts, BLOCK=256
+    )

python/freetoken/moe/offload_cache.py

@@ -634,13 +634,16 @@ def _init_prefill_overlap_buffers(self) -> None:
 
     def _invalidate_prefill_buffer(self, buffer_id: int) -> None:
         slot_start = buffer_id * self.num_experts
-        slot_end = slot_start + self.num_experts
-        old_ids = self.id_of_slot[slot_start:slot_end]
-        self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1
-        old_ids.fill_(-1)
-        # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in
-        # ensure_experts evicts them first.
-        self.usage[slot_start:slot_end].zero_()
+        # One fixed-shape launch. The previous boolean-mask index
+        # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every
+        # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48
+        # layers the host waited, per layer, for ALL enqueued GPU work (the cached-context
+        # attention included), serializing the prefill-overlap pipeline.
+        from freetoken.kernel.triton.moe import invalidate_prefill_slots
+
+        invalidate_prefill_slots(
+            self.id_of_slot, self.slot_for_id, self.usage, slot_start, self.num_experts
+        )
 
     def begin_prefill(self) -> None:
         if not self.prefill_overlap:

tests/moe/test_offload.py

@@ -277,6 +277,50 @@ def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots():
     assert torch.equal(cache.bank_caches["down"][:num_experts], down_source[0])
 
 
+def test_invalidate_prefill_slots_kernel_matches_eager_reference():
+    from freetoken.kernel.triton.moe import invalidate_prefill_slots
+
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the invalidation kernel")
+
+    for num_layers, num_experts, cache_size, slot_start in [
+        (3, 4, 8, 0),
+        (3, 4, 8, 4),
+        (2, 256, 1024, 512),
+    ]:
+        for trial in range(20):
+            torch.manual_seed(1234 + trial)
+            ids = torch.full((cache_size,), -1, dtype=torch.int32)
+            k = torch.randint(0, num_experts + 1, (1,)).item()
+            slots = torch.randperm(cache_size)[:k]
+            ids[slots] = torch.randint(0, num_layers * num_experts, (k,), dtype=torch.int32)
+            slot_for_id = torch.full((num_layers, num_experts), -1, dtype=torch.int32)
+            for slot, expert in enumerate(ids.tolist()):
+                if expert >= 0:
+                    slot_for_id.view(-1)[expert] = slot
+            usage = torch.randint(0, 999, (cache_size,), dtype=torch.int64)
+
+            ids_c, sfi_c, usage_c = ids.cuda(), slot_for_id.cuda(), usage.cuda()
+
+            ref_sfi = slot_for_id.clone()
+            ref_usage = usage.clone()
+            ref_ids = ids.clone()
+            buf = ref_ids[slot_start : slot_start + num_experts]
+            ref_sfi.view(-1)[buf[buf >= 0].long()] = -1
+            buf.fill_(-1)
+            ref_usage[slot_start : slot_start + num_experts].zero_()
+
+            invalidate_prefill_slots(ids_c, sfi_c, usage_c, slot_start, num_experts)
+            torch.cuda.synchronize()
+
+            assert torch.equal(ref_sfi, sfi_c.cpu())
+            assert torch.equal(ref_ids, ids_c.cpu())
+            assert torch.equal(ref_usage, usage_c.cpu())
+
+    with pytest.raises(ValueError):
+        invalidate_prefill_slots(ids.new_full((1,), -1), slot_for_id, usage.new_zeros(1), 0, 2)
+
+
 def test_prefill_overlap_waits_for_previous_prefill_release_after_begin(monkeypatch):
     from freetoken.moe.offload_cache import OffloadMoeCache