diff --git a/benchmarks/bench_offload_cache_copy.py b/benchmarks/bench_offload_cache_copy.py index 8374510ff..4eaaa9b7c 100644 --- a/benchmarks/bench_offload_cache_copy.py +++ b/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), } diff --git a/python/freetoken/control_cli.py b/python/freetoken/control_cli.py index 50a32aae6..bf25b6cac 100644 --- a/python/freetoken/control_cli.py +++ b/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) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 7ab792f9e..5be9d882d 100644 --- a/python/freetoken/engine/config.py +++ b/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 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 006c19089..86a52beac 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -18,7 +18,13 @@ from freetoken.moe import is_offload_moe_strategy from freetoken.moe.expert_banks import load_expert_banks from freetoken.moe.host_banks import PinFailed -from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache +from freetoken.moe.offload_cache import ( + OffloadMoeCache, + OwnerOffloadMoeCache, + attach_offload_moe_cache, + attach_owner_moe_cache, +) +from freetoken.moe.ownership import ExpertOwnership, OwnerCacheGeometry from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig @@ -47,6 +53,66 @@ def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: ) +def _owner_ep_enabled(config: EngineConfig) -> bool: + return config.moe_ep_size > 1 + + +def _validate_owner_ep_config(config: EngineConfig) -> None: + """Fail before model/bank allocation unless the initial owner topology is explicit.""" + if config.moe_ep_size == 1: + return + if config.moe_ep_size != config.tp_info.size or config.tp_info.size != 2: + raise ValueError( + "owner EP currently requires the initial same-group TP2+EP2 topology " + "(--tensor-parallel-size 2 --moe-ep-size 2)" + ) + if config.moe_strategy != "offload": + raise ValueError("owner EP currently requires --moe-strategy offload") + if config.moe_cache_rate is not None: + raise ValueError( + "owner EP sizes a LOCAL pool, so --moe-cache-rate (a fraction of the GLOBAL " + "expert count) has no owner-local meaning; use --moe-cache-size, or " + "--moe-cache-auto to fill whatever the KV pool leaves" + ) + if config.moe_cache_size <= 0 and not config.moe_cache_auto: + raise ValueError( + "owner EP needs an explicit --moe-cache-size, or --moe-cache-auto (which now " + "solves against the owner-local expert geometry)" + ) + if config.moe_cpu_layers: + raise ValueError( + "owner EP does not implement the CPU/hybrid expert path: the owner cache wraps " + "the GPU slot cache and _decode_owner is selected before the is_cpu_layer " + "branch, so --moe-cpu-layers would be accepted and then silently ignored" + ) + from freetoken.checkpoint.ftw import is_ftw_checkpoint + + if is_ftw_checkpoint(config.model_path): + raise ValueError( + "owner EP is not supported for FTW checkpoints: load_ftw_banks rebuilds " + "[num_experts, ...] GLOBAL expert rows with no ownership filter, so the banks " + "cannot bind to the owner-local geometry" + ) + # CUDA graphs are allowed: decode admission goes through the fixed-shape, sync-free + # OwnerOffloadMoeCache.ensure_route_graph when graphs are on (see _owner_graph_safe). + + +def _owner_graph_safe(config: EngineConfig) -> bool: + """Whether the owner decode route uses the fixed-shape graph-safe admission. + + Default (``auto``): follow the resolved CUDA-graph setting -- capture is only possible on + the sync-free path, so the two always agree. ``FREETOKEN_OWNER_GRAPH_SAFE=1/0`` forces + the admission implementation independently, which is what lets the graph-safe route be + A/B'd eagerly (graphs off) before trusting it inside a capture. + """ + forced = os.getenv("FREETOKEN_OWNER_GRAPH_SAFE", "auto").strip().lower() + if forced in ("0", "false", "no", "off"): + return False + if forced in ("1", "true", "yes", "on"): + return True + return bool(config.cuda_graph_max_bs) + + def _flashinfer_available() -> bool: from freetoken.kernel.backend import is_flashinfer_installed @@ -296,6 +362,7 @@ class ForwardOutput(NamedTuple): class Engine: def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() + _validate_owner_ep_config(config) set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend))) _ensure_expandable_segments() # before the first CUDA allocation below @@ -430,6 +497,11 @@ def __init__(self, config: EngineConfig): dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) + # NOTE: ``--moe-collect-decode-freq`` is CUDA-graph safe. The histogram lives on the + # device (``OffloadMoeCache.decode_freq``) and is accumulated by a device-side + # ``scatter_add_`` at the raw-ids point, so a captured decode graph replays the + # accumulation with every step. No warning is needed and graphs must not be disabled + # for it. if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() @@ -474,23 +546,48 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso config.model_path, self.device, include_moe_experts=not is_offload_moe_strategy(config.moe_strategy), + tp_shard=config.tp_info.size > 1, ), device=self.device, ) - - def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) -> tuple[int, int, bool]: + def _resolve_auto_moe_cache_size( + self, config: EngineConfig, banks, method=None, + ownership: ExpertOwnership | None = None, + ) -> tuple[int, int, bool]: """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap). Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable without a GPU. Reused by the Phase-2 runtime rebuild. + + ``ownership`` switches the expert geometry to the OWNER-LOCAL namespace: under EP2 a + rank's pool only ever holds its own rows, so the floor/cap and the coverage the plan + is solved against are ``local_num_experts`` and ``num_layers * local_num_experts``. + Solving against the global counts would under-fill (the cap ``total_experts`` is 2x + too large and the floor is wrong), which is why owner EP used to demand an explicit + ``--moe-cache-size``. + + A fixed KV pool (``--num-tokens``) is reserved EXACTLY, not at the + ``--kv-reserve-tokens`` floor: the caller has already pinned the KV geometry, so + every remaining byte belongs to the expert cache -- that is the whole point of + "pin KV, let MoE fill the rest". """ from freetoken.engine.cache_budget import expert_bytes_per_slot, resolve_moe_cache_auto cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config) fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed - num_experts = config.model_config.num_experts - total_experts = config.model_config.num_moe_layers * num_experts + if ownership is None: + num_experts = config.model_config.num_experts + total_experts = config.model_config.num_moe_layers * num_experts + else: + num_experts = ownership.local_num_experts + total_experts = config.model_config.num_moe_layers * num_experts + # getattr: duck-typed test configs may predate the --num-tokens knob + num_token_override = getattr(config, "num_token_override", None) + if num_token_override is not None: + kv_reserve_tokens = num_token_override + else: + kv_reserve_tokens = max(config.kv_reserve_tokens, min_reserve) return resolve_moe_cache_auto( baseline_free=self._baseline_free, weights_bytes=self._weights_bytes, @@ -501,12 +598,29 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) num_experts=num_experts, total_experts=total_experts, prefill_overlap=config.moe_prefill_overlap, - kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), + kv_reserve_tokens=kv_reserve_tokens, page_size=page_tokens, max_slots=method.slot_limit() if method is not None else None, ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: + owner_ep = _owner_ep_enabled(config) + ownership = None + owner_geometry = None + if owner_ep: + if config.model_config.model_type != "qwen4_exp": + raise NotImplementedError("owner EP is currently implemented only for Qwen4Exp") + if config.model_config.expert_quant != "nvfp4": + raise NotImplementedError("owner EP currently requires native NVFP4 expert banks") + ownership = ExpertOwnership( + global_num_experts=config.model_config.num_experts, + world_size=config.moe_ep_size, + rank=config.tp_info.rank, + ) + # owner_geometry is deliberately built LATER, once --moe-cache-auto has resolved + # the slot count: the geometry validates cache_size against the local expert + # count, so constructing it here with moe_cache_size == 0 (auto) would reject a + # perfectly valid request. method = shared_offload_method(self.model) num_moe_layers = config.model_config.num_moe_layers cpu_layer_ids = _resolve_cpu_layers(config, num_moe_layers, reserved=self._host_tables_bytes, method=method) @@ -570,11 +684,14 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: parallel=expert_parallel, decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"), layer_residency=requested_residency, + ownership=ownership, ) except PinFailed as exc: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc if config.moe_cache_auto: - size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + size, pages, overlap = self._resolve_auto_moe_cache_size( + config, banks, method, ownership=ownership + ) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: @@ -590,7 +707,28 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"--moe-cache-auto resolved moe_cache_size={size} " f"num_pages={pages} (prefill_overlap={overlap})" ) - _require_offload_cache_size(config.moe_cache_size, config.model_config.num_experts) + if ownership is not None: + # Built here (not at the top) so an auto-sized moe_cache_size is already in + # config. Validates the local floor and the 2*local overlap minimum before + # any allocation. + owner_geometry = OwnerCacheGeometry( + global_num_experts=config.model_config.num_experts, + world_size=config.moe_ep_size, + rank=config.tp_info.rank, + num_layers=config.model_config.num_moe_layers, + cache_size=config.moe_cache_size, + prefill_overlap=config.moe_prefill_overlap, + ) + if config.moe_prefill_overlap: + logger.info_rank0( + f"owner EP prefill overlap enabled: slots " + f"[0, {2 * ownership.local_num_experts}) of {config.moe_cache_size} " + f"are borrowed as the two-layer prefill buffer" + ) + _require_offload_cache_size( + config.moe_cache_size, + ownership.local_num_experts if ownership is not None else config.model_config.num_experts, + ) layout = max_slots = None if method is not None: if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name): @@ -600,22 +738,34 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: ) layout = method.layout() max_slots = method.slot_limit() - cache = OffloadMoeCache( - # Models with leading dense layers (GLM-4) only have experts on the MoE - # layers; num_moe_layers == num_layers when first_k_dense_replace == 0. - num_layers=config.model_config.num_moe_layers, - num_experts=config.model_config.num_experts, - cache_size=config.moe_cache_size, - device=self.device, - cache_policy=config.moe_cache_policy, - prefill_overlap=config.moe_prefill_overlap, - prefill_hit_d2d=config.moe_prefill_hit_d2d, - quant_format=banks.quant_format, - decode_target=decode_target, - hybrid_max_fetch=config.moe_hybrid_max_fetch, - layout=layout, - max_slots=max_slots, - ) + if owner_geometry is not None: + cache = OwnerOffloadMoeCache( + owner_geometry, + self.device, + cache_policy=config.moe_cache_policy, + prefill_hit_d2d=config.moe_prefill_hit_d2d, + quant_format=banks.quant_format, + graph_safe=_owner_graph_safe(config), + layout=layout, + max_slots=max_slots, + ) + else: + cache = OffloadMoeCache( + # Models with leading dense layers (GLM-4) only have experts on the MoE + # layers; num_moe_layers == num_layers when first_k_dense_replace == 0. + num_layers=config.model_config.num_moe_layers, + num_experts=config.model_config.num_experts, + cache_size=config.moe_cache_size, + device=self.device, + cache_policy=config.moe_cache_policy, + prefill_overlap=config.moe_prefill_overlap, + prefill_hit_d2d=config.moe_prefill_hit_d2d, + quant_format=banks.quant_format, + decode_target=decode_target, + hybrid_max_fetch=config.moe_hybrid_max_fetch, + layout=layout, + max_slots=max_slots, + ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency) @@ -625,7 +775,43 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # Must be set before CUDA graph capture so the (device-side) accumulation ops are # captured and re-run on every decode replay. cache.collect_stats = config.moe_collect_stats - layers = attach_offload_moe_cache(self.model, cache) + cache.collect_decode_freq = config.moe_collect_decode_freq + # Opt-in ordered route trace (moe/route_trace.py): records raw global expert + # ids per ensure_experts call for offline LRU/EP replay. Host-side (.cpu()), + # so it would break CUDA-graph capture -- refuse unless graphs are off, and do + # it here (before GraphRunner runs) so it fails fast, not mid-capture. + if config.moe_trace_route: + if config.cuda_graph_max_bs != 0: + raise ValueError( + "--moe-trace-route records host-side per call and is NOT CUDA-graph " + "safe; relaunch with --cuda-graph-max-bs 0 (diagnostic sampling only)." + ) + from freetoken.moe.route_trace import RouteTraceRecorder + + cache.route_recorder = RouteTraceRecorder( + config.moe_trace_route, + num_experts=( + cache.global_num_experts + if owner_geometry is not None + else cache.num_experts + ), + num_layers=cache.num_layers, + cache_size=cache.cache_size, + top_k=config.model_config.num_experts_per_tok, + model=config.model_path, + decode_target=cache.decode_target, + # One file per rank: every rank records the same configured path, so sharing + # it would let the TP writers truncate each other's trace. + rank=config.tp_info.rank if config.tp_info.size > 1 else None, + ) + logger.info_rank0( + f"--moe-trace-route: recording ordered decode route trace to " + f"{config.moe_trace_route} (graphs off; replay with " + f"tools/trace/replay_route_trace.py)" + ) + # attach_offload_moe_cache walks the model for OffloadMoELayers (model-specific + # subclasses included, e.g. DSV4's DSV4OffloadMoELayer). + layers = attach_owner_moe_cache(self.model, cache) if owner_geometry is not None else attach_offload_moe_cache(self.model, cache) assert len(layers) == config.model_config.num_moe_layers if cache.decode_target in ("cpu", "hybrid"): self._init_cpu_moe_executor(config, cache, layers) @@ -986,6 +1172,9 @@ def _warmup_prefill(self) -> None: ) def shutdown(self) -> None: + rec = getattr(self, "moe_offload_cache", None) + if rec is not None and getattr(rec, "route_recorder", None) is not None: + rec.route_recorder.close() # flush the ordered route trace (meta + body) self.graph_runner.destroy_cuda_graphs() torch.distributed.destroy_process_group() destroy_distributed() @@ -1571,6 +1760,9 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if is_moe: object.__setattr__(model_config, "moe_strategy", config.moe_strategy) object.__setattr__(model_config, "decode_target", _decode_target(config)) + # owner-local EP is decided by the engine, but the model builds its MoE layers + # first: expose the group size so they report an unsharded expert GEMM. + object.__setattr__(model_config, "moe_ep_size", getattr(config, "moe_ep_size", 1)) # Must stay LAST: page_size is only final here (_adjust_dsv4_config sets P=128, the # TRTLLM block sets 64). Also covers the programmatic LLM(...) path that bypasses parse_args. diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea57351..474deda3d 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/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 diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 10169d651..a16816343 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/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) diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index d8a3e8619..2fa039bbd 100644 --- a/python/freetoken/layers/linear.py +++ b/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 diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 42ad822b4..db268a797 100644 --- a/python/freetoken/layers/moe.py +++ b/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,6 +137,8 @@ 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, @@ -139,7 +146,8 @@ def forward( 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) diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index 89f13ed09..8768f9d82 100644 --- a/python/freetoken/layers/quantization/moe/base.py +++ b/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), diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b..bdea47ded 100644 --- a/python/freetoken/message/frontend.py +++ b/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 diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c785..5c7e78a7c 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -48,6 +48,10 @@ class DetokenizeMsg(BaseTokenizerMsg): swa_total_tokens: int = 0 # Bytes this engine process holds on the GPU (torch reserved pool). 0 on CPU. gpu_mem_bytes: int = 0 + # Throttled MoE slot-cache snapshot (miss/eviction/residency counters, see + # OffloadMoeCache.stats_snapshot) for /v1/stats. None on non-offload models or + # between sample intervals (the frontend keeps the last-known value). + moe_stats: dict | None = None @dataclass diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 07a16df5c..10f952f0f 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -263,6 +263,10 @@ class ModelConfig: moe_strategy: str = "fused" # where routed experts decode (gpu / cpu / hybrid); set by the engine from the flags, it gates which expert kernels can serve decode_target: str = "gpu" + # Expert-parallel group size (owner-local EP). 1 keeps the global-ID cache; >1 partitions + # the experts across the group, so the routed-expert GEMM is unsharded and the layer + # all-reduces its output once. Set by the engine before the model is built. + moe_ep_size: int = 1 # The QuantConfig the engine builds from the checkpoint; layers ask it for their method. quant: Any | None = None # ----- optional, model-specific extensions (default keeps other models intact) ----- diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 6b6245674..efe53c398 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -100,6 +100,11 @@ def _prefill_routed( # streaming buffers disown their borrowed slots on invalidation. cache = self.offload_cache assert cache is not None + if self.owner_cache is not None: + # The owner adapter must see the original global route. Its local-row remap, + # borrowed-buffer lifecycle, and remote zeroing are handled by the base owner + # implementation; the short-prefill optimization is global-cache-only. + return super()._prefill_routed(hidden_states, topk_weights, topk_ids) # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) if ( hidden_states.shape[0] * self.top_k >= self.num_experts diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 2dfdaaa35..248732de4 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -8,6 +8,7 @@ import safetensors import torch +from freetoken.moe.ownership import ExpertOwnership from freetoken.utils import download_hf_weight from tqdm import tqdm @@ -73,6 +74,7 @@ def iter_nvfp4_expert_pieces( chunk: int = 8 << 20, drop_page_cache: DropPageCache | None = None, primary: bool = True, + ownership: ExpertOwnership | None = None, ): """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale`` (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects, @@ -80,6 +82,10 @@ def iter_nvfp4_expert_pieces( Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land. + + ``ownership`` (owner-local TP+EP): keep only this rank's experts and renumber them into the + rank-local bank rows ``[0, local_num_experts)``, so the pieces land in the owner-local banks + the cache actually allocates. Without it every expert is loaded at its global row. """ from freetoken.models.loader import drop_page_cache as _drop from freetoken.models.loader import safetensors_weight_map @@ -89,6 +95,15 @@ def iter_nvfp4_expert_pieces( folder = download_hf_weight(model_path) weight_map = safetensors_weight_map(folder) + global_E = config.num_experts + if ownership is not None and ownership.global_num_experts != global_E: + raise ValueError( + f"expert ownership has global_num_experts={ownership.global_num_experts}, " + f"but checkpoint config has num_experts={global_E}" + ) + local_E = ownership.local_num_experts if ownership is not None else global_E + global_start = ownership.global_start if ownership is not None else 0 + wanted: dict[str, tuple[int, int, str]] = {} for name in weight_map: match = spec.key_pattern.match(name) @@ -97,14 +112,17 @@ def iter_nvfp4_expert_pieces( bank_layer = _bank_layer(spec, int(match.group("layer")), config) if bank_layer is None: continue + expert = int(match.group("expert")) + if ownership is not None and not ownership.owns(expert): + continue proj = match.group("proj") if proj not in spec.proj_to_role: raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}") kind = _canon_kind(spec, match.group("kind")) if kind not in ("weight", "weight_scale", "weight_scale_2"): raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") - wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind)) - expected = _num_moe_layers(config) * config.num_experts * 9 + wanted[name] = (bank_layer, expert - global_start, spec.proj_to_role[proj] + _kind_suffix(kind)) + expected = _num_moe_layers(config) * local_E * 9 if len(wanted) != expected: raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}") diff --git a/python/freetoken/models/qwen3_5_moe/moe.py b/python/freetoken/models/qwen3_5_moe/moe.py index 766059e5d..59324fdf5 100644 --- a/python/freetoken/models/qwen3_5_moe/moe.py +++ b/python/freetoken/models/qwen3_5_moe/moe.py @@ -76,9 +76,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # kernel may write into ``hidden_states`` in place, which would corrupt the # shared expert's input (HF also evaluates the shared expert first). router_logits = self.gate.forward(hidden_states) - shared = self.shared_expert.forward(hidden_states) + owner_ep = getattr(self.experts, "owner_cache", None) is not None + if owner_ep and not isinstance(self.shared_expert.down_proj, LinearRowParallel): + raise NotImplementedError( + "owner EP shared+routed fusion requires a row-parallel shared down projection" + ) + shared = self.shared_expert.down_proj.forward( + silu_and_mul(self.shared_expert.gate_up_proj.forward(hidden_states)), + reduce=not owner_ep, + ) if owner_ep else self.shared_expert.forward(hidden_states) shared = shared * torch.sigmoid(self.shared_expert_gate.forward(hidden_states)) - routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits) + routed = self.experts.forward( + hidden_states=hidden_states, router_logits=router_logits, reduce=not owner_ep + ) + if owner_ep: + routed = self.experts._maybe_all_reduce(routed + shared) + return routed.view(num_tokens, hidden_dim) return (routed + shared).view(num_tokens, hidden_dim) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 5fc12e747..2b238bffc 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -346,11 +346,13 @@ def _moe_dims(model_config): ) -def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): +def iter_expert_pieces(model_path, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20, ownership=None): """Block-fp8 routed experts, one piece per expert: ``{gate, up, down}`` fp8 codes and their ``_scale`` (block scale) companions, named as the checkpoint's dialect stores them. Other expert kinds use the generic readers.""" if kind is not QuantKind.FP8_BLOCK: return None + if ownership is not None: + raise NotImplementedError("block-FP8 expert banks do not support owner-local TP+EP") if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") from freetoken.models.weight import experts_scattered, iter_expert_tensors_parallel diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 66211c7c2..dfc6d78d5 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -19,7 +19,9 @@ import torch from freetoken.core import get_global_ctx -from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.distributed import get_tp_info +from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearOProj, LinearReplicated +from freetoken.models.qwen4_exp.config import qwen4_exp_tp_geometry from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -118,18 +120,39 @@ class Qwen4ExpAttention(BaseOP): def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> None: self.layer_id = layer_id - self.num_q = config.num_qo_heads - self.num_kv = config.num_kv_heads + geometry = qwen4_exp_tp_geometry(config) + self.num_q = geometry.num_q_heads + self.num_kv = geometry.num_kv_heads self.head_dim = config.head_dim - self.qo_attn_dim = self.num_q * self.head_dim - self.kv_attn_dim = self.num_kv * self.head_dim + self.qo_attn_dim = geometry.q_attn_dim + self.kv_attn_dim = geometry.kv_attn_dim self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + # ``LinearColParallelMerged`` shards each segment by TP, so it is built from the + # GLOBAL sizes; the forward splits the rank-local output by ``self._qkv_split``. + self._qkv_global_split = [ + 2 * config.num_qo_heads * config.head_dim, + config.num_kv_heads * config.head_dim, + config.num_kv_heads * config.head_dim, + ] + # q|k|v are all quantized together (or all bf16), so the merged GEMM stays a + # single kernel; a modelopt MIXED_PRECISION checkpoint declares them FP8_PB_WO. + # + # The bf16 pair is column-parallel in, row-parallel out: ``qkv_proj`` hands each + # rank its own slice of heads, so ``o_proj`` must take that sharded input and + # all-reduce the partial sums. ``LinearOProj`` does both and degenerates to + # ``LinearReplicated`` at TP=1 (same weight, same GEMM, reduction skipped). A + # replicated ``o_proj`` here fails quietly under TP: each rank's partial sum still + # decodes to fluent-looking text. + if get_tp_info().size > 1 and getattr(config, "attn_quant", "none") != "none": + raise NotImplementedError( + "qwen4_exp dense TP currently supports the BF16 attention path only" + ) self.qkv_proj = LinearColParallelMerged( - config.hidden_size, self._qkv_split, has_bias=False, + config.hidden_size, self._qkv_global_split, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.qkv_proj", ) - self.o_proj = LinearReplicated( - self.qo_attn_dim, config.hidden_size, has_bias=False, + self.o_proj = LinearOProj( + config.num_qo_heads * self.head_dim, config.hidden_size, has_bias=False, quant_config=config.quant, prefix=f"{prefix}.o_proj", ) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 4bcd28de3..e9a13cda9 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -1,11 +1,12 @@ from __future__ import annotations from dataclasses import dataclass +from fnmatch import fnmatch from typing import Any, Tuple import torch -from freetoken.layers.quantization import QuantConfig +from freetoken.distributed import try_get_tp_info from freetoken.models.config import ( FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, @@ -13,6 +14,7 @@ RotaryConfig, SlotStateSpec, ) +from freetoken.utils import div_even @dataclass(frozen=True) @@ -68,6 +70,85 @@ def ple_state_width(self) -> int: return self.hc_count * self.hidden_size +@dataclass(frozen=True) +class Qwen4ExpTPGeometry: + """Rank-local dense geometry derived from the global Qwen4Exp config.""" + + tp_size: int + rank: int + num_q_heads: int + num_kv_heads: int + num_key_heads: int + num_value_heads: int + head_dim: int + key_head_dim: int + value_head_dim: int + + @property + def q_attn_dim(self) -> int: + return self.num_q_heads * self.head_dim + + @property + def kv_attn_dim(self) -> int: + return self.num_kv_heads * self.head_dim + + @property + def key_dim(self) -> int: + return self.num_key_heads * self.key_head_dim + + @property + def value_dim(self) -> int: + return self.num_value_heads * self.value_head_dim + + @property + def conv_dim(self) -> int: + return 2 * self.key_dim + self.value_dim + + @property + def local_conv_dim(self) -> int: + """Rank-local GDN convolution width used by the linear-state pool.""" + return self.conv_dim + + @property + def local_recurrent_state_shape(self) -> tuple[int, int, int]: + """Rank-local ``(value_heads, key_dim, value_dim)`` recurrent-state shape.""" + return (self.num_value_heads, self.key_head_dim, self.value_head_dim) + + +def qwen4_exp_tp_geometry( + config: ModelConfig, *, tp_size: int | None = None, rank: int | None = None +) -> Qwen4ExpTPGeometry: + """Resolve local QSA/GDN dimensions without mutating global model config.""" + tp = try_get_tp_info() + tp_size = (1 if tp is None else tp.size) if tp_size is None else tp_size + rank = (0 if tp is None else tp.rank) if rank is None else rank + if tp_size < 1: + raise ValueError(f"TP size must be positive, got {tp_size}") + if not 0 <= rank < tp_size: + raise ValueError(f"TP rank {rank} is outside [0, {tp_size})") + + linear = config.linear_attention_group() + if linear is None: + num_key_heads = num_value_heads = key_head_dim = value_head_dim = 0 + else: + num_key_heads = div_even(linear.num_key_heads, tp_size, allow_replicate=True) + num_value_heads = div_even(linear.num_value_heads, tp_size, allow_replicate=True) + key_head_dim = linear.key_head_dim + value_head_dim = linear.value_head_dim + + return Qwen4ExpTPGeometry( + tp_size=tp_size, + rank=rank, + num_q_heads=div_even(config.num_qo_heads, tp_size), + num_kv_heads=div_even(config.num_kv_heads, tp_size, allow_replicate=True), + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + head_dim=config.head_dim, + key_head_dim=key_head_dim, + value_head_dim=value_head_dim, + ) + + PLE_CONV_STATE = "ple_conv" PLE_NGRAM_STATE = "ple_ngram_ctx" @@ -93,6 +174,17 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]: ) +def _quant_get(hf_config: Any): + quant = getattr(hf_config, "quantization_config", None) + if quant is None: + return None + return quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d)) + + +def _ignored(patterns, module_name: str) -> bool: + return any(fnmatch(module_name, pat) for pat in patterns) + + def _layer_types(text: Any) -> list[str]: layer_types = getattr(text, "layer_types", None) if layer_types is not None: @@ -109,6 +201,45 @@ def _layer_types(text: Any) -> list[str]: ] +# modelopt spellings for 128x128 per-block, weight-only FP8 on the dense modules. +_FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) + + +def dense_quant_mode(algo: str, quantized_layers: Any) -> str: + """The quantization mode the dense (non-expert) projections will actually be SERVED in. + + Single source of truth for the two sides that must agree: :func:`parse_config`, which + decides the modules the model BUILDS, and the weight loader, which decides the buffers + it EMITS. They previously derived this independently from the same declaration - safe + only while they cannot disagree, and they can: the block-FP8 linears have no + tensor-parallel variant, so a rank running under TP>1 has to fall back to bf16. If only + one side knew that, the loaded buffers would not match the built modules. + + Returns ``"fp8_block"`` only when the checkpoint declares per-block weight-only FP8 on a + non-expert module AND this rank can serve it; ``"none"`` otherwise (bf16, via the + dequantize-at-load path). A checkpoint carrying ``weight_scale_inv`` without declaring + the algo is "none" here and keeps the pre-existing dequant behaviour. + """ + if str(algo or "").lower() != "mixed_precision": + return "none" + declared = any( + ".mlp.experts" not in str(module) + and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS + for module, spec in (quantized_layers or {}).items() + ) + if not declared: + return "none" + # Resolved here, once, so both sides downgrade together. ``Engine.__init__`` sets TP + # info as its very first statement, before the model config or any weight is built, so + # a rank always knows its size by the time this matters. try_get_tp_info is used rather + # than get_tp_info because config parsing also happens with no engine at all (checkpoint + # conversion, tooling, tests), where get_tp_info raises; unset means a single rank. + tp = try_get_tp_info() + if tp is not None and tp.size > 1: + return "none" + return "fp8_block" + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) @@ -138,14 +269,65 @@ def parse_config(hf_config: Any) -> ModelConfig: else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))} ) + # NOTE: the per-module schemes the built layers actually serve are decided by the + # ``QuantConfig`` (engine/config.py injects it as ``ModelConfig.quant``); the flags + # derived here are the loader/engine-facing summary of the same declaration. + get = _quant_get(hf_config) + if get is None: + expert_quant = attn_quant = dense_quant = lm_head_quant = "none" + else: + algo = str(get("quant_algo") or get("quant_method") or "").lower() + block = get("weight_block_size") + if algo == "fp8" and block: + # Official FP8 build (DeepSeek-V3-style block-fp8): only the routed experts + # are quantized (fp8-e4m3 weights + per-block weight_scale_inv); attention, + # GDN, the shared expert, HC, PLE and lm_head stay bf16. + bs = tuple(int(x) for x in block) + assert bs == (128, 128), f"only 128x128 block-fp8 is supported, got {bs}" + expert_quant = "fp8_block" + attn_quant = dense_quant = lm_head_quant = "none" + elif algo == "mixed_precision": + # modelopt MIXED_PRECISION: the quant algo is declared per module in + # ``quantized_layers`` rather than once at the top level. The community + # NVFP4-FP8 build of Qwen3.8-Flash-Next quantizes the routed experts to NVFP4 + # (read natively by the offload cache) and the dense attn/GDN projections to + # 128x128 block-FP8, declared per module as ``FP8_PB_WO``. + quantized = get("quantized_layers") or {} + experts_nvfp4 = any( + ".mlp.experts" in str(module) + and str((spec or {}).get("quant_algo", "")).upper() == "NVFP4" + for module, spec in quantized.items() + ) + # The same map declares the dense attn/GDN projections as FP8_PB_WO + # (per-block, weight-only FP8 with a ``weight_scale_inv`` sibling). Serve + # them natively instead of dequantizing at load: the four-way in_proj fusion + # splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM (see gdn.py), which + # halves the dense bytes read on every decode step. Resolved through + # dense_quant_mode so the loader reaches the same answer, TP downgrade + # included - see that function. + expert_quant = "nvfp4" if experts_nvfp4 else "none" + attn_quant = dense_quant_mode(algo, quantized) + dense_quant = lm_head_quant = "none" + else: + is_fp4 = "fp4" in algo + ignore = list(get("ignore") or []) + + # The RadixArk NVFP4 build quantizes only the routed experts; attention/GDN, + # the shared expert, HC, PLE and lm_head all sit in the modelopt ignore list + # and stay bf16. Derive every flag from that list instead of assuming the split. + def _quant(probe: str) -> str: + return "nvfp4" if is_fp4 and not _ignored(ignore, probe) else "none" + + prefix = "model.language_model.layers.0" + expert_quant = _quant(f"{prefix}.mlp.experts.0.gate_proj") + dense_quant = _quant(f"{prefix}.mlp.shared_expert.gate_proj") + attn_quant = _quant(f"{prefix}.self_attn.q_proj") + lm_head_quant = _quant("lm_head") + layer_types = _layer_types(text) full_ids = tuple(i for i, t in enumerate(layer_types) if t == "full_attention") linear_ids = tuple(i for i, t in enumerate(layer_types) if t == "linear_attention") - # the engine reads this flag for its MoE strategy decisions; every module takes its own scheme from the QuantConfig when it is built - expert_scheme = QuantConfig.from_hf(hf_config).scheme_for_name("model.language_model.layers.0.mlp.experts.0.gate_proj") - expert_quant = "none" if expert_scheme is None else str(expert_scheme.kind) - # HF stores ple_layer_ids one-indexed (validated upstream as [1, num_layers]). ple_layer_ids = tuple(int(i) - 1 for i in (getattr(text, "ple_layer_ids", None) or ())) for lid in ple_layer_ids: @@ -245,9 +427,20 @@ def parse_config(hf_config: Any) -> ModelConfig: image_token_id=getattr(hf_config, "image_token_id", None), attention_groups=groups, expert_quant=expert_quant, + attn_quant=attn_quant, + dense_quant=dense_quant, + lm_head_quant=lm_head_quant, qwen4_args=qwen4_args, slot_states=ple_slot_states(qwen4_args), ) -__all__ = ["PLE_CONV_STATE", "PLE_NGRAM_STATE", "Qwen4ExpArgs", "parse_config", "ple_slot_states"] +__all__ = [ + "PLE_CONV_STATE", + "PLE_NGRAM_STATE", + "Qwen4ExpArgs", + "Qwen4ExpTPGeometry", + "parse_config", + "ple_slot_states", + "qwen4_exp_tp_geometry", +] diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index ed49dd42f..667081de8 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -3,9 +3,11 @@ import torch import torch.nn.functional as F from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen -from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj from freetoken.layers.quantization import QuantConfig +from freetoken.utils import div_even from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla @@ -42,23 +44,38 @@ def __init__( assert head_k_dim == head_v_dim, ( f"GatedDeltaNet requires head_k_dim == head_v_dim, got {head_k_dim} != {head_v_dim}" ) - self.num_k_heads = num_k_heads - self.num_v_heads = num_v_heads + tp_size = get_tp_info().size + self.num_k_heads = div_even(num_k_heads, tp_size, allow_replicate=True) + self.num_v_heads = div_even(num_v_heads, tp_size, allow_replicate=True) self.head_k_dim = head_k_dim self.head_v_dim = head_v_dim - self.key_dim = num_k_heads * head_k_dim - self.value_dim = num_v_heads * head_v_dim + self.key_dim = self.num_k_heads * head_k_dim + self.value_dim = self.num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim + global_key_dim = num_k_heads * head_k_dim + global_value_dim = num_v_heads * head_v_dim + global_conv_dim = 2 * global_key_dim + global_value_dim self.conv_kernel_size = conv_kernel_size - # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM) + # Quantized checkpoints quantize qkv|z but not b|a, so the four-way fusion splits + # into a qkvz GEMM and a bf16 ba GEMM with their own schemes (matches sglang/vLLM). + # A modelopt MIXED_PRECISION checkpoint declares the dense attn/GDN projections + # FP8_PB_WO while the experts are NVFP4, so the scheme - not a model flag - decides. self._split_in_proj = ( - quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None + quant_config is not None + and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None ) - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] + # The fused GEMM is built from the GLOBAL sizes (``LinearColParallelMerged`` shards + # each segment by TP); the forward splits its output by the rank-local sizes. + self._in_proj_global_split = [global_conv_dim, global_value_dim, num_v_heads, num_v_heads] + self._in_proj_split = [self.conv_dim, self.value_dim, self.num_v_heads, self.num_v_heads] + if self._split_in_proj and tp_size > 1: + raise NotImplementedError( + "qwen4_exp dense TP currently supports the BF16 GDN path only" + ) if self._split_in_proj: self.in_proj_qkvz = LinearColParallelMerged( - hidden_size, [self.conv_dim, self.value_dim], has_bias=False, + hidden_size, [global_conv_dim, global_value_dim], has_bias=False, quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz", ) self.in_proj_ba = LinearColParallelMerged( @@ -68,7 +85,7 @@ def __init__( else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. self.in_proj = LinearColParallelMerged( - hidden_size, self._in_proj_split, has_bias=False, + hidden_size, self._in_proj_global_split, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.in_proj", ) self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size) @@ -76,11 +93,14 @@ def __init__( # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a # per-call .float() upcast in the decode wrapper. The weight loader exempts # *.A_log / *.dt_bias from the model-dtype downcast. - self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) - self.A_log = torch.empty(num_v_heads, dtype=torch.float32) + self.dt_bias = torch.empty(self.num_v_heads, dtype=torch.float32) + self.A_log = torch.empty(self.num_v_heads, dtype=torch.float32) self.norm = GatedRMSNorm(head_v_dim, eps=rms_norm_eps, activation=output_gate) - self.out_proj = LinearReplicated( - self.value_dim, hidden_size, has_bias=False, + # out_proj follows the checkpoint quant; it is row-parallel so each rank takes its + # own shard of the GDN output and all-reduces the partial sums (a replicated copy + # would leave each rank with a partial sum that still decodes fluently). + self.out_proj = LinearOProj( + global_value_dim, hidden_size, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.out_proj", ) diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 3687f2803..597233343 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -4,6 +4,7 @@ import torch from freetoken.kernel.triton.moe_shared_gate import shared_gate_mul_add, shared_gate_sigmoid +from freetoken.layers import silu_and_mul from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE if TYPE_CHECKING: @@ -20,6 +21,28 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) router_logits = self.gate.forward(hidden_states) + owner_ep = getattr(self.experts, "owner_cache", None) is not None + if owner_ep: + if not hasattr(self.shared_expert.down_proj, "_tp_size"): + raise NotImplementedError( + "owner EP shared+routed fusion requires a row-parallel shared projection" + ) + # Compute the gate before routed experts: a fused routed kernel is allowed to + # mutate its hidden input in-place. The non-owner path has the same ordering + # contract; owner mode must not rely on the current NVFP4 kernel being benign. + gate = shared_gate_sigmoid( + hidden_states, self.shared_expert_gate.weight.view(-1) + ) + shared = self.shared_expert.down_proj.forward( + silu_and_mul(self.shared_expert.gate_up_proj.forward(hidden_states)), + reduce=False, + ) + routed = self.experts.forward( + hidden_states=hidden_states, router_logits=router_logits, reduce=False + ) + merged = shared_gate_mul_add(routed, shared, gate) + return self.experts._maybe_all_reduce(merged).view(num_tokens, hidden_dim) + shared = self.shared_expert.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 4a2a5b615..7a26d61c3 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -22,13 +22,14 @@ import torch from freetoken.distributed import get_tp_info from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.qwen4_exp.config import dense_quant_mode from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) from freetoken.layers.quantization import get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import cached_load_hf_config, download_hf_weight +from freetoken.utils import cached_load_hf_config, div_ceil, div_even, download_hf_weight from freetoken.utils.progress import byte_bar from tqdm import tqdm @@ -47,7 +48,8 @@ desc="Qwen3.8-Flash-Next NVFP4 experts", ) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# ``.weight_scale_inv`` is the 128x128 block-FP8 reciprocal scale (see _load_maybe_block_fp8). +_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".weight_scale_inv", ".input_scale") # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." @@ -81,15 +83,186 @@ _ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn} -def _rename(raw_name: str) -> str | None: - """Checkpoint key -> FreeToken state-dict key, or None to skip.""" +def _partition(size: int, rank: int, world_size: int, *, allow_replicate: bool = False): + local = div_even(size, world_size, allow_replicate=allow_replicate) + if world_size <= size: + start = rank * local + else: + # Replicate each global head across a consecutive group of ranks, matching the + # generic KV sharder used by Qwen3. This matters when TP exceeds KV/GDN heads. + start = (rank // (world_size // size)) * local + return start, local + + +def _shard_head_rows( + tensor: torch.Tensor, + *, + num_heads: int, + rows_per_head: int, + rank: int, + world_size: int, + allow_replicate: bool = False, +) -> torch.Tensor: + expected = num_heads * rows_per_head + if tensor.shape[0] != expected: + raise ValueError( + f"expected {expected} rows for {num_heads} heads, got {tuple(tensor.shape)}" + ) + start, local = _partition( + num_heads, rank, world_size, allow_replicate=allow_replicate + ) + view = tensor.reshape(num_heads, rows_per_head, *tensor.shape[1:]) + return view[start : start + local].reshape(local * rows_per_head, *tensor.shape[1:]).contiguous() + + +def _shard_dim1(tensor: torch.Tensor, *, rank: int, world_size: int) -> torch.Tensor: + if tensor.ndim < 2: + raise ValueError(f"dim-1 sharding needs a matrix, got {tuple(tensor.shape)}") + start, local = _partition(tensor.shape[1], rank, world_size) + return tensor.narrow(1, start, local).contiguous() + + +def shard_qwen4_exp_dense_tensor( + key: str, + tensor: torch.Tensor, + *, + config, + rank: int, + world_size: int, +) -> torch.Tensor: + """Shard one *raw, unfused* Qwen4Exp dense tensor for TP. + + Fusion happens after this function. QSA q/k/v and GDN q/k/v/z/b/a therefore keep + head boundaries, while row-parallel output projections are sliced on input columns. + Routed experts, PLE/HC/indexer tensors and router weights are intentionally replicated; + expert-ID ownership belongs to the later EP stage. + """ + if world_size == 1: + return tensor + + linear = config.linear_attention_group() + if key in {"model.embed_tokens.weight", "lm_head.weight"}: + rows = div_ceil(tensor.shape[0], world_size) + start = rank * rows + shard = tensor[start : min(start + rows, tensor.shape[0])] + if shard.shape[0] != rows: + # The vocabulary axis is padded, not truncated: VocabParallelEmbedding (and the + # row-parallel lm_head) always allocate ``div_ceil(vocab, tp)`` rows -- its + # ``finish_idx`` clamps the token-index range, not the allocation -- so when the + # vocabulary is not divisible by TP the final rank must still hand over a + # full-width shard or strict loading fails on shape. The padding rows are never + # reachable by a token id. + pad = shard.new_zeros((rows - shard.shape[0], *shard.shape[1:])) + shard = torch.cat((shard, pad), dim=0) + return shard.contiguous() + + if key.endswith(".self_attn.q_proj.weight"): + return _shard_head_rows( + tensor, num_heads=config.num_qo_heads, rows_per_head=2 * config.head_dim, + rank=rank, world_size=world_size, + ) + if key.endswith((".self_attn.k_proj.weight", ".self_attn.v_proj.weight")): + return _shard_head_rows( + tensor, num_heads=config.num_kv_heads, rows_per_head=config.head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ) + if key.endswith(".self_attn.o_proj.weight"): + return _shard_dim1(tensor, rank=rank, world_size=world_size) + + if linear is not None and key.endswith(".linear_attn.in_proj_qkv.weight"): + q, k, v = torch.split( + tensor, + [ + linear.num_key_heads * linear.key_head_dim, + linear.num_key_heads * linear.key_head_dim, + linear.num_value_heads * linear.value_head_dim, + ], + dim=0, + ) + return torch.cat( + [ + _shard_head_rows( + q, num_heads=linear.num_key_heads, rows_per_head=linear.key_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + _shard_head_rows( + k, num_heads=linear.num_key_heads, rows_per_head=linear.key_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + _shard_head_rows( + v, num_heads=linear.num_value_heads, rows_per_head=linear.value_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + ], + dim=0, + ).contiguous() + if linear is not None and key.endswith(".linear_attn.in_proj_z.weight"): + return _shard_head_rows( + tensor, num_heads=linear.num_value_heads, rows_per_head=linear.value_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ) + if linear is not None and key.endswith((".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight")): + return _shard_head_rows( + tensor, num_heads=linear.num_value_heads, rows_per_head=1, + rank=rank, world_size=world_size, allow_replicate=True, + ) + if linear is not None and key.endswith(".linear_attn.conv1d.weight"): + q, k, v = torch.split( + tensor, + [ + linear.num_key_heads * linear.key_head_dim, + linear.num_key_heads * linear.key_head_dim, + linear.num_value_heads * linear.value_head_dim, + ], + dim=0, + ) + return torch.cat( + [ + _shard_head_rows( + q, num_heads=linear.num_key_heads, rows_per_head=linear.key_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + _shard_head_rows( + k, num_heads=linear.num_key_heads, rows_per_head=linear.key_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + _shard_head_rows( + v, num_heads=linear.num_value_heads, rows_per_head=linear.value_head_dim, + rank=rank, world_size=world_size, allow_replicate=True, + ), + ], + dim=0, + ).contiguous() + if linear is not None and key.endswith((".linear_attn.A_log", ".linear_attn.dt_bias")): + return _shard_head_rows( + tensor, num_heads=linear.num_value_heads, rows_per_head=1, + rank=rank, world_size=world_size, allow_replicate=True, + ) + if linear is not None and key.endswith(".linear_attn.out_proj.weight"): + return _shard_dim1(tensor, rank=rank, world_size=world_size) + + if key.endswith((".mlp.shared_expert.gate_proj.weight", ".mlp.shared_expert.up_proj.weight")): + return tensor.chunk(world_size, dim=0)[rank].contiguous() + if key.endswith(".mlp.shared_expert.down_proj.weight"): + return _shard_dim1(tensor, rank=rank, world_size=world_size) + return tensor + + +def _rename(raw_name: str, keep_scale_inv: bool = False) -> str | None: + """Checkpoint key -> FreeToken state-dict key, or None to skip. + + ``keep_scale_inv`` retains the block-FP8 ``weight_scale_inv`` tensors, which the + fp8 linears need alongside their weight; they are dropped otherwise (a rank that + dequantized its dense weights has no use for the reciprocal scale).""" if raw_name.startswith(("mtp.", "model.visual.", "visual.")): return None if _PLE_TABLE_INFIX in raw_name: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks - if raw_name.endswith(_SCALE_SUFFIXES): + if raw_name.endswith(_SCALE_SUFFIXES) and not ( + keep_scale_inv and raw_name.endswith(".weight_scale_inv") + ): return None if raw_name.startswith("model.language_model."): return "model." + raw_name[len("model.language_model.") :] @@ -112,8 +285,11 @@ class _DenseFuser: The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from. """ - def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None: + def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...], *, dequantized: bool = False) -> None: self.quant = quant + # The reader normalized every dense weight to bf16 (a TP>1 rank cannot serve the + # block-FP8 kernels), so there is nothing left for the scheme to agree with. + self.dequantized = dequantized self.groups = {fused: parts for fused, parts in packed if fused != "experts"} # experts: bank reader self.by_part: dict[str, list[tuple[str, int]]] = {} for fused, parts in self.groups.items(): @@ -142,6 +318,8 @@ def _target(self, parent: str, leaf: str) -> tuple[str, int] | None: def check(self, module: str, name: str, tensor: torch.Tensor) -> None: """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from.""" + if self.dequantized: + return scheme = self.scheme(module) if name.endswith(".weight_scale_inv"): if scheme is None or not scheme.has("weight_scale_inv"): @@ -189,12 +367,35 @@ def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor] return [(fused + kind, torch.cat(rows, dim=0))] +def _load_maybe_block_fp8(f, raw_name: str, keyset: set[str]) -> torch.Tensor: + """Load ``raw_name``, dequantizing 128x128 block-FP8 to bf16 when a sibling + ``.weight_scale_inv`` is present in the same shard; pass plain bf16 through unchanged. + + Only the TP>1 rank-local path needs this: the block-FP8 dense kernels are + replicated-only, so a rank at TP>1 builds those projections in bf16 and the reader has + to dequantize to match (the same downgrade :func:`dense_quant_mode` makes for the + model). At TP=1 the dense side is served natively as block-FP8 instead, so the fp8 + codes and their ``weight_scale_inv`` travel through ``_DenseFuser`` untouched.""" + tensor = f.get_tensor(raw_name) + if raw_name.endswith(".weight"): + base = raw_name[: -len(".weight")] + if base + ".weight_scale_inv" in keyset: + from freetoken.kernel.triton.fp8_block_linear import dequant_block_fp8 + + return dequant_block_fp8( + tensor, f.get_tensor(base + ".weight_scale_inv") + ).to(torch.bfloat16) + return tensor + + def iter_weights( model_path: str, device: torch.device, *, include_moe_experts: bool, include_non_moe: bool, + tp_shard: bool = False, + config=None, ) -> Iterator[tuple[str, torch.Tensor]]: """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers. @@ -202,26 +403,64 @@ def iter_weights( A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections. Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader. + + ``tp_shard`` enables the rank-local TP path: each RAW tensor is sliced with + :func:`shard_qwen4_exp_dense_tensor` **before** fusion, so fused buffers keep head + boundaries (QSA q/k/v, GDN qkv/z/b/a) and row-parallel output projections are cut on + their input columns. It defaults to False, and TP>1 without it still fails fast, so no + existing TP1 caller changes behaviour. The emitted keys/shapes are exactly what a TP + rank's model builds. ``config`` optionally supplies the already-parsed + :class:`~freetoken.models.config.ModelConfig` (the engine has it); otherwise the + checkpoint config is parsed here. """ - if get_tp_info().size > 1: - raise NotImplementedError("qwen4_exp weight loading supports TP=1 only") if not include_non_moe: return + tp_info = get_tp_info() + if tp_info.size > 1 and not tp_shard: + raise NotImplementedError( + "qwen4_exp runtime TP requires iter_weights(tp_shard=True); the default " + "path is TP1-only. Pass tp_shard=True to load a rank-local shard." + ) + shard = tp_info.size > 1 + if shard and config is None: + from freetoken.models.qwen4_exp.config import parse_config + + config = parse_config(cached_load_hf_config(model_path)) hf_config = cached_load_hf_config(model_path) spec = get_model_spec(hf_config.architectures[0]) - fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping) + # A TP>1 rank builds the dense projections in bf16 (the block-FP8 kernels are + # replicated-only), so the reader dequantizes instead of carrying the fp8 codes. + serve_block_fp8 = not shard + fuser = _DenseFuser( + get_quant_config(), spec.packed_modules_mapping, dequantized=not serve_block_fp8 + ) for file in tqdm( iter_weight_files(model_path), desc="Loading weights", disable=not get_tp_info().is_primary(), ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: + keyset = set(f.keys()) for raw_name in f.keys(): - name = _rename(raw_name) + name = _rename(raw_name, keep_scale_inv=serve_block_fp8) if name is None: continue - tensor = f.get_tensor(raw_name) + tensor = ( + f.get_tensor(raw_name) + if serve_block_fp8 + else _load_maybe_block_fp8(f, raw_name, keyset) + ) + if shard: + # Slice the RAW name: fusion happens afterwards, so head-bearing + # groups are still separable and the fused order is preserved. + tensor = shard_qwen4_exp_dense_tensor( + name, + tensor, + config=config, + rank=tp_info.rank, + world_size=tp_info.size, + ) fused = fuser.fuse(name, tensor) if fused is None: fuser.check_unfused(name, tensor) @@ -378,5 +617,6 @@ def nvfp4_expert_spec(model_path: str, config): "nvfp4_expert_spec", "PleTable", "iter_weights", + "shard_qwen4_exp_dense_tensor", "load_ple_table", ] diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 0cff7eecf..8a47ab5db 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -12,6 +12,7 @@ from __future__ import annotations import glob +import inspect import json import mmap import os @@ -211,6 +212,8 @@ def load_weight( device: torch.device, *, include_moe_experts: bool = True, + tp_shard: bool = False, + tp_config=None, ) -> Iterator[Tuple[str, torch.Tensor]]: # FTW checkpoint: dense weights are stored post-iter_weights, so we replay them # model-agnostically instead of re-running the per-model reader. Which tensors exist is @@ -225,6 +228,10 @@ def load_weight( # stack. Vision is opt-in (default OFF, see vision_load_enabled): when it is off the # model never builds the tower, so replaying those tensors would trip load_state_dict's # strict unexpected-key check. Skip them here to match the model the engine built. + # + # ``tp_shard`` is a no-op on this path: the FTW dense shard is written POST-shard, so + # it is already rank-local. The engine asks for tp_shard on every TP>1 launch, so + # rejecting the combination here would break FTW checkpoints at TP>1. skip_vision = not vision_load_enabled() for name, tensor in iter_ftw_weights(model_path): if skip_vision and name.startswith(VISION_KEY_PREFIXES): @@ -234,12 +241,20 @@ def load_weight( _config, spec = _spec_for_model_path(model_path) iter_weights = _load_attr(spec.module, spec.iter_weights) - yield from iter_weights( - model_path, - device, - include_moe_experts=include_moe_experts, - include_non_moe=True, - ) + kwargs = dict(include_moe_experts=include_moe_experts, include_non_moe=True) + parameters = inspect.signature(iter_weights).parameters + if "tp_shard" in parameters: + # Readers that declare ``tp_shard`` slice the RAW checkpoint tensors themselves. + kwargs["tp_shard"] = tp_shard + if tp_config is not None and "config" in parameters: + kwargs["config"] = tp_config + # Readers that do NOT declare ``tp_shard`` already shard inside ``iter_weights`` (they call + # ``shard_tensor`` with ``tp_info.rank``/``tp_info.size``: llama, qwen2, qwen3, qwen3_moe, + # mistral, gpt_oss, minimax_m2), so TP>1 is handled and the flag must NOT be forwarded -- + # forwarding it raised and turned a working TP>1 launch into a startup failure. + # Architectures that cannot shard at all still fail loudly in ``load_state_dict``'s shape + # check, exactly as they did before this flag existed. + yield from iter_weights(model_path, device, **kwargs) def load_q4_0_moe_expert_sources( diff --git a/python/freetoken/moe/__init__.py b/python/freetoken/moe/__init__.py index 8794abce6..97d6a4f11 100644 --- a/python/freetoken/moe/__init__.py +++ b/python/freetoken/moe/__init__.py @@ -9,9 +9,20 @@ MOE_STRATEGIES = ("fused", "offload", "cpu", "hybrid") OFFLOAD_MOE_STRATEGIES = frozenset({"offload", "cpu", "hybrid"}) +# Owner-local expert placement (TP+EP): the expert-parallel group size, the per-rank bank +# geometry derived from it, and the incremental update a runtime resize hands back. +from .ownership import ExpertOwnership, OwnerCacheGeometry, OwnerCacheUpdate + def is_offload_moe_strategy(strategy: str) -> bool: return strategy in OFFLOAD_MOE_STRATEGIES -__all__ = ["MOE_STRATEGIES", "OFFLOAD_MOE_STRATEGIES", "is_offload_moe_strategy"] +__all__ = [ + "MOE_STRATEGIES", + "OFFLOAD_MOE_STRATEGIES", + "is_offload_moe_strategy", + "ExpertOwnership", + "OwnerCacheGeometry", + "OwnerCacheUpdate", +] diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d..9ed7c3142 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -75,6 +75,7 @@ def build_expert_banks( device: torch.device, layer_sink=None, dummy: bool = False, + num_experts: int | None = None, ) -> ExpertBanks: """Fill host banks in the kernel's layout from a stream of expert pieces. @@ -89,7 +90,9 @@ def build_expert_banks( kernel = method.kernel layout = method.layout() - E = method.cfg.num_experts + # Owner-local EP fills only this rank's rows, so the bank's E dim is the local expert + # count, not the layer's (global) routing count. ``num_experts`` overrides it. + E = method.cfg.num_experts if num_experts is None else num_experts specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident} hb = alloc_layer_banks(specs, num_layers) banks = {role: [b.tensor for b in hb[role]] for role in specs} @@ -178,30 +181,39 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, } -def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None, ownership=None) -> ExpertBanks: expert_quant = model_config.expert_quant if expert_quant not in _PROVIDERS: raise ValueError( f"{expert_quant!r} experts load through their MoE quant method; " f"only {sorted(_PROVIDERS)} still have a format provider" ) + provider_kwargs = dict( + parallel=parallel, workers=workers, chunk=chunk, + decode_target=decode_target, layer_sink=layer_sink, + ) + if ownership is not None: + provider_kwargs["ownership"] = ownership return _PROVIDERS[expert_quant]( - model_path, model_config, device, dtype, dummy, - parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, - layer_sink=layer_sink, + model_path, model_config, device, dtype, dummy, **provider_kwargs ) -def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks: +def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None, ownership=None) -> ExpertBanks: from freetoken.moe.expert_pieces import iter_expert_pieces num_layers = model_config.num_moe_layers + # owner-local EP: the banks hold only this rank's rows + E = ownership.local_num_experts if ownership is not None else None if dummy: - return build_expert_banks(method, num_layers, None, device=device, dummy=True) + return build_expert_banks(method, num_layers, None, device=device, dummy=True, num_experts=E) pieces = iter_expert_pieces( - model_path, model_config, method.kind, parallel=parallel, workers=workers, chunk=chunk + model_path, model_config, method.kind, parallel=parallel, workers=workers, chunk=chunk, + ownership=ownership, + ) + return build_expert_banks( + method, num_layers, pieces, device=device, layer_sink=layer_sink, num_experts=E ) - return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink) def _host_ram_fits_parallel(model_path: str) -> bool: @@ -286,6 +298,7 @@ def load_expert_banks( decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + ownership=None, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -314,6 +327,15 @@ def load_expert_banks( from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks if model_path and is_ftw_checkpoint(model_path) and not dummy: + if ownership is not None: + # ``load_ftw_banks`` rebuilds ``[num_experts, ...]`` GLOBAL rows and has no + # ownership filter, so the banks could not bind to the owner-local geometry. + # The engine rejects this combination up front; guard the loader too so a + # converter/tool call cannot reach the same inconsistent state. + raise NotImplementedError( + "owner-local expert banks are not supported for FTW checkpoints: the FTW " + "bank loader rebuilds global expert rows and does not filter by ownership" + ) banks = load_ftw_banks( model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, layer_residency=layer_residency, @@ -354,8 +376,8 @@ def load_expert_banks( def _build(par: bool) -> ExpertBanks: if method is not None: - return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink) - return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink) + return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, ownership) + return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink, ownership) with requested_residency(layer_residency) as residency_plan: try: diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py index 79f655a95..4bf63f5b0 100644 --- a/python/freetoken/moe/expert_pieces.py +++ b/python/freetoken/moe/expert_pieces.py @@ -8,6 +8,7 @@ from __future__ import annotations +import inspect from typing import Callable, Iterable, Iterator import torch @@ -38,7 +39,8 @@ def _model_hook(spec, name: str): def iter_expert_pieces( - model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20 + model_path: str, config, kind: QuantKind, *, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20, + ownership=None, ) -> Iterator[Piece]: """The pieces of ``model_path``'s routed experts, stored as ``kind``. @@ -47,14 +49,29 @@ def iter_expert_pieces( experts come from the family's stacked ``iter_weights`` and NVFP4 experts from its ``nvfp4_expert_spec``. The reader is resolved here, before any bank is allocated, so a missing parallel reader raises ``NotImplementedError`` while a serial fallback is still cheap. + + ``ownership`` (owner-local TP+EP) restricts the stream to this rank's experts and renumbers + them into rank-local bank rows; readers that cannot serve it raise ``NotImplementedError``. """ spec = get_model_spec(config.architectures[0]) hook = _model_hook(spec, "iter_expert_pieces") if hook is not None: - pieces = hook(model_path, config, kind, parallel=parallel, workers=workers, chunk=chunk) + params = inspect.signature(hook).parameters + kw = dict(parallel=parallel, workers=workers, chunk=chunk) + if "ownership" in params: + kw["ownership"] = ownership + elif ownership is not None: + raise NotImplementedError( + f"{spec.module}.iter_expert_pieces does not support owner-local expert banks" + ) + pieces = hook(model_path, config, kind, **kw) if pieces is not None: return pieces if kind is QuantKind.NONE: + if ownership is not None: + raise NotImplementedError( + f"{spec.module} has no owner-local reader for bf16 experts" + ) return _bf16_pieces(model_path, config, spec, parallel=parallel, workers=workers, chunk=chunk) if kind is QuantKind.NVFP4: spec_hook = _model_hook(spec, "nvfp4_expert_spec") @@ -63,7 +80,8 @@ def iter_expert_pieces( from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces return iter_nvfp4_expert_pieces( - model_path, config, spec_hook(model_path, config), parallel=parallel, workers=workers, chunk=chunk + model_path, config, spec_hook(model_path, config), + parallel=parallel, workers=workers, chunk=chunk, ownership=ownership, ) raise NotImplementedError(f"{spec.module} provides no expert reader for {kind!r} experts") diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 7abb9c800..c74fe0503 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,7 +2,7 @@ import math import os -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Iterator import torch @@ -27,6 +27,8 @@ from freetoken.utils import init_logger +from .ownership import OwnerCacheGeometry, OwnerCacheUpdate, same_device + logger = init_logger(__name__) # quant_format -> bank names, in registration order: the single place a format's bank @@ -102,6 +104,11 @@ @dataclass class OffloadMoeCache: + # Marks the global-ID cache. ``OwnerOffloadMoeCache`` overrides this so the MoE layer + # can dispatch to the owner-local route adapter without importing the wrapper (which + # would be a cycle: the wrapper wraps this class). + is_owner_local = False + num_layers: int num_experts: int cache_size: int @@ -141,11 +148,27 @@ class OffloadMoeCache: # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal # time (perfect overlap): fetched : cpu = pcie : cpu - pcie. hybrid_fetch_fraction: float = 0.0 + # Explicit owner-local geometry is opt-in. The legacy cache uses global expert IDs + # throughout its kernels and must not silently accept local owner rows. + owner_geometry: OwnerCacheGeometry | None = None # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots layout: dict | None = None max_slots: int | None = None def __post_init__(self) -> None: + if self.owner_geometry is not None: + self.owner_geometry.validate_cache_binding( + num_layers=self.num_layers, + num_experts=self.num_experts, + cache_size=self.cache_size, + prefill_overlap=self.prefill_overlap, + ) + raise NotImplementedError( + "owner-local cache geometry is validated but the OffloadMoeCache runtime " + "namespace mapping is not enabled; leave owner_geometry unset until " + "global/local/slot IDs are wired through every cache kernel" + ) + policy_ids = {"lru": 0} assert self.cache_policy in policy_ids assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target @@ -246,6 +269,11 @@ def __post_init__(self) -> None: # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the # captured graph would not re-run this host-side scatter on replay). self.collect_decode_freq = False + # Opt-in ordered route trace recorder (moe/route_trace.RouteTraceRecorder), + # attached by the engine when --moe-trace-route is set. Records RAW global + # expert ids in call order before lru_ensure rewrites them to slots. None = + # disabled (production default, zero overhead). + self.route_recorder = None self.decode_freq = torch.zeros( (self.num_layers, self.num_experts), dtype=torch.int64, device=self.device ) @@ -354,7 +382,11 @@ def set_bank_sources( name, layer_id, source.shape, source.dtype, ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.empty( + # Remote owner-route entries use slot zero with zero weight. Decode kernels still + # load the slot row before applying the weight, so an uninitialized row can turn + # ``0 * NaN`` into NaN and poison logits. Zero-fill the data plane; bookkeeping + # workspaces below remain uninitialized for graph/performance reasons. + self.bank_caches[name] = torch.zeros( (self.cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device, @@ -492,7 +524,7 @@ def rebuild(self, cache_size: int) -> None: # 3. Reallocate the slot cache from the retained host sources. for name in self.bank_schema: head = self.bank_sources[name][0] - self.bank_caches[name] = torch.empty( + self.bank_caches[name] = torch.zeros( (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] @@ -848,6 +880,9 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: # slot ids in place), so snapshot the routing histogram before that happens. ids = expert_ids.reshape(-1).long() self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) + if self.route_recorder is not None: + # same raw-ids point: append the ordered trace BEFORE the in-place rewrite. + self.route_recorder.record(layer_id, expert_ids, phase=0) self._pending_src_layer = layer_id self._pending_whole_layer = False ensure_experts(self, layer_id, expert_ids) @@ -867,6 +902,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None if self.collect_decode_freq: ids = expert_ids.reshape(-1).long() self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) + if self.route_recorder is not None: + self.route_recorder.record(layer_id, expert_ids, phase=0) self._pending_src_layer = layer_id self._pending_whole_layer = False ensure_experts_hybrid( @@ -993,6 +1030,13 @@ def decode_routing_stats(self) -> dict: C = max(1, int(round(slots_per_layer))) sorted_f, _ = torch.sort(freq, dim=1, descending=True) oracle_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() + # The realized cache is ONE unified LRU slot pool shared across layers, so the + # tight per-row bound is the top-cache_size rows of the flattened (layer, expert) + # activation distribution -- how often a perfect policy would find the expert + # already resident. The per-layer figure above assumes an even per-layer split. + flat = freq.reshape(-1) + top = torch.sort(flat, descending=True).values[: self.cache_size] + oracle_hit_global = (top.sum() / flat.sum().clamp(min=1)).item() ws = (freq > 0).sum(dim=1).float() cdf = torch.cumsum(sorted_f, dim=1) / total.clamp(min=1).unsqueeze(1) cover90 = ((cdf < 0.9).sum(dim=1).float() + 1)[valid] @@ -1005,15 +1049,73 @@ def decode_routing_stats(self) -> dict: "working_set_max": int(ws[valid].max().item()), "experts_for_90pct": cover90.mean().item(), "oracle_hit_at_slots": oracle_hit, + "oracle_hit_global": oracle_hit_global, "norm_entropy": norm_ent, } + def stats_snapshot(self) -> dict: + """Cross-process snapshot of the slot-cache counters for /v1/stats. + + Built for the scheduler's throttled stamp onto the reply stream: ONE device + sync per counter group (stack + tolist) instead of decode_miss_stats()'s + per-field .item()s. ``resident`` (slots holding a live expert, vs the + ``cache_size`` capacity and ``total_experts`` expert-layer pairs) is always + available; miss/fetch counters require ``collect_stats``, and the routing + concentration block requires ``collect_decode_freq``. All values are + since-process-start (or the last rebuild/reset), like the other getters.""" + out: dict = { + "cache_size": self.cache_size, + "total_experts": self.num_layers * self.num_experts, + "resident": int(self.usage.gt(0).sum().item()), + "target": self.decode_target, + } + if self.decode_target == "hybrid": + active, missing, calls, fetched = ( + int(x) + for x in torch.stack( + [self.stat_active, self.stat_missing, self.stat_calls, self.stat_fetched] + ).tolist() + ) + elif self.collect_stats: + active, missing, calls = (int(x) for x in self.lru_stats.sum(0).tolist()) + fetched = int(self.stat_fetched.item()) + else: + active = missing = calls = fetched = 0 + out.update( + { + "layer_calls": calls, + "active_per_layer": (active / calls) if calls else 0.0, + "missing_per_layer": (missing / calls) if calls else 0.0, + "miss_rate": (missing / active) if active else 0.0, + "fetched_per_layer": (fetched / calls) if calls else 0.0, + "fetch_rate": (fetched / missing) if missing else 0.0, + "prefill_hit_rows": self.prefill_hit_rows, + "prefill_rows": self.prefill_total_rows, + } + ) + if self.collect_decode_freq and int(self.decode_freq.sum().item()) > 0: + try: + out["routing"] = self.decode_routing_stats() + except Exception: # noqa: BLE001 -- stats must never break the reply path + pass + return out + def copy_missing(self) -> None: assert self.banks, "set_bank_sources must register the banks first" layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" + whole_layer = self._pending_whole_layer + # Consume the staged state exactly ONCE. ``src_indices``/``evict_slots`` are shared + # buffers overwritten by the next layer's staging, so leaving ``_pending_src_layer`` + # set makes it impossible for a caller to tell "nothing staged" from "staged two + # layers ago" -- and the owner wrapper's ``is None`` guard depends on that + # distinction. Clearing here (before the copies, which only need the captured locals + # plus the already-staged slot tensors) also means a mid-copy exception leaves the + # cache in a clean "nothing staged" state instead of a stale one. + self._pending_src_layer = None + self._pending_whole_layer = False if layer_id in self._unpinned_layers: - if not self._pending_whole_layer: + if not whole_layer: raise RuntimeError( f"layer {layer_id} is unpinned: its only copy is the whole-layer " f"pageable materialize (position == expert id); ensure_experts's " @@ -1053,6 +1155,442 @@ def copy_missing(self) -> None: ) +class OwnerOffloadMoeCache: + """Owner-local adapter over the existing GPU slot-cache implementation. + + The wrapped cache sees only ``local_num_experts`` rows. Callers must use + :meth:`ensure_route` rather than the legacy ``ensure_experts`` entry point: the adapter + compacts owned route entries before LRU admission, then restores the original route shape + with zero-weight slot-zero placeholders for remote entries. This class is opt-in and is + not attached by ``attach_offload_moe_cache``; prefill/collective/model wiring remains a + separate P3 task. + + Two admission paths share the same namespace boundary: + + * :meth:`ensure_route` -- eager: compacts the route to its owned positions and reports + miss/eviction diagnostics. The compaction needs a device->host read per layer + (``nonzero`` + ``num_indices.item()``), so it cannot be captured. + * :meth:`ensure_route_graph` -- fixed-shape, sync-free sentinel admission for CUDA-graph + decode (see its docstring). Selected by the ``graph_safe`` constructor flag. + """ + + is_owner_local = True # MoELayer dispatches to ensure_route() on this flag + + def __init__( + self, + geometry: OwnerCacheGeometry, + device: torch.device, + *, + cache_policy: str = "lru", + quant_format: str = "bf16", + prefill_hit_d2d: bool = False, + graph_safe: bool = False, + layout: dict | None = None, + max_slots: int | None = None, + ) -> None: + if layout is None and quant_format not in _BANK_SCHEMAS: + raise ValueError(f"unknown quant_format {quant_format!r}") + self.geometry = geometry + # The owner adapter is GPU-only: it wraps the GPU slot cache and `_decode_owner` is + # selected before the `is_cpu_layer` branch, so a CPU/hybrid target could not be + # honoured. `_validate_owner_ep_config` rejects `--moe-cpu-layers` under owner EP + # rather than accepting a configuration this class would silently ignore. + self._cache = OffloadMoeCache( + num_layers=geometry.num_layers, + num_experts=geometry.local_num_experts, + cache_size=geometry.cache_size, + device=device, + cache_policy=cache_policy, + prefill_overlap=geometry.prefill_overlap, + prefill_hit_d2d=prefill_hit_d2d, + quant_format=quant_format, + decode_target="gpu", + layout=layout, + max_slots=max_slots, + ) + self._pending_owned = False + # Decode admission implementation: True selects the fixed-shape, sync-free + # ``ensure_route_graph`` (required for CUDA-graph capture), False keeps the eager + # compacting ``ensure_route`` with its miss/eviction diagnostics. + self.graph_safe = graph_safe + + @property + def global_num_experts(self) -> int: + return self.geometry.global_num_experts + + @property + def num_experts(self) -> int: + """The row count visible to the wrapped local bank and slot kernels.""" + return self.geometry.local_num_experts + + @property + def cache_size(self) -> int: + return self.geometry.cache_size + + @property + def resident(self) -> int: + """Number of resident owner-local slots.""" + return int(self._cache.id_of_slot.ge(0).sum().item()) + + @property + def device(self) -> torch.device: + return self._cache.device + + def __getattr__(self, name): + # Keep the wrapper small while preserving the existing cache's read-only reports and + # bank-view helpers. Explicit route methods below prevent unsafe legacy admission. + cache = object.__getattribute__(self, "_cache") + return getattr(cache, name) + + # --- engine-assigned flags: forward BOTH directions to the wrapped cache ----------- + # ``__getattr__`` only handles reads; a plain assignment would land on the wrapper while + # the inner cache keeps its own default (False / empty), silently disabling stats, the + # route trace and CPU-layer routing. These properties keep the two objects in sync. + @property + def cpu_layer_ids(self): + return self._cache.cpu_layer_ids + + @cpu_layer_ids.setter + def cpu_layer_ids(self, value): + self._cache.cpu_layer_ids = value + + @property + def collect_stats(self): + return self._cache.collect_stats + + @collect_stats.setter + def collect_stats(self, value): + self._cache.collect_stats = value + + @property + def collect_decode_freq(self): + return self._cache.collect_decode_freq + + @collect_decode_freq.setter + def collect_decode_freq(self, value): + self._cache.collect_decode_freq = value + + @property + def route_recorder(self): + return self._cache.route_recorder + + @route_recorder.setter + def route_recorder(self, value): + self._cache.route_recorder = value + + @property + def decode_target(self): + return self._cache.decode_target + + @decode_target.setter + def decode_target(self, value): + self._cache.decode_target = value + + @property + def hybrid_max_fetch(self): + return self._cache.hybrid_max_fetch + + @hybrid_max_fetch.setter + def hybrid_max_fetch(self, value): + self._cache.hybrid_max_fetch = value + + @property + def hybrid_fetch_fraction(self): + return self._cache.hybrid_fetch_fraction + + @hybrid_fetch_fraction.setter + def hybrid_fetch_fraction(self, value): + self._cache.hybrid_fetch_fraction = value + + def set_bank_sources( + self, + sources: dict[str, list[torch.Tensor]], + layer_residency: list[str] | None = None, + ) -> None: + """Attach banks whose first dimension is the owner-local expert count.""" + self.geometry.validate_source_banks(sources) + self._cache.set_bank_sources(sources, layer_residency=layer_residency) + + def set_alphas( + self, gate_up_alpha: torch.Tensor | None, down_alpha: torch.Tensor | None + ) -> None: + """Attach owner-local per-layer alpha vectors for tiled quantized backends.""" + if gate_up_alpha is None and down_alpha is None: + return + if gate_up_alpha is None or down_alpha is None: + raise ValueError("gate_up_alpha and down_alpha must be provided together") + expected = (self.geometry.num_layers * self.num_experts,) + if gate_up_alpha.shape != expected or down_alpha.shape != expected: + raise ValueError( + f"owner alpha vectors must have shape {expected}, got " + f"{tuple(gate_up_alpha.shape)} and {tuple(down_alpha.shape)}" + ) + self._cache.set_alphas(gate_up_alpha, down_alpha) + + def ensure_experts(self, *_args, **_kwargs) -> None: + raise RuntimeError( + "owner-local cache requires ensure_route(weights, global_expert_ids); " + "raw ensure_experts would admit remote global IDs" + ) + + def ensure_experts_hybrid(self, *_args, **_kwargs) -> None: + raise NotImplementedError( + "owner-local hybrid admission is not implemented; use the GPU owner adapter" + ) + + def ensure_route( + self, layer_id: int, weights: torch.Tensor, global_expert_ids: torch.Tensor + ) -> OwnerCacheUpdate: + """Admit only owned route entries and return slot IDs safe for local GEMM. + + Admission is compacted to owned positions because the legacy LRU kernel has no mask + argument. The returned tensors retain the input route shape; remote positions use + local-row zero, slot zero, and zero weight. The caller must invoke ``copy_missing`` + before reading the slot bank views. + """ + if self._pending_owned: + raise RuntimeError("copy_missing must complete the previous owner route first") + if not same_device(weights.device, self.device) or not same_device( + global_expert_ids.device, self.device + ): + raise ValueError( + f"owner route tensors must be on {self.device}, got " + f"{weights.device} and {global_expert_ids.device}" + ) + if not 0 <= layer_id < self.geometry.num_layers: + raise ValueError( + f"layer_id {layer_id} is outside [0, {self.geometry.num_layers})" + ) + if self.route_recorder is not None: + # The recorder contract is global route IDs before any owner-local remap. + self.route_recorder.record(layer_id, global_expert_ids, phase=0) + route = self.geometry.partition_route(weights, global_expert_ids) + flat_ids, owned = self.geometry.global_to_local_flat(layer_id, global_expert_ids) + owned_positions = owned.reshape(-1).nonzero(as_tuple=False).flatten() + slot_ids_flat = torch.zeros( + (global_expert_ids.numel(),), dtype=torch.int32, device=self.device + ) + missing = torch.empty((0,), dtype=torch.int32, device=self.device) + evicted = torch.empty((0,), dtype=torch.int32, device=self.device) + + if owned_positions.numel(): + local_ids = route.local_ids.reshape(-1).index_select(0, owned_positions) + local_ids = local_ids.to(dtype=torch.int32).contiguous() + old_id_of_slot = self._cache.id_of_slot.clone() + # The inner cache must not record the compressed local IDs as if they were the + # raw global route. Owner mode records at this boundary, before local remapping. + recorder = self._cache.route_recorder + self._cache.route_recorder = None + try: + self._cache.ensure_experts(layer_id, local_ids) + finally: + self._cache.route_recorder = recorder + self._pending_owned = True + + count = int(self._cache.num_indices.item()) + missing = self._cache.src_indices[:count].clone() + victim_slots = self._cache.evict_slots[:count].long() + old_ids = old_id_of_slot.index_select(0, victim_slots) + evicted = old_ids[old_ids.ge(0)].to(dtype=torch.int32) + slot_ids_flat.index_copy_(0, owned_positions, local_ids) + + return OwnerCacheUpdate( + slot_ids=slot_ids_flat.reshape(global_expert_ids.shape), + local_ids=route.local_ids, + local_flat_ids=torch.where(owned, flat_ids, torch.zeros_like(flat_ids)), + weights=route.weights, + owned_mask=route.owned_mask, + missing_local_ids=missing, + evicted_flat_ids=evicted, + ) + + def ensure_route_graph( + self, layer_id: int, weights: torch.Tensor, global_expert_ids: torch.Tensor + ) -> OwnerCacheUpdate: + """CUDA-graph-safe owner admission: fixed shape, zero device->host reads. + + The eager :meth:`ensure_route` compacts the route onto its owned positions, which + both makes the admitted tensor's LENGTH data-dependent and forces a sync per layer + (``nonzero`` to find the positions, ``num_indices.item()`` to size the diagnostics). + Both are illegal inside ``torch.cuda.graph`` capture. + +This variant never changes the shape. Remote entries are remapped to a row that is + already owned by the SAME route instead of being dropped:: + + local_row = owned ? global_id - global_start : row_of_first_owned_position + + The flashlib LRU kernel accepts a full route unchanged (no masked-admission API is + needed) and its in-place rewrite yields a valid slot id for EVERY position; a remote + position simply points at an owned row's slot with ``weight == 0``, contributing + exactly nothing to the grouped GEMM. Those rows hold real expert weights (finite), so + the zero weighting never relies on ``0 * NaN``. + + Reusing an owned row keeps the admitted row set identical to the eager compaction's, + so the two paths place the same rows in the same slots and share one cache behaviour. + The counters do report the full top-k as active, since every position is submitted. + + Diagnostics that need a device->host read (``missing_local_ids`` / + ``evicted_flat_ids``) are returned empty; use the eager path when those are needed. + """ + if self._pending_owned: + raise RuntimeError("copy_missing must complete the previous owner route first") + if not same_device(weights.device, self.device) or not same_device( + global_expert_ids.device, self.device + ): + raise ValueError( + f"owner route tensors must be on {self.device}, got " + f"{weights.device} and {global_expert_ids.device}" + ) + if not 0 <= layer_id < self.geometry.num_layers: + raise ValueError( + f"layer_id {layer_id} is outside [0, {self.geometry.num_layers})" + ) + if weights.shape != global_expert_ids.shape: + raise ValueError( + f"route weights and expert IDs must have the same shape, got " + f"{tuple(weights.shape)} and {tuple(global_expert_ids.shape)}" + ) + if global_expert_ids.dtype not in (torch.int8, torch.int16, torch.int32, torch.int64): + raise TypeError( + f"router expert IDs must be an integer tensor, got {global_expert_ids.dtype}" + ) + if self.route_recorder is not None: + raise RuntimeError( + "--moe-trace-route records host-side and cannot run inside a CUDA graph; " + "disable it or serve with --cuda-graph-max-bs 0" + ) + + # Sync-free equivalents of geometry.partition_route / global_to_local_flat: those + # helpers validate with `if torch.any(...)`, which is itself a device->host read. + ownership = self.geometry.ownership + local_ids, owned = ownership.global_to_local(global_expert_ids) + zero_row = torch.zeros((), dtype=local_ids.dtype, device=self.device) + if owned.ndim == 0 or owned.shape[-1] == 0: + raise ValueError("owner route needs at least one candidate position per row") + # Remote entries reuse an OWNED row of the SAME route (the row of its first owned + # position) rather than a fixed sentinel row. Their weight is zeroed below, so which + # row they point at cannot affect the math -- but reusing an owned row keeps the + # admitted row set EXACTLY the owned rows, i.e. identical to what the eager + # compaction admits. A fixed sentinel would instead add one always-touched row per + # layer, and that row measurably churns (missing/layer 0.72 -> 1.35 measured), which + # is pure extra PCIe traffic. + first_owned = torch.argmax(owned.to(torch.int8), dim=-1, keepdim=True) + fallback = local_ids.gather(-1, first_owned) + # All-remote row (rare): there is no owned row to borrow, so fall back to row zero. + fallback = torch.where(owned.any(dim=-1, keepdim=True), fallback, zero_row) + safe_row = torch.where(owned, local_ids, fallback) + # ``admit`` is rewritten IN PLACE into slot ids by the LRU kernel, and ``.to(int32)`` + # returns the same tensor when the route is already int32 -- so every value that must + # survive as a bank ROW has to be materialized before the kernel runs. + local_rows = safe_row.clone() + # Computed before the rewrite as well (this allocates a fresh tensor, but keep the + # ordering explicit so a future in-place tweak cannot alias it). + local_flat = layer_id * self.geometry.local_num_experts + safe_row + admit = safe_row.to(dtype=torch.int32).contiguous() + + recorder = self._cache.route_recorder + self._cache.route_recorder = None + try: + # Admits row zero for remote-only routes too, so the copy plan is never empty and + # the captured kernel sequence is identical on every replay. + self._cache.ensure_experts(layer_id, admit) + finally: + self._cache.route_recorder = recorder + self._pending_owned = True + + zero_weight = torch.zeros((), dtype=weights.dtype, device=weights.device) + empty = torch.empty((0,), dtype=torch.int32, device=self.device) + return OwnerCacheUpdate( + # ``admit`` was rewritten in place: owned positions carry their slot id, remote + # positions carry row zero's slot id (harmless: their weight is zero). + slot_ids=admit.reshape(global_expert_ids.shape), + local_ids=local_rows, + local_flat_ids=local_flat, + weights=torch.where(owned, weights, zero_weight).contiguous(), + owned_mask=owned, + missing_local_ids=empty, + evicted_flat_ids=empty, + ) + + def copy_missing(self) -> None: + """Copy the pending owned misses; remote-only routes are a no-op.""" + if not self._pending_owned: + # materialize_layer stages a whole-layer copy without setting the decode + # admission flag; the inner cache still has pending state to complete. + if self._cache._pending_src_layer is None: + return + self._cache.copy_missing() + self._pending_owned = False + + def rebuild(self, cache_size: int) -> None: + """Resize the wrapped slot cache and keep the owner geometry in step. + + ``__getattr__`` would otherwise forward ``rebuild`` to the inner cache, whose + implementation disables ``prefill_overlap`` when the new size cannot hold two complete + local layers. ``geometry`` is frozen and would keep the old ``cache_size`` / + ``prefill_overlap``, so ``materialize_layer`` would still take the overlap path and + later call ``wait_prefill_layer`` against an inner cache that has overlap disabled -- + buffers that no longer exist. Re-derive the geometry from the inner cache instead. + """ + self._cache.rebuild(cache_size) + self.geometry = replace( + self.geometry, + cache_size=self._cache.cache_size, + prefill_overlap=self._cache.prefill_overlap, + ) + + def begin_prefill(self) -> None: + """Start the borrowed-buffer lifecycle for an owner-local prefill.""" + self._cache.begin_prefill() + + def prefetch_prefill_layer(self, layer_id: int) -> None: + self._cache.prefetch_prefill_layer(layer_id) + + def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: + return self._cache.wait_prefill_layer(layer_id) + + def release_prefill_layer(self, layer_id: int) -> None: + self._cache.release_prefill_layer(layer_id) + + def materialize_layer(self, layer_id: int, buffer_id: int = 0) -> torch.Tensor: + """Materialize all local rows using the legacy prefill choreography.""" + if self.geometry.prefill_overlap: + if buffer_id != layer_id % 2: + raise ValueError( + "owner prefill buffer_id must match layer_id % 2 for the legacy buffers" + ) + self._cache.prefetch_prefill_layer(layer_id) + else: + if buffer_id != 0: + raise ValueError("owner prefill overlap is disabled; buffer_id must be 0") + self._cache.materialize_layer(layer_id) + # materialize_layer only stages the whole-layer copy in the legacy cache; + # complete it before owner GEMM reads the local bank rows. + self.copy_missing() + return torch.arange( + buffer_id * self.num_experts, + (buffer_id + 1) * self.num_experts, + dtype=torch.int32, + device=self.device, + ) + + def reset(self) -> None: + self._cache.reset() + self._pending_owned = False + + def validate_invariants(self) -> None: + """Validate the wrapped cache maps in the owner-local flat-ID namespace.""" + self.geometry.validate_slot_maps(self._cache.slot_for_id, self._cache.id_of_slot) + for slot, flat in enumerate(self._cache.id_of_slot.tolist()): + if flat < 0: + continue + layer, local = divmod(flat, self.num_experts) + if int(self._cache.slot_for_id[layer, local].item()) != slot: + raise AssertionError( + f"owner cache reverse map mismatch at slot={slot}, flat_id={flat}" + ) + + def iter_offload_moe_layers(model) -> Iterator: from freetoken.layers import BaseOP, OffloadMoELayer @@ -1075,3 +1613,18 @@ def attach_offload_moe_cache(model, cache: OffloadMoeCache) -> list: for layer in layers: layer.offload_cache = cache return layers + + +def attach_owner_moe_cache(model, owner_cache) -> list: + """Attach an ``OwnerOffloadMoeCache`` to every offload MoE layer. + + Sets ``layer.owner_cache`` and leaves ``layer.offload_cache`` pointing at the adapter. + This prevents bespoke subclasses from bypassing global-to-local route remapping through + the inner global-ID cache. Opt-in: nothing calls this unless a caller explicitly builds + an owner cache, and the global-ID cache path is untouched when it is absent. + """ + layers = list(iter_offload_moe_layers(model)) + for layer in layers: + layer.owner_cache = owner_cache + layer.offload_cache = owner_cache + return layers diff --git a/python/freetoken/moe/ownership.py b/python/freetoken/moe/ownership.py new file mode 100644 index 000000000..63a109394 --- /dev/null +++ b/python/freetoken/moe/ownership.py @@ -0,0 +1,569 @@ +"""Global-to-local routed-expert ownership for tensor-parallel expert sharding. + +This module deliberately knows nothing about cache slots or expert weights. A global router +emits IDs in ``[0, global_num_experts)``; an owner maps those IDs to a local bank row, while the +LRU cache later maps local rows to cache slots. Keeping the three namespaces separate avoids +feeding a remote global ID or a slot ID into a bank/GEMM kernel. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import torch + + +def same_device(left: torch.device, right: torch.device) -> bool: + """Compare devices without spuriously failing on an unspecified index. + + ``torch.device("cuda") != torch.device("cuda:0")`` even though tensors allocated on + either resolve to the same physical device. Cache adapters are often constructed with + the bare ``"cuda"`` string while route tensors carry ``cuda:0``; treat an unspecified + index as a match, but never let ``cuda:0`` and ``cuda:1`` compare equal. + """ + if left.type != right.type: + return False + if left.index is None or right.index is None: + return True + return left.index == right.index + + +@dataclass(frozen=True) +class OwnedRoute: + """A router decision masked for one expert owner. + + ``local_ids`` is always safe to use as a local-bank row index. Remote entries use row zero + as a harmless placeholder and have zero ``weights``; callers must not infer ownership from + the placeholder. The original router weights are never renormalized per rank. + """ + + local_ids: torch.Tensor + weights: torch.Tensor + owned_mask: torch.Tensor + + @property + def owned_count(self) -> int: + return int(self.owned_mask.sum().item()) + + @property + def remote_count(self) -> int: + return int((~self.owned_mask).sum().item()) + + +@dataclass(frozen=True) +class ExpertOwnership: + """Contiguous expert ownership for one rank in an EP group. + + The first implementation intentionally requires an even partition. An interleaved or + history-weighted owner map can be added later, but it must preserve the same explicit + global/local/slot boundary and be tested independently. + """ + + global_num_experts: int + world_size: int + rank: int + + def __post_init__(self) -> None: + if self.global_num_experts <= 0: + raise ValueError("global_num_experts must be positive") + if self.world_size <= 0: + raise ValueError("world_size must be positive") + if not 0 <= self.rank < self.world_size: + raise ValueError(f"rank {self.rank} is outside [0, {self.world_size})") + if self.global_num_experts % self.world_size: + raise ValueError( + f"global_num_experts={self.global_num_experts} is not divisible by " + f"world_size={self.world_size}" + ) + + @property + def local_num_experts(self) -> int: + return self.global_num_experts // self.world_size + + @property + def global_start(self) -> int: + return self.rank * self.local_num_experts + + @property + def global_end(self) -> int: + return self.global_start + self.local_num_experts + + def owns(self, expert_id: int) -> bool: + return self.global_start <= expert_id < self.global_end + + def owner(self, expert_id: int) -> int: + if not 0 <= expert_id < self.global_num_experts: + raise ValueError( + f"global expert id {expert_id} is outside [0, {self.global_num_experts})" + ) + return expert_id // self.local_num_experts + + def global_to_local(self, expert_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(local_ids, owned_mask)`` without changing route weights. + + ``local_ids`` is ``-1`` for remote entries. Consumers must use ``owned_mask`` before + indexing a local bank; no invalid sentinel may reach a GPU kernel. The returned mask + has the same shape as ``expert_ids`` and works on CPU or CUDA tensors. + """ + ids = expert_ids.to(dtype=torch.int64) + owned = (ids >= self.global_start) & (ids < self.global_end) + local = torch.where(owned, ids - self.global_start, torch.full_like(ids, -1)) + return local.to(dtype=expert_ids.dtype), owned + + def local_to_global(self, local_ids: torch.Tensor) -> torch.Tensor: + """Convert valid local bank rows back to global router IDs.""" + ids = local_ids.to(dtype=torch.int64) + if torch.any((ids < 0) | (ids >= self.local_num_experts)): + raise ValueError("local expert IDs must be in [0, local_num_experts)") + return (ids + self.global_start).to(dtype=local_ids.dtype) + + def validate_global_ids(self, expert_ids: torch.Tensor) -> None: + """Fail fast on malformed router IDs before owner masking/remapping.""" + ids = expert_ids.to(dtype=torch.int64) + if torch.any((ids < 0) | (ids >= self.global_num_experts)): + raise ValueError( + f"router expert IDs must be in [0, {self.global_num_experts})" + ) + + def partition_route( + self, weights: torch.Tensor, expert_ids: torch.Tensor + ) -> OwnedRoute: + """Mask a global top-k route for this owner without changing its probabilities. + + The returned route is an adapter contract, not a cache implementation: ``local_ids`` + address the owner's source-bank rows, while a future cache layer must create its own + local-row-to-slot mapping. This method deliberately preserves duplicate route entries + and the global top-k weights. In particular, a rank with zero local experts receives a + valid all-zero route rather than an invalid ``-1`` index. + """ + if weights.shape != expert_ids.shape: + raise ValueError( + f"route weights and expert IDs must have the same shape, got " + f"{tuple(weights.shape)} and {tuple(expert_ids.shape)}" + ) + if expert_ids.dtype not in (torch.int8, torch.int16, torch.int32, torch.int64): + raise TypeError(f"router expert IDs must be an integer tensor, got {expert_ids.dtype}") + self.validate_global_ids(expert_ids) + local_ids, owned = self.global_to_local(expert_ids) + safe_local_ids = torch.where(owned, local_ids, torch.zeros_like(local_ids)) + local_weights = torch.where(owned, weights, torch.zeros_like(weights)) + return OwnedRoute( + local_ids=safe_local_ids, + weights=local_weights, + owned_mask=owned, + ) + + +@dataclass(frozen=True) +class OwnerCacheGeometry: + """Explicit geometry contract for a future owner-local expert cache. + + The current :class:`~freetoken.moe.offload_cache.OffloadMoeCache` uses one global expert + namespace for route IDs, source-bank rows, and cache IDs. An EP owner cache must not reuse + that scalar for ``local_num_experts``: global router IDs first become local bank rows, and + only then become cache slots. This contract describes those dimensions without enabling the + runtime path prematurely. + + ``cache_size`` is the number of local expert slots on this rank. It is deliberately checked + against the local expert count, not the global count. Prefill overlap borrows two complete + local layers from the unified pool, hence its separate ``2 * local_num_experts`` minimum. + """ + + global_num_experts: int + world_size: int + rank: int + num_layers: int + cache_size: int + prefill_overlap: bool = False + + def __post_init__(self) -> None: + owner = ExpertOwnership(self.global_num_experts, self.world_size, self.rank) + if self.num_layers <= 0: + raise ValueError("num_layers must be positive") + if self.cache_size < owner.local_num_experts: + raise ValueError( + f"cache_size={self.cache_size} is smaller than local_num_experts=" + f"{owner.local_num_experts}" + ) + if self.prefill_overlap and self.cache_size < 2 * owner.local_num_experts: + raise ValueError( + "prefill_overlap requires cache_size >= 2 * local_num_experts " + f"({2 * owner.local_num_experts}), got {self.cache_size}" + ) + + @property + def ownership(self) -> ExpertOwnership: + return ExpertOwnership(self.global_num_experts, self.world_size, self.rank) + + @property + def local_num_experts(self) -> int: + return self.ownership.local_num_experts + + @property + def global_start(self) -> int: + return self.ownership.global_start + + @property + def global_end(self) -> int: + return self.ownership.global_end + + def validate_cache_binding( + self, + *, + num_layers: int, + num_experts: int, + cache_size: int, + prefill_overlap: bool, + ) -> None: + """Validate a cache constructor's global geometry before any allocation. + + ``num_experts`` is intentionally compared with the *global* dimension. A caller that + passes ``local_num_experts`` to the legacy cache will fail here instead of silently + corrupting layer-flat IDs or allocating a bank with the wrong row count. + """ + expected = { + "num_layers": (num_layers, self.num_layers), + "num_experts": (num_experts, self.global_num_experts), + "cache_size": (cache_size, self.cache_size), + } + for name, (actual, wanted) in expected.items(): + if actual != wanted: + raise ValueError( + f"owner cache geometry mismatch for {name}: got {actual}, expected {wanted}" + ) + if bool(prefill_overlap) != self.prefill_overlap: + raise ValueError( + "owner cache geometry mismatch for prefill_overlap: " + f"got {bool(prefill_overlap)}, expected {self.prefill_overlap}" + ) + + def global_to_local(self, expert_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Map global route IDs to local rows and return the ownership mask.""" + return self.ownership.global_to_local(expert_ids) + + def partition_route( + self, weights: torch.Tensor, expert_ids: torch.Tensor + ) -> OwnedRoute: + """Return this rank's safe local-row route without local renormalization.""" + return self.ownership.partition_route(weights, expert_ids) + + def local_to_flat_id(self, layer_id: int, local_ids: torch.Tensor) -> torch.Tensor: + """Map a local bank row to a layer-local cache ID namespace. + + This ID is only a contract for an owner-aware cache implementation. It must not be + passed to today's global-ID cache kernels, whose stride is the global expert count. + """ + if not 0 <= layer_id < self.num_layers: + raise ValueError(f"layer_id {layer_id} is outside [0, {self.num_layers})") + ids = local_ids.to(dtype=torch.int64) + if torch.any((ids < 0) | (ids >= self.local_num_experts)): + raise ValueError( + f"local expert IDs must be in [0, {self.local_num_experts})" + ) + return (layer_id * self.local_num_experts + ids).to(dtype=local_ids.dtype) + + def flat_id_to_local(self, flat_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Decode a local flat cache ID into ``(layer_id, local_expert_id)``. + + The inverse is intentionally defined only for the owner-local namespace. It is a + future cache adapter contract and must not be used to decode today's global-stride + ``OffloadMoeCache.id_of_slot`` values. + """ + ids = flat_ids.to(dtype=torch.int64) + total = self.num_layers * self.local_num_experts + if torch.any((ids < 0) | (ids >= total)): + raise ValueError(f"owner flat IDs must be in [0, {total})") + layers = torch.div(ids, self.local_num_experts, rounding_mode="floor") + local = torch.remainder(ids, self.local_num_experts) + return layers.to(dtype=flat_ids.dtype), local.to(dtype=flat_ids.dtype) + + def global_to_local_flat( + self, layer_id: int, expert_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Map global route IDs to safe owner-local flat IDs plus an ownership mask. + + Remote entries use the layer's local row zero as a safe placeholder before flattening; + callers must apply the returned mask/weights and must never infer ownership from the + placeholder. The resulting IDs use ``layer * local_num_experts + local_row`` and are + not valid inputs to the legacy global-stride cache kernels. + """ + local, owned = self.global_to_local(expert_ids) + safe_local = torch.where(owned, local, torch.zeros_like(local)) + return self.local_to_flat_id(layer_id, safe_local), owned + + def validate_source_banks( + self, sources: Mapping[str, Sequence[torch.Tensor]] + ) -> None: + """Check that every bank has one local-row tensor per layer. + + This accepts arbitrary bank names and trailing dimensions so it can validate BF16, + NVFP4, and future layouts without coupling the ownership contract to a quantizer. + """ + if not sources: + raise ValueError("owner source banks must not be empty") + for name, per_layer in sources.items(): + if len(per_layer) != self.num_layers: + raise ValueError( + f"bank {name!r} has {len(per_layer)} layers, expected {self.num_layers}" + ) + for layer_id, bank in enumerate(per_layer): + if bank.ndim == 0 or bank.shape[0] != self.local_num_experts: + got = tuple(bank.shape) if hasattr(bank, "shape") else type(bank).__name__ + raise ValueError( + f"bank {name!r} layer {layer_id} has row shape {got}; " + f"expected first dimension {self.local_num_experts}" + ) + + def validate_slot_maps( + self, slot_for_id: torch.Tensor, id_of_slot: torch.Tensor + ) -> None: + """Check the expected local-row-to-slot and reverse-map shapes.""" + expected_forward = (self.num_layers, self.local_num_experts) + if tuple(slot_for_id.shape) != expected_forward: + raise ValueError( + f"owner slot_for_id shape {tuple(slot_for_id.shape)} does not match " + f"{expected_forward}" + ) + if tuple(id_of_slot.shape) != (self.cache_size,): + raise ValueError( + f"owner id_of_slot shape {tuple(id_of_slot.shape)} does not match " + f"({self.cache_size},)" + ) + integer_types = (torch.int8, torch.int16, torch.int32, torch.int64) + if slot_for_id.dtype not in integer_types or id_of_slot.dtype not in integer_types: + raise TypeError("owner slot maps must use integer tensors") + for name, values, upper in ( + ("slot_for_id", slot_for_id, self.cache_size), + ("id_of_slot", id_of_slot, self.num_layers * self.local_num_experts), + ): + values = values.to(dtype=torch.int64) + if torch.any((values < -1) | (values >= upper)): + raise ValueError( + f"owner {name} entries must be -1 or in [0, {upper})" + ) + + +@dataclass(frozen=True) +class OwnerCacheUpdate: + """Result of one owner-local cache admission in the reference adapter. + + ``slot_ids`` is safe for an owner-local GEMM: remote route entries use slot zero and have + zero ``weights``. ``local_flat_ids`` is the local namespace used by the future cache + bookkeeping, not the global-stride ID understood by today's legacy cache. + """ + + slot_ids: torch.Tensor + local_ids: torch.Tensor + local_flat_ids: torch.Tensor + weights: torch.Tensor + owned_mask: torch.Tensor + missing_local_ids: torch.Tensor + evicted_flat_ids: torch.Tensor + + +class OwnerCacheAdapter: + """Small deterministic owner-local cache reference implementation. + + This adapter deliberately runs ordinary Python/Torch bookkeeping instead of flashlib or + Triton kernels. It is the executable P3 contract for the three namespaces and is intended + for tests and later GPU-kernel bring-up; it is not wired into the serving path yet. + """ + + def __init__( + self, geometry: OwnerCacheGeometry, device: torch.device | str = "cpu" + ) -> None: + self.geometry = geometry + self.device = torch.device(device) + self.slot_for_id = torch.full( + (geometry.num_layers, geometry.local_num_experts), + -1, + dtype=torch.int32, + device=self.device, + ) + self.id_of_slot = torch.full( + (geometry.cache_size,), + -1, + dtype=torch.int32, + device=self.device, + ) + self.usage = torch.zeros( + (geometry.cache_size,), dtype=torch.int64, device=self.device + ) + self.step = 0 + + @property + def cache_size(self) -> int: + return self.geometry.cache_size + + @property + def local_num_experts(self) -> int: + return self.geometry.local_num_experts + + @property + def resident(self) -> int: + return int(self.id_of_slot.ge(0).sum().item()) + + def _check_layer(self, layer_id: int) -> None: + if not 0 <= layer_id < self.geometry.num_layers: + raise ValueError( + f"layer_id {layer_id} is outside [0, {self.geometry.num_layers})" + ) + + def _check_input_device(self, *tensors: torch.Tensor) -> None: + for tensor in tensors: + if not same_device(tensor.device, self.device): + raise ValueError( + f"owner cache input is on {tensor.device}, expected {self.device}" + ) + + def _flat_id(self, layer_id: int, local_id: int) -> int: + return layer_id * self.local_num_experts + local_id + + def _clear_slot(self, slot: int) -> int: + old_flat = int(self.id_of_slot[slot].item()) + if old_flat < 0: + return old_flat + old_layer, old_local = divmod(old_flat, self.local_num_experts) + self.slot_for_id[old_layer, old_local] = -1 + self.id_of_slot[slot] = -1 + self.usage[slot] = 0 + return old_flat + + def _admit(self, layer_id: int, local_ids: list[int]) -> tuple[list[int], list[int]]: + """Admit one layer's unique local rows and return misses and evicted flat IDs.""" + self.step += 1 + step = self.step + active_slots: set[int] = set() + misses: list[int] = [] + for local_id in local_ids: + slot = int(self.slot_for_id[layer_id, local_id].item()) + if slot < 0: + misses.append(local_id) + else: + self.usage[slot] = step + active_slots.add(slot) + + evicted: list[int] = [] + reserved = set(active_slots) + for local_id in misses: + candidates = [slot for slot in range(self.cache_size) if slot not in reserved] + if not candidates: + # A single layer can address at most local_num_experts and geometry guarantees + # cache_size >= local_num_experts, so this is defensive rather than a normal + # route. Failing is safer than evicting a hit from the same admission. + raise RuntimeError("owner cache has no slot for an active local route") + slot = min(candidates, key=lambda value: (int(self.usage[value].item()), value)) + old_flat = self._clear_slot(slot) + if old_flat >= 0: + evicted.append(old_flat) + flat = self._flat_id(layer_id, local_id) + self.id_of_slot[slot] = flat + self.slot_for_id[layer_id, local_id] = slot + self.usage[slot] = step + reserved.add(slot) + return misses, evicted + + def ensure_route( + self, layer_id: int, weights: torch.Tensor, expert_ids: torch.Tensor + ) -> OwnerCacheUpdate: + """Admit a global route and rewrite owned entries to owner-local cache slots. + + The route weights remain globally normalized. Remote entries are masked to zero and + use slot zero as a safe placeholder, so no ``-1`` sentinel can reach a future GEMM. + """ + self._check_layer(layer_id) + self._check_input_device(weights, expert_ids) + route = self.geometry.partition_route(weights, expert_ids) + local_flat_ids, owned = self.geometry.global_to_local_flat(layer_id, expert_ids) + local_flat_ids = torch.where( + owned, + local_flat_ids, + torch.zeros_like(local_flat_ids), + ) + owned_local = route.local_ids.reshape(-1)[owned.reshape(-1)] + unique_local = list(dict.fromkeys(int(value) for value in owned_local.tolist())) + missing, evicted = self._admit(layer_id, unique_local) + + slots = self.slot_for_id[layer_id][route.local_ids.long()] + safe_slots = torch.where(owned, slots, torch.zeros_like(slots)) + update = OwnerCacheUpdate( + slot_ids=safe_slots, + local_ids=route.local_ids, + local_flat_ids=local_flat_ids, + weights=route.weights, + owned_mask=route.owned_mask, + missing_local_ids=torch.tensor( + missing, dtype=torch.int32, device=self.device + ), + evicted_flat_ids=torch.tensor( + evicted, dtype=torch.int32, device=self.device + ), + ) + self.validate_invariants() + return update + + def materialize_layer(self, layer_id: int, buffer_id: int = 0) -> torch.Tensor: + """Materialize every local expert of one layer into a contiguous slot buffer.""" + self._check_layer(layer_id) + if self.geometry.prefill_overlap: + if buffer_id not in (0, 1): + raise ValueError("owner prefill buffer_id must be 0 or 1") + elif buffer_id != 0: + raise ValueError("owner prefill overlap is disabled; buffer_id must be 0") + + first = buffer_id * self.local_num_experts + target = list(range(first, first + self.local_num_experts)) + for local_id in range(self.local_num_experts): + old_slot = int(self.slot_for_id[layer_id, local_id].item()) + if old_slot >= 0 and old_slot not in target: + self._clear_slot(old_slot) + for slot in target: + self._clear_slot(slot) + + self.step += 1 + for local_id, slot in enumerate(target): + flat = self._flat_id(layer_id, local_id) + self.id_of_slot[slot] = flat + self.slot_for_id[layer_id, local_id] = slot + self.usage[slot] = self.step + self.validate_invariants() + return torch.tensor(target, dtype=torch.int32, device=self.device) + + def reset(self) -> None: + self.slot_for_id.fill_(-1) + self.id_of_slot.fill_(-1) + self.usage.zero_() + self.step = 0 + + def validate_invariants(self) -> None: + """Verify forward/reverse maps are a bijection for all resident entries.""" + self.geometry.validate_slot_maps(self.slot_for_id, self.id_of_slot) + for layer_id in range(self.geometry.num_layers): + for local_id in range(self.local_num_experts): + slot = int(self.slot_for_id[layer_id, local_id].item()) + if slot < 0: + continue + flat = self._flat_id(layer_id, local_id) + if int(self.id_of_slot[slot].item()) != flat: + raise AssertionError( + f"owner cache forward map mismatch at layer={layer_id}, " + f"local_id={local_id}, slot={slot}" + ) + for slot, flat in enumerate(self.id_of_slot.tolist()): + if flat < 0: + continue + layer, local = divmod(flat, self.local_num_experts) + if int(self.slot_for_id[layer, local].item()) != slot: + raise AssertionError( + f"owner cache reverse map mismatch at slot={slot}, flat_id={flat}" + ) + + +__all__ = [ + "ExpertOwnership", + "OwnedRoute", + "OwnerCacheGeometry", + "OwnerCacheUpdate", + "OwnerCacheAdapter", +] diff --git a/python/freetoken/moe/route_trace.py b/python/freetoken/moe/route_trace.py new file mode 100644 index 000000000..22ded4d5e --- /dev/null +++ b/python/freetoken/moe/route_trace.py @@ -0,0 +1,272 @@ +"""Ordered MoE route trace: capture + offline LRU replay (PLAN_TP_EP.md 6A). + +Why this exists +--------------- +``--moe-collect-decode-freq`` only yields a per-(layer, expert) HISTOGRAM. A +histogram cannot reproduce LRU behaviour, adjacent-step overlap, or a miss +forecast, because it has thrown away the ORDER of expert activations. The EP2 +capacity decision (does doubling unique slots actually cut miss, and by how much +on the slow rank) needs the ordered sequence. + +This module records the RAW global expert ids in call order -- captured at the +same point ``collect_decode_freq`` snapshots them, i.e. BEFORE ``lru_ensure`` +rewrites ``expert_ids`` into slot ids in place -- and replays them offline +through an LRU that mirrors ``flashlib.kernels.slot_cache.lru_ensure`` semantics +(batch hit-protection, ``(usage, slot)`` victim order, ascending-id miss/victim +pairing, dedup). See ``test_route_trace.py`` for the semantics lock. + +Capture is host-side (a ``.cpu()`` per call), so it is NOT CUDA-graph safe: the +engine refuses ``--moe-trace-route`` unless ``--cuda-graph-max-bs 0``. That is the +intended use -- a correctness/sampling run on an experiment instance, never the +production perf path (plan 6A/8). + +On-disk format (append-only binary, crash-safe; partial data is always usable): + meta : ``.meta.json`` (num_experts, num_layers, cache_size, top_k, ...) + body : ```` repeated records, each + ``struct ' None: + # Every TP rank would otherwise open the SAME configured path, and ``wb`` truncates + # while the sidecar is rewritten too -- two writers racing on one body + one meta + # leaves a truncated trace holding a single rank's records. A rank suffix keeps them + # separate; ``None`` (single rank) leaves the configured path byte-for-byte. + if rank is not None: + path = f"{path}.rank{rank}" + self.path = path + self.meta_path = path + ".meta.json" + self.max_records = max_records + self._meta = RouteTraceMeta( + num_experts=num_experts, num_layers=num_layers, cache_size=cache_size, + top_k=top_k, model=model, decode_target=decode_target, + ) + os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) + # truncate any previous run so the body and meta stay consistent + self._f = open(path, "wb") + self._buf = bytearray() + self._n = 0 + self._overflow = False + self._flush_every = 8192 # records per buffer flush (bounds memory + I/O size) + self._write_meta() + + def _write_meta(self) -> None: + self._meta.records = self._n + self._meta.overflow = self._overflow + with open(self.meta_path, "w") as f: + json.dump(self._meta.__dict__, f, indent=1) + + def _flush(self) -> None: + if self._buf: + self._f.write(self._buf) + self._buf.clear() + self._f.flush() + + def record(self, layer_id: int, expert_ids, phase: int = 0) -> None: + """``expert_ids``: the RAW global ids tensor for this call (pre-rewrite).""" + if self._overflow: + return + if self.max_records and self._n >= self.max_records: + self._overflow = True + self._flush() + self._write_meta() + return + # int32 little-endian, host sync (NOT graph-safe; engine gates on graphs off) + import torch + + ids = expert_ids.reshape(-1).to(dtype=torch.int32).cpu().numpy() + self._buf += _RECORD_HDR.pack(phase, int(layer_id), int(ids.size)) + self._buf += ids.tobytes() + self._n += 1 + if self._n % self._flush_every == 0: # periodic flush so a crash keeps most data + self._flush() + self._write_meta() + + def close(self) -> None: + try: + self._flush() + self._f.close() + finally: + self._write_meta() + + def __del__(self): # best-effort; append-binary means data is already on disk + try: + self.close() + except Exception: + pass + + +def read_trace(path: str): + """Yield ``(phase, layer_id, ids_tuple)`` per record; return ``(meta, records)``. + + ``records`` is a list so the caller can replay repeatedly (multiple capacities). + """ + with open(path + ".meta.json") as f: + meta = RouteTraceMeta(**json.load(f)) + records = [] + with open(path, "rb") as f: + while True: + hdr = f.read(_RECORD_HDR.size) + if len(hdr) < _RECORD_HDR.size: + break + phase, layer_id, n = _RECORD_HDR.unpack(hdr) + raw = f.read(4 * n) + if len(raw) < 4 * n: + break # truncated tail from a crash -- keep what parsed + ids = tuple(struct.unpack(f"<{n}i", raw)) if n else () + records.append((phase, layer_id, ids)) + return meta, records + + +# ----------------------------------------------------------------- LRU mirror +class LRU: + """Unified-pool LRU mirroring ``flashlib.lru_ensure`` (see test_route_trace). + + One ``ensure`` == one kernel call == one LRU step: + * dedup ids; every HIT bumps ``usage[slot] = step`` first (batch protection: + a slot touched this call is never this call's victim). + * victims = ascending ``(usage, slot)``; the i-th ascending miss id takes the + i-th coldest victim. + """ + + def __init__(self, cache_size: int, num_experts: int): + import heapq + + self.C = cache_size + self.E = num_experts + self.slot_of: dict[int, int] = {} + self.owner: list[int | None] = [None] * cache_size + self.usage: list[int] = [0] * cache_size + self.heap: list[tuple[int, int]] = [] + self.free = list(range(cache_size)) + self.step = 0 + self.miss = 0 + self.active = 0 + self._h = heapq + + def _evict_one(self) -> int: + if self.free: + return self.free.pop() + while True: + u, s = self._h.heappop(self.heap) + if self.owner[s] is not None and self.usage[s] == u: + del self.slot_of[self.owner[s]] + self.owner[s] = None + return s + + def ensure(self, layer: int, ids) -> None: + self.step += 1 + base = layer * self.E + uniq = sorted(set(ids)) + self.active += len(uniq) + misses = [] + for e in uniq: + fid = base + e + s = self.slot_of.get(fid) + if s is None: + misses.append(e) + else: + self.usage[s] = self.step + self._h.heappush(self.heap, (self.step, s)) + self.miss += len(misses) + for e in misses: + s = self._evict_one() + fid = base + e + self.slot_of[fid] = s + self.owner[s] = fid + self.usage[s] = self.step + self._h.heappush(self.heap, (self.step, s)) + + +def replay(records, cache_size: int, num_experts: int, *, phase: int = 0, + ep_owner: tuple[int, int] | None = None): + """Replay ``records`` through one LRU pool. + + ``ep_owner=(lo, hi)``: keep only ids in ``[lo, hi)`` (remapped to ``id-lo``) -- + models one EP rank's local cache. ``None`` replays all ids (today's TP1). + Returns ``(miss, active, miss_rate)``. + """ + lru = LRU(cache_size, num_experts if ep_owner is None else (ep_owner[1] - ep_owner[0])) + for ph, layer, ids in records: + if ph != phase: + continue + if ep_owner is not None: + lo, hi = ep_owner + ids = tuple(e - lo for e in ids if lo <= e < hi) + lru.ensure(layer, ids) # empty local set still advances the LRU clock + rate = lru.miss / lru.active if lru.active else 0.0 + return lru.miss, lru.active, rate + + +def replay_ep2(records, cache_sizes: tuple[int, int], num_experts: int, *, phase: int = 0): + """Replay both EP ranks; report per-rank miss and the per-step SLOW side. + + The decode step waits on max(rank0, rank1) per layer, so the slow-side miss + (not the average) is what sets the step time. Returns a dict. + """ + half = num_experts // 2 + ranks = [LRU(cache_sizes[r], half) for r in range(2)] + slow_miss = slow_active = 0 + for ph, layer, ids in records: + if ph != phase: + continue + m0 = [r.miss for r in ranks] + a0 = [r.active for r in ranks] + for r in range(2): + lo = r * half + loc = tuple(e - lo for e in ids if lo <= e < lo + half) + ranks[r].ensure(layer, loc) + dm = [ranks[i].miss - m0[i] for i in range(2)] + da = [ranks[i].active - a0[i] for i in range(2)] + slow_miss += max(dm) + slow_active += max(da) + return { + "rank_miss": [r.miss for r in ranks], + "rank_active": [r.active for r in ranks], + "rank_rate": [r.miss / r.active if r.active else 0.0 for r in ranks], + "slow_miss": slow_miss, + "slow_active": slow_active, + "slow_rate": slow_miss / slow_active if slow_active else 0.0, + } diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b0..e09665a74 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -1,5 +1,6 @@ from __future__ import annotations +import time from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias import torch @@ -397,6 +398,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: mem = self._gpu_mem_bytes() mamba_used, mamba_total = mamba_slots or (0, 0) swa_used, swa_total = swa_tokens or (0, 0) + moe_stats = self._moe_stats_snapshot() for m in reply: m.kv_used_pages = used m.kv_total_pages = total @@ -405,6 +407,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: m.swa_used_tokens = swa_used m.swa_total_tokens = swa_total m.gpu_mem_bytes = mem + if moe_stats is not None: + m.moe_stats = moe_stats self.status_reporter.report_batch( batch, running_reqs=len(self.decode_manager.running_reqs), @@ -474,6 +478,27 @@ def _gpu_mem_bytes(self) -> int: return 0 return torch.cuda.memory_reserved(self.device) + def _moe_stats_snapshot(self) -> dict | None: + """Throttled MoE slot-cache snapshot for /v1/stats (miss/eviction/residency). + + None when the model has no offload MoE cache. stats_snapshot() syncs the + device once per counter group, so sampling is rate-limited to ~1/s; between + samples the frontend keeps the last-known value (same semantics as kv/mamba). + A failing snapshot must never break the reply stream.""" + cache = getattr(self.engine, "moe_offload_cache", None) + if cache is None: + return None + now = time.monotonic() + if now - getattr(self, "_moe_stats_last_at", 0.0) < 1.0: + return None + try: + snap = cache.stats_snapshot() + except Exception as e: # noqa: BLE001 -- observability must not break serving + logger.warning(f"moe stats snapshot failed: {e!r}") + return None + self._moe_stats_last_at = now + return snap + def _process_one_msg(self, msg: BaseBackendMsg) -> None: if isinstance(msg, BatchBackendMsg): for msg in msg.data: diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index ffd717280..1bd5a6f26 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -134,10 +134,16 @@ class ModelCard(BaseModel): created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "FreeToken" root: str - # The model's own limit, not the KV budget in force. Two spellings of the same number: - # `max_model_len` is vLLM/SGLang's, `context_length` what most other clients look for. + # What the server will actually ADMIT, i.e. min(model max_position, KV pool tokens). This is + # what a client must size its own window against: the scheduler rejects a request whose + # prompt reaches this number and clamps the output budget to the remainder. Two spellings of + # the same number: `max_model_len` is vLLM/SGLang's, `context_length` what most other + # clients look for. max_model_len: int | None = None context_length: int | None = None + # The checkpoint's own ceiling (config max_position), for reference. Larger than + # max_model_len whenever the KV pool was configured below the model's maximum. + model_max_len: int | None = None # The checkpoint's probed effort vocabulary (freetoken.tokenizer.effort); None # (not []) when the model has no effort knob or the probe could not run. supported_reasoning_efforts: list[str] | None = None diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..98ab7eb5e 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -174,6 +174,11 @@ class FrontendManager: # "num_mamba_slots"}, from the same ack. Seeds geometry before the first generation reply # (the running snapshot channel) has anything. None until meta arrives. cache_pools: Dict[str, int] | None = None + # Effective context ceiling the scheduler enforces, min(model max_position, KV pool + # tokens), from the same ack. /v1/models and /v1/stats report this so a client cannot + # size its window above what the server will actually admit; 0 until meta arrives (the + # metadata routes then fall back to the model's own ceiling). + max_seq_len: int = 0 # one {index, name, uuid, total_bytes} per TP rank, from the same ack; /v1/stats gpus gpus: List[Dict[str, Any]] = field(default_factory=list) # Backend worker Process handles (TP schedulers + tokenizer/detokenizer), captured from the @@ -477,6 +482,32 @@ async def _record_request_middleware(request: Request, call_next): return response +# Minimal bearer auth for LAN exposure (`ft serve --host 0.0.0.0`): loopback callers +# (desktop app, ft shell/ctl on this machine) are trusted; everyone else must send +# `Authorization: Bearer $FREETOKEN_API_KEY` (OpenAI-style) or `x-api-key: +# $FREETOKEN_API_KEY` (Anthropic-style — Copilot's "messages" apiType sends this). +# Unset/empty env = auth disabled. +# The X-Forwarded-For guard matters: uvicorn's proxy-headers middleware rewrites +# request.client from that header, so a remote client could otherwise spoof 127.0.0.1. +_API_KEY = os.environ.get("FREETOKEN_API_KEY", "").strip() + + +@app.middleware("http") +async def _api_key_auth(request: Request, call_next): + if not _API_KEY: + return await call_next(request) + client = request.client.host if request.client else "" + loopback = "x-forwarded-for" not in request.headers and client in ("127.0.0.1", "::1") + authorized = ( + loopback + or request.headers.get("authorization", "") == f"Bearer {_API_KEY}" + or request.headers.get("x-api-key", "") == _API_KEY + ) + if authorized: + return await call_next(request) + return JSONResponse({"error": "unauthorized"}, status_code=401) + + class CacheRebuildRequest(BaseModel): moe_cache_size: int | None = None num_pages: int | None = None @@ -1010,6 +1041,7 @@ def _on_meta(meta: dict) -> None: _GLOBAL_STATE.free_vram_bytes = int(meta.pop("free_vram_bytes", 0) or 0) _GLOBAL_STATE.cache_floors = meta.pop("floors", None) _GLOBAL_STATE.cache_pools = meta.pop("pools", None) + _GLOBAL_STATE.max_seq_len = int(meta.pop("max_seq_len", 0) or 0) _GLOBAL_STATE.swa_full_tokens_ratio = float(meta.pop("swa_full_tokens_ratio", 0.0) or 0.0) _GLOBAL_STATE.cache_budget_bytes = int(meta.pop("cache_budget_bytes", 0) or 0) _GLOBAL_STATE.gpus = list(meta.pop("gpus", None) or []) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 32e8d1266..b178bf9d8 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -264,6 +264,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The tensor parallelism size.", ) + parser.add_argument( + "--moe-ep-size", + type=int, + default=1, + help=( + "Routed-expert owner group size. Default 1 keeps the legacy global-ID MoE cache; " + "values >1 are an explicit TP+EP opt-in and currently require the owner runtime " + "to be supported by the selected model." + ), + ) + parser.add_argument( "--gpu", type=_lazy_gpu_arg, @@ -677,6 +688,51 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-collect-stats", + action="store_true", + dest="moe_collect_stats", + default=ServerArgs.moe_collect_stats, + help=( + "Enable the MoE offload cache's decode miss-rate counters (experts " + "active/missing/fetched per layer per step, captured into the decode " + "CUDA graph so they keep counting during replay). Cheap (device-side, " + "no per-step host sync); surfaces under /v1/stats 'moe' with the " + "residency/miss-rate numbers that explain decode-speed dips. Off by " + "default." + ), + ) + + parser.add_argument( + "--moe-collect-decode-freq", + action="store_true", + dest="moe_collect_decode_freq", + default=ServerArgs.moe_collect_decode_freq, + help=( + "Also collect the per-(layer, expert) decode routing histogram, giving " + "/v1/stats 'moe.routing' (working-set size, experts-to-90%%-of-activations, " + "oracle hit rate at the current cache size). Answers 'is it always the " + "same experts firing?'. Only accurate with CUDA graphs off " + "(--cuda-graph-max-bs 0): a captured graph replays without the host-side " + "histogram scatter. Implies --moe-collect-stats is NOT required, set both " + "for the full picture." + ), + ) + + parser.add_argument( + "--moe-trace-route", + type=str, + default=ServerArgs.moe_trace_route, + help=( + "Path to write an ORDERED MoE route trace: every ensure_experts call appends " + "its raw global expert ids (pre slot-rewrite) for offline LRU/EP replay " + "(freetoken/moe/route_trace.py, and tools/trace/replay_route_trace.py in the " + "deployment repo). Host-side, so NOT CUDA-graph " + "safe -- requires --cuda-graph-max-bs 0. Sampling/diagnostic only; never the " + "production perf path. Off by default." + ), + ) + parser.add_argument( "--shell-mode", action="store_true", @@ -705,6 +761,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None: f"{kwargs['tensor_parallel_size']}; give one entry per TP rank" ) + if kwargs["moe_ep_size"] < 1: + parser.error("--moe-ep-size must be >= 1") + if kwargs["moe_ep_size"] > 1 and kwargs["moe_ep_size"] != kwargs["tensor_parallel_size"]: + parser.error( + "--moe-ep-size must equal --tensor-parallel-size for the initial same-group TP+EP " + "topology" + ) + # resolve some arguments run_shell |= kwargs.pop("shell_mode") kwargs["shell_mode"] = run_shell diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..14c92d7d9 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -139,6 +139,7 @@ async def v1_models(): root=state.config.model_path, max_model_len=ctx, context_length=ctx, + model_max_len=_checkpoint_context_length(state), supported_reasoning_efforts=efforts, default_reasoning_effort=default_effort, )]) @@ -677,9 +678,34 @@ def _served_model_name(state: Any) -> str: def _model_context_length(state: Any) -> int | None: - """The model ceiling, not `min(ceiling, KV budget)`: a rebuild moves the latter, and agents - read this once at startup.""" + """The context a client should budget for: min(model ceiling, KV pool tokens). + + Deliberately NOT the checkpoint's own ceiling. Clients size their window from this route, + so advertising more than the scheduler admits produces hard 400s + (``context_length_exceeded``: "prompt is too long: N tokens > M maximum (prompt + + generation)") on prompts the model card said were fine -- exactly what a KV pool configured + below the model's max_position used to cause. The engine publishes the enforced value in its + readiness meta; while that is still in flight, or on an engine that sends none, fall back to + the ceiling. ``model_max_len`` on the card keeps the raw ceiling visible. + """ + ceiling = 0 try: # never 500 a metadata route: max_seq_len walks into the HF config on some builds + ceiling = int(state.config.max_seq_len) + except Exception: # noqa: BLE001 + ceiling = 0 + enforced = 0 + try: + enforced = int(getattr(state, "max_seq_len", 0) or 0) + except (TypeError, ValueError): + enforced = 0 + value = enforced if enforced > 0 else ceiling + return value if value > 0 else None + + +def _checkpoint_context_length(state: Any) -> int | None: + """The checkpoint's OWN ceiling (config max_position), reported for reference next to the + enforced ``max_model_len``/``context_length``. None when it cannot be read.""" + try: value = int(state.config.max_seq_len) except Exception: # noqa: BLE001 return None diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 76c6c308a..57d271536 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -30,6 +30,14 @@ def __init__(self, window_s: float = 5.0) -> None: # "cost saved by running locally" accounting. Monotonic; resets when the process restarts. self.prompt_tokens_total = 0 self.completion_tokens_total = 0 + # Lifetime prefix-cache hits: tokens of admitted prompts served from the radix + # cache instead of recomputed. Arrives per request on the same reply as + # prompt_tokens_delta (UserReply.cached_tokens); hit_ratio = this / + # prompt_tokens_total (same tokens denominator as vLLM's hits/queries pair). + self.cached_tokens_total = 0 + # Last MoE slot-cache snapshot stamped by the scheduler (miss/residency/routing + # concentration); None until the first sample or on non-offload models. + self.moe_stats: dict | None = None self.kv_used_pages = 0 self.kv_total_pages = 0 self.mamba_used_slots = 0 @@ -63,6 +71,10 @@ def observe(self, reply: Any, now: float | None = None) -> None: if getattr(reply, "prompt_tokens_delta", 0) > 0: self._prefill.append((t, reply.prompt_tokens_delta)) self.prompt_tokens_total += reply.prompt_tokens_delta + if getattr(reply, "cached_tokens", 0) > 0: + self.cached_tokens_total += reply.cached_tokens + if getattr(reply, "moe_stats", None) is not None: + self.moe_stats = reply.moe_stats if getattr(reply, "kv_total_pages", 0) > 0: # ignore 0/0 (prompt reply, owned-KV) self.kv_used_pages = reply.kv_used_pages self.kv_total_pages = reply.kv_total_pages @@ -118,6 +130,24 @@ def derive_model_card(config: Any) -> dict: } +def _resolved_page_size(state: Any, config: Any) -> int: + """The engine's REAL KV page size (tokens per page). + + ``_adjust_config`` runs inside the scheduler process, so the frontend's ``config.page_size`` + can still hold the CLI default while the engine actually pages at, say, 64 (qsa_sparse). + Reporting that default made ``total_pages`` look like a token count (4096 "tokens" for a + 262144-token pool). Prefer the value the engine published in its readiness meta. + """ + pools = getattr(state, "cache_pools", None) or {} + try: + value = int(pools.get("page_size") or 0) + except (TypeError, ValueError): + value = 0 + if value <= 0: + value = int(getattr(config, "page_size", 1) or 1) + return max(1, value) + + def _swa_page_size(config: Any) -> int: """The window pool's own page unit: P (window_size) for DSV4, 1 token for radix-SWA. Mirrors compute_cache_pools' swa_page_size.""" @@ -139,7 +169,7 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: uptime_s = max(0, int(time.monotonic() - ready_at)) if ready_at is not None else 0 kv = ( {"used_pages": tr.kv_used_pages, "total_pages": tr.kv_total_pages, - "page_size": getattr(config, "page_size", 1)} + "page_size": _resolved_page_size(state, config)} if tr.kv_total_pages > 0 else None ) mamba = ( @@ -147,18 +177,45 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: if tr.mamba_total_slots > 0 else None ) sps = _swa_page_size(config) + try: + model_max_seq_len = int(getattr(config, "max_seq_len", 0) or 0) + except (TypeError, ValueError): + model_max_seq_len = 0 + # Same expression the scheduler's admission check uses; fall back to the model ceiling + # while the readiness meta is still in flight. + try: + effective_max_seq_len = int(getattr(state, "max_seq_len", 0) or 0) + except (TypeError, ValueError): + effective_max_seq_len = 0 + if effective_max_seq_len <= 0: + effective_max_seq_len = model_max_seq_len swa = ( {"used_pages": tr.swa_used_tokens // sps, "total_pages": tr.swa_total_tokens // sps, "page_size": sps} if tr.swa_total_tokens > 0 else None ) + model_card = derive_model_card(config) + # ``model.ctx`` must be the limit the scheduler ENFORCES, not the checkpoint ceiling: + # ``launch._stats_context_length`` reads exactly this field as its fallback when sizing a + # client's context window, so leaving the raw ceiling here would still let a client send + # prompts the scheduler rejects whenever the KV pool is smaller than max_position. The + # raw ceiling stays available in ``limits.model_max_seq_len``. + model_card["ctx"] = effective_max_seq_len return { "instance_id": getattr(state, "instance_id", None), - "model": derive_model_card(config), + "model": model_card, "uptime_s": uptime_s, "kv": kv, "mamba": mamba, "swa": swa, + # What the scheduler actually ADMITS. prompt_tokens must stay under max_seq_len (the + # output budget is clamped to the remainder), and it is min(model max_position, KV + # pool tokens) -- so it can be smaller than the model's own ceiling when the KV pool + # was configured below it. model_max_seq_len is that ceiling, for reference. + "limits": { + "max_seq_len": effective_max_seq_len, + "model_max_seq_len": model_max_seq_len, + }, "vram_bytes": tr.vram_bytes, "gpus": list(getattr(state, "gpus", None) or []), "throughput": { @@ -173,4 +230,14 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "prompt_tokens_total": tr.prompt_tokens_total, "completion_tokens_total": tr.completion_tokens_total, }, + "prefix_cache": { + "cached_tokens_total": tr.cached_tokens_total, + "prompt_tokens_total": tr.prompt_tokens_total, + "hit_ratio": ( + round(tr.cached_tokens_total / tr.prompt_tokens_total, 4) + if tr.prompt_tokens_total + else 0.0 + ), + }, + "moe": tr.moe_stats, } diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862d0..252af80ed 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -223,6 +223,7 @@ def tokenize_worker( swa_used_tokens=msg.swa_used_tokens, swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, + moe_stats=msg.moe_stats, ) for msg, reply in zip(detokenize_msg, replies, strict=True) ] diff --git a/tests/engine/test_owner_ep_config.py b/tests/engine/test_owner_ep_config.py new file mode 100644 index 000000000..67d0d948e --- /dev/null +++ b/tests/engine/test_owner_ep_config.py @@ -0,0 +1,90 @@ +"""``_validate_owner_ep_config`` must reject what the owner runtime cannot honour. + +Two settings used to be accepted and then silently ignored: + +* ``--moe-cpu-layers`` -- ``_decode_owner`` is selected before the ``is_cpu_layer`` + branch, and ``OwnerOffloadMoeCache`` hard-codes a GPU inner cache, so the flag + had no effect at all. +* FTW checkpoints -- ``load_ftw_banks`` rebuilds ``[num_experts, ...]`` GLOBAL + expert rows with no ownership filter, so the banks cannot bind to the + owner-local geometry. + +Fail-fast is the contract for owner EP (``moe_ep_size > 1`` is opt-in and +validated before any allocation), so both must raise rather than serve a +different configuration than the one requested. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from freetoken.engine.engine import _validate_owner_ep_config + + +def _config(**over): + base = dict( + moe_ep_size=2, + tp_info=SimpleNamespace(size=2, rank=0), + moe_strategy="offload", + moe_cache_rate=None, + moe_cache_size=4096, + moe_cache_auto=False, + moe_cpu_layers=None, + model_path="/models/unit-model", + ) + base.update(over) + return SimpleNamespace(**base) + + +@pytest.fixture(autouse=True) +def _not_ftw(monkeypatch): + monkeypatch.setattr("freetoken.checkpoint.ftw.is_ftw_checkpoint", lambda path: False) + + +def test_a_valid_owner_topology_passes(): + _validate_owner_ep_config(_config()) + + +def test_moe_cache_auto_satisfies_the_cache_requirement(): + _validate_owner_ep_config(_config(moe_cache_size=0, moe_cache_auto=True)) + + +def test_ep_size_one_is_a_no_op(monkeypatch): + monkeypatch.setattr( + "freetoken.checkpoint.ftw.is_ftw_checkpoint", + lambda path: pytest.fail("must not probe the checkpoint when EP is off"), + ) + _validate_owner_ep_config(_config(moe_ep_size=1)) + + +@pytest.mark.parametrize("layers", ["0", "0,1", "0-3"]) +def test_cpu_layers_are_rejected_instead_of_silently_ignored(layers): + with pytest.raises(ValueError, match="moe-cpu-layers"): + _validate_owner_ep_config(_config(moe_cpu_layers=layers)) + + +def test_ftw_checkpoints_are_rejected(monkeypatch): + monkeypatch.setattr("freetoken.checkpoint.ftw.is_ftw_checkpoint", lambda path: True) + with pytest.raises(ValueError, match="FTW"): + _validate_owner_ep_config(_config()) + + +def test_an_explicit_cache_size_is_still_required_without_auto(): + with pytest.raises(ValueError, match="moe-cache-size"): + _validate_owner_ep_config(_config(moe_cache_size=0)) + + +def test_moe_cache_rate_is_still_rejected(): + with pytest.raises(ValueError, match="moe-cache-rate"): + _validate_owner_ep_config(_config(moe_cache_rate=0.5)) + + +def test_a_non_offload_strategy_is_still_rejected(): + with pytest.raises(ValueError, match="offload"): + _validate_owner_ep_config(_config(moe_strategy="resident")) + + +def test_a_topology_other_than_tp2_ep2_is_still_rejected(): + with pytest.raises(ValueError, match="TP2"): + _validate_owner_ep_config(_config(tp_info=SimpleNamespace(size=4, rank=0))) diff --git a/tests/kvcache/test_linear_state_pool_alloc.py b/tests/kvcache/test_linear_state_pool_alloc.py index e9fe40542..9f59a3afa 100644 --- a/tests/kvcache/test_linear_state_pool_alloc.py +++ b/tests/kvcache/test_linear_state_pool_alloc.py @@ -189,3 +189,19 @@ def test_slot_state_bytes_for_the_real_geometry(): delta = linear_state_bytes_per_req(group, 1, torch.bfloat16, (spec,)) - \ linear_state_bytes_per_req(group, 1, torch.bfloat16) assert delta == 4 * 2560 * 9 * 2 == 180 * 1024 + + +def test_tp2_state_geometry_matches_byte_accounting(): + group = _group() + pool = LinearStatePool( + group=group, + num_slots=5, + dtype=torch.bfloat16, + device=torch.device("cpu"), + tp_size=2, + ) + # 2 key heads -> 1, 4 value heads -> 2 at TP2. + assert pool.conv_states.shape == (2, 5, 64, 3) + assert pool.recurrent_states.shape == (2, 5, 2, 16, 16) + expected = linear_state_bytes_per_req(group, 2, torch.bfloat16) + assert pool.bytes_per_slot() == expected diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index b81e02f4d..7e4e730e7 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -6,7 +6,7 @@ from freetoken.attention import AttnType from freetoken.models.config import FullAttentionGroupConfig, LinearGatedDeltaGroupConfig -from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.config import parse_config, qwen4_exp_tp_geometry from .common import LOVEDHEART_NVFP4_FP8, NVIDIA_NVFP4, QWEN_FP8, RADIXARK_NVFP4 @@ -132,6 +132,27 @@ def test_qwen4_args_payload(): assert args.ngram_boundary_token_id == 248044 +def test_tp2_geometry_is_local_but_model_config_stays_global(): + cfg = parse_config(_hf_config()) + local = qwen4_exp_tp_geometry(cfg, tp_size=2, rank=1) + assert (local.num_q_heads, local.num_kv_heads) == (12, 1) + assert (local.num_key_heads, local.num_value_heads) == (8, 24) + assert local.q_attn_dim == 12 * 256 + assert local.kv_attn_dim == 256 + assert local.conv_dim == 2 * 8 * 128 + 24 * 128 + assert local.local_conv_dim == local.conv_dim + assert local.local_recurrent_state_shape == (24, 128, 128) + assert cfg.num_qo_heads == 24 and cfg.num_kv_heads == 2 + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_tp_geometry_rejects_non_divisible_dense_heads(tp_size): + cfg = parse_config(_hf_config()) + geometry = qwen4_exp_tp_geometry(cfg, tp_size=tp_size, rank=0) + assert geometry.tp_size == tp_size + assert geometry.num_q_heads * tp_size == cfg.num_qo_heads + + def test_ple_on_full_attention_layer_rejected(): hf = _hf_config() hf.text_config.ple_layer_ids = [4] # one-indexed 4 == zero-based 3, a full_attention layer diff --git a/tests/models/qwen4_exp/test_gdn.py b/tests/models/qwen4_exp/test_gdn.py index 81dd0e767..873223c60 100644 --- a/tests/models/qwen4_exp/test_gdn.py +++ b/tests/models/qwen4_exp/test_gdn.py @@ -8,6 +8,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -68,7 +70,7 @@ def _make_layer(ratio: int, output_gate: str = "sigmoid", seed: int = 0): return op, ref -def _ctx(ratio: int, num_slots: int = 8) -> Context: +def _ctx(ratio: int, num_slots: int = 8, tp_size: int = 1) -> Context: import freetoken.core as core from freetoken.kvcache.linear_state_pool import LinearStatePool @@ -80,7 +82,9 @@ def _ctx(ratio: int, num_slots: int = 8) -> Context: ) core._GLOBAL_CTX = None ctx = Context(page_size=64) - ctx.linear_state_pool = LinearStatePool(group, num_slots, torch.bfloat16, DEV, tp_size=1) + ctx.linear_state_pool = LinearStatePool( + group, num_slots, torch.bfloat16, DEV, tp_size=tp_size + ) core.set_global_ctx(ctx) return ctx @@ -149,6 +153,119 @@ def test_ragged_prefill_then_decode(ratio): torch.testing.assert_close(dec[i].float(), full[-1], rtol=RTOL, atol=ATOL) +@pytest.mark.parametrize("ratio", (2, 3)) +def test_tp2_rank_partials_sum_to_the_tp1_output(ratio, monkeypatch): + """P2 numeric gate: a TP2 rank's LOCAL GDN forward, summed over both ranks, equals the + TP1 output. Each rank runs the real kernel with its own head slice and its own state + pool; only the row-parallel all-reduce is emulated (it is a plain SUM).""" + import freetoken.distributed.info as info + from freetoken.distributed import DistributedCommunicator + from freetoken.models.qwen4_exp.gdn import Qwen4ExpGatedDeltaNet + from freetoken.utils.torch_utils import torch_dtype + + num_k, num_v = HEADS[ratio] + torch.manual_seed(ratio) + hidden = torch.randn(37, HIDDEN, device=DEV, dtype=torch.bfloat16) + + # TP1 reference: one op holding the full head set. + ref_op, _ref = _make_layer(ratio, seed=ratio) + ref_ctx = _ctx(ratio, tp_size=1) + with ref_ctx.forward_batch(_batch([37])): + want = ref_op.forward(hidden) + + # TP2: the same checkpoint weights sharded per rank, then each rank's local forward. + full_state = {k: v.clone() for k, v in ref_op.state_dict().items()} + config = SimpleNamespace( + linear_attention_group=lambda: LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0,), num_key_heads=num_k, num_value_heads=num_v, + key_head_dim=HEAD_DIM, value_head_dim=HEAD_DIM, conv_kernel_dim=CONV_K, + output_gate="sigmoid", + ) + ) + from freetoken.models.qwen4_exp.weight import shard_qwen4_exp_dense_tensor + + # The row-parallel out_proj would all-reduce in production; here we keep each rank's + # UNREDUCED partial and sum them ourselves, which is exactly the TP contract. + monkeypatch.setattr(DistributedCommunicator, "all_reduce", lambda self, x: x) + + partials = [] + for rank in range(2): + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=rank, size=2)) + # Local dims: the op derives them from get_tp_info() at construction. + local_k = num_k // 2 + local_v = num_v // 2 + with torch.device("meta"), torch_dtype(torch.bfloat16): + op = Qwen4ExpGatedDeltaNet( + hidden_size=HIDDEN, num_k_heads=num_k, num_v_heads=num_v, + head_k_dim=HEAD_DIM, head_v_dim=HEAD_DIM, conv_kernel_size=CONV_K, + rms_norm_eps=EPS, layer_id=0, output_gate="sigmoid", + ) + # Shard every raw checkpoint tensor with the tested P2 contract. The op's fused + # in_proj is rebuilt by slicing its parts (qkv|z|b|a) separately. + sharded = {} + prefix = "model.layers.0." + # TP1 in_proj layout: q|k|v (2*K*D + V*D) | z (V*D) | b (V) | a (V). + qkv_rows = 2 * num_k * HEAD_DIM + num_v * HEAD_DIM + z_rows = num_v * HEAD_DIM + for key, value in full_state.items(): + if key == "in_proj.weight": + qkv = shard_qwen4_exp_dense_tensor( + prefix + "linear_attn.in_proj_qkv.weight", + value[:qkv_rows], + config=config, rank=rank, world_size=2, + ) + z = shard_qwen4_exp_dense_tensor( + prefix + "linear_attn.in_proj_z.weight", + value[qkv_rows : qkv_rows + z_rows], + config=config, rank=rank, world_size=2, + ) + b = shard_qwen4_exp_dense_tensor( + prefix + "linear_attn.in_proj_b.weight", + value[qkv_rows + z_rows : qkv_rows + z_rows + num_v], + config=config, rank=rank, world_size=2, + ) + a = shard_qwen4_exp_dense_tensor( + prefix + "linear_attn.in_proj_a.weight", + value[qkv_rows + z_rows + num_v :], + config=config, rank=rank, world_size=2, + ) + sharded[key] = torch.cat([qkv, z, b, a], dim=0) + continue + key_name = { + "conv1d.weight": "linear_attn.conv1d.weight", + "A_log": "linear_attn.A_log", + "dt_bias": "linear_attn.dt_bias", + "norm.weight": "linear_attn.norm.weight", + "out_proj.weight": "linear_attn.out_proj.weight", + }[key] + sharded[key] = shard_qwen4_exp_dense_tensor( + prefix + key_name, value, config=config, rank=rank, world_size=2 + ) + op.load_state_dict(sharded) + assert op.num_k_heads == local_k and op.num_v_heads == local_v + assert op.conv_dim == 2 * local_k * HEAD_DIM + local_v * HEAD_DIM + + ctx = _ctx(ratio, tp_size=2) + with ctx.forward_batch(_batch([37])): + partials.append(op.forward(hidden)) + + # Row-parallel out_proj emits an UNREDUCED local partial; the SUM is the TP contract. + merged = partials[0] + partials[1] + torch.testing.assert_close(merged.float(), want.float(), rtol=RTOL, atol=ATOL) + assert DistributedCommunicator.plugins[-1] is not None + + +def _batch(lengths): + reqs = [ + Req(input_ids=torch.zeros(n, dtype=torch.int32), table_idx=i + 1, cached_len=0, + output_len=1, uid=i, sampling_params=SamplingParams(), cache_handle=None) + for i, n in enumerate(lengths) + ] + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = reqs + return batch + + def test_chunk_and_recurrent_rules_agree(): """The chunked form (what the fla prefill kernel implements) against the sequential definition, both fp32: the chunk oracle is only worth anything if it reproduces the diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 1d3b944ce..3cea2a38e 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -248,3 +248,107 @@ def test_two_qsa_layers_keep_separate_slab_slots(monkeypatch): slab = fixture.pool.cmp_k_cache assert not torch.equal(slab(0), slab(1)) + + +def _unfuse_qkv(fused: torch.Tensor, config) -> dict[str, torch.Tensor]: + """Split the model's fused ``qkv_proj`` buffer back into the checkpoint's raw q/k/v keys. + + The fused order is ``[2*qo*head_dim | kv*head_dim | kv*head_dim]`` (q carries the gate), + so the split is exact for both TP1 and a rank-local buffer. + """ + qo = config.num_qo_heads * config.head_dim + kv = config.num_kv_heads * config.head_dim + q, k, v = fused.split([2 * qo, kv, kv], dim=0) + return { + f"layers.{QSA_LAYER}.self_attn.q_proj.weight": q, + f"layers.{QSA_LAYER}.self_attn.k_proj.weight": k, + f"layers.{QSA_LAYER}.self_attn.v_proj.weight": v, + } + + +@requires_cuda +def test_tp2_rank_partials_sum_to_the_tp1_output(monkeypatch): + """P2 numeric gate for QSA: two TP2 ranks' local layer outputs sum to the TP1 output, + and both ranks must select the SAME blocks (the indexer and its compressed slab are + replicated, so a split selection would corrupt the all-reduced result). + + Each rank runs the real backend over its own rank-local K/V slab; only the row-parallel + ``o_proj`` all-reduce is replaced by an identity so the unreduced partials can be summed + here, which is exactly the TP contract. + """ + import freetoken.distributed.info as info + from freetoken.distributed import DistributedCommunicator + from freetoken.models.qwen4_exp.weight import shard_qwen4_exp_dense_tensor + + lengths = [2051, 1000, 137] # dense regime: every complete block is selected + req_meta = [(i, 0, n) for i, n in enumerate(lengths)] + + # ---- TP1 reference ------------------------------------------------------------- + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=0, size=1)) + config1 = parsed_config() + fix1 = Fixture(config1, num_pages=128) + attn1 = fix1.layer(QSA_LAYER) + x = torch.cat(_inputs(fix1, lengths)) + seen1 = selection_spy(monkeypatch, fix1.backend) + batch1 = fix1.batch([fix1.req(*meta) for meta in req_meta], "prefill") + want = attn1.forward(x, batch1).clone() + want_idx = seen1["indices"].clone() + full_state = {k: v.detach().clone() for k, v in attn1.state_dict().items()} + assert config1.num_qo_heads == 4 and config1.num_kv_heads == 2 # global stays global + + # ---- TP2: two rank-local runs --------------------------------------------------- + monkeypatch.setattr(DistributedCommunicator, "all_reduce", lambda self, x: x) + partials, selections, slab_kv_heads = [], [], [] + for rank in range(2): + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=rank, size=2)) + config2 = parsed_config() # geometry resolves rank-local via get_tp_info() + fix2 = Fixture(config2, num_pages=128) + attn2 = fix2.layer(QSA_LAYER) + assert attn2.num_q == 2 and attn2.num_kv == 1, "rank-local heads not applied" + + # Rank-local weights, built the way the real loader does it: shard the RAW q/k/v + # separately (so head groups stay intact), then re-fuse them into `qkv_proj`. + q, k, v = _unfuse_qkv(full_state["qkv_proj.weight"], config1).values() + parts = [ + shard_qwen4_exp_dense_tensor( + name, tensor, config=config1, rank=rank, world_size=2 + ) + for name, tensor in ( + (f"layers.{QSA_LAYER}.self_attn.q_proj.weight", q), + (f"layers.{QSA_LAYER}.self_attn.k_proj.weight", k), + (f"layers.{QSA_LAYER}.self_attn.v_proj.weight", v), + ) + ] + raw = { + "qkv_proj.weight": torch.cat(parts, dim=0), + "o_proj.weight": shard_qwen4_exp_dense_tensor( + f"layers.{QSA_LAYER}.self_attn.o_proj.weight", full_state["o_proj.weight"], + config=config1, rank=rank, world_size=2, + ), + "q_norm.weight": full_state["q_norm.weight"].clone(), + "k_norm.weight": full_state["k_norm.weight"].clone(), + } + # The indexer is replicated: copy it verbatim (no sharding) and prove it is identical. + for leaf in ( + "indexer.index_qk_proj.weight", "indexer.q_layernorm.weight", + "indexer.k_layernorm.weight", + ): + raw[leaf] = full_state[leaf].clone() + assert torch.equal(raw[leaf], full_state[leaf]) + attn2.load_state_dict(raw) + + seen2 = selection_spy(monkeypatch, fix2.backend) + batch2 = fix2.batch([fix2.req(*meta) for meta in req_meta], "prefill") + partials.append(attn2.forward(x, batch2).clone()) + selections.append(seen2["indices"].clone()) + slab_kv_heads.append(fix2.pool.k_cache(QSA_LAYER).shape[2]) + + # (a) replicated indexer => both ranks select exactly the same tokens as TP1 + assert torch.equal(selections[0], selections[1]), "TP2 ranks selected different blocks" + assert torch.equal(selections[0], want_idx), "TP2 selection diverged from TP1" + # (b) the K/V slab really is rank-local (TP1 has 2 kv heads, each rank has 1) + assert slab_kv_heads == [1, 1] + # (c) numeric gate: unreduced row-parallel partials sum to the TP1 output + merged = partials[0] + partials[1] + torch.testing.assert_close(merged.float(), want.float(), rtol=2e-2, atol=2e-2) + diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index fa75d6ee6..eea5f4468 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -131,6 +131,38 @@ def test_hc_merged_gemm_layout_and_top_level_mixer(): assert torch.allclose(x, ref_x, rtol=1e-5, atol=1e-6) +def test_parallel_lm_head_gathers_vocab_shards_in_rank_order(monkeypatch): + """The TP lm-head must restore rank-major vocab rows and trim padding.""" + import freetoken.distributed.info as info + from freetoken.layers.embedding import ParallelLMHead + + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=0, size=2)) + head0 = ParallelLMHead(num_embeddings=7, embedding_dim=3) + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=1, size=2)) + head1 = ParallelLMHead(num_embeddings=7, embedding_dim=3) + + full_weight = torch.arange(21, dtype=torch.float32).view(7, 3) / 10 + head0.weight.copy_(full_weight[:4]) + head1.weight.copy_(torch.cat((full_weight[4:], torch.zeros(1, 3)))) + x = torch.tensor([[1.0, -2.0, 0.5]]) + local0 = F.linear(x, head0.weight) + local1 = F.linear(x, head1.weight) + expected = F.linear(x, full_weight) + + ctx = _fresh_ctx() + batch = SimpleNamespace(size=1, is_prefill=False) + monkeypatch.setattr( + type(head0._comm), "all_gather", lambda _self, _x: torch.cat((local0, local1)) + ) + with ctx.forward_batch(batch): + got0 = head0.forward(x) + torch.testing.assert_close(got0, expected) + + with ctx.forward_batch(batch): + got1 = head1.forward(x) + torch.testing.assert_close(got1, expected) + + # -------------------------------------------------------------------------------------- # PLE # -------------------------------------------------------------------------------------- diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1ddbf4250..e2ddc0420 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -21,6 +21,7 @@ _DenseFuser, iter_weights, load_ple_table, + shard_qwen4_exp_dense_tensor, ) from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into @@ -310,6 +311,171 @@ def test_gdn_in_proj_slices_round_trip(loaded, checkpoint): assert torch.equal(part, back) +def test_tp2_qsa_and_gdn_head_shards_reassemble(checkpoint): + _folder, raw = checkpoint + config = SimpleNamespace( + num_qo_heads=QH, + num_kv_heads=KVH, + head_dim=AHD, + linear_attention_group=lambda: SimpleNamespace( + num_key_heads=KH, + num_value_heads=VH, + key_head_dim=HD, + value_head_dim=HD, + ), + ) + + def shard(key): + return [ + shard_qwen4_exp_dense_tensor( + key, raw[f"model.language_model.{key}"], config=config, + rank=rank, world_size=2, + ) + for rank in range(2) + ] + + qsa = "layers.1.self_attn.q_proj.weight" + assert torch.equal(torch.cat(shard(qsa), dim=0), raw[f"model.language_model.{qsa}"]) + for proj in ("k_proj", "v_proj"): + key = f"layers.1.self_attn.{proj}.weight" + parts = shard(key) + assert torch.equal(torch.cat(parts, dim=0), raw[f"model.language_model.{key}"]) + + gdn_qkv = "layers.0.linear_attn.in_proj_qkv.weight" + qkv = raw[f"model.language_model.{gdn_qkv}"] + # The checkpoint's qkv part is [q, k, v] with q/k each KH*HD and v VH*HD. + q, k, v = torch.split(qkv, [KH * HD, KH * HD, VH * HD], dim=0) + shards = shard(gdn_qkv) + expected = [ + torch.cat([q[: KH * HD // 2], k[: KH * HD // 2], v[: VH * HD // 2]], dim=0), + torch.cat([q[KH * HD // 2 :], k[KH * HD // 2 :], v[VH * HD // 2 :]], dim=0), + ] + assert all(torch.equal(got, want) for got, want in zip(shards, expected)) + + conv = "layers.0.linear_attn.conv1d.weight" + conv_shards = shard(conv) + cq, ck, cv = torch.split( + raw[f"model.language_model.{conv}"], [KH * HD, KH * HD, VH * HD], dim=0 + ) + assert torch.equal( + conv_shards[0], + torch.cat([cq[: KH * HD // 2], ck[: KH * HD // 2], cv[: VH * HD // 2]], dim=0), + ) + assert torch.equal( + conv_shards[1], + torch.cat([cq[KH * HD // 2 :], ck[KH * HD // 2 :], cv[VH * HD // 2 :]], dim=0), + ) + + for key in ( + "layers.0.linear_attn.in_proj_z.weight", + "layers.0.linear_attn.in_proj_b.weight", + "layers.0.linear_attn.in_proj_a.weight", + "layers.0.linear_attn.A_log", + "layers.0.linear_attn.dt_bias", + ): + value = raw[f"model.language_model.{key}"] + shards = [ + shard_qwen4_exp_dense_tensor( + key, value, config=config, rank=rank, world_size=2 + ) + for rank in range(2) + ] + assert torch.equal(torch.cat(shards, dim=0), value), key + + +def test_tp2_row_parallel_dense_weights_reassemble(checkpoint): + _folder, raw = checkpoint + config = SimpleNamespace( + num_qo_heads=QH, + num_kv_heads=KVH, + head_dim=AHD, + linear_attention_group=lambda: SimpleNamespace( + num_key_heads=KH, num_value_heads=VH, key_head_dim=HD, value_head_dim=HD, + ), + ) + for key in ( + "layers.1.self_attn.o_proj.weight", + "layers.0.linear_attn.out_proj.weight", + "layers.0.mlp.shared_expert.gate_proj.weight", + "layers.0.mlp.shared_expert.up_proj.weight", + "layers.0.mlp.shared_expert.down_proj.weight", + ): + value = raw[f"model.language_model.{key}"] + shards = [ + shard_qwen4_exp_dense_tensor( + key, value, config=config, rank=rank, world_size=2 + ) + for rank in range(2) + ] + dim = 1 if key.endswith(("o_proj.weight", "out_proj.weight", "down_proj.weight")) else 0 + assert torch.equal(torch.cat(shards, dim=dim), value), key + + for key, raw_key in ( + ("model.embed_tokens.weight", "model.language_model.embed_tokens.weight"), + ("lm_head.weight", "lm_head.weight"), + ): + value = raw[raw_key] + shards = [ + shard_qwen4_exp_dense_tensor( + key, value, config=config, rank=rank, world_size=2 + ) + for rank in range(2) + ] + # The vocabulary axis is PADDED per rank to div_ceil(V, tp), not truncated: the + # model allocates that many rows (VocabParallelEmbedding.num_embeddings_tp) and its + # gather trims the padding -- see test_skeleton's parallel lm-head test. Only the + # real rows must reassemble. + rows = -(-value.shape[0] // 2) + assert [s.shape[0] for s in shards] == [rows, rows], key + assert torch.equal(torch.cat(shards, dim=0)[: value.shape[0]], value), key + + +def test_tp2_short_final_vocabulary_shard_is_zero_padded(): + """The vocabulary axis is padded, not truncated. + + ``VocabParallelEmbedding`` always allocates ``div_ceil(vocab, tp)`` rows -- its + ``finish_idx`` clamps the token-index range, not the allocation -- so when the + vocabulary is not divisible by TP the final rank must still hand over a + full-width shard or strict loading fails on shape. + """ + config = SimpleNamespace(linear_attention_group=lambda: None) + vocab, width = 7, 4 # 7 is not divisible by 2 + value = torch.arange(vocab * width, dtype=torch.float32).reshape(vocab, width) + rows = -(-vocab // 2) + + shards = [ + shard_qwen4_exp_dense_tensor( + "model.embed_tokens.weight", value, config=config, rank=rank, world_size=2 + ) + for rank in range(2) + ] + + assert [tuple(s.shape) for s in shards] == [(rows, width), (rows, width)] + assert torch.equal(shards[0], value[:rows]) + # rank 1 carries the real tail rows plus a zero row no token id can reach + assert torch.equal(shards[1][: vocab - rows], value[rows:]) + assert torch.count_nonzero(shards[1][vocab - rows :]) == 0 + # the real vocabulary still reassembles exactly + assert torch.equal(torch.cat(shards, dim=0)[:vocab], value) + + +def test_tp2_lm_head_short_final_shard_is_zero_padded_too(): + config = SimpleNamespace(linear_attention_group=lambda: None) + vocab, width = 5, 3 # 5 % 4 != 0, so three of four ranks pad + value = torch.arange(vocab * width, dtype=torch.float32).reshape(vocab, width) + rows = -(-vocab // 4) + + shards = [ + shard_qwen4_exp_dense_tensor( + "lm_head.weight", value, config=config, rank=rank, world_size=4 + ) + for rank in range(4) + ] + + assert all(tuple(s.shape) == (rows, width) for s in shards) + assert torch.equal(torch.cat(shards, dim=0)[:vocab], value) + + def test_shared_expert_gate_up_merge(loaded, checkpoint): _folder, raw = checkpoint base = "model.language_model.layers.1.mlp.shared_expert" @@ -425,6 +591,117 @@ def test_read_range_into_rejects_a_short_destination(blob): read_range_into(bank.memoryview(), path, file_offset=0, nbytes=1 << 20) +def test_iter_weights_tp_shard_reassembles_every_dense_buffer(checkpoint, monkeypatch): + """TP2 loader contract: `tp_shard=True` emits rank-local buffers whose concatenation + equals the TP1 loader's output for every dense tensor, and the fused groups keep their + head boundaries. No model is constructed; the comparison is against `iter_weights` TP1.""" + import freetoken.distributed.info as info + + folder, _raw = checkpoint + config = SimpleNamespace( + num_qo_heads=QH, + num_kv_heads=KVH, + head_dim=AHD, + linear_attention_group=lambda: SimpleNamespace( + num_key_heads=KH, num_value_heads=VH, key_head_dim=HD, value_head_dim=HD, + ), + ) + tp1 = { + name: tensor + for name, tensor in iter_weights( + folder, torch.device("cpu"), include_moe_experts=False, include_non_moe=True + ) + } + + shards: list[dict[str, torch.Tensor]] = [] + for rank in range(2): + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=rank, size=2)) + shards.append( + { + name: tensor + for name, tensor in iter_weights( + folder, torch.device("cpu"), include_moe_experts=False, + include_non_moe=True, tp_shard=True, config=config, + ) + } + ) + + assert set(shards[0]) == set(shards[1]) == set(tp1) + # dim-0-concatenated groups vs dim-1 (row-parallel output projections). + dim1 = ("o_proj.weight", "out_proj.weight", "shared_expert.down_proj.weight") + # Replicated (not sharded) buffers must be bit-identical on both ranks. + replicated = ( + ".q_norm.weight", ".k_norm.weight", ".hc_norm.weight", + "input_mix_weight_down.weight", "input_mix_weight_up.weight", + "input_mix_weight_down_block_inject.weight", "norm_key.weight", "norm_query.weight", + "norm_conv.weight", ".ple.", "layer_multipliers", "ngram_heads_offsets", + "ngram_heads_vocab_sizes", ".indexer.", ".mlp.gate.weight", "shared_expert_gate.weight", + ".linear_attn.norm.weight", + ) + # Head-group buffers are sharded PER GROUP, so rank rows interleave (q0,k0,v0 | q1,k1,v1) + # instead of splitting the global fused rows in half. Sizes are the GLOBAL head groups. + qkv = (KH * HD, KH * HD, VH * HD) + head_group = { + "self_attn.qkv_proj.weight": (2 * QH * AHD, KVH * AHD, KVH * AHD), + "linear_attn.in_proj.weight": (*qkv, VH * HD, VH, VH), # q|k|v|z|b|a + "linear_attn.conv1d.weight": qkv, + "shared_expert.gate_up_proj.weight": (I, I), # gate|up, each halved by rank + } + for name, full in tp1.items(): + if any(token in name for token in replicated): + assert torch.equal(shards[0][name], full) and torch.equal(shards[1][name], full), name + continue + group = next( + (sizes for suffix, sizes in head_group.items() if name.endswith(suffix)), None + ) + if group is not None: + # Split the GLOBAL fused rows into head groups; each group halves by rank. + parts = torch.split(full, group, dim=0) + assert [p.shape[0] for p in parts] == list(group), name + for rank in range(2): + got = torch.split(shards[rank][name], [s // 2 for s in group], dim=0) + for part, half, size in zip(parts, got, group): + expected = part[rank * (size // 2) : (rank + 1) * (size // 2)] + assert torch.equal(half, expected), (name, rank) + continue + dim = 1 if name.endswith(dim1) else 0 + if name in ("model.embed_tokens.weight", "lm_head.weight"): + # PADDED, not truncated: each rank holds div_ceil(V, tp) rows and the model's + # vocab gather trims the tail (see test_skeleton's parallel lm-head test), so + # only the real rows have to reassemble and the padding must be zero. + rows = -(-full.shape[0] // 2) + for rank in range(2): + got = shards[rank][name] + assert got.shape[0] == rows, name + real = full[rank * rows : (rank + 1) * rows] + assert torch.equal(got[: real.shape[0]], real), name + assert torch.count_nonzero(got[real.shape[0] :]) == 0, name + continue + merged = torch.cat([shards[0][name], shards[1][name]], dim=dim) + assert torch.equal(merged, full), name + assert shards[0][name].shape != full.shape or full.shape[dim] == 1, name + + # The fused QSA qkv keeps [2*qo | kv | kv] ordering on each rank, so a naive split of + # the global fused buffer would NOT reproduce it. + key = "model.layers.1.self_attn.qkv_proj.weight" + per_rank = shards[0][key].shape[0] + assert per_rank == (2 * (QH // 2) + 2 * (KVH // 2)) * AHD + assert shards[0][key].shape[0] + shards[1][key].shape[0] == tp1[key].shape[0] + + +def test_iter_weights_tp_shard_is_opt_in_and_fails_fast_without_it(checkpoint, monkeypatch): + import freetoken.distributed.info as info + + folder, _raw = checkpoint + monkeypatch.setattr(info, "_TP_INFO", info.DistributedInfo(rank=0, size=2)) + with pytest.raises(NotImplementedError, match="tp_shard=True"): + list( + iter_weights( + folder, torch.device("cpu"), include_moe_experts=False, include_non_moe=True + ) + ) + + # ====================================================================================== # AOT shape table # ====================================================================================== diff --git a/tests/models/test_weight_tp_shard.py b/tests/models/test_weight_tp_shard.py new file mode 100644 index 000000000..21a96851e --- /dev/null +++ b/tests/models/test_weight_tp_shard.py @@ -0,0 +1,120 @@ +"""``load_weight`` must forward ``tp_shard`` ONLY to readers that declare it. + +Regression: the engine asks for ``tp_shard`` on every TP>1 launch, and +``load_weight`` used to raise ``NotImplementedError`` for any reader without that +parameter. Seven readers shard *internally* instead (they call ``shard_tensor`` +with ``tp_info.rank``/``tp_info.size`` inside ``iter_weights``: llama, qwen2, +qwen3, qwen3_moe, mistral, gpt_oss, minimax_m2), so forwarding the flag turned a +working TP>1 launch into a startup failure. FTW checkpoints store POST-shard +weights, so the flag is a no-op there and must not be rejected either. +""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.models import weight as weight_mod + + +def _reader(declares: bool, seen: list): + if declares: + + def iter_weights( + model_path, device, *, include_moe_experts, include_non_moe, + tp_shard=False, config=None, + ): + seen.append({"tp_shard": tp_shard, "config": config}) + yield "w", torch.zeros(2, 2) + + else: + + def iter_weights(model_path, device, *, include_moe_experts, include_non_moe): + seen.append({"tp_shard": ""}) + yield "w", torch.zeros(2, 2) + + return iter_weights + + +def _patch(monkeypatch, reader): + monkeypatch.setattr( + weight_mod, + "_spec_for_model_path", + lambda path: ( + None, + SimpleNamespace(module="fake.mod", iter_weights="iter_weights"), + ), + ) + monkeypatch.setattr(weight_mod, "_load_attr", lambda module, name: reader) + monkeypatch.setattr("freetoken.checkpoint.ftw.is_ftw_checkpoint", lambda path: False) + + +def test_tp_shard_is_forwarded_to_a_reader_that_declares_it(monkeypatch): + seen: list = [] + _patch(monkeypatch, _reader(True, seen)) + list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=True)) + assert seen == [{"tp_shard": True, "config": None}] + + +def test_tp_config_rides_along_only_when_the_reader_accepts_it(monkeypatch): + seen: list = [] + _patch(monkeypatch, _reader(True, seen)) + sentinel = object() + list( + weight_mod.load_weight( + "/m", torch.device("cpu"), tp_shard=True, tp_config=sentinel + ) + ) + assert seen == [{"tp_shard": True, "config": sentinel}] + + +def test_tp_shard_is_not_forwarded_to_a_reader_that_shards_internally(monkeypatch): + """The regression: this raised, breaking TP>1 for the internally-sharded readers.""" + seen: list = [] + _patch(monkeypatch, _reader(False, seen)) + out = list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=True)) + assert [name for name, _ in out] == ["w"] + assert seen == [{"tp_shard": ""}] + + +def test_tp1_does_not_forward_tp_shard_to_such_a_reader_either(monkeypatch): + seen: list = [] + _patch(monkeypatch, _reader(False, seen)) + list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=False)) + assert seen == [{"tp_shard": ""}] + + +def test_ftw_checkpoints_accept_tp_shard(monkeypatch): + """FTW stores post-shard weights: tp_shard is a no-op there, not an error.""" + monkeypatch.setattr("freetoken.checkpoint.ftw.is_ftw_checkpoint", lambda path: True) + monkeypatch.setattr( + "freetoken.checkpoint.ftw.iter_ftw_weights", + lambda path: iter([("w", torch.zeros(2, 2))]), + ) + monkeypatch.setattr("freetoken.models.config.vision_load_enabled", lambda: False) + out = list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=True)) + assert [name for name, _ in out] == ["w"] + + +def test_a_reader_without_tp_shard_is_never_asked_for_it(monkeypatch): + """Guards the mechanism, not just the outcome: the kwarg must not be built at all.""" + seen: list = [] + _patch(monkeypatch, _reader(False, seen)) + list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=True)) + assert "" in seen[0]["tp_shard"], "tp_shard reached a reader that cannot take it" + + +def test_a_non_callable_reader_attribute_is_not_probed_as_a_signature(monkeypatch): + """``_load_attr`` may return a non-function; only a real signature decides.""" + monkeypatch.setattr( + weight_mod, + "_spec_for_model_path", + lambda path: (None, SimpleNamespace(module="fake.mod", iter_weights="iter_weights")), + ) + monkeypatch.setattr( + weight_mod, "_load_attr", lambda module, name: "not-callable" + ) + monkeypatch.setattr("freetoken.checkpoint.ftw.is_ftw_checkpoint", lambda path: False) + with pytest.raises(Exception): + list(weight_mod.load_weight("/m", torch.device("cpu"), tp_shard=True)) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index b29ce531d..a3154f351 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1,4 +1,5 @@ from contextlib import contextmanager +from types import SimpleNamespace import pytest import torch @@ -666,6 +667,71 @@ def test_offload_cache_rebuild_keeps_overlap_at_boundary(): assert cache.cache_size == 8 +def test_copy_missing_consumes_the_staged_layer_exactly_once(): + # Regression: _pending_src_layer was never cleared, so a later copy_missing() with + # nothing freshly staged replayed the PREVIOUS layer's src_indices/evict_slots and + # overwrote slots that had since been reassigned to another layer. The owner adapter's + # "nothing staged" guard (inner pending is None) depends on one-shot consumption. + cache, _ = _make_split_cache(num_layers=2, locked=(1,)) + + cache._pending_src_layer = 1 + cache._pending_whole_layer = True + cache.copy_missing() + + assert cache._pending_src_layer is None + assert cache._pending_whole_layer is False + # a second call is an explicit "nothing staged", never a silent replay of layer 1 + with pytest.raises(AssertionError, match="no staged misses"): + cache.copy_missing() + + +def test_owner_cache_rebuild_keeps_the_geometry_in_step(): + # Regression: __getattr__ forwarded rebuild() to the inner cache, whose implementation + # disables prefill_overlap when the new size cannot hold two complete local layers. + # The frozen geometry kept the old values, so materialize_layer() still took the + # overlap path and waited on buffers the inner cache no longer had. + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + _init_tp() + geometry = OwnerCacheGeometry( + global_num_experts=8, world_size=2, rank=0, num_layers=1, + cache_size=8, prefill_overlap=True, + ) + owner = OwnerOffloadMoeCache(geometry, torch.device("cpu")) + owner.set_bank_sources( + {"gate_up": [torch.randn(4, 32, 8)], "down": [torch.randn(4, 8, 16)]} + ) + assert owner.geometry.prefill_overlap is True + + owner.rebuild(5) # 5 < 2*local_num_experts (8) -> the inner cache drops overlap + + assert owner._cache.prefill_overlap is False + assert owner.geometry.prefill_overlap is False, "geometry must follow the inner cache" + assert owner.geometry.cache_size == 5 + assert owner._cache.cache_size == 5 + + +def test_owner_cache_rebuild_keeps_overlap_when_the_new_size_still_fits(): + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + _init_tp() + geometry = OwnerCacheGeometry( + global_num_experts=8, world_size=2, rank=0, num_layers=1, + cache_size=8, prefill_overlap=True, + ) + owner = OwnerOffloadMoeCache(geometry, torch.device("cpu")) + owner.set_bank_sources( + {"gate_up": [torch.randn(4, 32, 8)], "down": [torch.randn(4, 8, 16)]} + ) + + owner.rebuild(8) # exactly 2*local_num_experts -> overlap survives + + assert owner.geometry.prefill_overlap is True + assert owner.geometry.cache_size == 8 + + def test_offload_cache_validate_rebuild_enforces_marlin_cap_and_floor(): # The constructor caps nvfp4_marlin slots at 992; a runtime rebuild must enforce the # same upper cap (and the num_experts floor), else marlin decode kernels later break. @@ -876,3 +942,652 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def _owner_nvfp4_banks(num_layers, local_experts, out, inner, base): + """Owner-local nvfp4 banks; row `e` of layer `l` carries fingerprint base+l*E+e.""" + + def bank(o, i, dtype): + layers = [] + for l in range(num_layers): + t = torch.zeros(local_experts, o, i, dtype=dtype) + for e in range(local_experts): + t[e].view(torch.uint8).fill_(base + l * local_experts + e) + layers.append(t) + return layers + + def pinned(layers): + return [t.pin_memory() for t in layers] + + return { + "gate_up_packed": pinned(bank(out, inner // 2, torch.uint8)), + "gate_up_scale": pinned(bank(out, inner // 16, torch.float8_e4m3fn)), + "gate_up_global": pinned( + [t.squeeze(-1).contiguous() for t in bank(out, 1, torch.float16)] + ), + "down_packed": pinned(bank(out, inner // 2, torch.uint8)), + "down_scale": pinned(bank(out, inner // 16, torch.float8_e4m3fn)), + "down_global": pinned( + [t.squeeze(-1).contiguous() for t in bank(out, 1, torch.float16)] + ), + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_owner_offload_cache_cuda_route_copies_local_rows_through_real_kernels(): + """P3 namespace smoke on real GPU kernels: global route -> owner-local bank row -> + legacy slot cache. Remote entries must never read a bank row, and a cache hit must + not re-copy. No model is loaded; the banks are synthetic.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL, S = 2, 8, 4, 4 + OUT, IN = 64, 512 # rows >= 128B so the fast_index_copy JIT has a kernel + dev = torch.device("cuda") + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=1, num_layers=L, cache_size=S + ) + cache = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4") + cache.set_bank_sources(_owner_nvfp4_banks(L, E_LOCAL, OUT, IN, base=100)) + cache.reset() + + def fingerprint(slot): + packed = cache.bank_caches["gate_up_packed"] + return int(packed[slot].view(torch.uint8).flatten()[0].item()) + + # rank 1 owns global [4, 8) -> local rows [0, 4). + ids = torch.tensor([[0, 4, 7, 5, 4]], dtype=torch.int32, device=dev) + weights = torch.tensor([[0.1, 0.2, 0.3, 0.15, 0.25]], device=dev) + update = cache.ensure_route(0, weights, ids) + cache.copy_missing() + torch.cuda.synchronize() + + assert update.owned_mask.tolist() == [[False, True, True, True, True]] + assert update.local_ids.tolist() == [[0, 0, 3, 1, 0]] + # remote (global 0) stays a safe placeholder: slot 0, zero weight, never read. + assert int(update.slot_ids[0, 0].item()) == 0 + assert float(update.weights[0, 0].item()) == 0.0 + # flashlib emits the miss set in ascending local-row order (the CPU reference + # adapter emits route order) -- only the SET is part of the contract. + assert sorted(update.missing_local_ids.tolist()) == [0, 1, 3] + + # each owned position's slot holds ITS OWN local row's bytes (layer 0 -> base 100). + owned_local = update.local_ids[update.owned_mask].tolist() + owned_slots = update.slot_ids[update.owned_mask].tolist() + assert [fingerprint(s) for s in owned_slots] == [100 + e for e in owned_local] + cache.validate_invariants() + + # A repeated route is a pure hit: no miss, no eviction, identical bytes. + again = cache.ensure_route(0, weights, ids) + cache.copy_missing() + torch.cuda.synchronize() + assert again.missing_local_ids.numel() == 0 + assert again.evicted_flat_ids.numel() == 0 + assert again.slot_ids.tolist() == update.slot_ids.tolist() + assert [fingerprint(s) for s in owned_slots] == [100 + e for e in owned_local] + + # The pool is unified: layer 1 must evict layer-0 entries and serve layer-1 bytes + # (base 100 + 4) without ever mixing the two layers' rows. + ids1 = torch.tensor([[4, 6]], dtype=torch.int32, device=dev) + update1 = cache.ensure_route(1, torch.ones(1, 2, device=dev), ids1) + cache.copy_missing() + torch.cuda.synchronize() + slots1 = update1.slot_ids.reshape(-1).tolist() + assert [fingerprint(s) for s in slots1] == [104 + e for e in [0, 2]] + assert int(update1.evicted_flat_ids.numel()) > 0 # S=6 < 2 layers * 4 local rows + cache.validate_invariants() + + # Remote-only route: no admission, no copy, all placeholders. + remote_only = cache.ensure_route( + 1, torch.full((1, 2), 0.5, device=dev), + torch.tensor([[0, 1]], dtype=torch.int32, device=dev), + ) + assert remote_only.missing_local_ids.numel() == 0 + assert remote_only.slot_ids.tolist() == [[0, 0]] + assert remote_only.weights.tolist() == [[0.0, 0.0]] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_owner_route_graph_admission_matches_eager_slots_and_zeroes_remote(): + """The graph-safe admission must place every OWNED route entry on the same local row as + the eager compacting path, keep remote entries at zero weight, and keep the route shape + statically known (that is what makes it capturable). Remote entries point at the + sentinel row instead of being dropped, so they must be resident and finite -- the zero + weighting must never rely on ``0 * NaN``.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL, S = 2, 8, 4, 8 + OUT, IN = 64, 512 + dev = torch.device("cuda") + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=1, num_layers=L, cache_size=S + ) + + def build(graph_safe): + c = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4", graph_safe=graph_safe) + c.set_bank_sources(_owner_nvfp4_banks(L, E_LOCAL, OUT, IN, base=100)) + c.reset() + return c + + def fingerprint(cache, slot): + packed = cache.bank_caches["gate_up_packed"] + return int(packed[slot].view(torch.uint8).flatten()[0].item()) + + # rank 1 owns global [4, 8) -> local rows [0, 4); global 0/1 are remote. + # Route chosen so the FIRST owned position is local row 3, not row 0: that makes the + # remote fallback distinguishable from a fixed row-zero sentinel. + ids = torch.tensor([[0, 7, 5, 6, 4]], dtype=torch.int32, device=dev) + weights = torch.tensor([[0.1, 0.2, 0.3, 0.15, 0.25]], device=dev) + + eager = build(graph_safe=False) + up_eager = eager.ensure_route(0, weights, ids) + eager.copy_missing() + torch.cuda.synchronize() + + graph = build(graph_safe=True) + assert graph.graph_safe is True + up_graph = graph.ensure_route_graph(0, weights, ids) + graph.copy_missing() + torch.cuda.synchronize() + + # Weights agree exactly: remote positions are zero in both paths. + assert up_graph.weights.tolist() == up_eager.weights.tolist() + assert up_graph.weights[0, 0].item() == 0.0 + assert up_graph.owned_mask.tolist() == up_eager.owned_mask.tolist() + + # Shape/dtype are static, and the diagnostics contract is "empty, never a host read". + assert up_graph.slot_ids.shape == ids.shape + assert up_graph.slot_ids.dtype == torch.int32 + assert up_graph.missing_local_ids.numel() == 0 + assert up_graph.evicted_flat_ids.numel() == 0 + + # The remote entry borrows the first owned row (row 3), NOT a fixed row-zero sentinel. + local_row = up_graph.local_ids[0].tolist() + assert local_row == [3, 3, 1, 2, 0] + owned = up_graph.owned_mask[0].tolist() + slots = up_graph.slot_ids[0].tolist() + # The two admission paths must place every owned entry in the same slot: sharing the row + # set is what keeps the graph path from changing cache behaviour (extra misses). + assert [s for s, o in zip(slots, owned) if o] == [ + s for s, o in zip(up_eager.slot_ids[0].tolist(), owned) if o + ] + # Every owned position carries ITS OWN row's bytes; the remote position carries row 3's. + assert [fingerprint(graph, s) for s, o in zip(slots, owned) if o] == [ + 100 + r for r, o in zip(local_row, owned) if o + ] + assert fingerprint(graph, slots[0]) == 103 + assert torch.isfinite(graph.bank_caches["gate_up_packed"][slots[0]].float()).all() + graph.validate_invariants() + + # All-remote route (no owned row to borrow) falls back to row zero and contributes + # nothing -- the only case that admits an extra row. + only_remote = graph.ensure_route_graph( + 1, torch.full((1, 2), 0.5, device=dev), + torch.tensor([[0, 1]], dtype=torch.int32, device=dev), + ) + graph.copy_missing() + torch.cuda.synchronize() + assert only_remote.weights.tolist() == [[0.0, 0.0]] + assert set(only_remote.local_ids[0].tolist()) == {0} + assert fingerprint(graph, only_remote.slot_ids[0, 0].item()) == 104 # layer 1, row 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_owner_route_graph_admission_is_capturable_and_replays(): + """The point of the graph-safe path: capture admission + copy into a real CUDA graph and + replay it. The eager path's ``nonzero``/``num_indices.item()`` make this fail, which is + why owner EP used to be restricted to ``--cuda-graph-max-bs 0``.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL, S = 2, 8, 4, 8 + OUT, IN = 64, 512 + dev = torch.device("cuda") + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=1, num_layers=L, cache_size=S + ) + cache = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4", graph_safe=True) + cache.set_bank_sources(_owner_nvfp4_banks(L, E_LOCAL, OUT, IN, base=100)) + cache.reset() + + ids = torch.tensor([[0, 4, 7, 5, 4]], dtype=torch.int32, device=dev) + weights = torch.tensor([[0.1, 0.2, 0.3, 0.15, 0.25]], device=dev) + + def step(ids_buf, weights_buf): + update = cache.ensure_route_graph(0, weights_buf, ids_buf) + cache.copy_missing() + return update + + # Warm the kernels on a side stream so capture starts from a steady state. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + step(ids.clone(), weights.clone()) + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + ids_buf = ids.clone() + weights_buf = weights.clone() + with torch.cuda.graph(graph): + captured = step(ids_buf, weights_buf) + + graph.replay() + torch.cuda.synchronize() + + packed = cache.bank_caches["gate_up_packed"] + slots = captured.slot_ids[0].tolist() + rows = captured.local_ids[0].tolist() + owned = captured.owned_mask[0].tolist() + assert captured.weights[0, 0].item() == 0.0 + assert [int(packed[s].view(torch.uint8).flatten()[0].item()) + for s, o in zip(slots, owned) if o] == [ + 100 + r for r, o in zip(rows, owned) if o + ] + cache.validate_invariants() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_owner_prefill_materialize_copies_local_rows_and_does_not_leak_layers(): + """Owner prefill must move bytes, not just remap slots: materialize layer 0, check + every local row's fingerprint, then materialize layer 1 and prove no layer-0 bytes + remain in the materialized view.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL, S = 2, 8, 4, 4 + OUT, IN = 64, 512 + dev = torch.device("cuda") + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=1, num_layers=L, cache_size=S + ) + cache = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4") + cache.set_bank_sources(_owner_nvfp4_banks(L, E_LOCAL, OUT, IN, base=100)) + cache.reset() + + def fingerprint(slot): + packed = cache.bank_caches["gate_up_packed"] + return int(packed[slot].view(torch.uint8).flatten()[0].item()) + + cache.materialize_layer(0) + torch.cuda.synchronize() + assert [fingerprint(s) for s in range(E_LOCAL)] == [100 + e for e in range(E_LOCAL)] + + cache.materialize_layer(1) + torch.cuda.synchronize() + assert [fingerprint(s) for s in range(E_LOCAL)] == [104 + e for e in range(E_LOCAL)] + cache.validate_invariants() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_owner_prefill_overlap_keeps_two_layers_resident_at_once(): + """Owner prefill overlap must stream layer L+1 while layer L computes. The tell is + that the two borrowed buffers hold DIFFERENT layers SIMULTANEOUSLY after + ``prefetch(0) -> prefetch(1)`` and before any release -- a choreography that only + prefetched the current layer would leave buffer 1 untouched here.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL = 2, 8, 4 + S = 2 * E_LOCAL # the overlap floor: two full local layers + OUT, IN = 64, 512 + dev = torch.device("cuda") + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=1, num_layers=L, + cache_size=S, prefill_overlap=True, + ) + cache = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4") + cache.set_bank_sources(_owner_nvfp4_banks(L, E_LOCAL, OUT, IN, base=100)) + cache.reset() + assert cache.geometry.prefill_overlap is True + assert cache.prefill_overlap is True # forwarded to the wrapped cache + + def fingerprint(buf, row): + return int(buf[row].view(torch.uint8).flatten()[0].item()) + + cache.begin_prefill() + cache.prefetch_prefill_layer(0) # -> buffer 0 + cache.prefetch_prefill_layer(1) # -> buffer 1, issued while layer 0 will compute + torch.cuda.synchronize() + + buffers = cache.prefill_bank_buffers[0] # bank 0, [2, E_LOCAL, ...] + assert [fingerprint(buffers[0], r) for r in range(E_LOCAL)] == [ + 100 + r for r in range(E_LOCAL) + ] + # Layer 1 is already staged in the OTHER buffer: that is the overlap. + assert [fingerprint(buffers[1], r) for r in range(E_LOCAL)] == [ + 104 + r for r in range(E_LOCAL) + ] + + # The hand-off the layer performs: wait(cur) returns cur's buffer, release frees it. + views0 = cache.wait_prefill_layer(0) + assert [fingerprint(views0[0], r) for r in range(E_LOCAL)] == [ + 100 + r for r in range(E_LOCAL) + ] + cache.release_prefill_layer(0) + views1 = cache.wait_prefill_layer(1) + assert [fingerprint(views1[0], r) for r in range(E_LOCAL)] == [ + 104 + r for r in range(E_LOCAL) + ] + cache.release_prefill_layer(1) + + # A second prefill over the same buffers must not leak the previous layer's bytes. + cache.begin_prefill() + cache.prefetch_prefill_layer(1) + torch.cuda.synchronize() + assert [fingerprint(cache.prefill_bank_buffers[0][1], r) for r in range(E_LOCAL)] == [ + 104 + r for r in range(E_LOCAL) + ] + cache.validate_invariants() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize( + "route_ids", + [ + [0, 1, 2, 3, 4, 5, 6, 7, 0, 7], # 5-5 split across the two owners + [0] * 10, # rank 0 owns every entry + [4] * 10, # rank 1 owns every entry + ], +) +def test_owner_prefill_gemm_partials_sum_to_tp1(route_ids): + """Two owner-local prefill GEMMs must sum to the TP1 GEMM for the same route. + + Each rank materializes its own local layer, remaps the global route to local rows + with remote entries zero-weighted, and runs the real Triton NVFP4 prefill kernel. + The TP1 reference runs the same kernel over the full 8-expert layer.""" + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4 + from freetoken.moe.offload_cache import OffloadMoeCache, OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + L, E_GLOBAL, E_LOCAL = 1, 8, 4 + OUT, IN = 64, 512 + dev = torch.device("cuda") + hidden = torch.randn(1, OUT, dtype=torch.bfloat16, device=dev) / 4 + weights = torch.arange(1, 11, dtype=torch.float32, device=dev).reshape(1, 10) / 55 + ids = torch.tensor([route_ids], dtype=torch.int32, device=dev) + + def random_sources(num_experts, seed): + g = torch.Generator().manual_seed(seed) + total = L * num_experts + + def rand_u8(*shape): + return torch.randint(0, 256, shape, dtype=torch.uint8, generator=g) + + def rand_scale(*shape): + return (torch.rand(*shape, generator=g) * 1.5 + 0.25).to(torch.float8_e4m3fn) + + flat = { + "gate_up_packed": rand_u8(total, 2 * IN, OUT // 2), + "gate_up_scale": rand_scale(total, 2 * IN, OUT // 16), + "gate_up_global": torch.full((total, 2 * IN), 1.0, dtype=torch.float16), + "down_packed": rand_u8(total, OUT, IN // 2), + "down_scale": rand_scale(total, OUT, IN // 16), + "down_global": torch.full((total, OUT), 0.75, dtype=torch.float16), + } + return {name: list(t.pin_memory().split(num_experts)) for name, t in flat.items()} + + full_sources = random_sources(E_GLOBAL, seed=7) + tp1 = OffloadMoeCache( + num_layers=L, num_experts=E_GLOBAL, cache_size=E_GLOBAL, device=dev, + quant_format="nvfp4", + ) + tp1.set_bank_sources(full_sources) + tp1.reset() + tp1.materialize_layer(0) + tp1.copy_missing() + want = fused_experts_nvfp4( + hidden, *tp1.bank_views(E_GLOBAL), weights, ids, E_GLOBAL, "silu", False, + ) + + partials = [] + for rank in range(2): + geometry = OwnerCacheGeometry( + global_num_experts=E_GLOBAL, world_size=2, rank=rank, num_layers=L, + cache_size=E_LOCAL, + ) + owner = OwnerOffloadMoeCache(geometry, dev, quant_format="nvfp4") + local_sources = { + name: [ + full_sources[name][0][rank * E_LOCAL:(rank + 1) * E_LOCAL] + .clone() + .pin_memory() + ] + for name in full_sources + } + owner.set_bank_sources(local_sources) + owner.reset() + owner.materialize_layer(0) + torch.cuda.synchronize() + local_ids, owned = geometry.global_to_local(ids) + safe_ids = torch.where(owned, local_ids, torch.zeros_like(local_ids)).contiguous() + safe_weights = torch.where( + owned, weights, torch.zeros_like(weights) + ).contiguous() + partials.append( + fused_experts_nvfp4( + hidden, *owner.bank_views(E_LOCAL), safe_weights, safe_ids, E_LOCAL, + "silu", False, + ) + ) + got = (partials[0] + partials[1]).float() + ref = want.float() + tol = 0.03 * float(ref.abs().max()) + torch.testing.assert_close(got, ref, rtol=3e-2, atol=max(tol, 3e-2)) + + +class _RecordingMoEMethod: + """Stand-in for the layer's MoE quant method: records what ``_expert_gemm`` hands the + kernel (bank views + routing ids) and returns the hidden states unchanged.""" + + def __init__(self): + self.calls = [] + + def apply(self, hidden_states, topk_weights, topk_ids, view, *, layer, is_prefill): + self.calls.append( + SimpleNamespace( + weights=topk_weights, + ids=topk_ids, + views=view.tensors, + n=view.n, + is_prefill=is_prefill, + ) + ) + return hidden_states + + +def _make_owner_layer(quant_format="bf16", prefill_overlap=False): + """OffloadMoELayer wired to an OwnerOffloadMoeCache with tiny local banks.""" + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + from freetoken.moe.ownership import OwnerCacheGeometry + + _init_tp() + layer = _bf16_offload_layer(0, 8, 2, 8, 16) + geometry = OwnerCacheGeometry( + global_num_experts=8, world_size=2, rank=1, num_layers=1, + cache_size=8, prefill_overlap=prefill_overlap, + ) + owner = OwnerOffloadMoeCache(geometry, torch.device("cpu"), quant_format=quant_format) + if quant_format == "bf16": + owner.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8)], + "down": [torch.randn(4, 8, 16)], + }) + else: + owner.set_bank_sources(_owner_nvfp4_banks(1, 4, 64, 512, base=100)) + layer.owner_cache = owner + layer.offload_cache = owner._cache + # the owner path is exercised through the layer's kernel seam, so record there + layer.quant_method = _RecordingMoEMethod() + return layer, owner + + +def test_owner_layer_decode_uses_owner_route_and_never_the_global_ids(monkeypatch): + """P3 wiring: an attached owner cache must route decode through ``ensure_route`` and + feed the kernel the LOCAL slot ids + masked weights, never the raw global ids.""" + layer, owner = _make_owner_layer() + # rank 1 owns global [4, 8) -> local rows [0, 4). + topk_weights = torch.tensor([[0.1, 0.2, 0.3]], dtype=torch.float32) + topk_ids = torch.tensor([[0, 4, 7]], dtype=torch.int32) + hidden_states = torch.randn(1, 8) + calls = {} + + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids), + ) + # The global-ID entry point must NOT be reached on the owner path. + monkeypatch.setattr( + owner._cache, "ensure_experts", + lambda *a, **k: pytest.fail("owner path called the global-ID ensure_experts"), + ) + + def fake_update(layer_id, weights, expert_ids): + owned = expert_ids >= 4 + slots = torch.where(owned, expert_ids - 4, torch.zeros_like(expert_ids)) + return SimpleNamespace( + weights=torch.where(owned, weights, torch.zeros_like(weights)), + slot_ids=slots, + ) + + monkeypatch.setattr(owner, "ensure_route", fake_update) + monkeypatch.setattr(owner, "copy_missing", lambda: None) + + original_route = owner.ensure_route + + def fake_route(layer_id, weights, expert_ids): + calls["route_args"] = (layer_id, weights.clone(), expert_ids.clone()) + return original_route(layer_id, weights, expert_ids) + + monkeypatch.setattr(owner, "ensure_route", fake_route) + + out = layer.decode_forward(hidden_states, torch.randn(1, 8)) + + call = layer.quant_method.calls[-1] + assert out is hidden_states + assert calls["route_args"][0] == 0 + assert calls["route_args"][2].tolist() == [[0, 4, 7]] # raw global ids reached the adapter + # remote (global 0) is zero-weighted; owned entries keep their global weights. + assert torch.allclose(call.weights, torch.tensor([[0.0, 0.2, 0.3]])) + # ids handed to the kernel are LOCAL slots, strictly inside the local pool. + assert call.ids.dtype == torch.int32 + assert int(call.ids.min()) >= 0 + assert int(call.ids.max()) < owner.cache_size + # the banks the kernel reads are the owner-local ones (4 rows), not the global 8. + assert call.views["gate_up"].shape[0] == owner.cache_size + assert call.views["down"].shape[0] == owner.cache_size + + +def test_owner_layer_remote_only_route_zeroes_the_contribution(monkeypatch): + """A rank that owns none of the routed experts must emit an all-zero, in-range route.""" + layer, owner = _make_owner_layer() + topk_weights = torch.full((1, 2), 0.5, dtype=torch.float32) + topk_ids = torch.tensor([[0, 2]], dtype=torch.int32) # both owned by rank 0 + calls = {} + + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids), + ) + monkeypatch.setattr( + owner, + "ensure_route", + lambda layer_id, weights, expert_ids: SimpleNamespace( + weights=torch.zeros_like(weights), + slot_ids=torch.zeros_like(expert_ids), + ), + ) + monkeypatch.setattr(owner, "copy_missing", lambda: None) + + layer.decode_forward(torch.randn(1, 8), torch.randn(1, 8)) + + call = layer.quant_method.calls[-1] + assert torch.equal(call.weights, torch.zeros_like(topk_weights)) + assert call.ids.tolist() == [[0, 0]] + assert owner.resident == 0 # nothing was admitted + + +def test_owner_layer_prefill_remaps_global_ids_to_local_rows(monkeypatch): + """P3 prefill: bank row ids must be LOCAL rows with remote entries zero-weighted.""" + layer, owner = _make_owner_layer(prefill_overlap=False) + topk_weights = torch.tensor([[0.25, 0.75]], dtype=torch.float32) + topk_ids = torch.tensor([[1, 6]], dtype=torch.int32) # 1 = rank0, 6 = local row 2 + calls = {} + monkeypatch.setattr(owner, "materialize_layer", lambda layer_id, buffer_id=0: None) + monkeypatch.setattr(owner, "bank_views", lambda n=None: (torch.empty(8, 32, 8), torch.empty(8, 8, 16))) + + layer._prefill_routed(torch.randn(1, 8), topk_weights, topk_ids) + + call = layer.quant_method.calls[-1] + assert torch.allclose(call.weights, torch.tensor([[0.0, 0.75]])) + assert call.ids.tolist() == [[0, 2]] # global 6 -> local row 2 + assert int(call.ids.max()) < owner.num_experts + + +def test_owner_layer_prefill_overlap_waits_and_releases_borrowed_buffer(monkeypatch): + """The owner prefill path must use the same borrowed-buffer lifecycle as the global + cache, INCLUDING the one-layer lookahead: prefetch(cur) stages this layer, prefetch(next) + starts the following layer's H2D on the copy stream so it runs while this layer's GEMMs + run on the compute stream. Without the lookahead the copy is issued and immediately + waited on, which serializes the two and defeats the overlap.""" + layer, owner = _make_owner_layer(prefill_overlap=True) + topk_weights = torch.tensor([[0.25, 0.75]], dtype=torch.float32) + topk_ids = torch.tensor([[1, 6]], dtype=torch.int32) + calls = {} + lifecycle = [] + monkeypatch.setattr(owner, "begin_prefill", lambda: lifecycle.append("begin")) + monkeypatch.setattr(owner, "prefetch_prefill_layer", lambda layer_id: lifecycle.append(("prefetch", layer_id))) + monkeypatch.setattr( + owner, + "wait_prefill_layer", + lambda layer_id: (torch.empty(8, 32, 8), torch.empty(8, 8, 16)), + ) + monkeypatch.setattr(owner, "release_prefill_layer", lambda layer_id: lifecycle.append("release")) + + layer._prefill_routed(torch.randn(1, 8), topk_weights, topk_ids) + + call = layer.quant_method.calls[-1] + assert torch.allclose(call.weights, torch.tensor([[0.0, 0.75]])) + assert call.ids.tolist() == [[0, 2]] + assert call.n == owner.num_experts + # layer 0 -> the lookahead asks for layer 1 (a no-op past the last layer). + assert lifecycle == ["begin", ("prefetch", 0), ("prefetch", 1), "release"] + + +def test_owner_wrapper_forwards_engine_assigned_flags_to_the_inner_cache(): + """Engine sets these by assignment; without the properties they would land on the + wrapper and silently disable stats / the route trace / CPU-layer routing.""" + layer, owner = _make_owner_layer() + owner.collect_stats = True + owner.collect_decode_freq = True + owner.route_recorder = object() + owner.cpu_layer_ids = frozenset({1}) + assert owner._cache.collect_stats is True + assert owner._cache.collect_decode_freq is True + assert owner._cache.route_recorder is not None + assert owner._cache.cpu_layer_ids == frozenset({1}) + assert owner.collect_stats is True and owner.cpu_layer_ids == frozenset({1}) + + +def test_owner_cache_bank_slots_start_zero_and_finite(): + """Slot-zero remote placeholders must not read torch.empty/NaN data.""" + _layer, owner = _make_owner_layer() + for cache in owner._cache.bank_caches.values(): + assert torch.count_nonzero(cache).item() == 0 + assert torch.isfinite(cache).all().item() + + +def test_attach_owner_moe_cache_wires_layers_and_keeps_global_banks(): + from freetoken.layers import BaseOP + from freetoken.moe.offload_cache import attach_owner_moe_cache + + layer, owner = _make_owner_layer() + model = BaseOP() + model.block = layer + layers = attach_owner_moe_cache(model, owner) + + assert layers == [layer] + assert layer.owner_cache is owner + assert layer.offload_cache is owner diff --git a/tests/moe/test_ownership.py b/tests/moe/test_ownership.py new file mode 100644 index 000000000..860525ea8 --- /dev/null +++ b/tests/moe/test_ownership.py @@ -0,0 +1,482 @@ +"""CPU contracts for the first TP+EP ownership boundary. + +These tests intentionally stop before slot-cache admission: global router IDs, local bank rows, +and cache slots are separate namespaces until the EP runtime path is complete. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from freetoken.models.nvfp4_banks import ( + Nvfp4ExpertSourceSpec, + iter_nvfp4_expert_pieces, +) +from freetoken.moe.ownership import ( + ExpertOwnership, + OwnerCacheAdapter, + OwnerCacheGeometry, +) + + +_GENERIC_RE = re.compile( + r"^layer\.(?P\d+)\.expert\.(?P\d+)\." + r"(?Pgate|up|down)\.(?Pweight|weight_scale|weight_scale_2)$" +) +_GENERIC_SPEC = Nvfp4ExpertSourceSpec( + key_pattern=_GENERIC_RE, + proj_to_role={"gate": "gate", "up": "up", "down": "down"}, + layer_to_bank=lambda layer, _config: layer, + desc="test NVFP4 experts", +) + + +def test_contiguous_ownership_covers_global_experts_once(): + owners = [ExpertOwnership(8, 2, rank) for rank in range(2)] + assert [o.local_num_experts for o in owners] == [4, 4] + assert [o.global_start for o in owners] == [0, 4] + assert [o.global_end for o in owners] == [4, 8] + assert [owners[0].owner(i) for i in range(8)] == [0, 0, 0, 0, 1, 1, 1, 1] + + for rank, owner in enumerate(owners): + local, mask = owner.global_to_local(torch.arange(8, dtype=torch.int32)) + expected = list(range(4)) if rank == 0 else [-1] * 4 + list(range(4)) + if rank == 0: + expected = list(range(4)) + [-1] * 4 + assert local.tolist() == expected + assert int(mask.sum()) == 4 + assert owner.local_to_global(torch.arange(4, dtype=torch.int32)).tolist() == list( + range(owner.global_start, owner.global_end) + ) + + +def test_ownership_rejects_invalid_geometry_and_ids(): + with pytest.raises(ValueError, match="not divisible"): + ExpertOwnership(5, 2, 0) + with pytest.raises(ValueError, match="outside"): + ExpertOwnership(8, 2, 2) + + owner = ExpertOwnership(8, 2, 0) + with pytest.raises(ValueError, match="global expert id"): + owner.owner(8) + with pytest.raises(ValueError, match="local expert IDs"): + owner.local_to_global(torch.tensor([-1])) + with pytest.raises(ValueError, match="router expert IDs"): + owner.validate_global_ids(torch.tensor([0, 8])) + + +@pytest.mark.parametrize( + "route_ids", + [ + [0, 1, 2, 3, 4, 5, 6, 7, 0, 7], # 5/5, including duplicates + [0] * 10, # rank 0 owns all entries + [4] * 10, # rank 1 owns all entries + ], +) +def test_partition_route_masks_remote_entries_without_local_renormalization(route_ids): + ids = torch.tensor([route_ids], dtype=torch.int32) + weights = torch.arange(1, 11, dtype=torch.float32).reshape(1, 10) / 55 + owners = [ExpertOwnership(8, 2, rank) for rank in range(2)] + routes = [owner.partition_route(weights, ids) for owner in owners] + + assert torch.equal(routes[0].weights + routes[1].weights, weights) + assert torch.equal(routes[0].owned_mask | routes[1].owned_mask, torch.ones_like(ids, dtype=torch.bool)) + assert torch.equal(routes[0].owned_mask & routes[1].owned_mask, torch.zeros_like(ids, dtype=torch.bool)) + for route in routes: + assert torch.all(route.local_ids >= 0) + assert torch.all(route.local_ids < 4) + assert torch.equal(route.weights[~route.owned_mask], torch.zeros_like(route.weights[~route.owned_mask])) + + +def test_partition_route_zero_local_entries_uses_safe_placeholder(): + owner = ExpertOwnership(8, 2, 0) + ids = torch.tensor([[4, 5, 6, 7]], dtype=torch.int32) + weights = torch.full((1, 4), 0.25, dtype=torch.float32) + route = owner.partition_route(weights, ids) + assert route.owned_count == 0 and route.remote_count == 4 + assert torch.equal(route.local_ids, torch.zeros_like(ids)) + assert torch.equal(route.weights, torch.zeros_like(weights)) + + +def test_partition_route_rejects_shape_and_dtype_mismatches(): + owner = ExpertOwnership(8, 2, 0) + with pytest.raises(ValueError, match="same shape"): + owner.partition_route(torch.ones(2), torch.zeros(1, dtype=torch.int32)) + with pytest.raises(TypeError, match="integer tensor"): + owner.partition_route(torch.ones(1), torch.zeros(1, dtype=torch.float32)) + + +def test_owner_cache_geometry_separates_global_local_and_flat_namespaces(): + geometry = OwnerCacheGeometry( + global_num_experts=8, + world_size=2, + rank=1, + num_layers=3, + cache_size=8, + prefill_overlap=True, + ) + + assert geometry.local_num_experts == 4 + assert (geometry.global_start, geometry.global_end) == (4, 8) + ids = torch.tensor([[0, 4, 7]], dtype=torch.int32) + local, owned = geometry.global_to_local(ids) + assert local.tolist() == [[-1, 0, 3]] + assert owned.tolist() == [[False, True, True]] + flat, flat_owned = geometry.global_to_local_flat(2, ids) + assert flat.tolist() == [[8, 8, 11]] + assert torch.equal(flat_owned, owned) + decoded_layer, decoded_local = geometry.flat_id_to_local(flat) + assert decoded_layer.tolist() == [[2, 2, 2]] + assert decoded_local.tolist() == [[0, 0, 3]] + assert geometry.local_to_flat_id(2, torch.tensor([0, 3], dtype=torch.int32)).tolist() == [8, 11] + all_flat = torch.arange(geometry.num_layers * geometry.local_num_experts, dtype=torch.int32) + layers, local_rows = geometry.flat_id_to_local(all_flat) + rebuilt = torch.cat([ + geometry.local_to_flat_id(layer, local_rows[layers == layer]) + for layer in range(geometry.num_layers) + ]) + assert torch.equal(rebuilt, all_flat) + assert layers.tolist() == [0] * 4 + [1] * 4 + [2] * 4 + with pytest.raises(ValueError, match="owner flat IDs"): + geometry.flat_id_to_local(torch.tensor([12], dtype=torch.int32)) + + route = geometry.partition_route( + torch.tensor([[0.2, 0.3, 0.5]]), ids + ) + assert route.local_ids.tolist() == [[0, 0, 3]] + assert torch.allclose(route.weights, torch.tensor([[0.0, 0.3, 0.5]])) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"global_num_experts": 7, "world_size": 2}, "not divisible"), + ({"global_num_experts": 8, "world_size": 2, "rank": 2}, "outside"), + ({"global_num_experts": 8, "world_size": 2, "cache_size": 3}, "smaller"), + ({"global_num_experts": 8, "world_size": 2, "cache_size": 4, "prefill_overlap": True}, "2 \\* local"), + ], +) +def test_owner_cache_geometry_rejects_invalid_capacity_or_owner(kwargs, message): + params = { + "global_num_experts": 8, + "world_size": 2, + "rank": 0, + "num_layers": 2, + "cache_size": 8, + "prefill_overlap": False, + } + params.update(kwargs) + with pytest.raises(ValueError, match=message): + OwnerCacheGeometry(**params) + + +def test_owner_cache_geometry_validates_legacy_binding_and_local_bank_shapes(): + geometry = OwnerCacheGeometry(8, 2, 0, num_layers=2, cache_size=6) + geometry.validate_cache_binding( + num_layers=2, num_experts=8, cache_size=6, prefill_overlap=False + ) + with pytest.raises(ValueError, match="num_experts"): + geometry.validate_cache_binding( + num_layers=2, num_experts=4, cache_size=6, prefill_overlap=False + ) + + good_sources = { + "gate_up": [torch.empty(4, 8, 16) for _ in range(2)], + "down": [torch.empty(4, 16, 4) for _ in range(2)], + } + geometry.validate_source_banks(good_sources) + with pytest.raises(ValueError, match="first dimension 4"): + geometry.validate_source_banks( + {"gate_up": [torch.empty(8, 8, 16) for _ in range(2)]} + ) + + geometry.validate_slot_maps( + torch.full((2, 4), -1, dtype=torch.int32), + torch.full((6,), -1, dtype=torch.int32), + ) + with pytest.raises(ValueError, match="slot_for_id shape"): + geometry.validate_slot_maps( + torch.full((2, 8), -1, dtype=torch.int32), + torch.full((6,), -1, dtype=torch.int32), + ) + with pytest.raises(ValueError, match="entries"): + geometry.validate_slot_maps( + torch.tensor([[6, -1, -1, -1], [-1, -1, -1, -1]], dtype=torch.int32), + torch.full((6,), -1, dtype=torch.int32), + ) + + +def test_offload_cache_owner_geometry_is_default_off_and_fail_fast_when_explicit(): + from freetoken.moe.offload_cache import OffloadMoeCache + + # Existing callers do not pass owner_geometry and retain the global cache geometry. + cache = OffloadMoeCache( + num_layers=1, + num_experts=4, + cache_size=4, + device=torch.device("cpu"), + ) + assert cache.slot_for_id.shape == (1, 4) + + geometry = OwnerCacheGeometry(8, 2, 0, num_layers=1, cache_size=4) + with pytest.raises(NotImplementedError, match="namespace mapping"): + OffloadMoeCache( + num_layers=1, + num_experts=8, + cache_size=4, + device=torch.device("cpu"), + owner_geometry=geometry, + ) + + +def test_owner_cache_adapter_rewrites_owned_routes_and_preserves_weights(): + geometry = OwnerCacheGeometry(8, 2, 1, num_layers=2, cache_size=4) + cache = OwnerCacheAdapter(geometry) + ids = torch.tensor([[0, 4, 7, 5, 4]], dtype=torch.int32) + weights = torch.tensor([[0.1, 0.2, 0.3, 0.15, 0.25]]) + + update = cache.ensure_route(1, weights, ids) + + assert update.owned_mask.tolist() == [[False, True, True, True, True]] + assert update.local_ids.tolist() == [[0, 0, 3, 1, 0]] + assert update.local_flat_ids.tolist() == [[0, 4, 7, 5, 4]] + assert torch.allclose(update.weights, torch.tensor([[0.0, 0.2, 0.3, 0.15, 0.25]])) + assert update.missing_local_ids.tolist() == [0, 3, 1] + assert torch.all(update.slot_ids[update.owned_mask] >= 0) + assert torch.equal(update.slot_ids[~update.owned_mask], torch.zeros(1, dtype=torch.int32)) + assert cache.resident == 3 + + +def test_owner_cache_adapter_remote_only_route_is_safe_and_does_not_admit(): + geometry = OwnerCacheGeometry(8, 2, 0, num_layers=1, cache_size=4) + cache = OwnerCacheAdapter(geometry) + ids = torch.tensor([[4, 5, 6, 7]], dtype=torch.int32) + weights = torch.full((1, 4), 0.25) + + update = cache.ensure_route(0, weights, ids) + + assert update.missing_local_ids.numel() == 0 + assert update.evicted_flat_ids.numel() == 0 + assert torch.equal(update.slot_ids, torch.zeros_like(ids)) + assert torch.equal(update.weights, torch.zeros_like(weights)) + assert cache.resident == 0 + assert cache.step == 1 # an empty owner route still advances one logical LRU call + + +@pytest.mark.parametrize("rank, local_id, remote_id", [(0, 0, 4), (1, 4, 0)]) +def test_owner_cache_adapter_all_local_and_all_remote_have_no_cross_admission( + rank, local_id, remote_id +): + geometry = OwnerCacheGeometry(8, 2, rank, num_layers=1, cache_size=4) + cache = OwnerCacheAdapter(geometry) + weights = torch.full((1, 10), 0.1) + + local = cache.ensure_route(0, weights, torch.full((1, 10), local_id, dtype=torch.int32)) + assert local.missing_local_ids.tolist() == [0] + assert cache.resident == 1 + + remote = cache.ensure_route(0, weights, torch.full((1, 10), remote_id, dtype=torch.int32)) + assert remote.missing_local_ids.numel() == 0 + assert remote.evicted_flat_ids.numel() == 0 + assert torch.equal(remote.slot_ids, torch.zeros_like(remote.slot_ids)) + assert torch.equal(remote.weights, torch.zeros_like(remote.weights)) + assert cache.resident == 1 + + +def test_owner_cache_adapter_lru_is_layer_local_but_pool_is_unified(): + geometry = OwnerCacheGeometry(4, 2, 0, num_layers=2, cache_size=2) + cache = OwnerCacheAdapter(geometry) + cache.ensure_route(0, torch.ones(1, 2), torch.tensor([[0, 1]], dtype=torch.int32)) + update = cache.ensure_route(1, torch.ones(1, 1), torch.tensor([[0]], dtype=torch.int32)) + + assert update.evicted_flat_ids.tolist() == [0] + assert cache.slot_for_id.tolist() == [[-1, 1], [0, -1]] + cache.validate_invariants() + + +def test_owner_cache_adapter_materialize_uses_two_local_prefill_buffers(): + geometry = OwnerCacheGeometry(8, 2, 0, num_layers=2, cache_size=8, prefill_overlap=True) + cache = OwnerCacheAdapter(geometry) + assert cache.materialize_layer(0, buffer_id=0).tolist() == [0, 1, 2, 3] + assert cache.materialize_layer(1, buffer_id=1).tolist() == [4, 5, 6, 7] + assert cache.slot_for_id.tolist() == [[0, 1, 2, 3], [4, 5, 6, 7]] + + cache.materialize_layer(0, buffer_id=0) + assert cache.slot_for_id.tolist() == [[0, 1, 2, 3], [4, 5, 6, 7]] + with pytest.raises(ValueError, match="buffer_id"): + cache.materialize_layer(1, buffer_id=2) + + +def test_owner_offload_cache_compacts_routes_before_legacy_admission(monkeypatch): + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + + geometry = OwnerCacheGeometry(8, 2, 1, num_layers=1, cache_size=4) + cache = OwnerOffloadMoeCache(geometry, torch.device("cpu")) + calls = [] + + def fake_ensure(layer_id, local_ids): + calls.append((layer_id, local_ids.clone())) + assert layer_id == 0 + # The legacy kernel has already accepted a local-row route and rewrites it to slots. + slot_by_local = {0: 2, 1: 1, 3: 0} + slots = torch.tensor( + [slot_by_local[int(local)] for local in local_ids.tolist()], + dtype=torch.int32, + ) + cache._cache.num_indices.fill_(3) + cache._cache.src_indices[:3] = torch.tensor([0, 3, 1], dtype=torch.int32) + cache._cache.evict_slots[:3] = torch.tensor([2, 1, 0], dtype=torch.int32) + for local, slot in slot_by_local.items(): + cache._cache.slot_for_id[layer_id, local] = slot + cache._cache.id_of_slot[slot] = local + local_ids.copy_(slots) + + monkeypatch.setattr(cache._cache, "ensure_experts", fake_ensure) + copied = [] + monkeypatch.setattr(cache._cache, "copy_missing", lambda: copied.append(True)) + + ids = torch.tensor([[0, 4, 7, 5, 4]], dtype=torch.int32) + weights = torch.tensor([[0.1, 0.2, 0.3, 0.15, 0.25]]) + update = cache.ensure_route(0, weights, ids) + + assert len(calls) == 1 + assert calls[0][1].tolist() == [0, 3, 1, 0] + assert update.slot_ids.tolist() == [[0, 2, 0, 1, 2]] + assert update.local_ids.tolist() == [[0, 0, 3, 1, 0]] + assert torch.allclose( + update.weights, torch.tensor([[0.0, 0.2, 0.3, 0.15, 0.25]]) + ) + assert update.missing_local_ids.tolist() == [0, 3, 1] + assert update.evicted_flat_ids.tolist() == [] + + cache.copy_missing() + assert copied == [True] + with pytest.raises(RuntimeError, match="ensure_route"): + cache.ensure_experts(0, torch.tensor([0], dtype=torch.int32)) + + +def test_owner_offload_cache_remote_only_route_does_not_stage_copy(): + from freetoken.moe.offload_cache import OwnerOffloadMoeCache + + geometry = OwnerCacheGeometry(8, 2, 0, num_layers=1, cache_size=4) + cache = OwnerOffloadMoeCache(geometry, torch.device("cpu")) + ids = torch.tensor([[4, 5, 6, 7]], dtype=torch.int32) + weights = torch.full((1, 4), 0.25) + + update = cache.ensure_route(0, weights, ids) + + assert update.slot_ids.tolist() == [[0, 0, 0, 0]] + assert update.weights.tolist() == [[0.0, 0.0, 0.0, 0.0]] + assert update.missing_local_ids.numel() == 0 + assert cache._pending_owned is False + + +def _write_tiny_nvfp4_checkpoint(folder: Path, *, experts: int = 4) -> dict[str, torch.Tensor]: + """Create one native-bank-shaped layer with distinct data per global expert.""" + H = I = 16 # keep every bank dimension divisible by the native 16-byte scale blocks + tensors: dict[str, torch.Tensor] = {} + for expert in range(experts): + base = f"layer.0.expert.{expert}" + for proj, out, inn in (("gate", I, H), ("up", I, H), ("down", H, I)): + tensors[f"{base}.{proj}.weight"] = torch.full( + (out, inn // 2), expert + (1 if proj == "up" else 11 if proj == "down" else 101), + dtype=torch.uint8, + ) + tensors[f"{base}.{proj}.weight_scale"] = torch.full( + (out, inn // 16), expert + 1, dtype=torch.float8_e4m3fn + ) + # Give every expert/projection a different global scale. The test catches a + # local-ID lookup here: rank 1 row 0 must receive expert 2's scale, not expert 0's. + tensors[f"{base}.{proj}.weight_scale_2"] = torch.tensor( + 10.0 + expert * 3 + {"gate": 0, "up": 1, "down": 2}[proj], + dtype=torch.float32, + ) + + shard = folder / "model.safetensors" + save_file(tensors, str(shard)) + (folder / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {name: shard.name for name in tensors}}), + encoding="utf-8", + ) + return tensors + + +def test_owner_reader_yields_local_rows_and_keeps_global_scales(tmp_path, monkeypatch): + raw = _write_tiny_nvfp4_checkpoint(tmp_path) + monkeypatch.setattr( + "freetoken.models.nvfp4_banks.download_hf_weight", lambda _path: str(tmp_path) + ) + + config = SimpleNamespace( + num_experts=4, + hidden_size=16, + moe_intermediate_size=16, + num_moe_layers=1, + ) + owner = ExpertOwnership(global_num_experts=4, world_size=2, rank=1) + pieces = list( + iter_nvfp4_expert_pieces( + str(tmp_path), + config, + _GENERIC_SPEC, + drop_page_cache=lambda _path: None, + primary=False, + ownership=owner, + ) + ) + + # The owner has global experts [2, 4), renumbered into LOCAL rows [0, 2): no full-E + # read and no remote row may appear in the stream. + assert [(layer, e0, e1) for layer, e0, e1, _ in pieces] == [(0, 0, 1), (0, 1, 2)] + by_local = {e0: piece for _, e0, _, piece in pieces} + + for local, global_expert in enumerate((2, 3)): + piece = by_local[local] + for proj in ("gate", "up", "down"): + assert torch.equal( + piece[proj][0], raw[f"layer.0.expert.{global_expert}.{proj}.weight"] + ) + assert torch.equal( + piece[f"{proj}_scale"][0], + raw[f"layer.0.expert.{global_expert}.{proj}.weight_scale"], + ) + # Global scale lookup must use checkpoint ID 2/3 while the destination row is 0/1. + expected = raw[f"layer.0.expert.{global_expert}.{proj}.weight_scale_2"].to( + torch.float16 + ) + got = piece[f"{proj}_global"][0] + assert torch.equal(got, expected.expand_as(got)), (local, global_expert, proj) + + # No remote expert row is present. + assert not torch.equal(by_local[0]["gate"][0], raw["layer.0.expert.0.gate.weight"]) + + +def test_owner_reader_rejects_ownership_geometry_mismatch(tmp_path, monkeypatch): + _write_tiny_nvfp4_checkpoint(tmp_path) + monkeypatch.setattr( + "freetoken.models.nvfp4_banks.download_hf_weight", lambda _path: str(tmp_path) + ) + config = SimpleNamespace( + num_experts=4, + hidden_size=16, + moe_intermediate_size=16, + num_moe_layers=1, + ) + with pytest.raises(ValueError, match="global_num_experts"): + list( + iter_nvfp4_expert_pieces( + str(tmp_path), + config, + _GENERIC_SPEC, + drop_page_cache=lambda _path: None, + primary=False, + ownership=ExpertOwnership(global_num_experts=8, world_size=2, rank=0), + ) + ) diff --git a/tests/moe/test_route_trace.py b/tests/moe/test_route_trace.py new file mode 100644 index 000000000..c9147416b --- /dev/null +++ b/tests/moe/test_route_trace.py @@ -0,0 +1,214 @@ +"""Tests for moe/route_trace.py: capture round-trip, LRU-mirror semantics, EP2 slow-side. + +Pure CPU, no torch/CUDA needed for the replay/LRU parts (the recorder's ``record`` +takes a tensor, so it uses a tiny stub). Mirrors PLAN_TP_EP.md 6A's requirement that +the replay match flashlib lru_ensure semantics before it is trusted for the EP2 call. +""" +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +# Load route_trace.py directly by path: it is stdlib-only, so importing it without +# the freetoken.moe package __init__ (which pulls torch/transformers) keeps these +# tests runnable on a bare interpreter. Resolve it relative to THIS file so the test +# collects in any checkout (an absolute author path fails everywhere else). +_MOD = Path(__file__).resolve().parents[2] / "python" / "freetoken" / "moe" / "route_trace.py" +_spec = importlib.util.spec_from_file_location("route_trace", _MOD) +_rt = importlib.util.module_from_spec(_spec) +sys.modules["route_trace"] = _rt # @dataclass resolves cls.__module__ via sys.modules +_spec.loader.exec_module(_rt) +LRU = _rt.LRU +RouteTraceRecorder = _rt.RouteTraceRecorder +read_trace = _rt.read_trace +replay = _rt.replay +replay_ep2 = _rt.replay_ep2 + + +class _FakeTensor: + """Minimal stand-in for the raw expert-ids tensor ``record`` receives.""" + + def __init__(self, ids): + self._ids = list(ids) + + def reshape(self, _): + return self + + def tolist(self): + return self._ids + + +def test_roundtrip(tmp_path): + import struct + + path = str(tmp_path / "tr.bin") + rec = RouteTraceRecorder( + path, num_experts=512, num_layers=48, cache_size=2840, top_k=10, + model="flash", decode_target="gpu", + ) + # emulate record() without torch: write the same binary layout by hand + def put(phase, layer, ids): + rec._buf += struct.Struct(" expert 2 is a protected hit, only 0 misses (v1 bug = 2) + c = LRU(2, 512) + c.ensure(0, [1, 2]) + before = c.miss + c.ensure(0, [0, 2]) + assert c.miss - before == 1, f"expected 1 miss, got {c.miss - before}" + + +def test_lru_stack_property(): + # same access sequence, larger cache never increases miss (standard LRU property, + # the theoretical basis for the capacity-expansion main line) + seq = [(0, [i % 40 for i in range(s, s + 10)]) for s in range(0, 300, 3)] + counts = [] + for C in (8, 16, 32, 64): + c = LRU(C, 512) + for layer, ids in seq: + c.ensure(layer, ids) + counts.append(c.miss) + assert all(counts[i] >= counts[i + 1] for i in range(len(counts) - 1)), counts + + +def test_lru_dedup(): + c = LRU(8, 512) + c.ensure(0, [3, 3, 3, 7]) + assert c.miss == 2 and c.active == 2 + + +def test_replay_tp1_vs_ep2(tmp_path): + import struct + + path = str(tmp_path / "r.bin") + rec = RouteTraceRecorder( + path, num_experts=512, num_layers=2, cache_size=100, top_k=8, model="t", + decode_target="gpu", + ) + + def put(layer, ids): + rec._buf += struct.Struct("> pool (40), so TP1 thrashes; under + # EP2 each rank only caches its own 60, and a 40-slot pool covers 40/60 of it + # vs 40/120 for TP1 -> strictly fewer misses on the slow rank. + import random + rng = random.Random(0) + window = list(range(60)) + list(range(256, 316)) + for _ in range(400): + put(0, rng.sample(window, 8)) + put(1, rng.sample(window, 8)) + rec.close() + + meta, records = read_trace(path) + _, _, tp1 = replay(records, 40, 512) + ep = replay_ep2(records, (40, 40), 512) + assert ep["slow_rate"] < tp1, f"EP2 slow {ep['slow_rate']} should beat TP1 {tp1}" + # symmetric window -> roughly balanced load per rank (random draw, not exact) + a0, a1 = ep["rank_active"] + assert abs(a0 - a1) <= 0.1 * max(a0, a1), (a0, a1) + # slow-side miss >= each rank's miss (it is the per-step max), so it is the limiter + assert ep["slow_miss"] >= max(ep["rank_miss"]) + assert tp1 > 0.3, f"expected real cache pressure, TP1 miss={tp1}" + + +def test_overflow_flag(tmp_path): + import struct + + path = str(tmp_path / "o.bin") + rec = RouteTraceRecorder( + path, num_experts=512, num_layers=1, cache_size=10, top_k=2, model="t", + decode_target="gpu", max_records=3, + ) + + def put(ids): + if rec._overflow: + return + if rec._n >= rec.max_records: + rec._overflow = True + return + rec._buf += struct.Struct("