239 lines
11 KiB
Python
239 lines
11 KiB
Python
"""全流式 blob 专家源：并行 pread + F_NOCACHE + 即时物化为 MLX 量化数组。
|
||
|
||
字节布局由 prep/blob_layout.py 统一描述（单一真相源），支持两种格式：
|
||
- v1 affine（expert_blob_v1）：每 proj 段序 [weight, scales, biases]；weight uint32，
|
||
scales/biases 存 uint16 原始位、用时 .view(bfloat16)。
|
||
- v2 mxfp4（expert_blob_v2_mxfp4）：每 proj 段序 [weight, scales]（无 biases）；
|
||
weight uint32，scales 为 uint8 原始位、绝不 view(bf16)。
|
||
计算复用 MLX quantized_matmul / gather_qmm（不碰已 NO-GO 的 fused kernel）。
|
||
"""
|
||
import os
|
||
import threading
|
||
from collections import OrderedDict
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
|
||
import mlx.core as mx
|
||
import numpy as np
|
||
|
||
from mlx_streaming import config
|
||
|
||
# macOS fcntl：F_NOCACHE 提示内核不要把读过的页留在 page cache（支撑低内存目标）。
|
||
_F_NOCACHE = 48
|
||
|
||
|
||
class BlobExpertSource:
|
||
def __init__(self, blob_dir: str, hidden: int, inter: int, group: int, bits: int,
|
||
num_experts: int, workers: int = 8, nocache: bool = True,
|
||
quant_mode: str = "affine", blob_format: "str | None" = None):
|
||
self.dir = blob_dir
|
||
self.h, self.i, self.g, self.b, self.ne = hidden, inter, group, bits, num_experts
|
||
self.workers = workers
|
||
self.nocache = nocache
|
||
# blob 格式：mxfp4 默认 v2（无 biases、scales uint8），其余默认 v1 affine。
|
||
self.quant_mode = quant_mode
|
||
self.blob_format = blob_format or (
|
||
"expert_blob_v2_mxfp4" if quant_mode == "mxfp4" else "expert_blob_v1")
|
||
self._segs, self.stride = self._layout()
|
||
self._fds: "dict[int, int]" = {}
|
||
# 后台预取：只预读原始字节进 _pf_cache（不建 mx.array，可后台线程跑）；
|
||
# 滚动窗口只保留最近 window 层的预取字节，支撑低内存。
|
||
self.window = config.stream_blob_window()
|
||
self._pf_pool = ThreadPoolExecutor(max_workers=workers)
|
||
self._pf_cache: "OrderedDict[tuple, bytes]" = OrderedDict()
|
||
self._pf_futures: "dict[tuple, object]" = {}
|
||
self._pf_layers: "list[int]" = []
|
||
self._lock = threading.Lock()
|
||
self.preads = 0
|
||
self.prefetch_hits = 0
|
||
|
||
def _layout(self):
|
||
# 段表统一由 blob_layout.layout_for 给出（dtype 为字符串 "uint32"/"uint16"/"uint8"）。
|
||
from mlx_streaming.prep.blob_layout import layout_for
|
||
return layout_for(self.blob_format, self.h, self.i, self.b, self.g)
|
||
|
||
def _fd(self, layer: int) -> int:
|
||
fd = self._fds.get(layer)
|
||
if fd is None:
|
||
fd = os.open(os.path.join(self.dir, f"layer{layer:02d}.blob"), os.O_RDONLY)
|
||
if self.nocache:
|
||
try:
|
||
import fcntl
|
||
fcntl.fcntl(fd, _F_NOCACHE, 1)
|
||
except (OSError, ValueError):
|
||
pass
|
||
self._fds[layer] = fd
|
||
return fd
|
||
|
||
def _materialize(self, raw: bytes, view_bf16: bool = True) -> dict:
|
||
# view_bf16=False：scales/biases 保留 uint16（.view(bfloat16) 在后台线程会报
|
||
# no Stream(gpu,0)，故后台物化时关掉，由主线程消费时再 view）。
|
||
# 段表 dtype 为字符串：仅 v1 affine 的 uint16 scales/biases 是 bf16 位重解释；
|
||
# mxfp4 的 uint8 scales 保持原样，绝不 view(bf16)。
|
||
out = {}
|
||
off = 0
|
||
for proj, tensor, dt, shape, nb in self._segs:
|
||
np_dt = np.dtype(dt)
|
||
view = np.frombuffer(raw, dtype=np_dt, count=nb // np_dt.itemsize, offset=off).reshape(shape)
|
||
arr = mx.array(view)
|
||
if view_bf16 and dt == "uint16" and tensor in ("scales", "biases"):
|
||
arr = arr.view(mx.bfloat16)
|
||
out[f"{proj}.{tensor}"] = arr
|
||
off += nb
|
||
return out
|
||
|
||
def _pread(self, layer: int, e: int) -> bytes:
|
||
return os.pread(self._fd(layer), self.stride, e * self.stride)
|
||
|
||
def prefetch_async(self, layer: int, expert_ids) -> None:
|
||
"""后台并行预读预测专家的字节（仅 IO，不建 mx.array）。供跨层预测调用。"""
|
||
for e in (int(x) for x in expert_ids):
|
||
key = (layer, e)
|
||
with self._lock:
|
||
if key in self._pf_cache or key in self._pf_futures:
|
||
continue
|
||
self._pf_futures[key] = self._pf_pool.submit(self._pf_job, layer, e)
|
||
self._evict_window(layer)
|
||
|
||
def _pf_job(self, layer: int, e: int) -> bytes:
|
||
raw = self._pread(layer, e)
|
||
with self._lock:
|
||
self._pf_cache[(layer, e)] = raw
|
||
self._pf_futures.pop((layer, e), None)
|
||
return raw
|
||
|
||
def _evict_window(self, layer: int) -> None:
|
||
with self._lock:
|
||
if layer not in self._pf_layers:
|
||
self._pf_layers.append(layer)
|
||
while len(self._pf_layers) > self.window:
|
||
old = self._pf_layers.pop(0)
|
||
for key in [k for k in self._pf_cache if k[0] == old]:
|
||
del self._pf_cache[key]
|
||
|
||
def wait_prefetch(self) -> None:
|
||
with self._lock:
|
||
futs = list(self._pf_futures.values())
|
||
for f in futs:
|
||
f.result()
|
||
|
||
def read_raw(self, layer: int, expert_ids) -> "list[bytes]":
|
||
"""取原始字节：优先用预取缓存/在途结果，未命中的并行 pread（仅 IO）。"""
|
||
ids = [int(e) for e in expert_ids]
|
||
results: "dict[int, bytes]" = {}
|
||
misses = []
|
||
for e in ids:
|
||
key = (layer, e)
|
||
with self._lock:
|
||
if key in self._pf_cache:
|
||
results[e] = self._pf_cache[key]
|
||
self.prefetch_hits += 1
|
||
continue
|
||
fut = self._pf_futures.get(key)
|
||
if fut is not None:
|
||
results[e] = fut.result()
|
||
self.prefetch_hits += 1
|
||
else:
|
||
misses.append(e)
|
||
|
||
def rd(e):
|
||
return e, self._pread(layer, e)
|
||
|
||
if misses:
|
||
self.preads += len(misses)
|
||
if self.workers <= 1 or len(misses) <= 1:
|
||
for e in misses:
|
||
results[e] = self._pread(layer, e)
|
||
else:
|
||
# 复用持久线程池（之前每次调用新建/销毁 ThreadPoolExecutor → 大量线程创建 + GIL 争用）。
|
||
for e, raw in self._pf_pool.map(rd, misses):
|
||
results[e] = raw
|
||
return [results[e] for e in ids]
|
||
|
||
def load_experts(self, layer: int, expert_ids, view_bf16: bool = True) -> dict:
|
||
"""返回 {expert_id: {proj.tensor: mx.array}}。并行读字节，lazy 物化 mx.array。
|
||
|
||
view_bf16=False：scales/biases 留 uint16（供后台线程用，避免 .view 报 no Stream）。
|
||
注：零拷贝(preadv 进 MLX buffer)实测在真实 lazy-eval 流里更慢——为拿可写 buffer
|
||
需 per-expert mx.eval 同步，代价高于省掉的拷贝(见报告)。故保留 lazy frombuffer 路径。
|
||
"""
|
||
ids = [int(e) for e in expert_ids]
|
||
raws = self.read_raw(layer, ids)
|
||
return {e: self._materialize(raw, view_bf16=view_bf16) for e, raw in zip(ids, raws)}
|
||
|
||
def load_experts_stacked(self, layer: int, expert_ids, view_bf16: bool = True) -> dict:
|
||
"""批量物化：每段只建一个 (N,*shape) mx.array（np.stack 后单次构造），
|
||
取代逐专家 6N 次 frombuffer+mx.array。返回 {proj.tensor: (N,*shape)}，按 ids 顺序。
|
||
|
||
与 load_experts 同源（同段表、同 view_bf16 语义），只是把"逐专家逐段"折成
|
||
"逐段一次堆叠"——demand 多 miss 消费的碎 mx.array 构造从 6N 降到 6。
|
||
"""
|
||
ids = [int(e) for e in expert_ids]
|
||
raws = self.read_raw(layer, ids)
|
||
out = {}
|
||
off = 0
|
||
for proj, tensor, dt, shape, nb in self._segs:
|
||
np_dt = np.dtype(dt)
|
||
cnt = nb // np_dt.itemsize
|
||
views = [np.frombuffer(raw, dtype=np_dt, count=cnt, offset=off).reshape(shape)
|
||
for raw in raws]
|
||
arr = mx.array(np.stack(views, axis=0)) # (N, *shape) 单次构造
|
||
if view_bf16 and dt == "uint16" and tensor in ("scales", "biases"):
|
||
arr = arr.view(mx.bfloat16)
|
||
out[f"{proj}.{tensor}"] = arr
|
||
off += nb
|
||
return out
|
||
|
||
def load_experts_native(self, layer: int, expert_ids, view_bf16: bool = False) -> dict:
|
||
"""native 物化：用 C++ blob_load 把字节 pread 进 MLX 数组（eval 时在 C++ 跑、绕 GIL），
|
||
Python 侧只建惰性切片/view 图（无数据拷贝）。目的是把"拷贝"挪出 GIL，减少后台线程
|
||
对主线程 per-layer 派发的争用。
|
||
"""
|
||
import mlx_streaming.native_moe_ext as _N
|
||
ids = [int(e) for e in expert_ids]
|
||
path = os.path.join(self.dir, f"layer{layer:02d}.blob")
|
||
raw = _N.blob_load(path, mx.array(ids, dtype=mx.uint32), self.stride) # [n, stride] uint8（lazy）
|
||
out = {}
|
||
for i, e in enumerate(ids):
|
||
row = raw[i]
|
||
d = {}
|
||
off = 0
|
||
for proj, tensor, dt, shape, nb in self._segs:
|
||
seg = row[off:off + nb]
|
||
# 按段 dtype 重解释：weight→uint32；v1 affine 的 uint16 scales/biases
|
||
# 可再 view(bf16)；mxfp4 的 uint8 scales 保持 uint8。
|
||
if dt == "uint32":
|
||
arr = seg.view(mx.uint32).reshape(shape)
|
||
elif dt == "uint16":
|
||
arr = seg.view(mx.uint16).reshape(shape)
|
||
if view_bf16:
|
||
arr = arr.view(mx.bfloat16)
|
||
else: # uint8（mxfp4 scales）
|
||
arr = seg.reshape(shape)
|
||
d[f"{proj}.{tensor}"] = arr
|
||
off += nb
|
||
out[e] = d
|
||
return out
|
||
|
||
def keys(self) -> "list[str]":
|
||
return [f"{proj}.{tensor}" for proj, tensor, _, _, _ in self._segs]
|
||
|
||
def acquire(self, layer: int, expert_ids):
|
||
"""对齐 ResidentExpertPool.acquire：返回 (pool_arrays, slots)。
|
||
|
||
pool_arrays: {proj.tensor: (n_uniq,...) 堆叠数组}，喂给 _sub.forward。
|
||
slots: 每个 expert_id 在 unique 列表中的下标（= local，reshape 前）。
|
||
"""
|
||
ids = [int(e) for e in expert_ids]
|
||
uniq = list(dict.fromkeys(ids))
|
||
experts = self.load_experts(layer, uniq)
|
||
pool = {k: mx.stack([experts[e][k] for e in uniq], axis=0) for k in self.keys()}
|
||
pos = {e: i for i, e in enumerate(uniq)}
|
||
slots = [pos[e] for e in ids]
|
||
return pool, slots
|
||
|
||
def close(self):
|
||
self._pf_pool.shutdown(wait=True)
|
||
for fd in self._fds.values():
|
||
os.close(fd)
|
||
self._fds.clear()
|