432 lines
25 KiB
Python
432 lines
25 KiB
Python
"""MoE 热路径块：路由 → 选专家 → 取专家权重 → 计算 → 加权合并。
|
||
|
||
包含两种块：
|
||
- `StreamingMoeBlock`：包住原生 MoE 块，专家权重常驻（switch_mlp），只激活选中专家。
|
||
- `FileStreamingMoeBlock`：专家权重从磁盘按需加载（流式），是低内存推理的核心热路径，
|
||
集成常驻池 acquire、native-fused-prefetch 预取、共享专家叠加等逻辑。
|
||
"""
|
||
import os
|
||
import time
|
||
|
||
import mlx.core as mx
|
||
|
||
from mlx_streaming import config
|
||
from mlx_streaming.core.moe import native_moe
|
||
from mlx_streaming.core import route_trace
|
||
from mlx_streaming.core.profiling import (
|
||
PROF, WINDOW_PROF, PREDICT_RECALL_PROF, MISS_ATTRIB, note_miss_attrib,
|
||
note_tprof, TPROF_ON, _PROF_ON, _tick, note_union, UNION_ON)
|
||
|
||
# decode/verify 热路径判据:seq 短(单 token decode=1、MTP verify=K≤几)，与 prefill 长 seq 区分。
|
||
_DECODE_SEQ_MAX = 8
|
||
from mlx_streaming.core.moe.gate import _effective_top_k
|
||
from mlx_streaming.core.cache import autopin
|
||
from mlx_streaming.core.moe.compute import (
|
||
streaming_switch_glu_forward, PersistentSubGLU)
|
||
|
||
|
||
class StreamingMoeBlock:
|
||
"""包住原 Qwen3MoeSparseMoeBlock：路由器常驻，专家计算改为只算选中专家。
|
||
|
||
若提供 store（LruExpertStore），可进一步把专家权重从磁盘按需加载；否则直接在
|
||
常驻的 switch_mlp 上做 uniq 切片计算（仍只激活少数专家）。
|
||
"""
|
||
|
||
def __init__(self, orig_block, layer_idx: int, store=None):
|
||
self.gate = orig_block.gate # 路由器常驻（很小）
|
||
self.top_k = orig_block.top_k
|
||
self.norm_topk_prob = orig_block.norm_topk_prob
|
||
self.switch_mlp = orig_block.switch_mlp
|
||
self.store = store
|
||
self.layer_idx = layer_idx
|
||
|
||
def __call__(self, x: mx.array) -> mx.array:
|
||
gates = mx.softmax(self.gate(x), axis=-1, precise=True)
|
||
k = _effective_top_k(self.top_k)
|
||
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
|
||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||
if self.norm_topk_prob:
|
||
scores = scores / mx.sum(scores, axis=-1, keepdims=True)
|
||
mx.eval(inds) # 先物化路由结果，才能按结果取专家
|
||
y = streaming_switch_glu_forward(self.switch_mlp, x, inds)
|
||
return (y * scores[..., None]).sum(axis=-2)
|
||
|
||
|
||
class FileStreamingMoeBlock:
|
||
"""文件后端流式 MoE 块：路由器常驻，专家权重从磁盘按需加载（不持有堆叠 switch_mlp）。"""
|
||
|
||
def __init__(self, gate, top_k, norm_topk_prob, store, layer_idx,
|
||
hidden, moe_inter, group_size, bits,
|
||
proj_bits: dict | None = None,
|
||
shared_expert=None, shared_expert_gate=None):
|
||
self.gate = gate
|
||
self.top_k = top_k
|
||
self.norm_topk_prob = norm_topk_prob
|
||
self.store = store
|
||
self.layer_idx = layer_idx
|
||
self.hidden = hidden
|
||
self.moe_inter = moe_inter
|
||
self.group_size = group_size
|
||
self.bits = bits
|
||
# Qwen3-Next 等带共享专家的模型：共享专家恒激活、必须常驻（不流式），
|
||
# 输出叠加 sigmoid(shared_expert_gate(x)) * shared_expert(x)。None 则退化为纯路由（如 Qwen3-MoE）。
|
||
self.shared_expert = shared_expert
|
||
self.shared_expert_gate = shared_expert_gate
|
||
# 可选全流式 blob 源（STREAM_BLOB=1 时由 model_builder 注入），见 core/blob_loader.py。
|
||
self._blob = None
|
||
# 持久化子模块：跨 token 复用，避免每调用重建 QSL。
|
||
# proj_bits 非空时走混合精度（逐 proj 不同 bit）。
|
||
self._sub = PersistentSubGLU(hidden, moe_inter, group_size, bits,
|
||
proj_bits=proj_bits, layer_idx=layer_idx)
|
||
|
||
def __call__(self, x: mx.array) -> mx.array:
|
||
if _PROF_ON:
|
||
return self._call_prof(x)
|
||
gates = mx.softmax(self.gate(x), axis=-1, precise=True)
|
||
k = _effective_top_k(self.top_k)
|
||
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
|
||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||
if self.norm_topk_prob:
|
||
scores = scores / mx.sum(scores, axis=-1, keepdims=True)
|
||
if config.predict_recall_prof():
|
||
ps = getattr(self, "_predicted_set", None)
|
||
if ps is not None:
|
||
act = {int(i) for i in inds.reshape(-1).tolist()}
|
||
PREDICT_RECALL_PROF["hit"] += len(ps & act)
|
||
PREDICT_RECALL_PROF["routed"] += len(act)
|
||
PREDICT_RECALL_PROF["n"] += 1
|
||
if config.native_moe():
|
||
native_y = self._try_native_forward(x, inds, scores, gates.shape[-1])
|
||
if native_y is not None:
|
||
y = native_y
|
||
if self.shared_expert is not None:
|
||
y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
|
||
return y
|
||
if self._blob is not None and config.stream_blob():
|
||
# 全流式 blob 路径：每层按需并行读专家 → 复用 _sub.forward(MLX quantized_matmul)。
|
||
flat = [int(i) for i in inds.reshape(-1).tolist()]
|
||
pool_arrays, slots = self._blob.acquire(self.layer_idx, flat)
|
||
local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape)
|
||
y = self._sub.forward(pool_arrays, len(set(flat)), x, local)
|
||
y = (y * scores[..., None]).sum(axis=-2)
|
||
if self.shared_expert is not None:
|
||
y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
|
||
return y
|
||
if config.probe_perlayer_sync():
|
||
# 诊断：每层强制一次 host 同步（模拟预测 barrier），不做任何预取。
|
||
# 若单这个就拖慢 ~30% → 税是 eval barrier，C++ 读统一内存也救不了。
|
||
_ = int(inds.reshape(-1)[:1].item())
|
||
if (config.stream_blob_bg()
|
||
and getattr(self.store, "_bg", None) is not None):
|
||
# 后台预取已物化的专家在此（主线程、acquire 前）写进常驻池槽 → 转 miss 为 hit。
|
||
if config.window_prof():
|
||
t0 = getattr(self, "_submit_t", None)
|
||
if t0 is not None:
|
||
WINDOW_PROF["sum_s"] += time.perf_counter() - t0
|
||
WINDOW_PROF["n"] += 1
|
||
self.store.promote_prefetched(self.layer_idx)
|
||
# native-fused-prefetch 的 promote（把回调预读好的专家写进池槽）下移到各 acquire 分支前：
|
||
# host/verify 路径已算出真实路由 uniq_set，可"只 promote 命中本层路由的专家"，零额外同步地
|
||
# 丢弃假阳性 → 省掉无用 scatter + 池污染（这是 promote -3.2% 开销的主因）。
|
||
_stg_mgr = getattr(self.store, "_staging", None)
|
||
# 双源：投机专家留池侧区由 demand_dual 取，不 promote、不驱逐 → 关 promote。
|
||
_do_promote = (_stg_mgr is not None and not config.native_no_promote()
|
||
and not config.zerocopy_dual_source())
|
||
if config.route_trace_enabled():
|
||
flat_trace = [int(i) for i in inds.reshape(-1).tolist()]
|
||
resident = (self.store.resident_experts(self.layer_idx)
|
||
if hasattr(self.store, "resident_experts") else set())
|
||
rank = (self.store.resident_lru_scores(self.layer_idx)
|
||
if hasattr(self.store, "resident_lru_scores") else {})
|
||
route_trace.record(
|
||
self.layer_idx, flat_trace, set(flat_trace) - resident, resident, rank)
|
||
# native-fused-prefetch：在 seq 分支**之前**提交，使 decode(seq=1) 与 MTP verify(seq=K)
|
||
# 两条路径都触发预取。dummy 折进 inds(加 0)：GPU 路径靠 acquire_gpu 的 n_miss eval、
|
||
# host 路径靠后面的 .tolist() eval，都会触发完成回调里的 pread。verify 时 x 为 K 个
|
||
# token（verify_in=[x, d_1..d_{K-1}]），预测的是"下一层这 K 个 token 的专家并集"（recall≈0.96 的口径）。
|
||
# 双源双缓冲：每块开头先推进/检测前向边界（必须在本块预取提交之前，保证同前向内 fill_gen/read_gen 恒定）。
|
||
if config.zerocopy_dual_source() and getattr(self, "_vpool", None) is not None:
|
||
self._vpool.begin_forward(self.layer_idx)
|
||
if (getattr(self.store, "_staging", None) is not None
|
||
and not config.native_no_submit()):
|
||
_dummy = self._native_fused_prefetch(x)
|
||
if _dummy is not None:
|
||
inds = inds + (_dummy.reshape(()).astype(inds.dtype) * 0)
|
||
layer_cap = self.store.cap_for(self.layer_idx)
|
||
# verify(小 seq)可走 GPU 重映射;prefill(大 seq)唯一专家可能超 cap,须留在 host 路径(有超容量 fetch 回退)。
|
||
if config.zerocopy_dual_source() and getattr(self, "_vpool", None) is not None:
|
||
# 双源两级缓存本就是为 MTP verify 建的:verify 走 dual 路径才能读到侧区,否则落 host
|
||
# 路径会白填侧区(预取填了却不读)→ 命中骤降、读盘翻倍。
|
||
#
|
||
# 但判据必须按 demand_dual 真实区**可安放容量**算,不能按 cap+侧区行:侧区行是预取只读的,
|
||
# miss 只能落真实区的 cap 槽,其中 pinned(AUTOPIN)永不可驱逐。设本次唯一专家数 U、
|
||
# 真实区 cap=C、pinned=P,可用槽 = C − P − (本次命中的非 pinned 常驻数),推得
|
||
# 「U ≤ C − P」是安放得下的充分条件。超了 C++ 分不到槽会把专家落 0 号槽 → 逐位错算
|
||
# (prefill chunk 大时曾整段 system+tools 都算错)。U 无同步上界即 inds.size = seq×top_k。
|
||
_verify_gpu = (x.shape[1] * k <= self._dual_placeable(layer_cap))
|
||
else:
|
||
_verify_gpu = (config.verify_gpu_remap() and x.shape[1] * k <= layer_cap)
|
||
if (config.resident_pool_enabled()
|
||
and config.gpu_remap_enabled()
|
||
and (x.shape[1] == 1 or _verify_gpu)):
|
||
# decode 热路径:GPU 侧 slot 重映射,命中层零 host 往返(消除每层 .tolist 同步),
|
||
# 仅一次 miss 标志同步;真 miss 才回退读盘。decode top-k(≤cap)恒装得下。
|
||
# 专家总数 = gate 输出维(gates 末维),无需依赖 gate.weight。
|
||
if UNION_ON:
|
||
# GPU 路径本不 .tolist 真实路由;仅 UNION_PROF=1 时付一次同步取并集大小。
|
||
note_union(x.shape[1], len({int(i) for i in inds.reshape(-1).tolist()}),
|
||
self.layer_idx)
|
||
# decode/verify GPU 重映射路径:host 无现成真实路由。开关开时用 GPU membership 仅对
|
||
# 预读候选现算 used(drain≤budget)过滤假阳性;关时退回整批写入(回退/对照基线)。
|
||
if _do_promote:
|
||
if config.gpu_remap_promote_filter():
|
||
_stg_mgr.promote(self.layer_idx, self.store,
|
||
route_inds=inds, num_experts=gates.shape[-1])
|
||
else:
|
||
_stg_mgr.promote(self.layer_idx, self.store)
|
||
if config.miss_attrib():
|
||
# 诊断专用:GPU remap 路径本不 .tolist 真实路由(那是它消栅栏的关键),
|
||
# 仅 MISS_ATTRIB=1 时付一次 .tolist 取真实路由,统计 decode 热路径 A/B 构成。
|
||
_uniq = {int(i) for i in inds.reshape(-1).tolist()}
|
||
_ps = getattr(self, "_predicted_set", None)
|
||
_res = (self.store.resident_experts(self.layer_idx)
|
||
if hasattr(self.store, "resident_experts") else set())
|
||
_rdy = _stg_mgr.last_ready.get(self.layer_idx) if _stg_mgr else None
|
||
note_miss_attrib(_uniq, _ps, _res, x.shape[1] <= _DECODE_SEQ_MAX, _rdy)
|
||
if config.zerocopy_dual_source():
|
||
# 双源双缓冲：统一走 VirtualPool.acquire（呈现「所有专家都在」视角，
|
||
# 内部真实区表 ∪ 侧区(读代) 单次 gather，返回 (pool, local, n_experts)）。
|
||
pool_arrays, local, n_experts = self._vpool.acquire(
|
||
self.layer_idx, inds, gates.shape[-1],
|
||
seq_len=x.shape[1], layer_cap=layer_cap)
|
||
y = self._sub.forward(pool_arrays, n_experts, x, local)
|
||
else:
|
||
pool_arrays, local = self.store.acquire_gpu(
|
||
self.layer_idx, inds, gates.shape[-1])
|
||
if config.stg_verify():
|
||
# 诊断：消费侧字节校验（默认 off），定位混合 ahead 下池槽损坏。
|
||
self.store._resident.verify_acquire_bytes(
|
||
self.layer_idx, inds, _stg_mgr)
|
||
y = self._sub.forward(pool_arrays, layer_cap, x, local)
|
||
else:
|
||
# prefill/大批量(seq>1)或显式关闭 GPU_REMAP:走 host 路径。
|
||
# 只在此处对 inds 做一次 .tolist() 同步，uniq/local 全在 Python 里算。
|
||
flat = [int(i) for i in inds.reshape(-1).tolist()]
|
||
autopin.note(self.layer_idx, flat) # AUTOPIN 热度:搭已物化的 flat,零额外同步(AUTOPIN=0 立即返回)
|
||
uniq_set = set(flat)
|
||
if UNION_ON:
|
||
note_union(x.shape[1], len(uniq_set), self.layer_idx) # 本层路由专家并集(零额外同步,uniq_set 已算)
|
||
# promote 只写"预读好 ∩ 本层真实路由"的专家：复用已算的 uniq_set，零额外同步，
|
||
# 假阳性不进池 → 省 scatter、不污染池（acquire 前完成，命中转化生效）。
|
||
if _do_promote:
|
||
_stg_mgr.promote(self.layer_idx, self.store, used=uniq_set)
|
||
if config.miss_attrib():
|
||
# 在 promote 之后、acquire(读盘)之前统计：本层真实路由的命中构成。
|
||
# miss 拆成 A(预测到却没进池：budget丢/时序/驱逐) 与 B(没预测到：召回缺口)。
|
||
_ps = getattr(self, "_predicted_set", None)
|
||
_res = (self.store.resident_experts(self.layer_idx)
|
||
if hasattr(self.store, "resident_experts") else set())
|
||
_rdy = _stg_mgr.last_ready.get(self.layer_idx) if _stg_mgr else None
|
||
note_miss_attrib(uniq_set, _ps, _res, x.shape[1] <= _DECODE_SEQ_MAX, _rdy)
|
||
# 池只能同时容纳 ≤该层容量 个唯一专家；prefill 唯一数超容量时回退 stack。
|
||
# dual 模式（_vpool 存在）统一走 VirtualPool.acquire_host 收口 host+fetch；
|
||
# 非 dual（_vpool 为 None）保留内联逻辑（非目标路径，避免额外依赖）。
|
||
if (config.resident_pool_enabled()
|
||
and getattr(self, "_vpool", None) is not None):
|
||
pool_arrays, local, n_experts = self._vpool.acquire_host(
|
||
self.layer_idx, flat, inds.shape, inds.dtype, layer_cap)
|
||
y = self._sub.forward(pool_arrays, n_experts, x, local)
|
||
elif (config.resident_pool_enabled()
|
||
and len(uniq_set) <= layer_cap):
|
||
pool_arrays, slots = self.store.acquire(self.layer_idx, flat)
|
||
local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape)
|
||
y = self._sub.forward(pool_arrays, layer_cap, x, local)
|
||
else:
|
||
uniq_sorted = sorted(uniq_set)
|
||
remap = {g: i for i, g in enumerate(uniq_sorted)}
|
||
local = mx.array([remap[i] for i in flat], dtype=inds.dtype).reshape(inds.shape)
|
||
fetched = self.store.fetch(self.layer_idx, uniq_sorted)
|
||
y = self._sub.forward(fetched, len(uniq_sorted), x, local)
|
||
y = (y * scores[..., None]).sum(axis=-2)
|
||
if self.shared_expert is not None: # 共享专家（常驻）叠加
|
||
y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
|
||
return y
|
||
|
||
def _dual_placeable(self, layer_cap: int) -> int:
|
||
"""demand_dual 真实区本层能安放的唯一专家上限 = cap − pinned。
|
||
|
||
pinned 集合只在启动 AUTOPIN 预热时写入、生成期恒定,故首次查到后缓存
|
||
(decode 每层每步都要判,不能每次都进 C++ 取锁)。
|
||
"""
|
||
n = getattr(self, "_pinned_n", None)
|
||
if n is None:
|
||
try:
|
||
import mlx_streaming.native_moe_ext as _N
|
||
n = int(_N.real_pinned_count(int(self.layer_idx)))
|
||
except Exception: # noqa: BLE001 取不到就按最保守的「全钉死」算,逼回 host 路径
|
||
n = int(layer_cap)
|
||
self._pinned_n = n
|
||
return max(1, int(layer_cap) - n)
|
||
|
||
def _native_fused_prefetch(self, x: mx.array):
|
||
"""搭车式预取：用下 AHEAD 层 gate 对 x 算预测 inds(lazy)，挂 GPU 完成回调，
|
||
在 C++ 里(GPU 算完后)读 id + pread 预热下层专家字节。返回 dummy 张量(需被 eval 才触发)。
|
||
全程不加主线程 host 同步——靠折进 inds、搭 acquire_gpu 的 n_miss eval。
|
||
"""
|
||
model = getattr(self, "_prefetch_model_ref", None)
|
||
bl = getattr(self.store, "_blob_loader", None)
|
||
if model is None or bl is None:
|
||
return None
|
||
layers = getattr(model, "layers", [])
|
||
vp = getattr(self, "_vpool", None)
|
||
if vp is not None:
|
||
tgt = vp.target_for(self.layer_idx) # per-layer cutoff ahead
|
||
if not (self.layer_idx < tgt < len(layers)): # 0/越界/无前瞻 → 跳过
|
||
return None
|
||
else:
|
||
ahead = max(1, config.cross_layer_ahead(default=1))
|
||
tgt = self.layer_idx + ahead
|
||
if not (0 <= tgt < len(layers)):
|
||
return None
|
||
tmlp = getattr(layers[tgt], "mlp", None)
|
||
if not isinstance(tmlp, FileStreamingMoeBlock):
|
||
return None
|
||
try:
|
||
from mlx_streaming import native_moe_ext as _N
|
||
except Exception:
|
||
return None
|
||
# 方案B：预测"宽集合"(top-N, N=predict_width)只为高 recall，不占内存（仅一次 gate argpartition）；
|
||
# 真正占 staging 的是 C++ 回调过滤常驻后、按分截到 staging budget 的"缺口"子集。
|
||
_tp0 = time.perf_counter() if TPROF_ON else 0.0
|
||
predict_width = config.cross_layer_predict_width()
|
||
# 默认(PREDICT_USE_X=1)：用本层 MoE 输入 x 喂目标层 gate。x 含了本层 attention，离 L+1
|
||
# 近半层 → 比未归一化输入更新鲜，实测 recall 0.812→0.847(+3.6pp)、还省一次 norm。
|
||
# PREDICT_USE_X=0 回退旧路径：目标层 post_attention_layernorm 作用在本层未归一化输入上。
|
||
h = getattr(self, "_unnormed_input", None)
|
||
if config.predict_use_x():
|
||
g = tmlp.gate(x)
|
||
elif h is None:
|
||
h = x # 退化：无未归一化输入时用 x（norm 偏差，覆盖会差）
|
||
g = tmlp.gate(h)
|
||
else:
|
||
g = tmlp.gate(layers[tgt].post_attention_layernorm(h))
|
||
# 按 seq 维聚合成"该批 K 个 token 的专家并集"近似。聚合方式 PREDICT_AGG：
|
||
# - max（默认）：任一 token 强烈想要即入选；mean：各 token 平均偏好；
|
||
# - union：每 token 各取 top-kk 再并集（与真实路由"并集"结构最一致，候选更多，由 handler 去重+截断）。
|
||
_agg = config.predict_agg()
|
||
kk = min(g.shape[-1], predict_width)
|
||
if g.ndim == 3 and _agg == "union":
|
||
# union 候选数 = seq × union_k。每 token 只取 top-union_k（独立于 predict_width），
|
||
# 调小 union_k 缩候选 → 减少缺口、缓解 staging/读盘洪水（每 token 真实只路由 top_k 个）。
|
||
uk = min(g.shape[-1], config.predict_union_k())
|
||
pred = mx.argpartition(g, kth=-uk, axis=-1)[..., -uk:].reshape(-1).astype(mx.uint32)
|
||
else:
|
||
if g.ndim == 3:
|
||
g = g.mean(axis=1) if _agg == "mean" else g.max(axis=1)
|
||
# 取 top-kk（argpartition，O(E) 不全排序）。实测：加宽 kk 不提升命中——staging 更多 = 后台
|
||
# pread 更多、到位更晚 + 抢带宽 → 反而更低；故 width=16 即最优，cap 截断极少触发、无需排序。
|
||
pred = mx.argpartition(g, kth=-kk, axis=-1)[..., -kk:].reshape(-1).astype(mx.uint32)
|
||
if config.predict_recall_prof() or config.miss_attrib():
|
||
try:
|
||
tmlp._predicted_set = {int(e) for e in pred.tolist()}
|
||
except Exception:
|
||
pass
|
||
if TPROF_ON:
|
||
note_tprof("predict_s", time.perf_counter() - _tp0, count_key="predict_n")
|
||
stg = getattr(self.store, "_staging", None)
|
||
if stg is not None:
|
||
resident = (self.store.resident_experts(tgt)
|
||
if hasattr(self.store, "resident_experts") else None)
|
||
if config.zerocopy_dual_source():
|
||
# 双源双缓冲：经 VirtualPool 向填代 submit,C++ 回调把预读段散写进该代侧区行。
|
||
rp = self.store._resident
|
||
if tgt not in rp._pools:
|
||
return None # 目标层池未建(首 token 预热) → 跳过本次
|
||
segs = stg.src._segs # (proj, tensor, dt, shape, nb)，与池 key 顺序一致
|
||
pool_list = [rp._pools[tgt][f"{p}.{t}"] for p, t, *_ in segs]
|
||
import os as _os
|
||
if _os.environ.get("POOL_PTR_TRACE") is not None and tgt == int(_os.environ["POOL_PTR_TRACE"]):
|
||
try:
|
||
import mlx_streaming.native_moe_ext as _N
|
||
_fk = f"{segs[0][0]}.{segs[0][1]}"
|
||
print(f"[POOL_PTR] submit layer={tgt} key={_fk} obj_id={id(rp._pools[tgt][_fk])} "
|
||
f"ptr={hex(_N.array_data_ptr(rp._pools[tgt][_fk]))}", flush=True)
|
||
except Exception as _e:
|
||
print(f"[POOL_PTR] err {_e}", flush=True)
|
||
return self._vpool.prefetch(tgt, pred, resident, pool_list)
|
||
# miss→hit:回调按目标层常驻快照过滤，只把缺口 pread 进 staging（≤budget 行），promote 时写池。
|
||
if TPROF_ON:
|
||
_ts0 = time.perf_counter()
|
||
_r = stg.submit(tgt, pred, resident)
|
||
note_tprof("submit_s", time.perf_counter() - _ts0, count_key="submit_n")
|
||
return _r
|
||
return stg.submit(tgt, pred, resident)
|
||
# 仅预热字节（page cache）的轻量版。
|
||
path = os.path.join(bl.dir, f"layer{tgt:02d}.blob")
|
||
return _N.prefetch_on_complete(pred, path, int(bl.stride), True)
|
||
|
||
def _try_native_forward(self, x: mx.array, inds: mx.array, scores: mx.array, num_experts: int):
|
||
# native 第一版只支持三投影同 bit 的专家；其他情况保持现有 MLX 回退。
|
||
proj_bits = getattr(self._sub, "proj_bits", {})
|
||
bits_set = {int(proj_bits.get(name, self.bits)) for name in ("gate_proj", "up_proj", "down_proj")}
|
||
if len(bits_set) != 1:
|
||
return None
|
||
bits = bits_set.pop()
|
||
if not native_moe.can_native_moe(
|
||
self.layer_idx, self.hidden, self.moe_inter,
|
||
self.group_size, bits, num_experts):
|
||
return None
|
||
t_sync = time.perf_counter()
|
||
mx.eval(inds)
|
||
flat = [int(i) for i in inds.reshape(-1).tolist()]
|
||
native_moe.note_route_sync(time.perf_counter() - t_sync)
|
||
return native_moe.try_native_moe(
|
||
self.layer_idx, flat, x, scores,
|
||
self.hidden, self.moe_inter, self.group_size, bits, num_experts)
|
||
|
||
def _call_prof(self, x: mx.array) -> mx.array:
|
||
t = time.perf_counter()
|
||
gates = mx.softmax(self.gate(x), axis=-1, precise=True)
|
||
k = _effective_top_k(self.top_k)
|
||
inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
|
||
scores = mx.take_along_axis(gates, inds, axis=-1)
|
||
if self.norm_topk_prob:
|
||
scores = scores / mx.sum(scores, axis=-1, keepdims=True)
|
||
mx.eval(inds, scores)
|
||
_tick("route", t); t = time.perf_counter()
|
||
|
||
flat = [int(i) for i in inds.reshape(-1).tolist()]
|
||
uniq_set = set(flat)
|
||
layer_cap = self.store.cap_for(self.layer_idx)
|
||
use_pool = (config.resident_pool_enabled()
|
||
and len(uniq_set) <= layer_cap)
|
||
if not use_pool:
|
||
uniq_sorted = sorted(uniq_set)
|
||
remap = {g: i for i, g in enumerate(uniq_sorted)}
|
||
local = mx.array([remap[i] for i in flat], dtype=inds.dtype).reshape(inds.shape)
|
||
mx.eval(local)
|
||
_tick("pyremap", t); t = time.perf_counter()
|
||
|
||
if use_pool:
|
||
pool_arrays, slots = self.store.acquire(self.layer_idx, flat)
|
||
local = mx.array(slots, dtype=inds.dtype).reshape(inds.shape)
|
||
mx.eval(local)
|
||
fetched, n_experts = pool_arrays, layer_cap
|
||
else:
|
||
fetched = self.store.fetch(self.layer_idx, uniq_sorted)
|
||
n_experts = len(uniq_sorted)
|
||
mx.eval(list(fetched.values()))
|
||
_tick("fetch", t); t = time.perf_counter()
|
||
|
||
y = self._sub.forward(fetched, n_experts, x, local)
|
||
mx.eval(y)
|
||
_tick("matmul", t); t = time.perf_counter()
|
||
|
||
y = (y * scores[..., None]).sum(axis=-2)
|
||
if self.shared_expert is not None:
|
||
y = y + mx.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
|
||
mx.eval(y)
|
||
_tick("combine", t)
|
||
PROF["n_calls"] += 1
|
||
return y
|