deltafin 本地 2.8T MoE 推理实验

121DEV = _resolve_device()                                      # cpu | mps | cuda[:N]
122SPINE = os.environ.get("K3_SPINE") or _auto_spine()           # bf16 | int8 | mixed
123if SPINE == "bf16" and "K3_SPINE" not in os.environ:
124    print("[config] int8 spine not found — using bf16 (2x the per-token I/O). "
125          "Build it with: python tools/convert_spine_int8.py", flush=True)
126
127# --- mixed-precision spine (K3_SPINE=mixed; DEFAULT IS UNCHANGED) ------------
128# The 53 GB int8 spine is re-read every token, so bytes are seconds. A mixed
129# spine keeps int8 where the error lands on token identity and drops selected
130# ROLES to 6 or 4 bits; tools/mixed_spine.py reads the mix (falling back to the
131# int8 spine for every tensor the policy left alone) and dequantizes each codec
132# with its own Metal kernel. Build one with:
133#     python tools/convert_spine_int4.py --policy moe4
134# Per-role error measurements that justify a policy: tools/spine_sensitivity.py.
135MIXED = SPINE == "mixed"
136QUANT = SPINE in ("int8", "mixed")      # "a quantized spine of some kind"
137if MIXED:
138    import mixed_spine  # noqa: E402
139    if not mixed_spine.available():
140        sys.exit(f"K3_SPINE=mixed but no mixed spine at {mixed_spine.MIXED_DIR}. "
141                 f"Build one with: python tools/convert_spine_int4.py --policy moe4")
142    print(f"[spine] mixed: {mixed_spine.describe()}", flush=True)
143    if os.environ.get("K3_PROFILE", "0") == "1":
144        import atexit
145        atexit.register(lambda: print(mixed_spine.phase_report(), flush=True))
146# K3_APPROX=1 = "approx mode": approximate numerics (fp16 weights) + n-gram
147# speculation. Output stays coherent but near-tie tokens may differ from the
148# fp32 reference — never use for oracle runs. Speed effect is unproven until a
149# quiet-machine A/B; if it measures faster it can earn a faster name.
150APPROX = os.environ.get("K3_APPROX", "0") == "1"
151DT = torch.float16 if (APPROX or os.environ.get("K3_DTYPE", "fp32") == "fp16") else torch.float32
152
153

1332def build_layers():
1333    global _INT8_KDA_QKV_CONTROLLERS
1334    layers = []
1335    if not TEMPLATES:
1336        with torch.device("meta"):
1337            for i in range(NL):
1338                layers.append(ml.KimiDecoderLayer(config, i).eval())
1339        return layers
1340    with torch.device("meta"):
1341        l0 = ml.KimiDecoderLayer(config, 0).eval()        # dense KDA (unique shape)
1342        tpl_kda = ml.KimiDecoderLayer(config, 1).eval()   # KDA + MoE class
1343        tpl_mla = ml.KimiDecoderLayer(config, 3).eval()   # MLA + MoE class
1344    if INT8_KDA_QKV and _INT8_KDA_QKV_STATE.enabled:
1345        installed = []
1346        try:
1347            # The reusable KDA template executes 68 layers per pass. Layer 0
1348            # executes once and stays on the dense fallback: giving it another
1349            # 337.7 MiB packed arena is a poor memory/I/O trade, while attempting
1350            # to register one MPS arena on both modules caused the original
1351            # full-model candidate to disable itself before its first call.
1352            # A future backend with an explicitly supported shared-storage
1353            # contract can reconsider layer 0 independently.
1354            if tpl_kda.is_linear_attn:
1355                installed.append(
1356                    spine_fast.install_dynamic_q8_qkv(
1357                        tpl_kda,
1358                        DEV,
1359                        _INT8_KDA_QKV_STATE,
1360                        _PACKED_Q8_BACKEND,
1361                        storage_mode=_INT8_KDA_STORAGE_MODE,
1362                        stage_sync_mode=_INT8_KDA_STAGE_SYNC_MODE,
1363                    )
1364                )
1365        except (RuntimeError, ValueError, MemoryError) as exc:
1366            for controller in reversed(installed):
1367                controller.uninstall()
1368            _INT8_KDA_QKV_STATE.disable(exc)
1369        else:
1370            _INT8_KDA_QKV_CONTROLLERS = installed
1371            packed_bytes = sum(
1372                controller.nbytes for controller in installed
1373            )
1374            print(
1375                f"[kda-bundle] native packed-int8 active for "
1376                f"{len(installed)} KDA template(s), "
1377                f"{packed_bytes/2**20:.1f} MiB persistent storage, "
1378                f"device={DEV.type}, max_T={_PACKED_Q8_BACKEND.max_tokens}, "
1379                f"fused_bundle={int(_PACKED_Q8_BACKEND.fuse_qkv)}, "
1380                f"storage={_INT8_KDA_STORAGE_MODE}, "
1381                f"stage_sync={_INT8_KDA_STAGE_SYNC_MODE}",
1382                flush=True,
1383            )
1384    if TEMPLATE_ARENA:
1385        _bind_template_arena(((0, l0), (1, tpl_kda), (3, tpl_mla)))
1386    for i in range(NL):
1387        layers.append(l0 if i == 0 else
1388                      tpl_kda if config.is_kda_layer(i) else tpl_mla)
1389    return layers
1390

1605def forward_pass(layers, cache, hidden, step, verbose=True):
1606    """hidden: [1, T, H] fp32. Returns logits [1, T, vocab]."""
1607    T = hidden.shape[1]
1608    if pilot.PILOT:
1609        pilot.init(config, DEV, _pilot_load, PFX,
1610                   load_packed=_load_int8_packed if QUANT else None,
1611                   native_int8=(
1612                       _PACKED_Q8_BACKEND.matmul
1613                       if (
1614                           _PACKED_Q8_BACKEND.available
1615                           and DEV.type == "mps"
1616                       )
1617                       else False
1618                   ))
1619        pilot.begin_pass(fetch_v2)
1620        if PREAD:
1621            fetch_v2.drop_prefetch()    # nothing from the previous pass is valid
1622            if pilot.ASYNC_DRAIN:
1623                pilot.install_async_drain(fetch_v2.reader())
1624    elif EXPERT_PREFETCH:
1625        global _PREV_SEL
1626        _PREV_SEL = dict(_LAST_SEL)     # last token's routing = this pass's oracle
1627        fetch_v2.drop_prefetch()        # nothing from the previous pass is valid
1628    past = cache.get_seq_length() or 0
1629    mask = causal_mask(T, past) if T > 1 else None
1630    block_residual = hidden.new_zeros(T, 0, H)
1631
1632    def _next_unpinned(j):
1633        while j < NL and j < PIN_N and getattr(layers[j], "_k3_res", False):
1634            j += 1
1635        return j
1636
1637    nxt = _next_unpinned(0)
1638    fut = (_PRELOADER.submit(_spine_read, layers[nxt], f"{PFX}layers.{nxt}.")
1639           if PRELOAD and nxt < NL else None)
1640    for i, layer in enumerate(layers):
1641        _step_ctx["layer"] = i
1642        if TEMPLATES:
1643            layer.layer_idx = i
1644            layer.self_attn.layer_idx = i
1645        pinned = i < PIN_N and getattr(layer, "_k3_res", False)
1646        if not pinned:
1647            if PRELOAD and fut is not None and i == nxt:
1648                _tw = time.time()
1649                blobs = fut.result()
1650                TIMES["preload_wait"] += time.time() - _tw
1651                j = _next_unpinned(i + 1)
1652                fut = (_PRELOADER.submit(_spine_read, layers[j], f"{PFX}layers.{j}.")
1653                       if j < NL else None)
1654                nxt = j
1655                _spine_apply(layer, f"{PFX}layers.{i}.", blobs)
1656            else:
1657                if TEMPLATES or spine_fast.PACK:
1658                    pfx = f"{PFX}layers.{i}."
1659                    _spine_apply(layer, pfx, _spine_read(layer, pfx))
1660                else:
1661                    materialize_resident(layer, f"{PFX}layers.{i}.")
1662            if i < PIN_N:
1663                layer._k3_res = True   # pinned from now on
1664        if pilot.enabled():
1665            pilot.arm(layer)
1666        if PROFILE and DEV.type in ("mps", "cuda"):
1667            _device_synchronize()
1668        t0 = time.time()
1669        hidden, block_residual = layer(
1670            hidden, attention_mask=mask, position_ids=None,
1671            past_key_values=cache, use_cache=True, block_residual=block_residual)
1672        if PROFILE and DEV.type in ("mps", "cuda"):
1673            _device_synchronize()
1674        dt_layer = time.time() - t0
1675        TIMES["compute"] += dt_layer
1676        if PROFILE:
1677            k = "kda" if layer.is_linear_attn else "mla"
1678            PROF[k] += dt_layer
1679            PROF["n_" + k] += 1
1680        if not TEMPLATES and not (i < PIN_N):
1681            dematerialize(layer)
1682        if verbose and (i % 10 == 0 or i == NL - 1):
1683            print(f"    layer {i:2d}/92 done  (res_io {TIMES['resident_io']:.0f}s "
1684                  f"exp {TIMES['expert_fetch']:.0f}s comp {TIMES['compute']:.0f}s)",
1685                  flush=True)
1686    TRACE.end_pass()

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)

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

128    def do_POST(self):
129        try:
130            body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)) or 0) or b"{}")
131        except json.JSONDecodeError:
132            return self._err(400, "invalid JSON body")
133        chat = self.path in ("/v1/chat/completions", "/chat/completions")
134        comp = self.path in ("/v1/completions", "/completions")
135        if not (chat or comp):
136            return self._err(404, f"no route {self.path}")
137
138        if chat:
139            messages = body.get("messages")
140            if not messages:
141                return self._err(400, "messages required")
142            ids = _tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
143        else:
144            prompt = body.get("prompt")
145            if not isinstance(prompt, str):
146                return self._err(400, "prompt (string) required")
147            ids = _tok.encode(prompt)
148        # OpenAI semantics: omitted max_tokens means the model decides. Chat ends
149        # naturally at EOS; raw completions have no terminator, so only THEY get
150        # a default cap (256) — an explicit max_tokens is always honored.
151        req_max = body.get("max_tokens")
152        if req_max:
153            max_new = int(req_max)
154        else:
155            max_new = 1_000_000 if chat else 256
156        if MAX_TOKENS_CAP:
157            max_new = min(max_new, MAX_TOKENS_CAP)
158        stream = bool(body.get("stream"))
159        rid = ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex[:20]
160        created = int(time.time())
161
162        if not _lock.acquire(timeout=5):
163            return self._err(429, "a generation is already running (Deltafin serves one at a time)")
164        try:
165            mode = "chat" if chat else "completion"
166            cached = _memo.get(mode, ids, max_new)
167            if cached is not None:
168                print(f"[serve] deterministic response memo hit "
169                      f"({len(cached.token_ids)} tokens)", flush=True)
170            if stream:

23class DeterministicResponseMemo:
24    def __init__(self, max_entries: int):
25        if max_entries < 0:
26            raise ValueError("max_entries must be non-negative")
27        self.max_entries = max_entries
28        self._items: collections.OrderedDict[
29            tuple[str, tuple[int, ...], int], MemoizedResponse
30        ] = collections.OrderedDict()
31        self._lock = threading.Lock()
32        self.hits = 0
33        self.misses = 0
34
35    @staticmethod
36    def _key(
37            mode: str, prompt_ids: Iterable[int], max_new: int,
38    ) -> tuple[str, tuple[int, ...], int]:
39        if max_new <= 0:
40            raise ValueError("max_new must be positive")
41        return str(mode), tuple(int(value) for value in prompt_ids), int(max_new)
42
43    def get(
44            self, mode: str, prompt_ids: Iterable[int], max_new: int,
45    ) -> MemoizedResponse | None:
46        if self.max_entries == 0:
47            return None
48        key = self._key(mode, prompt_ids, max_new)
49        with self._lock:
50            value = self._items.pop(key, None)
51            if value is None:
52                self.misses += 1
53                return None
54            self._items[key] = value
55            self.hits += 1
56            return value
57
58    def put(
59            self, mode: str, prompt_ids: Iterable[int], max_new: int,
60            token_ids: Iterable[int], text: str, finish_reason: str,
61    ) -> None:
62        if self.max_entries == 0:
63            return
64        key = self._key(mode, prompt_ids, max_new)
65        value = MemoizedResponse(
66            tuple(int(token) for token in token_ids),
67            str(text),
68            str(finish_reason),
69        )
70        with self._lock:
71            self._items.pop(key, None)
72            self._items[key] = value
73            while len(self._items) > self.max_entries:
74                self._items.popitem(last=False)
75

1885def generate(layers, cache, embed, ids, max_new, spec=None, on_token=None,
1886             verbose_prefill=False, log=lambda *a: None):
1887    """Greedy generation (+ certified-lossless n-gram speculation).
1888
1889    Shared by the CLI and the OpenAI-compatible server. Calls on_token(token_id)
1890    as each token is emitted. Returns at most ``max_new`` tokens and never
1891    streams anything after EOS_ID, including inside a speculative burst."""
1892    if spec is None:
1893        spec = os.environ.get("K3_SPEC", "1") == "1"
1894    if max_new <= 0:
1895        return []
1896    generated = []
1897
1898    def emit(t):
1899        generated.append(t)
1900        if on_token:
1901            on_token(t)
1902
1903    _step_ctx["step"] = 0
1904    logits = forward_pass(layers, cache, embed(ids), step=0, verbose=verbose_prefill)
1905    first = int(logits[0, -1].argmax())
1906    emit(first)
1907    if first == EOS_ID:
1908        return generated
1909    for _k in EXPERT_SEL:      # the union factor that matters is the decode one
1910        EXPERT_SEL[_k] = 0
1911    s = 1
1912    deep = spec and spec_decode.enabled()
1913    while len(generated) < max_new:
1914        _step_ctx["step"] = s
1915        if os.environ.get("K3_PREFETCH", "1") == "1":
1916            prefetch_prev_token()
1917        t0 = time.perf_counter_ns()
1918        tag = ""
1919        if deep:
1920            new, tag = _spec_step_deep(layers, cache, embed,
1921                                       ids + generated, generated[-1], s)
1922            if EOS_ID in new:   # never stream past EOS (a pass emits up to D+1)
1923                new = new[:new.index(EOS_ID) + 1]
1924        else:
1925            draft = ngram_draft(ids + generated) if spec else None
1926            if draft is not None:
1927                snap = snapshot_states(cache)
1928                logits = forward_pass(layers, cache, embed([generated[-1], draft]),
1929                                      step=s, verbose=False)
1930                # Read both verifier decisions in one device-to-host transfer.
1931                # The first position decides accept/rollback; on acceptance the
1932                # second ID is already present instead of forcing another sync.
1933                verified = logits[0].argmax(-1).tolist()
1934                n1 = verified[0]
1935                if n1 == draft:
1936                    new = verified[:2]
1937                    tag = " spec+2"
1938                else:
1939                    restore_states(cache, snap)
1940                    logits = forward_pass(layers, cache, embed([generated[-1]]),
1941                                          step=s, verbose=False)
1942                    new = [int(logits[0, -1].argmax())]
1943                    tag = " spec-miss"
1944            else:
1945                logits = forward_pass(layers, cache, embed([generated[-1]]),
1946                                      step=s, verbose=False)
1947                new = [int(logits[0, -1].argmax())]
1948        # A verifier can accept several tokens in one pass. Bound the burst
1949        # before invoking callbacks or structured logging so max_new remains a
1950        # real API contract and throughput cannot be inflated by over-emission.
1951        new = new[:max_new - len(generated)]
1952        for t in new:
1953            emit(t)
1954        log(s, tag, t0, list(generated))
1955        s += 1
1956        if EOS_ID in new:
1957            break
1958    return generated

1"""Self-tuning, pressure-aware RAM cache for the int8 spine.
2
3WHY THE OLD ONE WAS OFF BY DEFAULT
4----------------------------------
5kimi_run.py's first attempt sized a spine RAM cache from "free + inactive" and
6took 29.7 GB on this 64 GB box.  A decode token went from 14 s to 122 s: the
7inactive pages it counted as spare were the page cache that was already serving
8the spine and the expert blobs, so claiming them as anonymous heap did not add
9capacity, it MOVED capacity from a class the kernel can drop for free (clean
10file pages) into a class it can only reclaim by compressing or swapping.  It
11then compressed, and every token paid for it.  Two things follow, and this
12module implements both:
13
14  1. THE HEADROOM SIGNAL MUST BE `free`, NOT `free + inactive`.  `free` is the
15     only class that can be taken without displacing something that is already
16     doing useful work.  We additionally hold a reserve (K3_SPINE_CACHE_FLOOR_GB,
17     default 8 GB) below which we stop growing, because the page cache still has
18     to hold the expert payloads.
19
20  2. GROWTH MUST BE WATCHED, NOT PREDICTED.  No static budget can know what else
21     the machine will do while we run. Every admission re-reads the kernel's own
22     accounting (Mach host_statistics64 on macOS; /proc on Linux); if the
23     compressor or swapper starts moving, we stop growing and give memory back
24     layer by layer. That is the difference between the 8x cliff and a cache
25     that just quietly becomes smaller.
26
27WHAT THE CACHE IS WORTH, MEASURED ON THIS MACHINE
28-------------------------------------------------
29Reading a layer's blobs cold off the NVMe:            ~7.0 GB/s
30Reading the same blobs back out of the page cache:   ~27-38 GB/s
31So an in-RAM layer is ~5x cheaper, and the 53 GB spine's ~7.5 s/token of read
32collapses toward ~1.5 s for whatever fraction is resident.  The fraction is what
33matters: access is a strict 0..91 cycle, so under LRU a working set larger than
34RAM gets a ~0% hit rate (each layer is evicted just before it is needed again).
35Caching a FIXED SUBSET instead converts that into a hit rate equal to the
36fraction that fits, which is why a partial cache is still worth having -- but
37only if the layers outside the subset can be prevented from evicting it, which
38is what K3_SPINE_STREAM_NOCACHE (see spine_io) is for.
39
40FLAGS (defaults reproduce today's behaviour exactly: no cache at all)
41    K3_SPINE_CACHE_GB=<n>       hard ceiling on the cache, 0 = disabled
42    K3_SPINE_CACHE_AUTO=1       grow toward the ceiling only while the kernel
43                                says there is room; without _GB, cap growth by
44                                queried physical/recommended working-set limits
45    K3_SPINE_CACHE_FLOOR_GB=<n> stop growing with less than this much `free`
46                                (default 8)
47    K3_SPINE_CACHE_TOL_GB=<n>   how much new compressor traffic to tolerate
48                                before shrinking (default 0.5)
49    K3_SPINE_CACHE_SHRINK=<n>   layers to drop per shrink event (default 4)
50
51ON A 128 GB MAC set K3_SPINE_CACHE_AUTO=1 and leave the ceiling alone (or
52K3_SPINE_CACHE_GB=44 K3_SPINE_CACHE_AUTO=1 to be explicit): 53 GB of spine plus
53~12 GB of hot expert page cache plus the model's own working set fits with room
54to spare, and the watchdog will still claw it back if something else on the box
55wants memory.  On 64 GB, K3_SPINE_CACHE_GB=6 K3_SPINE_CACHE_AUTO=1 is the
56conservative setting -- it caches the first ~10 layers and cannot reach the
57pressure point.
58"""