运行时层模板与生成循环 深度解析

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()