diff --git a/python/freetoken/kernel/triton/moe/__init__.py b/python/freetoken/kernel/triton/moe/__init__.py new file mode 100644 index 0000000000..3f620fec6e --- /dev/null +++ b/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"] diff --git a/python/freetoken/kernel/triton/moe/invalidate.py b/python/freetoken/kernel/triton/moe/invalidate.py new file mode 100644 index 0000000000..24c5fbf3f9 --- /dev/null +++ b/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 + ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 719ebcbe52..a369e4b66b 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/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: diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc773..aaf3a4c097 100644 --- a/tests/moe/test_offload.py +++ b/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