238 lines
14 KiB
Python
238 lines
14 KiB
Python
"""VirtualPool：预取调度器 + 双源双缓冲协调器（统一收口）。
|
||
|
||
一个对象承担两职，因为 block.py 用同一个 `self._vpool` 属性：
|
||
|
||
1. ahead 调度（两种模式都用）：`ahead_for` / `target_for` 决定「第 L 层应预读哪层」——
|
||
早层用小 ahead 保召回、cutoff 起用大 ahead 抢时序。`_native_fused_prefetch` 靠它选目标层。
|
||
|
||
2. 双源双缓冲协调（仅 ZEROCOPY_DUAL_SOURCE 模式）：`begin_forward` / `read_gen` /
|
||
`fill_gen` / `acquire` / `prefetch`。侧区分两「代(gen)」物理行，每前向读上一代、fill 写
|
||
另一代，根除「本前向消费 gather 读的物理行在本前向 eval 期被 fill 覆盖」的竞态
|
||
（见 spec 2026-06-25-qwen-virtual-pool-double-buffer）。对外只暴露「专家→物理行」单次
|
||
gather 接口，消费者每层零 host 同步。
|
||
|
||
构造两种签名并存（互斥使用）：
|
||
- 调度器：VirtualPool(num_layers=.., cutoff=.., ahead_lo=.., ahead_hi=..)
|
||
- 协调器：VirtualPool(resident, staging, spec_slots)，dual-source 下再补调度参数即可两职合一。
|
||
"""
|
||
import mlx.core as mx
|
||
|
||
from mlx_streaming.core.cache import autopin
|
||
|
||
# 方案B STG_VERIFY 校验累计态（诊断用，默认路径不触及）。
|
||
_stg_verify_state = {"ok": 0, "bad": 0, "printed": 0, "calls": 0}
|
||
|
||
|
||
class VirtualPool:
|
||
def __init__(self, resident=None, staging=None, spec_slots=None, *,
|
||
num_layers=None, cutoff=None, ahead_lo=None, ahead_hi=None, store=None):
|
||
# --- 双源协调（resident/staging 存在时启用）---
|
||
self._rp = resident
|
||
self._stg = staging
|
||
self._store = store # acquire_host 超容量 fetch 回退用(blob 感知);调度器形态可为 None
|
||
self._spec = int(spec_slots) if spec_slots is not None else 0
|
||
self._gen = 0
|
||
self._last_layer = -1 # 前向边界检测：层号回绕(<= 上次) 即新前向
|
||
# --- ahead 调度 ---
|
||
self._num_layers = int(num_layers) if num_layers is not None else 0
|
||
self._cutoff = int(cutoff) if cutoff is not None else 0
|
||
self._a_lo = max(1, int(ahead_lo)) if ahead_lo is not None else 1
|
||
self._a_hi = max(1, int(ahead_hi)) if ahead_hi is not None else 1
|
||
|
||
# ---- ahead 调度 ----
|
||
def ahead_for(self, src_layer: int) -> int:
|
||
# cutoff：早层用 lo（保召回），cutoff 起用 hi（保时序）。
|
||
return self._a_lo if int(src_layer) < self._cutoff else self._a_hi
|
||
|
||
def target_for(self, src_layer: int) -> int:
|
||
# 目标层 = src+ahead。末层无可预读 → 返回 0（跳过）。
|
||
# 不再 clamp 到末层：clamp 会让多个源层同时预读同一末层，对该层每前向 submit 次数 >
|
||
# staging ring，其环形 buffer 在惰性切片被本前向 eval 消费前就被后续 submit 完成回调覆盖成
|
||
# 别的专家字节 → 池槽装错字节。越界即跳过（交给 demand 读真值），保证每层每前向至多 1 次 submit。
|
||
L = int(src_layer)
|
||
if L >= self._num_layers - 1:
|
||
return 0
|
||
tgt = L + self.ahead_for(L)
|
||
if tgt > self._num_layers - 1: # 越界（本会触发 clamp 堆叠）→ 跳过
|
||
return 0
|
||
return tgt
|
||
|
||
# ---- 双源双缓冲协调 ----
|
||
def begin_forward(self, layer_idx: int):
|
||
"""每个 MoE 块 __call__ 开头调；层号回绕(<= 上次) 判为新前向 → 代 +1。
|
||
稳健：不依赖首个 MoE 层是 layer 0、也不要求 MoE 层连续。"""
|
||
if layer_idx <= self._last_layer:
|
||
self._gen += 1
|
||
# AUTOPIN:dual decode 无 Python 计频点,用前向边界驱动周期落盘(AUTOPIN=0 立即返回)。
|
||
autopin.tick()
|
||
# 新前向开头排空上一前向提交的侧区 fill：C++-owned 池 buffer 由后台异步直写，
|
||
# 若消费前 fill 未写完，GPU gather 会读到半写侧区行（DUAL_VERIFY BAD）。
|
||
# drain 阻塞到在途 fill 全部写完 → 本前向要消费的侧区行字节必已就绪。
|
||
if self._spec > 0:
|
||
import mlx_streaming.native_moe_ext as _N
|
||
_N.sideregion_drain()
|
||
self._last_layer = layer_idx
|
||
|
||
def _gens(self) -> int:
|
||
# 代数取自常驻池;单代(=1)时读=填=0(持久 LFU 单区),双代(=2)时交替(%2==&1)。
|
||
g = getattr(self._rp, "spec_gens", 2) if self._rp is not None else 2
|
||
return max(1, int(g))
|
||
|
||
def read_gen(self) -> int:
|
||
return (self._gen - 1) % self._gens() # 读上一前向填好的代;单代恒 0
|
||
|
||
def fill_gen(self) -> int:
|
||
return self._gen % self._gens() # fill 写本代;单代恒 0
|
||
|
||
def acquire(self, layer, inds, num_experts, *, seq_len=None, layer_cap=None):
|
||
"""统一取用入口（GPU-remap 路径）：对外呈现「所有专家都在」的视角。
|
||
|
||
返回 (pool_arrays, local, n_experts)，计算侧零分支：
|
||
- dual（有侧区 staging 且 spec>0）：C++ demand_dual 唯一权威（真实区 ∪ 侧区读代单次 gather）；
|
||
n_experts = layer_cap + spec_gens*spec_slots。
|
||
- 非 dual GPU-remap：acquire_gpu；n_experts = layer_cap。
|
||
（host/fetch 路径见 acquire_host，其输入是 host 侧已 .tolist 的 flat，语义不同故分开。）
|
||
"""
|
||
cap = int(layer_cap) if layer_cap is not None else self._rp.cap_for(layer)
|
||
if self._stg is not None and self._spec > 0:
|
||
side_gen = self.read_gen()
|
||
# C++ demand_dual 是双源 decode 真实区的唯一权威（每层 1 次 inds 同步、零主线程落池/记账）。
|
||
# 无 Python 退路：native 未编译 demand_dual → 明确报错（decode 依赖 native 全套能力）。
|
||
if not getattr(self._rp, "_native_demand", False):
|
||
raise RuntimeError(
|
||
"双源 decode 需要 native demand_dual（已成唯一权威），但未检出。请编译 native_moe_ext。")
|
||
pool, local = self._acquire_native(layer, inds, side_gen, cap)
|
||
n_exp = cap + self._rp.spec_gens * self._rp.spec_slots
|
||
return pool, local, n_exp
|
||
pool, local = self._rp.acquire_gpu(layer, inds, num_experts)
|
||
return pool, local, cap
|
||
|
||
def _native_meta(self, layer):
|
||
"""缓存每层 demand_dual 的不变入参（pool_list/seg_nbytes/path），避免逐层重建 host 胶水。"""
|
||
cache = getattr(self, "_nd_meta", None)
|
||
if cache is None:
|
||
cache = self._nd_meta = {}
|
||
m = cache.get(layer)
|
||
if m is None:
|
||
stg = self._stg
|
||
segs = stg.src._segs # (proj, tensor, dt, shape, nb)，与池 key 同序
|
||
pool_list = [self._rp._pools[layer][f"{p}.{t}"] for p, t, *_ in segs]
|
||
m = (pool_list, [int(nb) for *_, nb in segs],
|
||
f"{stg.src.dir}/layer{int(layer):02d}.blob", int(stg.stride))
|
||
cache[layer] = m
|
||
return m
|
||
|
||
def _acquire_native(self, layer, inds, side_gen, cap):
|
||
"""方案B 取用：委派 C++ demand_dual（每层 1 次 inds 同步 + 并行 worker pread 落池），
|
||
更新 rp 统计计数（供报告口径一致）。"""
|
||
import mlx_streaming.native_moe_ext as N
|
||
rp = self._rp
|
||
rp._bootstrap_dual_pool(layer) # 首次建池 + real_init（幂等）
|
||
# 方案B 容量前提：miss 只能落真实区的 cap 槽，其中 pinned 永不可驱逐 → 可安放上限 cap−pinned。
|
||
# 超了 C++ 会把专家落 0 号槽（拿别的专家权重算）→ 逐位不正确。调用方（block.py 的分流判据）
|
||
# 负责在超容量前把该前向送去 host/fetch；这里只做一次性告警兜底。
|
||
_placeable = int(cap) - int(N.real_pinned_count(int(layer)))
|
||
if int(inds.size) > _placeable and not getattr(self, "_overcap_warned", False):
|
||
self._overcap_warned = True
|
||
import sys
|
||
print(f"[DEMAND_DUAL] 警告：inds.size={int(inds.size)} > 可安放槽 {_placeable}"
|
||
f"(cap={int(cap)} − pinned={int(cap) - _placeable})，真实区超容量，逐位将不正确。"
|
||
f"请调小 PREFILL_CHUNK，或调大 EXPERT_SLOTS / 调小 AUTOPIN_BUDGET_FRAC。",
|
||
file=sys.stderr, flush=True)
|
||
pool_list, seg_nbytes, path, stride = self._native_meta(layer)
|
||
local = N.demand_dual(inds, pool_list, seg_nbytes, int(layer), int(side_gen), path,
|
||
stride, int(cap),
|
||
rp.eviction_policy == "lfu", int(rp.lfu_decay_interval))
|
||
st = N.demand_last_stats() # [hitpos, misspos, loads, fallback01, unplaced]
|
||
rp.hits += st[0]
|
||
rp.misses += st[2]
|
||
if len(st) > 4:
|
||
rp.unplaced += st[4] # >0 即有前向落了 0 号槽，输出不可信（/api/stats 暴露）
|
||
if st[3] == 0:
|
||
rp.gpu_fastpath += 1
|
||
else:
|
||
rp.gpu_fallback += 1
|
||
from mlx_streaming import config
|
||
if config.stg_verify(): # 诊断:方案B 池字节逐 key 真值校验(默认关)
|
||
self._verify_native_bytes(layer, inds, local)
|
||
return rp._pools[layer], local
|
||
|
||
def _verify_native_bytes(self, layer, inds, local):
|
||
"""诊断(STG_VERIFY，方案B)：校验「字节落池不变量」——真实区每个占用槽的池字节 == 该槽当前
|
||
C++ 属主专家(g_real)的 blob 真值。这是 C++ 接管落池的字节等价铁证；发现不一致即池装错字节。
|
||
|
||
注：不以 local→expert 为判据（local 可能因跨调用/多模型共享 g_real 而滞后于 g_real，属路由级
|
||
问题、非落池字节问题；逐位权威信号是 e2e n_mismatch）。
|
||
"""
|
||
st = _stg_verify_state
|
||
st["calls"] += 1
|
||
pool = self._rp._pools.get(layer)
|
||
if pool is None:
|
||
return
|
||
import mlx_streaming.native_moe_ext as N
|
||
stg = self._stg
|
||
path = f"{stg.src.dir}/layer{int(layer):02d}.blob"
|
||
segs = stg.src._segs
|
||
flat = N.real_region_contents(int(layer)) # [expert0,slot0,expert1,slot1,...]
|
||
for j in range(0, len(flat), 2):
|
||
e, slot = flat[j], flat[j + 1]
|
||
raw = N.blob_load(path, mx.array([e], dtype=mx.uint32), int(stg.stride))[0]
|
||
bad, off = None, 0
|
||
for p, t, dt, shape, nb in segs:
|
||
k = f"{p}.{t}"
|
||
pv = pool[k][slot].reshape(-1).view(mx.uint8)
|
||
if not bool(mx.all(pv == raw[off:off + nb])):
|
||
bad = k
|
||
break
|
||
off += nb
|
||
if bad is None:
|
||
st["ok"] += 1
|
||
else:
|
||
st["bad"] += 1
|
||
if st["printed"] < 12:
|
||
st["printed"] += 1
|
||
print(f"[STG_VERIFY-DUAL] BAD 落池字节错 call={st['calls']} layer={layer} "
|
||
f"expert={e} slot={slot} key={bad} (ok={st['ok']} bad={st['bad']})", flush=True)
|
||
|
||
def acquire_host(self, layer, flat, inds_shape, inds_dtype, layer_cap):
|
||
"""host/fetch 路径收口（prefill/大 seq 或关 GPU-remap）：flat 为 host 侧路由 id 列表。
|
||
|
||
返回 (pool_arrays, local, n_experts)，与 block.py 原 host/fetch 分支逐元素等价：
|
||
- uniq <= cap：acquire(flat)，local 为槽位；n_experts = layer_cap。
|
||
- uniq > cap：fetch(uniq_sorted)，local 为 remap 到 [0,uniq) 连续索引；n_experts = uniq 数。
|
||
|
||
native demand(dual)下真实区槽状态归 C++ g_real，Python 的 _slot_of/_free 从不填充，
|
||
走 self._rp.acquire 必然在 _choose_victim 抛 "no evictable non-current slot"。
|
||
故 dual 下一律走 fetch 旁路（不进池、正确但慢）——这条路径本就是超容量前向的兜底。
|
||
"""
|
||
import mlx.core as mx
|
||
cap = int(layer_cap)
|
||
uniq_set = set(flat)
|
||
# 入池的充分条件是 uniq ≤ cap − pinned:acquire 把本次全部唯一专家列为不可驱逐(current),
|
||
# 腾槽时只能挑「非 pinned 且非 current」的受害者,少一个都会在 _choose_victim 抛
|
||
# "no evictable non-current slot"。pinned 是那层实际钉死数,不是拍出来的固定余量——
|
||
# 写死余量会在 cap 小于余量时把所有请求推进 fetch 慢路径(每次重读盘)。
|
||
_pinned_n = len(getattr(self._rp, "_pinned", {}).get(layer, ()) or ())
|
||
if (not getattr(self._rp, "_native_demand", False)
|
||
and len(uniq_set) <= max(1, cap - _pinned_n)):
|
||
pool, slots = self._rp.acquire(layer, flat)
|
||
local = mx.array(slots, dtype=inds_dtype).reshape(inds_shape)
|
||
return pool, local, cap
|
||
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)
|
||
# 超容量(或 dual:真实区槽状态归 C++ g_real,Python _slot_of 从不填充 → acquire 必抛)
|
||
# 走 fetch 旁路:不进池、按 uniq 现堆一份参与计算,正确但慢。这条路径本就是兜底。
|
||
# fetch 由 blob 感知的 FileExpertStore 提供;ResidentExpertPool 本身没有,故优先用注入的 store。
|
||
src = self._store if self._store is not None else self._rp
|
||
if not hasattr(src, "fetch"):
|
||
raise RuntimeError("VirtualPool.acquire_host 超容量回退需要带 fetch 的 store(构造时传入)")
|
||
fetched = src.fetch(layer, uniq_sorted)
|
||
return fetched, local, len(uniq_sorted)
|
||
|
||
def prefetch(self, layer, pred, resident, pool_list):
|
||
"""向 fill 代 submit 预读：base_row = cap_for(layer) + fill_gen*spec。"""
|
||
g = self.fill_gen()
|
||
base = self._rp.cap_for(layer) + g * self._spec
|
||
return self._stg.submit_pool_sideregion(layer, pred, resident, pool_list, base, gen=g)
|