feat(moe): owner-local expert parallelism and tensor parallelism

benchmarks/bench_offload_cache_copy.py

@@ -47,6 +47,9 @@ class ModelProfile:
     # gpt-oss MXFP4 (block-32 e2m1 codes + e8m0 scales), H == I == 2880, top-4 routing
     "gpt-oss-20b": ModelProfile(24, 32, 4, "mxfp4_triton", 2880, 2880),
     "gpt-oss-120b": ModelProfile(36, 128, 4, "mxfp4_triton", 2880, 2880),
+    # Qwen3.8-Flash-Next-NVFP4: 48 MoE layers, 512 experts, top-10; H=2560, moe_inter=640.
+    # 6-bank triton NVFP4 -> 2,772,480 B/expert (matches the served cache/status unit_bytes).
+    "qwen3.8-flash-next": ModelProfile(48, 512, 10, "nvfp4", 2560, 640),
 }
 
 

python/freetoken/control_cli.py

@@ -230,6 +230,28 @@ def _format_stats(doc: dict[str, Any]) -> str:
         lines.append(f"mamba={mamba.get('used_slots', 0)}/{mamba.get('total_slots', 0)} slots")
     else:
         lines.append("mamba=none")
+    pcache = doc.get("prefix_cache")
+    if isinstance(pcache, dict):
+        lines.append(
+            f"prefix_cache cached_tokens={pcache.get('cached_tokens_total', 0)} "
+            f"hit_ratio={pcache.get('hit_ratio', 0) * 100:.1f}%"
+        )
+    moe = doc.get("moe")
+    if isinstance(moe, dict):
+        lines.append(
+            f"moe resident={moe.get('resident', 0)}/{moe.get('cache_size', 0)} slots "
+            f"(of {moe.get('total_experts', 0)} expert-layers) "
+            f"miss_rate={moe.get('miss_rate', 0) * 100:.1f}% "
+            f"miss/layer/step={moe.get('missing_per_layer', 0):.1f} "
+            f"fetched={moe.get('fetched_per_layer', 0):.1f}"
+        )
+        routing = moe.get("routing")
+        if isinstance(routing, dict):
+            lines.append(
+                f"moe routing working_set={routing.get('working_set_mean', 0):.1f} "
+                f"experts_for_90pct={routing.get('experts_for_90pct', 0):.1f} "
+                f"oracle_hit={routing.get('oracle_hit_global', 0) * 100:.1f}%"
+            )
     return "\n".join(lines)
 
 

python/freetoken/engine/config.py

@@ -21,6 +21,10 @@ class EngineConfig:
     model_path: str
     tp_info: DistributedInfo
     dtype: torch.dtype
+    # Opt-in routed-expert owner group size.  1 preserves the legacy global-ID cache;
+    # values >1 require an explicit TP+EP runtime implementation and fail fast unless the
+    # model/quantizer supports the owner-local bank path.
+    moe_ep_size: int = 1
     max_running_req: int = 4
     attention_backend: str = "auto"
     moe_strategy: str = "auto"
@@ -46,6 +50,17 @@ class EngineConfig:
     # (cudaMemcpyBatchAsync); no-op unless moe_cache_size > 2 * num_experts.
     moe_prefill_hit_d2d: bool = False
     moe_collect_stats: bool = False  # capture decode miss-rate counters into the cuda graph
+    # Per-(layer, expert) decode routing histogram (working-set / oracle-hit analysis).
+    # Accumulated on the DEVICE by a ``scatter_add_`` at the raw-ids point, so a captured
+    # decode graph replays it with every step -- this flag is CUDA-graph safe and does not
+    # require disabling graphs.
+    moe_collect_decode_freq: bool = False
+    # Ordered MoE route trace path (--moe-trace-route): when set, every
+    # ``ensure_experts`` call appends its RAW global expert ids (pre slot-rewrite)
+    # to this file for offline LRU/EP replay (moe/route_trace.py). Host-side, so it
+    # is NOT CUDA-graph safe -- the engine refuses it unless --cuda-graph-max-bs 0.
+    # None (default) = no recorder, zero overhead on the production path.
+    moe_trace_route: str | None = None
     # CPU MoE backend (--moe-strategy cpu): number of CPU worker threads computing
     # the decode experts. 0 = auto (physical cores). Ignored by other backends.
     moe_cpu_threads: int = 0

python/freetoken/kernel/pynccl.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import functools
+import os
+import sys
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
 from freetoken.env import ENV
@@ -27,7 +30,23 @@ def get_buffer(self) -> int: ...
 
 @functools.cache
 def _load_nccl_module() -> Module:
-    return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"])
+    # NVIDIA's pip/conda NCCL wheel intentionally ships only the SONAME file
+    # ``libnccl.so.2`` (no development symlink ``libnccl.so``), while a system NCCL
+    # development package normally provides the latter. ``-lnccl`` therefore works on
+    # the system layout but fails in zl_freetoken even when LD_LIBRARY_PATH is correct.
+    # Pass the wheel's absolute path when present; retain ``-lnccl`` for system installs.
+    candidates = []
+    env_nccl = os.environ.get("NCCL_LIB")
+    if env_nccl:
+        candidates.append(Path(env_nccl) / "libnccl.so.2")
+    candidates.extend(
+        Path(entry) / "nvidia" / "nccl" / "lib" / "libnccl.so.2"
+        for entry in sys.path
+        if entry
+    )
+    nccl = next((path for path in candidates if path.is_file()), None)
+    ldflags = [str(nccl)] if nccl is not None else ["-lnccl"]
+    return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=ldflags)
 
 
 @functools.cache

python/freetoken/kvcache/cache_status.py

@@ -194,6 +194,16 @@ def compute_cache_status_meta(engine: "Engine") -> Dict[str, Any]:
     meta["free_vram_bytes"] = _pool_budget_free_vram_bytes(engine)
     meta["floors"] = compute_cache_floors(engine)
     meta["pools"] = compute_cache_pools(engine)
+    # The context ceiling the SCHEDULER actually enforces: min(model max_position, KV pool
+    # tokens). /v1/models must report this rather than the checkpoint's own ceiling, otherwise
+    # a client sizes its window from the model card and then gets a hard 400
+    # (context_length_exceeded) on prompts the card said were fine -- which is exactly what
+    # happened whenever the KV pool was configured below the model's max_position.
+    # 0 when it could not be read (fake engines in tests); consumers then keep their fallback.
+    try:
+        meta["max_seq_len"] = int(engine.max_seq_len or 0)
+    except Exception:  # noqa: BLE001 -- best-effort; readiness must not depend on this
+        meta["max_seq_len"] = 0
     # Current window/full reuse ratio (the tunable knob), for DSV4 and radix-SWA; 0.0 otherwise.
     cfg = engine.config
     has_swa_ratio = cfg is not None and _supports_swa_ratio(cfg)

python/freetoken/layers/linear.py

@@ -147,9 +147,9 @@ def __init__(
             quant_config=quant_config, prefix=prefix,
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, *, reduce: bool = True) -> torch.Tensor:
         y = self.quant_method.apply(self, x)
-        if self._tp_size > 1:
+        if self._tp_size > 1 and reduce:
             y = self._comm.all_reduce(y)
         return y
 
@@ -174,8 +174,8 @@ def __init__(
             quant_config=quant_config, prefix=prefix,
         )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, *, reduce: bool = True) -> torch.Tensor:
         y = self.quant_method.apply(self, x)
-        if self._tp_size > 1:
+        if self._tp_size > 1 and reduce:
             y = self._comm.all_reduce(y)
         return y

python/freetoken/layers/moe.py

@@ -55,6 +55,7 @@ def __init__(
         layer_id: int | None = None,
         strategy: str = "resident",
         decode_target: str = "gpu",
+        expert_tp_size: int | None = None,
         quant_config: QuantConfig | None = None,
         prefix: str = "",
     ):
@@ -69,6 +70,10 @@ def __init__(
         tp_info = get_tp_info()
         self.tp_rank = tp_info.rank
         self.tp_size = tp_size = tp_info.size
+        # The routed-expert GEMM is tensor-parallel only for the resident/fused path; under
+        # owner-local EP every rank holds whole, disjoint experts, so the kernel sees an
+        # UNSHARDED intermediate and the layer all-reduces once at the output.
+        self.expert_tp_size = expert_tp_size if expert_tp_size is not None else tp_size
         self.renormalize = renormalize
         self.activation = activation
         self.apply_router_weight_on_input = apply_router_weight_on_input
@@ -132,14 +137,17 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         router_logits: torch.Tensor | None = None,
+        *,
+        reduce: bool = True,
     ):
         topk_weights, topk_ids = fused_topk(
             hidden_states=hidden_states,
             gating_output=router_logits,
             topk=self.top_k,
             renormalize=self.renormalize,
         )
-        return self._maybe_all_reduce(self._resident_gemm(hidden_states, topk_weights, topk_ids))
+        out = self._resident_gemm(hidden_states, topk_weights, topk_ids)
+        return self._maybe_all_reduce(out) if reduce else out
 
 
 class OffloadMoELayer(MoELayer):
@@ -161,6 +169,7 @@ def __init__(
         has_bias: bool = False,
         strategy: str = "offload",
         decode_target: str = "gpu",
+        expert_tp_size: int | None = None,
         quant_config: QuantConfig | None = None,
         prefix: str = "",
     ):
@@ -181,28 +190,36 @@ def __init__(
             layer_id=layer_id,
             strategy=strategy,
             decode_target=decode_target,
+            expert_tp_size=expert_tp_size,
             quant_config=quant_config,
             prefix=prefix,
         )
         self.offload_cache: OffloadMoeCache | None = None
+        # Owner-local EP cache (``OwnerOffloadMoeCache``). When attached, decode routes through
+        # ``ensure_route`` and the local-row slot ids, never the global-ID ``ensure_experts``.
+        self.owner_cache = None
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         router_logits: torch.Tensor | None = None,
+        *,
+        reduce: bool = True,
     ):
         ctx = get_global_ctx()
         if ctx.batch.is_prefill:
             final_hidden_states = self.prefill_forward(hidden_states, router_logits)
         else:
             final_hidden_states = self.decode_forward(hidden_states, router_logits)
-        return self._maybe_all_reduce(final_hidden_states)
+        return self._maybe_all_reduce(final_hidden_states) if reduce else final_hidden_states
 
     def routed_forward(
         self,
         hidden_states: torch.Tensor,
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
+        *,
+        reduce: bool = True,
     ) -> torch.Tensor:
         """Expert compute for an externally computed routing decision (``TopK``).
 
@@ -216,7 +233,7 @@ def routed_forward(
             out = self._prefill_routed(hidden_states, topk_weights, topk_ids)
         else:
             out = self._decode_routed(hidden_states, topk_weights, topk_ids)
-        return self._maybe_all_reduce(out)
+        return self._maybe_all_reduce(out) if reduce else out
 
     def decode_forward(
         self,
@@ -270,6 +287,8 @@ def _decode_routed(
         ids), so no ``ensure_experts``/``copy_missing`` here."""
         cache = self.offload_cache
         assert cache is not None
+        if self.owner_cache is not None:
+            return self._decode_owner(hidden_states, topk_weights, topk_ids)
         if cache.is_cpu_layer(self.layer_id):
             executor = cache.cpu_executor
             assert executor is not None, "CPU MoE executor was not initialized"
@@ -289,6 +308,49 @@ def _decode_routed(
             is_prefill=False,
         )
 
+    def _decode_owner(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Owner-local EP decode: global route -> local bank row -> owner cache slot -> GEMM.
+
+        The wrapped cache only ever sees ``local_num_experts`` rows, so the route MUST go
+        through the owner admission path: it masks remote entries to zero weight, rewrites
+        owned entries to owner-local slot ids, and returns tensors the existing
+        ``_expert_gemm`` already accepts unchanged (decode kernels index slots directly).
+        Passing the raw global ids here would read another rank's bank rows -- the kernels
+        do not range-check.
+
+        Two implementations, selected by ``owner.graph_safe``: the sync-free fixed-shape
+        ``ensure_route_graph`` under CUDA-graph capture, else the eager ``ensure_route``.
+        Both admit the same rows and mask remote entries identically.
+        """
+        owner = self.owner_cache
+        if owner.graph_safe:
+            update = owner.ensure_route_graph(self.layer_id, topk_weights, topk_ids)
+        else:
+            update = owner.ensure_route(self.layer_id, topk_weights, topk_ids)
+        owner.copy_missing()
+        out = self._expert_gemm(
+            owner,
+            hidden_states,
+            update.weights,
+            update.slot_ids,
+            views=owner.bank_views(),
+            n=None,
+            alphas=owner.alphas_for_slots(self.layer_id),
+            is_prefill=False,
+        )
+        if __debug__ and not owner.graph_safe and update.slot_ids.numel():
+            # Guardrail: a slot id outside the local pool would be an unchecked OOB read.
+            # Skipped when graph_safe: this read is a device->host sync, illegal in capture.
+            assert int(update.slot_ids.max()) < owner.cache_size, (
+                "owner-local slot id escaped the local pool"
+            )
+        return out
+
     def _decode_hybrid(
         self,
         cache: OffloadMoeCache,
@@ -350,6 +412,8 @@ def _prefill_routed(
         pass through unmapped."""
         cache = self.offload_cache
         assert cache is not None
+        if self.owner_cache is not None:
+            return self._prefill_owner(hidden_states, topk_weights, topk_ids)
         if cache.prefill_overlap:
             views = self._wait_prefill_overlap(cache)
             out = self._expert_gemm(
@@ -377,6 +441,54 @@ def _prefill_routed(
             is_prefill=True,
         )
 
+    def _prefill_owner(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Owner-local EP prefill: materialize the local layer, then route into local rows.
+
+        Prefill kernels take bank-row ids (``position == expert id`` for the global cache);
+        for the owner wrapper that namespace is the LOCAL row, so the global route is
+        remapped here and remote entries are zero-weighted. The global fused buffer is never
+        used, which is why this cannot reuse ``_prefill_routed``.
+        """
+        owner = self.owner_cache
+        if owner.geometry.prefill_overlap:
+            # Same begin -> prefetch(current) -> prefetch(next) -> wait -> release
+            # choreography as the global-ID path (_wait_prefill_overlap): the NEXT layer's
+            # H2D runs on the copy stream while THIS layer's GEMMs run on the compute
+            # stream.  Prefetching only the current layer would serialize copy and compute
+            # and lose the whole point of overlap.  prefetch_prefill_layer is a no-op past
+            # the last layer, so the lookahead needs no bounds check here.
+            if self.layer_id == 0:
+                owner.begin_prefill()
+            owner.prefetch_prefill_layer(self.layer_id)
+            owner.prefetch_prefill_layer(self.layer_id + 1)
+            views = owner.wait_prefill_layer(self.layer_id)
+        else:
+            owner.materialize_layer(self.layer_id, buffer_id=0)
+            views = owner.bank_views(owner.num_experts)
+        local_ids, owned = owner.geometry.global_to_local(topk_ids)
+        safe_ids = torch.where(owned, local_ids, torch.zeros_like(local_ids)).contiguous()
+        safe_weights = torch.where(
+            owned, topk_weights, torch.zeros_like(topk_weights)
+        ).contiguous()
+        out = self._expert_gemm(
+            owner,
+            hidden_states,
+            safe_weights,
+            safe_ids,
+            views=views,
+            n=owner.num_experts,
+            alphas=owner.alphas_for_layer(self.layer_id),
+            is_prefill=True,
+        )
+        if owner.geometry.prefill_overlap:
+            owner.release_prefill_layer(self.layer_id)
+        return out
+
     def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]:
         """Double-buffer choreography for this layer's overlap prefill: kick off the
         next layer's full-layer H2D copy, then return this layer's bank views (in
@@ -483,4 +595,7 @@ def make_moe_layer(
         kwargs["layer_id"] = layer_id
         kwargs["strategy"] = config.moe_strategy
         kwargs["decode_target"] = config.decode_target
+    if getattr(config, "moe_ep_size", 1) > 1:
+        # owner-local EP: whole disjoint experts per rank -> the expert GEMM is not sharded
+        kwargs["expert_tp_size"] = 1
     return layer_cls(**kwargs)

python/freetoken/layers/quantization/moe/base.py

@@ -35,13 +35,15 @@ class MoEConfig:
 
     @classmethod
     def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig":
+        # owner-local EP (TP+EP): experts are partitioned, never tensor-sharded, so the
+        # expert GEMM sees an unsharded intermediate (the layer all-reduces its output).
         return cls(
             num_experts=layer.num_experts,
             hidden=layer.hidden_size,
             intermediate=layer.intermediate_size,
             top_k=layer.top_k,
             tp_rank=layer.tp_rank,
-            tp_size=layer.tp_size,
+            tp_size=getattr(layer, "expert_tp_size", layer.tp_size),
             scheme=scheme,
             activation=layer.activation,
             alpha=float(layer.alpha),

python/freetoken/message/frontend.py

@@ -45,6 +45,10 @@ class UserReply(BaseFrontendMsg):
     swa_total_tokens: int = 0
     # Bytes the engine process holds on the GPU (torch reserved pool). 0 when not reported.
     gpu_mem_bytes: int = 0
+    # MoE slot-cache snapshot (miss_rate / residency / routing concentration), passed
+    # through from DetokenizeMsg on the scheduler's throttled interval. None when the
+    # model has no offload cache or the sample hasn't arrived yet.
+    moe_stats: dict | None = None
     # Set (with finished=True) when a request failed before producing output — e.g. a chat
     # template that the tokenizer cannot render, or a prompt that exceeds the KV budget the
     # scheduler can serve. Carries a human-readable reason. Without this, such a request would