Expert cache 与 fetch_v2 深度解析

78def _range_fetch(name, retries=6):
79    t = INV[name]
80    start = 8 + t["hlen"] + t["offsets"][0]
81    size = t["offsets"][1] - t["offsets"][0]
82    req = urllib.request.Request(
83        BASE + t["shard"], headers={"Range": f"bytes={start}-{start+size-1}"})
84    for attempt in range(retries):
85        try:
86            t0 = time.time()
87            with urllib.request.urlopen(req, timeout=120) as r:
88                buf = r.read()
89            if len(buf) != size:
90                raise IOError(f"short read {len(buf)}/{size}")
91            with _stats_lock:
92                stats["http_bytes"] += size
93                stats["http_s"] += time.time() - t0
94            return buf
95        except Exception:
96            if attempt == retries - 1:
97                raise
98            time.sleep(1.5 * (attempt + 1))
99

118def fetch_expert_raw(layer, eid):
119    """Return dict wname -> (packed U8 np, scale U8 np). Disk-cached as raw packed bytes (~17.5MB)."""
120    path = os.path.join(ECACHE, f"L{layer}-E{eid}.npz")
121    if os.path.exists(path):
122        with _stats_lock:
123            stats["expert_disk"] += 1
124        z = np.load(path)
125        return {w: (z[w + "_p"], z[w + "_s"]) for w in ("w1", "w2", "w3")}
126    pfx = _expert_prefix(layer, eid)
127    out = {}
128    for w in ("w1", "w2", "w3"):
129        tp = INV[pfx + f".{w}.weight_packed"]
130        ts = INV[pfx + f".{w}.weight_scale"]
131        p = np.frombuffer(_range_fetch(pfx + f".{w}.weight_packed"), dtype=np.uint8).reshape(tp["shape"])
132        s = np.frombuffer(_range_fetch(pfx + f".{w}.weight_scale"), dtype=np.uint8).reshape(ts["shape"])
133        out[w] = (p, s)
134    tmp = path + f".tmp{os.getpid()}"
135    np.savez(tmp, **{w + suf: arr for w, (p, s) in out.items()
136                     for suf, arr in (("_p", p), ("_s", s))})
137    os.replace(tmp + ".npz" if os.path.exists(tmp + ".npz") else tmp, path)
138    register_cache_file(path)
139    with _stats_lock:
140        stats["expert_http"] += 1
141    return out
142
143
144def fetch_experts(layer, eids, workers=12, dequant=True):
145    """Parallel fetch of a routed set. dequant=True -> {eid: {w: fp32 torch tensor}};
146    dequant=False -> {eid: {w: (packed u8, scale u8) contiguous np arrays}}."""
147    raw = {}
148    with concurrent.futures.ThreadPoolExecutor(workers) as ex:
149        futs = {ex.submit(fetch_expert_raw, layer, e): e for e in eids}
150        for f in concurrent.futures.as_completed(futs):
151            raw[futs[f]] = f.result()
152    if not dequant:
153        return {e: {w: (np.ascontiguousarray(p), np.ascontiguousarray(s))
154                    for w, (p, s) in ws.items()} for e, ws in raw.items()}
155    out = {}
156    for e, ws in raw.items():
157        out[e] = {w: torch.from_numpy(dequant_mxfp4(p, s)) for w, (p, s) in ws.items()}
158    return out

56# ---- fixed intra-expert layout (offset, nbytes, shape) --------------------
57EXPERT_SPAN = 17547264
58_P, _S = 5505024, 344064
59LAYOUT = [
60    ("w1_p", 0,            _P, (3072, 1792)),
61    ("w1_s", _P,           _S, (3072, 112)),
62    ("w2_p", _P + _S,      _P, (3584, 1536)),
63    ("w2_s", 2 * _P + _S,  _S, (3584, 96)),
64    ("w3_p", 2 * (_P + _S), _P, (3072, 1792)),
65    ("w3_s", 3 * _P + 2 * _S, _S, (3072, 112)),
66]
67
68USER_AGENT = "k3-lazy-fetch/2.0"
69MAX_CONNS = int(os.environ.get("K3_FETCH_CONNS", "4"))
70BACKEND = os.environ.get("K3_FETCH_BACKEND", "httpclient")  # or "httpx"
71MAX_COALESCE = int(os.environ.get("K3_FETCH_COALESCE", "4"))  # experts per merged range
72
73# ---- local read path (K3_EXPERT_READ) -------------------------------------
74# "mmap"  : np.memmap views over the .bin cache; the pages are demand-faulted
75#           later, inside the GEMV kernel. Measured 0.87-0.93 GB/s.
76# "pread" : the whole layer's selected experts are pread() in parallel by a
77#           persistent worker pool into reused page-aligned buffers, so the
78#           kernel only ever touches resident memory. Measured ~7.0 GB/s.
79EXPERT_READ = os.environ.get("K3_EXPERT_READ", "pread")
80# Real F_NOCACHE grid on the 8 TB M1 Max SSD: 8 workers sustained 7.51
81# logical GB/s versus 6.91/6.69/6.91 for 4/6/12.  Keep the env override: newer
82# chips and different SSD widths must rerun tools/bench_expert_io_grid.py.
83PREAD_WORKERS = runtime_platform.configured_cpu_workers(
84    "K3_PREAD_WORKERS", 8
85)