332 lines
15 KiB
C++
332 lines
15 KiB
C++
// [6] Phase 2 方案B：真实区槽状态 C++ 全接管（1 次同步版）。
|
||
// 复刻 Python ResidentExpertPool 的 _slot_of/_free/_freq + _choose_victim/_alloc_slot 语义：
|
||
// - free 优先(front pop，free 初始 [0,cap))；free 空 → LFU 驱逐，受害者槽直接复用(不回 free)。
|
||
// - _choose_victim：candidates=插入序中 ∉current 者；victim=min(freq, 候选下标)，驱逐不删 freq。
|
||
// - dual 语义：真实区命中不移动插入序；仅新放入 miss 追加序尾。
|
||
#include "demand.h"
|
||
#include "side_region.h"
|
||
#include "../io/bg_reader.h"
|
||
|
||
#include <chrono>
|
||
#include <cstdio>
|
||
#include <cstdlib>
|
||
#include <unordered_map>
|
||
#include <unordered_set>
|
||
|
||
struct RealLayer {
|
||
std::vector<int> order; // 插入序(LRU tie-break)，与 e2r 同步维护
|
||
std::unordered_map<int, int> e2r; // expert -> slot [0,cap)
|
||
std::vector<int> free_rows; // 空闲槽(front pop，仿 free.pop(0))
|
||
std::unordered_map<int, uint32_t> freq; // LFU 频次(驱逐不删，与 Python 一致)
|
||
std::unordered_set<int> pinned; // AUTOPIN 钉死集：驱逐永不选(见 real_pin)
|
||
int cap = 0;
|
||
long access = 0; // 累计访问(decay 用)
|
||
bool inited = false;
|
||
};
|
||
static std::mutex g_real_mutex;
|
||
static std::map<int, RealLayer> g_real;
|
||
|
||
// demand 统计：累计 + 本次(供 Python 更新 rp.hits/misses/gpu_fastpath/gpu_fallback)。
|
||
static std::mutex g_dstat_mutex;
|
||
// 本次 [hitpos, misspos, loads, fallback01, unplaced]
|
||
static long g_d_last[5] = {0, 0, 0, 0, 0};
|
||
static std::atomic<long> g_demand_ticket{1000000000}; // demand 并行 pread 用的独立 ticket 段
|
||
|
||
// 诊断计时(DEMAND_TIMING=1)：累计各段主线程微秒，供定位结构性开销。默认关。
|
||
static bool g_dt_on = false;
|
||
static double g_dt[6] = {0, 0, 0, 0, 0, 0}; // [inds_eval, pool_eval, side_snap, real_lock, core, build]
|
||
static inline double dt_now_us() {
|
||
return std::chrono::duration<double, std::micro>(
|
||
std::chrono::steady_clock::now().time_since_epoch()).count();
|
||
}
|
||
std::vector<double> demand_timings() { return {g_dt[0], g_dt[1], g_dt[2], g_dt[3], g_dt[4], g_dt[5]}; }
|
||
void demand_timing_enable(bool on) {
|
||
g_dt_on = on;
|
||
for (int i = 0; i < 6; ++i) g_dt[i] = 0;
|
||
}
|
||
|
||
static void real_ensure_locked(RealLayer& c, int cap) {
|
||
if (c.inited) return;
|
||
c.cap = cap;
|
||
c.free_rows.clear();
|
||
for (int r = 0; r < cap; ++r) c.free_rows.push_back(r); // free 初始 [0,cap)
|
||
c.inited = true;
|
||
}
|
||
|
||
void real_init(int layer, int cap) {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
real_ensure_locked(g_real[layer], cap);
|
||
}
|
||
|
||
std::vector<int> real_region_contents(int layer) {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
std::vector<int> out;
|
||
auto it = g_real.find(layer);
|
||
if (it != g_real.end())
|
||
for (auto& p : it->second.e2r) { out.push_back(p.first); out.push_back(p.second); }
|
||
return out;
|
||
}
|
||
|
||
int real_region_count(int layer) {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
auto it = g_real.find(layer);
|
||
return it == g_real.end() ? 0 : static_cast<int>(it->second.e2r.size());
|
||
}
|
||
|
||
int real_pinned_count(int layer) {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
auto it = g_real.find(layer);
|
||
return it == g_real.end() ? 0 : static_cast<int>(it->second.pinned.size());
|
||
}
|
||
|
||
void real_reset() {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
g_real.clear();
|
||
}
|
||
|
||
// 复刻 _choose_victim：遍历插入序(order)选 ∉current 且 freq 最小者，并列取最早(候选下标最小)。
|
||
// pinned(AUTOPIN 钉死)与 current 同样永不选。返回 expert id；-1 表示无可驱逐。调用方须持 g_real_mutex。
|
||
static int choose_victim_locked(RealLayer& c, const std::unordered_set<int>& current) {
|
||
int victim = -1;
|
||
uint32_t best = 0;
|
||
for (int e : c.order) {
|
||
if (!c.e2r.count(e) || current.count(e) || c.pinned.count(e)) continue;
|
||
uint32_t f = c.freq.count(e) ? c.freq[e] : 0;
|
||
if (victim < 0 || f < best) { victim = e; best = f; } // 并列不更新 → 保留更早者
|
||
}
|
||
return victim;
|
||
}
|
||
|
||
// AUTOPIN 预热钉死：把历史热专家注册进真实区并标记 pinned(此后驱逐永不选它)。
|
||
// 已驻 → 仅补 pinned 标记、返回原槽；未驻 → free 头取槽登记(池行字节由调用方负责写入)。
|
||
// 返回与 experts 平行的槽位；空闲耗尽分不到槽给 -1(调用方跳过该专家)。幂等。
|
||
// 前置：调用方须先 real_init(否则 free_rows 为空、全返 -1)。
|
||
std::vector<int> real_pin(int layer, const std::vector<int>& experts) {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
RealLayer& c = g_real[layer];
|
||
std::vector<int> out;
|
||
out.reserve(experts.size());
|
||
for (int e : experts) {
|
||
auto it = c.e2r.find(e);
|
||
if (it != c.e2r.end()) { c.pinned.insert(e); out.push_back(it->second); continue; }
|
||
if (c.free_rows.empty()) { out.push_back(-1); continue; }
|
||
int slot = c.free_rows.front();
|
||
c.free_rows.erase(c.free_rows.begin()); // free.pop(0)
|
||
c.e2r[e] = slot;
|
||
c.order.push_back(e);
|
||
c.pinned.insert(e);
|
||
out.push_back(slot);
|
||
}
|
||
return out;
|
||
}
|
||
|
||
// AUTOPIN 热度持久化：导出各层 LFU 累计频次，扁平 [layer, expert, count, ...]。
|
||
// 仅 lfu 策略下 demand 路径计频；驱逐不删 freq，故即该进程全生命周期的路由热度，
|
||
// 由 Python 侧按增量差分合并进 usage 文件(dual decode 无 Python 物化点，这是唯一计频来源)。
|
||
std::vector<long> real_freq_dump() {
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
std::vector<long> out;
|
||
for (auto& kv : g_real)
|
||
for (auto& p : kv.second.freq) {
|
||
out.push_back(static_cast<long>(kv.first));
|
||
out.push_back(static_cast<long>(p.first));
|
||
out.push_back(static_cast<long>(p.second));
|
||
}
|
||
return out;
|
||
}
|
||
|
||
// 复刻 _alloc_slot（e 为 miss，不会已在 e2r）：free 优先，否则 LFU 驱逐复用受害者槽。
|
||
// 返回 slot；-1 表示无可驱逐(超容量，不该在 dual 发生)。调用方须持 g_real_mutex。
|
||
static int alloc_slot_locked(RealLayer& c, int e, const std::unordered_set<int>& current) {
|
||
auto it = c.e2r.find(e);
|
||
if (it != c.e2r.end()) return it->second;
|
||
int slot;
|
||
if (!c.free_rows.empty()) {
|
||
slot = c.free_rows.front();
|
||
c.free_rows.erase(c.free_rows.begin()); // free.pop(0)
|
||
} else {
|
||
int victim = choose_victim_locked(c, current);
|
||
if (victim < 0) return -1;
|
||
slot = c.e2r[victim];
|
||
c.e2r.erase(victim);
|
||
for (auto oit = c.order.begin(); oit != c.order.end(); ++oit)
|
||
if (*oit == victim) { c.order.erase(oit); break; }
|
||
}
|
||
c.e2r[e] = slot;
|
||
c.order.push_back(e);
|
||
return slot;
|
||
}
|
||
|
||
// LFU 频次 bump（canonical：本次全部唯一专家各 +1）+ decay。调用方须持 g_real_mutex。
|
||
static void note_access_locked(RealLayer& c, const std::vector<int>& uniq_access,
|
||
bool lfu, int decay_interval) {
|
||
if (!lfu) return;
|
||
for (int e : uniq_access) c.freq[e] += 1;
|
||
c.access += static_cast<long>(uniq_access.size());
|
||
if (decay_interval > 0 && c.access >= decay_interval) {
|
||
for (auto it = c.freq.begin(); it != c.freq.end();) {
|
||
it->second /= 2;
|
||
if (it->second == 0) it = c.freq.erase(it);
|
||
else ++it;
|
||
}
|
||
c.access = 0;
|
||
}
|
||
}
|
||
|
||
// 核心状态机（纯 CPU、无 I/O）：给定 host inds(ip[n]) + 侧区快照(side)，算 local、分配 miss 槽，
|
||
// 把需要落盘的新放入 (expert, slot) 追加到 placements（字节由调用方在锁外并行 pread 落池）。
|
||
// 返回 local(int32 vector)。stats: [hitpos, misspos, loads(=placements 数), unplaced]。
|
||
static std::vector<int32_t> demand_core_locked(
|
||
RealLayer& c, const uint32_t* ip, size_t n, const std::unordered_map<int, int>& side,
|
||
bool lfu, int decay_interval, std::vector<std::pair<int, int>>& placements, long stats[4]) {
|
||
std::vector<int32_t> local(n, -1);
|
||
std::vector<int> uniq_miss, access_order;
|
||
std::unordered_set<int> miss_seen, access_seen;
|
||
int hitpos = 0;
|
||
// pass1：算命中(侧区覆盖真实区)、收集 miss(首见序)、收集唯一访问(freq)。
|
||
for (size_t i = 0; i < n; ++i) {
|
||
int e = static_cast<int>(ip[i]);
|
||
if (access_seen.insert(e).second) access_order.push_back(e);
|
||
auto sit = side.find(e);
|
||
if (sit != side.end()) { local[i] = sit->second; ++hitpos; continue; }
|
||
auto rit = c.e2r.find(e);
|
||
if (rit != c.e2r.end()) { local[i] = rit->second; ++hitpos; continue; }
|
||
if (miss_seen.insert(e).second) uniq_miss.push_back(e);
|
||
}
|
||
note_access_locked(c, access_order, lfu, decay_interval);
|
||
// pass2：miss 分配槽（不落盘）。current = 本前向全部唯一路由专家(命中+miss)：绝不驱逐本前向要读的
|
||
// 任何专家的槽（否则真实区命中专家的槽被 miss 复写 → 脏字节。比 Python 仅护 miss 更严格、更正确）。
|
||
std::unordered_map<int, int> new_slot;
|
||
long unplaced = 0;
|
||
for (int e : uniq_miss) {
|
||
int slot = alloc_slot_locked(c, e, access_seen);
|
||
if (slot < 0) {
|
||
// 超容量：可驱逐槽(= cap − pinned − 本前向命中数)不够安放全部 miss。只能落 0 号槽，
|
||
// 即拿别的专家的权重参与计算 → 该前向逐位错算。调用方须靠 stats[3] 拦截并回退到
|
||
// host/fetch 路径；这里只保证「绝不静默」。
|
||
new_slot[e] = 0;
|
||
++unplaced;
|
||
continue;
|
||
}
|
||
new_slot[e] = slot;
|
||
placements.emplace_back(e, slot);
|
||
}
|
||
// pass3：回填 miss 位置。
|
||
int misspos = 0;
|
||
for (size_t i = 0; i < n; ++i) {
|
||
if (local[i] < 0) {
|
||
auto it = new_slot.find(static_cast<int>(ip[i]));
|
||
local[i] = (it != new_slot.end()) ? it->second : 0;
|
||
++misspos;
|
||
}
|
||
}
|
||
stats[0] = hitpos; stats[1] = misspos; stats[2] = static_cast<long>(placements.size());
|
||
stats[3] = unplaced;
|
||
return local;
|
||
}
|
||
|
||
// demand 全接管：inds 惰性(内部 eval 一次=1 次同步)；side_gen 指定侧区代；pool_list 为 _segs 顺序
|
||
// 的 per-key 池数组(已 eval、指针稳定)。返回 local(int32, inds.shape)。
|
||
mx::array demand_dual(
|
||
const mx::array& inds, const std::vector<mx::array>& pool_list,
|
||
const std::vector<int>& seg_nbytes, int layer, int side_gen, const std::string& path,
|
||
int stride, int cap, bool lfu, int decay_interval, mx::StreamOrDevice s = {}) {
|
||
if (seg_nbytes.size() != pool_list.size())
|
||
throw std::invalid_argument("demand_dual: seg_nbytes.size() != pool_list.size()");
|
||
double t0 = g_dt_on ? dt_now_us() : 0;
|
||
// 关键:inds 常是 argpartition(...)[..., -k:] 切片 → 非连续 strided 视图(每行按父数组 stride E
|
||
// 偏移)。若直接按连续读 data(),仅首行(token0)正确,后续 token 读到错位内存 → 装错专家、seq≥2 全错
|
||
// (投机 verify 全体受害)。contiguous() 强制物化为连续,读取才逐位正确。
|
||
mx::array ids = mx::contiguous(inds);
|
||
ids.eval(); // 唯一同步
|
||
size_t n = ids.size();
|
||
const uint32_t* ip = ids.data<uint32_t>();
|
||
double t1 = g_dt_on ? dt_now_us() : 0;
|
||
std::vector<uint8_t*> ptrs;
|
||
ptrs.reserve(pool_list.size());
|
||
for (auto& a0 : pool_list) { mx::array a = a0; a.eval(); ptrs.push_back(a.data<uint8_t>()); }
|
||
double t2 = g_dt_on ? dt_now_us() : 0;
|
||
std::unordered_map<int, int> side = sideregion_snapshot(layer, side_gen); // 侧区快照(该代)
|
||
double ta = g_dt_on ? dt_now_us() : 0;
|
||
long stats[4];
|
||
std::vector<int32_t> local;
|
||
std::vector<std::pair<int, int>> placements; // (expert, slot)：锁外并行落盘
|
||
double tb;
|
||
{
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
tb = g_dt_on ? dt_now_us() : 0;
|
||
RealLayer& c = g_real[layer];
|
||
real_ensure_locked(c, cap);
|
||
local = demand_core_locked(c, ip, n, side, lfu, decay_interval, placements, stats);
|
||
}
|
||
// 锁外：把 miss 专家的字节 pread 落真实区槽。复刻基线并发模型——多 miss 派给 BgReader worker
|
||
// 并行 pread（高优队列、直写池段行，无主线程 tmp/memcpy），本层等待其完成（并行 → 远快于串行）。
|
||
static const bool kSkipIO = []() {
|
||
const char* e = std::getenv("DEMAND_SKIP_IO");
|
||
return e && e[0] == '1';
|
||
}();
|
||
if (!placements.empty() && !kSkipIO) {
|
||
std::vector<long> seg_off(seg_nbytes.size()), seg_nb(seg_nbytes.size());
|
||
long acc = 0;
|
||
for (size_t k = 0; k < seg_nbytes.size(); ++k) {
|
||
seg_off[k] = acc; seg_nb[k] = seg_nbytes[k]; acc += seg_nbytes[k];
|
||
}
|
||
std::vector<long> tickets;
|
||
tickets.reserve(placements.size());
|
||
for (auto& pr : placements) {
|
||
long tk = g_demand_ticket.fetch_add(1);
|
||
bg_pread_into_pool(pool_list, seg_off, seg_nb, pr.second, pr.first, path,
|
||
static_cast<long>(stride), tk, /*prio=*/1,
|
||
/*nocache=*/false); // demand route 读走 cache：段偏移非页对齐
|
||
tickets.push_back(tk);
|
||
}
|
||
for (long tk : tickets) bg_reader_wait(tk); // 并行 pread 完成 → 池槽字节就绪
|
||
}
|
||
double t3 = g_dt_on ? dt_now_us() : 0;
|
||
{
|
||
std::lock_guard<std::mutex> lk(g_dstat_mutex);
|
||
g_d_last[0] = stats[0]; g_d_last[1] = stats[1]; g_d_last[2] = stats[2];
|
||
g_d_last[3] = (stats[1] == 0) ? 0 : 1;
|
||
g_d_last[4] = stats[3];
|
||
}
|
||
if (stats[3] > 0) {
|
||
static std::atomic<long> warned{0};
|
||
long k = warned.fetch_add(1);
|
||
if (k < 8 || k % 1000 == 0)
|
||
fprintf(stderr,
|
||
"[DEMAND_DUAL] 错算：layer=%d 有 %ld 个专家分不到槽(落 0 号槽)。"
|
||
"inds=%zu cap=%d pinned=%d —— 该前向输出不可信,请调小 PREFILL_CHUNK "
|
||
"或调大 EXPERT_SLOTS / 调小 AUTOPIN_BUDGET_FRAC。\n",
|
||
layer, stats[3], n, cap, real_pinned_count(layer));
|
||
}
|
||
mx::array out = mx::array(local.data(), ids.shape(), mx::int32);
|
||
if (g_dt_on) {
|
||
double t4 = dt_now_us();
|
||
g_dt[0] += t1 - t0; g_dt[1] += t2 - t1; g_dt[2] += ta - t2;
|
||
g_dt[3] += tb - ta; g_dt[4] += t3 - tb; g_dt[5] += t4 - t3;
|
||
}
|
||
return out;
|
||
}
|
||
|
||
// 取本次 demand 统计 [hitpos, misspos, loads, fallback01, unplaced]（主线程串行，安全）。
|
||
std::vector<long> demand_last_stats() {
|
||
std::lock_guard<std::mutex> lk(g_dstat_mutex);
|
||
return {g_d_last[0], g_d_last[1], g_d_last[2], g_d_last[3], g_d_last[4]};
|
||
}
|
||
|
||
// 测试壳：纯状态推进(不 pread/不侧区)，把 experts_flat 当一次 demand 的 inds，返回 local 槽位。
|
||
// 供 LFU 驱逐语义逐步等价单测(与 Python 参考对拍)。
|
||
std::vector<int> real_debug_place(int layer, const std::vector<int>& experts_flat, int cap,
|
||
bool lfu, int decay_interval) {
|
||
std::vector<uint32_t> u(experts_flat.begin(), experts_flat.end());
|
||
long stats[4];
|
||
std::vector<std::pair<int, int>> placements;
|
||
std::lock_guard<std::mutex> lk(g_real_mutex);
|
||
RealLayer& c = g_real[layer];
|
||
real_ensure_locked(c, cap);
|
||
std::unordered_map<int, int> empty_side;
|
||
auto local = demand_core_locked(c, u.data(), u.size(), empty_side, lfu, decay_interval,
|
||
placements, stats);
|
||
return std::vector<int>(local.begin(), local.end());
|
||
}
|