Project-Chronos 0.1.1__py3-none-any.whl

chronos/router/prefill_scheduler.py

@@ -0,0 +1,248 @@
+"""
+chronos/router/prefill_scheduler.py
+
+Orchestrates the dual-layer pipeline:
+
+  Prefill phase:
+    1. Run IntentClassifier on full prompt → IntentVector
+    2. ExpertPredictor converts IntentVector → ExpertSet
+    3. Batch-load all predicted experts to VRAM (or Metal) BEFORE
+       generation starts — so generation sees zero IO latency
+
+  Generation phase:
+    4. Pass per-layer avail masks to model.forward()
+    5. After generation, update ExpertPredictor frequency table
+
+This is the key architectural difference from per-token LookaheadRouter:
+IO is fully front-loaded into prefill, not spread across generation steps.
+
+Usage (with IntentClassifier trained):
+    scheduler = PrefillScheduler(
+        classifier=intent_clf,
+        predictor=expert_pred,
+        expert_store=store,           # ExpertStore or MLXExpertStore
+    )
+
+    # Called once before generation
+    expert_set = scheduler.prepare(prompt_ids, device)
+
+    # Get masks to pass to model.forward()
+    masks = scheduler.avail_masks(device)
+
+    # After generation
+    scheduler.record_activation(activated_expert_ids)
+
+Usage (without IntentClassifier — fallback mode):
+    scheduler = PrefillScheduler(
+        classifier=None,              # No classifier
+        predictor=expert_pred,
+        expert_store=store,
+    )
+    # Falls back to frequency-based heuristic selection
+"""
+import time
+import threading
+from typing import List, Optional, Set, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from chronos.router.expert_predictor import ExpertSet
+
+
+class PrefillScheduler:
+    """
+    Dual-layer prefill: intent classify → expert predict → batch preload.
+
+    classifier: IntentClassifier or None (disables Layer 1, uses heuristic)
+    predictor:  ExpertPredictor (always required)
+    expert_store: ExpertStore (PyTorch) or MLXExpertStore
+    async_load: if True, preloads in a background thread while the main
+                model is still running the LM head on the prompt tokens
+    """
+
+    def __init__(
+        self,
+        classifier,           # IntentClassifier | None
+        predictor,            # ExpertPredictor
+        expert_store,
+        async_load: bool = True,
+    ):
+        self.classifier = classifier
+        self.predictor = predictor
+        self.store = expert_store
+        self.async_load = async_load
+
+        self._current_expert_set = None
+        self._load_thread: Optional[threading.Thread] = None
+        self._load_done = threading.Event()
+
+    # ── Public API ────────────────────────────────────────────────
+
+    def prepare(
+        self,
+        prompt_ids: torch.Tensor,    # [1, S] token ids
+        device: str = "cpu",
+        classifier_device: str = "cpu",
+    ) -> "ExpertSet":
+        """
+        Run prefill scheduling. Call this BEFORE model.generate().
+
+        1. Run IntentClassifier (or use heuristic if classifier=None)
+        2. Determine ExpertSet
+        3. Kick off background loading (or synchronous if async=False)
+
+        Returns ExpertSet immediately (loading may still be in progress
+        if async_load=True — call wait() before first forward pass).
+        """
+        t0 = time.monotonic()
+
+        if self.classifier is not None:
+            self.classifier.eval()
+            with torch.no_grad():
+                ids = prompt_ids.to(classifier_device)
+                intent = self.classifier(ids)
+            expert_set = self.predictor.predict(intent)
+        else:
+            # No classifier: build a synthetic IntentVector from frequency table
+            from chronos.router.intent_classifier import IntentVector
+            freq = self.predictor._frequency
+            per_layer = [freq.clone()
+                         for _ in range(self.predictor.num_moe_layers)]
+            intent = IntentVector(per_layer_probs=per_layer, global_probs=freq)
+            expert_set = self.predictor.predict(intent)
+
+        self._current_expert_set = expert_set
+        elapsed_classify = (time.monotonic() - t0) * 1000
+
+        # Trigger expert loading
+        self._load_done.clear()
+        if self.async_load:
+            self._load_thread = threading.Thread(
+                target=self._load_experts,
+                args=(expert_set.expert_ids,),
+                daemon=True,
+                name="prefill-loader",
+            )
+            self._load_thread.start()
+        else:
+            self._load_experts(expert_set.expert_ids)
+
+        mode = "async" if self.async_load else "sync"
+        print(
+            f"[PrefillScheduler] classify={elapsed_classify:.1f}ms  "
+            f"experts={sorted(expert_set.expert_ids)}  "
+            f"confidence={expert_set.confidence:.2f}  "
+            f"fallback={expert_set.fallback}  load={mode}"
+        )
+        return expert_set
+
+    def wait(self):
+        """
+        Block until background expert loading is complete.
+        Call this just before the first generation forward pass.
+        """
+        self._load_done.wait()
+        # Sync H2D stream on CUDA backends
+        if hasattr(self.store, "sync_h2d"):
+            self.store.sync_h2d()
+
+    def avail_masks(self, device=None) -> List[torch.Tensor]:
+        """
+        Return per-layer float availability masks to pass to model.forward().
+        Returns all-ones masks if no expert set has been computed yet
+        (safe default: treat all experts as available).
+        """
+        if self._current_expert_set is None:
+            num_layers = self.predictor.num_moe_layers
+            num_experts = self.predictor.num_experts
+            return [torch.ones(num_experts).to(device or "cpu")
+                    for _ in range(num_layers)]
+        return self.predictor.avail_masks_float(
+            self._current_expert_set, device=device
+        )
+
+    def record_activation(self, activated_expert_ids: List[int]):
+        """
+        Call after generation completes with the list of expert IDs
+        that were actually used. Updates the frequency heuristic.
+        """
+        self.predictor.update_frequency(activated_expert_ids)
+
+    # ── Internal ─────────────────────────────────────────────────
+
+    def _load_experts(self, expert_ids: Set[int]):
+        """
+        SSD → RAM → VRAM for the predicted expert set.
+        Runs in background thread when async_load=True.
+        """
+        try:
+            id_list = list(expert_ids)
+            # Step 1: SSD → pinned RAM (bulk)
+            self.store.prefetch_to_ram(id_list)
+            # Step 2: RAM → VRAM (promote each)
+            for eid in id_list:
+                self.store.promote_to_vram(eid)
+        finally:
+            self._load_done.set()
+
+    # ── Integration helpers ───────────────────────────────────────
+
+    @classmethod
+    def build(
+        cls,
+        main_model,
+        config,
+        expert_store,
+        classifier_path: Optional[str] = None,
+        device: str = "cpu",
+        async_load: bool = True,
+    ) -> "PrefillScheduler":
+        """
+        Factory: build a PrefillScheduler from a main model + config.
+        Loads IntentClassifier from disk if path provided; otherwise
+        creates an untrained one (will use heuristic fallback until trained).
+
+        classifier_path: path to .pt checkpoint, or None to skip classifier
+        """
+        from chronos.router.intent_classifier import IntentClassifier
+        from chronos.router.expert_predictor import ExpertPredictor
+
+        from chronos.model.moe_chronos import ChronosMOEFeedForward
+        num_moe_layers = sum(
+            1 for l in main_model.model.layers
+            if isinstance(l.mlp, ChronosMOEFeedForward)
+        )
+
+        predictor = ExpertPredictor(
+            num_experts=config.num_experts,
+            num_moe_layers=num_moe_layers,
+            vram_capacity=expert_store.vram_capacity
+            if hasattr(expert_store, "vram_capacity") else config.num_experts,
+        )
+
+        classifier = None
+        if classifier_path is not None and __import__("os").path.exists(classifier_path):
+            classifier = IntentClassifier.load(
+                classifier_path,
+                vocab_size=config.vocab_size,
+                num_experts=config.num_experts,
+                num_moe_layers=num_moe_layers,
+            ).to(device)
+            print(f"[PrefillScheduler] Loaded IntentClassifier from {classifier_path}")
+        elif classifier_path is not None:
+            # Path provided but file doesn't exist yet → create untrained
+            classifier = IntentClassifier(
+                vocab_size=config.vocab_size,
+                num_experts=config.num_experts,
+                num_moe_layers=num_moe_layers,
+            ).to(device)
+            print("[PrefillScheduler] IntentClassifier not found — "
+                  "using heuristic fallback (train with chronos train --mode intent)")
+
+        return cls(
+            classifier=classifier,
+            predictor=predictor,
+            expert_store=expert_store,
+            async_load=async_load,
+        )

chronos/runtime/__init__.py

@@ -0,0 +1,2 @@
+from .cache_manager import CacheManager
+from .inference_engine import ChronosInferenceEngine

chronos/runtime/cache_manager.py

@@ -0,0 +1,117 @@
+"""
+chronos/runtime/cache_manager.py
+
+Unified VRAM/RAM cache manager for Project Chronos inference.
+
+Wraps ExpertStore + AsyncPrefetcher into a single interface used by
+the inference engine. Provides:
+- availability_mask() for soft gating
+- step() to advance prefetch schedule after each decode step
+- stats() for monitoring
+"""
+from typing import List, Optional
+import os
+import torch
+
+from chronos.io.expert_store import ExpertStore
+from chronos.io.async_prefetcher import AsyncPrefetcher, PrefetchScheduler
+from chronos.io.storage import ClusterStorage, MANIFEST_FILENAME
+
+
+class CacheManager:
+    """
+    Single entry point for all expert caching operations during inference.
+
+    Lifecycle:
+        mgr = CacheManager(model, config, ssd_dir="./expert_cache")
+        mgr.start()
+        for each decode step:
+            mask = mgr.availability_mask()          # pass to model.forward
+            outputs, lp = model(x, available_expert_masks=[mask]*L)
+            mgr.step(lp, current_expert_ids)        # schedule next prefetch
+        mgr.stop()
+    """
+
+    def __init__(self, model, config, ssd_dir: str = "./expert_cache"):
+        self.config = config
+        self.expert_store = ExpertStore(model, config, ssd_dir=ssd_dir)
+        self.prefetcher = AsyncPrefetcher(self.expert_store)
+        self.scheduler = PrefetchScheduler(self.prefetcher, self.expert_store)
+        self._num_layers = len(self.expert_store.moe_layers)
+
+        # Auto-attach any cluster manifest present on disk. Also supports an
+        # explicit override via config.cluster_manifest_path pointing to a
+        # different layout (e.g. a post-calibration Louvain repack).
+        manifest_override = getattr(config, "cluster_manifest_path", None)
+        if manifest_override and ClusterStorage.has_manifest(manifest_override):
+            self.expert_store.attach_cluster_manifest(manifest_override)
+        elif ClusterStorage.has_manifest(ssd_dir):
+            self.expert_store.attach_cluster_manifest(ssd_dir)
+
+    def start(self):
+        """Start background prefetch thread."""
+        self.prefetcher.start()
+
+    def stop(self):
+        """Stop background prefetch thread."""
+        self.prefetcher.stop()
+
+    def warm_up(self, initial_expert_ids: Optional[List[int]] = None):
+        """
+        Pre-load a set of experts into VRAM before generation starts.
+        Defaults to experts 0..vram_capacity-1.
+        """
+        if initial_expert_ids is None:
+            initial_expert_ids = list(range(
+                min(self.config.num_experts, self.expert_store.vram_capacity)
+            ))
+        self.expert_store.prefetch_to_ram(initial_expert_ids)
+        for eid in initial_expert_ids:
+            self.expert_store.promote_to_vram(eid)
+        # Ensure all H2D copies are complete before model sees the weights
+        self.expert_store.sync_h2d()
+
+    def availability_mask(self) -> torch.Tensor:
+        """[num_experts] bool — True if expert is currently in VRAM."""
+        return self.expert_store.vram_availability_mask()
+
+    def availability_masks_all_layers(self) -> List[torch.Tensor]:
+        """Returns the same mask replicated for each MoE layer."""
+        mask = self.availability_mask()
+        return [mask] * self._num_layers
+
+    def step(self, lookahead_probs, current_expert_ids: List[int]):
+        """
+        Legacy single-call API (M1+M2 compatible): blocking promote of
+        current experts, then async submit of future predictions, then
+        full sync. Use the M3 split below to get pipeline overlap.
+        """
+        self.scheduler.step(lookahead_probs, current_expert_ids)
+        self.expert_store.sync_h2d()
+
+    # ── M3: split prefetch from blocking sync ────────────────────
+
+    def prefetch_for_next_step(self, lookahead_probs) -> List[int]:
+        """Fire off async prefetch for the LookaheadRouter's predictions.
+        Non-blocking. Should be called BEFORE the next forward so H2D and
+        compute overlap. Returns the submitted expert IDs."""
+        return self.scheduler.prefetch_only(lookahead_probs)
+
+    def ensure_resident(self, expert_ids: List[int]) -> List[int]:
+        """Promote `expert_ids` to VRAM (non-blocking H2D) and then have the
+        compute stream wait on *only those* events. Returns the list newly
+        promoted. After this call returns, the named experts are guaranteed
+        visible to the next forward without globally syncing other in-flight
+        H2D copies."""
+        newly = self.scheduler.promote_current(expert_ids, blocking=False)
+        # The compute stream waits on each newly-promoted expert's event;
+        # already-resident experts contribute no event, so this only blocks
+        # on the truly necessary transfers.
+        self.expert_store.wait_for_experts(newly)
+        return newly
+
+    def stats(self) -> dict:
+        return {
+            **self.expert_store.stats(),
+            **self.prefetcher.stats,
+        }

chronos/runtime/inference_engine.py

@@ -0,0 +1,295 @@
+"""
+chronos/runtime/inference_engine.py
+
+End-to-end Chronos inference engine with:
+- Async prefetch driven by LookaheadRouter
+- Soft gating fallback via CacheManager availability mask
+- Token/s throughput measurement
+- Optional streamer support
+
+Usage:
+    engine = ChronosInferenceEngine(model, config, ssd_dir="./expert_cache")
+    engine.setup()  # offload experts to SSD, warm up cache
+    output_ids = engine.generate(input_ids, max_new_tokens=200)
+    print(engine.last_stats)
+"""
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import time
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.moe_chronos import ChronosMOEFeedForward
+from chronos.runtime.cache_manager import CacheManager
+from chronos.router.prefill_scheduler import PrefillScheduler
+
+
+class ChronosInferenceEngine:
+    """
+    Wraps ChronosForCausalLM with the full async prefetch + soft gating pipeline.
+
+    The decode loop:
+    1. Get availability mask from CacheManager
+    2. Forward pass with mask → soft gating handles cache misses
+    3. Extract LookaheadRouter probs from output
+    4. Submit future expert predictions to prefetcher (non-blocking)
+    5. Repeat
+    """
+
+    def __init__(
+        self,
+        model: ChronosForCausalLM,
+        config: ChronosConfig,
+        ssd_dir: str = "./expert_cache",
+    ):
+        self.model = model
+        self.config = config
+        self.device = next(model.parameters()).device
+        self.cache_manager = CacheManager(model, config, ssd_dir=ssd_dir)
+        self.last_stats: dict = {}
+
+    def setup(self, warm_expert_ids: Optional[List[int]] = None):
+        """
+        One-time setup: serialize experts to SSD, start prefetcher, warm cache.
+        Call once before the first generate().
+        """
+        print("[Chronos] Offloading expert weights to SSD...")
+        self.cache_manager.expert_store.offload_all_to_ssd()
+        self.cache_manager.start()
+        print("[Chronos] Warming up VRAM cache...")
+        self.cache_manager.warm_up(warm_expert_ids)
+        print(f"[Chronos] Ready. Cache stats: {self.cache_manager.stats()}")
+
+    def setup_from_state_dict(
+        self,
+        state_dict: dict,
+        warm_expert_ids: Optional[List[int]] = None,
+    ):
+        """Setup path for real lazy expert loading on CPU/MPS.
+
+        The caller loads base/non-expert weights into the live model, then
+        hands the full checkpoint here so ExpertStore can shard only expert
+        tensors to SSD and replace live experts with placeholders.
+        """
+        print("[Chronos] Offloading expert weights to SSD from checkpoint...")
+        self.cache_manager.expert_store.offload_from_state_dict(state_dict)
+        self.cache_manager.start()
+        print("[Chronos] Warming up VRAM cache...")
+        self.cache_manager.warm_up(warm_expert_ids)
+        print(f"[Chronos] Ready. Cache stats: {self.cache_manager.stats()}")
+
+    def teardown(self):
+        self.cache_manager.stop()
+
+    def _get_current_expert_ids(self) -> List[int]:
+        """Extract top-1 expert IDs from the last forward pass (layer 0)."""
+        moe_layers = [
+            l.mlp for l in self.model.model.layers
+            if isinstance(l.mlp, ChronosMOEFeedForward)
+        ]
+        if not moe_layers or moe_layers[0].last_router_probs is None:
+            return []
+        probs = moe_layers[0].last_router_probs  # [B, S, E]
+        return probs[:, -1, :].argmax(dim=-1).unique().cpu().tolist()
+
+    def _plan_prefill_experts(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> List[int]:
+        """Probe prompt routing with shared fallback, then load hot experts.
+
+        The live model contains placeholder experts after lazy setup. A cold
+        first prefill would therefore run entirely through shared fallback and
+        can diverge from the real model before decode starts. This probe pass
+        does not materialize expert modules; it only records gate choices, then
+        returns the most frequent expert IDs bounded by the active cache budget.
+        The caller reruns the real prefill after those experts are promoted.
+        """
+        if not getattr(self.config, "use_moe", True):
+            return []
+        num_experts = int(getattr(self.config, "num_experts", 0) or 0)
+        if num_experts <= 0:
+            return []
+
+        cold_mask = torch.zeros(num_experts, dtype=torch.bool, device=self.device)
+        masks = [cold_mask] * len(self.cache_manager.expert_store.moe_layers)
+        self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            past_key_values=None,
+            use_cache=False,
+            available_expert_masks=masks,
+        )
+
+        counts = torch.zeros(num_experts, dtype=torch.long)
+        top_k = int(getattr(self.config, "num_experts_per_tok", 1) or 1)
+        top_k = max(1, min(top_k, num_experts))
+        for layer in self.model.model.layers:
+            moe = getattr(layer, "mlp", None)
+            if not isinstance(moe, ChronosMOEFeedForward) or moe.last_router_probs is None:
+                continue
+            probs = moe.last_router_probs.reshape(-1, num_experts)
+            expert_ids = torch.topk(probs, k=top_k, dim=-1).indices.reshape(-1).cpu()
+            counts += torch.bincount(expert_ids, minlength=num_experts)
+
+        if not torch.any(counts):
+            return []
+        capacity = max(1, min(
+            int(getattr(self.cache_manager.expert_store, "vram_capacity", 1) or 1),
+            num_experts,
+        ))
+        ordered = torch.argsort(counts, descending=True).tolist()
+        return [int(eid) for eid in ordered if int(counts[eid]) > 0][:capacity]
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 256,
+        temperature: float = 0.85,
+        top_p: float = 0.85,
+        top_k: int = 50,
+        eos_token_id: int = 2,
+        streamer=None,
+        do_sample: bool = True,
+        repetition_penalty: float = 1.0,
+        scheduler: Optional[PrefillScheduler] = None,
+    ) -> torch.Tensor:
+        input_ids = input_ids.to(self.device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(self.device)
+
+        # ── Prefill phase: front-load expert IO before generation ──────────
+        if scheduler is not None:
+            scheduler.prepare(input_ids, device=str(self.device))
+            scheduler.wait()  # block until SSD→VRAM load complete
+            prefill_expert_ids: List[int] = []
+        else:
+            prefill_expert_ids = self._plan_prefill_experts(input_ids, attention_mask)
+            if prefill_expert_ids:
+                self.cache_manager.ensure_resident(prefill_expert_ids)
+
+        past_key_values = None
+        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=self.device)
+        if streamer:
+            streamer.put(input_ids.cpu())
+
+        t0 = time.monotonic()
+        tokens_generated = 0
+        cache_hits, cache_misses = 0, 0
+        activated_expert_ids: List[int] = []
+        # M3: holds last step's lookahead predictions; we submit them as
+        # prefetch BEFORE this step's forward so the H2D stream overlaps
+        # with compute. None on the first iteration (no predictions yet).
+        prev_lookahead_probs = None
+        prev_current_ids: List[int] = []
+
+        for step in range(max_new_tokens):
+            # ── M3 pipeline: fire async prefetch BEFORE forward ──
+            if scheduler is None:
+                if prev_lookahead_probs is not None:
+                    # Submit next-token predictions to the async SSD→RAM thread.
+                    self.cache_manager.prefetch_for_next_step(prev_lookahead_probs)
+                if prev_current_ids:
+                    # Ensure the experts we actually need NOW are resident in
+                    # VRAM, waiting only on their own H2D events.
+                    self.cache_manager.ensure_resident(prev_current_ids)
+
+            # Get availability masks AFTER ensure_resident so the mask
+            # reflects experts that just became available.
+            if scheduler is not None:
+                masks = scheduler.avail_masks(device=str(self.device))
+            else:
+                masks = self.cache_manager.availability_masks_all_layers()
+
+            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
+            outputs, lookahead_probs = self.model(
+                input_ids[:, past_len:],
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                use_cache=True,
+                available_expert_masks=masks,
+            )
+
+            if attention_mask is not None:
+                attention_mask = torch.cat(
+                    [attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1
+                )
+
+            # Track cache hit/miss from mask
+            mask_tensor = masks[0] if masks else None
+            current_ids = self._get_current_expert_ids()
+            if mask_tensor is not None:
+                for eid in current_ids:
+                    if eid < len(mask_tensor) and mask_tensor[eid]:
+                        cache_hits += 1
+                    else:
+                        cache_misses += 1
+            activated_expert_ids.extend(current_ids)
+
+            # Stash predictions for the NEXT iteration's prefetch.
+            prev_lookahead_probs = lookahead_probs
+            prev_current_ids = current_ids
+
+            # Sampling
+            logits = outputs.logits[:, -1, :] / temperature
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+            if top_k > 0:
+                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                mask_p = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                mask_p[..., 1:] = mask_p[..., :-1].clone()
+                mask_p[..., 0] = 0
+                logits[mask_p.scatter(1, sorted_indices, mask_p)] = -float('inf')
+
+            next_token = (
+                torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+                if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+            )
+            if eos_token_id is not None:
+                next_token = torch.where(
+                    finished.unsqueeze(-1),
+                    next_token.new_full((next_token.shape[0], 1), eos_token_id),
+                    next_token,
+                )
+
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            past_key_values = outputs.past_key_values
+            tokens_generated += 1
+
+            if streamer:
+                streamer.put(next_token.cpu())
+
+            if eos_token_id is not None:
+                finished |= next_token.squeeze(-1).eq(eos_token_id)
+                if finished.all():
+                    break
+
+        # ── Post-generation: update frequency heuristic ───────────────────
+        if scheduler is not None:
+            scheduler.record_activation(list(set(activated_expert_ids)))
+
+        elapsed = time.monotonic() - t0
+        total_cache = cache_hits + cache_misses
+        self.last_stats = {
+            "tokens_generated": tokens_generated,
+            "elapsed_s": round(elapsed, 3),
+            "tokens_per_sec": round(tokens_generated / max(elapsed, 1e-6), 2),
+            "cache_hit_rate": round(cache_hits / max(total_cache, 1), 4),
+            "prefill_expert_ids": prefill_expert_ids,
+            **self.cache_manager.stats(),
+        }
+
+        if streamer:
+            streamer.end()
+        return input_ids