"""共享装配层:把"文件后端流式加载主模型 + 取 hidden + 贪心生成"等跨入口复用的逻辑 集中在这里,而不是藏在某个入口脚本(原 validate_mtp)里。 依赖关系:本模块位于 core(mem/expert_store/streaming_moe)与 mtp 之上, 是把它们粘合成可运行模型的装配层,被 cli/ 各入口与测试复用。 环境变量: MODEL 主模型路径(MLX 量化) EXPERT_DIR 拆分/重量化后的 per-expert safetensors 目录 EXPERT_SLOTS 每层常驻池容量 EXPERT_POOL_PROFILE 每层池预算 JSON(无损省内存,可选) HIDDEN_VARIANT pre_final_norm(默认)| post_final_norm(排错时切换) """ import json import os import mlx.core as mx from mlx_lm import load from mlx_lm.models.base import create_attention_mask, create_ssm_mask from mlx_streaming import config from mlx_streaming.core.cache.expert_store import FileExpertStore from mlx_streaming.core.prefetch.patch import patch_model_filebacked MODEL = config.model_path() EXPERT_DIR = config.expert_dir() EXPERT_SLOTS = config.expert_slots() # pre_final_norm(默认)| post_final_norm(排错时切换) HIDDEN_VARIANT = config.hidden_variant() # 默认 profile 文件名:放在 EXPERT_DIR 下随专家目录一起走,常用跑法自动启用(无损省内存) DEFAULT_PROFILE_NAME = "pool_profile.json" def load_pool_profile(expert_dir: str) -> "dict[int, int] | None": """解析每层池预算 profile,返回 layer_caps 或 None。 优先级:环境变量 EXPERT_POOL_PROFILE 显式指定路径 > {expert_dir}/pool_profile.json 默认。 EXPERT_POOL_PROFILE=none/0/off 显式关闭(回到 uniform capacity)。 profile 无损:仅按各层真实工作集分配,命中率/输出/吞吐不变(caps 仍被 capacity 上限钳制)。 """ p = config.expert_pool_profile() if p.lower() in ("none", "0", "off"): return None if not p: # 未显式指定 → 默认找专家目录下的 profile cand = os.path.join(expert_dir, DEFAULT_PROFILE_NAME) p = cand if os.path.exists(cand) else "" if p and os.path.exists(p): with open(p) as f: caps = json.load(f).get("layer_caps", {}) return {int(k): int(v) for k, v in caps.items()} return None def build_streaming_model(): """用文件后端流式 patch 加载主模型(32GB 机器装不下 41GB 非流式)。""" model, tok = load(MODEL, lazy=True) # 取首个 MoE 维度 dims = None for layer in model.layers: mlp = getattr(layer, "mlp", None) if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"): gp = mlp.switch_mlp.gate_proj dims = {"hidden": gp.input_dims, "moe_inter": gp.output_dims, "group_size": getattr(gp, "group_size", 64), "bits": getattr(gp, "bits", 4)} break bits, group, proj_bits, layer_proj_bits = ( dims["bits"], dims["group_size"], None, None) meta_path = os.path.join(EXPERT_DIR, "_split_meta.json") if os.path.exists(meta_path): with open(meta_path) as f: meta = json.load(f) ed = meta.get("dims", {}) bits = ed.get("bits", bits) group = ed.get("group_size", group) proj_bits = ed.get("proj_bits") if "per_layer_proj_bits" in ed: layer_proj_bits = {int(k): v for k, v in ed["per_layer_proj_bits"].items()} # 每层池预算 profile(pool_footprint 产出):默认从 {EXPERT_DIR}/pool_profile.json 自动启用, # 无损省内存(命中率/输出/吞吐不变，仅不再为低占用层预留满 capacity)。 layer_caps = load_pool_profile(EXPERT_DIR) store = FileExpertStore(EXPERT_DIR, capacity=EXPERT_SLOTS, layer_caps=layer_caps) if config.zerocopy_dual_source(): # 零拷贝双源双缓冲：常驻池换成侧区模式（预分配 cap+2*spec_slots 行、禁 grow），复用原池 loader/cap/profile。 from mlx_streaming.core.cache.resident_pool import ResidentExpertPool _old = store._resident # 默认单缓冲(持久 LFU,一份工作集,省一半侧区内存=生产路径);仅显式 legacy(SIDEREGION_LFU=0)用双缓冲。 _spec_gens = 1 if config.sideregion_lfu() else 2 store._resident = ResidentExpertPool( _old.capacity, loader=_old.loader, layer_caps=_old.layer_caps, spec_slots=config.pool_spec_slots(), spec_gens=_spec_gens) if config.stream_blob_loader(): # blob 接入常驻池 miss-loader：复用 GPU-remap 快路径，小 EXPERT_SLOTS 即低内存。 store._blob_loader = _make_blob_source(dims, group, bits) # 主动预取（native-fused-prefetch miss→hit）：opt-in（NATIVE_FUSED_PREFETCH=1）。 # 经"promote 只写真实路由命中专家"修正后已是净正：易缓存基座上 +15.5% tok/s # （demand 11.86→13.70，hit 0.731→0.851，读盘 −45%；见 active-prefetch-turnaround-2026-06-17.md）。 # 默认关只因收益依赖场景（基座可缓存性/是否磁盘受限）且只影响速度不影响质量、多占少量 staging 内存， # 故作 opt-in 而非默认路径，落地配方见上述报告 §6。 if config.native_fused_prefetch() and getattr(store, "_blob_loader", None) is not None: try: import mlx_streaming.native_moe_ext # noqa: F401 确认扩展已编译 from mlx_streaming.core.prefetch.native_staging import NativeStagingManager _budget = (config.pool_spec_slots() if config.zerocopy_dual_source() else config.stream_blob_bg_budget(default=16)) store._staging = NativeStagingManager(store._blob_loader, budget=_budget) except Exception: store._staging = None # 扩展不可用 → 关闭，不影响主路径 # 零拷贝双源不变量：staging 侧区行数必须等于池 spec_slots，否则 C++ 会越界写池（静默损坏）。 if config.zerocopy_dual_source() and getattr(store, "_staging", None) is not None: assert store._staging.budget == store._resident.spec_slots, ( f"零拷贝双源要求 staging.budget({store._staging.budget}) " f"== 池 spec_slots({store._resident.spec_slots})") if config.stream_blob_bg(): # 后台预取池预填：bg 在独立 stream 物化预测专家，promote 写进池槽（需 CROSS_LAYER_PREFETCH=1）。 from mlx_streaming.core.prefetch.bg_prefetch import BackgroundExpertPrefetcher src = _make_blob_source(dims, group, bits) store._blob_loader = src store._bg = BackgroundExpertPrefetcher( src, window=config.stream_blob_window()) patch_model_filebacked(model, store, dims["hidden"], dims["moe_inter"], group, bits, proj_bits=proj_bits, layer_proj_bits=layer_proj_bits) # 双源双缓冲：构造一个共享 VirtualPool（gen 跨层全局、每前向 +1）挂到每个流式 MoE 块。 if config.zerocopy_dual_source() and getattr(store, "_staging", None) is not None: from mlx_streaming.core.cache.virtual_pool import VirtualPool from mlx_streaming.core.moe.block import FileStreamingMoeBlock # 双源模式仍需 ahead 调度：block._native_fused_prefetch 靠 target_for 选目标层， # 不传调度参数会让 target_for 恒返回 0（_num_layers=0）→ 预取全跳过、侧区永远空。 _vpool = VirtualPool(store._resident, store._staging, config.pool_spec_slots(), num_layers=len(model.layers), cutoff=config.cross_layer_cutoff(), ahead_lo=config.cross_layer_ahead_lo(), ahead_hi=config.cross_layer_ahead_hi(), store=store) for layer in model.layers: mlp = getattr(layer, "mlp", None) if isinstance(mlp, FileStreamingMoeBlock): mlp._vpool = _vpool # 主动预取（非 zerocopy）：挂 per-layer ahead 调度器 vpool（cutoff），让晚层预读更早发起。 if (config.native_fused_prefetch() and not config.zerocopy_dual_source() and getattr(store, "_staging", None) is not None): from mlx_streaming.core.cache.virtual_pool import VirtualPool from mlx_streaming.core.moe.block import FileStreamingMoeBlock _sched = VirtualPool(num_layers=len(model.layers), cutoff=config.cross_layer_cutoff(), ahead_lo=config.cross_layer_ahead_lo(), ahead_hi=config.cross_layer_ahead_hi()) for layer in model.layers: mlp = getattr(layer, "mlp", None) if isinstance(mlp, FileStreamingMoeBlock): mlp._vpool = _sched if config.stream_blob(): _attach_blob_source(model, dims, group, bits) # KV 量化(IsoQuant K4/V3 + SO(4) 旋转):仅作用于 12 个全注意力层,128k KV 3.0→~0.68 GiB。 if config.kv_quant(): from mlx_streaming.core.cache.kv_quant_patch import patch_kv_quant patch_kv_quant(model, group_size=config.kv_group_size(), k_bits=config.kv_k_bits(), v_bits=config.kv_v_bits(), rotate=config.kv_rotate(), seed=config.kv_rot_seed()) # AUTOPIN 预热钉死(池 + blob_loader + staging 就绪后、return 前):按历史路由热度把每层 # top-N 热专家预填进常驻池并 pin 住(不参与任何驱逐),消灭冷启动慢热;usage 缺失则跳过 # (会话中计数仍在累计,落盘供下次启动)。AUTOPIN=0 零行为。 if config.autopin(): from mlx_streaming.core.cache import autopin as _autopin _ap = _autopin.warm_start_pin(store) if _ap["pinned"]: print(f"[AUTOPIN] 预热 pin: {_ap['layers']} 层共 {_ap['pinned']} 专家 " f"(frac={config.autopin_budget_frac()}, 耗时 {_ap['seconds']}s)" + (f", 跳过 {_ap['skipped_layers']} 层" if _ap["skipped_layers"] else ""), flush=True) # 把 store 挂到 model 上，供只拿到 model 的调用方(server /api/stats 读专家池健康度)取用。 # 非 array/dict/list/tuple 走 object 属性，不进 nn.Module 参数树。 model._expert_store = store return model, tok, store def _make_blob_source(dims, group, bits): from mlx_streaming.core.cache.blob_loader import BlobExpertSource blob_dir = config.blob_dir() or os.path.join(EXPERT_DIR, "blobs") workers = config.stream_blob_workers() nocache = config.stream_blob_nocache(default="0") num_experts = 512 idx_path = os.path.join(blob_dir, "blob_index.json") if os.path.exists(idx_path): with open(idx_path) as f: num_experts = int(json.load(f).get("num_experts", num_experts)) return BlobExpertSource(blob_dir, dims["hidden"], dims["moe_inter"], group, bits, num_experts=num_experts, workers=workers, nocache=nocache) def _attach_blob_source(model, dims, group, bits): """STREAM_BLOB=1：给每个流式 MoE 块注入共享 BlobExpertSource（全流式低内存路径）。""" from mlx_streaming.core.cache.blob_loader import BlobExpertSource from mlx_streaming.core.moe.block import FileStreamingMoeBlock blob_dir = config.blob_dir() or os.path.join(EXPERT_DIR, "blobs") workers = config.stream_blob_workers() nocache = config.stream_blob_nocache(default="1") num_experts = 512 idx_path = os.path.join(blob_dir, "blob_index.json") if os.path.exists(idx_path): with open(idx_path) as f: num_experts = int(json.load(f).get("num_experts", num_experts)) src = BlobExpertSource(blob_dir, dims["hidden"], dims["moe_inter"], group, bits, num_experts=num_experts, workers=workers, nocache=nocache) for layer in model.layers: mlp = getattr(layer, "mlp", None) if isinstance(mlp, FileStreamingMoeBlock): mlp._blob = src def capture_prenorm_hidden(model, input_ids: mx.array) -> mx.array: """跑主模型层循环但跳过最后的 model.norm,返回 last-layer hidden(norm 前)。 HIDDEN_VARIANT=post_final_norm 时返回 norm 之后(用于消歧排错)。 """ inner = model.model h = inner.embed_tokens(input_ids) layers = inner.layers if not layers: return h cache = model.make_cache() fa_idx = next((i for i, l in enumerate(layers) if not l.is_linear), 0) ssm_idx = next((i for i, l in enumerate(layers) if l.is_linear), 0) fa_mask = create_attention_mask(h, cache[fa_idx]) ssm_mask = create_ssm_mask(h, cache[ssm_idx]) for layer, c in zip(layers, cache): mask = ssm_mask if layer.is_linear else fa_mask h = layer(h, mask=mask, cache=c) if HIDDEN_VARIANT == "post_final_norm": h = inner.norm(h) return h def greedy(model, input_ids: mx.array, n: int) -> mx.array: """主模型贪心生成 n 个 token,返回拼接后的完整序列(用作自投机参考)。""" cache = model.make_cache() cur = input_ids out = [] for _ in range(n): logits = model(cur, cache=cache) nxt = mx.argmax(logits[:, -1, :], axis=-1, keepdims=True) out.append(nxt) cur = nxt mx.eval(nxt) return mx.concatenate([input_ids] + out, axis=1)