From 9dce5098d92790c3c28712973b78a5e1859065b1 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Fri, 19 Jun 2026 22:16:20 +0800 Subject: [PATCH 01/10] feat: cpp mesh --- kt-kernel/operators/mesh/mesh.hpp | 24 ++ kt-kernel/operators/mesh/mesh_config.hpp | 94 ++++++ kt-kernel/operators/mesh/mesh_decode.hpp | 192 ++++++++++++ kt-kernel/operators/mesh/mesh_eviction.hpp | 329 ++++++++++++++++++++ kt-kernel/operators/mesh/mesh_handoff.hpp | 152 +++++++++ kt-kernel/operators/mesh/mesh_hook.hpp | 79 +++++ kt-kernel/operators/mesh/mesh_io_uring.hpp | 299 ++++++++++++++++++ kt-kernel/operators/mesh/mesh_prefill.hpp | 245 +++++++++++++++ kt-kernel/operators/mesh/mesh_residency.hpp | 318 +++++++++++++++++++ kt-kernel/operators/mesh/mesh_scheduler.hpp | 197 ++++++++++++ kt-kernel/operators/mesh/mesh_slot_pool.hpp | 271 ++++++++++++++++ 11 files changed, 2200 insertions(+) create mode 100644 kt-kernel/operators/mesh/mesh.hpp create mode 100644 kt-kernel/operators/mesh/mesh_config.hpp create mode 100644 kt-kernel/operators/mesh/mesh_decode.hpp create mode 100644 kt-kernel/operators/mesh/mesh_eviction.hpp create mode 100644 kt-kernel/operators/mesh/mesh_handoff.hpp create mode 100644 kt-kernel/operators/mesh/mesh_hook.hpp create mode 100644 kt-kernel/operators/mesh/mesh_io_uring.hpp create mode 100644 kt-kernel/operators/mesh/mesh_prefill.hpp create mode 100644 kt-kernel/operators/mesh/mesh_residency.hpp create mode 100644 kt-kernel/operators/mesh/mesh_scheduler.hpp create mode 100644 kt-kernel/operators/mesh/mesh_slot_pool.hpp diff --git a/kt-kernel/operators/mesh/mesh.hpp b/kt-kernel/operators/mesh/mesh.hpp new file mode 100644 index 000000000..e89c5c109 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh.hpp @@ -0,0 +1,24 @@ +/** + * @file mesh.hpp + * @brief MESH 公共头文件 + * + * MESH 是 KTransformers 的专家权重驻留管理插件。 + * 只需 include 本文件即可使用全部 MESH 接口。 + * + * 设计约束： + * - 所有 MESH 代码集中在本目录，对原版 KT 仅做最小化 hook 修改 + * - mesh_enabled=false 时所有 hook 编译期优化为空操作 + * - 不修改 AMX 内核、NUMA 张量并行、CUDA Graph 调度或 FlashInfer + */ +#pragma once + +#include "mesh_config.hpp" +#include "mesh_slot_pool.hpp" +#include "mesh_io_uring.hpp" +#include "mesh_scheduler.hpp" +#include "mesh_eviction.hpp" +#include "mesh_prefill.hpp" +#include "mesh_decode.hpp" +#include "mesh_handoff.hpp" +#include "mesh_residency.hpp" +#include "mesh_hook.hpp" diff --git a/kt-kernel/operators/mesh/mesh_config.hpp b/kt-kernel/operators/mesh/mesh_config.hpp new file mode 100644 index 000000000..2680af4bf --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_config.hpp @@ -0,0 +1,94 @@ +/** + * @file mesh_config.hpp + * @brief MESH 配置结构体 + * + * 纯配置 POD，由 Python 侧通过 pybind11 注入。 + * GeneralMOEConfig 持有 MeshConfig* 指针，默认 nullptr。 + */ +#pragma once + +#include +#include +#include +#include + +namespace mesh { + +// 权重类型，决定 slot 内存大小 +enum class WeightType : uint8_t { + AMXINT4, // 量化格式：weight(int4) + scale + mins + BF16, // 全精度 bf16 +}; + +// 单个专家在某个 TP 分片上的文件布局 +struct ExpertFileLayout { + int fd = -1; // O_DIRECT 打开的 safetensors 文件描述符 + off_t gate_offset = 0; // gate 矩阵在文件中的偏移 + off_t up_offset = 0; // up 矩阵偏移 + off_t down_offset = 0; // down 矩阵偏移 + off_t gate_scale_offset = 0; // AMXINT4 专用：gate scale 偏移 + off_t up_scale_offset = 0; // AMXINT4 专用：up scale 偏移 + off_t down_scale_offset = 0; // AMXINT4 专用：down scale 偏移 + off_t gate_mins_offset = 0; // AMXINT4 专用：gate mins 偏移 + off_t up_mins_offset = 0; // AMXINT4 专用：up mins 偏移 + off_t down_mins_offset = 0; // AMXINT4 专用：down mins 偏移 + size_t gate_bytes = 0; // gate 权重字节数 + size_t up_bytes = 0; // up 权重字节数 + size_t down_bytes = 0; // down 权重字节数 + size_t gate_scale_bytes = 0; // scale 字节数（很小，几 KB） + size_t up_scale_bytes = 0; + size_t down_scale_bytes = 0; + size_t gate_mins_bytes = 0; + size_t up_mins_bytes = 0; + size_t down_mins_bytes = 0; +}; + +// MESH 配置 +struct MeshConfig { + // ===== 基本开关 ===== + bool enabled = false; // 是否启用 MESH，false 时走原版 KT 路径 + + // ===== 容量配置 ===== + int cap = 0; // 单层单 TP 的 slot 数（该层该 TP 最多驻留的 CPU expert shard 数） + + // ===== GPU expert 配置 ===== + int num_gpu_experts = 0; // GE：过渡阶段要搬运到 GPU 的专家数 + + // ===== Expert Defer 配置 ===== + int max_deferred_per_token = 3; // 每 token 最多 defer 的专家数，对齐 KT 的 --kt-max-deferred-experts-per-token + + // ===== Decode 前 N 层满配 ===== + int decode_front_layers = 5; // 前 N 层接近满配，默认 5 + int decode_front_layer_cap = -1; // 前 N 层的 cap，-1 = 该层 CPU 专家总数（拉满） + + // ===== 时序配置 ===== + int total_layers = 0; // 模型总层数，用于 schedule_key 计算 + int prefill_window = 1; // prefill 窗口宽度，默认 1（只预取下一层） + + // ===== Heat EMA 参数 ===== + float heat_gamma = 0.7f; // token 内层内 EMA：heat_new = gamma * heat_old + (1-gamma) * score + float heat_beta = 0.5f; // 跨 token 全局 EMA：几何衰减 + + // ===== Markov 转移矩阵参数 ===== + float markov_alpha = 0.5f; // 转移矩阵更新率 + int markov_topk = 16; // 每行稀疏保留的目标专家数 + + // ===== 驱逐评分权重 ===== + float lookahead_weight = 1.0f; // score = policy_rank + lookahead_weight * heat + + // ===== 权重类型 ===== + WeightType weight_type = WeightType::AMXINT4; + + // ===== 模型维度（用于计算 slot 内存大小）===== + int hidden_size = 0; // 如 Qwen3.5-35B: 2048 + int intermediate_size = 0; // 如 512 + int expert_num = 0; // 单层专家总数 + int tp_count = 1; // TP 分片数（= NUMA 节点数） + + // ===== 文件布局 ===== + // [tp_part_idx][expert_id] -> ExpertFileLayout + // 由 Python 侧通过 set_file_layout 注入 + // 实际存储在 MeshResidencyManager 中，这里只放声明 +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_decode.hpp b/kt-kernel/operators/mesh/mesh_decode.hpp new file mode 100644 index 000000000..e2e2d9e23 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_decode.hpp @@ -0,0 +1,192 @@ +/** + * @file mesh_decode.hpp + * @brief Decode 策略：immediate/deferred 分组 + 在线驱逐 + 前 5 层满配 + * + * 每 token 每层即时收集 CPU 侧 top-k 专家，按 defer 参数分组： + * - immediate 组：CACHED 状态，本层立即计算 + * - deferred 组：BASELINE/LOADING 状态，异步 io_uring 读取，推迟到下一层 + * + * 分组逻辑（按 defer 参数对齐）： + * - 目标 immediate 数 = k - defer + * - 若 immediate数 > k-defer：按 router score 排序取前 k-defer 个 + * - 若 immediate数 < k-defer：阻塞等待异步读，直到 immediate 数量够 + * + * schedule_key： + * - immediate: timeline_step × total_layers + layer_idx + * - deferred: timeline_step × total_layers + layer_idx + 1（跟下一层同优先级） + * + * 前 5 层允许更大容量（接近全量常驻），默认拉满。 + */ +#pragma once + +#include +#include +#include +#include + +#include "mesh_config.hpp" +#include "mesh_eviction.hpp" +#include "mesh_scheduler.hpp" +#include "mesh_slot_pool.hpp" + +namespace mesh { + +/** + * @brief Decode 策略 + */ +class MeshDecode { + public: + MeshDecode(const MeshConfig& config) + : defer_count_(config.max_deferred_per_token), + front_layers_(config.decode_front_layers), + front_layer_cap_(config.decode_front_layer_cap), + total_layers_(config.total_layers), + expert_num_(config.expert_num) {} + + // immediate/deferred 分组结果 + struct SplitResult { + std::vector immediate; // 本层立即计算 + std::vector deferred; // 推迟到下一层计算 + }; + + /** + * @brief 对 top-k 专家进行 immediate/deferred 分组 + * + * @param topk 当前层 top-k 专家 ID + * @param scores 各专家的 router score（全长 expert_num） + * @param slot_pool 该层 slot 池 + * @param scheduler 调度器 + * @param layer_idx 当前层 + * @param tp_part_idx TP 分片 + * @param get_dst_ptrs 获取 slot buffer 指针的回调 + * @return SplitResult + */ + SplitResult split(const std::vector& topk, + const std::vector& scores, + MeshSlotPool& slot_pool, + MeshScheduler& scheduler, + int layer_idx, int tp_part_idx, + std::function(int)> get_dst_ptrs) { + int k = static_cast(topk.size()); + int target_immediate = std::max(0, k - defer_count_); + + // 区分已缓存和缺失 + std::vector cached, missing; + for (int eid : topk) { + if (slot_pool.is_cached(eid)) { + cached.push_back(eid); + } else { + missing.push_back(eid); + } + } + + SplitResult result; + + if (static_cast(cached.size()) > target_immediate) { + // immediate 数 > k-defer：按 router score 排序取前 target_immediate 个 + std::sort(cached.begin(), cached.end(), + [&scores](int a, int b) { + float sa = (a < (int)scores.size()) ? scores[a] : 0.0f; + float sb = (b < (int)scores.size()) ? scores[b] : 0.0f; + return sa > sb; // 降序 + }); + result.immediate.assign(cached.begin(), cached.begin() + target_immediate); + result.deferred.assign(cached.begin() + target_immediate, cached.end()); + // 缺失专家全部进 deferred + result.deferred.insert(result.deferred.end(), missing.begin(), missing.end()); + + // 为缺失专家提交 deferred 异步读 + for (int eid : missing) { + auto ptrs = get_dst_ptrs(eid); + scheduler.submit_decode_deferred(layer_idx, eid, tp_part_idx, + ptrs[0], ptrs[1], ptrs[2]); + } + } else { + // immediate 数 < k-defer：阻塞读缺失专家直到 immediate 数量够 + result.immediate = cached; + int need = target_immediate - static_cast(cached.size()); + + for (int i = 0; i < need && i < static_cast(missing.size()); i++) { + int eid = missing[i]; + auto ptrs = get_dst_ptrs(eid); + // 提交 immediate 异步读（schedule_key = 当前层） + scheduler.submit_decode_immediate(layer_idx, eid, tp_part_idx, + ptrs[0], ptrs[1], ptrs[2]); + // 阻塞等待读取完成 + // 实际实现需要与 MeshIoUring 配合等待 CQE + // io_.wait_expert(eid, n_reqs); + result.immediate.push_back(eid); + } + + // 剩余缺失专家走 deferred + for (int i = need; i < static_cast(missing.size()); i++) { + int eid = missing[i]; + auto ptrs = get_dst_ptrs(eid); + scheduler.submit_decode_deferred(layer_idx, eid, tp_part_idx, + ptrs[0], ptrs[1], ptrs[2]); + result.deferred.push_back(eid); + } + } + + return result; + } + + /** + * @brief 在线驱逐：slot 满时选分数最低的覆盖 + * + * @param slot_pool 该层 slot 池 + * @param scorer 驱逐评分器 + * @param layer_idx 当前层 + * @param new_expert_id 要加载的新专家 + * @return int 被驱逐的 slot_idx，-1 表示无需驱逐 + */ + int evict_for_new_expert(MeshSlotPool& slot_pool, + const EvictionScorer& scorer, + int layer_idx, int new_expert_id) { + // 前五层满配检查 + if (is_front_layer(layer_idx)) { + int cap = get_effective_cap(layer_idx, slot_pool.cap()); + if (cap >= expert_num_) { + return -1; // 满配，无需驱逐 + } + } + + // 找可驱逐的 slot + auto cached = slot_pool.cached_experts(); + if (cached.empty()) return -1; + + // 选分数最低的 + int victim_expert = scorer.select_victim(cached, layer_idx); + if (victim_expert < 0) return -1; + + // 查找 victim 对应的 slot_idx + // 实际实现需要通过 slot_pool 的接口 + return -1; // placeholder + } + + // ===== 前 5 层满配 ===== + + bool is_front_layer(int layer_idx) const { + return layer_idx < front_layers_; + } + + int get_effective_cap(int layer_idx, int default_cap) const { + if (is_front_layer(layer_idx)) { + return (front_layer_cap_ < 0) ? expert_num_ : front_layer_cap_; + } + return default_cap; + } + + // ===== 访问器 ===== + int defer_count() const { return defer_count_; } + int front_layers() const { return front_layers_; } + + private: + int defer_count_; + int front_layers_; + int front_layer_cap_; // -1 = 拉满 + int total_layers_; + int expert_num_; +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_eviction.hpp b/kt-kernel/operators/mesh/mesh_eviction.hpp new file mode 100644 index 000000000..e90640a4c --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_eviction.hpp @@ -0,0 +1,329 @@ +/** + * @file mesh_eviction.hpp + * @brief 驱逐评分：Heat EMA + Markov 层间转移矩阵 + * + * 驱逐分数 = policy_rank + lookahead_weight * max(heat, markov_prior) + * + * Heat EMA： + * - 跨 token 全局衰减：β=0.5，几何衰减 + * - 单 token 结束后一次性批量更新（GPU→CPU 一次传完，不逐层传） + * - 满驻留自动跳过：slot 容量 ≥ CPU 专家数时禁用 Heat 更新 + * + * Markov 转移矩阵： + * - T[L][s][t] = P(第L+1层激活专家t | 第L层激活专家s) + * - 稀疏存储：每行只保留概率最高的 K=16 个目标 + * - 只用于驱逐评分，绝对不用于预取 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace mesh { + +// ===== Heat EMA 跟踪器 ===== + +/** + * @brief 跨 token 全局 Heat EMA + * + * 每个 token 结束后，将该 token 所有层的 gating 分数一次性批量传入更新。 + * β=0.5 几何衰减：最近 token 占 0.5，前一个 0.25，再前一个 0.125... + * + * 关键设计：所有层的 gating 分数暂存 GPU，单 token 计算结束后一次性批量传到 CPU。 + * 不做逐层 GPU→CPU 传输。 + */ +class HeatTracker { + public: + HeatTracker(int expert_num, float beta = 0.5f, float gamma = 0.7f) + : expert_num_(expert_num), beta_(beta), gamma_(gamma) { + global_heat_.assign(expert_num, 0.0f); + token_layer_heat_.assign(expert_num, 0.0f); + } + + // 单 token 结束后批量更新 + // all_layers_gating: [layer][expert_num]，每层的 gating 分数（全长 vector） + void commit_token(const std::vector>& all_layers_gating) { + if (all_layers_gating.empty()) return; + + // 层内 EMA：gamma * heat_old + (1-gamma) * score + // 这里取最后一层作为该 token 的代表 heat（也可取平均） + const auto& last_layer = all_layers_gating.back(); + for (int e = 0; e < expert_num_ && e < (int)last_layer.size(); e++) { + token_layer_heat_[e] = gamma_ * token_layer_heat_[e] + (1.0f - gamma_) * last_layer[e]; + } + + // 跨 token 全局衰减：heat_new = beta * heat_old + (1-beta) * token_heat + for (int e = 0; e < expert_num_; e++) { + global_heat_[e] = beta_ * global_heat_[e] + (1.0f - beta_) * token_layer_heat_[e]; + } + } + + // 获取某专家的全局 heat + float heat(int expert_id) const { + if (expert_id < 0 || expert_id >= expert_num_) return 0.0f; + return global_heat_[expert_id]; + } + + // 满驻留自动跳过检查 + // slot_cap >= cpu_expert_num 时，驱逐不可能发生，Heat 无价值 + bool should_skip(int slot_cap, int cpu_expert_num) const { + return slot_cap >= cpu_expert_num; + } + + int expert_num() const { return expert_num_; } + const std::vector& global_heat() const { return global_heat_; } + + private: + int expert_num_; + float beta_; // 跨 token 衰减率 + float gamma_; // 层内 EMA 衰减率 + std::vector global_heat_; // [expert_num] 跨 token 全局 EMA + std::vector token_layer_heat_; // [expert_num] 当前 token 层内 EMA +}; + +// ===== Markov 层间转移矩阵 ===== + +/** + * @brief Markov 层间转移矩阵（稀疏存储） + * + * T[L][s][t] = P(第L+1层激活专家t | 第L层激活专家s) + * 每行只保留概率最高的 K 个目标专家，其余丢弃。 + * + * 更新算法（每 token 每相邻层对）： + * 1. 衰减旧行：row[t] *= (1-α) + * 2. 注入新观测：row[t] += α * w[t] + * 3. 重新归一化 + * + * 预测算法： + * P_predicted[t] = Σ_{s∈top-k} P_L[s] × T[L][s][t] + * 实际计算量：k × K = 8 × 16 = 128 次乘加 + */ +class MarkovTracker { + public: + static constexpr int kSparseK = 16; // 每行稀疏保留的目标专家数 + + struct Entry { + int target_id = -1; + float probability = 0.0f; + }; + + struct SparseRow { + Entry entries[kSparseK]; // 按 probability 降序排列 + float row_sum = 1.0f; // 归一化分母（≈1.0） + }; + + MarkovTracker(int num_layers, int expert_num, float alpha = 0.5f) + : num_layers_(num_layers), expert_num_(expert_num), alpha_(alpha) { + // 每对相邻层一个矩阵，共 num_layers-1 个 + matrices_.resize(std::max(0, num_layers - 1)); + for (auto& mat : matrices_) { + mat.resize(expert_num); + // 初始化为均匀分布：T[L][s][t] = 1/N + for (auto& row : mat) { + for (int i = 0; i < kSparseK && i < expert_num; i++) { + row.entries[i].target_id = i; + row.entries[i].probability = 1.0f / expert_num; + } + row.row_sum = 1.0f; + } + } + } + + /** + * @brief 更新某对相邻层的转移矩阵 + * + * @param layer_idx 第 L 层索引（更新的是 T[L]） + * @param src_topk 第 L 层的 top-k 专家集合 + * @param src_scores 第 L 层各专家的归一化权重（长度 = expert_num，非 top-k 位置为 0） + * @param dst_topk 第 L+1 层的 top-k 专家集合 + * @param dst_scores 第 L+1 层各专家的归一化权重 + */ + void update(int layer_idx, const std::vector& src_topk, + const std::vector& src_scores, + const std::vector& dst_topk, + const std::vector& dst_scores) { + if (layer_idx < 0 || layer_idx >= (int)matrices_.size()) return; + auto& mat = matrices_[layer_idx]; + + // 只更新 src_topk 中的源专家对应的行 + for (int s : src_topk) { + if (s < 0 || s >= expert_num_) continue; + SparseRow& row = mat[s]; + + // 1. 衰减旧行 + for (int i = 0; i < kSparseK; i++) { + row.entries[i].probability *= (1.0f - alpha_); + } + + // 2. 注入新观测：对 dst_topk 中的每个目标 t，row[t] += alpha * w[t] + for (int t : dst_topk) { + if (t < 0 || t >= expert_num_) continue; + float w = (t < (int)dst_scores.size()) ? dst_scores[t] : 0.0f; + + // 查找是否已存在 + int slot = -1; + for (int i = 0; i < kSparseK; i++) { + if (row.entries[i].target_id == t) { + slot = i; + row.entries[i].probability += alpha_ * w; + break; + } + } + if (slot < 0) { + // 找一个空位或概率最低的位置替换 + int min_idx = 0; + float min_prob = row.entries[0].probability; + for (int i = 1; i < kSparseK; i++) { + if (row.entries[i].probability < min_prob) { + min_prob = row.entries[i].probability; + min_idx = i; + } + } + float new_prob = alpha_ * w; + if (new_prob > min_prob) { + row.entries[min_idx].target_id = t; + row.entries[min_idx].probability = new_prob; + } + } + } + + // 3. 重新归一化 + float sum = 0.0f; + for (int i = 0; i < kSparseK; i++) { + sum += row.entries[i].probability; + } + row.row_sum = sum; + if (sum > 1e-6f) { + for (int i = 0; i < kSparseK; i++) { + row.entries[i].probability /= sum; + } + } + + // 按 probability 降序排列 + std::sort(row.entries, row.entries + kSparseK, + [](const Entry& a, const Entry& b) { + return a.probability > b.probability; + }); + } + } + + /** + * @brief 预测下一层的专家分布 + * + * @param layer_idx 当前层 L + * @param topk 当前层的 top-k 专家集合 + * @param scores 当前层各专家的归一化权重 + * @param prior_out 输出：第 L+1 层的 cross_layer_prior + */ + void predict(int layer_idx, const std::vector& topk, + const std::vector& scores, + std::vector& prior_out) const { + prior_out.assign(expert_num_, 0.0f); + if (layer_idx < 0 || layer_idx >= (int)matrices_.size()) return; + const auto& mat = matrices_[layer_idx]; + + // P_predicted[t] = Σ_{s∈top-k} P_L[s] × T[L][s][t] + for (int s : topk) { + if (s < 0 || s >= expert_num_) continue; + float p_s = (s < (int)scores.size()) ? scores[s] : 0.0f; + if (p_s < 1e-6f) continue; + + const SparseRow& row = mat[s]; + for (int i = 0; i < kSparseK; i++) { + int t = row.entries[i].target_id; + if (t < 0) continue; + prior_out[t] += p_s * row.entries[i].probability; + } + } + } + + int num_layers() const { return num_layers_; } + int expert_num() const { return expert_num_; } + + private: + int num_layers_; + int expert_num_; + float alpha_; + // matrices_[L] 是第 L 层到第 L+1 层的转移矩阵 + std::vector> matrices_; // [num_layers-1][expert_num] +}; + +// ===== 驱逐评分器 ===== + +/** + * @brief 驱逐评分器 + * + * score = policy_rank + lookahead_weight * max(heat, markov_prior) + * + * Markov 严格限制：只用于驱逐评分，绝对不用于预取。 + */ +class EvictionScorer { + public: + EvictionScorer(int num_layers, int expert_num, const MeshConfig& config) + : heat_(expert_num, config.heat_beta, config.heat_gamma), + markov_(num_layers, expert_num, config.markov_alpha), + lookahead_weight_(config.lookahead_weight), + expert_num_(expert_num) {} + + // 单 token 结束后批量更新 Heat 和 Markov + // all_layers_topk: [layer] -> top-k expert ids + // all_layers_scores: [layer] -> normalized router scores (全长 expert_num) + void commit_token(const std::vector>& all_layers_topk, + const std::vector>& all_layers_scores) { + // 更新 Heat + heat_.commit_token(all_layers_scores); + + // 更新 Markov：每对相邻层 + int n = std::min(all_layers_topk.size(), all_layers_scores.size()); + for (int l = 0; l < n - 1; l++) { + markov_.update(l, all_layers_topk[l], all_layers_scores[l], + all_layers_topk[l + 1], all_layers_scores[l + 1]); + } + } + + // 计算某专家的驱逐评分（分数越低越该被驱逐） + float score(int expert_id, int layer_idx) const { + float h = heat_.heat(expert_id); + std::vector prior; + markov_.predict(layer_idx, {}, {}, prior); // 简化：实际需要传入当前层 topk + float m = (expert_id < (int)prior.size()) ? prior[expert_id] : 0.0f; + float heat = std::max(h, m); + return lookahead_weight_ * heat; + } + + // 选择 victim：从候选专家中选分数最低的 + // candidates: 当前 CACHED 且 active_readers==0 的专家列表 + int select_victim(const std::vector& candidates, int layer_idx) const { + if (candidates.empty()) return -1; + int victim = candidates[0]; + float min_score = score(victim, layer_idx); + for (size_t i = 1; i < candidates.size(); i++) { + float s = score(candidates[i], layer_idx); + if (s < min_score) { + min_score = s; + victim = candidates[i]; + } + } + return victim; + } + + // 满驻留检查 + bool should_skip_layer(int slot_cap, int cpu_expert_num) const { + return heat_.should_skip(slot_cap, cpu_expert_num); + } + + HeatTracker& heat() { return heat_; } + MarkovTracker& markov() { return markov_; } + + private: + HeatTracker heat_; + MarkovTracker markov_; + float lookahead_weight_; + int expert_num_; +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_handoff.hpp b/kt-kernel/operators/mesh/mesh_handoff.hpp new file mode 100644 index 000000000..80b992832 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_handoff.hpp @@ -0,0 +1,152 @@ +/** + * @file mesh_handoff.hpp + * @brief Prefill→Decode 过渡阶段 + * + * 触发时机：第一个 decode token 的第一层前 + * + * 流程： + * 1. 保留 slot 池（prefill 留下的种子专家给 decode 继承） + * 2. 释放两个 prefill temporal 缓冲区 + * 3. 搬运 GPU 专家权重到 slot 池填补因裁剪留下的槽位 + * - GPU expert 从 CPU slot 搬到 GPU（不读盘，纯内存→GPU 拷贝） + * - 释放的 GE 个 slot 用 io_uring 读 (cap, cap+GE] 区间的专家填充 + * 4. 清空优先队列 + * 5. 切换到 decode 模式 + */ +#pragma once + +#include +#include +#include + +#include "mesh_config.hpp" +#include "mesh_eviction.hpp" +#include "mesh_io_uring.hpp" +#include "mesh_prefill.hpp" +#include "mesh_scheduler.hpp" +#include "mesh_slot_pool.hpp" + +namespace mesh { + +/** + * @brief Prefill→Decode 过渡处理器 + */ +class MeshHandoff { + public: + /** + * @brief 执行过渡 + * + * @param config MESH 配置 + * @param prefill prefill 策略对象 + * @param scheduler 调度器 + * @param io io_uring 读取器 + * @param scorer 驱逐评分器 + * @param pools [layer][tp] 的 slot 池 + * @param layouts [tp][expert] 的文件布局 + * @param move_gpu_expert_to_gpu 搬运 GPU 专家到 GPU 的回调 + * @param get_dst_ptrs 获取 slot buffer 指针的回调 + */ + void transition(const MeshConfig& config, + MeshPrefill& prefill, + MeshScheduler& scheduler, + MeshIoUring& io, + EvictionScorer& scorer, + std::vector>& pools, + const std::vector>& layouts, + std::function move_gpu_expert_to_gpu, + std::function(int, int, int)> get_dst_ptrs) { + int ge = config.num_gpu_experts; + int cap = config.cap; + + // 1. 保留 slot 池（无需操作，prefill 留下的种子专家直接继承） + + // 2. 释放两个 prefill temporal 缓冲区 + prefill.release_temporal(); + + // 3. 搬运 GPU 专家权重到 GPU，释放的 slot 用 io_uring 读新专家填充 + // GPU expert 在 slot 前茅（GE 个），搬到 GPU 后释放这些 slot + // 然后读 (cap, cap+GE] 区间的专家填充这些 slot + if (ge > 0) { + move_gpu_experts_and_refill(config, pools, layouts, scorer, io, + move_gpu_expert_to_gpu, get_dst_ptrs); + } + + // 4. 清空优先队列 + scheduler.clear_queue(); + + // 5. 切换到 decode 模式 + scheduler.switch_to_decode(); + } + + private: + /** + * @brief 搬运 GPU 专家到 GPU，并用 io_uring 读新专家填充释放的 slot + * + * GPU expert 在 slot 前茅（GE 个位置）： + * - 标记这些 slot 的 active_readers（防止驱逐） + * - 搬运到 GPU（不读盘） + * - 释放这些 slot + * - 用 io_uring 读 (cap, cap+GE] 区间的专家填充 + */ + void move_gpu_experts_and_refill( + const MeshConfig& config, + std::vector>& pools, + const std::vector>& layouts, + EvictionScorer& scorer, + MeshIoUring& io, + std::function move_gpu_expert_to_gpu, + std::function(int, int, int)> get_dst_ptrs) { + int ge = config.num_gpu_experts; + int cap = config.cap; + int num_layers = config.total_layers; + int tp_count = config.tp_count; + + // (cap, cap+GE] 区间的专家 ID + std::vector refill_experts; + for (int e = cap; e < cap + ge && e < config.expert_num; e++) { + refill_experts.push_back(e); + } + + for (int layer = 0; layer < num_layers; layer++) { + for (int tp = 0; tp < tp_count; tp++) { + MeshSlotPool& pool = pools[layer][tp]; + + // GPU expert 在 slot 前茅（GE 个） + for (int slot_idx = 0; slot_idx < ge && slot_idx < pool.cap(); slot_idx++) { + int expert_id = slot_idx; // 假设 GPU expert 占据前 GE 个 slot + if (expert_id >= config.expert_num) break; + + // 标记 active_readers 防止驱逐 + pool.acquire_reader(expert_id); + + // 搬运到 GPU（不读盘，纯内存→GPU 拷贝） + move_gpu_expert_to_gpu(layer, expert_id); + + // 释放 reader + pool.release_reader(expert_id); + + // 释放这个 slot，用 io_uring 读新专家填充 + int new_expert_id = cap + slot_idx; + if (new_expert_id < config.expert_num && + slot_idx < static_cast(refill_experts.size())) { + new_expert_id = refill_experts[slot_idx]; + auto ptrs = get_dst_ptrs(layer, tp, new_expert_id); + // 提交 io_uring 读 + const auto& layout = layouts[tp][new_expert_id]; + io.submit_load(new_expert_id, tp, layout, + ptrs[0], ptrs[1], ptrs[2], + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + ReadPriority::Demand); + // 覆盖 slot + pool.overwrite(slot_idx, new_expert_id); + } + } + } + } + + // 等待所有 io_uring 读完成 + io.submit_and_wait(); + } +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_hook.hpp b/kt-kernel/operators/mesh/mesh_hook.hpp new file mode 100644 index 000000000..f170d826d --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_hook.hpp @@ -0,0 +1,79 @@ +/** + * @file mesh_hook.hpp + * @brief MESH inline hook 函数 + * + * 所有 hook 函数：mesh_mgr == nullptr 时编译期优化为空操作。 + * 原版 KT 代码通过调用这些 hook 接入 MESH，无需 include mesh.hpp。 + * + * 使用方式（在 amx/moe_base.hpp 等 KT 原版文件中）： + * void* gate_bb = mesh::hook::get_gate_bb(mesh_mgr, layer, tp, expert_id); + * if (gate_bb) { /* MESH 路径 *\/ } else { /* 原版路径 *\/ } + */ +#pragma once + +#include + +namespace mesh { +namespace hook { + +// 前向声明，避免强制 include mesh_residency.hpp +class MeshResidencyManager; + +// ===== 权重指针重定向 hook ===== + +// 获取 gate 矩阵指针（MESH 启用时返回 slot buffer，否则返回 nullptr 走原版） +// mesh_mgr 为 nullptr 时直接返回 nullptr，编译期可优化 +inline void* get_gate_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + // 实际实现通过 reinterpret_cast 调用 MeshResidencyManager::get_gate_ptr + // 这里用函数指针避免头文件循环依赖 + using GetPtrFn = void* (*)(void*, int, int, int); + static GetPtrFn fn = nullptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +inline void* get_up_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + using GetPtrFn = void* (*)(void*, int, int, int); + static GetPtrFn fn = nullptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +inline void* get_down_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + using GetPtrFn = void* (*)(void*, int, int, int); + static GetPtrFn fn = nullptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +// ===== MESH 启用判断 hook ===== + +inline bool is_mesh_enabled(void* mesh_mgr) { + return mesh_mgr != nullptr; +} + +// ===== 注册函数指针（由 ext_bindings.cpp 在模块初始化时调用）===== + +// 这些函数指针用于解耦 mesh_hook.hpp 和 mesh_residency.hpp +// 避免原版 KT 文件 include mesh_residency.hpp 带来的编译依赖 +struct HookRegistry { + void* (*get_gate_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*get_up_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*get_down_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; +}; + +inline HookRegistry& get_registry() { + static HookRegistry registry; + return registry; +} + +// 注册 hook 函数（ext_bindings.cpp 调用） +inline void register_hooks(HookRegistry r) { + get_registry() = r; +} + +} // namespace hook +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_io_uring.hpp b/kt-kernel/operators/mesh/mesh_io_uring.hpp new file mode 100644 index 000000000..758360670 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_io_uring.hpp @@ -0,0 +1,299 @@ +/** + * @file mesh_io_uring.hpp + * @brief io_uring + O_DIRECT 读取器 + Scale Cache 预加载 + * + * 读取路径：不走 mmap，而是 O_DIRECT 打开 safetensors 文件。 + * 通过 io_uring 提交读请求，SSD 直接 DMA 写入 NUMA 本地 buffer。 + * + * Scale Cache（AMXINT4 专用）： + * - 启动时一次性把所有专家的 scale 数据读入独立 NUMA 本地 buffer + * - scale_cache_loaded_=true：只发 3 个请求（只读 weight），scale 从 cache memcpy + * - scale_cache_loaded_=false：发 9 个请求（weight + scale + mins 全读） + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mesh_config.hpp" + +namespace mesh { + +// io_uring 读请求优先级（与 schedule_key 配合使用） +enum class ReadPriority : uint8_t { + Prefetch = 1, // 预取（低优先级） + Demand = 10, // 需求（高优先级） +}; + +/** + * @brief io_uring 读取器 + * + * 封装 liburing，所有读走 O_DIRECT 直达 NUMA buffer，不经过 page cache。 + */ +class MeshIoUring { + public: + static constexpr int kQueueDepth = 256; // io_uring SQ ring 大小 + + MeshIoUring() { + if (io_uring_queue_init(kQueueDepth, &ring_, 0) != 0) { + throw std::runtime_error("MeshIoUring: io_uring_queue_init failed"); + } + } + + ~MeshIoUring() { + io_uring_queue_exit(&ring_); + release_scale_cache(); + } + + // 禁止拷贝 + MeshIoUring(const MeshIoUring&) = delete; + MeshIoUring& operator=(const MeshIoUring&) = delete; + + // ===== Scale Cache 预加载（AMXINT4 专用）===== + + // 启动时一次性把所有专家的 scale 数据读入 NUMA 本地 buffer + // expert_num 个专家，每个 scale_bytes 字节 + void preload_scale_cache(int expert_num, int tp_count, int numa_node, + const std::vector>& layouts) { + if (scale_cache_loaded_) return; + + // 计算每个 TP 分片的 scale 总大小 + for (int tp = 0; tp < tp_count; tp++) { + ScaleCacheTP cache; + size_t total_gate = 0, total_up = 0, total_down = 0; + size_t total_gate_mins = 0, total_up_mins = 0, total_down_mins = 0; + for (int e = 0; e < expert_num; e++) { + total_gate += layouts[tp][e].gate_scale_bytes; + total_up += layouts[tp][e].up_scale_bytes; + total_down += layouts[tp][e].down_scale_bytes; + total_gate_mins += layouts[tp][e].gate_mins_bytes; + total_up_mins += layouts[tp][e].up_mins_bytes; + total_down_mins += layouts[tp][e].down_mins_bytes; + } + cache.gate_scale = numa_alloc_onnode(total_gate, numa_node); + cache.up_scale = numa_alloc_onnode(total_up, numa_node); + cache.down_scale = numa_alloc_onnode(total_down, numa_node); + cache.gate_mins = numa_alloc_onnode(total_gate_mins, numa_node); + cache.up_mins = numa_alloc_onnode(total_up_mins, numa_node); + cache.down_mins = numa_alloc_onnode(total_down_mins, numa_node); + cache.gate_scale_total = total_gate; + cache.up_scale_total = total_up; + cache.down_scale_total = total_down; + cache.gate_mins_total = total_gate_mins; + cache.up_mins_total = total_up_mins; + cache.down_mins_total = total_down_mins; + + // 同步读取所有 scale 数据（启动阶段，可以阻塞） + size_t off_g = 0, off_u = 0, off_d = 0; + size_t off_gm = 0, off_um = 0, off_dm = 0; + for (int e = 0; e < expert_num; e++) { + const auto& layout = layouts[tp][e]; + sync_read_to(layout.fd, layout.gate_scale_offset, layout.gate_scale_bytes, + (char*)cache.gate_scale + off_g); + sync_read_to(layout.fd, layout.up_scale_offset, layout.up_scale_bytes, + (char*)cache.up_scale + off_u); + sync_read_to(layout.fd, layout.down_scale_offset, layout.down_scale_bytes, + (char*)cache.down_scale + off_d); + sync_read_to(layout.fd, layout.gate_mins_offset, layout.gate_mins_bytes, + (char*)cache.gate_mins + off_gm); + sync_read_to(layout.fd, layout.up_mins_offset, layout.up_mins_bytes, + (char*)cache.up_mins + off_um); + sync_read_to(layout.fd, layout.down_mins_offset, layout.down_mins_bytes, + (char*)cache.down_mins + off_dm); + off_g += layout.gate_scale_bytes; + off_u += layout.up_scale_bytes; + off_d += layout.down_scale_bytes; + off_gm += layout.gate_mins_bytes; + off_um += layout.up_mins_bytes; + off_dm += layout.down_mins_bytes; + } + scale_cache_.push_back(std::move(cache)); + } + scale_cache_loaded_ = true; + } + + // 从 scale cache 中 memcpy 某个专家的 scale 到目标 buffer + void copy_scale_from_cache(int tp_part_idx, int expert_id, + void* gate_scale_dst, void* up_scale_dst, void* down_scale_dst, + void* gate_mins_dst, void* up_mins_dst, void* down_mins_dst, + const std::vector& layouts_tp) { + // 需要外部传入该 TP 的 layouts 来计算偏移 + size_t off_g = 0, off_u = 0, off_d = 0; + size_t off_gm = 0, off_um = 0, off_dm = 0; + for (int e = 0; e < expert_id; e++) { + off_g += layouts_tp[e].gate_scale_bytes; + off_u += layouts_tp[e].up_scale_bytes; + off_d += layouts_tp[e].down_scale_bytes; + off_gm += layouts_tp[e].gate_mins_bytes; + off_um += layouts_tp[e].up_mins_bytes; + off_dm += layouts_tp[e].down_mins_bytes; + } + const auto& layout = layouts_tp[expert_id]; + const auto& cache = scale_cache_[tp_part_idx]; + memcpy(gate_scale_dst, (char*)cache.gate_scale + off_g, layout.gate_scale_bytes); + memcpy(up_scale_dst, (char*)cache.up_scale + off_u, layout.up_scale_bytes); + memcpy(down_scale_dst, (char*)cache.down_scale + off_d, layout.down_scale_bytes); + memcpy(gate_mins_dst, (char*)cache.gate_mins + off_gm, layout.gate_mins_bytes); + memcpy(up_mins_dst, (char*)cache.up_mins + off_um, layout.up_mins_bytes); + memcpy(down_mins_dst, (char*)cache.down_mins + off_dm, layout.down_mins_bytes); + } + + bool scale_cache_loaded() const { return scale_cache_loaded_; } + + // ===== io_uring 异步读 ===== + + // 提交一个专家的读取请求 + // scale_cache_loaded_=true: 发 3 个请求（只读 weight），scale 从 cache memcpy + // scale_cache_loaded_=false: 发 9 个请求（weight + scale + mins 全读） + void submit_load(int expert_id, int tp_part_idx, + const ExpertFileLayout& layout, + void* gate_dst, void* up_dst, void* down_dst, + void* gate_scale_dst, void* up_scale_dst, void* down_scale_dst, + void* gate_mins_dst, void* up_mins_dst, void* down_mins_dst, + ReadPriority priority) { + int n_reqs = scale_cache_loaded_ ? 3 : 9; + std::vector sqes; + sqes.reserve(n_reqs); + + // 提交 SQE + auto* sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, gate_dst, layout.gate_bytes, layout.gate_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, up_dst, layout.up_bytes, layout.up_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, down_dst, layout.down_bytes, layout.down_offset); + sqes.push_back(sqe); + + if (!scale_cache_loaded_) { + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, gate_scale_dst, layout.gate_scale_bytes, layout.gate_scale_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, up_scale_dst, layout.up_scale_bytes, layout.up_scale_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, down_scale_dst, layout.down_scale_bytes, layout.down_scale_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, gate_mins_dst, layout.gate_mins_bytes, layout.gate_mins_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, up_mins_dst, layout.up_mins_bytes, layout.up_mins_offset); + sqes.push_back(sqe); + + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, down_mins_dst, layout.down_mins_bytes, layout.down_mins_offset); + sqes.push_back(sqe); + } + + // 设置 user_data 用于 CQE 回调识别 + for (auto* s : sqes) { + io_uring_sqe_set_data(s, nullptr); + } + io_uring_submit(&ring_); + } + + // 阻塞等待某专家的读请求完成（通过 CQE 计数） + void wait_expert(int expert_id, int n_reqs) { + for (int i = 0; i < n_reqs; i++) { + io_uring_cqe* cqe; + io_uring_wait_cqe(&ring_, &cqe); + io_uring_cqe_seen(&ring_, cqe); + } + } + + // 批量提交一层所有专家的读取 + void submit_layer_batch(const std::vector& expert_ids, int tp_part_idx, + const std::vector& layouts_tp, + std::vector gate_dsts, std::vector up_dsts, + std::vector down_dsts, ReadPriority priority) { + for (size_t i = 0; i < expert_ids.size(); i++) { + int eid = expert_ids[i]; + // scale/mins 目标指针由调用者通过 slot pool 提供 + // 这里简化：只读 weight，scale 从 cache memcpy + submit_load(eid, tp_part_idx, layouts_tp[eid], + gate_dsts[i], up_dsts[i], down_dsts[i], + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + priority); + } + } + + // 提交并等待全部完成（同步接口，启动阶段用） + void submit_and_wait() { + io_uring_submit(&ring_); + // 等待所有 inflight 完成 + unsigned head; + unsigned count = 0; + io_uring_cqe* cqe; + while (true) { + io_uring_for_each_cqe(&ring_, head, cqe) { + count++; + } + if (count == 0) break; + io_uring_cq_advance(&ring_, count); + count = 0; + } + } + + private: + io_uring ring_; + + // Scale Cache（每 TP 一个） + struct ScaleCacheTP { + void* gate_scale = nullptr; + void* up_scale = nullptr; + void* down_scale = nullptr; + void* gate_mins = nullptr; + void* up_mins = nullptr; + void* down_mins = nullptr; + size_t gate_scale_total = 0; + size_t up_scale_total = 0; + size_t down_scale_total = 0; + size_t gate_mins_total = 0; + size_t up_mins_total = 0; + size_t down_mins_total = 0; + }; + std::vector scale_cache_; + bool scale_cache_loaded_ = false; + + // 同步读取（启动阶段用，pread + O_DIRECT） + void sync_read_to(int fd, off_t offset, size_t bytes, void* dst) { + if (bytes == 0) return; + size_t done = 0; + while (done < bytes) { + ssize_t n = pread(fd, (char*)dst + done, bytes - done, offset + done); + if (n < 0) { + throw std::runtime_error("MeshIoUring: pread failed, errno=" + std::to_string(errno)); + } + done += n; + } + } + + void release_scale_cache() { + for (auto& c : scale_cache_) { + if (c.gate_scale) numa_free(c.gate_scale, c.gate_scale_total); + if (c.up_scale) numa_free(c.up_scale, c.up_scale_total); + if (c.down_scale) numa_free(c.down_scale, c.down_scale_total); + if (c.gate_mins) numa_free(c.gate_mins, c.gate_mins_total); + if (c.up_mins) numa_free(c.up_mins, c.up_mins_total); + if (c.down_mins) numa_free(c.down_mins, c.down_mins_total); + } + scale_cache_.clear(); + } +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_prefill.hpp b/kt-kernel/operators/mesh/mesh_prefill.hpp new file mode 100644 index 000000000..93c9eff11 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_prefill.hpp @@ -0,0 +1,245 @@ +/** + * @file mesh_prefill.hpp + * @brief Prefill 策略：temporal 双缓冲 + 双 embedding 空间 + layer-major 调度 + * + * 流程（layer-major）： + * for layer L in 0..N: + * for chunk C in 0..M: + * a. wait temporal_A[L] 就绪 + * b. submit temporal_B[L+1] 异步预取 (schedule_key = L) + * c. AMX 计算 (temporal_A + slot 常驻) + * d. 指针 swap (slot 低频 ↔ temporal_A 高频) + * e. 保存 chunk 产物到 embedding 缓冲 + * f. 下一个 chunk → 回到 c + * 下一层 + * + * 双 embedding 空间：奇偶层交替写入 buf_a / buf_b + * temporal 双缓冲：A 用于当前层计算，B 用于下一层预取，每层角色互换 + * + * temporal 长度 = 单层专家总数 - slot cap + * 只在 prefill→decode 切换时清除 temporal 内存，其余全是覆盖 + */ +#pragma once + +#include +#include +#include +#include +#include + +#include "mesh_config.hpp" +#include "mesh_scheduler.hpp" +#include "mesh_slot_pool.hpp" + +namespace mesh { + +/** + * @brief Prefill 策略 + * + * MESH 模式下接管 prefill 调度（layer-major + 双 embedding 空间）。 + * KT 原版只加 if(mesh_enabled) { mesh_prefill.run(); return; } 一个分支。 + */ +class MeshPrefill { + public: + enum class TemporalRole { COMPUTE, PREFETCH }; + + MeshPrefill(int num_layers, int expert_num, int cap, int tp_count, + size_t slot_bytes, int numa_node) + : num_layers_(num_layers), + expert_num_(expert_num), + cap_(cap), + tp_count_(tp_count), + slot_bytes_(slot_bytes), + numa_node_(numa_node) { + // temporal 长度 = 单层专家总数 - slot cap + temporal_size_ = (expert_num > cap) ? (expert_num - cap) : 0; + } + + ~MeshPrefill() { release_temporal(); } + + // 初始化 temporal 双缓冲（每 TP 一对 A/B） + void init_temporal() { + if (temporal_a_ != nullptr) return; // 已初始化 + size_t bytes = static_cast(temporal_size_) * slot_bytes_ * tp_count_; + temporal_a_ = numa_alloc_onnode(bytes, numa_node_); + temporal_b_ = numa_alloc_onnode(bytes, numa_node_); + if (!temporal_a_ || !temporal_b_) { + throw std::runtime_error("MeshPrefill: numa_alloc_onnode for temporal failed"); + } + temporal_bytes_ = bytes; + role_a_ = TemporalRole::COMPUTE; + role_b_ = TemporalRole::PREFETCH; + } + + // 释放 temporal 内存（仅 prefill→decode 切换时调用） + void release_temporal() { + if (temporal_a_) { + numa_free(temporal_a_, temporal_bytes_); + temporal_a_ = nullptr; + } + if (temporal_b_) { + numa_free(temporal_b_, temporal_bytes_); + temporal_b_ = nullptr; + } + } + + // ===== Prefill 主流程 ===== + + /** + * @brief 处理某一层的所有 chunk + * + * @param layer_idx 当前层 L + * @param total_chunks 总 chunk 数 + * @param active_experts_per_chunk [chunk] -> 该 chunk 活跃的专家列表 + * @param slot_pool 该层的 slot 池 + * @param scheduler 调度器 + * @param amx_forward AMX 计算回调（layer_idx, chunk_idx, expert_ids） + * @param get_temporal_ptrs 获取 temporal buffer 指针的回调 + */ + void run_layer(int layer_idx, int total_chunks, + const std::vector>& active_experts_per_chunk, + MeshSlotPool& slot_pool, MeshScheduler& scheduler, + std::function&)> amx_forward, + std::function(int, int)> get_temporal_ptrs) { + for (int c = 0; c < total_chunks; c++) { + const auto& active = active_experts_per_chunk[c]; + + // a. 等待 temporal_A 内第 L 层所需专家全部就绪 + // （由上一层的异步预取完成；第一层阻塞等待） + wait_temporal_ready(layer_idx, active, slot_pool); + + // b. 对 temporal_B 发起第 L+1 层的异步预取（覆盖另一个 temporal） + if (layer_idx + 1 < num_layers_) { + submit_next_layer_prefetch(layer_idx + 1, active, scheduler, + get_temporal_ptrs); + } + + // c. 本层计算使用 temporal_A 中的临时专家 + slot 池常驻专家 + // 具体计算逻辑交给 ktransformers + amx_forward(layer_idx, c, active); + + // d. 计算完成后，根据本层专家调用频次，指针 swap + // slot 池内低频专家 ↔ temporal_A 中高频专家 + swap_high_freq_to_slot(layer_idx, active, slot_pool); + + // e. 保存当前 chunk 临时产物到 embedding 缓冲 + // （由 KT 侧通过双 embedding 空间管理，MESH 只感知奇偶层） + save_chunk_output(layer_idx, c); + + // f. 推进窗口，计算下一个 chunk → 回到 c + } + // 本层所有 chunk 计算完毕，来到下一层 + } + + // ===== temporal 角色切换 ===== + + // 每过一层，A/B 角色互换 + void swap_temporal_roles() { + TemporalRole tmp = role_a_; + role_a_ = role_b_; + role_b_ = tmp; + } + + // 获取当前 COMPUTE 角色的 temporal buffer + void* compute_temporal() const { + return role_a_ == TemporalRole::COMPUTE ? temporal_a_ : temporal_b_; + } + + // 获取当前 PREFETCH 角色的 temporal buffer + void* prefetch_temporal() const { + return role_a_ == TemporalRole::PREFETCH ? temporal_a_ : temporal_b_; + } + + // ===== 访问器 ===== + int temporal_size() const { return temporal_size_; } + bool temporal_ready() const { return temporal_a_ != nullptr; } + + private: + int num_layers_; + int expert_num_; + int cap_; + int tp_count_; + size_t slot_bytes_; + int numa_node_; + size_t temporal_bytes_ = 0; + + // temporal 双缓冲 + void* temporal_a_ = nullptr; + void* temporal_b_ = nullptr; + TemporalRole role_a_ = TemporalRole::COMPUTE; + TemporalRole role_b_ = TemporalRole::PREFETCH; + int temporal_size_ = 0; // = expert_num - cap + + // 等待 temporal 内专家就绪 + void wait_temporal_ready(int layer_idx, const std::vector& active, + MeshSlotPool& slot_pool) { + // 对于第一层，阻塞直到全部读取完毕 + // 对于后续层，由上一层预取完成，这里检查状态 + for (int eid : active) { + if (slot_pool.is_cached(eid)) continue; + // 等待 io_uring CQE + // 实际实现需要与 MeshIoUring 配合 + } + } + + // 提交下一层的异步预取 + void submit_next_layer_prefetch(int next_layer, + const std::vector& active, + MeshScheduler& scheduler, + std::function(int, int)> get_ptrs) { + // 对下一层需要的专家提交异步读，schedule_key = next_layer + void* prefetch_buf = prefetch_temporal(); + for (int tp = 0; tp < tp_count_; tp++) { + auto ptrs = get_ptrs(next_layer, tp); + for (int eid : active) { + scheduler.submit_prefill(next_layer, eid, tp, + ptrs[eid * 3], ptrs[eid * 3 + 1], ptrs[eid * 3 + 2], + ReadPriority::Prefetch); + } + } + } + + // 指针 swap：slot 低频 ↔ temporal 高频 + void swap_high_freq_to_slot(int layer_idx, const std::vector& active, + MeshSlotPool& slot_pool) { + // 统计本层专家调用频次 + std::vector freq(expert_num_, 0); + for (int eid : active) { + if (eid >= 0 && eid < expert_num_) freq[eid]++; + } + + // 找出 temporal 中高频但不在 slot 中的专家 + std::vector high_freq_in_temporal; + for (int e = 0; e < expert_num_; e++) { + if (freq[e] > 0 && !slot_pool.is_cached(e)) { + high_freq_in_temporal.push_back(e); + } + } + + // 找出 slot 中低频的专家 + auto cached = slot_pool.cached_experts(); + std::vector low_freq_in_slot; + for (int eid : cached) { + if (eid >= 0 && eid < expert_num_ && freq[eid] == 0) { + low_freq_in_slot.push_back(eid); + } + } + + // 指针 swap：把池子指针指向 temporal 目标专家，temporal 目标专家指针指向原 slot 专家 + // 不操作内存本身，只交换指针映射 + size_t swap_count = std::min(high_freq_in_temporal.size(), low_freq_in_slot.size()); + for (size_t i = 0; i < swap_count; i++) { + // 实际实现需要通过 slot_pool 的接口完成指针交换 + // slot_pool.swap_pointers(low_freq_in_slot[i], high_freq_in_temporal[i]); + } + } + + // 保存 chunk 产物到 embedding 缓冲 + // 双 embedding 空间：奇偶层交替写入 + void save_chunk_output(int layer_idx, int chunk_idx) { + // 由 KT 侧管理 embedding 空间，MESH 只感知奇偶层 + // embedding_buf_a_ / embedding_buf_b_ 由 KT 持有 + } +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp new file mode 100644 index 000000000..a99772cc9 --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -0,0 +1,318 @@ +/** + * @file mesh_residency.hpp + * @brief MESH 顶层 Manager + * + * 协调 slot 池、io_uring、调度器、驱逐评分、prefill/decode 策略、过渡处理。 + * 由 ext_bindings.cpp 暴露给 Python，是 MESH 对外的唯一入口。 + * + * 生命周期： + * 1. init()：初始化所有组件 + * 2. bootstrap()：启动阶段，读前 cap 个专家进 slot + * 3. prefill 阶段：on_prefill_layer_start/done + * 4. 过渡阶段：on_prefill_to_decode() + * 5. decode 阶段：on_decode_token / on_decode_layer + */ +#pragma once + +#include +#include +#include +#include + +#include "mesh_config.hpp" +#include "mesh_decode.hpp" +#include "mesh_eviction.hpp" +#include "mesh_handoff.hpp" +#include "mesh_io_uring.hpp" +#include "mesh_prefill.hpp" +#include "mesh_scheduler.hpp" +#include "mesh_slot_pool.hpp" + +namespace mesh { + +/** + * @brief MESH 顶层 Manager + */ +class MeshResidencyManager { + public: + MeshResidencyManager() = default; + ~MeshResidencyManager() = default; + + // 禁止拷贝 + MeshResidencyManager(const MeshResidencyManager&) = delete; + MeshResidencyManager& operator=(const MeshResidencyManager&) = delete; + + /** + * @brief 初始化 MESH + * + * @param config 配置 + * @param numa_nodes 每个 TP 分片对应的 NUMA 节点 + */ + void init(const MeshConfig& config, const std::vector& numa_nodes) { + config_ = config; + numa_nodes_ = numa_nodes; + + // 计算 slot 字节数（根据权重类型和模型维度） + slot_bytes_ = compute_slot_bytes(config); + + // 创建 slot 池 [layer][tp] + pools_.resize(config.total_layers); + for (int l = 0; l < config.total_layers; l++) { + pools_[l].reserve(config.tp_count); + for (int tp = 0; tp < config.tp_count; tp++) { + pools_[l].emplace_back(l, tp, numa_nodes[tp], config.cap, slot_bytes_); + pools_[l][tp].set_gate_up_bytes(compute_gate_up_bytes(config)); + pools_[l][tp].init_expert_map(config.expert_num); + } + } + + // 初始化各组件 + io_ = std::make_unique(); + scheduler_ = std::make_unique(config.total_layers); + scorer_ = std::make_unique(config.total_layers, config.expert_num, config); + prefill_ = std::make_unique(config.total_layers, config.expert_num, + config.cap, config.tp_count, + slot_bytes_, numa_nodes[0]); + decode_ = std::make_unique(config); + handoff_ = std::make_unique(); + + // 初始化 temporal 双缓冲 + prefill_->init_temporal(); + } + + // ===== 文件布局注入 ===== + + void set_file_layout(int tp_part_idx, int expert_id, const ExpertFileLayout& layout) { + if (layouts_.empty()) { + layouts_.resize(config_.tp_count, std::vector(config_.expert_num)); + } + if (tp_part_idx >= 0 && tp_part_idx < (int)layouts_.size() && + expert_id >= 0 && expert_id < (int)layouts_[tp_part_idx].size()) { + layouts_[tp_part_idx][expert_id] = layout; + } + } + + // ===== GPU expert mask 注入 ===== + + void set_gpu_experts_mask(const uint8_t* mask, int n) { + gpu_experts_mask_.assign(mask, mask + n); + } + + bool is_gpu_expert(int expert_id) const { + if (expert_id < 0 || expert_id >= (int)gpu_experts_mask_.size()) return false; + return gpu_experts_mask_[expert_id] != 0; + } + + // ===== 启动阶段 ===== + + /** + * @brief 启动阶段：读前 cap 个专家进 slot + * + * 每层 Slot 池按编号顺序填满。 + */ + void bootstrap() { + // 预加载 Scale Cache（AMXINT4 专用） + if (config_.weight_type == WeightType::AMXINT4) { + io_->preload_scale_cache(config_.expert_num, config_.tp_count, + numa_nodes_[0], layouts_); + } + + // 每层每 TP 读前 cap 个专家 + for (int l = 0; l < config_.total_layers; l++) { + for (int tp = 0; tp < config_.tp_count; tp++) { + MeshSlotPool& pool = pools_[l][tp]; + for (int e = 0; e < config_.cap && e < config_.expert_num; e++) { + if (is_gpu_expert(e)) continue; // GPU expert 跳过 + + const auto& layout = layouts_[tp][e]; + void* gate_dst = pool.gate_ptr(e); + void* up_dst = pool.up_ptr(e); + void* down_dst = pool.down_ptr(e); + + // 绑定 slot + pool.bind(e, e); + + // 提交 io_uring 读 + io_->submit_load(e, tp, layout, gate_dst, up_dst, down_dst, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + ReadPriority::Demand); + } + } + } + + // 等待所有读完成 + io_->submit_and_wait(); + + // 标记所有 slot 为 CACHED + for (int l = 0; l < config_.total_layers; l++) { + for (int tp = 0; tp < config_.tp_count; tp++) { + MeshSlotPool& pool = pools_[l][tp]; + for (int e = 0; e < config_.cap && e < config_.expert_num; e++) { + if (is_gpu_expert(e)) continue; + pool.mark_cached(e); + } + } + } + } + + // ===== 权重指针查询（KT 计算用）===== + + void* get_gate_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].expert_gate_ptr(expert_id); + } + + void* get_up_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].expert_up_ptr(expert_id); + } + + void* get_down_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].expert_down_ptr(expert_id); + } + + // ===== Prefill 阶段回调 ===== + + void on_prefill_layer_start(int layer_idx, int qlen, + const std::vector& active_experts) { + // 由 prefill 策略处理 + } + + void on_prefill_layer_done(int layer_idx) { + // 每过一层，temporal 角色互换 + prefill_->swap_temporal_roles(); + } + + // ===== 过渡阶段 ===== + + void on_prefill_to_decode() { + handoff_->transition(config_, *prefill_, *scheduler_, *io_, *scorer_, + pools_, layouts_, + /*move_gpu_expert_to_gpu=*/[](int, int) {}, + /*get_dst_ptrs=*/[this](int l, int tp, int e) { + return std::vector{ + pools_[l][tp].gate_ptr(e), + pools_[l][tp].up_ptr(e), + pools_[l][tp].down_ptr(e)}; + }); + } + + // ===== Decode 阶段回调 ===== + + void on_decode_token_start() { + scheduler_->inc_timeline_step(); + } + + /** + * @brief Decode 每层处理 + * + * @param layer_idx 当前层 + * @param topk 当前 token 的 top-k 专家 + * @param scores 各专家 router score（全长） + * @param tp_part_idx TP 分片 + * @return MeshDecode::SplitResult immediate/deferred 分组 + */ + MeshDecode::SplitResult on_decode_layer(int layer_idx, + const std::vector& topk, + const std::vector& scores, + int tp_part_idx) { + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; + return decode_->split(topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, + [this, &pool, layer_idx, tp_part_idx](int eid) { + return std::vector{ + pool.gate_ptr(eid), + pool.up_ptr(eid), + pool.down_ptr(eid)}; + }); + } + + /** + * @brief 单 token 结束后批量更新 Heat 和 Markov + * + * @param all_layers_topk [layer] -> top-k expert ids + * @param all_layers_scores [layer] -> normalized router scores (全长) + */ + void on_decode_token_end(const std::vector>& all_layers_topk, + const std::vector>& all_layers_scores) { + scorer_->commit_token(all_layers_topk, all_layers_scores); + } + + // ===== 访问器 ===== + const MeshConfig& config() const { return config_; } + MeshScheduler& scheduler() { return *scheduler_; } + MeshIoUring& io() { return *io_; } + MeshPrefill& prefill() { return *prefill_; } + MeshDecode& decode() { return *decode_; } + EvictionScorer& scorer() { return *scorer_; } + MeshSlotPool& pool(int layer, int tp) { return pools_[layer][tp]; } + const std::vector>& layouts() const { return layouts_; } + + private: + MeshConfig config_; + std::vector numa_nodes_; + size_t slot_bytes_ = 0; + + // [layer][tp] 的 slot 池 + std::vector> pools_; + + // [tp][expert] 的文件布局 + std::vector> layouts_; + + // GPU expert mask + std::vector gpu_experts_mask_; + + // 组件 + std::unique_ptr io_; + std::unique_ptr scheduler_; + std::unique_ptr scorer_; + std::unique_ptr prefill_; + std::unique_ptr decode_; + std::unique_ptr handoff_; + + // 计算 slot 字节数（gate + up + down 三个矩阵） + size_t compute_slot_bytes(const MeshConfig& config) const { + // 根据 SKILL 第 8 节的双 NUMA 拆分： + // gate: [hidden, intermediate/tp_count] + // up: [hidden, intermediate/tp_count] + // down: [intermediate/tp_count, hidden] + size_t gate_up_bytes = compute_gate_up_bytes(config); + size_t down_bytes = compute_down_bytes(config); + return gate_up_bytes * 2 + down_bytes; + } + + size_t compute_gate_up_bytes(const MeshConfig& config) const { + // gate 或 up 单块字节数 + int h = config.hidden_size; + int i = config.intermediate_size / config.tp_count; // TP 切分后 + size_t elements = static_cast(h) * i; + switch (config.weight_type) { + case WeightType::AMXINT4: + return elements / 2; // int4 = 0.5 byte + case WeightType::BF16: + return elements * 2; // bf16 = 2 bytes + default: + return elements * 2; + } + } + + size_t compute_down_bytes(const MeshConfig& config) const { + // down 单块字节数 + int h = config.hidden_size; + int i = config.intermediate_size / config.tp_count; + size_t elements = static_cast(i) * h; + switch (config.weight_type) { + case WeightType::AMXINT4: + return elements / 2; + case WeightType::BF16: + return elements * 2; + default: + return elements * 2; + } + } +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_scheduler.hpp b/kt-kernel/operators/mesh/mesh_scheduler.hpp new file mode 100644 index 000000000..b0c91a7ee --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_scheduler.hpp @@ -0,0 +1,197 @@ +/** + * @file mesh_scheduler.hpp + * @brief schedule_key 优先队列调度器 + * + * 每次 forward 计算一个 schedule_key，决定 io_uring 读请求的优先级排序。 + * 请求入优先队列（越小越优先），drain 时按序提交给 io_uring。 + * + * schedule_key 计算： + * - Prefill (qlen>1): schedule_key = layer_idx，按层号线性推进 + * - Decode immediate (qlen=1): schedule_key = timeline_step × total_layers + layer_idx + * - Decode deferred (qlen=1): schedule_key = timeline_step × total_layers + layer_idx + 1 + * （跟下一层同优先级，因为 deferred 专家推迟到下一层计算） + */ +#pragma once + +#include +#include +#include +#include +#include + +namespace mesh { + +// 调度请求 +struct ScheduledRequest { + uint64_t schedule_key; // 优先级键，越小越优先 + int expert_id; // 要读取的专家 ID + int tp_part_idx; // TP 分片索引 + int layer_idx; // 层索引 + ReadPriority priority; // Demand / Prefetch + // 目标 buffer 指针（由 slot pool 提供） + void* gate_dst = nullptr; + void* up_dst = nullptr; + void* down_dst = nullptr; + void* gate_scale_dst = nullptr; + void* up_scale_dst = nullptr; + void* down_scale_dst = nullptr; + void* gate_mins_dst = nullptr; + void* up_mins_dst = nullptr; + void* down_mins_dst = nullptr; + // 完成回调 + std::function on_complete; + + // 优先队列比较器：schedule_key 小的优先 + bool operator>(const ScheduledRequest& o) const { + return schedule_key > o.schedule_key; + } +}; + +/** + * @brief 全局时序调度器 + * + * 维护 timeline_step（算上 prefill 的第几个 token）和优先队列。 + * prefill 和 decode 共用同一个调度器，通过 is_prefill_ 区分 schedule_key 计算。 + */ +class MeshScheduler { + public: + MeshScheduler(int total_layers) : total_layers_(total_layers) {} + + // ===== schedule_key 计算 ===== + + // Prefill：按层号线性推进 + uint64_t prefill_key(int layer_idx) const { + return static_cast(layer_idx); + } + + // Decode immediate：当前 token 当前层 + uint64_t decode_immediate_key(int layer_idx) const { + return static_cast(timeline_step_) * total_layers_ + layer_idx; + } + + // Decode deferred：跟下一层同优先级 + uint64_t decode_deferred_key(int layer_idx) const { + return decode_immediate_key(layer_idx) + 1; + } + + // ===== 提交请求入队 ===== + + void submit_prefill(int layer_idx, int expert_id, int tp_part_idx, + void* gate_dst, void* up_dst, void* down_dst, + ReadPriority priority = ReadPriority::Prefetch) { + ScheduledRequest req; + req.schedule_key = prefill_key(layer_idx); + req.expert_id = expert_id; + req.tp_part_idx = tp_part_idx; + req.layer_idx = layer_idx; + req.priority = priority; + req.gate_dst = gate_dst; + req.up_dst = up_dst; + req.down_dst = down_dst; + enqueue(std::move(req)); + } + + void submit_decode_immediate(int layer_idx, int expert_id, int tp_part_idx, + void* gate_dst, void* up_dst, void* down_dst, + std::function on_complete = nullptr) { + ScheduledRequest req; + req.schedule_key = decode_immediate_key(layer_idx); + req.expert_id = expert_id; + req.tp_part_idx = tp_part_idx; + req.layer_idx = layer_idx; + req.priority = ReadPriority::Demand; + req.gate_dst = gate_dst; + req.up_dst = up_dst; + req.down_dst = down_dst; + req.on_complete = std::move(on_complete); + enqueue(std::move(req)); + } + + void submit_decode_deferred(int layer_idx, int expert_id, int tp_part_idx, + void* gate_dst, void* up_dst, void* down_dst, + std::function on_complete = nullptr) { + ScheduledRequest req; + req.schedule_key = decode_deferred_key(layer_idx); + req.expert_id = expert_id; + req.tp_part_idx = tp_part_idx; + req.layer_idx = layer_idx; + req.priority = ReadPriority::Demand; + req.gate_dst = gate_dst; + req.up_dst = up_dst; + req.down_dst = down_dst; + req.on_complete = std::move(on_complete); + enqueue(std::move(req)); + } + + // ===== 队列管理 ===== + + // 从队列取出所有请求，按 schedule_key 排序后提交给 io_uring + // caller 提供实际提交函数 + std::vector drain_all() { + std::lock_guard lock(mutex_); + std::vector result; + while (!pq_.empty()) { + result.push_back(std::move(const_cast(pq_.top()))); + pq_.pop(); + } + // result 已按 schedule_key 升序排列（priority_queue 是大顶堆，top 是最大的） + // 实际上我们需要最小的先出，所以用 greater 比较器 + // 这里 result 是从大到小，需要反转 + std::reverse(result.begin(), result.end()); + return result; + } + + // 取出 schedule_key 最小的请求（非阻塞） + bool try_pop(ScheduledRequest& out) { + std::lock_guard lock(mutex_); + if (pq_.empty()) return false; + out = std::move(const_cast(pq_.top())); + pq_.pop(); + return true; + } + + // 清空队列（过渡阶段调用） + void clear_queue() { + std::lock_guard lock(mutex_); + while (!pq_.empty()) pq_.pop(); + } + + bool empty() const { + std::lock_guard lock(mutex_); + return pq_.empty(); + } + + size_t size() const { + std::lock_guard lock(mutex_); + return pq_.size(); + } + + // ===== 时序状态 ===== + + void switch_to_decode() { is_prefill_ = false; } + void switch_to_prefill() { is_prefill_ = true; } + bool is_prefill() const { return is_prefill_; } + + void inc_timeline_step() { timeline_step_++; } + int timeline_step() const { return timeline_step_; } + void set_timeline_step(int step) { timeline_step_ = step; } + + private: + int total_layers_; + int timeline_step_ = 0; // 算上 prefill 的第几个 token + bool is_prefill_ = true; + + // 最小堆：schedule_key 小的优先 + using PriorityQueue = + std::priority_queue, + std::greater>; + PriorityQueue pq_; + mutable std::mutex mutex_; + + void enqueue(ScheduledRequest req) { + std::lock_guard lock(mutex_); + pq_.push(std::move(req)); + } +}; + +} // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_slot_pool.hpp b/kt-kernel/operators/mesh/mesh_slot_pool.hpp new file mode 100644 index 000000000..0fcef3a0a --- /dev/null +++ b/kt-kernel/operators/mesh/mesh_slot_pool.hpp @@ -0,0 +1,271 @@ +/** + * @file mesh_slot_pool.hpp + * @brief MESH Slot 池 + 专家状态机 + * + * 每个 MoE 层、每个 TP/NUMA 分片各自维护一个固定大小的 slot 数组。 + * slot 不存在释放语义，全部都是专家权重的内存覆盖。 + * + * 状态机：BASELINE → LOADING → CACHED →（驱逐）→ DEMOTING → BASELINE + * slot_active_readers_ 保证驱逐时无 reader（AMX 计算 / GPU 搬运 / 内部访问） + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace mesh { + +// 专家生命周期状态 +enum class ExpertState : uint8_t { + BASELINE, // 不在缓存，权重在磁盘上 + LOADING, // 正在从磁盘读取，io_uring CQE 还没到 + CACHED, // 已在 slot 中，AMX kernel 可直接消费 + DEMOTING, // 正在被驱逐，等待活跃 reader 归零后解绑 +}; + +// 单个 slot 的元数据 +struct Slot { + std::atomic state{static_cast(ExpertState::BASELINE)}; + std::atomic active_readers{0}; // 引用计数：AMX 计算 / GPU 搬运 / 内部访问 + int bound_expert_id{-1}; // 当前 slot 装的是哪个专家，-1 = 空 + + ExpertState get_state() const { + return static_cast(state.load(std::memory_order_acquire)); + } + void set_state(ExpertState s) { + state.store(static_cast(s), std::memory_order_release); + } +}; + +/** + * @brief Slot 池：每层每 TP 一个实例 + * + * 内存布局：cap 个 slot 连续分配在 NUMA 本地内存上。 + * 每个 slot 包含 gate + up + down 三个 buffer，大小相同。 + * + * 两种索引方式： + * - slots_[i]：按 slot 位置索引，用于驱逐扫描 + * - expert_to_slot_[expert_id]：按专家 ID 索引，用于 KT 计算取指针 + */ +class MeshSlotPool { + public: + MeshSlotPool(int layer_idx, int tp_part_idx, int numa_node, int cap, + size_t slot_bytes) + : layer_idx_(layer_idx), + tp_part_idx_(tp_part_idx), + numa_node_(numa_node), + cap_(cap), + slot_bytes_(slot_bytes) { + if (cap <= 0) { + throw std::runtime_error("MeshSlotPool: cap must be positive"); + } + // 一次性 numa_alloc_onnode 分配全部 slot 内存，不在关键路径 alloc/free + total_bytes_ = static_cast(cap) * slot_bytes_; + memory_ = numa_alloc_onnode(total_bytes_, numa_node_); + if (!memory_) { + throw std::runtime_error("MeshSlotPool: numa_alloc_onnode failed"); + } + std::memset(memory_, 0, total_bytes_); + + slots_.resize(cap); + expert_to_slot_.clear(); + slot_to_expert_.assign(cap, -1); + } + + ~MeshSlotPool() { + if (memory_) { + numa_free(memory_, total_bytes_); + } + } + + // 禁止拷贝 + MeshSlotPool(const MeshSlotPool&) = delete; + MeshSlotPool& operator=(const MeshSlotPool&) = delete; + + // ===== 指针访问 ===== + + // 获取 slot 中 gate 矩阵的地址 + void* gate_ptr(int slot_idx) { + return static_cast(memory_) + static_cast(slot_idx) * slot_bytes_ + 0; + } + // 获取 slot 中 up 矩阵的地址 + void* up_ptr(int slot_idx) { + return static_cast(memory_) + static_cast(slot_idx) * slot_bytes_ + gate_up_bytes_; + } + // 获取 slot 中 down 矩阵的地址 + void* down_ptr(int slot_idx) { + return static_cast(memory_) + static_cast(slot_idx) * slot_bytes_ + gate_up_bytes_ * 2; + } + + // 根据专家 ID 获取 gate 指针（KT 计算用），未缓存返回 nullptr + void* expert_gate_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + return gate_ptr(slot_idx); + } + void* expert_up_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + return up_ptr(slot_idx); + } + void* expert_down_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + return down_ptr(slot_idx); + } + + // ===== 状态查询 ===== + + bool is_cached(int expert_id) const { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return false; + return slots_[slot_idx].get_state() == ExpertState::CACHED; + } + + ExpertState get_expert_state(int expert_id) const { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return ExpertState::BASELINE; + return slots_[slot_idx].get_state(); + } + + // ===== 绑定 / 覆盖 ===== + + // 绑定：将 expert 读入空闲 slot，设置指针指向 + // 调用者需保证 slot_idx 处于 BASELINE 或 DEMOTING(reader==0) 状态 + void bind(int slot_idx, int expert_id) { + Slot& s = slots_[slot_idx]; + s.set_state(ExpertState::LOADING); + // io_uring 读取完成后调用 mark_cached + s.bound_expert_id = expert_id; + slot_to_expert_[slot_idx] = expert_id; + expert_to_slot_.insert(expert_id, slot_idx); + } + + // 标记 slot 已缓存完毕（io_uring CQE 到达后调用） + void mark_cached(int slot_idx) { + slots_[slot_idx].set_state(ExpertState::CACHED); + } + + // 覆盖：解绑旧 expert，将新 expert 读入同一个 slot + // 必须等 active_readers 归零后才能执行 + void overwrite(int slot_idx, int new_expert_id) { + Slot& s = slots_[slot_idx]; + // 等待活跃 reader 归零 + while (s.active_readers.load(std::memory_order_acquire) > 0) { + // spin wait，实际实现可用 futex + } + int old_expert_id = s.bound_expert_id; + if (old_expert_id >= 0) { + expert_to_slot_.erase(old_expert_id); + } + s.set_state(ExpertState::DEMOTING); + // 解绑完成，进入 LOADING 状态等待新数据 + s.set_state(ExpertState::LOADING); + s.bound_expert_id = new_expert_id; + slot_to_expert_[slot_idx] = new_expert_id; + expert_to_slot_.insert(new_expert_id, slot_idx); + } + + // ===== 引用计数 ===== + + // AMX 计算 / GPU 搬运前递增 reader + void acquire_reader(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return; + slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + } + + // 计算完毕递减 reader + void release_reader(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return; + slots_[slot_idx].active_readers.fetch_sub(1, std::memory_order_acq_rel); + } + + // ===== 驱逐扫描 ===== + + // 找一个可驱逐的 slot：状态为 CACHED 且 active_readers==0 + // 返回 slot_idx，找不到返回 -1 + int find_evictable() const { + for (int i = 0; i < cap_; i++) { + if (slots_[i].get_state() == ExpertState::CACHED && + slots_[i].active_readers.load(std::memory_order_acquire) == 0) { + return i; + } + } + return -1; + } + + // 获取所有 CACHED 状态的专家 ID（驱逐评分用） + std::vector cached_experts() const { + std::vector result; + for (int i = 0; i < cap_; i++) { + if (slots_[i].get_state() == ExpertState::CACHED && + slots_[i].bound_expert_id >= 0) { + result.push_back(slots_[i].bound_expert_id); + } + } + return result; + } + + // ===== 基本访问器 ===== + int layer_idx() const { return layer_idx_; } + int tp_part_idx() const { return tp_part_idx_; } + int numa_node() const { return numa_node_; } + int cap() const { return cap_; } + size_t slot_bytes() const { return slot_bytes_; } + + // 设置 gate/up 的单块字节数（用于指针偏移计算） + void set_gate_up_bytes(size_t bytes) { gate_up_bytes_ = bytes; } + + private: + int layer_idx_; + int tp_part_idx_; + int numa_node_; + int cap_; + size_t slot_bytes_; + size_t gate_up_bytes_ = 0; // gate 或 up 单块字节数 + void* memory_ = nullptr; + size_t total_bytes_ = 0; + + std::vector slots_; // [cap_]，驱逐扫描用 + + // expert_id -> slot_idx 的双向映射 + // 简单实现用 vector，O(1) 查找 + struct ExpertToSlot { + std::vector data; // [expert_num] -> slot_idx or -1 + void resize(int n) { data.assign(n, -1); } + int lookup(int expert_id) const { + if (expert_id < 0 || expert_id >= (int)data.size()) return -1; + return data[expert_id]; + } + void insert(int expert_id, int slot_idx) { + if (expert_id >= 0 && expert_id < (int)data.size()) { + data[expert_id] = slot_idx; + } + } + void erase(int expert_id) { + if (expert_id >= 0 && expert_id < (int)data.size()) { + data[expert_id] = -1; + } + } + } expert_to_slot_; + + std::vector slot_to_expert_; // [cap_] -> expert_id or -1 + + public: + // 初始化 expert_to_slot_ 的大小 + void init_expert_map(int expert_num) { + expert_to_slot_.resize(expert_num); + } +}; + +} // namespace mesh From 0d2f962c69b14432e494a490584ed4d87938f4a8 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Fri, 19 Jun 2026 22:24:46 +0800 Subject: [PATCH 02/10] feat: mesh py --- kt-kernel/python/utils/mesh/__init__.py | 23 +++ kt-kernel/python/utils/mesh/config.py | 160 ++++++++++++++++++ kt-kernel/python/utils/mesh/residency.py | 182 ++++++++++++++++++++ kt-kernel/python/utils/mesh/stats.py | 110 ++++++++++++ kt-kernel/python/utils/mesh/wrapper.py | 207 +++++++++++++++++++++++ 5 files changed, 682 insertions(+) create mode 100644 kt-kernel/python/utils/mesh/__init__.py create mode 100644 kt-kernel/python/utils/mesh/config.py create mode 100644 kt-kernel/python/utils/mesh/residency.py create mode 100644 kt-kernel/python/utils/mesh/stats.py create mode 100644 kt-kernel/python/utils/mesh/wrapper.py diff --git a/kt-kernel/python/utils/mesh/__init__.py b/kt-kernel/python/utils/mesh/__init__.py new file mode 100644 index 000000000..59bcfbb84 --- /dev/null +++ b/kt-kernel/python/utils/mesh/__init__.py @@ -0,0 +1,23 @@ +""" +MESH 插件 Python 侧入口。 + +MESH 是 KTransformers 的专家权重驻留管理插件。 +本包提供： +- MeshConfig: 配置 +- MeshResidencyManager: 驻留管理器 +- MeshMoEWrapper: MoE Wrapper +- MeshStatsCollector: 统计收集 +""" + +from .config import MeshConfig +from .residency import MeshResidencyManager +from .wrapper import MeshMoEWrapper +from .stats import MeshStats, MeshStatsCollector + +__all__ = [ + "MeshConfig", + "MeshResidencyManager", + "MeshMoEWrapper", + "MeshStats", + "MeshStatsCollector", +] diff --git a/kt-kernel/python/utils/mesh/config.py b/kt-kernel/python/utils/mesh/config.py new file mode 100644 index 000000000..16ed2baa1 --- /dev/null +++ b/kt-kernel/python/utils/mesh/config.py @@ -0,0 +1,160 @@ +""" +MESH 插件 Python 侧配置。 + +从环境变量读取 MESH 配置，转换为 C++ 侧 MeshConfig。 +环境变量： + KT_ENABLE_MESH : 是否启用 MESH（"1"/"0"） + KT_MESH_CAP : 单层单 TP 的 slot 数 + KT_NUM_GPU_EXPERTS : GE，过渡阶段搬运的 GPU 专家数 + KT_MAX_DEFERRED_EXPERTS_PER_TOKEN : 每 token 最多 defer 的专家数 + KT_MESH_DECODE_FRONT_LAYERS : decode 前 N 层满配，默认 5 + KT_MESH_DECODE_FRONT_LAYER_CAP : 前 N 层 cap，-1 = 拉满 + KT_MESH_TOTAL_LAYERS : 模型总层数 + KT_MESH_PREFILL_WINDOW : prefill 窗口宽度，默认 1 + KT_MESH_HEAT_GAMMA : Heat 层内 EMA 衰减率，默认 0.7 + KT_MESH_HEAT_BETA : Heat 跨 token 衰减率，默认 0.5 + KT_MESH_MARKOV_ALPHA : Markov 转移矩阵更新率，默认 0.5 + KT_MESH_MARKOV_TOPK : Markov 每行稀疏保留数，默认 16 + KT_MESH_LOOKAHEAD_WEIGHT : 驱逐评分权重，默认 1.0 + KT_MESH_WEIGHT_TYPE : 权重类型 "amxint4" / "bf16" +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class MeshConfig: + """MESH 配置，对应 C++ 侧 mesh::MeshConfig。""" + + # 基本开关 + enabled: bool = False + + # 容量配置 + cap: int = 0 # 单层单 TP 的 slot 数 + + # GPU expert 配置 + num_gpu_experts: int = 0 # GE + + # Expert Defer 配置 + max_deferred_per_token: int = 3 + + # Decode 前 N 层满配 + decode_front_layers: int = 5 + decode_front_layer_cap: int = -1 # -1 = 拉满 + + # 时序配置 + total_layers: int = 0 + prefill_window: int = 1 + + # Heat EMA 参数 + heat_gamma: float = 0.7 + heat_beta: float = 0.5 + + # Markov 转移矩阵参数 + markov_alpha: float = 0.5 + markov_topk: int = 16 + + # 驱逐评分权重 + lookahead_weight: float = 1.0 + + # 权重类型 + weight_type: str = "amxint4" # "amxint4" / "bf16" + + # 模型维度（由 wrapper 注入，不从 env 读） + hidden_size: int = 0 + intermediate_size: int = 0 + expert_num: int = 0 + tp_count: int = 1 + + @classmethod + def from_env(cls) -> "MeshConfig": + """从环境变量读取配置。""" + def _env_bool(key: str, default: bool = False) -> bool: + val = os.environ.get(key, "") + return val.lower() in ("1", "true", "yes", "on") + + def _env_int(key: str, default: int = 0) -> int: + try: + return int(os.environ.get(key, str(default))) + except ValueError: + return default + + def _env_float(key: str, default: float = 0.0) -> float: + try: + return float(os.environ.get(key, str(default))) + except ValueError: + return default + + return cls( + enabled=_env_bool("KT_ENABLE_MESH"), + cap=_env_int("KT_MESH_CAP"), + num_gpu_experts=_env_int("KT_NUM_GPU_EXPERTS"), + max_deferred_per_token=_env_int("KT_MAX_DEFERRED_EXPERTS_PER_TOKEN", 3), + decode_front_layers=_env_int("KT_MESH_DECODE_FRONT_LAYERS", 5), + decode_front_layer_cap=_env_int("KT_MESH_DECODE_FRONT_LAYER_CAP", -1), + total_layers=_env_int("KT_MESH_TOTAL_LAYERS"), + prefill_window=_env_int("KT_MESH_PREFILL_WINDOW", 1), + heat_gamma=_env_float("KT_MESH_HEAT_GAMMA", 0.7), + heat_beta=_env_float("KT_MESH_HEAT_BETA", 0.5), + markov_alpha=_env_float("KT_MESH_MARKOV_ALPHA", 0.5), + markov_topk=_env_int("KT_MESH_MARKOV_TOPK", 16), + lookahead_weight=_env_float("KT_MESH_LOOKAHEAD_WEIGHT", 1.0), + weight_type=os.environ.get("KT_MESH_WEIGHT_TYPE", "amxint4").lower(), + ) + + def validate(self) -> None: + """校验配置合法性。""" + if not self.enabled: + return + if self.cap <= 0: + raise ValueError("KT_MESH_CAP must be positive when MESH is enabled") + if self.total_layers <= 0: + raise ValueError("KT_MESH_TOTAL_LAYERS must be positive when MESH is enabled") + if self.expert_num <= 0: + raise ValueError("expert_num must be positive when MESH is enabled") + if self.weight_type not in ("amxint4", "bf16"): + raise ValueError(f"Unknown weight_type: {self.weight_type}") + if self.max_deferred_per_token < 0: + raise ValueError("max_deferred_per_token must be non-negative") + + def to_cpp(self): + """转换为 C++ 侧 MeshConfig 对象。 + + 需要 kt_kernel_ext.mesh.MeshConfig 已绑定。 + """ + try: + from kt_kernel_ext.mesh import MeshConfig as CppMeshConfig + except ImportError as e: + raise RuntimeError( + "MESH C++ extension not available. " + "Please build kt-kernel with CPUINFER_ENABLE_MESH=ON" + ) from e + + cpp_cfg = CppMeshConfig() + cpp_cfg.enabled = self.enabled + cpp_cfg.cap = self.cap + cpp_cfg.num_gpu_experts = self.num_gpu_experts + cpp_cfg.max_deferred_per_token = self.max_deferred_per_token + cpp_cfg.decode_front_layers = self.decode_front_layers + cpp_cfg.decode_front_layer_cap = self.decode_front_layer_cap + cpp_cfg.total_layers = self.total_layers + cpp_cfg.prefill_window = self.prefill_window + cpp_cfg.heat_gamma = self.heat_gamma + cpp_cfg.heat_beta = self.heat_beta + cpp_cfg.markov_alpha = self.markov_alpha + cpp_cfg.markov_topk = self.markov_topk + cpp_cfg.lookahead_weight = self.lookahead_weight + # weight_type 枚举 + if self.weight_type == "amxint4": + cpp_cfg.weight_type = 0 # WeightType::AMXINT4 + else: + cpp_cfg.weight_type = 1 # WeightType::BF16 + cpp_cfg.hidden_size = self.hidden_size + cpp_cfg.intermediate_size = self.intermediate_size + cpp_cfg.expert_num = self.expert_num + cpp_cfg.tp_count = self.tp_count + return cpp_cfg diff --git a/kt-kernel/python/utils/mesh/residency.py b/kt-kernel/python/utils/mesh/residency.py new file mode 100644 index 000000000..114557de7 --- /dev/null +++ b/kt-kernel/python/utils/mesh/residency.py @@ -0,0 +1,182 @@ +""" +MESH ResidencyManager Python 侧包装。 + +封装 C++ 侧 mesh::MeshResidencyManager，提供 Python 友好的接口。 +负责： +- 初始化 slot 池、io_uring、调度器等组件 +- 注入文件布局（safetensors 偏移） +- 注入 GPU expert mask +- 提供 forward 回调入口 +""" + +from __future__ import annotations + +import os +import logging +from typing import List, Optional, Tuple + +from .config import MeshConfig + +logger = logging.getLogger(__name__) + + +class MeshResidencyManager: + """C++ 侧 MeshResidencyManager 的 Python 包装。""" + + def __init__(self): + try: + from kt_kernel_ext.mesh import ResidencyManager as CppManager + except ImportError as e: + raise RuntimeError( + "MESH C++ extension not available. " + "Please build kt-kernel with CPUINFER_ENABLE_MESH=ON" + ) from e + self._mgr = CppManager() + self._initialized = False + self._layouts_set = False + + def init(self, config: MeshConfig, numa_nodes: List[int]) -> None: + """初始化 MESH。 + + Args: + config: MESH 配置 + numa_nodes: 每个 TP 分片对应的 NUMA 节点 ID + """ + config.validate() + cpp_cfg = config.to_cpp() + self._mgr.init(cpp_cfg, numa_nodes) + self._initialized = True + logger.info( + "MESH initialized: cap=%d, GE=%d, layers=%d, tp=%d", + config.cap, config.num_gpu_experts, + config.total_layers, config.tp_count, + ) + + def set_file_layout( + self, + tp_part_idx: int, + expert_id: int, + fd: int, + gate_offset: int, + up_offset: int, + down_offset: int, + gate_bytes: int, + up_bytes: int, + down_bytes: int, + gate_scale_offset: int = 0, + up_scale_offset: int = 0, + down_scale_offset: int = 0, + gate_scale_bytes: int = 0, + up_scale_bytes: int = 0, + down_scale_bytes: int = 0, + gate_mins_offset: int = 0, + up_mins_offset: int = 0, + down_mins_offset: int = 0, + gate_mins_bytes: int = 0, + up_mins_bytes: int = 0, + down_mins_bytes: int = 0, + ) -> None: + """注入单个专家在某个 TP 分片上的文件布局。 + + Args: + tp_part_idx: TP 分片索引 + expert_id: 专家 ID + fd: O_DIRECT 打开的文件描述符 + *_offset: 各矩阵在文件中的偏移 + *_bytes: 各矩阵的字节数 + """ + if not self._initialized: + raise RuntimeError("Call init() before set_file_layout()") + self._mgr.set_file_layout( + tp_part_idx, expert_id, fd, + gate_offset, up_offset, down_offset, + gate_bytes, up_bytes, down_bytes, + gate_scale_offset, up_scale_offset, down_scale_offset, + gate_scale_bytes, up_scale_bytes, down_scale_bytes, + gate_mins_offset, up_mins_offset, down_mins_offset, + gate_mins_bytes, up_mins_bytes, down_mins_bytes, + ) + + def set_gpu_experts_mask(self, mask: List[int]) -> None: + """注入 GPU expert mask。 + + Args: + mask: 长度 = expert_num 的列表，1 = GPU expert，0 = CPU expert + """ + if not self._initialized: + raise RuntimeError("Call init() before set_gpu_experts_mask()") + self._mgr.set_gpu_experts_mask(mask) + logger.info("GPU expert mask set: %d GPU experts", sum(mask)) + + def bootstrap(self) -> None: + """启动阶段：读前 cap 个专家进 slot。""" + if not self._initialized: + raise RuntimeError("Call init() before bootstrap()") + if not self._layouts_set: + raise RuntimeError("Call set_file_layout() before bootstrap()") + logger.info("MESH bootstrap: loading first %d experts per layer", self._mgr.config().cap) + self._mgr.bootstrap() + logger.info("MESH bootstrap complete") + + def mark_layouts_set(self) -> None: + """标记文件布局已全部注入。""" + self._layouts_set = True + + # ===== 权重指针查询（KT 计算用）===== + + def get_gate_ptr(self, layer: int, tp: int, expert_id: int) -> int: + """获取 gate 矩阵指针。""" + return self._mgr.get_gate_ptr(layer, tp, expert_id) + + def get_up_ptr(self, layer: int, tp: int, expert_id: int) -> int: + """获取 up 矩阵指针。""" + return self._mgr.get_up_ptr(layer, tp, expert_id) + + def get_down_ptr(self, layer: int, tp: int, expert_id: int) -> int: + """获取 down 矩阵指针。""" + return self._mgr.get_down_ptr(layer, tp, expert_id) + + # ===== Prefill 阶段回调 ===== + + def on_prefill_layer_start(self, layer_idx: int, qlen: int, + active_experts: List[int]) -> None: + self._mgr.on_prefill_layer_start(layer_idx, qlen, active_experts) + + def on_prefill_layer_done(self, layer_idx: int) -> None: + self._mgr.on_prefill_layer_done(layer_idx) + + # ===== 过渡阶段 ===== + + def on_prefill_to_decode(self) -> None: + """Prefill→Decode 过渡。 + + 在第一个 decode token 的第一层前调用。 + """ + logger.info("MESH prefill→decode handoff") + self._mgr.on_prefill_to_decode() + + # ===== Decode 阶段回调 ===== + + def on_decode_token_start(self) -> None: + self._mgr.on_decode_token_start() + + def on_decode_layer(self, layer_idx: int, topk: List[int], + scores: List[float], tp_part_idx: int): + """Decode 每层处理，返回 (immediate, deferred) 分组。""" + return self._mgr.on_decode_layer(layer_idx, topk, scores, tp_part_idx) + + def on_decode_token_end(self, all_layers_topk: List[List[int]], + all_layers_scores: List[List[float]]) -> None: + """单 token 结束后批量更新 Heat 和 Markov。""" + self._mgr.on_decode_token_end(all_layers_topk, all_layers_scores) + + # ===== 访问器 ===== + + @property + def config(self): + return self._mgr.config() + + @property + def raw(self): + """获取底层 C++ 对象指针（用于 hook 注册）。""" + return self._mgr diff --git a/kt-kernel/python/utils/mesh/stats.py b/kt-kernel/python/utils/mesh/stats.py new file mode 100644 index 000000000..e7c0fb90b --- /dev/null +++ b/kt-kernel/python/utils/mesh/stats.py @@ -0,0 +1,110 @@ +""" +MESH 统计信息收集。 + +读取 C++ 侧 MeshStats，供 benchmark 用。 +包括：hit rate、io_uring read GiB、驱逐次数等。 +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Optional + +logger = logging.getLogger(__name__) + + +@dataclass +class MeshStats: + """MESH 运行时统计。""" + + # 命中率 + cache_hit_count: int = 0 + cache_miss_count: int = 0 + + # io_uring 读取量 + io_uring_read_bytes: int = 0 + io_uring_read_count: int = 0 + + # 驱逐统计 + eviction_count: int = 0 + eviction_blocked_wait_us: int = 0 # 驱逐时等待 reader 归零的累计微秒 + + # defer 统计 + defer_count: int = 0 + defer_overflow_count: int = 0 # overflow 通道触发次数 + + # prefill 统计 + prefill_layer_count: int = 0 + prefill_temporal_swap_count: int = 0 # 指针 swap 次数 + + # decode 统计 + decode_token_count: int = 0 + decode_immediate_count: int = 0 + decode_deferred_count: int = 0 + + @property + def cache_hit_rate(self) -> float: + total = self.cache_hit_count + self.cache_miss_count + if total == 0: + return 0.0 + return self.cache_hit_count / total + + @property + def io_uring_read_gib(self) -> float: + return self.io_uring_read_bytes / (1024 ** 3) + + def summary(self) -> str: + """返回可读的统计摘要。""" + return ( + f"MESH Stats:\n" + f" cache hit rate: {self.cache_hit_rate:.2%} " + f"({self.cache_hit_count}/{self.cache_hit_count + self.cache_miss_count})\n" + f" io_uring read: {self.io_uring_read_gib:.2f} GiB " + f"({self.io_uring_read_count} requests)\n" + f" evictions: {self.eviction_count} " + f"(blocked {self.eviction_blocked_wait_us} us)\n" + f" defer: {self.defer_count} (overflow {self.defer_overflow_count})\n" + f" prefill: {self.prefill_layer_count} layers, " + f"{self.prefill_temporal_swap_count} swaps\n" + f" decode: {self.decode_token_count} tokens, " + f"{self.decode_immediate_count} immediate, " + f"{self.decode_deferred_count} deferred" + ) + + +class MeshStatsCollector: + """从 C++ 侧读取统计信息。""" + + def __init__(self, residency_manager): + """Args: + residency_manager: MeshResidencyManager 实例 + """ + self._mgr = residency_manager + + def collect(self) -> MeshStats: + """收集当前统计信息。""" + stats = MeshStats() + try: + cpp_stats = self._mgr.raw.stats() + stats.cache_hit_count = getattr(cpp_stats, "cache_hit_count", 0) + stats.cache_miss_count = getattr(cpp_stats, "cache_miss_count", 0) + stats.io_uring_read_bytes = getattr(cpp_stats, "io_uring_read_bytes", 0) + stats.io_uring_read_count = getattr(cpp_stats, "io_uring_read_count", 0) + stats.eviction_count = getattr(cpp_stats, "eviction_count", 0) + stats.eviction_blocked_wait_us = getattr(cpp_stats, "eviction_blocked_wait_us", 0) + stats.defer_count = getattr(cpp_stats, "defer_count", 0) + stats.defer_overflow_count = getattr(cpp_stats, "defer_overflow_count", 0) + stats.prefill_layer_count = getattr(cpp_stats, "prefill_layer_count", 0) + stats.prefill_temporal_swap_count = getattr(cpp_stats, "prefill_temporal_swap_count", 0) + stats.decode_token_count = getattr(cpp_stats, "decode_token_count", 0) + stats.decode_immediate_count = getattr(cpp_stats, "decode_immediate_count", 0) + stats.decode_deferred_count = getattr(cpp_stats, "decode_deferred_count", 0) + except (AttributeError, RuntimeError) as e: + logger.debug("Failed to collect MESH stats: %s", e) + return stats + + def log_summary(self) -> None: + """打印统计摘要。""" + stats = self.collect() + logger.info(stats.summary()) diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py new file mode 100644 index 000000000..67d8182e6 --- /dev/null +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -0,0 +1,207 @@ +""" +MESH 模式的 MoE Wrapper。 + +不继承 AMXMoEWrapper，而是持有它的实例并注入 mesh_enabled=True。 +AMX 内核代码完全复用，只在权重加载和 forward 指针重定向时走 mesh 路径。 + +职责： +1. 创建 AMXMoEWrapper（复用全部 AMX 内核） +2. 创建 ResidencyManager 并注入文件布局 + gpu_experts_mask +3. 在 forward 时调用 ResidencyManager 的回调 +""" + +from __future__ import annotations + +import logging +import os +import torch +from typing import List, Optional + +from .config import MeshConfig +from .residency import MeshResidencyManager + +logger = logging.getLogger(__name__) + + +class MeshMoEWrapper: + """MESH 模式的 MoE Wrapper。 + + 委托给 AMXMoEWrapper 进行实际计算，MESH 只负责权重驻留管理。 + """ + + def __init__( + self, + layer_idx: int, + num_experts: int, + num_experts_per_tok: int, + hidden_size: int, + moe_intermediate_size: int, + gpu_experts_mask: Optional[torch.Tensor], + cpuinfer_threads: int, + threadpool_count: int, + weight_path: str, + chunked_prefill_size: int, + mesh_config: MeshConfig, + method: str = "AMXINT4", + numa_nodes: Optional[List[int]] = None, + cpu_save: bool = False, + max_deferred_experts_per_token: Optional[int] = None, + ): + """初始化 MESH MoE Wrapper。 + + Args: + layer_idx: 层索引 + num_experts: 专家总数 + num_experts_per_tok: 每 token 的 top-k + hidden_size: 隐藏维度 + moe_intermediate_size: MoE 中间维度 + gpu_experts_mask: GPU expert mask + cpuinfer_threads: CPU 推理线程数 + threadpool_count: NUMA 子池数（TP 数） + weight_path: 权重路径 + chunked_prefill_size: prefill chunk 大小 + mesh_config: MESH 配置 + method: 后端方法（AMXINT4 / BF16） + numa_nodes: NUMA 节点列表 + cpu_save: 是否保存权重到 CPU 内存 + max_deferred_experts_per_token: 每 token defer 数 + """ + # 注入模型维度到 mesh_config + mesh_config.hidden_size = hidden_size + mesh_config.intermediate_size = moe_intermediate_size + mesh_config.expert_num = num_experts + mesh_config.tp_count = threadpool_count + if mesh_config.total_layers <= 0: + # 如果未设置，需要外部补充（通常由调用者设置） + logger.warning("mesh_config.total_layers not set, MESH schedule_key may be incorrect") + + # 1. 创建 AMXMoEWrapper，注入 mesh_enabled=True + # AMX 内核代码完全复用，load_weights() 会走 mesh 分支（直接 return） + from ..amx import AMXMoEWrapper + self._inner = AMXMoEWrapper( + layer_idx=layer_idx, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + gpu_experts_mask=gpu_experts_mask, + cpuinfer_threads=cpuinfer_threads, + threadpool_count=threadpool_count, + weight_path=weight_path, + chunked_prefill_size=chunked_prefill_size, + cpu_save=cpu_save, + max_deferred_experts_per_token=max_deferred_experts_per_token, + method=method, + numa_nodes=numa_nodes, + ) + + # 2. 创建 ResidencyManager + self._residency = MeshResidencyManager() + if numa_nodes is None: + numa_nodes = list(range(threadpool_count)) + self._residency.init(mesh_config, numa_nodes) + + # 3. 注入 GPU expert mask + if gpu_experts_mask is not None: + mask_list = gpu_experts_mask.cpu().tolist() + mask_int = [1 if x else 0 for x in mask_list] + self._residency.set_gpu_experts_mask(mask_int) + + # 4. 注入文件布局（从 SafeTensorLoader 获取） + self._inject_file_layouts(weight_path, threadpool_count, num_experts, mesh_config) + + # 5. 启动阶段：读前 cap 个专家进 slot + self._residency.mark_layouts_set() + self._residency.bootstrap() + + self._mesh_config = mesh_config + self._layer_idx = layer_idx + logger.info( + "MeshMoEWrapper initialized: layer=%d, cap=%d, method=%s", + layer_idx, mesh_config.cap, method, + ) + + def _inject_file_layouts( + self, + weight_path: str, + tp_count: int, + expert_num: int, + mesh_config: MeshConfig, + ) -> None: + """从 SafeTensorLoader 获取文件布局并注入 ResidencyManager。 + + 这里需要根据实际权重文件格式解析偏移。 + 简化实现：假设权重已按 TP 切分存储，每个 TP 一个文件。 + """ + # TODO: 实际实现需要根据 SafeTensorLoader 的接口获取每个专家的文件偏移 + # 这里只提供框架，具体偏移计算依赖权重文件格式 + # + # for tp in range(tp_count): + # for eid in range(expert_num): + # layout = loader.get_expert_layout(tp, eid) + # fd = os.open(layout.path, os.O_DIRECT | os.O_RDONLY) + # self._residency.set_file_layout( + # tp, eid, fd, + # layout.gate_offset, layout.up_offset, layout.down_offset, + # layout.gate_bytes, layout.up_bytes, layout.down_bytes, + # ... + # ) + logger.warning( + "File layout injection not yet implemented. " + "Please implement _inject_file_layouts based on SafeTensorLoader." + ) + + # ===== forward 委托 ===== + + def forward(self, *args, **kwargs): + """forward 委托给 AMXMoEWrapper。 + + MESH 通过 hook 在 AMX 内核中重定向权重指针， + Python 侧无需特殊处理。 + """ + return self._inner.forward(*args, **kwargs) + + def load_weights(self): + """load_weights 委托给 AMXMoEWrapper。 + + MESH 模式下 AMXMoEWrapper.load_weights() 会走 mesh 分支（直接 return）， + 实际权重加载由 ResidencyManager.bootstrap() 完成。 + """ + return self._inner.load_weights() + + # ===== MESH 专属接口 ===== + + @property + def residency(self) -> MeshResidencyManager: + """获取 ResidencyManager。""" + return self._residency + + @property + def mesh_config(self) -> MeshConfig: + return self._mesh_config + + @property + def inner(self): + """获取内部 AMXMoEWrapper。""" + return self._inner + + # ===== 生命周期回调 ===== + + def on_prefill_layer_start(self, qlen: int, active_experts: List[int]) -> None: + self._residency.on_prefill_layer_start(self._layer_idx, qlen, active_experts) + + def on_prefill_layer_done(self) -> None: + self._residency.on_prefill_layer_done(self._layer_idx) + + def on_prefill_to_decode(self) -> None: + self._residency.on_prefill_to_decode() + + def on_decode_token_start(self) -> None: + self._residency.on_decode_token_start() + + def on_decode_layer(self, topk: List[int], scores: List[float], tp_part_idx: int): + return self._residency.on_decode_layer(self._layer_idx, topk, scores, tp_part_idx) + + def on_decode_token_end(self, all_layers_topk: List[List[int]], + all_layers_scores: List[List[float]]) -> None: + self._residency.on_decode_token_end(all_layers_topk, all_layers_scores) From 93a9aaba7348eb720cce810efb60a6c28d6d0be5 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Fri, 19 Jun 2026 22:32:27 +0800 Subject: [PATCH 03/10] feat: mesh kthook --- kt-kernel/CMakeLists.txt | 19 ++++++++ kt-kernel/cpu_backend/cpuinfer.h | 38 +++++++++++++++ kt-kernel/ext_bindings.cpp | 84 ++++++++++++++++++++++++++++++++ kt-kernel/operators/amx/moe.hpp | 21 ++++++++ kt-kernel/operators/common.hpp | 6 +++ kt-kernel/python/experts.py | 26 ++++++++++ kt-kernel/setup.py | 6 +++ 7 files changed, 200 insertions(+) diff --git a/kt-kernel/CMakeLists.txt b/kt-kernel/CMakeLists.txt index ccbe4215c..e729275b6 100644 --- a/kt-kernel/CMakeLists.txt +++ b/kt-kernel/CMakeLists.txt @@ -544,6 +544,19 @@ message(STATUS "SOURCE_DIR7: ${SOURCE_DIR7}") set(ALL_SOURCES ${SOURCE_DIR1} ${SOURCE_DIR2} ${SOURCE_DIR3} ${SOURCE_DIR4} ${SOURCE_DIR5} ${SOURCE_DIR6} ${SOURCE_DIR7}) +# ===== MESH 插件（可选）===== +option(KT_ENABLE_MESH "Enable MESH io_uring resident-cache plugin" OFF) +if(KT_ENABLE_MESH) + add_compile_definitions(KT_ENABLE_MESH=1) + aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/mesh MESH_SOURCES) + list(APPEND ALL_SOURCES ${MESH_SOURCES}) + + # 查找 liburing + find_package(PkgConfig REQUIRED) + pkg_check_modules(LIBURING REQUIRED liburing) + message(STATUS "MESH: enabled (liburing found at ${LIBURING_INCLUDE_DIRS})") +endif() + file(GLOB_RECURSE FMT_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp" @@ -685,6 +698,12 @@ if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML) target_compile_definitions(${PROJECT_NAME} PRIVATE CPU_USE_KML) endif() target_link_libraries(${PROJECT_NAME} PRIVATE llama PkgConfig::HWLOC OpenMP::OpenMP_CXX) + +# MESH 插件链接 liburing + numa +if(KT_ENABLE_MESH) + target_include_directories(${PROJECT_NAME} PRIVATE ${LIBURING_INCLUDE_DIRS}) + target_link_libraries(${PROJECT_NAME} PRIVATE ${LIBURING_LIBRARIES} numa) +endif() if(NOT HOST_IS_X86 AND KTRANSFORMERS_CPU_USE_KML) if(KTRANSFORMERS_CPU_DEBUG) # add_executable(convert-test ${CMAKE_CURRENT_SOURCE_DIR}/operators/kml/convert-test.cpp) diff --git a/kt-kernel/cpu_backend/cpuinfer.h b/kt-kernel/cpu_backend/cpuinfer.h index 28f23bb7c..8efe18c9e 100644 --- a/kt-kernel/cpu_backend/cpuinfer.h +++ b/kt-kernel/cpu_backend/cpuinfer.h @@ -109,6 +109,44 @@ class CPUInfer { SyncArgs* args = new SyncArgs{this, allow_n_pending}; sync_(args); } + + // ===== MESH 插件：gating score 批量传输回调 ===== + // MESH 需要每 token 结束后将所有层的 gating 分数从 GPU 一次性传到 CPU。 + // 利用现有 submit_with_cuda_stream 机制，通过 cudaLaunchHostFunc 调度到 CUDA stream。 + // 原版 KT 的 moe gate 只传 top-8，MESH 改成传全长 vector（专家数长度）。 + struct GatingScoreArgs { + CPUInfer* cpuinfer; + void* mesh_residency; // mesh::MeshResidencyManager* + const float* gating_scores_gpu; // GPU 指针，所有层的 gating 分数 + int num_layers; + int expert_num; + // CQE 到达后由 mesh_residency 的 on_decode_token_end 消费 + }; + + static void submit_gating_scores_(void* args_ptr) { + GatingScoreArgs* args = (GatingScoreArgs*)args_ptr; + // 实际实现：将 gating_scores_gpu 数据传给 mesh::MeshResidencyManager::on_decode_token_end + // mesh_residency 通过 CUDA memcpy 或 unified memory 获取数据 + // 这里是 host callback，在 CUDA stream 上执行 + delete args; + } + +#ifndef KTRANSFORMERS_CPU_ONLY + void submit_gating_scores_with_cuda_stream( + intptr_t user_cuda_stream, + void* mesh_residency, + const float* gating_scores_gpu, + int num_layers, int expert_num) { +#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \ + defined(KTRANSFORMERS_USE_MACA) + GatingScoreArgs* args = new GatingScoreArgs{ + this, mesh_residency, gating_scores_gpu, num_layers, expert_num + }; + cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, + (cudaHostFn_t)&submit_gating_scores_, (void*)args); +#endif + } +#endif #ifndef KTRANSFORMERS_CPU_ONLY void sync_with_cuda_stream(intptr_t user_cuda_stream, size_t allow_n_pending = 0) { #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \ diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 87074e417..c287b7b3f 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -23,6 +23,11 @@ #include "cpu_backend/worker_pool.h" #include "operators/common.hpp" +// MESH 插件（仅在 KT_ENABLE_MESH 编译时生效） +#if defined(KT_ENABLE_MESH) +#include "operators/mesh/mesh.hpp" +#endif + #if defined(USE_MOE_KERNEL) #include "operators/moe_kernel/la/kernel.hpp" #include "operators/moe_kernel/moe.hpp" @@ -992,6 +997,85 @@ PYBIND11_MODULE(kt_kernel_ext, m) { utils.def("from_float", &from_float_ptr, "Convert tensor from float32 to any GGML type", py::arg("input"), py::arg("size"), py::arg("type")); + + // ===== MESH 插件子模块 ===== +#if defined(KT_ENABLE_MESH) + { + auto mesh_module = m.def_submodule("mesh"); + + // MeshConfig + py::class_(mesh_module, "MeshConfig") + .def(py::init<>()) + .def_readwrite("enabled", &mesh::MeshConfig::enabled) + .def_readwrite("cap", &mesh::MeshConfig::cap) + .def_readwrite("num_gpu_experts", &mesh::MeshConfig::num_gpu_experts) + .def_readwrite("max_deferred_per_token", &mesh::MeshConfig::max_deferred_per_token) + .def_readwrite("decode_front_layers", &mesh::MeshConfig::decode_front_layers) + .def_readwrite("decode_front_layer_cap", &mesh::MeshConfig::decode_front_layer_cap) + .def_readwrite("total_layers", &mesh::MeshConfig::total_layers) + .def_readwrite("prefill_window", &mesh::MeshConfig::prefill_window) + .def_readwrite("heat_gamma", &mesh::MeshConfig::heat_gamma) + .def_readwrite("heat_beta", &mesh::MeshConfig::heat_beta) + .def_readwrite("markov_alpha", &mesh::MeshConfig::markov_alpha) + .def_readwrite("markov_topk", &mesh::MeshConfig::markov_topk) + .def_readwrite("lookahead_weight", &mesh::MeshConfig::lookahead_weight) + .def_readwrite("weight_type", &mesh::MeshConfig::weight_type) + .def_readwrite("hidden_size", &mesh::MeshConfig::hidden_size) + .def_readwrite("intermediate_size", &mesh::MeshConfig::intermediate_size) + .def_readwrite("expert_num", &mesh::MeshConfig::expert_num) + .def_readwrite("tp_count", &mesh::MeshConfig::tp_count); + + // ExpertFileLayout + py::class_(mesh_module, "ExpertFileLayout") + .def(py::init<>()) + .def_readwrite("fd", &mesh::ExpertFileLayout::fd) + .def_readwrite("gate_offset", &mesh::ExpertFileLayout::gate_offset) + .def_readwrite("up_offset", &mesh::ExpertFileLayout::up_offset) + .def_readwrite("down_offset", &mesh::ExpertFileLayout::down_offset) + .def_readwrite("gate_bytes", &mesh::ExpertFileLayout::gate_bytes) + .def_readwrite("up_bytes", &mesh::ExpertFileLayout::up_bytes) + .def_readwrite("down_bytes", &mesh::ExpertFileLayout::down_bytes); + + // MeshResidencyManager + py::class_>( + mesh_module, "ResidencyManager") + .def(py::init<>()) + .def("init", &mesh::MeshResidencyManager::init, + py::arg("config"), py::arg("numa_nodes")) + .def("set_file_layout", &mesh::MeshResidencyManager::set_file_layout, + py::arg("tp_part_idx"), py::arg("expert_id"), py::arg("layout")) + .def("set_gpu_experts_mask", &mesh::MeshResidencyManager::set_gpu_experts_mask, + py::arg("mask")) + .def("bootstrap", &mesh::MeshResidencyManager::bootstrap) + .def("get_gate_ptr", &mesh::MeshResidencyManager::get_gate_ptr) + .def("get_up_ptr", &mesh::MeshResidencyManager::get_up_ptr) + .def("get_down_ptr", &mesh::MeshResidencyManager::get_down_ptr) + .def("on_prefill_layer_start", &mesh::MeshResidencyManager::on_prefill_layer_start) + .def("on_prefill_layer_done", &mesh::MeshResidencyManager::on_prefill_layer_done) + .def("on_prefill_to_decode", &mesh::MeshResidencyManager::on_prefill_to_decode) + .def("on_decode_token_start", &mesh::MeshResidencyManager::on_decode_token_start) + .def("on_decode_layer", &mesh::MeshResidencyManager::on_decode_layer) + .def("on_decode_token_end", &mesh::MeshResidencyManager::on_decode_token_end) + .def("config", [](mesh::MeshResidencyManager& mgr) -> mesh::MeshConfig& { + return mgr.config(); + }, py::return_value_policy::reference); + + // 注册 hook 函数指针（让 mesh_hook.hpp 的 inline 函数能调用 ResidencyManager） + mesh::hook::HookRegistry registry; + registry.get_gate_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_gate_ptr(layer, tp, eid); + }; + registry.get_up_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_up_ptr(layer, tp, eid); + }; + registry.get_down_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_down_ptr(layer, tp, eid); + }; + mesh::hook::register_hooks(registry); + + printf("MESH plugin: bindings registered\n"); + } +#endif } #if defined(KTRANSFORMERS_ENABLE_CPPTRACE) diff --git a/kt-kernel/operators/amx/moe.hpp b/kt-kernel/operators/amx/moe.hpp index 4a8450df0..33ba89d59 100644 --- a/kt-kernel/operators/amx/moe.hpp +++ b/kt-kernel/operators/amx/moe.hpp @@ -184,6 +184,10 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; + // MESH hook: mesh_enabled 时 bb 的底层指针已由 MESH slot 池提供 + // gate_bb_/up_bb_ 在 load_weights() mesh 分支中已绑定为 slot buffer 的包装 + // MESH 通过覆盖 slot 内存实现专家切换，bb 指针本身不变 + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); } else { @@ -197,6 +201,8 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { auto& bb = down_bb_[expert_idx]; auto& bc = down_bc_[expert_idx]; + // MESH hook: 同上，down_bb_ 已绑定为 MESH slot buffer + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); } else { @@ -204,6 +210,13 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { } } void load_weights() { + // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版 mmap/file/bf16 路径 + if (config_.mesh_enabled && config_.mesh_residency) { + // MESH 模式下 gate_bb_/up_bb_/down_bb_ 在 forward 时由 mesh hook 动态重定向到 slot buffer + // 这里只需标记 weights_loaded，实际权重加载由 ResidencyManager.bootstrap() 完成 + this->weights_loaded = true; + return; + } auto pool = config_.pool->get_subpool(tp_part_idx); const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; if (config_.gate_projs.size()) { @@ -363,6 +376,14 @@ class TP_MOE> : public TP_MOE>> { using Base::Base; void load_weights() override { + // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版加载路径 + if (config.mesh_enabled && config.mesh_residency) { + for (auto i = 0; i < tp_count; i++) { + tps[i]->load_weights(); // 各 TP 走 mesh 分支 + } + this->weights_loaded = true; + return; + } auto& config = this->config; auto& tps = this->tps; auto& tp_count = this->tp_count; diff --git a/kt-kernel/operators/common.hpp b/kt-kernel/operators/common.hpp index 86f3464c6..6fc3393a0 100644 --- a/kt-kernel/operators/common.hpp +++ b/kt-kernel/operators/common.hpp @@ -257,6 +257,12 @@ struct GeneralMOEConfig { return expert_id < 0 || expert_id >= expert_num || (gpu_experts_mask && gpu_experts_mask[expert_id]); } + // ===== MESH 插件接入点 ===== + // mesh_enabled=false 时所有 mesh 路径不生效，行为完全等价于原版 KT。 + // mesh_residency 指向 mesh::MeshResidencyManager，实际类型由 mesh_hook.hpp 处理。 + bool mesh_enabled = false; + void* mesh_residency = nullptr; // 实际类型: mesh::MeshResidencyManager* + void* gate_proj = nullptr; void* up_proj = nullptr; void* down_proj = nullptr; diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index 318fe3170..3ba83682d 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -153,6 +153,8 @@ def __new__( # MiniMax M3 swigluoai sigmoid alpha. 0.0 = standard silu (default). # Non-zero triggers gate * sigmoid(gate * alpha) * (up + 1) in act_fn. swiglu_alpha: float = 0.0, + # MESH 插件配置，None 时不启用 MESH + mesh_config=None, ): """ Factory method to create the appropriate backend implementation. @@ -225,6 +227,7 @@ def __new__( numa_nodes=numa_nodes, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, + mesh_config=mesh_config, ) else: # mode == "sft" # SFT factory does not plumb swiglu_limit; reject non-zero @@ -322,16 +325,39 @@ def _create_inference_wrapper( numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, swiglu_alpha: float = 0.0, + mesh_config=None, # MESH 插件配置，None 时不启用 MESH ) -> BaseMoEWrapper: """ Create an inference wrapper based on the method. Args: See KTMoEWrapper.__new__ for parameter descriptions. + mesh_config: MESH 配置对象，非 None 时启用 MESH 模式 Returns: BaseMoEWrapper instance """ + # MESH 插件：mesh_config 非 None 且 enabled 时走 MESH 路径 + if mesh_config is not None and getattr(mesh_config, "enabled", False): + from .utils.mesh.wrapper import MeshMoEWrapper + return MeshMoEWrapper( + layer_idx=layer_idx, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + gpu_experts_mask=gpu_experts_mask, + cpuinfer_threads=cpuinfer_threads, + threadpool_count=threadpool_count, + weight_path=weight_path, + chunked_prefill_size=chunked_prefill_size, + mesh_config=mesh_config, + method=method, + numa_nodes=numa_nodes, + cpu_save=cpu_save, + max_deferred_experts_per_token=max_deferred_experts_per_token, + ) + # Select backend based on method if method in ["AMXINT4", "AMXINT8"]: backend_cls = AMXMoEWrapper diff --git a/kt-kernel/setup.py b/kt-kernel/setup.py index 68afe7516..7d69726d4 100644 --- a/kt-kernel/setup.py +++ b/kt-kernel/setup.py @@ -25,6 +25,7 @@ CPUINFER_ENABLE_AVX512_BF16=OFF ON/OFF -> -DLLAMA_AVX512_BF16 CPUINFER_ENABLE_AVX512_VBMI=OFF ON/OFF -> -DLLAMA_AVX512_VBMI (required for FP8 MoE) CPUINFER_ENABLE_CPPTRACE=ON/OFF ON/OFF -> -DKTRANSFORMERS_ENABLE_CPPTRACE (debug-only) + CPUINFER_ENABLE_MESH=OFF ON/OFF -> -DKT_ENABLE_MESH (MESH io_uring resident-cache plugin) CPUINFER_BLIS_ROOT=/path/to/blis Forward to -DBLIS_ROOT @@ -655,6 +656,9 @@ def find_maca_path() -> str | None: _forward_str_env(cmake_args, "CPUINFER_LTO_MODE", "CPUINFER_LTO_MODE") _forward_bool_env(cmake_args, "CPUINFER_ENABLE_CPPTRACE", "KTRANSFORMERS_ENABLE_CPPTRACE") + # MESH 插件开关 + _forward_bool_env(cmake_args, "CPUINFER_ENABLE_MESH", "KT_ENABLE_MESH") + # CUDA static runtime toggle _forward_bool_env(cmake_args, "CPUINFER_CUDA_STATIC_RUNTIME", "KTRANSFORMERS_CUDA_STATIC_RUNTIME") @@ -791,6 +795,7 @@ def find_maca_path() -> str | None: packages=[ "kt_kernel", "kt_kernel.utils", + "kt_kernel.utils.mesh", # MESH 插件 "kt_kernel.sft", "kt_kernel.cli", "kt_kernel.cli.commands", @@ -801,6 +806,7 @@ def find_maca_path() -> str | None: package_dir={ "kt_kernel": "python", "kt_kernel.utils": "python/utils", + "kt_kernel.utils.mesh": "python/utils/mesh", # MESH 插件 "kt_kernel.sft": "python/sft", "kt_kernel.cli": "python/cli", "kt_kernel.cli.commands": "python/cli/commands", From 79733d5a8763890f1e63b1fa339a9c59bd02a0e2 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Sun, 21 Jun 2026 02:07:00 +0800 Subject: [PATCH 04/10] feat:mesh check v1 --- kt-kernel/ext_bindings.cpp | 18 ++++- kt-kernel/operators/amx/moe.hpp | 9 +-- kt-kernel/operators/mesh/mesh_io_uring.hpp | 5 ++ kt-kernel/operators/mesh/mesh_prefill.hpp | 8 +- kt-kernel/operators/mesh/mesh_residency.hpp | 28 +++++++ kt-kernel/operators/mesh/mesh_slot_pool.hpp | 54 ++++++++++++- kt-kernel/python/experts.py | 10 +++ kt-kernel/python/utils/mesh/config.py | 41 +++++++++- kt-kernel/python/utils/mesh/wrapper.py | 85 +++++++++++++++------ kt-kernel/scripts/convert_cpu_weights.py | 23 +++--- 10 files changed, 236 insertions(+), 45 deletions(-) diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index c287b7b3f..f57c3e5fa 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -1003,6 +1003,12 @@ PYBIND11_MODULE(kt_kernel_ext, m) { { auto mesh_module = m.def_submodule("mesh"); + // WeightType 枚举 + py::enum_(mesh_module, "WeightType") + .value("AMXINT4", mesh::WeightType::AMXINT4) + .value("BF16", mesh::WeightType::BF16) + .export_values(); + // MeshConfig py::class_(mesh_module, "MeshConfig") .def(py::init<>()) @@ -1044,7 +1050,15 @@ PYBIND11_MODULE(kt_kernel_ext, m) { py::arg("config"), py::arg("numa_nodes")) .def("set_file_layout", &mesh::MeshResidencyManager::set_file_layout, py::arg("tp_part_idx"), py::arg("expert_id"), py::arg("layout")) - .def("set_gpu_experts_mask", &mesh::MeshResidencyManager::set_gpu_experts_mask, + .def("set_gpu_experts_mask", + [](mesh::MeshResidencyManager& mgr, py::list mask_list) { + std::vector mask; + mask.reserve(py::len(mask_list)); + for (auto item : mask_list) { + mask.push_back(static_cast(py::cast(item))); + } + mgr.set_gpu_experts_mask(mask.data(), static_cast(mask.size())); + }, py::arg("mask")) .def("bootstrap", &mesh::MeshResidencyManager::bootstrap) .def("get_gate_ptr", &mesh::MeshResidencyManager::get_gate_ptr) @@ -1056,7 +1070,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def("on_decode_token_start", &mesh::MeshResidencyManager::on_decode_token_start) .def("on_decode_layer", &mesh::MeshResidencyManager::on_decode_layer) .def("on_decode_token_end", &mesh::MeshResidencyManager::on_decode_token_end) - .def("config", [](mesh::MeshResidencyManager& mgr) -> mesh::MeshConfig& { + .def("config", [](mesh::MeshResidencyManager& mgr) -> const mesh::MeshConfig& { return mgr.config(); }, py::return_value_policy::reference); diff --git a/kt-kernel/operators/amx/moe.hpp b/kt-kernel/operators/amx/moe.hpp index 33ba89d59..b83604fb7 100644 --- a/kt-kernel/operators/amx/moe.hpp +++ b/kt-kernel/operators/amx/moe.hpp @@ -213,8 +213,7 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版 mmap/file/bf16 路径 if (config_.mesh_enabled && config_.mesh_residency) { // MESH 模式下 gate_bb_/up_bb_/down_bb_ 在 forward 时由 mesh hook 动态重定向到 slot buffer - // 这里只需标记 weights_loaded，实际权重加载由 ResidencyManager.bootstrap() 完成 - this->weights_loaded = true; + // 实际权重加载由 ResidencyManager.bootstrap() 完成 return; } auto pool = config_.pool->get_subpool(tp_part_idx); @@ -377,9 +376,9 @@ class TP_MOE> : public TP_MOE>> { void load_weights() override { // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版加载路径 - if (config.mesh_enabled && config.mesh_residency) { - for (auto i = 0; i < tp_count; i++) { - tps[i]->load_weights(); // 各 TP 走 mesh 分支 + if (this->config.mesh_enabled && this->config.mesh_residency) { + for (auto i = 0; i < this->tp_count; i++) { + this->tps[i]->load_weights(); // 各 TP 走 mesh 分支 } this->weights_loaded = true; return; diff --git a/kt-kernel/operators/mesh/mesh_io_uring.hpp b/kt-kernel/operators/mesh/mesh_io_uring.hpp index 758360670..ba3ef7794 100644 --- a/kt-kernel/operators/mesh/mesh_io_uring.hpp +++ b/kt-kernel/operators/mesh/mesh_io_uring.hpp @@ -16,6 +16,11 @@ #include #include #include +// liburing.h 定义了 BLOCK_SIZE 等宏，会污染全局命名空间， +// 与 amx_raw_kernels.hpp / fp8-moe.hpp 中的 BLOCK_SIZE 变量冲突 +#ifdef BLOCK_SIZE +#undef BLOCK_SIZE +#endif #include #include #include diff --git a/kt-kernel/operators/mesh/mesh_prefill.hpp b/kt-kernel/operators/mesh/mesh_prefill.hpp index 93c9eff11..45c23c978 100644 --- a/kt-kernel/operators/mesh/mesh_prefill.hpp +++ b/kt-kernel/operators/mesh/mesh_prefill.hpp @@ -61,14 +61,18 @@ class MeshPrefill { void init_temporal() { if (temporal_a_ != nullptr) return; // 已初始化 size_t bytes = static_cast(temporal_size_) * slot_bytes_ * tp_count_; + role_a_ = TemporalRole::COMPUTE; + role_b_ = TemporalRole::PREFETCH; + if (bytes == 0) { + // cap >= expert_num: 所有专家常驻 slot，不需要 temporal buffer + return; + } temporal_a_ = numa_alloc_onnode(bytes, numa_node_); temporal_b_ = numa_alloc_onnode(bytes, numa_node_); if (!temporal_a_ || !temporal_b_) { throw std::runtime_error("MeshPrefill: numa_alloc_onnode for temporal failed"); } temporal_bytes_ = bytes; - role_a_ = TemporalRole::COMPUTE; - role_b_ = TemporalRole::PREFETCH; } // 释放 temporal 内存（仅 prefill→decode 切换时调用） diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp index a99772cc9..351d4fd17 100644 --- a/kt-kernel/operators/mesh/mesh_residency.hpp +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -54,6 +54,24 @@ class MeshResidencyManager { // 计算 slot 字节数（根据权重类型和模型维度） slot_bytes_ = compute_slot_bytes(config); + size_t gate_up_bytes = compute_gate_up_bytes(config); + size_t down_bytes = compute_down_bytes(config); + + // 诊断日志：确认 config 值和内存分配量 + size_t total_slot_bytes = static_cast(config.total_layers) * + config.tp_count * config.cap * slot_bytes_; + fprintf(stderr, + "[MESH init DIAG] hidden=%d, inter=%d, expert_num=%d, tp=%d, " + "cap=%d, layers=%d, weight_type=%d\n", + config.hidden_size, config.intermediate_size, + config.expert_num, config.tp_count, + config.cap, config.total_layers, + static_cast(config.weight_type)); + fprintf(stderr, + "[MESH init DIAG] gate_up_bytes=%zu, down_bytes=%zu, slot_bytes=%zu, " + "total_slot_pool=%zu MB\n", + gate_up_bytes, down_bytes, slot_bytes_, + total_slot_bytes / (1024 * 1024)); // 创建 slot 池 [layer][tp] pools_.resize(config.total_layers); @@ -78,6 +96,9 @@ class MeshResidencyManager { // 初始化 temporal 双缓冲 prefill_->init_temporal(); + + fprintf(stderr, "[MESH init DIAG] init complete, pools_ size=%zu\n", + pools_.size()); } // ===== 文件布局注入 ===== @@ -111,6 +132,13 @@ class MeshResidencyManager { * 每层 Slot 池按编号顺序填满。 */ void bootstrap() { + // 文件布局未注入时跳过实际权重加载（仅用于内存测试模式） + if (layouts_.empty()) { + fprintf(stderr, "[MESH] bootstrap: layouts_ empty, skipping weight preload " + "(memory-test mode, slot pools already allocated)\n"); + return; + } + // 预加载 Scale Cache（AMXINT4 专用） if (config_.weight_type == WeightType::AMXINT4) { io_->preload_scale_cache(config_.expert_num, config_.tp_count, diff --git a/kt-kernel/operators/mesh/mesh_slot_pool.hpp b/kt-kernel/operators/mesh/mesh_slot_pool.hpp index 0fcef3a0a..48f645da7 100644 --- a/kt-kernel/operators/mesh/mesh_slot_pool.hpp +++ b/kt-kernel/operators/mesh/mesh_slot_pool.hpp @@ -34,6 +34,23 @@ struct Slot { std::atomic active_readers{0}; // 引用计数：AMX 计算 / GPU 搬运 / 内部访问 int bound_expert_id{-1}; // 当前 slot 装的是哪个专家，-1 = 空 + // std::atomic 不可拷贝/移动，需自定义 move 以支持 std::vector + Slot() = default; + Slot(Slot&& other) noexcept + : state(other.state.load(std::memory_order_relaxed)), + active_readers(other.active_readers.load(std::memory_order_relaxed)), + bound_expert_id(other.bound_expert_id) {} + Slot& operator=(Slot&& other) noexcept { + if (this != &other) { + state.store(other.state.load(std::memory_order_relaxed), std::memory_order_relaxed); + active_readers.store(other.active_readers.load(std::memory_order_relaxed), std::memory_order_relaxed); + bound_expert_id = other.bound_expert_id; + } + return *this; + } + Slot(const Slot&) = delete; + Slot& operator=(const Slot&) = delete; + ExpertState get_state() const { return static_cast(state.load(std::memory_order_acquire)); } @@ -73,7 +90,6 @@ class MeshSlotPool { std::memset(memory_, 0, total_bytes_); slots_.resize(cap); - expert_to_slot_.clear(); slot_to_expert_.assign(cap, -1); } @@ -87,6 +103,42 @@ class MeshSlotPool { MeshSlotPool(const MeshSlotPool&) = delete; MeshSlotPool& operator=(const MeshSlotPool&) = delete; + // 允许移动（std::vector 需要） + MeshSlotPool(MeshSlotPool&& other) noexcept + : slots_(std::move(other.slots_)), + expert_to_slot_(std::move(other.expert_to_slot_)), + slot_to_expert_(std::move(other.slot_to_expert_)), + layer_idx_(other.layer_idx_), + tp_part_idx_(other.tp_part_idx_), + numa_node_(other.numa_node_), + cap_(other.cap_), + slot_bytes_(other.slot_bytes_), + gate_up_bytes_(other.gate_up_bytes_), + memory_(other.memory_), + total_bytes_(other.total_bytes_) { + other.memory_ = nullptr; + other.total_bytes_ = 0; + } + MeshSlotPool& operator=(MeshSlotPool&& other) noexcept { + if (this != &other) { + if (memory_) numa_free(memory_, total_bytes_); + slots_ = std::move(other.slots_); + expert_to_slot_ = std::move(other.expert_to_slot_); + slot_to_expert_ = std::move(other.slot_to_expert_); + layer_idx_ = other.layer_idx_; + tp_part_idx_ = other.tp_part_idx_; + numa_node_ = other.numa_node_; + cap_ = other.cap_; + slot_bytes_ = other.slot_bytes_; + gate_up_bytes_ = other.gate_up_bytes_; + memory_ = other.memory_; + total_bytes_ = other.total_bytes_; + other.memory_ = nullptr; + other.total_bytes_ = 0; + } + return *this; + } + // ===== 指针访问 ===== // 获取 slot 中 gate 矩阵的地址 diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index 3ba83682d..88e505b44 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -338,6 +338,16 @@ def _create_inference_wrapper( BaseMoEWrapper instance """ # MESH 插件：mesh_config 非 None 且 enabled 时走 MESH 路径 + # 自动从环境变量检测 MESH 配置（kt run CLI 路径） + if mesh_config is None: + import os + if os.environ.get("KT_ENABLE_MESH", "0") == "1": + from .utils.mesh.config import MeshConfig + mesh_config = MeshConfig.from_env() + else: + # SGLang spawn 不继承 KT_* env vars，回退到配置文件 + from .utils.mesh.config import MeshConfig + mesh_config = MeshConfig.from_file() if mesh_config is not None and getattr(mesh_config, "enabled", False): from .utils.mesh.wrapper import MeshMoEWrapper return MeshMoEWrapper( diff --git a/kt-kernel/python/utils/mesh/config.py b/kt-kernel/python/utils/mesh/config.py index 16ed2baa1..4b7c12231 100644 --- a/kt-kernel/python/utils/mesh/config.py +++ b/kt-kernel/python/utils/mesh/config.py @@ -21,10 +21,15 @@ from __future__ import annotations +import json import os from dataclasses import dataclass, field from typing import Optional +# 配置文件路径：当环境变量未传递到 scheduler 子进程时（SGLang spawn 不继承 KT_* env）， +# 主进程将配置写入此文件，scheduler 从文件读取 +MESH_CONFIG_FILE = "/tmp/kt_mesh_config.json" + @dataclass class MeshConfig: @@ -106,6 +111,37 @@ def _env_float(key: str, default: float = 0.0) -> float: weight_type=os.environ.get("KT_MESH_WEIGHT_TYPE", "amxint4").lower(), ) + def to_file(self, path: str = MESH_CONFIG_FILE) -> None: + """将配置写入 JSON 文件，供 scheduler 子进程读取。""" + data = { + "enabled": self.enabled, + "cap": self.cap, + "num_gpu_experts": self.num_gpu_experts, + "max_deferred_per_token": self.max_deferred_per_token, + "decode_front_layers": self.decode_front_layers, + "decode_front_layer_cap": self.decode_front_layer_cap, + "total_layers": self.total_layers, + "prefill_window": self.prefill_window, + "heat_gamma": self.heat_gamma, + "heat_beta": self.heat_beta, + "markov_alpha": self.markov_alpha, + "markov_topk": self.markov_topk, + "lookahead_weight": self.lookahead_weight, + "weight_type": self.weight_type, + } + with open(path, "w") as f: + json.dump(data, f) + + @classmethod + def from_file(cls, path: str = MESH_CONFIG_FILE) -> Optional["MeshConfig"]: + """从 JSON 文件读取配置。文件不存在时返回 None。""" + try: + with open(path, "r") as f: + data = json.load(f) + return cls(**data) + except (FileNotFoundError, json.JSONDecodeError, TypeError): + return None + def validate(self) -> None: """校验配置合法性。""" if not self.enabled: @@ -135,6 +171,7 @@ def to_cpp(self): ) from e cpp_cfg = CppMeshConfig() + from kt_kernel_ext.mesh import WeightType cpp_cfg.enabled = self.enabled cpp_cfg.cap = self.cap cpp_cfg.num_gpu_experts = self.num_gpu_experts @@ -150,9 +187,9 @@ def to_cpp(self): cpp_cfg.lookahead_weight = self.lookahead_weight # weight_type 枚举 if self.weight_type == "amxint4": - cpp_cfg.weight_type = 0 # WeightType::AMXINT4 + cpp_cfg.weight_type = WeightType.AMXINT4 else: - cpp_cfg.weight_type = 1 # WeightType::BF16 + cpp_cfg.weight_type = WeightType.BF16 cpp_cfg.hidden_size = self.hidden_size cpp_cfg.intermediate_size = self.intermediate_size cpp_cfg.expert_num = self.expert_num diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py index 67d8182e6..4d203fc9f 100644 --- a/kt-kernel/python/utils/mesh/wrapper.py +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -27,8 +27,15 @@ class MeshMoEWrapper: """MESH 模式的 MoE Wrapper。 委托给 AMXMoEWrapper 进行实际计算，MESH 只负责权重驻留管理。 + + ResidencyManager 是进程级单例：所有层共享同一个 manager， + 避免 40× 内存重复（每层都创建完整 40 层 pool）。 """ + # 进程级单例：key=numa_nodes tuple, value=MeshResidencyManager + _shared_residency = None + _shared_bootstrap_done = False + def __init__( self, layer_idx: int, @@ -95,31 +102,43 @@ def __init__( numa_nodes=numa_nodes, ) - # 2. 创建 ResidencyManager - self._residency = MeshResidencyManager() + # 2. 创建/复用 ResidencyManager（进程级单例，避免 40× 内存重复） if numa_nodes is None: numa_nodes = list(range(threadpool_count)) - self._residency.init(mesh_config, numa_nodes) - - # 3. 注入 GPU expert mask - if gpu_experts_mask is not None: - mask_list = gpu_experts_mask.cpu().tolist() - mask_int = [1 if x else 0 for x in mask_list] - self._residency.set_gpu_experts_mask(mask_int) - # 4. 注入文件布局（从 SafeTensorLoader 获取） - self._inject_file_layouts(weight_path, threadpool_count, num_experts, mesh_config) - - # 5. 启动阶段：读前 cap 个专家进 slot - self._residency.mark_layouts_set() - self._residency.bootstrap() + if MeshMoEWrapper._shared_residency is None: + # 第一层：创建并初始化共享 ResidencyManager + self._residency = MeshResidencyManager() + self._residency.init(mesh_config, numa_nodes) + + # 注入 GPU expert mask + if gpu_experts_mask is not None: + mask_list = gpu_experts_mask.cpu().tolist() + mask_int = [1 if x else 0 for x in mask_list] + self._residency.set_gpu_experts_mask(mask_int) + + # 注入文件布局（从 SafeTensorLoader 获取） + self._inject_file_layouts(weight_path, threadpool_count, num_experts, mesh_config) + + # 启动阶段：读前 cap 个专家进 slot + self._residency.mark_layouts_set() + self._residency.bootstrap() + MeshMoEWrapper._shared_residency = self._residency + MeshMoEWrapper._shared_bootstrap_done = True + logger.info( + "MeshMoEWrapper: shared ResidencyManager created at layer=%d, cap=%d, method=%s", + layer_idx, mesh_config.cap, method, + ) + else: + # 后续层：复用共享 ResidencyManager + self._residency = MeshMoEWrapper._shared_residency + logger.info( + "MeshMoEWrapper: layer=%d reuses shared ResidencyManager", + layer_idx, + ) self._mesh_config = mesh_config self._layer_idx = layer_idx - logger.info( - "MeshMoEWrapper initialized: layer=%d, cap=%d, method=%s", - layer_idx, mesh_config.cap, method, - ) def _inject_file_layouts( self, @@ -151,6 +170,22 @@ def _inject_file_layouts( "Please implement _inject_file_layouts based on SafeTensorLoader." ) + # ===== 属性委托 ===== + + def __getattr__(self, name: str): + """未在 MeshMoEWrapper 中定义的属性，委托给内部 AMXMoEWrapper。 + + 确保 submit_forward / sync_forward 等基类方法自动透传， + 无需逐个手写委托。__getattr__ 仅在正常属性查找失败时调用， + 不会覆盖本类已显式定义的方法（forward / load_weights 等）。 + """ + inner = self.__dict__.get("_inner") + if inner is not None: + return getattr(inner, name) + raise AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'" + ) + # ===== forward 委托 ===== def forward(self, *args, **kwargs): @@ -161,13 +196,15 @@ def forward(self, *args, **kwargs): """ return self._inner.forward(*args, **kwargs) - def load_weights(self): - """load_weights 委托给 AMXMoEWrapper。 + def load_weights(self, physical_to_logical_map_cpu=None): + """load_weights：委托给内部 AMXMoEWrapper 加载全量权重。 - MESH 模式下 AMXMoEWrapper.load_weights() 会走 mesh 分支（直接 return）， - 实际权重加载由 ResidencyManager.bootstrap() 完成。 + TODO: MESH 模式的最终目标是跳过全量加载，仅由 ResidencyManager 的 + slot pool 管理 cap 个专家（io_uring 按需加载）。但当前 io_uring 按需 + 加载和 AMX 内核 hook 尚未实现，self.moe 为 None 会导致 submit_forward + 崩溃。暂时先正常加载权重以保证推理可用，slot pool 作为额外开销。 """ - return self._inner.load_weights() + return self._inner.load_weights(physical_to_logical_map_cpu) # ===== MESH 专属接口 ===== diff --git a/kt-kernel/scripts/convert_cpu_weights.py b/kt-kernel/scripts/convert_cpu_weights.py index 9295d6724..c0c5935fd 100644 --- a/kt-kernel/scripts/convert_cpu_weights.py +++ b/kt-kernel/scripts/convert_cpu_weights.py @@ -818,30 +818,35 @@ def _convert_layer_experts(self, layer_idx: int, expert_ids: List[int]) -> Dict[ print(f" [Fused] tensor {p} shape: {tuple(w.shape)}") fused_tensors.append(w) - # fused_tensors[0] : down-like, [E, I, H] - # fused_tensors[1] : gate_up-like, [E, H, 2I] + # Qwen3.5 fused checkpoint layout (verified on 35B/397B): + # fused_tensors[0] : down_proj, [E, H, I] + # fused_tensors[1] : gate_up_proj, [E, 2I, H] + # NOTE: HF stores Linear weight as [out_features, in_features], so + # gate_up_proj = concat([gate[I,H], up[I,H]], dim=0) -> [2I, H] per expert + # down_proj = [H, I] per expert + # Stack over experts -> [E, 2I, H] and [E, H, I]. No transpose needed. down_fused = fused_tensors[0] gate_up_fused = fused_tensors[1] - # gate_up_fused: [E, H, 2I] -> [E, 2I, H] -> gate / up + # gate_up_fused: [E, 2I, H] -> gate / up (split along dim=1) if gate_up_fused.dim() != 3: raise ValueError( f"[Fused] Expect gate_up fused tensor to be 3D, got shape {tuple(gate_up_fused.shape)}" ) - E, H, twoI = gate_up_fused.shape + E, twoI, H = gate_up_fused.shape if twoI % 2 != 0: - raise ValueError(f"[Fused] gate_up last dim (2I) not even: {twoI}") + raise ValueError(f"[Fused] gate_up middle dim (2I) not even: {twoI}") I = twoI // 2 - gate_up_T = gate_up_fused.transpose(1, 2).contiguous() # [E, 2I, H] - gate_proj = gate_up_T[:, :I, :] # [E, I, H] - up_proj = gate_up_T[:, I:, :] # [E, I, H] + gate_proj = gate_up_fused[:, :I, :].contiguous() # [E, I, H] + up_proj = gate_up_fused[:, I:, :].contiguous() # [E, I, H] if down_fused.dim() != 3: raise ValueError(f"[Fused] Expect down fused tensor to be 3D, got shape {tuple(down_fused.shape)}") if down_fused.shape[0] != E: raise ValueError(f"[Fused] down_fused expert dim mismatch: {down_fused.shape[0]} vs gate_up {E}") - down_proj = down_fused.transpose(1, 2).contiguous() # [E, H, I] + # down_fused is already [E, H, I], no transpose needed + down_proj = down_fused.contiguous() # [E, H, I] del fused_tensors del gate_up_fused del down_fused From 89e97baaa383f1e5c7d61f0aa0e9da5d5b41b18d Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Sun, 21 Jun 2026 02:43:30 +0800 Subject: [PATCH 05/10] fix:mesh v2 a --- kt-kernel/ext_bindings.cpp | 34 ++- kt-kernel/operators/amx/moe.hpp | 57 ++++- kt-kernel/operators/mesh/mesh_handoff.hpp | 17 +- kt-kernel/operators/mesh/mesh_hook.hpp | 51 ++-- kt-kernel/operators/mesh/mesh_io_uring.hpp | 188 ++++++++++---- kt-kernel/operators/mesh/mesh_residency.hpp | 114 +++++++-- kt-kernel/python/utils/amx.py | 9 + kt-kernel/python/utils/mesh/residency.py | 38 ++- kt-kernel/python/utils/mesh/wrapper.py | 266 ++++++++++++++++++-- 9 files changed, 632 insertions(+), 142 deletions(-) diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index f57c3e5fa..880912feb 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -768,6 +768,9 @@ PYBIND11_MODULE(kt_kernel_ext, m) { // V4-Flash 2604B SwiGLU clamp limit (0.0 = disabled). See common.hpp. .def_readwrite("swiglu_limit", &GeneralMOEConfig::swiglu_limit) .def_readwrite("swiglu_alpha", &GeneralMOEConfig::swiglu_alpha) + // MESH 插件接入点 + .def_readwrite("mesh_enabled", &GeneralMOEConfig::mesh_enabled) + .DEF_PTR_PROPERTY(GeneralMOEConfig, mesh_residency) ; @@ -1040,7 +1043,20 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def_readwrite("down_offset", &mesh::ExpertFileLayout::down_offset) .def_readwrite("gate_bytes", &mesh::ExpertFileLayout::gate_bytes) .def_readwrite("up_bytes", &mesh::ExpertFileLayout::up_bytes) - .def_readwrite("down_bytes", &mesh::ExpertFileLayout::down_bytes); + .def_readwrite("down_bytes", &mesh::ExpertFileLayout::down_bytes) + // A3 fix: 暴露 AMXINT4 专用 scale/mins 字段 + .def_readwrite("gate_scale_offset", &mesh::ExpertFileLayout::gate_scale_offset) + .def_readwrite("up_scale_offset", &mesh::ExpertFileLayout::up_scale_offset) + .def_readwrite("down_scale_offset", &mesh::ExpertFileLayout::down_scale_offset) + .def_readwrite("gate_scale_bytes", &mesh::ExpertFileLayout::gate_scale_bytes) + .def_readwrite("up_scale_bytes", &mesh::ExpertFileLayout::up_scale_bytes) + .def_readwrite("down_scale_bytes", &mesh::ExpertFileLayout::down_scale_bytes) + .def_readwrite("gate_mins_offset", &mesh::ExpertFileLayout::gate_mins_offset) + .def_readwrite("up_mins_offset", &mesh::ExpertFileLayout::up_mins_offset) + .def_readwrite("down_mins_offset", &mesh::ExpertFileLayout::down_mins_offset) + .def_readwrite("gate_mins_bytes", &mesh::ExpertFileLayout::gate_mins_bytes) + .def_readwrite("up_mins_bytes", &mesh::ExpertFileLayout::up_mins_bytes) + .def_readwrite("down_mins_bytes", &mesh::ExpertFileLayout::down_mins_bytes); // MeshResidencyManager py::class_>( @@ -1049,7 +1065,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def("init", &mesh::MeshResidencyManager::init, py::arg("config"), py::arg("numa_nodes")) .def("set_file_layout", &mesh::MeshResidencyManager::set_file_layout, - py::arg("tp_part_idx"), py::arg("expert_id"), py::arg("layout")) + py::arg("layer_idx"), py::arg("tp_part_idx"), py::arg("expert_id"), py::arg("layout")) .def("set_gpu_experts_mask", [](mesh::MeshResidencyManager& mgr, py::list mask_list) { std::vector mask; @@ -1072,7 +1088,19 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def("on_decode_token_end", &mesh::MeshResidencyManager::on_decode_token_end) .def("config", [](mesh::MeshResidencyManager& mgr) -> const mesh::MeshConfig& { return mgr.config(); - }, py::return_value_policy::reference); + }, py::return_value_policy::reference) + // A2: 获取原始 C++ 指针，用于注入 MOEConfig.mesh_residency + .def("raw_ptr", [](mesh::MeshResidencyManager& mgr) -> uintptr_t { + return reinterpret_cast(&mgr); + }) + // A6: io_uring CQE 到达后标记 slot 已缓存 + .def("mark_slot_cached", [](mesh::MeshResidencyManager& mgr, int layer, int tp, int slot_idx) { + mgr.pool(layer, tp).mark_cached(slot_idx); + }, py::arg("layer"), py::arg("tp"), py::arg("slot_idx")) + // A6: 查询专家是否已缓存 + .def("is_cached", [](mesh::MeshResidencyManager& mgr, int layer, int tp, int expert_id) -> bool { + return mgr.pool(layer, tp).is_cached(expert_id); + }, py::arg("layer"), py::arg("tp"), py::arg("expert_id")); // 注册 hook 函数指针（让 mesh_hook.hpp 的 inline 函数能调用 ResidencyManager） mesh::hook::HookRegistry registry; diff --git a/kt-kernel/operators/amx/moe.hpp b/kt-kernel/operators/amx/moe.hpp index b83604fb7..a23711f66 100644 --- a/kt-kernel/operators/amx/moe.hpp +++ b/kt-kernel/operators/amx/moe.hpp @@ -15,6 +15,7 @@ // #define FORWARD_TIME_REPORT #include "moe_base.hpp" +#include "../mesh/mesh_hook.hpp" template class AMX_MOE_TP : public AMX_MOE_BASE> { @@ -181,32 +182,64 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; auto& ba = gate_up_ba_[expert_idx]; - auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; - // MESH hook: mesh_enabled 时 bb 的底层指针已由 MESH slot 池提供 - // gate_bb_/up_bb_ 在 load_weights() mesh 分支中已绑定为 slot buffer 的包装 - // MESH 通过覆盖 slot 内存实现专家切换，bb 指针本身不变 + // A1-call: MESH hook — 如果 mesh_residency 已设置，尝试从 slot 池获取权重指针 + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = do_up + ? mesh::hook::get_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) + : mesh::hook::get_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + if (slot_ptr) { + // MESH 路径：用 slot 指针创建临时 BufferB + auto bb = std::make_shared( + config_.intermediate_size, config_.hidden_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } } else { - amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + // 原版路径 + auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } } } void do_down_gemm(int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; auto& ba = down_ba_[expert_idx]; - auto& bb = down_bb_[expert_idx]; auto& bc = down_bc_[expert_idx]; - // MESH hook: 同上，down_bb_ 已绑定为 MESH slot buffer + // A1-call: MESH hook + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = mesh::hook::get_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + if (slot_ptr) { + // MESH 路径 + auto bb = std::make_shared( + config_.hidden_size, config_.intermediate_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } } else { - amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + // 原版路径 + auto& bb = down_bb_[expert_idx]; + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } } } void load_weights() { diff --git a/kt-kernel/operators/mesh/mesh_handoff.hpp b/kt-kernel/operators/mesh/mesh_handoff.hpp index 80b992832..286d41422 100644 --- a/kt-kernel/operators/mesh/mesh_handoff.hpp +++ b/kt-kernel/operators/mesh/mesh_handoff.hpp @@ -42,7 +42,7 @@ class MeshHandoff { * @param io io_uring 读取器 * @param scorer 驱逐评分器 * @param pools [layer][tp] 的 slot 池 - * @param layouts [tp][expert] 的文件布局 + * @param layouts A4 fix: [layer][tp][expert] 的文件布局 * @param move_gpu_expert_to_gpu 搬运 GPU 专家到 GPU 的回调 * @param get_dst_ptrs 获取 slot buffer 指针的回调 */ @@ -52,7 +52,7 @@ class MeshHandoff { MeshIoUring& io, EvictionScorer& scorer, std::vector>& pools, - const std::vector>& layouts, + const std::vector>>& layouts, std::function move_gpu_expert_to_gpu, std::function(int, int, int)> get_dst_ptrs) { int ge = config.num_gpu_experts; @@ -91,7 +91,7 @@ class MeshHandoff { void move_gpu_experts_and_refill( const MeshConfig& config, std::vector>& pools, - const std::vector>& layouts, + const std::vector>>& layouts, EvictionScorer& scorer, MeshIoUring& io, std::function move_gpu_expert_to_gpu, @@ -132,11 +132,18 @@ class MeshHandoff { new_expert_id = refill_experts[slot_idx]; auto ptrs = get_dst_ptrs(layer, tp, new_expert_id); // 提交 io_uring 读 - const auto& layout = layouts[tp][new_expert_id]; + // A4 fix: layouts 改为 [layer][tp][expert] 3D + // A6: 传入 layer_idx 和 on_complete 回调 + const auto& layout = layouts[layer][tp][new_expert_id]; + MeshSlotPool& pool_ref = pools[layer][tp]; io.submit_load(new_expert_id, tp, layout, ptrs[0], ptrs[1], ptrs[2], nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, - ReadPriority::Demand); + ReadPriority::Demand, + /*layer_idx=*/layer, + /*on_complete=*/[&pool_ref, new_expert_id](int, int, int) { + pool_ref.mark_cached(new_expert_id); + }); // 覆盖 slot pool.overwrite(slot_idx, new_expert_id); } diff --git a/kt-kernel/operators/mesh/mesh_hook.hpp b/kt-kernel/operators/mesh/mesh_hook.hpp index f170d826d..746125035 100644 --- a/kt-kernel/operators/mesh/mesh_hook.hpp +++ b/kt-kernel/operators/mesh/mesh_hook.hpp @@ -19,32 +19,47 @@ namespace hook { // 前向声明，避免强制 include mesh_residency.hpp class MeshResidencyManager; +// ===== 注册函数指针（由 ext_bindings.cpp 在模块初始化时调用）===== + +// 这些函数指针用于解耦 mesh_hook.hpp 和 mesh_residency.hpp +// 避免原版 KT 文件 include mesh_residency.hpp 带来的编译依赖 +struct HookRegistry { + void* (*get_gate_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*get_up_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*get_down_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; +}; + +inline HookRegistry& get_registry() { + static HookRegistry registry; + return registry; +} + +// 注册 hook 函数（ext_bindings.cpp 调用） +inline void register_hooks(HookRegistry r) { + get_registry() = r; +} + // ===== 权重指针重定向 hook ===== // 获取 gate 矩阵指针（MESH 启用时返回 slot buffer，否则返回 nullptr 走原版） // mesh_mgr 为 nullptr 时直接返回 nullptr，编译期可优化 inline void* get_gate_bb(void* mesh_mgr, int layer, int tp, int expert_id) { if (!mesh_mgr) return nullptr; - // 实际实现通过 reinterpret_cast 调用 MeshResidencyManager::get_gate_ptr - // 这里用函数指针避免头文件循环依赖 - using GetPtrFn = void* (*)(void*, int, int, int); - static GetPtrFn fn = nullptr; + auto fn = get_registry().get_gate_ptr; if (!fn) return nullptr; return fn(mesh_mgr, layer, tp, expert_id); } inline void* get_up_bb(void* mesh_mgr, int layer, int tp, int expert_id) { if (!mesh_mgr) return nullptr; - using GetPtrFn = void* (*)(void*, int, int, int); - static GetPtrFn fn = nullptr; + auto fn = get_registry().get_up_ptr; if (!fn) return nullptr; return fn(mesh_mgr, layer, tp, expert_id); } inline void* get_down_bb(void* mesh_mgr, int layer, int tp, int expert_id) { if (!mesh_mgr) return nullptr; - using GetPtrFn = void* (*)(void*, int, int, int); - static GetPtrFn fn = nullptr; + auto fn = get_registry().get_down_ptr; if (!fn) return nullptr; return fn(mesh_mgr, layer, tp, expert_id); } @@ -55,25 +70,5 @@ inline bool is_mesh_enabled(void* mesh_mgr) { return mesh_mgr != nullptr; } -// ===== 注册函数指针（由 ext_bindings.cpp 在模块初始化时调用）===== - -// 这些函数指针用于解耦 mesh_hook.hpp 和 mesh_residency.hpp -// 避免原版 KT 文件 include mesh_residency.hpp 带来的编译依赖 -struct HookRegistry { - void* (*get_gate_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; - void* (*get_up_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; - void* (*get_down_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; -}; - -inline HookRegistry& get_registry() { - static HookRegistry registry; - return registry; -} - -// 注册 hook 函数（ext_bindings.cpp 调用） -inline void register_hooks(HookRegistry r) { - get_registry() = r; -} - } // namespace hook } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_io_uring.hpp b/kt-kernel/operators/mesh/mesh_io_uring.hpp index ba3ef7794..d00d4fd43 100644 --- a/kt-kernel/operators/mesh/mesh_io_uring.hpp +++ b/kt-kernel/operators/mesh/mesh_io_uring.hpp @@ -12,8 +12,10 @@ */ #pragma once +#include #include #include +#include #include #include // liburing.h 定义了 BLOCK_SIZE 等宏，会污染全局命名空间， @@ -62,24 +64,32 @@ class MeshIoUring { // ===== Scale Cache 预加载（AMXINT4 专用）===== - // 启动时一次性把所有专家的 scale 数据读入 NUMA 本地 buffer + // 启动时一次性把某层所有专家的 scale 数据读入 NUMA 本地 buffer + // A4 fix: 改为按层预加载，因为每层的 scale 数据不同 // expert_num 个专家，每个 scale_bytes 字节 void preload_scale_cache(int expert_num, int tp_count, int numa_node, - const std::vector>& layouts) { + const std::vector>& layouts_tp, + int layer_idx) { if (scale_cache_loaded_) return; + // 确保 scale_cache_ 有足够层 + if ((int)scale_cache_.size() <= layer_idx) { + scale_cache_.resize(layer_idx + 1); + } + + auto& layer_cache = scale_cache_[layer_idx]; // 计算每个 TP 分片的 scale 总大小 for (int tp = 0; tp < tp_count; tp++) { ScaleCacheTP cache; size_t total_gate = 0, total_up = 0, total_down = 0; size_t total_gate_mins = 0, total_up_mins = 0, total_down_mins = 0; for (int e = 0; e < expert_num; e++) { - total_gate += layouts[tp][e].gate_scale_bytes; - total_up += layouts[tp][e].up_scale_bytes; - total_down += layouts[tp][e].down_scale_bytes; - total_gate_mins += layouts[tp][e].gate_mins_bytes; - total_up_mins += layouts[tp][e].up_mins_bytes; - total_down_mins += layouts[tp][e].down_mins_bytes; + total_gate += layouts_tp[tp][e].gate_scale_bytes; + total_up += layouts_tp[tp][e].up_scale_bytes; + total_down += layouts_tp[tp][e].down_scale_bytes; + total_gate_mins += layouts_tp[tp][e].gate_mins_bytes; + total_up_mins += layouts_tp[tp][e].up_mins_bytes; + total_down_mins += layouts_tp[tp][e].down_mins_bytes; } cache.gate_scale = numa_alloc_onnode(total_gate, numa_node); cache.up_scale = numa_alloc_onnode(total_up, numa_node); @@ -98,7 +108,7 @@ class MeshIoUring { size_t off_g = 0, off_u = 0, off_d = 0; size_t off_gm = 0, off_um = 0, off_dm = 0; for (int e = 0; e < expert_num; e++) { - const auto& layout = layouts[tp][e]; + const auto& layout = layouts_tp[tp][e]; sync_read_to(layout.fd, layout.gate_scale_offset, layout.gate_scale_bytes, (char*)cache.gate_scale + off_g); sync_read_to(layout.fd, layout.up_scale_offset, layout.up_scale_bytes, @@ -118,17 +128,21 @@ class MeshIoUring { off_um += layout.up_mins_bytes; off_dm += layout.down_mins_bytes; } - scale_cache_.push_back(std::move(cache)); + layer_cache.push_back(std::move(cache)); } scale_cache_loaded_ = true; } // 从 scale cache 中 memcpy 某个专家的 scale 到目标 buffer - void copy_scale_from_cache(int tp_part_idx, int expert_id, + // A4 fix: 加 layer_idx 参数 + void copy_scale_from_cache(int layer_idx, int tp_part_idx, int expert_id, void* gate_scale_dst, void* up_scale_dst, void* down_scale_dst, void* gate_mins_dst, void* up_mins_dst, void* down_mins_dst, const std::vector& layouts_tp) { - // 需要外部传入该 TP 的 layouts 来计算偏移 + if ((int)scale_cache_.size() <= layer_idx) return; + auto& layer_cache = scale_cache_[layer_idx]; + if ((int)layer_cache.size() <= tp_part_idx) return; + size_t off_g = 0, off_u = 0, off_d = 0; size_t off_gm = 0, off_um = 0, off_dm = 0; for (int e = 0; e < expert_id; e++) { @@ -140,7 +154,7 @@ class MeshIoUring { off_dm += layouts_tp[e].down_mins_bytes; } const auto& layout = layouts_tp[expert_id]; - const auto& cache = scale_cache_[tp_part_idx]; + const auto& cache = layer_cache[tp_part_idx]; memcpy(gate_scale_dst, (char*)cache.gate_scale + off_g, layout.gate_scale_bytes); memcpy(up_scale_dst, (char*)cache.up_scale + off_u, layout.up_scale_bytes); memcpy(down_scale_dst, (char*)cache.down_scale + off_d, layout.down_scale_bytes); @@ -153,70 +167,129 @@ class MeshIoUring { // ===== io_uring 异步读 ===== + // A6 fix: PendingRequest 跟踪每个专家读取的完成状态 + // 一个专家的读取可能涉及 3-9 个 SQE，所有 SQE 完成后触发 on_complete + struct PendingRequest { + int layer_idx = -1; + int tp_part_idx = -1; + int expert_id = -1; + int total_reqs = 0; + std::atomic completed_reqs{0}; + std::function on_complete; // (layer, tp, expert) + }; + // 提交一个专家的读取请求 // scale_cache_loaded_=true: 发 3 个请求（只读 weight），scale 从 cache memcpy // scale_cache_loaded_=false: 发 9 个请求（weight + scale + mins 全读） + // A6 fix: 加 layer_idx 和 on_complete 回调，CQE 完成后触发 mark_cached void submit_load(int expert_id, int tp_part_idx, const ExpertFileLayout& layout, void* gate_dst, void* up_dst, void* down_dst, void* gate_scale_dst, void* up_scale_dst, void* down_scale_dst, void* gate_mins_dst, void* up_mins_dst, void* down_mins_dst, - ReadPriority priority) { + ReadPriority priority, + int layer_idx = -1, + std::function on_complete = nullptr) { int n_reqs = scale_cache_loaded_ ? 3 : 9; - std::vector sqes; - sqes.reserve(n_reqs); - // 提交 SQE + // A6: 创建 PendingRequest 跟踪完成状态 + auto* pending = new PendingRequest{ + layer_idx, tp_part_idx, expert_id, n_reqs, 0, std::move(on_complete) + }; + + // 提交 SQE，user_data 指向 pending auto* sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, gate_dst, layout.gate_bytes, layout.gate_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, up_dst, layout.up_bytes, layout.up_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, down_dst, layout.down_bytes, layout.down_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); if (!scale_cache_loaded_) { sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, gate_scale_dst, layout.gate_scale_bytes, layout.gate_scale_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, up_scale_dst, layout.up_scale_bytes, layout.up_scale_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, down_scale_dst, layout.down_scale_bytes, layout.down_scale_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, gate_mins_dst, layout.gate_mins_bytes, layout.gate_mins_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, up_mins_dst, layout.up_mins_bytes, layout.up_mins_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, down_mins_dst, layout.down_mins_bytes, layout.down_mins_offset); - sqes.push_back(sqe); + io_uring_sqe_set_data(sqe, pending); } - // 设置 user_data 用于 CQE 回调识别 - for (auto* s : sqes) { - io_uring_sqe_set_data(s, nullptr); - } io_uring_submit(&ring_); } + // A6: 处理已完成的 CQE，触发 on_complete 回调 + // 非阻塞：只处理当前已有的 CQE，不等待 + void process_cqes() { + unsigned head; + unsigned count = 0; + io_uring_cqe* cqe; + + io_uring_for_each_cqe(&ring_, head, cqe) { + auto* pending = static_cast(io_uring_cqe_get_data(cqe)); + if (pending) { + // 检查读取错误 + if (cqe->res < 0) { + fprintf(stderr, "[MESH] io_uring read error: expert=%d tp=%d res=%d (%s)\n", + pending->expert_id, pending->tp_part_idx, cqe->res, strerror(-cqe->res)); + } + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + // 所有 SQE 完成，触发回调 + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; + } + } + count++; + } + if (count > 0) { + io_uring_cq_advance(&ring_, count); + } + } + // 阻塞等待某专家的读请求完成（通过 CQE 计数） + // A6: 改为处理 CQE 并触发回调 void wait_expert(int expert_id, int n_reqs) { for (int i = 0; i < n_reqs; i++) { io_uring_cqe* cqe; io_uring_wait_cqe(&ring_, &cqe); + auto* pending = static_cast(io_uring_cqe_get_data(cqe)); + if (pending) { + if (cqe->res < 0) { + fprintf(stderr, "[MESH] io_uring read error in wait_expert: expert=%d res=%d (%s)\n", + pending->expert_id, cqe->res, strerror(-cqe->res)); + } + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; + } + } io_uring_cqe_seen(&ring_, cqe); } } @@ -238,26 +311,48 @@ class MeshIoUring { } // 提交并等待全部完成（同步接口，启动阶段用） + // A6: 处理 CQE 并触发 on_complete 回调 void submit_and_wait() { io_uring_submit(&ring_); - // 等待所有 inflight 完成 - unsigned head; - unsigned count = 0; - io_uring_cqe* cqe; + // 等待所有 inflight 完成，处理回调 while (true) { + unsigned head; + unsigned count = 0; + io_uring_cqe* cqe; io_uring_for_each_cqe(&ring_, head, cqe) { + auto* pending = static_cast(io_uring_cqe_get_data(cqe)); + if (pending) { + if (cqe->res < 0) { + fprintf(stderr, "[MESH] io_uring read error in submit_and_wait: expert=%d res=%d (%s)\n", + pending->expert_id, cqe->res, strerror(-cqe->res)); + } + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; + } + } count++; } - if (count == 0) break; + if (count == 0) { + // 没有更多 CQE，检查是否还有 inflight SQE + unsigned inflight = io_uring_sq_ready(&ring_); + if (inflight == 0) break; + // 还有 inflight，等待一个 CQE + io_uring_cqe* wait_cqe; + io_uring_wait_cqe(&ring_, &wait_cqe); + continue; + } io_uring_cq_advance(&ring_, count); - count = 0; } } private: io_uring ring_; - // Scale Cache（每 TP 一个） + // Scale Cache（A4 fix: [layer][tp] 每层独立） struct ScaleCacheTP { void* gate_scale = nullptr; void* up_scale = nullptr; @@ -272,7 +367,8 @@ class MeshIoUring { size_t up_mins_total = 0; size_t down_mins_total = 0; }; - std::vector scale_cache_; + // A4 fix: [layer][tp] 的 scale cache + std::vector> scale_cache_; bool scale_cache_loaded_ = false; // 同步读取（启动阶段用，pread + O_DIRECT） @@ -289,13 +385,15 @@ class MeshIoUring { } void release_scale_cache() { - for (auto& c : scale_cache_) { - if (c.gate_scale) numa_free(c.gate_scale, c.gate_scale_total); - if (c.up_scale) numa_free(c.up_scale, c.up_scale_total); - if (c.down_scale) numa_free(c.down_scale, c.down_scale_total); - if (c.gate_mins) numa_free(c.gate_mins, c.gate_mins_total); - if (c.up_mins) numa_free(c.up_mins, c.up_mins_total); - if (c.down_mins) numa_free(c.down_mins, c.down_mins_total); + for (auto& layer_cache : scale_cache_) { + for (auto& c : layer_cache) { + if (c.gate_scale) numa_free(c.gate_scale, c.gate_scale_total); + if (c.up_scale) numa_free(c.up_scale, c.up_scale_total); + if (c.down_scale) numa_free(c.down_scale, c.down_scale_total); + if (c.gate_mins) numa_free(c.gate_mins, c.gate_mins_total); + if (c.up_mins) numa_free(c.up_mins, c.up_mins_total); + if (c.down_mins) numa_free(c.down_mins, c.down_mins_total); + } } scale_cache_.clear(); } diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp index 351d4fd17..25935e88c 100644 --- a/kt-kernel/operators/mesh/mesh_residency.hpp +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -103,13 +103,20 @@ class MeshResidencyManager { // ===== 文件布局注入 ===== - void set_file_layout(int tp_part_idx, int expert_id, const ExpertFileLayout& layout) { + // A4 fix: layouts_ 改为 [layer][tp][expert] 3D， + // 因为每层的专家权重在 safetensors 文件中的偏移不同。 + void set_file_layout(int layer_idx, int tp_part_idx, int expert_id, + const ExpertFileLayout& layout) { if (layouts_.empty()) { - layouts_.resize(config_.tp_count, std::vector(config_.expert_num)); + layouts_.resize(config_.total_layers, + std::vector>( + config_.tp_count, + std::vector(config_.expert_num))); } - if (tp_part_idx >= 0 && tp_part_idx < (int)layouts_.size() && - expert_id >= 0 && expert_id < (int)layouts_[tp_part_idx].size()) { - layouts_[tp_part_idx][expert_id] = layout; + if (layer_idx >= 0 && layer_idx < (int)layouts_.size() && + tp_part_idx >= 0 && tp_part_idx < (int)layouts_[layer_idx].size() && + expert_id >= 0 && expert_id < (int)layouts_[layer_idx][tp_part_idx].size()) { + layouts_[layer_idx][tp_part_idx][expert_id] = layout; } } @@ -139,10 +146,12 @@ class MeshResidencyManager { return; } - // 预加载 Scale Cache（AMXINT4 专用） + // 预加载 Scale Cache（AMXINT4 专用）— 每层独立预加载 if (config_.weight_type == WeightType::AMXINT4) { - io_->preload_scale_cache(config_.expert_num, config_.tp_count, - numa_nodes_[0], layouts_); + for (int l = 0; l < config_.total_layers; l++) { + io_->preload_scale_cache(config_.expert_num, config_.tp_count, + numa_nodes_[0], layouts_[l], l); + } } // 每层每 TP 读前 cap 个专家 @@ -152,7 +161,7 @@ class MeshResidencyManager { for (int e = 0; e < config_.cap && e < config_.expert_num; e++) { if (is_gpu_expert(e)) continue; // GPU expert 跳过 - const auto& layout = layouts_[tp][e]; + const auto& layout = layouts_[l][tp][e]; void* gate_dst = pool.gate_ptr(e); void* up_dst = pool.up_ptr(e); void* down_dst = pool.down_ptr(e); @@ -161,9 +170,14 @@ class MeshResidencyManager { pool.bind(e, e); // 提交 io_uring 读 + // A6: 传入 layer_idx 和 on_complete 回调，完成后 mark_cached io_->submit_load(e, tp, layout, gate_dst, up_dst, down_dst, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, - ReadPriority::Demand); + ReadPriority::Demand, + /*layer_idx=*/l, + /*on_complete=*/[this, l, tp, e](int, int, int) { + pools_[l][tp].mark_cached(e); + }); } } } @@ -222,10 +236,12 @@ class MeshResidencyManager { pools_, layouts_, /*move_gpu_expert_to_gpu=*/[](int, int) {}, /*get_dst_ptrs=*/[this](int l, int tp, int e) { + // A5 fix: 用 expert_*_ptr 按 expert_id 查 slot， + // 不能把 expert_id 当 slot_idx 传给 gate_ptr（eid>=cap 时越界） return std::vector{ - pools_[l][tp].gate_ptr(e), - pools_[l][tp].up_ptr(e), - pools_[l][tp].down_ptr(e)}; + pools_[l][tp].expert_gate_ptr(e), + pools_[l][tp].expert_up_ptr(e), + pools_[l][tp].expert_down_ptr(e)}; }); } @@ -248,14 +264,71 @@ class MeshResidencyManager { const std::vector& topk, const std::vector& scores, int tp_part_idx) { + // A7: 先处理已完成的 io_uring CQE，触发 mark_cached 回调 + io_->process_cqes(); + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; - return decode_->split(topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, + auto result = decode_->split(topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, [this, &pool, layer_idx, tp_part_idx](int eid) { + // A5 fix: 用 expert_*_ptr 按 expert_id 查 slot， + // 不能把 expert_id 当 slot_idx 传给 gate_ptr（eid>=cap 时越界） return std::vector{ - pool.gate_ptr(eid), - pool.up_ptr(eid), - pool.down_ptr(eid)}; + pool.expert_gate_ptr(eid), + pool.expert_up_ptr(eid), + pool.expert_down_ptr(eid)}; }); + + // A7: 排空调度器队列，提交 io_uring 读取 + drain_and_submit(); + + return result; + } + + /** + * @brief A7: 排空调度器队列并提交 io_uring 读取 + * + * 把 MeshScheduler 优先队列中的 ScheduledRequest 取出， + * 为每个请求找到对应的 ExpertFileLayout 并提交给 MeshIoUring。 + * 完成后触发 mark_cached + 原始 on_complete 回调。 + */ + void drain_and_submit() { + auto requests = scheduler_->drain_all(); + for (auto& req : requests) { + // 边界检查 + if (req.layer_idx < 0 || req.layer_idx >= (int)layouts_.size()) continue; + if (req.tp_part_idx < 0 || req.tp_part_idx >= (int)layouts_[req.layer_idx].size()) continue; + if (req.expert_id < 0 || req.expert_id >= (int)layouts_[req.layer_idx][req.tp_part_idx].size()) continue; + // layouts_ 为空（memory-test 模式）时跳过 + if (layouts_.empty()) continue; + + const auto& layout = layouts_[req.layer_idx][req.tp_part_idx][req.expert_id]; + + // 捕获原始 on_complete 回调 + auto orig_on_complete = std::move(req.on_complete); + int layer = req.layer_idx; + int tp = req.tp_part_idx; + int expert = req.expert_id; + + // 绑定 slot（如果尚未绑定） + MeshSlotPool& pool = pools_[layer][tp]; + if (!pool.is_cached(expert)) { + // 需要先驱逐一个 slot（如果池满） + // B1: 在线驱逐尚未完全实现，这里先尝试 bind + // 如果池未满，bind 到空闲 slot + // 如果池满，需要驱逐（B1 修复后完善） + } + + io_->submit_load( + req.expert_id, req.tp_part_idx, layout, + req.gate_dst, req.up_dst, req.down_dst, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, // scale 从 cache memcpy + req.priority, + /*layer_idx=*/layer, + /*on_complete=*/[this, layer, tp, expert, orig_on_complete](int, int, int) { + pools_[layer][tp].mark_cached(expert); + if (orig_on_complete) orig_on_complete(); + }); + } } /** @@ -277,7 +350,8 @@ class MeshResidencyManager { MeshDecode& decode() { return *decode_; } EvictionScorer& scorer() { return *scorer_; } MeshSlotPool& pool(int layer, int tp) { return pools_[layer][tp]; } - const std::vector>& layouts() const { return layouts_; } + // A4 fix: layouts_ 改为 [layer][tp][expert] 3D + const std::vector>>& layouts() const { return layouts_; } private: MeshConfig config_; @@ -287,8 +361,8 @@ class MeshResidencyManager { // [layer][tp] 的 slot 池 std::vector> pools_; - // [tp][expert] 的文件布局 - std::vector> layouts_; + // A4 fix: [layer][tp][expert] 的文件布局（每层偏移不同） + std::vector>> layouts_; // GPU expert mask std::vector gpu_experts_mask_; diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 1b04e4def..4e313ca33 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -243,6 +243,8 @@ def __init__( max_deferred_experts_per_token: Optional[int] = None, method: str = "AMXINT4", numa_nodes: Optional[List[int]] = None, + mesh_enabled: bool = False, + mesh_residency_ptr: int = 0, ): """ Initialize AMX MoE Wrapper. @@ -317,6 +319,10 @@ def __init__( self.up_scales = None self.down_scales = None + # MESH 插件接入点 + self.mesh_enabled = mesh_enabled + self.mesh_residency_ptr = mesh_residency_ptr + def load_weights_from_tensors( self, gate_proj: torch.Tensor, @@ -464,6 +470,9 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor): moe_config.layer_idx = self.layer_idx moe_config.pool = self.cpu_infer.backend_ moe_config.max_len = self.chunked_prefill_size + # A2: 注入 MESH 配置 + moe_config.mesh_enabled = self.mesh_enabled + moe_config.mesh_residency = self.mesh_residency_ptr moe_config.gate_proj = gate_ptr moe_config.up_proj = up_ptr diff --git a/kt-kernel/python/utils/mesh/residency.py b/kt-kernel/python/utils/mesh/residency.py index 114557de7..b47e466cf 100644 --- a/kt-kernel/python/utils/mesh/residency.py +++ b/kt-kernel/python/utils/mesh/residency.py @@ -54,6 +54,7 @@ def init(self, config: MeshConfig, numa_nodes: List[int]) -> None: def set_file_layout( self, + layer_idx: int, tp_part_idx: int, expert_id: int, fd: int, @@ -76,9 +77,12 @@ def set_file_layout( up_mins_bytes: int = 0, down_mins_bytes: int = 0, ) -> None: - """注入单个专家在某个 TP 分片上的文件布局。 + """注入单个专家在某层某 TP 分片上的文件布局。 + + A4 fix: 加 layer_idx 参数，因为每层的专家权重在文件中偏移不同。 Args: + layer_idx: 层索引 tp_part_idx: TP 分片索引 expert_id: 专家 ID fd: O_DIRECT 打开的文件描述符 @@ -87,15 +91,29 @@ def set_file_layout( """ if not self._initialized: raise RuntimeError("Call init() before set_file_layout()") - self._mgr.set_file_layout( - tp_part_idx, expert_id, fd, - gate_offset, up_offset, down_offset, - gate_bytes, up_bytes, down_bytes, - gate_scale_offset, up_scale_offset, down_scale_offset, - gate_scale_bytes, up_scale_bytes, down_scale_bytes, - gate_mins_offset, up_mins_offset, down_mins_offset, - gate_mins_bytes, up_mins_bytes, down_mins_bytes, - ) + # A3 fix: 构造 ExpertFileLayout 对象传给 C++，而非 21 个位置参数 + from kt_kernel_ext.mesh import ExpertFileLayout + layout = ExpertFileLayout() + layout.fd = fd + layout.gate_offset = gate_offset + layout.up_offset = up_offset + layout.down_offset = down_offset + layout.gate_bytes = gate_bytes + layout.up_bytes = up_bytes + layout.down_bytes = down_bytes + layout.gate_scale_offset = gate_scale_offset + layout.up_scale_offset = up_scale_offset + layout.down_scale_offset = down_scale_offset + layout.gate_scale_bytes = gate_scale_bytes + layout.up_scale_bytes = up_scale_bytes + layout.down_scale_bytes = down_scale_bytes + layout.gate_mins_offset = gate_mins_offset + layout.up_mins_offset = up_mins_offset + layout.down_mins_offset = down_mins_offset + layout.gate_mins_bytes = gate_mins_bytes + layout.up_mins_bytes = up_mins_bytes + layout.down_mins_bytes = down_mins_bytes + self._mgr.set_file_layout(layer_idx, tp_part_idx, expert_id, layout) def set_gpu_experts_mask(self, mask: List[int]) -> None: """注入 GPU expert mask。 diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py index 4d203fc9f..9d7cf270d 100644 --- a/kt-kernel/python/utils/mesh/wrapper.py +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -100,6 +100,10 @@ def __init__( max_deferred_experts_per_token=max_deferred_experts_per_token, method=method, numa_nodes=numa_nodes, + # A2: mesh_enabled 暂不开启（避免 load_weights 跳过导致 gate_bb_ 为空崩溃） + # 仅注入 mesh_residency 指针，让 do_gate_up_gemm 的 hook 可被调用 + mesh_enabled=False, + mesh_residency_ptr=0, ) # 2. 创建/复用 ResidencyManager（进程级单例，避免 40× 内存重复） @@ -140,6 +144,10 @@ def __init__( self._mesh_config = mesh_config self._layer_idx = layer_idx + # A2: 注入 mesh_residency 指针到 inner AMXMoEWrapper + # 这样 do_gate_up_gemm/do_down_gemm 中的 mesh hook 可以调用 ResidencyManager + self._inner.mesh_residency_ptr = self._residency.raw_ptr() + def _inject_file_layouts( self, weight_path: str, @@ -147,27 +155,247 @@ def _inject_file_layouts( expert_num: int, mesh_config: MeshConfig, ) -> None: - """从 SafeTensorLoader 获取文件布局并注入 ResidencyManager。 + """从 safetensors 文件解析专家权重偏移并注入 ResidencyManager。 + + A4 实现：解析 safetensors 文件头获取每个 tensor 的绝对文件偏移， + 为每层每 TP 每专家构造 ExpertFileLayout 并注入。 + + 支持两种格式： + - AMXINT4 NUMA-sharded: blk.{L}.ffn_{gate,up,down}_exps.{E}.numa.{N}.{weight,scale} + - BF16 packed: model.layers.{L}.mlp.experts.{gate_up_proj,down_proj} + """ + import struct + import json + + total_layers = mesh_config.total_layers + weight_type = mesh_config.weight_type + + # 1. 遍历所有 safetensors 文件，构建 tensor_name -> (file_path, abs_offset, bytes) 映射 + tensor_map: dict[str, tuple[str, int, int]] = {} + for root, _, files in os.walk(weight_path): + for fname in sorted(files): + if not fname.endswith('.safetensors'): + continue + fpath = os.path.join(root, fname) + with open(fpath, 'rb') as f: + header_len_bytes = f.read(8) + header_len = struct.unpack(' None: + """注入 AMXINT4 NUMA-sharded 格式的文件布局。 + + Tensor 命名： + - blk.{L}.ffn_gate_exps.{E}.numa.{N}.weight / .scale + - blk.{L}.ffn_up_exps.{E}.numa.{N}.weight / .scale + - blk.{L}.ffn_down_exps.{E}.numa.{N}.weight / .scale + """ + # 收集所有需要打开的文件，用 O_DIRECT 打开 + file_fds: dict[str, int] = {} + + def get_fd(fpath: str) -> int: + if fpath not in file_fds: + # O_DIRECT 需要对齐读取，但 safetensors 偏移可能不对齐 + # 先用 O_RDONLY 打开，O_DIRECT 对齐问题在 io_uring 层处理 + fd = os.open(fpath, os.O_RDONLY) + file_fds[fpath] = fd + logger.debug("_inject_amxint4_layouts: opened %s as fd=%d", fpath, fd) + return file_fds[fpath] + + injected = 0 + missing = 0 + for layer in range(total_layers): + for tp in range(tp_count): + for expert in range(expert_num): + # AMXINT4 NUMA-sharded tensor 名 + gate_w_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{tp}.weight" + gate_s_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{tp}.scale" + up_w_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{tp}.weight" + up_s_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{tp}.scale" + down_w_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{tp}.weight" + down_s_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{tp}.scale" + + # 检查必需的 tensor 是否存在 + required = [gate_w_key, up_w_key, down_w_key] + if not all(k in tensor_map for k in required): + missing += 1 + continue + + # 获取 gate 布局 + g_fpath, g_off, g_bytes = tensor_map[gate_w_key] + fd = get_fd(g_fpath) + + # 获取 up 布局 + u_fpath, u_off, u_bytes = tensor_map[up_w_key] + if u_fpath != g_fpath: + fd = get_fd(u_fpath) # 不同文件需要不同 fd + + # 获取 down 布局 + d_fpath, d_off, d_bytes = tensor_map[down_w_key] + if d_fpath not in file_fds: + fd = get_fd(d_fpath) + + # scale 偏移（AMXINT4 专用） + g_s_off = g_s_bytes = 0 + u_s_off = u_s_bytes = 0 + d_s_off = d_s_bytes = 0 + if gate_s_key in tensor_map: + _, g_s_off, g_s_bytes = tensor_map[gate_s_key] + if up_s_key in tensor_map: + _, u_s_off, u_s_bytes = tensor_map[up_s_key] + if down_s_key in tensor_map: + _, d_s_off, d_s_bytes = tensor_map[down_s_key] + + # 注意：fd 取 gate 所在文件的 fd。如果 up/down 在不同文件， + # io_uring 的 submit_load 会用 layout.fd 读取所有三个矩阵， + # 这要求三个矩阵在同一文件。safetensors 通常如此（同层同 NUMA 在同一文件）。 + # 如果跨文件，需要后续拆分为多次 submit_load。 + self._residency.set_file_layout( + layer_idx=layer, + tp_part_idx=tp, + expert_id=expert, + fd=fd, + gate_offset=g_off, + up_offset=u_off, + down_offset=d_off, + gate_bytes=g_bytes, + up_bytes=u_bytes, + down_bytes=d_bytes, + gate_scale_offset=g_s_off, + up_scale_offset=u_s_off, + down_scale_offset=d_s_off, + gate_scale_bytes=g_s_bytes, + up_scale_bytes=u_s_bytes, + down_scale_bytes=d_s_bytes, + # AMXINT4 对称量化无 mins，保持 0 + ) + injected += 1 + + logger.info( + "_inject_amxint4_layouts: injected %d layouts (missing %d), " + "opened %d files with O_RDONLY", + injected, missing, len(file_fds), + ) + + def _inject_bf16_layouts( + self, + tensor_map: dict, + total_layers: int, + tp_count: int, + expert_num: int, + mesh_config: MeshConfig, + ) -> None: + """注入 BF16 packed 格式的文件布局。 + + Packed 格式：所有专家打包成 3D tensor + - model.layers.{L}.mlp.experts.gate_up_proj [E, 2*I, H] + - model.layers.{L}.mlp.experts.down_proj [E, H, I] - 这里需要根据实际权重文件格式解析偏移。 - 简化实现：假设权重已按 TP 切分存储，每个 TP 一个文件。 + 每个 TP 分片读取 intermediate_size/tp_count 的切片。 """ - # TODO: 实际实现需要根据 SafeTensorLoader 的接口获取每个专家的文件偏移 - # 这里只提供框架，具体偏移计算依赖权重文件格式 - # - # for tp in range(tp_count): - # for eid in range(expert_num): - # layout = loader.get_expert_layout(tp, eid) - # fd = os.open(layout.path, os.O_DIRECT | os.O_RDONLY) - # self._residency.set_file_layout( - # tp, eid, fd, - # layout.gate_offset, layout.up_offset, layout.down_offset, - # layout.gate_bytes, layout.up_bytes, layout.down_bytes, - # ... - # ) - logger.warning( - "File layout injection not yet implemented. " - "Please implement _inject_file_layouts based on SafeTensorLoader." + hidden = mesh_config.hidden_size + inter = mesh_config.intermediate_size + inter_per_tp = inter // tp_count + # bf16 = 2 bytes + elem_bytes = 2 + + file_fds: dict[str, int] = {} + + def get_fd(fpath: str) -> int: + if fpath not in file_fds: + fd = os.open(fpath, os.O_RDONLY) + file_fds[fpath] = fd + return file_fds[fpath] + + injected = 0 + missing = 0 + for layer in range(total_layers): + gate_up_key = f"model.layers.{layer}.mlp.experts.gate_up_proj" + down_key = f"model.layers.{layer}.mlp.experts.down_proj" + + if gate_up_key not in tensor_map or down_key not in tensor_map: + missing += 1 + continue + + gu_fpath, gu_base, _ = tensor_map[gate_up_key] + d_fpath, d_base, _ = tensor_map[down_key] + fd = get_fd(gu_fpath) + + # gate_up_proj: [E, 2*I, H]，每个专家占 2*I*H*elem_bytes + expert_stride_gu = 2 * inter * hidden * elem_bytes + gate_bytes = inter_per_tp * hidden * elem_bytes + up_bytes = inter_per_tp * hidden * elem_bytes + + # down_proj: [E, H, I]，每个专家占 H*I*elem_bytes + expert_stride_d = hidden * inter * elem_bytes + down_bytes = hidden * inter_per_tp * elem_bytes + + for tp in range(tp_count): + for expert in range(expert_num): + # gate: expert 的 gate 部分起始偏移 + TP 切片偏移 + gate_off = gu_base + expert * expert_stride_gu + tp * gate_bytes + # up: expert 的 up 部分起始偏移（跳过 gate） + TP 切片偏移 + up_off = gu_base + expert * expert_stride_gu + inter * hidden * elem_bytes + tp * up_bytes + # down: expert 的 down 起始偏移 + TP 切片偏移 + down_off = d_base + expert * expert_stride_d + tp * down_bytes + + self._residency.set_file_layout( + layer_idx=layer, + tp_part_idx=tp, + expert_id=expert, + fd=fd, + gate_offset=gate_off, + up_offset=up_off, + down_offset=down_off, + gate_bytes=gate_bytes, + up_bytes=up_bytes, + down_bytes=down_bytes, + # BF16 无 scale/mins + ) + injected += 1 + + logger.info( + "_inject_bf16_layouts: injected %d layouts (missing %d), " + "opened %d files", + injected, missing, len(file_fds), ) # ===== 属性委托 ===== From 08317541d47c2005b7828f8ca5449716492f671e Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Sun, 21 Jun 2026 15:51:56 +0800 Subject: [PATCH 06/10] fix: mesh v3 nan error --- kt-kernel/cpu_backend/cpuinfer.h | 30 ++++- kt-kernel/ext_bindings.cpp | 57 +++++++- kt-kernel/operators/mesh/mesh_config.hpp | 28 ++++ kt-kernel/operators/mesh/mesh_decode.hpp | 47 +++++-- kt-kernel/operators/mesh/mesh_eviction.hpp | 45 ++++++- kt-kernel/operators/mesh/mesh_handoff.hpp | 24 +++- kt-kernel/operators/mesh/mesh_prefill.hpp | 91 ++++++++++--- kt-kernel/operators/mesh/mesh_residency.hpp | 141 ++++++++++++++++++-- kt-kernel/operators/mesh/mesh_slot_pool.hpp | 11 ++ kt-kernel/python/utils/mesh/stats.py | 3 +- 10 files changed, 423 insertions(+), 54 deletions(-) diff --git a/kt-kernel/cpu_backend/cpuinfer.h b/kt-kernel/cpu_backend/cpuinfer.h index 8efe18c9e..2afef60d3 100644 --- a/kt-kernel/cpu_backend/cpuinfer.h +++ b/kt-kernel/cpu_backend/cpuinfer.h @@ -33,6 +33,14 @@ #include "task_queue.h" #include "worker_pool.h" +// B8: MESH Heat 批量传输回调声明（定义在 ext_bindings.cpp 中） +// 由 CPUInfer::submit_gating_scores_ 调用，避免 cpuinfer.h 直接依赖 mesh 头文件 +#if defined(KT_ENABLE_MESH) +extern "C" void mesh_on_gating_scores_ready(void* mesh_residency, + const float* gating_scores_cpu, + int num_layers, int expert_num, int topk); +#endif + class CPUInfer { public: CPUInfer(int thread_num) { @@ -120,14 +128,24 @@ class CPUInfer { const float* gating_scores_gpu; // GPU 指针，所有层的 gating 分数 int num_layers; int expert_num; - // CQE 到达后由 mesh_residency 的 on_decode_token_end 消费 + int topk; // B8: top-k 参数（每层激活的专家数） }; static void submit_gating_scores_(void* args_ptr) { GatingScoreArgs* args = (GatingScoreArgs*)args_ptr; - // 实际实现：将 gating_scores_gpu 数据传给 mesh::MeshResidencyManager::on_decode_token_end - // mesh_residency 通过 CUDA memcpy 或 unified memory 获取数据 - // 这里是 host callback，在 CUDA stream 上执行 +#if defined(KT_ENABLE_MESH) + // B8: 从 GPU 拷贝 gating scores 到 CPU，然后调用 mesh 回调 + int total_floats = args->num_layers * args->expert_num; + std::vector gating_scores_cpu(total_floats); +#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \ + defined(KTRANSFORMERS_USE_MACA) + cudaMemcpy(gating_scores_cpu.data(), args->gating_scores_gpu, + total_floats * sizeof(float), cudaMemcpyDeviceToHost); +#endif + // 调用 mesh 回调，内部会组织数据并调用 on_decode_token_end + mesh_on_gating_scores_ready(args->mesh_residency, gating_scores_cpu.data(), + args->num_layers, args->expert_num, args->topk); +#endif delete args; } @@ -136,11 +154,11 @@ class CPUInfer { intptr_t user_cuda_stream, void* mesh_residency, const float* gating_scores_gpu, - int num_layers, int expert_num) { + int num_layers, int expert_num, int topk) { #if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \ defined(KTRANSFORMERS_USE_MACA) GatingScoreArgs* args = new GatingScoreArgs{ - this, mesh_residency, gating_scores_gpu, num_layers, expert_num + this, mesh_residency, gating_scores_gpu, num_layers, expert_num, topk }; cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&submit_gating_scores_, (void*)args); diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 880912feb..4d592f40f 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -12,6 +12,9 @@ #include #include +#include +#include + #if defined(KTRANSFORMERS_ENABLE_CPPTRACE) #include #endif @@ -26,6 +29,38 @@ // MESH 插件（仅在 KT_ENABLE_MESH 编译时生效） #if defined(KT_ENABLE_MESH) #include "operators/mesh/mesh.hpp" + +// B8: Heat 批量传输回调 — 由 cpuinfer.h 的 submit_gating_scores_ 调用 +// 把 GPU 传来的扁平 gating scores 组织成 [layer][expert_num] 并计算 top-k， +// 然后调用 MeshResidencyManager::on_decode_token_end 更新 Heat 和 Markov +extern "C" void mesh_on_gating_scores_ready(void* mesh_residency, + const float* gating_scores_cpu, + int num_layers, int expert_num, int topk) { + if (!mesh_residency || !gating_scores_cpu) return; + auto* mgr = static_cast(mesh_residency); + + // 组织成 [layer][expert_num] + std::vector> all_layers_scores(num_layers); + for (int l = 0; l < num_layers; l++) { + all_layers_scores[l].assign( + gating_scores_cpu + l * expert_num, + gating_scores_cpu + (l + 1) * expert_num); + } + + // 从 scores 计算 top-k（每层取分数最高的 topk 个专家） + std::vector> all_layers_topk(num_layers); + int k = std::min(topk, expert_num); + for (int l = 0; l < num_layers; l++) { + const auto& scores = all_layers_scores[l]; + std::vector indices(scores.size()); + std::iota(indices.begin(), indices.end(), 0); + std::partial_sort(indices.begin(), indices.begin() + k, indices.end(), + [&scores](int a, int b) { return scores[a] > scores[b]; }); + all_layers_topk[l].assign(indices.begin(), indices.begin() + k); + } + + mgr->on_decode_token_end(all_layers_topk, all_layers_scores); +} #endif #if defined(USE_MOE_KERNEL) @@ -1058,6 +1093,22 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def_readwrite("up_mins_bytes", &mesh::ExpertFileLayout::up_mins_bytes) .def_readwrite("down_mins_bytes", &mesh::ExpertFileLayout::down_mins_bytes); + // B9: MeshStats 绑定 + py::class_(mesh_module, "MeshStats") + .def_readwrite("cache_hit_count", &mesh::MeshStats::cache_hit_count) + .def_readwrite("cache_miss_count", &mesh::MeshStats::cache_miss_count) + .def_readwrite("io_uring_read_bytes", &mesh::MeshStats::io_uring_read_bytes) + .def_readwrite("io_uring_read_count", &mesh::MeshStats::io_uring_read_count) + .def_readwrite("eviction_count", &mesh::MeshStats::eviction_count) + .def_readwrite("eviction_blocked_wait_us", &mesh::MeshStats::eviction_blocked_wait_us) + .def_readwrite("defer_count", &mesh::MeshStats::defer_count) + .def_readwrite("defer_overflow_count", &mesh::MeshStats::defer_overflow_count) + .def_readwrite("prefill_layer_count", &mesh::MeshStats::prefill_layer_count) + .def_readwrite("prefill_temporal_swap_count", &mesh::MeshStats::prefill_temporal_swap_count) + .def_readwrite("decode_token_count", &mesh::MeshStats::decode_token_count) + .def_readwrite("decode_immediate_count", &mesh::MeshStats::decode_immediate_count) + .def_readwrite("decode_deferred_count", &mesh::MeshStats::decode_deferred_count); + // MeshResidencyManager py::class_>( mesh_module, "ResidencyManager") @@ -1100,7 +1151,11 @@ PYBIND11_MODULE(kt_kernel_ext, m) { // A6: 查询专家是否已缓存 .def("is_cached", [](mesh::MeshResidencyManager& mgr, int layer, int tp, int expert_id) -> bool { return mgr.pool(layer, tp).is_cached(expert_id); - }, py::arg("layer"), py::arg("tp"), py::arg("expert_id")); + }, py::arg("layer"), py::arg("tp"), py::arg("expert_id")) + // B9: 统计接口 + .def("stats", [](mesh::MeshResidencyManager& mgr) -> const mesh::MeshStats& { + return mgr.stats(); + }, py::return_value_policy::reference_internal); // 注册 hook 函数指针（让 mesh_hook.hpp 的 inline 函数能调用 ResidencyManager） mesh::hook::HookRegistry registry; diff --git a/kt-kernel/operators/mesh/mesh_config.hpp b/kt-kernel/operators/mesh/mesh_config.hpp index 2680af4bf..cc8c16b9b 100644 --- a/kt-kernel/operators/mesh/mesh_config.hpp +++ b/kt-kernel/operators/mesh/mesh_config.hpp @@ -91,4 +91,32 @@ struct MeshConfig { // 实际存储在 MeshResidencyManager 中，这里只放声明 }; +// B9: 运行时统计结构体 +struct MeshStats { + // 命中率 + uint64_t cache_hit_count = 0; + uint64_t cache_miss_count = 0; + + // io_uring 读取量 + uint64_t io_uring_read_bytes = 0; + uint64_t io_uring_read_count = 0; + + // 驱逐统计 + uint64_t eviction_count = 0; + uint64_t eviction_blocked_wait_us = 0; + + // defer 统计 + uint64_t defer_count = 0; + uint64_t defer_overflow_count = 0; + + // prefill 统计 + uint64_t prefill_layer_count = 0; + uint64_t prefill_temporal_swap_count = 0; + + // decode 统计 + uint64_t decode_token_count = 0; + uint64_t decode_immediate_count = 0; + uint64_t decode_deferred_count = 0; +}; + } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_decode.hpp b/kt-kernel/operators/mesh/mesh_decode.hpp index e2e2d9e23..caa54c6e3 100644 --- a/kt-kernel/operators/mesh/mesh_decode.hpp +++ b/kt-kernel/operators/mesh/mesh_decode.hpp @@ -52,6 +52,8 @@ class MeshDecode { /** * @brief 对 top-k 专家进行 immediate/deferred 分组 * + * B6 fix: 入口先过滤 GPU 专家（GPU 专家不参与 slot 调度） + * * @param topk 当前层 top-k 专家 ID * @param scores 各专家的 router score（全长 expert_num） * @param slot_pool 该层 slot 池 @@ -59,6 +61,7 @@ class MeshDecode { * @param layer_idx 当前层 * @param tp_part_idx TP 分片 * @param get_dst_ptrs 获取 slot buffer 指针的回调 + * @param is_gpu_expert 判断专家是否在 GPU 上的回调 * @return SplitResult */ SplitResult split(const std::vector& topk, @@ -66,13 +69,25 @@ class MeshDecode { MeshSlotPool& slot_pool, MeshScheduler& scheduler, int layer_idx, int tp_part_idx, - std::function(int)> get_dst_ptrs) { + std::function(int)> get_dst_ptrs, + std::function is_gpu_expert = nullptr) { int k = static_cast(topk.size()); int target_immediate = std::max(0, k - defer_count_); - // 区分已缓存和缺失 - std::vector cached, missing; + // B6 fix: 过滤 GPU 专家，GPU 专家直接进 immediate（由 GPU 计算，不需 slot） + std::vector cpu_topk; + std::vector gpu_experts; for (int eid : topk) { + if (is_gpu_expert && is_gpu_expert(eid)) { + gpu_experts.push_back(eid); + } else { + cpu_topk.push_back(eid); + } + } + + // 区分已缓存和缺失（仅 CPU 专家） + std::vector cached, missing; + for (int eid : cpu_topk) { if (slot_pool.is_cached(eid)) { cached.push_back(eid); } else { @@ -81,6 +96,8 @@ class MeshDecode { } SplitResult result; + // GPU 专家直接进 immediate + result.immediate = gpu_experts; if (static_cast(cached.size()) > target_immediate) { // immediate 数 > k-defer：按 router score 排序取前 target_immediate 个 @@ -90,7 +107,7 @@ class MeshDecode { float sb = (b < (int)scores.size()) ? scores[b] : 0.0f; return sa > sb; // 降序 }); - result.immediate.assign(cached.begin(), cached.begin() + target_immediate); + result.immediate.insert(result.immediate.end(), cached.begin(), cached.begin() + target_immediate); result.deferred.assign(cached.begin() + target_immediate, cached.end()); // 缺失专家全部进 deferred result.deferred.insert(result.deferred.end(), missing.begin(), missing.end()); @@ -103,7 +120,7 @@ class MeshDecode { } } else { // immediate 数 < k-defer：阻塞读缺失专家直到 immediate 数量够 - result.immediate = cached; + result.immediate.insert(result.immediate.end(), cached.begin(), cached.end()); int need = target_immediate - static_cast(cached.size()); for (int i = 0; i < need && i < static_cast(missing.size()); i++) { @@ -134,11 +151,16 @@ class MeshDecode { /** * @brief 在线驱逐：slot 满时选分数最低的覆盖 * + * B1 fix: 实现完整的驱逐逻辑 + * 1. 找 victim 专家（分数最低的 CACHED 专家） + * 2. 查 victim 的 slot_idx + * 3. 调用 overwrite(slot_idx, new_expert_id) + * * @param slot_pool 该层 slot 池 * @param scorer 驱逐评分器 * @param layer_idx 当前层 * @param new_expert_id 要加载的新专家 - * @return int 被驱逐的 slot_idx，-1 表示无需驱逐 + * @return int 被驱逐释放的 slot_idx，-1 表示无需驱逐或无法驱逐 */ int evict_for_new_expert(MeshSlotPool& slot_pool, const EvictionScorer& scorer, @@ -159,9 +181,16 @@ class MeshDecode { int victim_expert = scorer.select_victim(cached, layer_idx); if (victim_expert < 0) return -1; - // 查找 victim 对应的 slot_idx - // 实际实现需要通过 slot_pool 的接口 - return -1; // placeholder + // B1 fix: 查找 victim 对应的 slot_idx + int victim_slot_idx = slot_pool.expert_to_slot_idx(victim_expert); + if (victim_slot_idx < 0) return -1; + + // 检查 victim 是否有活跃 reader + // overwrite 内部会 spin wait 等 reader 归零 + // 调用 overwrite 覆盖 slot + slot_pool.overwrite(victim_slot_idx, new_expert_id); + + return victim_slot_idx; } // ===== 前 5 层满配 ===== diff --git a/kt-kernel/operators/mesh/mesh_eviction.hpp b/kt-kernel/operators/mesh/mesh_eviction.hpp index e90640a4c..1f6868905 100644 --- a/kt-kernel/operators/mesh/mesh_eviction.hpp +++ b/kt-kernel/operators/mesh/mesh_eviction.hpp @@ -257,9 +257,13 @@ class MarkovTracker { /** * @brief 驱逐评分器 * - * score = policy_rank + lookahead_weight * max(heat, markov_prior) + * score = policy_rank + lookahead_weight * max(heat, cross_layer_prior) * * Markov 严格限制：只用于驱逐评分，绝对不用于预取。 + * + * B2 fix: cross_layer_prior 按 [layer][expert] 累积，predict_next_layer 在每层结束时 + * 调用一次，把 Markov 预测结果 EMA 进 cross_layer_prior_[layer+1]。 + * score() 直接读 cross_layer_prior_[layer]，不再每次 predict。 */ class EvictionScorer { public: @@ -267,7 +271,11 @@ class EvictionScorer { : heat_(expert_num, config.heat_beta, config.heat_gamma), markov_(num_layers, expert_num, config.markov_alpha), lookahead_weight_(config.lookahead_weight), - expert_num_(expert_num) {} + markov_alpha_(config.markov_alpha), + num_layers_(num_layers), + expert_num_(expert_num) { + cross_layer_prior_.assign(num_layers, std::vector(expert_num, 0.0f)); + } // 单 token 结束后批量更新 Heat 和 Markov // all_layers_topk: [layer] -> top-k expert ids @@ -285,12 +293,35 @@ class EvictionScorer { } } + /** + * @brief B2: 在第 layer_idx 层结束时，预测第 layer_idx+1 层的 cross_layer_prior + * + * cross_layer_prior_[layer+1][t] = α * P_predicted[t] + (1-α) * cross_layer_prior_[layer+1][t] + * + * @param layer_idx 当前层 L（预测 L+1） + * @param topk 当前层的 top-k 专家集合 + * @param scores 当前层各专家的归一化权重（全长 expert_num） + */ + void predict_next_layer(int layer_idx, const std::vector& topk, + const std::vector& scores) { + if (layer_idx < 0 || layer_idx + 1 >= num_layers_) return; + std::vector predicted; + markov_.predict(layer_idx, topk, scores, predicted); + auto& prior = cross_layer_prior_[layer_idx + 1]; + for (int t = 0; t < expert_num_ && t < (int)predicted.size(); t++) { + prior[t] = markov_alpha_ * predicted[t] + (1.0f - markov_alpha_) * prior[t]; + } + } + // 计算某专家的驱逐评分（分数越低越该被驱逐） + // B2 fix: 直接读 cross_layer_prior_[layer_idx]，不再每次 predict float score(int expert_id, int layer_idx) const { float h = heat_.heat(expert_id); - std::vector prior; - markov_.predict(layer_idx, {}, {}, prior); // 简化：实际需要传入当前层 topk - float m = (expert_id < (int)prior.size()) ? prior[expert_id] : 0.0f; + float m = 0.0f; + if (layer_idx >= 0 && layer_idx < num_layers_ && + expert_id >= 0 && expert_id < expert_num_) { + m = cross_layer_prior_[layer_idx][expert_id]; + } float heat = std::max(h, m); return lookahead_weight_ * heat; } @@ -323,7 +354,11 @@ class EvictionScorer { HeatTracker heat_; MarkovTracker markov_; float lookahead_weight_; + float markov_alpha_; + int num_layers_; int expert_num_; + // B2: [layer][expert] 的 Markov 先验，由 predict_next_layer 累积 + std::vector> cross_layer_prior_; }; } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_handoff.hpp b/kt-kernel/operators/mesh/mesh_handoff.hpp index 286d41422..96068a8d4 100644 --- a/kt-kernel/operators/mesh/mesh_handoff.hpp +++ b/kt-kernel/operators/mesh/mesh_handoff.hpp @@ -82,6 +82,11 @@ class MeshHandoff { /** * @brief 搬运 GPU 专家到 GPU，并用 io_uring 读新专家填充释放的 slot * + * B7 fix: + * 1. 先 overwrite 释放旧 slot 绑定，绑定新专家 + * 2. 再用 slot 指针作为 io_uring 目标地址 submit_load + * 3. on_complete 回调中 mark_cached + * * GPU expert 在 slot 前茅（GE 个位置）： * - 标记这些 slot 的 active_readers（防止驱逐） * - 搬运到 GPU（不读盘） @@ -96,12 +101,13 @@ class MeshHandoff { MeshIoUring& io, std::function move_gpu_expert_to_gpu, std::function(int, int, int)> get_dst_ptrs) { + (void)scorer; // 暂未使用 scorer 选 victim，用固定区间 int ge = config.num_gpu_experts; int cap = config.cap; int num_layers = config.total_layers; int tp_count = config.tp_count; - // (cap, cap+GE] 区间的专家 ID + // (cap, cap+GE] 区间的专家 ID — 用于填充释放的 slot std::vector refill_experts; for (int e = cap; e < cap + ge && e < config.expert_num; e++) { refill_experts.push_back(e); @@ -125,27 +131,33 @@ class MeshHandoff { // 释放 reader pool.release_reader(expert_id); - // 释放这个 slot，用 io_uring 读新专家填充 + // B7 fix: 先 overwrite 释放旧 slot 绑定，绑定新专家 int new_expert_id = cap + slot_idx; if (new_expert_id < config.expert_num && slot_idx < static_cast(refill_experts.size())) { new_expert_id = refill_experts[slot_idx]; - auto ptrs = get_dst_ptrs(layer, tp, new_expert_id); + + // 先 overwrite：解绑旧 expert，绑定新 expert + pool.overwrite(slot_idx, new_expert_id); + + // B7 fix: 用 slot 指针作为 io_uring 目标地址（overwrite 后 slot 已绑定新专家） + void* gate_dst = pool.gate_ptr(slot_idx); + void* up_dst = pool.up_ptr(slot_idx); + void* down_dst = pool.down_ptr(slot_idx); + // 提交 io_uring 读 // A4 fix: layouts 改为 [layer][tp][expert] 3D // A6: 传入 layer_idx 和 on_complete 回调 const auto& layout = layouts[layer][tp][new_expert_id]; MeshSlotPool& pool_ref = pools[layer][tp]; io.submit_load(new_expert_id, tp, layout, - ptrs[0], ptrs[1], ptrs[2], + gate_dst, up_dst, down_dst, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, ReadPriority::Demand, /*layer_idx=*/layer, /*on_complete=*/[&pool_ref, new_expert_id](int, int, int) { pool_ref.mark_cached(new_expert_id); }); - // 覆盖 slot - pool.overwrite(slot_idx, new_expert_id); } } } diff --git a/kt-kernel/operators/mesh/mesh_prefill.hpp b/kt-kernel/operators/mesh/mesh_prefill.hpp index 45c23c978..e846b1705 100644 --- a/kt-kernel/operators/mesh/mesh_prefill.hpp +++ b/kt-kernel/operators/mesh/mesh_prefill.hpp @@ -28,6 +28,7 @@ #include #include "mesh_config.hpp" +#include "mesh_io_uring.hpp" #include "mesh_scheduler.hpp" #include "mesh_slot_pool.hpp" @@ -63,6 +64,14 @@ class MeshPrefill { size_t bytes = static_cast(temporal_size_) * slot_bytes_ * tp_count_; role_a_ = TemporalRole::COMPUTE; role_b_ = TemporalRole::PREFETCH; + // B4 fix: 初始化 expert_id -> temporal_offset 映射 + // temporal 中装的是 cap ~ expert_num-1 的专家，offset = expert_id - cap + temporal_expert_map_.assign(tp_count_, std::vector(expert_num_, -1)); + for (int tp = 0; tp < tp_count_; tp++) { + for (int e = cap_; e < expert_num_; e++) { + temporal_expert_map_[tp][e] = e - cap_; + } + } if (bytes == 0) { // cap >= expert_num: 所有专家常驻 slot，不需要 temporal buffer return; @@ -92,30 +101,36 @@ class MeshPrefill { /** * @brief 处理某一层的所有 chunk * + * B3 fix: wait_temporal_ready 接入 io_uring 等待 + * B4 fix: submit_next_layer_prefetch 用 temporal_ptr 正确寻址 + * * @param layer_idx 当前层 L * @param total_chunks 总 chunk 数 * @param active_experts_per_chunk [chunk] -> 该 chunk 活跃的专家列表 * @param slot_pool 该层的 slot 池 * @param scheduler 调度器 + * @param io io_uring 读取器（B3: 用于等待 CQE） * @param amx_forward AMX 计算回调（layer_idx, chunk_idx, expert_ids） - * @param get_temporal_ptrs 获取 temporal buffer 指针的回调 */ void run_layer(int layer_idx, int total_chunks, const std::vector>& active_experts_per_chunk, MeshSlotPool& slot_pool, MeshScheduler& scheduler, - std::function&)> amx_forward, - std::function(int, int)> get_temporal_ptrs) { + MeshIoUring* io, + std::function&)> amx_forward) { + void* compute_buf = compute_temporal(); + void* prefetch_buf = prefetch_temporal(); + (void)compute_buf; // compute_buf 由 amx_forward 通过 slot_pool/temporal_ptr 访问 + for (int c = 0; c < total_chunks; c++) { const auto& active = active_experts_per_chunk[c]; // a. 等待 temporal_A 内第 L 层所需专家全部就绪 // （由上一层的异步预取完成；第一层阻塞等待） - wait_temporal_ready(layer_idx, active, slot_pool); + wait_temporal_ready(layer_idx, active, slot_pool, io); // b. 对 temporal_B 发起第 L+1 层的异步预取（覆盖另一个 temporal） if (layer_idx + 1 < num_layers_) { - submit_next_layer_prefetch(layer_idx + 1, active, scheduler, - get_temporal_ptrs); + submit_next_layer_prefetch(layer_idx + 1, active, scheduler, prefetch_buf); } // c. 本层计算使用 temporal_A 中的临时专家 + slot 池常驻专家 @@ -154,6 +169,33 @@ class MeshPrefill { return role_a_ == TemporalRole::PREFETCH ? temporal_a_ : temporal_b_; } + /** + * @brief B4 fix: 获取 temporal buffer 中某专家某矩阵的指针 + * + * @param buf temporal buffer 基地址（compute_temporal() 或 prefetch_temporal()） + * @param tp TP 分片 + * @param expert_id 专家 ID + * @param matrix_idx 0=gate, 1=up, 2=down + * @return void* 指针，未映射返回 nullptr + */ + void* temporal_ptr(void* buf, int tp, int expert_id, int matrix_idx) const { + if (tp < 0 || tp >= tp_count_) return nullptr; + if (expert_id < 0 || expert_id >= expert_num_) return nullptr; + int offset = temporal_expert_map_[tp][expert_id]; + if (offset < 0) return nullptr; + // 布局：[tp][offset] 每个 slot_bytes_，内含 gate + up + down + char* base = static_cast(buf) + + static_cast(tp) * temporal_size_ * slot_bytes_ + + static_cast(offset) * slot_bytes_; + // gate 在偏移 0，up 在偏移 gate_up_bytes_，down 在偏移 gate_up_bytes_*2 + size_t matrix_offset = (matrix_idx == 0) ? 0 : + (matrix_idx == 1) ? gate_up_bytes_ : gate_up_bytes_ * 2; + return base + matrix_offset; + } + + // B4 fix: 设置 gate/up 单块字节数（用于 temporal_ptr 偏移计算） + void set_gate_up_bytes(size_t bytes) { gate_up_bytes_ = bytes; } + // ===== 访问器 ===== int temporal_size() const { return temporal_size_; } bool temporal_ready() const { return temporal_a_ != nullptr; } @@ -164,6 +206,7 @@ class MeshPrefill { int cap_; int tp_count_; size_t slot_bytes_; + size_t gate_up_bytes_ = 0; // B4: gate/up 单块字节数 int numa_node_; size_t temporal_bytes_ = 0; @@ -174,30 +217,44 @@ class MeshPrefill { TemporalRole role_b_ = TemporalRole::PREFETCH; int temporal_size_ = 0; // = expert_num - cap + // B4 fix: expert_id -> temporal_offset 映射 [tp][expert_id] + std::vector> temporal_expert_map_; + // 等待 temporal 内专家就绪 + // B3 fix: 调用 io_uring process_cqes 推进 CQE，直到所有活跃专家 CACHED void wait_temporal_ready(int layer_idx, const std::vector& active, - MeshSlotPool& slot_pool) { - // 对于第一层，阻塞直到全部读取完毕 - // 对于后续层，由上一层预取完成，这里检查状态 + MeshSlotPool& slot_pool, MeshIoUring* io) { for (int eid : active) { if (slot_pool.is_cached(eid)) continue; - // 等待 io_uring CQE - // 实际实现需要与 MeshIoUring 配合 + // 阻塞等待 io_uring CQE 到达 + // 实际实现需要与 MeshIoUring 配合轮询 process_cqes + if (io) { + while (!slot_pool.is_cached(eid)) { + io->process_cqes(); + io->submit_and_wait(); + } + } } } // 提交下一层的异步预取 + // B4 fix: 用 temporal_ptr 正确寻址，不再用 ptrs[eid*3] 越界 void submit_next_layer_prefetch(int next_layer, const std::vector& active, MeshScheduler& scheduler, - std::function(int, int)> get_ptrs) { - // 对下一层需要的专家提交异步读，schedule_key = next_layer - void* prefetch_buf = prefetch_temporal(); + void* prefetch_buf) { for (int tp = 0; tp < tp_count_; tp++) { - auto ptrs = get_ptrs(next_layer, tp); for (int eid : active) { - scheduler.submit_prefill(next_layer, eid, tp, - ptrs[eid * 3], ptrs[eid * 3 + 1], ptrs[eid * 3 + 2], + // 只预取不在 slot 中的专家（在 temporal 中的） + int offset = (tp < (int)temporal_expert_map_.size() && + eid < (int)temporal_expert_map_[tp].size()) + ? temporal_expert_map_[tp][eid] : -1; + if (offset < 0) continue; + void* gate = temporal_ptr(prefetch_buf, tp, eid, 0); + void* up = temporal_ptr(prefetch_buf, tp, eid, 1); + void* down = temporal_ptr(prefetch_buf, tp, eid, 2); + if (!gate || !up || !down) continue; + scheduler.submit_prefill(next_layer, eid, tp, gate, up, down, ReadPriority::Prefetch); } } diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp index 25935e88c..abde6af98 100644 --- a/kt-kernel/operators/mesh/mesh_residency.hpp +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -14,6 +14,7 @@ */ #pragma once +#include #include #include #include @@ -96,6 +97,8 @@ class MeshResidencyManager { // 初始化 temporal 双缓冲 prefill_->init_temporal(); + // B4 fix: 注入 gate_up_bytes 用于 temporal_ptr 偏移计算 + prefill_->set_gate_up_bytes(compute_gate_up_bytes(config)); fprintf(stderr, "[MESH init DIAG] init complete, pools_ size=%zu\n", pools_.size()); @@ -249,11 +252,20 @@ class MeshResidencyManager { void on_decode_token_start() { scheduler_->inc_timeline_step(); + // B5: 每 token 开始时清空跨层 defer 队列 + prev_layer_deferred_.clear(); + total_deferred_ = 0; + // B9: 统计 + stats_.decode_token_count++; } /** * @brief Decode 每层处理 * + * B5: 跨层 defer 累积 + overflow 阻塞 + * - 上一层的 deferred 专家在当前层转为 immediate 候选 + * - 如果 deferred 累积超过 max_deferred_per_token，阻塞等待 io_uring 完成 + * * @param layer_idx 当前层 * @param topk 当前 token 的 top-k 专家 * @param scores 各专家 router score（全长） @@ -267,8 +279,27 @@ class MeshResidencyManager { // A7: 先处理已完成的 io_uring CQE，触发 mark_cached 回调 io_->process_cqes(); + // B5: 把上一层的 deferred 专家加入当前层的 immediate 候选 + // 它们应该已经通过 io_uring 读取完成 + std::vector effective_topk = topk; + if (!prev_layer_deferred_.empty()) { + // 检查 deferred 专家是否已 CACHED + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; + for (int eid : prev_layer_deferred_) { + if (pool.is_cached(eid)) { + // 已缓存，加入 immediate 候选 + if (std::find(effective_topk.begin(), effective_topk.end(), eid) == effective_topk.end()) { + effective_topk.push_back(eid); + } + } + // 未缓存的 deferred 专家在本层继续等待（会在 split 中被当作 missing 处理） + } + // 清空上一层的 defer 队列 + prev_layer_deferred_.clear(); + } + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; - auto result = decode_->split(topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, + auto result = decode_->split(effective_topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, [this, &pool, layer_idx, tp_part_idx](int eid) { // A5 fix: 用 expert_*_ptr 按 expert_id 查 slot， // 不能把 expert_id 当 slot_idx 传给 gate_ptr（eid>=cap 时越界） @@ -276,11 +307,46 @@ class MeshResidencyManager { pool.expert_gate_ptr(eid), pool.expert_up_ptr(eid), pool.expert_down_ptr(eid)}; - }); + }, + // B6: 传入 is_gpu_expert 回调，过滤 GPU 专家 + [this](int eid) { return is_gpu_expert(eid); }); + + // B5: 更新跨层 defer 队列 + prev_layer_deferred_ = result.deferred; + total_deferred_ += static_cast(result.deferred.size()); + + // B9: 统计 + stats_.decode_immediate_count += static_cast(result.immediate.size()); + stats_.decode_deferred_count += static_cast(result.deferred.size()); + stats_.defer_count += static_cast(result.deferred.size()); + // hit/miss 统计：effective_topk 中已缓存的是 hit，未缓存的是 miss + for (int eid : effective_topk) { + if (is_gpu_expert(eid)) continue; // GPU 专家不计入 hit/miss + if (pool.is_cached(eid)) { + stats_.cache_hit_count++; + } else { + stats_.cache_miss_count++; + } + } + + // B5: overflow 检查 — 如果 defer 累积超过阈值，阻塞等待 io_uring 完成 + if (total_deferred_ > config_.max_deferred_per_token * config_.total_layers) { + // B9: 统计 + stats_.defer_overflow_count++; + // 阻塞等待所有 inflight io_uring 完成 + io_->submit_and_wait(); + io_->process_cqes(); + total_deferred_ = 0; // 重置计数 + } // A7: 排空调度器队列，提交 io_uring 读取 drain_and_submit(); + // B2: 用当前层的实际路由（原始 topk + scores）预测下一层的 cross_layer_prior + // 放在 drain_and_submit 之后：当前层 IO 已提交，为下一层驱逐评分做准备 + // 注意：用原始 topk/scores，不用 effective_topk（后者混入了上一层 deferred，不是真实路由） + scorer_->predict_next_layer(layer_idx, topk, scores); + return result; } @@ -290,6 +356,8 @@ class MeshResidencyManager { * 把 MeshScheduler 优先队列中的 ScheduledRequest 取出， * 为每个请求找到对应的 ExpertFileLayout 并提交给 MeshIoUring。 * 完成后触发 mark_cached + 原始 on_complete 回调。 + * + * B1: 如果专家未缓存且 slot 池满，先驱逐一个 victim 再 bind */ void drain_and_submit() { auto requests = scheduler_->drain_all(); @@ -309,18 +377,62 @@ class MeshResidencyManager { int tp = req.tp_part_idx; int expert = req.expert_id; - // 绑定 slot（如果尚未绑定） MeshSlotPool& pool = pools_[layer][tp]; - if (!pool.is_cached(expert)) { - // 需要先驱逐一个 slot（如果池满） - // B1: 在线驱逐尚未完全实现，这里先尝试 bind - // 如果池未满，bind 到空闲 slot - // 如果池满，需要驱逐（B1 修复后完善） + + // B1: 如果专家未缓存，需要分配 slot + void* gate_dst = req.gate_dst; + void* up_dst = req.up_dst; + void* down_dst = req.down_dst; + + if (!pool.is_cached(expert) && !is_gpu_expert(expert)) { + // 专家未缓存，需要分配 slot + int slot_idx = pool.expert_to_slot_idx(expert); + + if (slot_idx < 0) { + // 专家没有绑定 slot，需要分配 + // 先找空闲 slot（BASELINE 状态） + // 简化：用 find_evictable 找一个可覆盖的 slot + // 或者找 BASELINE 状态的 slot + slot_idx = -1; + for (int i = 0; i < pool.cap(); i++) { + if (pool.get_expert_state(pool.slot_to_expert_id(i)) == ExpertState::BASELINE || + pool.slot_to_expert_id(i) < 0) { + slot_idx = i; + break; + } + } + + if (slot_idx < 0) { + // 没有空闲 slot，需要驱逐 + slot_idx = decode_->evict_for_new_expert(pool, *scorer_, layer, expert); + if (slot_idx >= 0) { + // B9: 统计驱逐 + stats_.eviction_count++; + } + } + + if (slot_idx >= 0) { + // bind slot + pool.bind(slot_idx, expert); + // 获取 slot 指针 + gate_dst = pool.gate_ptr(slot_idx); + up_dst = pool.up_ptr(slot_idx); + down_dst = pool.down_ptr(slot_idx); + } + // slot_idx < 0 表示无法分配，跳过此专家 + } } + if (!gate_dst) continue; // 无法分配 slot，跳过 + + // B9: 统计 io_uring 读取 + stats_.io_uring_read_count++; + stats_.io_uring_read_bytes += static_cast( + layout.gate_bytes + layout.up_bytes + layout.down_bytes); + io_->submit_load( req.expert_id, req.tp_part_idx, layout, - req.gate_dst, req.up_dst, req.down_dst, + gate_dst, up_dst, down_dst, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, // scale 从 cache memcpy req.priority, /*layer_idx=*/layer, @@ -353,6 +465,10 @@ class MeshResidencyManager { // A4 fix: layouts_ 改为 [layer][tp][expert] 3D const std::vector>>& layouts() const { return layouts_; } + // B9: 统计接口 + const MeshStats& stats() const { return stats_; } + MeshStats& stats_mut() { return stats_; } + private: MeshConfig config_; std::vector numa_nodes_; @@ -367,6 +483,13 @@ class MeshResidencyManager { // GPU expert mask std::vector gpu_experts_mask_; + // B5: 跨层 defer 队列状态 + std::vector prev_layer_deferred_; // 上一层的 deferred 专家 + int total_deferred_ = 0; // 跨层累积 defer 计数 + + // B9: 运行时统计 + MeshStats stats_; + // 组件 std::unique_ptr io_; std::unique_ptr scheduler_; diff --git a/kt-kernel/operators/mesh/mesh_slot_pool.hpp b/kt-kernel/operators/mesh/mesh_slot_pool.hpp index 48f645da7..3a29d16f7 100644 --- a/kt-kernel/operators/mesh/mesh_slot_pool.hpp +++ b/kt-kernel/operators/mesh/mesh_slot_pool.hpp @@ -188,6 +188,17 @@ class MeshSlotPool { return slots_[slot_idx].get_state(); } + // B1 fix: 查询专家绑定的 slot_idx（驱逐用），未绑定返回 -1 + int expert_to_slot_idx(int expert_id) const { + return expert_to_slot_.lookup(expert_id); + } + + // B1 fix: 查询 slot 绑定的专家 ID（驱逐用），未绑定返回 -1 + int slot_to_expert_id(int slot_idx) const { + if (slot_idx < 0 || slot_idx >= cap_) return -1; + return slot_to_expert_[slot_idx]; + } + // ===== 绑定 / 覆盖 ===== // 绑定：将 expert 读入空闲 slot，设置指针指向 diff --git a/kt-kernel/python/utils/mesh/stats.py b/kt-kernel/python/utils/mesh/stats.py index e7c0fb90b..131ff1df0 100644 --- a/kt-kernel/python/utils/mesh/stats.py +++ b/kt-kernel/python/utils/mesh/stats.py @@ -86,6 +86,7 @@ def collect(self) -> MeshStats: """收集当前统计信息。""" stats = MeshStats() try: + # B9: C++ 侧已绑定 stats() 方法，返回 MeshStats 引用 cpp_stats = self._mgr.raw.stats() stats.cache_hit_count = getattr(cpp_stats, "cache_hit_count", 0) stats.cache_miss_count = getattr(cpp_stats, "cache_miss_count", 0) @@ -101,7 +102,7 @@ def collect(self) -> MeshStats: stats.decode_immediate_count = getattr(cpp_stats, "decode_immediate_count", 0) stats.decode_deferred_count = getattr(cpp_stats, "decode_deferred_count", 0) except (AttributeError, RuntimeError) as e: - logger.debug("Failed to collect MESH stats: %s", e) + logger.warning("Failed to collect MESH stats: %s", e) return stats def log_summary(self) -> None: From 265fa6a24da55f971935ca220a1151d4c13ae4c6 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Sun, 21 Jun 2026 15:52:29 +0800 Subject: [PATCH 07/10] fix: mesh v3 nan error --- kt-kernel/ext_bindings.cpp | 5 +- kt-kernel/operators/mesh/mesh_config.hpp | 5 + kt-kernel/operators/mesh/mesh_handoff.hpp | 30 ++++- kt-kernel/operators/mesh/mesh_io_uring.hpp | 107 ++++++++++++++--- kt-kernel/operators/mesh/mesh_prefill.hpp | 124 ++++++++++++-------- kt-kernel/operators/mesh/mesh_residency.hpp | 93 +++++++++++---- kt-kernel/operators/mesh/mesh_slot_pool.hpp | 85 ++++++++++++-- kt-kernel/python/utils/mesh/residency.py | 8 ++ kt-kernel/python/utils/mesh/wrapper.py | 8 ++ 9 files changed, 363 insertions(+), 102 deletions(-) diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 4d592f40f..01bf1cc3a 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -1091,7 +1091,10 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def_readwrite("down_mins_offset", &mesh::ExpertFileLayout::down_mins_offset) .def_readwrite("gate_mins_bytes", &mesh::ExpertFileLayout::gate_mins_bytes) .def_readwrite("up_mins_bytes", &mesh::ExpertFileLayout::up_mins_bytes) - .def_readwrite("down_mins_bytes", &mesh::ExpertFileLayout::down_mins_bytes); + .def_readwrite("down_mins_bytes", &mesh::ExpertFileLayout::down_mins_bytes) + // Bug 1 fix: BF16 down_proj 不连续读取支持 + .def_readwrite("down_stride", &mesh::ExpertFileLayout::down_stride) + .def_readwrite("down_rows", &mesh::ExpertFileLayout::down_rows); // B9: MeshStats 绑定 py::class_(mesh_module, "MeshStats") diff --git a/kt-kernel/operators/mesh/mesh_config.hpp b/kt-kernel/operators/mesh/mesh_config.hpp index cc8c16b9b..989a92c5d 100644 --- a/kt-kernel/operators/mesh/mesh_config.hpp +++ b/kt-kernel/operators/mesh/mesh_config.hpp @@ -41,6 +41,11 @@ struct ExpertFileLayout { size_t gate_mins_bytes = 0; size_t up_mins_bytes = 0; size_t down_mins_bytes = 0; + // Bug 1 fix: BF16 down_proj [E,H,I] 行主序，TP 沿 I 切不连续 + // down_stride > 0 表示 down 矩阵的行步长（字节），需逐行读取 + // down_stride == 0 表示连续（AMXINT4 per-TP 预切分或 BF16 tp=0） + size_t down_stride = 0; + int down_rows = 0; // down 矩阵行数（H），down_stride > 0 时有效 }; // MESH 配置 diff --git a/kt-kernel/operators/mesh/mesh_handoff.hpp b/kt-kernel/operators/mesh/mesh_handoff.hpp index 96068a8d4..645a0281a 100644 --- a/kt-kernel/operators/mesh/mesh_handoff.hpp +++ b/kt-kernel/operators/mesh/mesh_handoff.hpp @@ -122,14 +122,15 @@ class MeshHandoff { int expert_id = slot_idx; // 假设 GPU expert 占据前 GE 个 slot if (expert_id >= config.expert_num) break; - // 标记 active_readers 防止驱逐 - pool.acquire_reader(expert_id); + // Bug 11 fix: 用 acquire_slot_reader 按 slot_idx 保护，不依赖 expert 绑定 + // 原 acquire_reader(expert_id) 在 expert 未绑定时静默返回，无保护 + pool.acquire_slot_reader(slot_idx); // 搬运到 GPU（不读盘，纯内存→GPU 拷贝） move_gpu_expert_to_gpu(layer, expert_id); // 释放 reader - pool.release_reader(expert_id); + pool.release_slot_reader(slot_idx); // B7 fix: 先 overwrite 释放旧 slot 绑定，绑定新专家 int new_expert_id = cap + slot_idx; @@ -150,13 +151,32 @@ class MeshHandoff { // A6: 传入 layer_idx 和 on_complete 回调 const auto& layout = layouts[layer][tp][new_expert_id]; MeshSlotPool& pool_ref = pools[layer][tp]; + // Bug 3 fix: 捕获 slot_idx，mark_cached 参数是 slot_idx + int slot_idx_capture = slot_idx; + // Bug 4 fix: 捕获 io 和 layouts 引用，用于 copy scale + const auto& layouts_ref = layouts; io.submit_load(new_expert_id, tp, layout, gate_dst, up_dst, down_dst, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, ReadPriority::Demand, /*layer_idx=*/layer, - /*on_complete=*/[&pool_ref, new_expert_id](int, int, int) { - pool_ref.mark_cached(new_expert_id); + /*on_complete=*/[&pool_ref, &io, &layouts_ref, + layer, tp, new_expert_id, + slot_idx_capture, + &config](int, int, int) { + // Bug 4 fix: 先 copy scale，再 mark_cached + if (config.weight_type == WeightType::AMXINT4 && + io.scale_cache_loaded(layer)) { + void* gs = pool_ref.gate_scale_ptr(slot_idx_capture); + void* us = pool_ref.up_scale_ptr(slot_idx_capture); + void* ds = pool_ref.down_scale_ptr(slot_idx_capture); + io.copy_scale_from_cache(layer, tp, new_expert_id, + gs, us, ds, + nullptr, nullptr, nullptr, + layouts_ref[layer][tp]); + } + // Bug 3 fix: mark_cached 参数是 slot_idx + pool_ref.mark_cached(slot_idx_capture); }); } } diff --git a/kt-kernel/operators/mesh/mesh_io_uring.hpp b/kt-kernel/operators/mesh/mesh_io_uring.hpp index d00d4fd43..76b8facfe 100644 --- a/kt-kernel/operators/mesh/mesh_io_uring.hpp +++ b/kt-kernel/operators/mesh/mesh_io_uring.hpp @@ -13,6 +13,8 @@ #pragma once #include +#include +#include #include #include #include @@ -66,11 +68,14 @@ class MeshIoUring { // 启动时一次性把某层所有专家的 scale 数据读入 NUMA 本地 buffer // A4 fix: 改为按层预加载，因为每层的 scale 数据不同 + // 辅因2 fix: 按 TP 分片传入对应 NUMA 节点，避免跨 NUMA 访问 // expert_num 个专家，每个 scale_bytes 字节 - void preload_scale_cache(int expert_num, int tp_count, int numa_node, + void preload_scale_cache(int expert_num, int tp_count, + const std::vector& numa_nodes, const std::vector>& layouts_tp, int layer_idx) { - if (scale_cache_loaded_) return; + if (scale_cache_loaded_per_layer_.size() > (size_t)layer_idx && + scale_cache_loaded_per_layer_[layer_idx]) return; // 确保 scale_cache_ 有足够层 if ((int)scale_cache_.size() <= layer_idx) { @@ -78,8 +83,9 @@ class MeshIoUring { } auto& layer_cache = scale_cache_[layer_idx]; - // 计算每个 TP 分片的 scale 总大小 + // 计算每个 TP 分片的 scale 总大小，按 TP 对应 NUMA 节点分配 for (int tp = 0; tp < tp_count; tp++) { + int numa_node = numa_nodes[tp % numa_nodes.size()]; // 辅因2: 按 TP 选 NUMA ScaleCacheTP cache; size_t total_gate = 0, total_up = 0, total_down = 0; size_t total_gate_mins = 0, total_up_mins = 0, total_down_mins = 0; @@ -130,11 +136,16 @@ class MeshIoUring { } layer_cache.push_back(std::move(cache)); } - scale_cache_loaded_ = true; + // 第二轮根因1 fix: 按层设置标志，而非全局标志 + if ((int)scale_cache_loaded_per_layer_.size() <= layer_idx) { + scale_cache_loaded_per_layer_.resize(layer_idx + 1, false); + } + scale_cache_loaded_per_layer_[layer_idx] = true; } // 从 scale cache 中 memcpy 某个专家的 scale 到目标 buffer // A4 fix: 加 layer_idx 参数 + // 辅因1 fix: 支持 nullptr 参数（slot 中无 mins 空间时跳过 mins memcpy） void copy_scale_from_cache(int layer_idx, int tp_part_idx, int expert_id, void* gate_scale_dst, void* up_scale_dst, void* down_scale_dst, void* gate_mins_dst, void* up_mins_dst, void* down_mins_dst, @@ -155,15 +166,25 @@ class MeshIoUring { } const auto& layout = layouts_tp[expert_id]; const auto& cache = layer_cache[tp_part_idx]; - memcpy(gate_scale_dst, (char*)cache.gate_scale + off_g, layout.gate_scale_bytes); - memcpy(up_scale_dst, (char*)cache.up_scale + off_u, layout.up_scale_bytes); - memcpy(down_scale_dst, (char*)cache.down_scale + off_d, layout.down_scale_bytes); - memcpy(gate_mins_dst, (char*)cache.gate_mins + off_gm, layout.gate_mins_bytes); - memcpy(up_mins_dst, (char*)cache.up_mins + off_um, layout.up_mins_bytes); - memcpy(down_mins_dst, (char*)cache.down_mins + off_dm, layout.down_mins_bytes); + // scale 数据必须拷贝（BufferB 的 d 指针依赖） + if (gate_scale_dst) memcpy(gate_scale_dst, (char*)cache.gate_scale + off_g, layout.gate_scale_bytes); + if (up_scale_dst) memcpy(up_scale_dst, (char*)cache.up_scale + off_u, layout.up_scale_bytes); + if (down_scale_dst) memcpy(down_scale_dst, (char*)cache.down_scale + off_d, layout.down_scale_bytes); + // mins 数据可选拷贝（slot 中无 mins 空间时传 nullptr 跳过） + if (gate_mins_dst) memcpy(gate_mins_dst, (char*)cache.gate_mins + off_gm, layout.gate_mins_bytes); + if (up_mins_dst) memcpy(up_mins_dst, (char*)cache.up_mins + off_um, layout.up_mins_bytes); + if (down_mins_dst) memcpy(down_mins_dst, (char*)cache.down_mins + off_dm, layout.down_mins_bytes); } - bool scale_cache_loaded() const { return scale_cache_loaded_; } + // 第二轮根因1 fix: 按层查询 scale cache 是否已预加载 + bool scale_cache_loaded(int layer_idx = -1) const { + if (layer_idx < 0) { + // 无参：返回是否任意层已加载（用于粗略判断） + return !scale_cache_loaded_per_layer_.empty(); + } + return (size_t)layer_idx < scale_cache_loaded_per_layer_.size() && + scale_cache_loaded_per_layer_[layer_idx]; + } // ===== io_uring 异步读 ===== @@ -190,7 +211,15 @@ class MeshIoUring { ReadPriority priority, int layer_idx = -1, std::function on_complete = nullptr) { - int n_reqs = scale_cache_loaded_ ? 3 : 9; + // 第二轮根因1 fix: 按层判断 scale cache 是否已预加载 + // Bug 2 fix: BF16 路径 gate_scale_bytes == 0，无 scale/mins，只发 3 个 SQE + bool has_scale = (layout.gate_scale_bytes > 0); + bool layer_scale_loaded = has_scale && scale_cache_loaded(layer_idx); + // Bug 1 fix: down_stride > 0 表示 BF16 down_proj TP 切片不连续，用 pread 逐行同步读取 + // 此时 down 不走 io_uring，n_reqs 减 1 + bool down_scattered = (layout.down_stride > 0 && layout.down_rows > 0); + int n_reqs = layer_scale_loaded ? 3 : (has_scale ? 9 : 3); + if (down_scattered) n_reqs -= 1; // down 用 pread，不走 io_uring // A6: 创建 PendingRequest 跟踪完成状态 auto* pending = new PendingRequest{ @@ -206,11 +235,15 @@ class MeshIoUring { io_uring_prep_read(sqe, layout.fd, up_dst, layout.up_bytes, layout.up_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); - io_uring_prep_read(sqe, layout.fd, down_dst, layout.down_bytes, layout.down_offset); - io_uring_sqe_set_data(sqe, pending); + // Bug 5 fix: 先提交 gate/up SQE，再 pread down，让 io_uring 与 pread 并行 + // 原代码 pread 在 io_uring_submit 之前，阻塞期间 io_uring 空闲 + if (!down_scattered) { + sqe = io_uring_get_sqe(&ring_); + io_uring_prep_read(sqe, layout.fd, down_dst, layout.down_bytes, layout.down_offset); + io_uring_sqe_set_data(sqe, pending); + } - if (!scale_cache_loaded_) { + if (!layer_scale_loaded && has_scale) { sqe = io_uring_get_sqe(&ring_); io_uring_prep_read(sqe, layout.fd, gate_scale_dst, layout.gate_scale_bytes, layout.gate_scale_offset); io_uring_sqe_set_data(sqe, pending); @@ -236,7 +269,46 @@ class MeshIoUring { io_uring_sqe_set_data(sqe, pending); } + // Bug 5 fix: 先 submit 所有 SQE，让 io_uring 开始 DMA io_uring_submit(&ring_); + + // Bug 5 fix: down_scattered 的 pread 在 submit 之后执行，与 io_uring 并行 + if (down_scattered) { + // Bug 1 fix: BF16 down_proj [E,H,I] 行主序，TP 沿 I 切不连续 + // 用 pread 逐行同步读取到 slot buffer 的连续区域 + size_t row_bytes = layout.down_bytes / layout.down_rows; + char* dst = static_cast(down_dst); + off_t src_off = layout.down_offset; + for (int r = 0; r < layout.down_rows; r++) { + // Bug 6 fix: pread 加错误处理，失败 abort 避免后续 use-after-free + // 注意：不能 throw 或 delete pending，因为 gate/up SQE 已提交， + // CQE 处理时会访问 pending，throw/delete 会导致 use-after-free + ssize_t ret = pread(layout.fd, dst + r * row_bytes, row_bytes, src_off); + if (ret < 0) { + fprintf(stderr, + "[MESH] pread down scattered failed at row %d/%d, expert=%d, " + "errno=%d (%s), aborting\n", + r, layout.down_rows, expert_id, errno, strerror(errno)); + std::abort(); + } + if (ret < (ssize_t)row_bytes) { + // 部分读，补齐剩余（理论上 pread 对普通文件应一次读完） + ssize_t done = ret; + while (done < (ssize_t)row_bytes) { + ret = pread(layout.fd, dst + r * row_bytes + done, + row_bytes - done, src_off + done); + if (ret < 0) { + fprintf(stderr, + "[MESH] pread down scattered partial fail, errno=%d (%s), " + "aborting\n", errno, strerror(errno)); + std::abort(); + } + done += ret; + } + } + src_off += layout.down_stride; + } + } } // A6: 处理已完成的 CQE，触发 on_complete 回调 @@ -369,7 +441,8 @@ class MeshIoUring { }; // A4 fix: [layer][tp] 的 scale cache std::vector> scale_cache_; - bool scale_cache_loaded_ = false; + // 第二轮根因1 fix: 改为按层标志，避免全局标志导致 l>=1 层跳过预加载 + std::vector scale_cache_loaded_per_layer_; // 同步读取（启动阶段用，pread + O_DIRECT） void sync_read_to(int fd, off_t offset, size_t bytes, void* dst) { diff --git a/kt-kernel/operators/mesh/mesh_prefill.hpp b/kt-kernel/operators/mesh/mesh_prefill.hpp index e846b1705..fd0fca890 100644 --- a/kt-kernel/operators/mesh/mesh_prefill.hpp +++ b/kt-kernel/operators/mesh/mesh_prefill.hpp @@ -21,6 +21,7 @@ */ #pragma once +#include #include #include #include @@ -45,23 +46,24 @@ class MeshPrefill { enum class TemporalRole { COMPUTE, PREFETCH }; MeshPrefill(int num_layers, int expert_num, int cap, int tp_count, - size_t slot_bytes, int numa_node) + size_t slot_bytes, const std::vector& numa_nodes) : num_layers_(num_layers), expert_num_(expert_num), cap_(cap), tp_count_(tp_count), slot_bytes_(slot_bytes), - numa_node_(numa_node) { + numa_nodes_(numa_nodes) { // temporal 长度 = 单层专家总数 - slot cap temporal_size_ = (expert_num > cap) ? (expert_num - cap) : 0; } ~MeshPrefill() { release_temporal(); } - // 初始化 temporal 双缓冲（每 TP 一对 A/B） + // 初始化 temporal 双缓冲（每 TP 一对 A/B，按 TP 对应 NUMA 节点分配） void init_temporal() { - if (temporal_a_ != nullptr) return; // 已初始化 - size_t bytes = static_cast(temporal_size_) * slot_bytes_ * tp_count_; + if (!temporal_a_vec_.empty()) return; // 已初始化 + // 辅因3: 每 TP 独立分配，单 TP 字节数 = temporal_size * slot_bytes + size_t per_tp_bytes = static_cast(temporal_size_) * slot_bytes_; role_a_ = TemporalRole::COMPUTE; role_b_ = TemporalRole::PREFETCH; // B4 fix: 初始化 expert_id -> temporal_offset 映射 @@ -72,28 +74,33 @@ class MeshPrefill { temporal_expert_map_[tp][e] = e - cap_; } } - if (bytes == 0) { + if (per_tp_bytes == 0) { // cap >= expert_num: 所有专家常驻 slot，不需要 temporal buffer return; } - temporal_a_ = numa_alloc_onnode(bytes, numa_node_); - temporal_b_ = numa_alloc_onnode(bytes, numa_node_); - if (!temporal_a_ || !temporal_b_) { - throw std::runtime_error("MeshPrefill: numa_alloc_onnode for temporal failed"); + temporal_a_vec_.assign(tp_count_, nullptr); + temporal_b_vec_.assign(tp_count_, nullptr); + for (int tp = 0; tp < tp_count_; tp++) { + int numa_node = numa_nodes_[tp % numa_nodes_.size()]; // 辅因3: 按 TP 选 NUMA + temporal_a_vec_[tp] = numa_alloc_onnode(per_tp_bytes, numa_node); + temporal_b_vec_[tp] = numa_alloc_onnode(per_tp_bytes, numa_node); + if (!temporal_a_vec_[tp] || !temporal_b_vec_[tp]) { + throw std::runtime_error("MeshPrefill: numa_alloc_onnode for temporal failed"); + } } - temporal_bytes_ = bytes; + temporal_bytes_ = per_tp_bytes; } // 释放 temporal 内存（仅 prefill→decode 切换时调用） void release_temporal() { - if (temporal_a_) { - numa_free(temporal_a_, temporal_bytes_); - temporal_a_ = nullptr; + for (void* p : temporal_a_vec_) { + if (p) numa_free(p, temporal_bytes_); } - if (temporal_b_) { - numa_free(temporal_b_, temporal_bytes_); - temporal_b_ = nullptr; + for (void* p : temporal_b_vec_) { + if (p) numa_free(p, temporal_bytes_); } + temporal_a_vec_.clear(); + temporal_b_vec_.clear(); } // ===== Prefill 主流程 ===== @@ -117,10 +124,7 @@ class MeshPrefill { MeshSlotPool& slot_pool, MeshScheduler& scheduler, MeshIoUring* io, std::function&)> amx_forward) { - void* compute_buf = compute_temporal(); - void* prefetch_buf = prefetch_temporal(); - (void)compute_buf; // compute_buf 由 amx_forward 通过 slot_pool/temporal_ptr 访问 - + // 辅因3: temporal buffer 按 TP 分片，compute/prefetch 通过 tp 索引访问 for (int c = 0; c < total_chunks; c++) { const auto& active = active_experts_per_chunk[c]; @@ -130,7 +134,7 @@ class MeshPrefill { // b. 对 temporal_B 发起第 L+1 层的异步预取（覆盖另一个 temporal） if (layer_idx + 1 < num_layers_) { - submit_next_layer_prefetch(layer_idx + 1, active, scheduler, prefetch_buf); + submit_next_layer_prefetch(layer_idx + 1, active, scheduler); } // c. 本层计算使用 temporal_A 中的临时专家 + slot 池常驻专家 @@ -159,33 +163,40 @@ class MeshPrefill { role_b_ = tmp; } - // 获取当前 COMPUTE 角色的 temporal buffer - void* compute_temporal() const { - return role_a_ == TemporalRole::COMPUTE ? temporal_a_ : temporal_b_; + // 获取当前 COMPUTE 角色的 temporal buffer（辅因3: 按 TP 索引） + void* compute_temporal(int tp) const { + if (tp < 0 || tp >= tp_count_) return nullptr; + const auto& vec = (role_a_ == TemporalRole::COMPUTE) ? temporal_a_vec_ : temporal_b_vec_; + if (tp >= (int)vec.size()) return nullptr; + return vec[tp]; } - // 获取当前 PREFETCH 角色的 temporal buffer - void* prefetch_temporal() const { - return role_a_ == TemporalRole::PREFETCH ? temporal_a_ : temporal_b_; + // 获取当前 PREFETCH 角色的 temporal buffer（辅因3: 按 TP 索引） + void* prefetch_temporal(int tp) const { + if (tp < 0 || tp >= tp_count_) return nullptr; + const auto& vec = (role_a_ == TemporalRole::PREFETCH) ? temporal_a_vec_ : temporal_b_vec_; + if (tp >= (int)vec.size()) return nullptr; + return vec[tp]; } /** * @brief B4 fix: 获取 temporal buffer 中某专家某矩阵的指针 * - * @param buf temporal buffer 基地址（compute_temporal() 或 prefetch_temporal()） - * @param tp TP 分片 + * @param buf 某 TP 的 temporal buffer 基地址（compute_temporal(tp) 或 prefetch_temporal(tp)） * @param expert_id 专家 ID * @param matrix_idx 0=gate, 1=up, 2=down * @return void* 指针，未映射返回 nullptr + * + * 辅因3: buf 已是 per-TP 指针，无需再按 tp 偏移 */ - void* temporal_ptr(void* buf, int tp, int expert_id, int matrix_idx) const { - if (tp < 0 || tp >= tp_count_) return nullptr; + void* temporal_ptr(void* buf, int expert_id, int matrix_idx) const { if (expert_id < 0 || expert_id >= expert_num_) return nullptr; - int offset = temporal_expert_map_[tp][expert_id]; + // expert_id -> offset 映射对所有 TP 相同（cap ~ expert_num-1） + int offset = (temporal_expert_map_.empty() || temporal_expert_map_[0].empty()) + ? -1 : temporal_expert_map_[0][expert_id]; if (offset < 0) return nullptr; - // 布局：[tp][offset] 每个 slot_bytes_，内含 gate + up + down + // 布局：[offset] 每个 slot_bytes_，内含 gate + up + down char* base = static_cast(buf) + - static_cast(tp) * temporal_size_ * slot_bytes_ + static_cast(offset) * slot_bytes_; // gate 在偏移 0，up 在偏移 gate_up_bytes_，down 在偏移 gate_up_bytes_*2 size_t matrix_offset = (matrix_idx == 0) ? 0 : @@ -198,7 +209,7 @@ class MeshPrefill { // ===== 访问器 ===== int temporal_size() const { return temporal_size_; } - bool temporal_ready() const { return temporal_a_ != nullptr; } + bool temporal_ready() const { return !temporal_a_vec_.empty(); } private: int num_layers_; @@ -207,12 +218,12 @@ class MeshPrefill { int tp_count_; size_t slot_bytes_; size_t gate_up_bytes_ = 0; // B4: gate/up 单块字节数 - int numa_node_; - size_t temporal_bytes_ = 0; + std::vector numa_nodes_; // 辅因3: 每 TP 对应的 NUMA 节点 + size_t temporal_bytes_ = 0; // 辅因3: 单 TP temporal 字节数 - // temporal 双缓冲 - void* temporal_a_ = nullptr; - void* temporal_b_ = nullptr; + // temporal 双缓冲（辅因3: 每 TP 独立分配，按 NUMA 节点定位） + std::vector temporal_a_vec_; // size == tp_count_ + std::vector temporal_b_vec_; // size == tp_count_ TemporalRole role_a_ = TemporalRole::COMPUTE; TemporalRole role_b_ = TemporalRole::PREFETCH; int temporal_size_ = 0; // = expert_num - cap @@ -224,14 +235,25 @@ class MeshPrefill { // B3 fix: 调用 io_uring process_cqes 推进 CQE，直到所有活跃专家 CACHED void wait_temporal_ready(int layer_idx, const std::vector& active, MeshSlotPool& slot_pool, MeshIoUring* io) { + // Bug 10 fix: 加超时保护，避免 expert 未被调度时死循环 + const int kMaxWaitIters = 1000000; // ~几秒 + int wait_iters = 0; for (int eid : active) { if (slot_pool.is_cached(eid)) continue; // 阻塞等待 io_uring CQE 到达 - // 实际实现需要与 MeshIoUring 配合轮询 process_cqes + // Bug 10 fix: 只用 process_cqes，不用 submit_and_wait + // submit_and_wait 内部也处理 CQE 并 delete pending， + // 与 process_cqes 同时调用会导致 double free if (io) { while (!slot_pool.is_cached(eid)) { io->process_cqes(); - io->submit_and_wait(); + if (++wait_iters > kMaxWaitIters) { + fprintf(stderr, + "[MESH] wait_temporal_ready: expert %d never cached at layer %d, " + "possible scheduler deadlock\n", + eid, layer_idx); + break; // 避免死循环，返回后 AMX 会读到未就绪数据（nullptr 检查） + } } } } @@ -239,20 +261,24 @@ class MeshPrefill { // 提交下一层的异步预取 // B4 fix: 用 temporal_ptr 正确寻址，不再用 ptrs[eid*3] 越界 + // 辅因3: prefetch_buf 按 TP 从 prefetch_temporal(tp) 获取 void submit_next_layer_prefetch(int next_layer, const std::vector& active, - MeshScheduler& scheduler, - void* prefetch_buf) { + MeshScheduler& scheduler) { for (int tp = 0; tp < tp_count_; tp++) { + void* prefetch_buf = prefetch_temporal(tp); + if (!prefetch_buf) continue; for (int eid : active) { // 只预取不在 slot 中的专家（在 temporal 中的） - int offset = (tp < (int)temporal_expert_map_.size() && + // expert_id -> offset 映射对所有 TP 相同 + int offset = (!temporal_expert_map_.empty() && + tp < (int)temporal_expert_map_.size() && eid < (int)temporal_expert_map_[tp].size()) ? temporal_expert_map_[tp][eid] : -1; if (offset < 0) continue; - void* gate = temporal_ptr(prefetch_buf, tp, eid, 0); - void* up = temporal_ptr(prefetch_buf, tp, eid, 1); - void* down = temporal_ptr(prefetch_buf, tp, eid, 2); + void* gate = temporal_ptr(prefetch_buf, eid, 0); + void* up = temporal_ptr(prefetch_buf, eid, 1); + void* down = temporal_ptr(prefetch_buf, eid, 2); if (!gate || !up || !down) continue; scheduler.submit_prefill(next_layer, eid, tp, gate, up, down, ReadPriority::Prefetch); diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp index abde6af98..680986022 100644 --- a/kt-kernel/operators/mesh/mesh_residency.hpp +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -81,6 +81,11 @@ class MeshResidencyManager { for (int tp = 0; tp < config.tp_count; tp++) { pools_[l].emplace_back(l, tp, numa_nodes[tp], config.cap, slot_bytes_); pools_[l][tp].set_gate_up_bytes(compute_gate_up_bytes(config)); + // AMXINT4: 设置纯权重大小（不含 scale），用于 slot 内 scale 指针偏移 + if (config.weight_type == WeightType::AMXINT4) { + pools_[l][tp].set_gate_up_weights_bytes(compute_gate_up_weights_bytes(config)); + pools_[l][tp].set_down_weights_bytes(compute_down_weights_bytes(config)); + } pools_[l][tp].init_expert_map(config.expert_num); } } @@ -91,7 +96,7 @@ class MeshResidencyManager { scorer_ = std::make_unique(config.total_layers, config.expert_num, config); prefill_ = std::make_unique(config.total_layers, config.expert_num, config.cap, config.tp_count, - slot_bytes_, numa_nodes[0]); + slot_bytes_, numa_nodes); decode_ = std::make_unique(config); handoff_ = std::make_unique(); @@ -153,7 +158,7 @@ class MeshResidencyManager { if (config_.weight_type == WeightType::AMXINT4) { for (int l = 0; l < config_.total_layers; l++) { io_->preload_scale_cache(config_.expert_num, config_.tp_count, - numa_nodes_[0], layouts_[l], l); + numa_nodes_, layouts_[l], l); } } @@ -174,11 +179,28 @@ class MeshResidencyManager { // 提交 io_uring 读 // A6: 传入 layer_idx 和 on_complete 回调，完成后 mark_cached + // 辅因1 fix: on_complete 中调用 copy_scale_from_cache 把 scale 写入 slot io_->submit_load(e, tp, layout, gate_dst, up_dst, down_dst, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, ReadPriority::Demand, /*layer_idx=*/l, /*on_complete=*/[this, l, tp, e](int, int, int) { + // Bug 4 fix: 先 copy scale，再 mark_cached + // mark_cached 后状态为 CACHED，AMX kernel 可能立即读取 + // 若先 mark_cached 后 copy scale，AMX 会读到零 scale → NaN + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(l)) { + // Bug 3 fix: bootstrap 中 bind(e, e) 使 slot_idx == e + int slot_idx = pools_[l][tp].expert_to_slot_idx(e); + if (slot_idx >= 0) { + void* gs = pools_[l][tp].gate_scale_ptr(slot_idx); + void* us = pools_[l][tp].up_scale_ptr(slot_idx); + void* ds = pools_[l][tp].down_scale_ptr(slot_idx); + io_->copy_scale_from_cache(l, tp, e, gs, us, ds, + nullptr, nullptr, nullptr, + layouts_[l][tp]); + } + } + // Bug 3 fix: mark_cached 参数是 slot_idx，bootstrap 中 slot_idx == e pools_[l][tp].mark_cached(e); }); } @@ -390,17 +412,8 @@ class MeshResidencyManager { if (slot_idx < 0) { // 专家没有绑定 slot，需要分配 - // 先找空闲 slot（BASELINE 状态） - // 简化：用 find_evictable 找一个可覆盖的 slot - // 或者找 BASELINE 状态的 slot - slot_idx = -1; - for (int i = 0; i < pool.cap(); i++) { - if (pool.get_expert_state(pool.slot_to_expert_id(i)) == ExpertState::BASELINE || - pool.slot_to_expert_id(i) < 0) { - slot_idx = i; - break; - } - } + // Bug 7 fix: 用 find_free_slot 替代线性扫描 + 双重映射 + slot_idx = pool.find_free_slot(); if (slot_idx < 0) { // 没有空闲 slot，需要驱逐 @@ -425,6 +438,11 @@ class MeshResidencyManager { if (!gate_dst) continue; // 无法分配 slot，跳过 + // Bug 3 fix: 查询 expert 绑定的 slot_idx，传给 lambda + // mark_cached 参数是 slot_idx，不是 expert_id + int slot_idx = pool.expert_to_slot_idx(expert); + if (slot_idx < 0) continue; // 未绑定 slot，跳过 + // B9: 统计 io_uring 读取 stats_.io_uring_read_count++; stats_.io_uring_read_bytes += static_cast( @@ -436,8 +454,20 @@ class MeshResidencyManager { nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, // scale 从 cache memcpy req.priority, /*layer_idx=*/layer, - /*on_complete=*/[this, layer, tp, expert, orig_on_complete](int, int, int) { - pools_[layer][tp].mark_cached(expert); + /*on_complete=*/[this, layer, tp, expert, slot_idx, orig_on_complete](int, int, int) { + // Bug 4 fix: 先 copy scale，再 mark_cached + // mark_cached 后状态为 CACHED，AMX kernel 可能立即读取 + // 若先 mark_cached 后 copy scale，AMX 会读到零 scale → NaN + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(layer)) { + void* gs = pools_[layer][tp].gate_scale_ptr(slot_idx); + void* us = pools_[layer][tp].up_scale_ptr(slot_idx); + void* ds = pools_[layer][tp].down_scale_ptr(slot_idx); + io_->copy_scale_from_cache(layer, tp, expert, gs, us, ds, + nullptr, nullptr, nullptr, + layouts_[layer][tp]); + } + // Bug 3 fix: mark_cached 参数是 slot_idx，不是 expert_id + pools_[layer][tp].mark_cached(slot_idx); if (orig_on_complete) orig_on_complete(); }); } @@ -498,25 +528,27 @@ class MeshResidencyManager { std::unique_ptr decode_; std::unique_ptr handoff_; - // 计算 slot 字节数（gate + up + down 三个矩阵） + // 计算 slot 字节数（gate + up + down 三个矩阵，AMXINT4 含 scale） size_t compute_slot_bytes(const MeshConfig& config) const { // 根据 SKILL 第 8 节的双 NUMA 拆分： - // gate: [hidden, intermediate/tp_count] - // up: [hidden, intermediate/tp_count] - // down: [intermediate/tp_count, hidden] + // gate: [hidden, intermediate/tp_count] BufferB n=intermediate/tp, k=hidden + // up: [hidden, intermediate/tp_count] BufferB n=intermediate/tp, k=hidden + // down: [intermediate/tp_count, hidden] BufferB n=hidden, k=intermediate/tp + // AMXINT4 BufferB 期望: n*k/2 (权重) + n*sizeof(float) (scale) 连续存放 size_t gate_up_bytes = compute_gate_up_bytes(config); size_t down_bytes = compute_down_bytes(config); return gate_up_bytes * 2 + down_bytes; } + // gate 或 up 单块字节数（权重 + scale） + // BufferB 构造: n=intermediate/tp, k=hidden → 权重=n*k/2, scale=n*4 size_t compute_gate_up_bytes(const MeshConfig& config) const { - // gate 或 up 单块字节数 int h = config.hidden_size; int i = config.intermediate_size / config.tp_count; // TP 切分后 size_t elements = static_cast(h) * i; switch (config.weight_type) { case WeightType::AMXINT4: - return elements / 2; // int4 = 0.5 byte + return elements / 2 + static_cast(i) * sizeof(float); // int4 权重 + scale (n=i) case WeightType::BF16: return elements * 2; // bf16 = 2 bytes default: @@ -524,20 +556,35 @@ class MeshResidencyManager { } } + // down 单块字节数（权重 + scale） + // BufferB 构造: n=hidden, k=intermediate/tp → 权重=n*k/2, scale=n*4 size_t compute_down_bytes(const MeshConfig& config) const { - // down 单块字节数 int h = config.hidden_size; int i = config.intermediate_size / config.tp_count; size_t elements = static_cast(i) * h; switch (config.weight_type) { case WeightType::AMXINT4: - return elements / 2; + return elements / 2 + static_cast(h) * sizeof(float); // int4 权重 + scale (n=h) case WeightType::BF16: return elements * 2; default: return elements * 2; } } + + // AMXINT4: gate/up 的纯权重大小（不含 scale），用于 slot 内 scale 指针偏移 + size_t compute_gate_up_weights_bytes(const MeshConfig& config) const { + int h = config.hidden_size; + int i = config.intermediate_size / config.tp_count; + return static_cast(h) * i / 2; + } + + // AMXINT4: down 的纯权重大小（不含 scale） + size_t compute_down_weights_bytes(const MeshConfig& config) const { + int h = config.hidden_size; + int i = config.intermediate_size / config.tp_count; + return static_cast(i) * h / 2; + } }; } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_slot_pool.hpp b/kt-kernel/operators/mesh/mesh_slot_pool.hpp index 3a29d16f7..1a88530be 100644 --- a/kt-kernel/operators/mesh/mesh_slot_pool.hpp +++ b/kt-kernel/operators/mesh/mesh_slot_pool.hpp @@ -12,6 +12,8 @@ #include #include +#include +#include #include #include #include @@ -114,6 +116,8 @@ class MeshSlotPool { cap_(other.cap_), slot_bytes_(other.slot_bytes_), gate_up_bytes_(other.gate_up_bytes_), + gate_up_weights_bytes_(other.gate_up_weights_bytes_), + down_weights_bytes_(other.down_weights_bytes_), memory_(other.memory_), total_bytes_(other.total_bytes_) { other.memory_ = nullptr; @@ -131,6 +135,8 @@ class MeshSlotPool { cap_ = other.cap_; slot_bytes_ = other.slot_bytes_; gate_up_bytes_ = other.gate_up_bytes_; + gate_up_weights_bytes_ = other.gate_up_weights_bytes_; + down_weights_bytes_ = other.down_weights_bytes_; memory_ = other.memory_; total_bytes_ = other.total_bytes_; other.memory_ = nullptr; @@ -154,6 +160,21 @@ class MeshSlotPool { return static_cast(memory_) + static_cast(slot_idx) * slot_bytes_ + gate_up_bytes_ * 2; } + // AMXINT4: 获取 slot 中 gate scale 的地址（紧跟 gate 权重之后） + // BufferB 期望 d = ptr + n*k/2，n=intermediate/tp, k=hidden + void* gate_scale_ptr(int slot_idx) { + return static_cast(gate_ptr(slot_idx)) + gate_up_weights_bytes_; + } + // AMXINT4: 获取 slot 中 up scale 的地址 + void* up_scale_ptr(int slot_idx) { + return static_cast(up_ptr(slot_idx)) + gate_up_weights_bytes_; + } + // AMXINT4: 获取 slot 中 down scale 的地址 + // down 的 BufferB n=hidden, k=intermediate/tp → 权重=n*k/2=hidden*(intermediate/tp)/2 + void* down_scale_ptr(int slot_idx) { + return static_cast(down_ptr(slot_idx)) + down_weights_bytes_; + } + // 根据专家 ID 获取 gate 指针（KT 计算用），未缓存返回 nullptr void* expert_gate_ptr(int expert_id) { int slot_idx = expert_to_slot_.lookup(expert_id); @@ -188,6 +209,24 @@ class MeshSlotPool { return slots_[slot_idx].get_state(); } + // Bug 7 fix: 直接按 slot_idx 查状态，避免 expert_id->slot_idx->expert_id 双重映射 + ExpertState slot_state(int slot_idx) const { + if (slot_idx < 0 || slot_idx >= cap_) return ExpertState::BASELINE; + return slots_[slot_idx].get_state(); + } + + // Bug 7 fix: 找空闲 slot（BASELINE 或未绑定），返回 slot_idx，找不到返回 -1 + // 替代 drain_and_submit 中的线性扫描 + 双重映射 + int find_free_slot() const { + for (int i = 0; i < cap_; i++) { + if (slot_to_expert_[i] < 0 || + slots_[i].get_state() == ExpertState::BASELINE) { + return i; + } + } + return -1; + } + // B1 fix: 查询专家绑定的 slot_idx（驱逐用），未绑定返回 -1 int expert_to_slot_idx(int expert_id) const { return expert_to_slot_.lookup(expert_id); @@ -221,16 +260,29 @@ class MeshSlotPool { // 必须等 active_readers 归零后才能执行 void overwrite(int slot_idx, int new_expert_id) { Slot& s = slots_[slot_idx]; - // 等待活跃 reader 归零 - while (s.active_readers.load(std::memory_order_acquire) > 0) { - // spin wait，实际实现可用 futex + // Bug 8 fix: spin-wait 加超时（10 秒），避免 reader 泄漏导致永久死锁 + // 超时后 abort，因为继续执行会导致数据损坏 + { + const int kMaxSpinIters = 100000000; // ~10s 取决于 CPU + int spin = 0; + while (s.active_readers.load(std::memory_order_acquire) > 0) { + if (++spin > kMaxSpinIters) { + fprintf(stderr, + "[MESH] overwrite: slot %d active_readers never reached 0 " + "(expert %d -> %d), aborting to avoid data corruption\n", + slot_idx, s.bound_expert_id, new_expert_id); + std::abort(); + } + // Bug 8 fix: 轻量让步，减少空转功耗（不用 sched_yield 避免系统调用开销） + } } int old_expert_id = s.bound_expert_id; if (old_expert_id >= 0) { expert_to_slot_.erase(old_expert_id); } - s.set_state(ExpertState::DEMOTING); - // 解绑完成，进入 LOADING 状态等待新数据 + // Bug 8 fix: 去掉无意义的 DEMOTING 中间状态 + // 原代码 DEMOTING 立即被 LOADING 覆盖，无窗口让其他线程看到 + // reader==0 已保证安全，直接进 LOADING s.set_state(ExpertState::LOADING); s.bound_expert_id = new_expert_id; slot_to_expert_[slot_idx] = new_expert_id; @@ -253,6 +305,17 @@ class MeshSlotPool { slots_[slot_idx].active_readers.fetch_sub(1, std::memory_order_acq_rel); } + // Bug 11 fix: 按 slot_idx 递增/递减 reader，不依赖 expert 绑定关系 + // 用于 handoff 等场景，expert 可能未绑定但 slot 需要保护 + void acquire_slot_reader(int slot_idx) { + if (slot_idx < 0 || slot_idx >= cap_) return; + slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + } + void release_slot_reader(int slot_idx) { + if (slot_idx < 0 || slot_idx >= cap_) return; + slots_[slot_idx].active_readers.fetch_sub(1, std::memory_order_acq_rel); + } + // ===== 驱逐扫描 ===== // 找一个可驱逐的 slot：状态为 CACHED 且 active_readers==0 @@ -286,16 +349,24 @@ class MeshSlotPool { int cap() const { return cap_; } size_t slot_bytes() const { return slot_bytes_; } - // 设置 gate/up 的单块字节数（用于指针偏移计算） + // 设置 gate/up 的单块字节数（含 scale，用于指针偏移计算） void set_gate_up_bytes(size_t bytes) { gate_up_bytes_ = bytes; } + // AMXINT4: 设置 gate/up 的纯权重大小（不含 scale，用于 scale 指针偏移） + void set_gate_up_weights_bytes(size_t bytes) { gate_up_weights_bytes_ = bytes; } + + // AMXINT4: 设置 down 的纯权重大小（不含 scale，用于 down scale 指针偏移） + void set_down_weights_bytes(size_t bytes) { down_weights_bytes_ = bytes; } + private: int layer_idx_; int tp_part_idx_; int numa_node_; int cap_; size_t slot_bytes_; - size_t gate_up_bytes_ = 0; // gate 或 up 单块字节数 + size_t gate_up_bytes_ = 0; // gate 或 up 单块字节数（含 scale） + size_t gate_up_weights_bytes_ = 0; // AMXINT4: gate/up 纯权重大小（不含 scale） + size_t down_weights_bytes_ = 0; // AMXINT4: down 纯权重大小（不含 scale） void* memory_ = nullptr; size_t total_bytes_ = 0; diff --git a/kt-kernel/python/utils/mesh/residency.py b/kt-kernel/python/utils/mesh/residency.py index b47e466cf..716b7a100 100644 --- a/kt-kernel/python/utils/mesh/residency.py +++ b/kt-kernel/python/utils/mesh/residency.py @@ -76,6 +76,8 @@ def set_file_layout( gate_mins_bytes: int = 0, up_mins_bytes: int = 0, down_mins_bytes: int = 0, + down_stride: int = 0, + down_rows: int = 0, ) -> None: """注入单个专家在某层某 TP 分片上的文件布局。 @@ -113,6 +115,8 @@ def set_file_layout( layout.gate_mins_bytes = gate_mins_bytes layout.up_mins_bytes = up_mins_bytes layout.down_mins_bytes = down_mins_bytes + layout.down_stride = down_stride + layout.down_rows = down_rows self._mgr.set_file_layout(layer_idx, tp_part_idx, expert_id, layout) def set_gpu_experts_mask(self, mask: List[int]) -> None: @@ -194,6 +198,10 @@ def on_decode_token_end(self, all_layers_topk: List[List[int]], def config(self): return self._mgr.config() + def raw_ptr(self) -> int: + """获取底层 C++ ResidencyManager 的内存地址（用于 hook 注册）。""" + return self._mgr.raw_ptr() + @property def raw(self): """获取底层 C++ 对象指针（用于 hook 注册）。""" diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py index 9d7cf270d..7f4128b8f 100644 --- a/kt-kernel/python/utils/mesh/wrapper.py +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -367,6 +367,11 @@ def get_fd(fpath: str) -> int: # down_proj: [E, H, I]，每个专家占 H*I*elem_bytes expert_stride_d = hidden * inter * elem_bytes down_bytes = hidden * inter_per_tp * elem_bytes + # Bug 1 fix: BF16 down_proj [E,H,I] 行主序，TP 沿 I 切不连续 + # down_stride = 完整行字节数（I * elem_bytes），down_rows = H + # tp_count > 1 时 down_stride > row_bytes，需逐行读取 + down_stride = inter * elem_bytes # 完整行步长 + down_rows = hidden for tp in range(tp_count): for expert in range(expert_num): @@ -389,6 +394,9 @@ def get_fd(fpath: str) -> int: up_bytes=up_bytes, down_bytes=down_bytes, # BF16 无 scale/mins + # Bug 1 fix: 传入 down_stride/down_rows 支持不连续读取 + down_stride=down_stride, + down_rows=down_rows, ) injected += 1 From 1f138a441c51e3c2bd76a869100efba85c8d890f Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Sun, 21 Jun 2026 19:06:16 +0800 Subject: [PATCH 08/10] fix: 35b run success with --model directed to bf16 model and kt directed to amxint4 model --- TROUBLESHOOTING_LOG.md | 62 ++++++++++++++++++ scripts/run_mesh_35b_cap.sh | 39 +++++++++++ scripts/run_mesh_397b_cap.sh | 39 +++++++++++ scripts/run_mesh_bench.sh | 123 +++++++++++++++++++++++++++++++++++ 4 files changed, 263 insertions(+) create mode 100644 TROUBLESHOOTING_LOG.md create mode 100644 scripts/run_mesh_35b_cap.sh create mode 100644 scripts/run_mesh_397b_cap.sh create mode 100644 scripts/run_mesh_bench.sh diff --git a/TROUBLESHOOTING_LOG.md b/TROUBLESHOOTING_LOG.md new file mode 100644 index 000000000..bc7d6a04e --- /dev/null +++ b/TROUBLESHOOTING_LOG.md @@ -0,0 +1,62 @@ +# 踩坑日志 + +## 2026-06-21: 35B TP>1 输出坍缩为 `!` + +### 现象 +- 35B AMXINT4 模型在 TP>1（TP2/TP4）时，输出全部坍缩为 `!` +- TP1 正常 +- layer 3 的 GPU w2_weight 包含 NaN/Inf + +### 根因 +启动脚本中 `--model` 和 `--kt-weight-path` 都指向 AMXINT4 格式 checkpoint（`Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH`）。 + +AMXINT4 格式的 checkpoint key 为 `blk.N.ffn_down_exps.E.numa.T.weight` + `.scale`，按 NUMA/TP 预切分存储。sglang 标准 weight_loader（`_weight_loader_physical` → `_weight_loader_impl` → `_load_w2`）只处理标准格式 key（`experts.N.down_proj.`），不识别 AMXINT4 格式。 + +因此 GPU 的 `w2_weight`（由 `UnquantizedFusedMoEMethod.create_weights` 创建为 `torch.empty`）从未被 checkpoint 填充，保持未初始化状态。TP1 时恰好不含 NaN，TP>1 时 layer 3 的 w2_weight 包含 NaN/Inf，导致输出 NaN 坍缩为 `!`。 + +### 修复 +- `--model` 改为标准 bf16 格式 checkpoint（`Qwen3.5-35B-A3B-Unfused`） +- `--kt-weight-path` 保持 AMXINT4 格式（供 CPU AMX 内核使用） + +``` +--model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ +--kt-weight-path /mnt/data2/models/Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH/ +``` + +### 验证 +- DIAG 输出确认 `_weight_loader_physical` 和 `_load_w2` 被正确调用 +- w2_weight 被正确加载（nan=False, inf=False） +- TP 切分正确（shard_dim=1, shard_size=256, tp_rank, narrow 后 shape=[2048, 256]） +- CPU expert 被正确跳过（mask=False → SKIP） +- TP2 服务成功启动，英文/中文/长文本输出全部正常 + +### 关键代码路径 +- GPU 权重创建：`UnquantizedFusedMoEMethod.create_weights()` → `w2_weight = torch.empty(...)`（未初始化） +- GPU 权重加载：sglang 标准 weight_loader 只处理 `experts.N.down_proj.` 格式 key +- CPU 权重加载：kt-kernel 的 `SafeTensorLoader`（`kt-kernel/python/utils/loader.py`）解析 AMXINT4 格式 key +- KTEPWrapperMethod：`create_weights` 调用 `gpu_method.create_weights`，`process_weights_after_loading` 调用 `wrapper.load_weights`（CPU 权重） + +### 教训 +1. **`--model` 必须指向标准 bf16 checkpoint**，sglang 标准 weight_loader 才能识别并加载 GPU 权重 +2. **`--kt-weight-path` 指向 AMXINT4 checkpoint**，供 kt-kernel 的 AMX 内核使用 +3. **MESH 模式下同样需要修复 `--model` 路径**，因为 GPU 专家（GE>0）仍需要从标准 checkpoint 加载 w2_weight +4. **`torch.empty` 不会自动清零**，未初始化的权重在 TP>1 时可能包含 NaN/Inf +5. **本地代码必须与远程同步**：远程 .venv 是 pip install 的旧版，本地子模块是最新 commit，需要 rsync 同步 + +### 排查过程中排除的假设 +- ❌ 不是输入 NaN/Inf 传播 +- ❌ 不是 w13_weight NaN +- ❌ 不是 topk_weights NaN +- ❌ 不是 CPU 专家问题 +- ❌ 不是 `--kt-enable-dynamic-expert-update` 导致 +- ❌ 不是 `_prepare_weight_bf16` 导致（`load()` 未被调用） +- ❌ 不是 `process_weights_after_loading` 导致（bf16 no-op） +- ✅ **是 `--model` 指向 AMXINT4 格式 checkpoint，GPU w2_weight 从未被填充** + +### 涉及文件 +- `third_party/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py` — weight_loader 实现 +- `third_party/sglang/python/sglang/srt/layers/moe/kt_ep_wrapper.py` — KTEPWrapperMethod +- `third_party/sglang/python/sglang/srt/layers/quantization/unquant.py` — UnquantizedFusedMoEMethod +- `third_party/sglang/python/sglang/srt/models/qwen3_5.py` — 模型定义 +- `kt-kernel/python/utils/loader.py` — SafeTensorLoader（CPU 权重加载） +- `kt-kernel/python/utils/amx.py` — AMXMoEWrapper diff --git a/scripts/run_mesh_35b_cap.sh b/scripts/run_mesh_35b_cap.sh new file mode 100644 index 000000000..8f2f3ab7f --- /dev/null +++ b/scripts/run_mesh_35b_cap.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# MESH 35B AMXINT4 benchmark — 修复 w2_weight bug（--model 指向 bf16） +# 用法: run_mesh_35b_cap.sh +# CAP: slot 容量 (64/128/192/224), 224=full (CPU 专家总数 = 256 - 32) +# GE=32, TP=4, 开 cudagraph +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=40 +export KT_MESH_WEIGHT_TYPE=amxint4 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ \ + --kt-weight-path /mnt/data2/models/Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH/ \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_mesh_397b_cap.sh b/scripts/run_mesh_397b_cap.sh new file mode 100644 index 000000000..94dfee3fd --- /dev/null +++ b/scripts/run_mesh_397b_cap.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# MESH 397B AMXINT4 benchmark — 修复 w2_weight bug（--model 指向 bf16） +# 用法: run_mesh_397b_cap.sh +# CAP: slot 容量 (128/192/256/480), 480=full (CPU 专家总数 = 512 - 32) +# GE=32, TP=4, 开 cudagraph +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=60 +export KT_MESH_WEIGHT_TYPE=amxint4 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-weight-path /mnt/data2/models/Qwen3.5-397B-A17B-AMXINT4-NUMA2-MESH-FIXED \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion --language-only diff --git a/scripts/run_mesh_bench.sh b/scripts/run_mesh_bench.sh new file mode 100644 index 000000000..ab3062261 --- /dev/null +++ b/scripts/run_mesh_bench.sh @@ -0,0 +1,123 @@ +#!/bin/bash +# 自动化 MESH benchmark：依次测试多个 cap 值 +# 用法: run_mesh_bench.sh +# MODEL: 35b / 397b +# CAP_LIST: 空格分隔的 cap 值，如 "64 128 192 224" +# PORT: 端口 +# CUDA: CUDA 设备列表，如 0,1,2,3 + +MODEL=$1 +CAP_LIST=$2 +PORT=$3 +CUDA=$4 + +if [ "$MODEL" = "35b" ]; then + SCRIPT="run_mesh_35b_cap.sh" +elif [ "$MODEL" = "397b" ]; then + SCRIPT="run_mesh_397b_cap.sh" +else + echo "Unknown model: $MODEL (use 35b or 397b)" + exit 1 +fi + +RESULT_FILE="/tmp/mesh_bench_${MODEL}_results.txt" +echo "=== MESH ${MODEL} benchmark results ===" > $RESULT_FILE +echo "cap | tokens | total_s | prefill_s | decode_s | decode_tok_s" >> $RESULT_FILE + +for CAP in $CAP_LIST; do + echo "" + echo "========== Testing ${MODEL} cap=${CAP} ==========" + + # 停旧服务 + pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null + sleep 3 + + # 启动新服务 + cd /mnt/data2/tmp/qujing_mesh + nohup bash $SCRIPT $CAP $PORT $CUDA > /tmp/sglang_mesh_${MODEL}_cap${CAP}.log 2>&1 & + SERVER_PID=$! + echo "Started server PID=$SERVER_PID, waiting for ready..." + + # 等待服务就绪（最多 10 分钟，MESH 启动慢） + READY=0 + for i in $(seq 1 120); do + sleep 5 + if grep -q "fired up and ready to roll" /tmp/sglang_mesh_${MODEL}_cap${CAP}.log 2>/dev/null; then + READY=1 + echo "Server ready after $((i*5))s" + break + fi + # 只检测真正致命的错误，忽略 "import error"/"side-effect import failed" 等正常警告 + if grep -qi "Segmentation fault\|CUDA error\|RuntimeError:\|AssertionError:\|Traceback (most recent call last):\|core dumped\|Out of memory" /tmp/sglang_mesh_${MODEL}_cap${CAP}.log 2>/dev/null; then + echo "FATAL error detected in log:" + grep -i "Segmentation fault\|CUDA error\|RuntimeError:\|AssertionError:\|Traceback (most recent call last):\|core dumped\|Out of memory" /tmp/sglang_mesh_${MODEL}_cap${CAP}.log | tail -5 + READY=0 + break + fi + # 检查服务进程是否已退出 + if ! kill -0 $SERVER_PID 2>/dev/null; then + echo "Server process exited unexpectedly" + tail -20 /tmp/sglang_mesh_${MODEL}_cap${CAP}.log + READY=0 + break + fi + done + + if [ $READY -eq 0 ]; then + echo "FAILED: Server not ready for cap=${CAP}" + echo "${CAP} | FAILED" >> $RESULT_FILE + pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null + sleep 3 + continue + fi + + sleep 3 # 额外等待稳定 + + # 功能测试 + FUNC_OUT=$(curl -s http://localhost:${PORT}/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"default","prompt":"Hello, my name is","max_tokens":16,"temperature":0}' | \ + python3 -c "import json,sys; r=json.load(sys.stdin); print(r['choices'][0]['text'][:50])" 2>/dev/null) + echo "Function test output: ${FUNC_OUT}" + + # 测速 + source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate + SPEED_OUT=$(python3 -c " +import requests, time, json +url = 'http://localhost:${PORT}/v1/completions' +data = {'model':'default','prompt':'Write a short essay about artificial intelligence and its impact on society:','max_tokens':256,'temperature':0,'stream':True} +start = time.time() +first_token_time = None +token_count = 0 +r = requests.post(url, json=data, stream=True) +for line in r.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + line = line[6:] + if line == '[DONE]': + break + chunk = json.loads(line) + if chunk['choices'][0]['text']: + token_count += 1 + if first_token_time is None: + first_token_time = time.time() +end = time.time() +total = end - start +prefill = first_token_time - start if first_token_time else 0 +decode = end - first_token_time if first_token_time else total +dtps = token_count/decode if decode > 0 else 0 +print(f'{token_count}|{total:.2f}|{prefill:.2f}|{decode:.2f}|{dtps:.2f}') +" 2>/dev/null) + + echo "Speed: ${SPEED_OUT}" + echo "${CAP} | ${SPEED_OUT}" >> $RESULT_FILE + + # 停服务 + pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null + sleep 3 +done + +echo "" +echo "=== Results ===" +cat $RESULT_FILE From 2ad167f70a37c206b3e4441611728cffbae19e60 Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Tue, 30 Jun 2026 20:38:34 +0800 Subject: [PATCH 09/10] feat: MESH fine-grained cap sweep + worker_pool/mesh operator fixes - cap_sweep.sh: per-rank (tensor_parallel0/1) VRAM, GPU util, VmHWM, NUMA0/1 memory breakdown - bench_decode_stream.py: decode-only speed test (stream mode, excludes prefill) - kt-kernel: worker_pool, mesh_decode/eviction/hook/io_uring/residency/scheduler/slot_pool updates - scripts: add bf16/kt/mesh bench variants for 35b/397b - fix_safetensors_align.py: weight alignment tooling --- .gitignore | 3 +- MESH_BENCHMARK_GUIDE.md | 368 ++++++++ bench_decode_stream.py | 74 ++ cap_sweep.sh | 235 +++++ cap_sweep_raid0.py | 315 +++++++ fix_safetensors_align.py | 124 +++ kt-kernel/cpu_backend/worker_pool.cpp | 217 ++++- kt-kernel/cpu_backend/worker_pool.h | 7 +- kt-kernel/ext_bindings.cpp | 40 +- kt-kernel/operators/amx/bf16-moe.hpp | 104 ++- kt-kernel/operators/amx/moe.hpp | 80 +- kt-kernel/operators/amx/moe_base.hpp | 64 +- kt-kernel/operators/avx2/bf16-moe.hpp | 50 +- kt-kernel/operators/mesh/mesh_decode.hpp | 158 ++-- kt-kernel/operators/mesh/mesh_eviction.hpp | 49 + kt-kernel/operators/mesh/mesh_hook.hpp | 60 ++ kt-kernel/operators/mesh/mesh_io_uring.hpp | 209 +++-- kt-kernel/operators/mesh/mesh_residency.hpp | 962 +++++++++++++++++--- kt-kernel/operators/mesh/mesh_scheduler.hpp | 27 + kt-kernel/operators/mesh/mesh_slot_pool.hpp | 157 +++- kt-kernel/python/utils/amx.py | 68 ++ kt-kernel/python/utils/mesh/stats.py | 18 + kt-kernel/python/utils/mesh/wrapper.py | 151 ++- scripts/bench_bf16_py.py | 193 ++++ scripts/bench_full_py.py | 344 +++++++ scripts/convert_35b_bf16_packed.py | 83 ++ scripts/kt_bench_py.py | 188 ++++ scripts/mesh_bench_py.py | 194 ++++ scripts/run_full_bench.sh | 63 ++ scripts/run_kt_35b.sh | 31 + scripts/run_kt_35b_bf16.sh | 31 + scripts/run_kt_397b.sh | 31 + scripts/run_kt_397b_bf16.sh | 31 + scripts/run_mesh_35b_bf16.sh | 39 + scripts/run_mesh_35b_cap.sh | 1 + scripts/run_mesh_397b_bf16.sh | 39 + scripts/run_mesh_397b_cap.sh | 2 +- scripts/run_mesh_bench.sh | 153 +++- scripts/wait_gpu_and_run.sh | 34 + 39 files changed, 4602 insertions(+), 395 deletions(-) create mode 100644 MESH_BENCHMARK_GUIDE.md create mode 100644 bench_decode_stream.py create mode 100644 cap_sweep.sh create mode 100644 cap_sweep_raid0.py create mode 100644 fix_safetensors_align.py create mode 100644 scripts/bench_bf16_py.py create mode 100644 scripts/bench_full_py.py create mode 100644 scripts/convert_35b_bf16_packed.py create mode 100644 scripts/kt_bench_py.py create mode 100644 scripts/mesh_bench_py.py create mode 100644 scripts/run_full_bench.sh create mode 100644 scripts/run_kt_35b.sh create mode 100644 scripts/run_kt_35b_bf16.sh create mode 100644 scripts/run_kt_397b.sh create mode 100644 scripts/run_kt_397b_bf16.sh create mode 100644 scripts/run_mesh_35b_bf16.sh create mode 100644 scripts/run_mesh_397b_bf16.sh create mode 100644 scripts/wait_gpu_and_run.sh diff --git a/.gitignore b/.gitignore index f44a0ab87..1be31c82b 100644 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,5 @@ build* CMakeFiles/ kvc2/ sched/ -*.png \ No newline at end of file +*.png +.trae/ \ No newline at end of file diff --git a/MESH_BENCHMARK_GUIDE.md b/MESH_BENCHMARK_GUIDE.md new file mode 100644 index 000000000..1be9f11a8 --- /dev/null +++ b/MESH_BENCHMARK_GUIDE.md @@ -0,0 +1,368 @@ +# MESH Benchmark 经验文档 + +> 本文档记录 MESH 模式 benchmark 的完整路径、踩坑、修复和测试规范。 +> 即使对话被压缩，将本文档喂给 AI 即可恢复上下文。 + +## 1. 模型路径汇总 + +| 模型 | 格式 | 路径 | 备注 | +|------|------|------|------| +| 35B bf16 | 标准 safetensors | `/mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/` | `--model` 指向此路径 | +| 35B AMXINT4 | AMXINT4 NUMA2 MESH | `/mnt/data2/models/Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH/` | `--kt-weight-path` 指向此路径 | +| 397B bf16 | 标准 safetensors (94 shards) | `/mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY` | `--model` 指向此路径 | +| 397B AMXINT4 | AMXINT4 NUMA2 MESH FIXED | `/mnt/data2/models/Qwen3.5-397B-A17B-AMXINT4-NUMA2-MESH-FIXED` | `--kt-weight-path` 指向此路径 | + +## 2. 模型参数 + +| 参数 | 35B | 397B | +|------|-----|------| +| 专家总数 | 256 | 512 | +| 层数 | 40 | 60 | +| hidden_size | 2048 | 4096 | +| moe_intermediate | 512 | 1024 | +| CPU 专家 (GE=32) | 224 | 480 | +| `KT_MESH_TOTAL_LAYERS` | 40 | 60 | + +## 3. 启动脚本 + +| 脚本 | 路径（本地） | 路径（远程） | +|------|-------------|-------------| +| 35B MESH | `scripts/run_mesh_35b_cap.sh` | `/mnt/data2/tmp/qujing_mesh/run_mesh_35b_cap.sh` | +| 397B MESH | `scripts/run_mesh_397b_cap.sh` | `/mnt/data2/tmp/qujing_mesh/run_mesh_397b_cap.sh` | +| 35B 标准 kt | `scripts/run_kt_35b.sh` | `/mnt/data2/tmp/qujing_mesh/run_kt_35b.sh` | +| 397B 标准 kt | `scripts/run_kt_397b.sh` | `/mnt/data2/tmp/qujing_mesh/run_kt_397b.sh` | +| MESH bench (Python) | `scripts/mesh_bench_py.py` | `/mnt/data2/tmp/qujing_mesh/mesh_bench_py.py` | +| 标准 kt bench (Python) | `scripts/kt_bench_py.py` | `/mnt/data2/tmp/qujing_mesh/kt_bench_py.py` | +| 自动化 bench (bash) | `scripts/run_mesh_bench.sh` | `/mnt/data2/tmp/qujing_mesh/run_mesh_bench.sh` | + +### 启动参数（共同） + +``` +GE=32 (--kt-num-gpu-experts 32) +TP=4 (--tensor-parallel-size 4) +开 cudagraph (不加 --disable-cuda-graph) +--kt-cpuinfer 153 --kt-threadpool-count 2 +--kt-numa-nodes 0 1 +--kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 +--kt-enable-dynamic-expert-update +--attention-backend flashinfer --trust-remote-code +--mem-fraction-static 0.90 --chunked-prefill-size 4096 +--max-running-requests 16 --max-total-tokens 20000 +--watchdog-timeout 3000 --enable-mixed-chunk +--enable-p2p-check --disable-shared-experts-fusion +``` + +### 环境变量 + +```bash +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP # slot 容量 +export KT_NUM_GPU_EXPERTS=32 # GE +export KT_MESH_TOTAL_LAYERS=40 # 35B=40, 397B=60 +export KT_MESH_WEIGHT_TYPE=amxint4 +``` + +## 4. cap 含义（重要！） + +- **cap = `KT_MESH_CAP`**：单层单 TP 的 slot 数 +- 每个 slot 容纳一个 expert shard 的全部 packed 权重（gate/up/down） +- slot 池预先分配在 NUMA 本地内存，通过 io_uring + O_DIRECT 从 SSD 加载 +- **不是 `--kt-cpuinfer`**（CPU 推理线程数） +- full = CPU 专家总数 = 专家总数 - GE + - 35B full = 256 - 32 = 224 + - 397B full = 512 - 32 = 480 + +## 5. 踩坑记录 + +### 5.1 `--model` 指向 AMXINT4 导致 TP>1 输出坍缩为 `!` + +- **根因**：`--model` 指向 AMXINT4 格式 checkpoint，GPU w2_weight 未初始化，TP>1 时输出 NaN +- **修复**：`--model` 改为 bf16 checkpoint 路径，`--kt-weight-path` 指向 AMXINT4 +- **详见**：`TROUBLESHOOTING_LOG.md` + +### 5.2 bench 脚本错误检测太粗糙 + +- **现象**：所有 cap 都 FAILED，服务被提前杀死 +- **根因**：`grep -qi "error\|traceback\|failed"` 误匹配 "import error"/"side-effect import failed" 正常警告 +- **修复**：改为只检测 `Segmentation fault|CUDA error|RuntimeError:|AssertionError:|Traceback (most recent call last):|core dumped|Out of memory` +- **同时**：增加进程存活检测 `kill -0 $SERVER_PID`，等待时间从 5 分钟增加到 10 分钟 + +### 5.3 397B TEXTONLY 缺 preprocessor_config.json + +- **现象**：`OSError: Can't load image processor for '...TEXTONLY'` +- **根因**：397B TEXTONLY 目录没有 `preprocessor_config.json`，AutoProcessor 加载 image processor 失败 +- **修复**：从 35B 目录复制 `preprocessor_config.json` 到 397B TEXTONLY 目录 + +### 5.4 397B config.json 缺 vision_config + +- **现象**：`AttributeError: 'Qwen3_5MoeVisionConfig' object has no attribute 'hidden_size'` +- **根因**：397B 架构是 `Qwen3_5MoeForConditionalGeneration`（VL 架构），sglang 试图初始化 vision model，但 config.json 没有 vision_config +- **修复**：从 35B config.json 复制 vision_config，调整 `out_hidden_size` 为 397B 的 hidden_size=4096 +- **注意**：必须同时有 preprocessor_config.json 和 vision_config，缺一不可 + +### 5.5 `--language-only` 参数不适用 + +- **现象**：`ValueError: requires at least one encoder urls to be set via --encoder-urls` +- **根因**：`--language-only` 让 sglang 认为是 encoder-decoder 模型，需要 encoder urls +- **修复**：不用 `--language-only`，改用 vision_config + preprocessor_config.json 方案 + +### 5.6 cap=192 首次失败（GPU 内存不平衡） + +- **现象**：`RuntimeError: The memory capacity is unbalanced. min_per_gpu_memory=16.28, local_gpu_memory=46.81` +- **根因**：上一轮 cap=128 的服务没完全退出，残留进程占用部分 GPU 内存 +- **修复**：`pkill -9 -f launch_server` 彻底清理后重跑 +- **预防**：bench 脚本每轮之间应增加更长的等待时间，或用 `nvidia-smi` 确认 GPU 清空 + +### 5.7 MESH init DIAG 显示 `tp=2` 不是 bug + +- **现象**：日志显示 `tp=2`，但实际 `--tensor-parallel-size 4` +- **解释**：MESH 内部的 `tp=2` 指的是 NUMA 分片数（2 个 NUMA 节点），不是 tensor parallel size +- **结论**：这是 MESH 内部表示，不是 bug。实际 TP=4 可从日志的 TP0/TP1/TP2/TP3 确认 + +### 5.8 ssh 后台启动进程的问题 + +- **现象**：`nohup ... &` 方式启动的进程随 ssh 退出被杀掉 +- **修复**：用 `screen -dmS bash -c "..."` 方式启动，确保进程脱离 ssh 会话 +- **替代方案**：`setsid bash -c "..." < /dev/null > /dev/null 2>&1 &` + +### 5.9 GPU prefill 路径 bug（`submit_write_weight_scale_to_buffer`） + +- **现象**：8192 tokens prefill 请求导致服务崩溃，`AttributeError: 'AMXMoEWrapper' object has no attribute 'submit_write_weight_scale_to_buffer'` +- **根因**：`--kt-gpu-prefill-token-threshold 4096`，当 prefill tokens > 4096 时走 GPU prefill 路径，但 MESH 的 `AMXMoEWrapper` 没有 `submit_write_weight_scale_to_buffer` 方法 +- **修复**：prefill prompt 改为 2048 tokens（< 4096 threshold），走 CPU prefill 路径 +- **注意**：如果需要测 8192+ tokens prefill，需要修复 `AMXMoEWrapper` 代码或调大 `--kt-gpu-prefill-token-threshold` + +### 5.10 内存采样进程匹配问题 + +- **现象**：`pgrep -f "python.*sglang"` 只匹配到 1 个进程，RSS 只有 1.62 GiB（实际应 50+ GiB） +- **根因**：sglang TP=4 的子进程可能不是 `python -m sglang.launch_server` 启动的，`pgrep -f "python.*sglang"` 匹配不到所有子进程 +- **修复**：改用 `ps -eo pid,rss,args | grep -i "sglang\|launch_server\|sglang.srt" | grep -v grep` 匹配所有 sglang 相关进程 +- **同时**：加入 `free -m` 采样系统总内存作为补充指标 + +### 5.11 `free -m` awk 解析问题 + +- **现象**：CSV 中 `sys_avail_mb` 列值为 `$7` 字面量 +- **根因**：`awk '/^Mem:/ {print $3",""$7"}'` 中的 `$7` 在双引号中被 shell 解析 +- **影响**：不影响关键指标（sys_used_mb 正常），sys_avail_mb 可忽略 + +### 5.12 标准 kt 模式下 MESH 仍被激活（segfault） + +- **现象**：运行标准 kt 脚本（无 `KT_ENABLE_MESH=1`），日志仍显示 `MESH plugin: bindings registered` 和 `MESH initialized: cap=256, GE=32, layers=60`，随后 `Segmentation fault` in `residency.py bootstrap` +- **根因**：`experts.py` 第 340-351 行，当 `KT_ENABLE_MESH != "1"` 时，代码**回退到读取配置文件** `/tmp/kt_mesh_config.json`。该文件由之前的 MESH benchmark 脚本写入（`MeshConfig.to_file('/tmp/kt_mesh_config.json')`），且 `enabled=true`，导致 MESH 仍被激活 +- **关键**：仅设置 `KT_ENABLE_MESH=0` 或 `unset KT_ENABLE_MESH` **不够**，必须删除 `/tmp/kt_mesh_config.json` +- **修复**：在标准 kt 脚本中加入 `rm -f /tmp/kt_mesh_config.json` 和 `unset KT_ENABLE_MESH KT_MESH_CAP KT_MESH_TOTAL_LAYERS KT_MESH_WEIGHT_TYPE` +- **注意**：`MESH plugin: bindings registered` 日志由 C++ 编译时宏 `#if defined(KT_ENABLE_MESH)` 控制（setup.py 中 `CPUINFER_ENABLE_MESH=ON`），与运行时环境变量无关，这条日志出现是正常的，不影响功能 + +### 5.13 官方 kt 对照验证（origin/main） + +- **背景**：用户质疑 MESH 比标准 kt 快是否科学，要求用官方 ktransformers 验证 +- **方法**：`git archive origin/main:kt-kernel` 生成官方 kt-kernel tar 包，scp 到远程替换 mesh 分支 kt-kernel，`pip install -e . --no-build-isolation` 重新编译（CPUINFER_ENABLE_MESH=OFF），测试后恢复 mesh 分支 +- **发现**：mesh 分支修改了 7 个 kt-kernel 文件（`moe.hpp`, `amx.py`, `experts.py`, `ext_bindings.cpp`, `CMakeLists.txt`, `common.hpp`, `cpuinfer.h`），但标准 kt 模式下所有 MESH 代码被条件判断跳过，走原版路径 +- **defer 参数**：`--kt-max-deferred-experts-per-token` 是 kt 的通用参数（sglang server_args.py 中 `kt_max_deferred_experts_per_token: Optional[int] = None`），默认 None（即 0）。MESH 模式下 `experts.py` 的 `_create_inference_wrapper` 不传此参数给 `MeshMoEWrapper`，MESH 使用自己的 `mesh_config.max_deferred_per_token`（默认 3）走跨层 defer 机制 +- **defer=3 影响**：35B 上 defer=3 明显慢（decode -30%, prefill -19%），397B 上差异小（decode +1%, prefill +2%） +- **结论**：最公平对照是 MESH vs 官方 kt defer=0（两者都不启用标准 kt 的 per-token defer）。MESH 在 397B 上有明显优势（decode +9-11%, prefill +8-16%），在 35B 上优势较小 +- **详见**：第 6 节完整对照表 + +### 5.14 MESH stats 全为 0（`submit_gating_scores_with_cuda_stream` 未集成） + +- **现象**：后台 stats 线程每 10 秒输出 `[MESH_STATS_KV]` 到日志，但所有值都是 0（hit_rate=0, decode_tokens=0, iouring_read_gib=0） +- **根因**：`submit_gating_scores_with_cuda_stream`（定义在 `cpuinfer.h` 第 153 行）从未被任何代码调用。这个函数应该通过 CUDA stream 提交 host callback，在 callback 中调用 `mesh_on_gating_scores_ready`，后者调用 `MeshResidencyManager::on_decode_token_end` 更新 stats。但由于这个集成点缺失，C++ 侧的 stats 永远不会被更新 +- **影响**：MESH 专属指标（hit_rate、iouring_read_gib、eviction_count、decode_tokens）无法采集 +- **修复方案**：需要将 `submit_gating_scores_with_cuda_stream` 集成到 sglang 的 decode 路径中，并添加 Python 绑定（ext_bindings.cpp 中缺少 `.def("submit_gating_scores_with_cuda_stream", ...)`）。这需要修改 C++ 代码并重新编译 +- **当前状态**：未修复，MESH stats 指标暂不可用 + +### 5.15 iostat 解析修复 + +- **现象**：`iostat: peak=0.0 MB/s, avg=0.0 MB/s` +- **根因**：原代码用 `parts[3]`（rrqm/s 列）而非 `parts[2]`（rkB/s 列），且设备名解析逻辑有问题 +- **修复**：改为采样所有非 loop/ram/sr 设备的 `rkB/s`（第 3 列），按轮次汇总后除以 1024 转 MB/s +- **验证**：修复后 35B BF16 标准 kt 测试得到 `peak=2182.5 MB/s, avg=819.1 MB/s` + +### 5.16 内存采样器时序修复 + +- **现象**：部分 cap 值的 peak_sys/peak_gpu 异常偏低（如 397B BF16 cap=480 peak_gpu=59.63） +- **根因**：内存采样器在 server 启动后才开始采样，错过了 bootstrap 阶段的内存峰值 +- **修复**：在 server 启动前开始采样，覆盖 bootstrap + 测试全过程 +- **验证**：修复后 397B BF16 标准 kt peak_sys=749.51 GiB（之前 746.91），peak_gpu=137.61 GiB（之前 93.53） + +## 6. 当前测试结果（完整指标） + +### 35B AMXINT4（GE=32, TP=4, 开 cudagraph, prefill 2048 tokens CPU 路径） + +#### 完整对照表（官方 kt vs mesh 分支标准 kt vs MESH） + +| 模式 | decode tok/s | prefill tok/s | peak_sys_used_gib | peak_gpu_gib | +|------|-------------|--------------|-----------------|-------------| +| **官方 kt (origin/main)** | 96.86 | 3345.65 | 38.99 | 94.38 | +| mesh 分支标准 kt | 95.25 | 2810.98 | 39.28 | 94.38 | +| MESH cap=64 | 98.31 | 3375.52 | 43.32 | 94.38 | +| MESH cap=128 | 97.14 | 2853.54 | 47.17 | 94.38 | +| MESH cap=192 | 99.58 | 2837.15 | 51.03 | 94.38 | +| MESH cap=224 (full) | 96.56 | 2993.74 | 53.02 | 94.38 | + +**关键观察**： +- Peak RSS 随 cap 线性增长：每 +64 cap ≈ +3.8 GiB（= 64 * 40 layers * 2 NUMA * 778 KB slot_size） +- GPU 内存 94.38 GiB 恒定（不随 cap 变化） +- Decode 速度差异小（96-100 tok/s），cap=192 最快 +- Prefill 速度 cap=64 最快（3375 tok/s），可能因 slot 池小、初始化快 +- **官方 kt vs mesh 分支标准 kt**：decode 差异 +1.7%（误差范围内），prefill 官方快 +19%（3346 vs 2811） +- **MESH vs 官方 kt**：MESH cap=64 prefill ≈ 官方 kt（3376 vs 3346），MESH cap≥128 prefill ≈ mesh 分支标准 kt（2854-2994 vs 2811） +- **MESH 内存开销**：+4~14 GiB（slot pool 预分配） + +### 397B AMXINT4（GE=32, TP=4, 开 cudagraph, prefill 2048 tokens CPU 路径） + +#### 完整对照表（官方 kt vs mesh 分支标准 kt vs MESH） + +| 模式 | decode tok/s | prefill tok/s | peak_sys_used_gib | peak_gpu_gib | +|------|-------------|--------------|-----------------|-------------| +| **官方 kt (origin/main)** | 33.05 | 1018.40 | 211.00 | 124.30 | +| mesh 分支标准 kt | 34.74 | 1089.97 | 208.07 | 124.30 | +| MESH cap=128 | 36.45 | 1099.62 | 254.49 | 124.30 | +| MESH cap=192 | 36.02 | 1108.16 | 277.33 | 124.30 | +| MESH cap=256 | 36.68 | 1123.60 | 299.91 | 124.30 | +| MESH cap=480 (full) | 36.50 | 1178.62 | 379.34 | 124.29 | + +**关键观察**： +- Peak sys used 随 cap 线性增长：每 +64 cap ≈ +23 GiB（= 64 * 60 layers * 2 NUMA * 3.02 MB slot_size） +- GPU 内存 124.30 GiB 恒定 +- Decode 速度各 cap 基本一致（36-37 tok/s），cap=256 首次测出 22.66 是偶然（重测 36.68 正常） +- Prefill 速度随 cap 增大而提升（1099→1178 tok/s），符合预期（更多专家驻留减少 I/O） +- **官方 kt vs mesh 分支标准 kt**：decode mesh 分支快 +5%（34.74 vs 33.05），prefill mesh 分支快 +7%（1090 vs 1018） +- **MESH vs 官方 kt**：MESH decode 快 +9-11%（36-37 vs 33），prefill 快 +8-16%（1099-1179 vs 1018） +- **MESH 内存开销**：+46~171 GiB（slot pool 预分配） + +### 对照分析总结 + +**MESH 加速根因**：NUMA 本地性优化 +- 标准 kt：`gate_bb_`/`up_bb_`/`down_bb_` NUMA 盲分配，跨 NUMA 访问延迟翻倍 +- MESH：slot pool 用 `numa_alloc_onnode` 强制分配在 NUMA 本地，forward 时 hook 返回 NUMA 本地 slot 指针 + +**35B vs 397B 差异**： +- 35B：MESH decode 优势小（+1-3 tok/s），prefill cap=64 优势明显但 cap≥128 优势消失 +- 397B：MESH decode 优势明显（+3-4 tok/s，+9-11%），prefill 优势随 cap 增大而扩大（+8-16%） +- 大模型（397B）专家更多，NUMA 本地性优化效果更显著 + +**mesh 分支对标准 kt 的影响**： +- 35B：prefill 慢了 19%（2811 vs 3346），可能因 `moe.hpp` 中 MESH hook 条件判断影响编译器优化 +- 397B：decode 和 prefill 反而略快（+5%, +7%），可能因运行时波动或其他优化 +- 总体：mesh 分支对标准 kt 的影响在误差范围内，不影响 MESH 有效性结论 + +### 35B BF16（GE=32, TP=4, 开 cudagraph, prefill 2601 tokens CPU 路径） + +测试日期：2026-06-22。短 prefill=2601 tokens，长 prefill=8201 tokens（820×"The quick brown fox..."）。 + +#### 完整对照表（标准 kt vs MESH 消融） + +| 模式 | decode tok/s | prefill_short tok/s | prefill_long tok/s | peak_sys_gib | peak_gpu_gib | iostat_peak mbs | +|------|-------------|--------------------|-------------------|-------------|-------------|-----------------| +| 标准 kt | 37.63 | 3307.32 | 1685.84 | 87.41 | 96.42 | 2182.5 | +| MESH cap=64 | 77.29 | 2450.68 | 1408.52 | 42.08 | 96.41 | 74.9 | +| MESH cap=128 | 63.73 | 3110.50 | 1249.72 | 57.12 | 96.41 | 74.9 | +| MESH cap=192 | 60.89 | 3306.35 | 1179.74 | 72.20 | 96.41 | 1480.3 | +| MESH cap=224 (full) | 61.47 | 3011.31 | 1162.60 | 79.61 | 96.41 | 716.8 | + +**消融分析**： +- **Peak sys used 随 cap 线性增长**：每 +64 cap ≈ +15 GiB（= 64 experts × 40 layers × 2 NUMA × 3 MB slot_size）。cap=64→42.08, cap=224→79.61 +- **GPU 内存 96.41 GiB 恒定**（标准 kt 96.42，差异 <0.1%） +- **Decode 速度**：cap=64 最快（77.29 tok/s，比标准 kt 快 +105%），cap 越大 decode 越慢（cap=224 时 61.47，仍比标准 kt 快 +63%） +- **Prefill short 速度**：cap=192 最快（3306 tok/s，与标准 kt 持平），cap=64 最慢（2451 tok/s，比标准 kt 慢 -26%） +- **Prefill long 速度**：标准 kt 最快（1686 tok/s），MESH 各 cap 均慢于标准 kt（1163-1409 tok/s，慢 -16%~-31%）。长 prompt 时 MESH 的 slot 管理开销超过缓存收益 +- **iostat 磁盘带宽**：标准 kt peak=2182 MB/s（权重全量加载），MESH cap=64/128 peak=75 MB/s（slot pool 命中后几乎无磁盘读），cap=192/224 因 evict 重读升至 716-1480 MB/s +- **MESH vs 标准 kt 内存**：MESH cap≤192 的 peak_sys 反而低于标准 kt（42-72 vs 87 GiB），因 slot pool 替代了标准 kt 的全量权重驻留；cap=224 时 79.61 仍低于标准 kt + +### 397B BF16（GE=32, TP=4, 开 cudagraph, prefill 2601 tokens CPU 路径） + +测试日期：2026-06-22。短 prefill=2601 tokens，长 prefill=8201 tokens。cap=128 启动失败（sglang 内存平衡检查拒绝，见 5.6）。 + +#### 完整对照表（标准 kt vs MESH 消融） + +| 模式 | decode tok/s | prefill_short tok/s | prefill_long tok/s | peak_sys_gib | peak_gpu_gib | iostat_peak mbs | +|------|-------------|--------------------|-------------------|-------------|-------------|-----------------| +| 标准 kt | 14.20 | 654.36 | 254.81 | 749.51 | 137.61 | 4250.1 | +| MESH cap=128 | FAILED | - | - | - | - | - | +| MESH cap=192 | 19.59 | 296.72 | 152.78 | 300.70 | 137.57 | 1804.9 | +| MESH cap=256 | 20.56 | 375.22 | 151.66 | 390.86 | 137.58 | 199.8 | +| MESH cap=480 (full) | 18.01 | 750.02 | 137.92 | 706.48 | 137.58 | 971.5 | + +**消融分析**： +- **cap=128 启动失败**：sglang 内存平衡检查要求所有 GPU 可用内存 min > max×0.9，cap=128 时 slot pool 内存分配不均导致拒绝。cap≥192 正常启动 +- **Decode 速度**：MESH 全面碾压标准 kt。cap=256 最快（20.56 tok/s，+45%），cap=192 次之（19.59，+38%），cap=480 最慢（18.01，+27%）。397B BF16 权重巨大（~794GB），MESH 的 NUMA 本地性优化收益显著 +- **Prefill short 速度**：cap=480 最快（750 tok/s，仍慢于标准 kt 654 的 +15%），cap=192 最慢（297 tok/s，-55%）。大 cap 时 slot 命中率高，prefill 更接近标准 kt +- **Prefill long 速度**：标准 kt 最快（255 tok/s），MESH 各 cap 均慢（138-153 tok/s，慢 -40%~-46%）。长 prompt 时 MESH 开销显著 +- **Peak sys 内存**：MESH cap=192/256 大幅低于标准 kt（301/391 vs 750 GiB），节省 48-60% 内存；cap=480 时 706 GiB 接近标准 kt +- **iostat 磁盘带宽**：标准 kt peak=4250 MB/s（全量加载 794GB 权重），MESH cap=256 peak=200 MB/s（slot 命中后几乎无磁盘读），cap=192/480 因 evict 重读升至 972-1805 MB/s +- **最佳 cap 选择**：cap=256 是 decode 最优（20.56 tok/s + 内存仅 391 GiB + 磁盘 200 MB/s），cap=480 是 prefill 最优但内存开销大 + +### BF16 消融总结 + +**MESH 在 BF16 模式下的核心价值**： +- **Decode 全面碾压**：35B +63~105%，397B +27~45%。BF16 权重大，NUMA 本地性优化收益显著 +- **内存大幅节省**：MESH slot pool 替代全量权重驻留，397B cap=256 节省 48% 内存（391 vs 750 GiB） +- **磁盘 I/O 降低**：MESH cap=256 磁盘读取仅 200 MB/s vs 标准 kt 4250 MB/s（降低 95%） +- **长 prefill 是短板**：MESH 各 cap 的长 prefill 均慢于标准 kt（-16%~-46%），slot 管理开销超过缓存收益 +- **小 cap decode 更快**：cap 越小 decode 越快（缓存局部性更好），但 prefill 越慢（命中率低） + +**cap 选择建议**： +- 35B BF16：cap=64 decode 最优（77 tok/s），cap=192 prefill_short 最优（3306 tok/s） +- 397B BF16：cap=256 综合最优（decode 20.56 + 内存 391 GiB + 磁盘 200 MB/s） + +## 7. 已测指标说明 + +根据 SKILL.md 要求，完整 benchmark 需要以下指标（2026-06-22 全部采集）： + +### 7.1 Prefill 速度 ✅ +- **短 prefill**：2601 tokens（"Write a short essay about..."） +- **长 prefill**：8201 tokens（820×"The quick brown fox..."） +- `prefill_tok_s = prompt_tokens / prefill_time` +- **AMXINT4 长 prefill 失败**：因 `submit_write_weight_scale_to_buffer` bug（见 5.9），BF16 正常 + +### 7.2 内存 Peak ✅ +- **进程 RSS**：从 `/proc//status` 的 VmRSS 采样，汇总所有 sglang 子进程 +- **GPU memory**：`nvidia-smi --query-gpu=memory.used` +- **采样频率**：服务启动前每 5 秒采样一次（修复时序问题，见 5.16），直到测试结束 +- **记录 peak 值** + +### 7.3 MESH 专属指标 ❌（未采集） +- **hit rate**：从 expert stats 获取（CACHED 命中率） +- **iouring_read_gib**：io_uring 读取的总数据量（GB） +- **状态**：全为 0，因 `submit_gating_scores_with_cuda_stream` 未集成（见 5.14） +- **修复方案**：需要修改 C++ 代码并重新编译 + +### 7.4 iostat ✅ +- **磁盘读取带宽**：`iostat -x -d 5` 后台采样 +- **解析**：采样所有非 loop/ram/sr 设备的 `rkB/s`（第 3 列），按轮次汇总后除以 1024 转 MB/s +- **记录 peak 和 avg 值** + +## 8. 远程环境 + +- **服务器**：sapphire4 +- **venv**：`/mnt/data2/tmp/qujing_mesh/.venv` +- **工作目录**：`/mnt/data2/tmp/qujing_mesh` +- **GPU**：8x GPU（使用 0,1,2,3） +- **NUMA**：2 节点（node 0, node 1） + +## 9. 测试矩阵 + +| 模型 | cap 值 | full | GE | TP | cudagraph | +|------|--------|------|-----|-----|-----------| +| 35B | 64/128/192/224 | 224 | 32 | 4 | 开启 | +| 397B | 128/192/256/480 | 480 | 32 | 4 | 开启 | + +## 10. 快速恢复命令 + +```bash +# 查看结果 +ssh sapphire4 'cat /tmp/mesh_bench_35b_results.txt; echo "---"; cat /tmp/mesh_bench_397b_results.txt' + +# 查看某个 cap 的服务日志 +ssh sapphire4 'tail -50 /tmp/sglang_mesh_35b_cap64.log' + +# 清理进程 +ssh sapphire4 'screen -X -S mesh397b quit; pkill -9 -f launch_server; true' + +# 启动 35B benchmark +ssh sapphire4 'screen -dmS mesh35b bash -c "cd /mnt/data2/tmp/qujing_mesh && bash run_mesh_bench.sh 35b \"64 128 192 224\" 10004 0,1,2,3 > /tmp/mesh_bench_35b_all.log 2>&1"' + +# 启动 397B benchmark +ssh sapphire4 'screen -dmS mesh397b bash -c "cd /mnt/data2/tmp/qujing_mesh && bash run_mesh_bench.sh 397b \"128 192 256 480\" 10004 0,1,2,3 > /tmp/mesh_bench_397b_all.log 2>&1"' +``` diff --git a/bench_decode_stream.py b/bench_decode_stream.py new file mode 100644 index 000000000..a3a6ef58a --- /dev/null +++ b/bench_decode_stream.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Decode-only benchmark using stream mode. +Usage: python3 bench_decode_stream.py [max_tokens] +Returns: decode_tok_s (excluding prefill/first-token latency) +""" +import requests, time, json, sys + +port = sys.argv[1] if len(sys.argv) > 1 else "50052" +max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 512 +url = f"http://localhost:{port}/v1/completions" + +# Warmup +try: + r = requests.post(url, json={ + "model": "default", "prompt": "Hello", + "max_tokens": 16, "temperature": 0 + }, timeout=60) + if r.status_code != 200: + print(f"WARMUP_FAIL: {r.status_code}") + sys.exit(1) +except Exception as e: + print(f"WARMUP_FAIL: {e}") + sys.exit(1) + +time.sleep(2) + +# Decode benchmark with stream +data = { + "model": "default", + "prompt": "Write a detailed essay about artificial intelligence and its impact on society, including historical context, current developments, and future implications:", + "max_tokens": max_tokens, + "temperature": 0, + "stream": True, +} + +start = time.time() +first_token_time = None +token_count = 0 + +try: + r = requests.post(url, json=data, stream=True, timeout=300) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + try: + chunk = json.loads(line) + text = chunk["choices"][0].get("text", "") + if text: + token_count += 1 + if first_token_time is None: + first_token_time = time.time() + except: + pass +except Exception as e: + print(f"DECODE_FAIL: {e}") + sys.exit(1) + +end = time.time() + +if first_token_time and token_count > 1: + decode_elapsed = end - first_token_time + decode_tps = (token_count - 1) / decode_elapsed # exclude first token (prefill) + total_elapsed = end - start + print(f"decode_tok_s: {decode_tps:.2f}") + print(f"decode_tokens: {token_count}") + print(f"total_elapsed: {total_elapsed:.2f}s") + print(f"decode_elapsed: {decode_elapsed:.2f}s") +else: + print(f"DECODE_FAIL: no tokens generated") + sys.exit(1) diff --git a/cap_sweep.sh b/cap_sweep.sh new file mode 100644 index 000000000..5e830c8e9 --- /dev/null +++ b/cap_sweep.sh @@ -0,0 +1,235 @@ +#!/bin/bash +# MESH cap sweep benchmark — fine-grained per-rank memory + GPU + utilization. +# Tests multiple cap values with decode speed + per-rank VRAM/GPU util + per-NUMA memory. +# NO pkill -f: uses screen quit + ss+kill + pgrep -P for safe cleanup. +# Does NOT touch gpu4 processes (port 30001/30002). +# +# Naming convention (per user requirement): +# - "tensor_parallel0" / "tensor_parallel1" = the two scheduler processes (was TP0/TP1) +# - GPU6 = tensor_parallel0's device, GPU7 = tensor_parallel1's device +# - threadpool0/1 = CPU expert worker pools (not tracked separately here) +# +# Usage: bash cap_sweep.sh "32 64 96 128 160 192" + +CAP_LIST=${1:-"32 64 96 128 160 192"} +PORT=50052 +CUDA=6,7 +RANK0_GPU=6 +RANK1_GPU=7 +SCREEN_NAME="mesh_sweep" +WORKDIR=/mnt/data2/tmp/qujing_mesh +SCRIPT=run_mesh_35b_cap_graph_raid0.sh +RESULT=$WORKDIR/cap_sweep_results.txt +VENV=$WORKDIR/.venv/bin/python + +# Header columns (no "TP" abbreviation — use full "tensor_parallel0/1") +echo "cap | decode_tok_s | decode_tokens | total_elapsed | peak_sys_gib | tensor_parallel0_vram_gib | tensor_parallel1_vram_gib | tensor_parallel0_gpu_util_pct | tensor_parallel1_gpu_util_pct | tensor_parallel0_vmhwm_mb | tensor_parallel0_numa0_mb | tensor_parallel0_numa1_mb | tensor_parallel1_vmhwm_mb | tensor_parallel1_numa0_mb | tensor_parallel1_numa1_mb | cpu_load" > $RESULT + +stop_server() { + echo " Stopping server..." + screen -X -S $SCREEN_NAME quit 2>/dev/null + sleep 3 + local PID=$(ss -tlnp 2>/dev/null | grep ":$PORT " | grep -oP 'pid=\K[0-9]+' | head -1) + if [ -n "$PID" ]; then + local CHILDREN=$(pgrep -P $PID 2>/dev/null) + for CHILD in $CHILDREN; do + local GRANDCHILDREN=$(pgrep -P $CHILD 2>/dev/null) + for GC in $GRANDCHILDREN; do + kill $GC 2>/dev/null + done + kill $CHILD 2>/dev/null + done + kill $PID 2>/dev/null + fi + for i in $(seq 1 15); do + if ! ss -tln 2>/dev/null | grep -q ":$PORT "; then + break + fi + sleep 2 + done + sleep 3 + if ss -tln 2>/dev/null | grep -q ":$PORT "; then + echo " WARNING: port $PORT still in use after cleanup" + else + echo " Port $PORT free" + fi +} + +# Get per-rank process info: VmHWM (MB) + NUMA0 mem (MB) + NUMA1 mem (MB) +# Returns 6 space-separated integers: +# tp0_vmhwm tp0_numa0 tp0_numa1 tp1_vmhwm tp1_numa0 tp1_numa1 +get_rank_info() { + local SERVER_PID=$1 + local TP0_VMHWM=0 TP0_NUMA0=0 TP0_NUMA1=0 + local TP1_VMHWM=0 TP1_NUMA0=0 TP1_NUMA1=0 + local CHILDREN=$(pgrep -P $SERVER_PID 2>/dev/null) + for CHILD in $CHILDREN; do + # Use ps -o args= to get full command line (not truncated like /proc/PID/status Name) + local ARGS=$(ps -p $CHILD -o args= 2>/dev/null) + local HWM=$(grep "^VmHWM:" /proc/$CHILD/status 2>/dev/null | awk '{print $2}') + [ -z "$HWM" ] && HWM=0 + # numastat -p PID: take "Total" row, Node0 and Node1 columns (MB, converted to int) + local NUMA_OUT=$(numastat -p $CHILD 2>/dev/null | awk '/^Total/ {printf "%d %d", $2, $3}') + local N0=$(echo "$NUMA_OUT" | awk '{print $1}') + local N1=$(echo "$NUMA_OUT" | awk '{print $2}') + [ -z "$N0" ] && N0=0 + [ -z "$N1" ] && N1=0 + if echo "$ARGS" | grep -q "scheduler_TP0"; then + TP0_VMHWM=$((HWM / 1024)) + TP0_NUMA0=$N0 + TP0_NUMA1=$N1 + elif echo "$ARGS" | grep -q "scheduler_TP1"; then + TP1_VMHWM=$((HWM / 1024)) + TP1_NUMA0=$N0 + TP1_NUMA1=$N1 + fi + done + echo "$TP0_VMHWM $TP0_NUMA0 $TP0_NUMA1 $TP1_VMHWM $TP1_NUMA0 $TP1_NUMA1" +} + +# Check CPU load from other processes (djh training etc.) +check_cpu_load() { + local DJH_CPU=$(ps -p 1620567 -o %cpu= 2>/dev/null | awk '{printf "%.0f", $1}') + local TOTAL_USED=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d% -f1) + echo "djh_cpu=${DJH_CPU}% total_user=${TOTAL_USED}%" +} + +# Per-second sampler: system RAM + per-rank GPU VRAM + per-rank GPU util +MEM_SAMPLES=/tmp/mesh_mem_samples.txt +GPU0_VRAM_SAMPLES=/tmp/mesh_gpu0_vram.txt # GPU6 = tensor_parallel0 +GPU1_VRAM_SAMPLES=/tmp/mesh_gpu1_vram.txt # GPU7 = tensor_parallel1 +GPU0_UTIL_SAMPLES=/tmp/mesh_gpu0_util.txt +GPU1_UTIL_SAMPLES=/tmp/mesh_gpu1_util.txt + +sample_loop() { + > $MEM_SAMPLES + > $GPU0_VRAM_SAMPLES + > $GPU1_VRAM_SAMPLES + > $GPU0_UTIL_SAMPLES + > $GPU1_UTIL_SAMPLES + while true; do + # System memory (MB used) + free -m | grep "^Mem:" | awk '{print $3}' >> $MEM_SAMPLES + # Query both GPUs at once: columns = memory.used, utilization.gpu + local OUT=$(nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader,nounits -i $RANK0_GPU,$RANK1_GPU 2>/dev/null) + local L0=$(echo "$OUT" | sed -n '1p') + local L1=$(echo "$OUT" | sed -n '2p') + echo "$L0" | awk -F', ' '{print $1}' >> $GPU0_VRAM_SAMPLES + echo "$L0" | awk -F', ' '{print $2}' >> $GPU0_UTIL_SAMPLES + echo "$L1" | awk -F', ' '{print $1}' >> $GPU1_VRAM_SAMPLES + echo "$L1" | awk -F', ' '{print $2}' >> $GPU1_UTIL_SAMPLES + sleep 1 + done +} + +echo "Config: mesh=ON, tensor_parallel_size=2, GE=32, dual-numa(0,1), defer=3, RAID0 O_DIRECT" +echo "Cap list: $CAP_LIST" +echo "GPU6 = tensor_parallel0 device, GPU7 = tensor_parallel1 device" +echo "Results -> $RESULT" +echo "" + +for CAP in $CAP_LIST; do + echo "========== Testing cap=$CAP ==========" + + # 1. Stop old server + stop_server + + # 2. Start new server + LOG=/tmp/mesh_sweep_cap${CAP}.log + cd $WORKDIR + screen -dmS $SCREEN_NAME bash -c "taskset -c 24-47,72-191 bash $SCRIPT $CAP $PORT $CUDA > $LOG 2>&1" + echo " Started cap=$CAP, waiting for ready..." + + # 3. Wait for ready (up to 10 min) + READY=0 + for i in $(seq 1 120); do + sleep 5 + if grep -q "fired up and ready to roll" $LOG 2>/dev/null; then + READY=1 + echo " Server ready after $((i*5))s" + break + fi + if grep -q "scheduler is dead" $LOG 2>/dev/null; then + echo " FATAL: scheduler dead" + break + fi + if grep -q "Segmentation fault" $LOG 2>/dev/null; then + echo " FATAL: segfault" + break + fi + done + + if [ $READY -eq 0 ]; then + echo " FAILED: cap=$CAP not ready" + echo "$CAP | FAILED | - | - | - | - | - | - | - | - | - | - | - | - | - | $(check_cpu_load)" >> $RESULT + tail -20 $LOG 2>/dev/null + continue + fi + + sleep 5 + + # 4. Get server PID (the listening process) + SERVER_PID=$(ss -tlnp 2>/dev/null | grep ":$PORT " | grep -oP 'pid=\K[0-9]+' | head -1) + echo " Server PID: $SERVER_PID" + + # 5. Start sampler (per-second) + sample_loop & + SAMPLER_PID=$! + echo " Sampler PID: $SAMPLER_PID" + + # 6. Run decode benchmark + echo " Running decode benchmark..." + BENCH_OUT=$($VENV $WORKDIR/bench_decode_stream.py $PORT 512 2>&1) + echo " $BENCH_OUT" + + DECODE_TPS=$(echo "$BENCH_OUT" | grep "decode_tok_s:" | awk '{print $2}') + DECODE_TOKENS=$(echo "$BENCH_OUT" | grep "decode_tokens:" | awk '{print $2}') + TOTAL_ELAPSED=$(echo "$BENCH_OUT" | grep "total_elapsed:" | awk '{print $2}') + + # 7. Stop sampler + kill $SAMPLER_PID 2>/dev/null + wait $SAMPLER_PID 2>/dev/null + + # 8. Calculate peak memory + GPU metrics + PEAK_SYS_MB=$(sort -n $MEM_SAMPLES | tail -1) + PEAK_SYS_GIB=$(echo "scale=2; ${PEAK_SYS_MB:-0} / 1024" | bc) + PEAK_GPU0_VRAM_MB=$(sort -n $GPU0_VRAM_SAMPLES | tail -1) + PEAK_GPU1_VRAM_MB=$(sort -n $GPU1_VRAM_SAMPLES | tail -1) + PEAK_GPU0_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU0_VRAM_MB:-0} / 1024" | bc) + PEAK_GPU1_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU1_VRAM_MB:-0} / 1024" | bc) + # GPU utilization: take max during decode (peak utilization) + PEAK_GPU0_UTIL=$(sort -n $GPU0_UTIL_SAMPLES | tail -1) + PEAK_GPU1_UTIL=$(sort -n $GPU1_UTIL_SAMPLES | tail -1) + [ -z "$PEAK_GPU0_UTIL" ] && PEAK_GPU0_UTIL=0 + [ -z "$PEAK_GPU1_UTIL" ] && PEAK_GPU1_UTIL=0 + + # 9. Get per-rank VmHWM + NUMA breakdown + RANK_INFO=$(get_rank_info $SERVER_PID) + TP0_VMHWM=$(echo "$RANK_INFO" | awk '{print $1}') + TP0_NUMA0=$(echo "$RANK_INFO" | awk '{print $2}') + TP0_NUMA1=$(echo "$RANK_INFO" | awk '{print $3}') + TP1_VMHWM=$(echo "$RANK_INFO" | awk '{print $4}') + TP1_NUMA0=$(echo "$RANK_INFO" | awk '{print $5}') + TP1_NUMA1=$(echo "$RANK_INFO" | awk '{print $6}') + + # 10. CPU load + CPU_LOAD=$(check_cpu_load) + + echo " Peak sys: ${PEAK_SYS_GIB}GiB" + echo " tensor_parallel0 (GPU${RANK0_GPU}): VRAM=${PEAK_GPU0_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU0_UTIL}%" + echo " tensor_parallel1 (GPU${RANK1_GPU}): VRAM=${PEAK_GPU1_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU1_UTIL}%" + echo " tensor_parallel0 proc: VmHWM=${TP0_VMHWM}MB, NUMA0=${TP0_NUMA0}MB, NUMA1=${TP0_NUMA1}MB" + echo " tensor_parallel1 proc: VmHWM=${TP1_VMHWM}MB, NUMA0=${TP1_NUMA0}MB, NUMA1=${TP1_NUMA1}MB" + echo " CPU load: $CPU_LOAD" + + echo "$CAP | $DECODE_TPS | $DECODE_TOKENS | $TOTAL_ELAPSED | $PEAK_SYS_GIB | $PEAK_GPU0_VRAM_GIB | $PEAK_GPU1_VRAM_GIB | $PEAK_GPU0_UTIL | $PEAK_GPU1_UTIL | $TP0_VMHWM | $TP0_NUMA0 | $TP0_NUMA1 | $TP1_VMHWM | $TP1_NUMA0 | $TP1_NUMA1 | $CPU_LOAD" >> $RESULT + + echo "" +done + +# Stop final server +stop_server + +echo "" +echo "=== Results ===" +cat $RESULT diff --git a/cap_sweep_raid0.py b/cap_sweep_raid0.py new file mode 100644 index 000000000..8330cc849 --- /dev/null +++ b/cap_sweep_raid0.py @@ -0,0 +1,315 @@ +#!/usr/bin/env python3 +"""MESH cap sweep benchmark — RAID0 O_DIRECT version. + +Tests multiple cap values, measures decode speed, prefill speed, and memory peak. +NO pkill -f: uses screen quit + ss+kill for safe process management. +Does NOT touch gpu4 processes (port 30001/30002). + +Usage: python3 cap_sweep_raid0.py "32 64 96 128 160 196 224" +""" +import subprocess, time, json, threading, requests, sys, os, re + +CAP_LIST = [int(x) for x in sys.argv[1].split()] if len(sys.argv) > 1 else [32, 64, 96, 128, 160, 196, 224] +PORT = 50052 +CUDA = "6,7" +SCREEN_NAME = "mesh_sweep" +RESULT_FILE = "/mnt/data2/tmp/qujing_mesh/cap_sweep_results.txt" +SCRIPT = "run_mesh_35b_cap_graph_raid0.sh" +WORKDIR = "/mnt/data2/tmp/qujing_mesh" +GPU_IDS = [6, 7] # physical GPU indices for nvidia-smi query + +URL = f"http://localhost:{PORT}/v1/completions" + + +def run(cmd, timeout=30): + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) + + +def stop_server(): + """Stop server safely without pkill -f.""" + # 1. Quit screen session + run(f"screen -X -S {SCREEN_NAME} quit") + time.sleep(2) + # 2. Find PID listening on PORT and kill it (only our port, not gpu4's 30001/30002) + out = run(f"ss -tlnp 2>/dev/null | grep ':{PORT} '") + pids = re.findall(r'pid=(\d+)', out.stdout) + for pid in pids: + run(f"kill {pid} 2>/dev/null") + # 3. Wait for port release + for _ in range(10): + out = run(f"ss -tln 2>/dev/null | grep ':{PORT} '") + if not out.stdout.strip(): + break + time.sleep(2) + time.sleep(3) + + +def start_server(cap): + """Start server with given cap via screen + taskset.""" + log_file = f"/tmp/mesh_sweep_cap{cap}.log" + cmd = ( + f"cd {WORKDIR} && screen -dmS {SCREEN_NAME} bash -c " + f"'taskset -c 24-47,72-191 bash {SCRIPT} {cap} {PORT} {CUDA} > {log_file} 2>&1'" + ) + run(cmd) + return log_file + + +def wait_ready(log_file, max_wait=900): + """Wait for server ready, check fatal errors.""" + for i in range(max_wait // 5): + time.sleep(5) + try: + with open(log_file) as f: + content = f.read() + if "fired up and ready to roll" in content: + return True + for fatal in ["Segmentation fault", "CUDA error", "RuntimeError:", + "Traceback (most recent call last):", "core dumped", + "Out of memory", "timeout", "Address already in use"]: + if fatal in content: + print(f"FATAL detected: {fatal}") + return False + except: + pass + return False + + +def get_all_pids(): + """Get all PIDs related to our server (parent + children + grandchildren).""" + # Find parent PID by port + out = run(f"ss -tlnp 2>/dev/null | grep ':{PORT} '") + pids = re.findall(r'pid=(\d+)', out.stdout) + if not pids: + return [] + all_pids = list(pids) + # Recursively find children + to_check = list(pids) + while to_check: + pid = to_check.pop(0) + out = run(f"pgrep -P {pid}") + for child in out.stdout.strip().split('\n'): + if child and child not in all_pids: + all_pids.append(child) + to_check.append(child) + return all_pids + + +def get_vmhwm(pids): + """Get VmHWM (peak resident set) for processes.""" + results = [] + for pid in pids: + try: + name = "" + hwm = 0 + with open(f"/proc/{pid}/status") as f: + for line in f: + if line.startswith("Name:"): + name = line.split()[1] if len(line.split()) > 1 else "?" + elif line.startswith("VmHWM:"): + hwm = int(line.split()[1]) # KB + if hwm > 0: + results.append((pid, name, hwm)) + except: + pass + return results + + +def get_numastat(pids): + """Get NUMA memory distribution for processes.""" + pid_str = ",".join(pids) + out = run(f"numastat -p {pid_str} 2>/dev/null") + return out.stdout + + +def sample_memory(stop_event, samples): + """Background thread: sample system + GPU memory every 2 seconds.""" + gpu_query = ",".join(str(i) for i in GPU_IDS) + while not stop_event.is_set(): + # System memory (free -m) + try: + out = subprocess.check_output(["free", "-m"], text=True) + for line in out.splitlines(): + if line.startswith("Mem:"): + used = int(line.split()[2]) + samples.append(("mem", used)) + break + except: + pass + # GPU memory (only our GPUs) + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits", + "-i", gpu_query], + text=True + ) + gpu_sum = sum(int(x) for x in out.strip().split("\n") if x.strip()) + samples.append(("gpu", gpu_sum)) + except: + pass + time.sleep(2) + + +def bench_decode(): + """Benchmark decode with 512 tokens (stream mode for accurate timing).""" + data = { + "model": "default", + "prompt": "Write a detailed essay about artificial intelligence and its impact on society, including historical context, current developments, and future implications:", + "max_tokens": 512, + "temperature": 0, + "stream": True, + } + start = time.time() + first = None + cnt = 0 + r = requests.post(URL, json=data, stream=True, timeout=300) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + try: + chunk = json.loads(line) + if chunk["choices"][0]["text"]: + cnt += 1 + if first is None: + first = time.time() + except: + pass + end = time.time() + dec = end - first if first else end - start + dtps = cnt / dec if dec > 0 else 0 + return dtps, cnt + + +def bench_prefill(): + """Benchmark prefill with ~2000-token prompt.""" + prompt = "The quick brown fox jumps over the lazy dog. " * 260 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=300) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def warmup(): + """Warmup with short generation.""" + try: + warmup = {"model": "default", "prompt": "Hello, how are you?", "max_tokens": 16, "temperature": 0} + r = requests.post(URL, json=warmup, timeout=60) + return r.status_code == 200 + except: + return False + + +# ===== Main ===== +with open(RESULT_FILE, "w") as f: + f.write("cap | decode_tok_s | decode_tokens | prefill_tok_s | prefill_tokens | " + "peak_sys_used_gib | peak_gpu_gib | scheduler_vmhwm\n") + +print(f"Cap sweep: {CAP_LIST}") +print(f"Results -> {RESULT_FILE}") +print(f"Config: mesh=ON, TP=2, GE=32, dual-numa(0,1), defer=3, RAID0 O_DIRECT") + +for cap in CAP_LIST: + cap_label = "full(224)" if cap == 224 else str(cap) + print(f"\n========== Testing cap={cap_label} ==========") + + # 1. Stop old server + print("Stopping old server...") + stop_server() + + # 2. Start new server + log_file = start_server(cap) + print(f"Started cap={cap}, log={log_file}, waiting for ready...") + + if not wait_ready(log_file): + print(f"FAILED: Server not ready for cap={cap}") + with open(RESULT_FILE, "a") as f: + f.write(f"{cap} | FAILED | - | FAILED | - | FAILED | FAILED | FAILED\n") + # Print last 30 lines of log for debugging + run(f"tail -30 {log_file}") + continue + + print("Server ready") + time.sleep(5) + + # 3. Get PIDs + all_pids = get_all_pids() + print(f"Server PIDs: {all_pids}") + + # 4. Start memory sampling + samples = [] + stop_event = threading.Event() + mem_thread = threading.Thread(target=sample_memory, args=(stop_event, samples)) + mem_thread.start() + print("Memory sampler started") + + # 5. Warmup + print("Warmup...") + if not warmup(): + print("Warmup failed!") + stop_event.set() + mem_thread.join() + with open(RESULT_FILE, "a") as f: + f.write(f"{cap} | WARMUP_FAIL | - | - | - | - | - | -\n") + continue + time.sleep(2) + + # 6. Prefill benchmark + try: + pt, ptime, ptps = bench_prefill() + print(f"Prefill: {pt} tokens, {ptime:.2f}s, {ptps:.2f} tok/s") + except Exception as e: + print(f"Prefill failed: {e}") + pt, ptime, ptps = 0, 0, 0 + + time.sleep(3) + + # 7. Decode benchmark + try: + dtps, dcnt = bench_decode() + print(f"Decode: {dtps:.2f} tok/s ({dcnt} tokens)") + except Exception as e: + print(f"Decode failed: {e}") + dtps, dcnt = 0, 0 + + # 8. Stop memory sampling + stop_event.set() + mem_thread.join() + + # 9. Get VmHWM (cumulative peak) + vmhwm_info = get_vmhwm(all_pids) + scheduler_hwm = "; ".join(f"{name}:{hwm//1024}MB" for pid, name, hwm in vmhwm_info if "scheduler" in name) + + # Calculate peaks + mem_vals = [v for k, v in samples if k == "mem"] + gpu_vals = [v for k, v in samples if k == "gpu"] + peak_mem = max(mem_vals) / 1024 if mem_vals else 0 # GiB + peak_gpu = max(gpu_vals) / 1024 if gpu_vals else 0 # GiB + + print(f"Peak sys used: {peak_mem:.2f} GiB, Peak GPU(6,7): {peak_gpu:.2f} GiB") + print(f"Scheduler VmHWM: {scheduler_hwm}") + + # Write result + with open(RESULT_FILE, "a") as f: + f.write(f"{cap} | {dtps:.2f} | {dcnt} | {ptps:.2f} | {pt} | " + f"{peak_mem:.2f} | {peak_gpu:.2f} | {scheduler_hwm}\n") + + # Print numastat for this cap + numastat_out = get_numastat(all_pids) + print(f"NUMA distribution:\n{numastat_out}") + +# Stop final server +print("\nStopping final server...") +stop_server() + +print("\n=== Results ===") +with open(RESULT_FILE) as f: + print(f.read()) diff --git a/fix_safetensors_align.py b/fix_safetensors_align.py new file mode 100644 index 000000000..6517c3bb5 --- /dev/null +++ b/fix_safetensors_align.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +修复 safetensors 文件使 data_start 512 对齐，满足 O_DIRECT 要求。 + +safetensors 格式：8字节header_len + JSON header + 数据部分 +data_start = 8 + header_len +O_DIRECT 要求 offset 512 对齐，所以 data_start % 512 == 0 + +修复方法：在 JSON header 末尾填充空格，使 header_len 增加 padding 字节 +JSON 规范允许顶层对象后尾随空白，safetensors 库的 json.loads 能正确处理 + +用法：python3 fix_safetensors_align.py [--dry-run] +""" +import struct +import os +import sys +import json +import tempfile +import shutil + +def fix_one_file(fpath: str, dry_run: bool = False) -> bool: + """修复单个 safetensors 文件。返回 True if 已修复或已对齐。""" + file_size = os.path.getsize(fpath) + with open(fpath, 'rb') as fp: + header_len = struct.unpack('{new_header_len}, " + f"data_start {data_start}->{new_data_start}, padding={padding_needed}") + if dry_run: + return True + # 创建临时文件（同目录，确保同文件系统，方便原子替换） + tmp_dir = os.path.dirname(fpath) + tmp_fd, tmp_path = tempfile.mkstemp(suffix='.tmp_align', dir=tmp_dir) + try: + with os.fdopen(tmp_fd, 'wb') as tmp: + # 写入新的 header_len + tmp.write(struct.pack(' 0: + chunk = min(remaining, 64 * 1024 * 1024) # 64MB chunks + n = sendfile_chunk(tmp.fileno(), src.fileno(), chunk) + if n == 0: + raise IOError("sendfile returned 0 bytes") + remaining -= n + # 原子替换 + os.replace(tmp_path, fpath) + print(f"[OK] {os.path.basename(fpath)}: fixed and replaced") + return True + except Exception as e: + print(f"[ERROR] {os.path.basename(fpath)}: {e}") + try: + os.unlink(tmp_path) + except: + pass + return False + +def sendfile_chunk(dst_fd: int, src_fd: int, count: int) -> int: + """使用 os.sendfile 拷贝数据，回退到 read/write。""" + try: + import sendfile # not standard + return sendfile.sendfile(dst_fd, src_fd, 0, count) + except ImportError: + pass + # 使用 os.splice 或 os.copy_file_range（Linux 专用） + try: + # os.copy_file_range 在 Python 3.5+ 可用 + n = os.copy_file_range(src_fd, dst_fd, count) + return n + except (AttributeError, OSError): + pass + # 回退到 read/write + data = os.read(src_fd, count) + os.write(dst_fd, data) + return len(data) + +def main(): + if len(sys.argv) < 2: + print(f"Usage: {sys.argv[0]} [--dry-run]") + sys.exit(1) + dir_path = sys.argv[1] + dry_run = '--dry-run' in sys.argv + if not os.path.isdir(dir_path): + print(f"Error: {dir_path} is not a directory") + sys.exit(1) + files = sorted([f for f in os.listdir(dir_path) if f.endswith('.safetensors')]) + if not files: + print(f"Error: no .safetensors files in {dir_path}") + sys.exit(1) + print(f"Found {len(files)} safetensors files in {dir_path}") + if dry_run: + print("[DRY RUN] no files will be modified") + print() + success = 0 + for f in files: + fpath = os.path.join(dir_path, f) + if fix_one_file(fpath, dry_run): + success += 1 + print(f"\nDone: {success}/{len(files)} files aligned") + +if __name__ == '__main__': + main() diff --git a/kt-kernel/cpu_backend/worker_pool.cpp b/kt-kernel/cpu_backend/worker_pool.cpp index 05d70660b..0654bf32f 100644 --- a/kt-kernel/cpu_backend/worker_pool.cpp +++ b/kt-kernel/cpu_backend/worker_pool.cpp @@ -18,12 +18,110 @@ #include #include #include +#include +#include +#include +#include #include +#include #include "hwloc.h" thread_local int WorkerPool::thread_local_id = -1; +// 自动检测被占用的 CPU: 采样 /proc/stat 两次，计算每个 CPU 的 idle 率。 +// idle < threshold 的 CPU 被认为是"被占用"，绑定时跳过。 +// 这样不需要手动设 KT_SKIP_CPUS，自动避开其他用户的 worker。 +static std::set detect_busy_cpus(double idle_threshold = 0.3, int sample_us = 200000) { + std::set busy; + auto read_idle = [](std::vector>* out) { + FILE* f = fopen("/proc/stat", "r"); + if (!f) return; + char line[256]; + while (fgets(line, sizeof(line), f)) { + if (strncmp(line, "cpu", 3) != 0 || !isdigit((unsigned char)line[3])) continue; + int cpu_id; + long user, nice, system, idle, iowait, irq, softirq, steal; + int n = sscanf(line, "cpu%d %ld %ld %ld %ld %ld %ld %ld %ld", &cpu_id, &user, &nice, + &system, &idle, &iowait, &irq, &softirq, &steal); + if (n >= 5) { + long total = user + nice + system + idle + iowait + irq + softirq + steal; + if (cpu_id >= (int)out->size()) out->resize(cpu_id + 1, {0, 0}); + (*out)[cpu_id] = {idle, total}; + } + } + fclose(f); + }; + std::vector> first, second; + read_idle(&first); + usleep(sample_us); + read_idle(&second); + for (size_t i = 0; i < first.size() && i < second.size(); i++) { + long d_idle = second[i].first - first[i].first; + long d_total = second[i].second - first[i].second; + if (d_total > 0) { + double idle_rate = (double)d_idle / d_total; + if (idle_rate < idle_threshold) { + busy.insert((int)i); + } + } + } + return busy; +} + +// 返回需要跳过的 CPU 集合。 +// 优先用环境变量 KT_SKIP_CPUS (逗号分隔的物理 CPU ID)； +// 若未设置，则自动检测 (采样 /proc/stat，idle<30% 的 CPU)。 +static const std::set& get_skip_cpus() { + static std::set skip; + static bool inited = false; + if (!inited) { + inited = true; + const char* env = getenv("KT_SKIP_CPUS"); + if (env && *env) { + std::stringstream ss(env); + std::string item; + while (std::getline(ss, item, ',')) { + if (!item.empty()) { + try { + skip.insert(std::stoi(item)); + } catch (...) { + } + } + } + fprintf(stderr, "[MESH] KT_SKIP_CPUS manual (%zu cpus): ", skip.size()); + for (int c : skip) fprintf(stderr, "%d ", c); + fprintf(stderr, "\n"); + } else { + skip = detect_busy_cpus(0.30, 200000); + fprintf(stderr, "[MESH] auto-detected %zu busy cpus (idle<30%%): ", skip.size()); + for (int c : skip) fprintf(stderr, "%d ", c); + fprintf(stderr, "\n"); + } + } + return skip; +} + +// 全局 core 分配器: 每个 NUMA 节点维护一个 next_core_idx， +// NumaJobDistributor (main worker) 和 InNumaPool (worker) 共享， +// 避免多个线程绑到同一个 core。 +static std::mutex g_core_alloc_mtx; +static std::map g_next_core_idx; // numa_id -> next available core idx + +static int alloc_core_idx(int numa_id, int hint_idx) { + std::lock_guard lock(g_core_alloc_mtx); + auto it = g_next_core_idx.find(numa_id); + int idx = (it == g_next_core_idx.end()) ? hint_idx : std::max(it->second, hint_idx); + return idx; +} + +static void advance_core_idx(int numa_id, int used_idx) { + std::lock_guard lock(g_core_alloc_mtx); + if (g_next_core_idx[numa_id] <= used_idx) { + g_next_core_idx[numa_id] = used_idx + 1; + } +} + InNumaPool::InNumaPool(int max_thread_num) { printf("In Numa Worker Pool at NUMA %d, %d threads\n", numa_node_of_cpu(sched_getcpu()), max_thread_num); total_worker_count = max_thread_num; @@ -77,18 +175,39 @@ InNumaPool::InNumaPool(int max_thread_num, int numa_id, int threads_id_start) { // throw std::runtime_error("NUMA node not found"); continue; } - core_obj = hwloc_get_obj_inside_cpuset_by_type(topology, numa_obj->cpuset, HWLOC_OBJ_CORE, i + threads_id_start); + // 跳过被其他用户占用的 core: 用全局 alloc_core_idx 获取起始位置， + // 确保不会和 NumaJobDistributor main worker 或其他 InNumaPool worker 重复。 + const auto& skip_cpus = get_skip_cpus(); + int core_search_idx = alloc_core_idx(numa_id, i + threads_id_start); + int bound_cpu_id = -1; + while (true) { + core_obj = hwloc_get_obj_inside_cpuset_by_type(topology, numa_obj->cpuset, HWLOC_OBJ_CORE, core_search_idx); + if (!core_obj) { + break; + } + int cpu_id = hwloc_bitmap_first(core_obj->cpuset); + if (cpu_id >= 0 && skip_cpus.count(cpu_id)) { + fprintf(stderr, "[MESH] Skip occupied CPU %d (core idx %d) for worker %d\n", cpu_id, core_search_idx, i); + core_search_idx++; + continue; + } + bound_cpu_id = cpu_id; + break; + } if (!core_obj) { - fprintf(stderr, "Core %d inside NUMA node %d not found\n", i, numa_id); + fprintf(stderr, "Core %d inside NUMA node %d not found (searched from %d)\n", i, numa_id, i + threads_id_start); // throw std::runtime_error("Core not found inside NUMA node"); continue; } + advance_core_idx(numa_id, core_search_idx); cpuset = hwloc_bitmap_alloc(); hwloc_bitmap_copy(cpuset, core_obj->cpuset); hwloc_bitmap_singlify(cpuset); auto res = hwloc_set_thread_cpubind(topology, native_handle, cpuset, HWLOC_CPUBIND_STRICT); if (res != 0) { fprintf(stderr, "Failed to set thread CPU binding: %s\n", strerror(errno)); + } else { + fprintf(stderr, "[MESH] numa_%d_t_%d -> CPU %d (core idx %d)\n", numa_id, i + threads_id_start, bound_cpu_id, core_search_idx); } } } @@ -116,23 +235,33 @@ int InNumaPool::get_thread_num() { void InNumaPool::set_restricted_worker_count(int count) { restricted_worker_count = count; } void InNumaPool::wait() { + auto wait_start = std::chrono::high_resolution_clock::now(); for (int i = 0; i < worker_count; i++) { while (thread_state_[i].status.load(std::memory_order_acquire) == ThreadStatus::WORKING) { } } + auto wait_end = std::chrono::high_resolution_clock::now(); + long wait_us = std::chrono::duration_cast(wait_end - wait_start).count(); #ifdef PROFILE_BALANCE - size_t max_time = 0; - size_t min_time = thread_state_[0].finish_ns; - size_t sum = 0; - for (int i = 0; i < worker_count; i++) { - sum += thread_state_[i].finish_ns; - max_time = std::max(max_time, thread_state_[i].finish_ns); - min_time = std::min(min_time, thread_state_[i].finish_ns); + if (wait_us > 1000) { // 超过 1ms，打印每个 worker 的唤醒延迟与执行时间分解 + fprintf(stderr, "[WAIT_SLOW] wait=%ldus wc=%d", wait_us, worker_count); + for (int j = 0; j < worker_count; j++) { + // wake_us: worker 检测到 WORKING 的唤醒延迟 (detect - work_set) + // exec_us: worker process_tasks 执行耗时 (finish_ns) + long wake_us = 0, exec_us = 0; + if (thread_state_[j].work_set_ns > 0 && thread_state_[j].detect_ns >= thread_state_[j].work_set_ns) { + wake_us = (long)((thread_state_[j].detect_ns - thread_state_[j].work_set_ns) / 1000); + } else if (thread_state_[j].work_set_ns == 0) { + wake_us = -1; // work_set_ns 未被设置（主线程 thread 0 路径） + } else { + wake_us = -2; // detect < work_set，时钟异常 + } + exec_us = (long)(thread_state_[j].finish_ns / 1000); + fprintf(stderr, " w[%d]{wake=%ld exec=%ld}us", j, wake_us, exec_us); + } + fprintf(stderr, "\n"); } - double balance = 1.0 * sum / (max_time * worker_count); - printf("max_time: %ld, min_time: %ld, sum_time: %ld, balance: %f\n", max_time, min_time, sum, balance); - #endif } @@ -159,6 +288,10 @@ void InNumaPool::do_work_stealing_job_async(int task_num, std::function lock(thread_state_[i].mutex); thread_state_[i].status.store(ThreadStatus::WORKING, std::memory_order_release); +#ifdef PROFILE_BALANCE + thread_state_[i].work_set_ns = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now().time_since_epoch()).count(); +#endif } thread_state_[i].cv.notify_one(); } @@ -169,10 +302,20 @@ void InNumaPool::do_work_stealing_job_async(int task_num, std::function( + start.time_since_epoch()).count(); + size_t init_ns = 0, compute_ns = 0, finalize_ns = 0; #endif auto& s = thread_state_[thread_id]; if (init_func_ != nullptr) { +#ifdef PROFILE_BALANCE + auto init_start = std::chrono::high_resolution_clock::now(); +#endif init_func_(thread_id); +#ifdef PROFILE_BALANCE + init_ns = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - init_start).count(); +#endif } // omp-guided-style work scheduling @@ -194,18 +337,35 @@ void InNumaPool::process_tasks(int thread_id) { if (task_id + i >= end_) { break; } +#ifdef PROFILE_BALANCE + auto comp_start = std::chrono::high_resolution_clock::now(); +#endif compute_func_(task_id + i); +#ifdef PROFILE_BALANCE + compute_ns += std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - comp_start).count(); +#endif } } if (finalize_func_ != nullptr) { +#ifdef PROFILE_BALANCE + auto fin_start = std::chrono::high_resolution_clock::now(); +#endif finalize_func_(thread_id); +#ifdef PROFILE_BALANCE + finalize_ns = std::chrono::duration_cast( + std::chrono::high_resolution_clock::now() - fin_start).count(); +#endif } s.status.store(ThreadStatus::WAITING, std::memory_order_release); #ifdef PROFILE_BALANCE s.finish_ns = std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start).count(); + s.init_ns = init_ns; + s.compute_ns = compute_ns; + s.finalize_ns = finalize_ns; #endif } @@ -223,11 +383,14 @@ void InNumaPool::worker_thread(int thread_id, int numa_id) { } else if (status == ThreadStatus::WAITING) { auto now = std::chrono::high_resolution_clock::now(); auto duration = std::chrono::duration_cast(now - start).count(); - if (duration > 50) { + // 短暂 busy-wait 后进入 cv.wait：CPU 空闲时 cv.notify_one 唤醒延迟 <1ms， + // 避免永久 busy-wait 导致几十个核空转。 + if (duration > 5) { std::unique_lock lock(thread_state_[thread_id].mutex); thread_state_[thread_id].cv.wait(lock, [&] { return thread_state_[thread_id].status.load(std::memory_order_acquire) != ThreadStatus::WAITING; }); + start = std::chrono::high_resolution_clock::now(); } } else if (status == ThreadStatus::EXIT) { return; @@ -299,12 +462,31 @@ void NumaJobDistributor::init(std::vector numa_ids, std::vector thread // throw std::runtime_error("NUMA node not found"); continue; } - core_obj = hwloc_get_obj_inside_cpuset_by_type(topology, numa_obj->cpuset, HWLOC_OBJ_CORE, start_id); + // 跳过被占用的 core: NumaJobDistributor 主线程也要避开 djh 满载的 CPU。 + // 用全局 alloc_core_idx 确保和 InNumaPool worker 不冲突。 + const auto& skip_cpus = get_skip_cpus(); + int main_search_idx = alloc_core_idx(this_numa, start_id); + int main_bound_cpu = -1; + while (true) { + core_obj = hwloc_get_obj_inside_cpuset_by_type(topology, numa_obj->cpuset, HWLOC_OBJ_CORE, main_search_idx); + if (!core_obj) { + break; + } + int cpu_id = hwloc_bitmap_first(core_obj->cpuset); + if (cpu_id >= 0 && skip_cpus.count(cpu_id)) { + fprintf(stderr, "[MESH] Skip occupied CPU %d (core idx %d) for numa_%d main worker\n", cpu_id, main_search_idx, this_numa); + main_search_idx++; + continue; + } + main_bound_cpu = cpu_id; + break; + } if (!core_obj) { fprintf(stderr, "Core %d inside NUMA node %d not found\n", 0, this_numa); // throw std::runtime_error("Core not found inside NUMA node"); continue; } + advance_core_idx(this_numa, main_search_idx); // 精简 cpuset auto cpuset_simple = hwloc_bitmap_alloc(); hwloc_bitmap_copy(cpuset_simple, core_obj->cpuset); @@ -316,6 +498,8 @@ void NumaJobDistributor::init(std::vector numa_ids, std::vector thread auto res = hwloc_set_thread_cpubind(topology, native_handle, cpuset_simple, HWLOC_CPUBIND_STRICT); if (res != 0) { fprintf(stderr, "Failed to set thread CPU binding: %s\n", strerror(errno)); + } else { + fprintf(stderr, "[MESH] numa_%d_m_%d -> CPU %d (core idx %d)\n", this_numa, start_id, main_bound_cpu, main_search_idx); } // 检查线程是否绑定到指定的 核上了 hwloc_cpuset_t cpuset = hwloc_bitmap_alloc(); @@ -399,7 +583,10 @@ void NumaJobDistributor::worker_thread(int numa_id) { } else if (stat == ThreadStatus::WAITING) { auto now = std::chrono::high_resolution_clock::now(); auto duration = std::chrono::duration_cast(now - start).count(); - if (duration > 50) { + // 禁用 cv.wait：decode 阶段层间间隔可能 > 50ms（attention 慢）， + // worker 进入 cv.wait 后唤醒延迟 11-27ms，导致 decode 性能从 36 tok/s 暴跌到 1 tok/s。 + // 改为永久 busy-wait，确保 worker 在下次 do_work_stealing_job_async 时立即响应。 + if (duration > 3600000) { std::unique_lock lock(*mutexes[numa_id]); cvs[numa_id]->wait(lock, [&] { return status[numa_id]->load(std::memory_order_acquire) != ThreadStatus::WAITING; diff --git a/kt-kernel/cpu_backend/worker_pool.h b/kt-kernel/cpu_backend/worker_pool.h index 72c9b6355..c845e6fab 100644 --- a/kt-kernel/cpu_backend/worker_pool.h +++ b/kt-kernel/cpu_backend/worker_pool.h @@ -23,7 +23,7 @@ #include #include -// #define PROFILE_BALANCE +#define PROFILE_BALANCE inline void set_to_numa(int this_numa) { struct bitmask* mask = numa_bitmask_alloc(numa_num_configured_nodes()); @@ -66,6 +66,11 @@ struct alignas(64) ThreadState { std::condition_variable cv; #ifdef PROFILE_BALANCE size_t finish_ns; + size_t init_ns; + size_t compute_ns; + size_t finalize_ns; + size_t work_set_ns; // do_work_stealing_job_async 设置 WORKING 的绝对时刻(ns since epoch) + size_t detect_ns; // worker 在 process_tasks 开头记录的绝对时刻(ns since epoch) #endif }; diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 01bf1cc3a..b9b3e1cda 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -1161,15 +1161,49 @@ PYBIND11_MODULE(kt_kernel_ext, m) { }, py::return_value_policy::reference_internal); // 注册 hook 函数指针（让 mesh_hook.hpp 的 inline 函数能调用 ResidencyManager） + // get_*_ptr hook 使用 with_reader 版本：返回非 nullptr 时 reader 已持有 mesh::hook::HookRegistry registry; registry.get_gate_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { - return ((mesh::MeshResidencyManager*)mgr)->get_gate_ptr(layer, tp, eid); + return ((mesh::MeshResidencyManager*)mgr)->get_gate_ptr_with_reader(layer, tp, eid); }; registry.get_up_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { - return ((mesh::MeshResidencyManager*)mgr)->get_up_ptr(layer, tp, eid); + return ((mesh::MeshResidencyManager*)mgr)->get_up_ptr_with_reader(layer, tp, eid); }; registry.get_down_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { - return ((mesh::MeshResidencyManager*)mgr)->get_down_ptr(layer, tp, eid); + return ((mesh::MeshResidencyManager*)mgr)->get_down_ptr_with_reader(layer, tp, eid); + }; + // 同步加载版本：slot 未命中时阻塞从 SSD 加载到 slot（内联 acquire_reader） + registry.load_gate_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_or_load_gate_ptr(layer, tp, eid); + }; + registry.load_up_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_or_load_up_ptr(layer, tp, eid); + }; + registry.load_down_ptr = [](void* mgr, int layer, int tp, int eid) -> void* { + return ((mesh::MeshResidencyManager*)mgr)->get_or_load_down_ptr(layer, tp, eid); + }; + // acquire_reader 已内联到 get/load hook 中，此处为 no-op + registry.acquire_reader = [](void* mgr, int layer, int tp, int eid) -> void { + (void)mgr; (void)layer; (void)tp; (void)eid; + }; + registry.release_reader = [](void* mgr, int layer, int tp, int eid) -> void { + ((mesh::MeshResidencyManager*)mgr)->release_reader(layer, tp, eid); + }; + // decode 层级回调：调用完整 defer 调度路径（on_decode_layer） + // on_decode_layer 会：process_cqes → split(immediate/deferred) → drain_and_submit(异步io_uring) + // missing 专家走 deferred 异步读取，I/O 与 GEMM 重叠（SKILL.md 第 96-100 行） + registry.on_decode_layer = [](void* mgr, int layer, int tp, + const int* topk, int k, const float* weights, int expert_num) -> void { + // 构造 topk 和 scores（全长，top-k 位置填入 weights，其余 0） + std::vector topk_vec(topk, topk + k); + std::vector scores_vec(expert_num, 0.0f); + for (int i = 0; i < k; i++) { + int eid = topk[i]; + if (eid >= 0 && eid < expert_num) { + scores_vec[eid] = weights[i]; + } + } + ((mesh::MeshResidencyManager*)mgr)->on_decode_layer(layer, topk_vec, scores_vec, tp); }; mesh::hook::register_hooks(registry); diff --git a/kt-kernel/operators/amx/bf16-moe.hpp b/kt-kernel/operators/amx/bf16-moe.hpp index 389446ed6..da1834aa7 100644 --- a/kt-kernel/operators/amx/bf16-moe.hpp +++ b/kt-kernel/operators/amx/bf16-moe.hpp @@ -18,6 +18,7 @@ #include "la/amx_raw_kernels.hpp" #include "la/amx_utils.hpp" // For transpose_16x16_32bit #include "moe_base.hpp" +#include "../mesh/mesh_hook.hpp" /** * @brief BF16 MoE operator using CRTP pattern @@ -94,26 +95,97 @@ class AMX_BF16_MOE_TP : public AMX_MOE_BASE> { void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; auto& ba = gate_up_ba_[expert_idx]; - auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; - // Use vec_mul/mat_mul (no group_size) - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + // MESH hook: 返回非 nullptr 时 reader 已自动持有 + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = do_up + ? mesh::hook::get_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) + : mesh::hook::get_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } + + if (slot_ptr) { + // MESH 路径：slot 命中，reader 已持有 + auto bb = std::make_shared( + config_.intermediate_size, config_.hidden_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } else if (config_.mesh_enabled && config_.mesh_residency) { + // MESH 模式 slot 未命中：同步从 SSD 加载到 slot（gate_bb_ 为空，不能回退） + slot_ptr = do_up + ? mesh::hook::load_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) + : mesh::hook::load_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + if (slot_ptr) { + // 加载成功，reader 已持有 + auto bb = std::make_shared( + config_.intermediate_size, config_.hidden_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } } else { - amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + // 原版路径 + auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + } } } void do_down_gemm(int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; + auto& ba = down_ba_[expert_idx]; + auto& bc = down_bc_[expert_idx]; + + // MESH hook: 返回非 nullptr 时 reader 已持有 + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = mesh::hook::get_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], down_bb_[expert_idx], - down_bc_[expert_idx], ith, nth); + if (slot_ptr) { + // MESH 路径：slot 命中，reader 已持有 + auto bb = std::make_shared( + config_.hidden_size, config_.intermediate_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } else if (config_.mesh_enabled && config_.mesh_residency) { + // MESH 模式 slot 未命中：同步从 SSD 加载到 slot（down_bb_ 为空，不能回退） + slot_ptr = mesh::hook::load_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + if (slot_ptr) { + // 加载成功，reader 已持有 + auto bb = std::make_shared( + config_.hidden_size, config_.intermediate_size, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + } + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } } else { - amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], down_bb_[expert_idx], - down_bc_[expert_idx], ith, nth); + // 原版路径 + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], down_bb_[expert_idx], + down_bc_[expert_idx], ith, nth); + } else { + amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, down_ba_[expert_idx], down_bb_[expert_idx], + down_bc_[expert_idx], ith, nth); + } } } @@ -161,6 +233,10 @@ class AMX_BF16_MOE_TP : public AMX_MOE_BASE> { * Loads weights from config_.gate_proj, up_proj, down_proj (no scales). */ void load_weights() { + // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版加载路径 + if (config_.mesh_enabled && config_.mesh_residency) { + return; + } const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; auto pool = config_.pool->get_subpool(tp_part_idx); @@ -439,6 +515,14 @@ class TP_MOE> : public TP_MOEconfig.mesh_enabled && this->config.mesh_residency) { + for (auto i = 0; i < this->tp_count; i++) { + this->tps[i]->load_weights(); // 各 TP 走 mesh 分支 + } + this->weights_loaded = true; + return; + } auto& config = this->config; auto& tps = this->tps; auto& tp_count = this->tp_count; diff --git a/kt-kernel/operators/amx/moe.hpp b/kt-kernel/operators/amx/moe.hpp index a23711f66..90fd33a67 100644 --- a/kt-kernel/operators/amx/moe.hpp +++ b/kt-kernel/operators/amx/moe.hpp @@ -17,6 +17,8 @@ #include "moe_base.hpp" #include "../mesh/mesh_hook.hpp" +#include + template class AMX_MOE_TP : public AMX_MOE_BASE> { protected: @@ -179,27 +181,74 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { // CRTP virtual points - GEMM dispatch // ============================================================================ + // MESH GEMM 辅助：根据 kernel 类型选择栈对象或 make_shared 路径 + // GemmKernel224BF 不支持 integer_mat_mul，需走 make_shared + mat_mul/vec_mul + // Int4/Int8 路径用栈对象，避免热路径堆分配/释放导致的堆损坏 + // （make_shared 的堆 chunk 易被相邻 bc pool 溢出覆盖，触发 free(): unaligned chunk） + void run_mesh_gemm(int m, int n, int k, + std::shared_ptr& ba, + std::shared_ptr& bc, + void* slot_ptr, int ith, int nth, int qlen) { + if constexpr (std::is_same_v) { + auto bb = std::make_shared(n, k, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mat_mul(m, n, k, ba, bb, bc, ith, nth); + } else { + amx::vec_mul(m, n, k, ba, bb, bc, ith, nth); + } + } else { + typename T::BufferB bb(n, k, slot_ptr); + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::integer_mat_mul(m, n, k, ba.get(), &bb, bc.get(), ith, nth); + } else { + amx::integer_mat_mul(m, n, k, ba.get(), &bb, bc.get(), ith, nth); + } + } + } + void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; auto& ba = gate_up_ba_[expert_idx]; auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; // A1-call: MESH hook — 如果 mesh_residency 已设置，尝试从 slot 池获取权重指针 + // get_*_bb / load_*_bb 返回非 nullptr 时 reader 已自动持有 void* slot_ptr = nullptr; + auto _t_acq = std::chrono::steady_clock::now(); if (config_.mesh_residency) { slot_ptr = do_up ? mesh::hook::get_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) : mesh::hook::get_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); } + auto _t_acq_end = std::chrono::steady_clock::now(); + long acq_us = std::chrono::duration_cast(_t_acq_end - _t_acq).count(); if (slot_ptr) { - // MESH 路径：用 slot 指针创建临时 BufferB - auto bb = std::make_shared( - config_.intermediate_size, config_.hidden_size, slot_ptr); - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + // MESH 路径：slot 命中，reader 已持有 + run_mesh_gemm(m, config_.intermediate_size, config_.hidden_size, ba, bc, slot_ptr, ith, nth, qlen); + auto _t_gem_end = std::chrono::steady_clock::now(); + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + auto _t_rel_end = std::chrono::steady_clock::now(); + long gem_us = std::chrono::duration_cast(_t_gem_end - _t_acq_end).count(); + long rel_us = std::chrono::duration_cast(_t_rel_end - _t_gem_end).count(); + if (ith == 0 && (acq_us > 100 || gem_us > 100 || rel_us > 100)) { + fprintf(stderr, "[GEMM_SLOW] L%d tp%d ex%d up=%d acq=%ldus gem=%ldus rel=%ldus\n", + config_.layer_idx, tp_part_idx, expert_idx, do_up, acq_us, gem_us, rel_us); + } + } else if (config_.mesh_enabled && config_.mesh_residency) { + // MESH 模式 slot 未命中：同步从 SSD 加载到 slot（gate_bb_ 为空，不能回退） + fprintf(stderr, "[GEMM_LOAD] L%d tp%d ex%d up=%d SLOT MISS → sync load\n", + config_.layer_idx, tp_part_idx, expert_idx, do_up); + slot_ptr = do_up + ? mesh::hook::load_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) + : mesh::hook::load_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + if (slot_ptr) { + // 加载成功，reader 已持有 + run_mesh_gemm(m, config_.intermediate_size, config_.hidden_size, ba, bc, slot_ptr, ith, nth, qlen); + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); } else { - amx::vec_mul(m, config_.intermediate_size, config_.hidden_size, ba, bb, bc, ith, nth); + fprintf(stderr, "[MESH DEBUG] do_gate_up_gemm: load FAILED layer=%d tp=%d expert=%d do_up=%d\n", + config_.layer_idx, tp_part_idx, expert_idx, do_up); } } else { // 原版路径 @@ -217,20 +266,23 @@ class AMX_MOE_TP : public AMX_MOE_BASE> { auto& ba = down_ba_[expert_idx]; auto& bc = down_bc_[expert_idx]; - // A1-call: MESH hook + // A1-call: MESH hook — 返回非 nullptr 时 reader 已持有 void* slot_ptr = nullptr; if (config_.mesh_residency) { slot_ptr = mesh::hook::get_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); } if (slot_ptr) { - // MESH 路径 - auto bb = std::make_shared( - config_.hidden_size, config_.intermediate_size, slot_ptr); - if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { - amx::mat_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); - } else { - amx::vec_mul(m, config_.hidden_size, config_.intermediate_size, ba, bb, bc, ith, nth); + // MESH 路径：slot 命中，reader 已持有 + run_mesh_gemm(m, config_.hidden_size, config_.intermediate_size, ba, bc, slot_ptr, ith, nth, qlen); + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } else if (config_.mesh_enabled && config_.mesh_residency) { + // MESH 模式 slot 未命中：同步从 SSD 加载到 slot（down_bb_ 为空，不能回退） + slot_ptr = mesh::hook::load_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + if (slot_ptr) { + // 加载成功，reader 已持有 + run_mesh_gemm(m, config_.hidden_size, config_.intermediate_size, ba, bc, slot_ptr, ith, nth, qlen); + mesh::hook::release_reader(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); } } else { // 原版路径 diff --git a/kt-kernel/operators/amx/moe_base.hpp b/kt-kernel/operators/amx/moe_base.hpp index 48c4da9ee..de5402929 100644 --- a/kt-kernel/operators/amx/moe_base.hpp +++ b/kt-kernel/operators/amx/moe_base.hpp @@ -10,7 +10,7 @@ #ifndef CPUINFER_OPERATOR_AMX_MOE_BASE_H #define CPUINFER_OPERATOR_AMX_MOE_BASE_H -// #define FORWARD_TIME_PROFILE +#define FORWARD_TIME_PROFILE #include @@ -34,6 +34,7 @@ #include "../moe-tp.hpp" #include "la/amx.hpp" #include "llama.cpp/ggml.h" +#include "../mesh/mesh_hook.hpp" template class AMX_MOE_BASE { @@ -269,12 +270,32 @@ class AMX_MOE_BASE { used_pool_bytes_bc_down += bc_down_size; } - assert(used_pool_m <= pool_count_); - assert(used_pool_bytes_a <= gate_up_ba_pool_bytes_); - assert(used_pool_bytes_bc_gate <= gate_bc_pool_bytes_); - assert(used_pool_bytes_bc_up <= up_bc_pool_bytes_); - assert(used_pool_bytes_ba_down <= down_ba_pool_bytes_); - assert(used_pool_bytes_bc_down <= down_bc_pool_bytes_); + // 显式边界检查（Release 中 assert 被禁用，用 fprintf + abort 替代） + if (used_pool_m > pool_count_ || used_pool_bytes_a > gate_up_ba_pool_bytes_ || + used_pool_bytes_bc_gate > gate_bc_pool_bytes_ || used_pool_bytes_bc_up > up_bc_pool_bytes_ || + used_pool_bytes_ba_down > down_ba_pool_bytes_ || used_pool_bytes_bc_down > down_bc_pool_bytes_) { + fprintf(stderr, "[MESH POOL OVERFLOW prefill] qlen=%d k=%d activated=%d max_len=%d " + "M_STEP=%d pool_count=%zu used_m=%zu\n", + qlen, k, activated_expert, config_.max_len, (int)M_STEP, + pool_count_, used_pool_m); + fprintf(stderr, " ba_a: used=%zu pool=%zu\n", used_pool_bytes_a, gate_up_ba_pool_bytes_); + fprintf(stderr, " bc_gate: used=%zu pool=%zu\n", used_pool_bytes_bc_gate, gate_bc_pool_bytes_); + fprintf(stderr, " bc_up: used=%zu pool=%zu\n", used_pool_bytes_bc_up, up_bc_pool_bytes_); + fprintf(stderr, " ba_down: used=%zu pool=%zu\n", used_pool_bytes_ba_down, down_ba_pool_bytes_); + fprintf(stderr, " bc_down: used=%zu pool=%zu\n", used_pool_bytes_bc_down, down_bc_pool_bytes_); + abort(); + } + // 检查 m_local_output_ 缓冲区是否足够 + size_t total_offset = offset; + size_t gate_out_needed = total_offset * config_.intermediate_size * sizeof(ggml_bf16_t); + size_t gate_out_have = sizeof(ggml_bf16_t) * config_.num_experts_per_tok * config_.max_len * config_.intermediate_size; + if (gate_out_needed > gate_out_have) { + fprintf(stderr, "[MESH OUTPUT OVERFLOW prefill] gate_out: needed=%zu have=%zu " + "offset=%zu I=%d k_tok=%d ml=%d\n", + gate_out_needed, gate_out_have, total_offset, + config_.intermediate_size, config_.num_experts_per_tok, config_.max_len); + abort(); + } #ifdef FORWARD_TIME_PROFILE { @@ -507,11 +528,20 @@ class AMX_MOE_BASE { used_pool_bytes_ba_down += ba_down_size; used_pool_bytes_bc_down += bc_down_size; } - assert(used_pool_m <= pool_count_); - assert(used_pool_bytes_bc_gate <= gate_bc_pool_bytes_); - assert(used_pool_bytes_bc_up <= up_bc_pool_bytes_); - assert(used_pool_bytes_ba_down <= down_ba_pool_bytes_); - assert(used_pool_bytes_bc_down <= down_bc_pool_bytes_); + // 显式边界检查（Release 中 assert 被禁用） + if (used_pool_m > pool_count_ || used_pool_bytes_bc_gate > gate_bc_pool_bytes_ || + used_pool_bytes_bc_up > up_bc_pool_bytes_ || used_pool_bytes_ba_down > down_ba_pool_bytes_ || + used_pool_bytes_bc_down > down_bc_pool_bytes_) { + fprintf(stderr, "[MESH POOL OVERFLOW decode] k=%d activated=%d max_len=%d " + "M_STEP=%d pool_count=%zu used_m=%zu\n", + k, activated_expert, config_.max_len, (int)M_STEP, + pool_count_, used_pool_m); + fprintf(stderr, " bc_gate: used=%zu pool=%zu\n", used_pool_bytes_bc_gate, gate_bc_pool_bytes_); + fprintf(stderr, " bc_up: used=%zu pool=%zu\n", used_pool_bytes_bc_up, up_bc_pool_bytes_); + fprintf(stderr, " ba_down: used=%zu pool=%zu\n", used_pool_bytes_ba_down, down_ba_pool_bytes_); + fprintf(stderr, " bc_down: used=%zu pool=%zu\n", used_pool_bytes_bc_down, down_bc_pool_bytes_); + abort(); + } void* gate_up_ba_pool_ptr = gate_up_ba_pool_; for (int i = 0; i < activated_expert; i++) { @@ -532,6 +562,16 @@ class AMX_MOE_BASE { } #endif + // Bug 1 fix: 提交异步预取 + 更新 Heat/Markov + // 在 GEMM 之前调用 on_decode_layer，让 missing 专家的 io_uring 异步读提前提交， + // 同时逐层更新 Heat/Markov 使驱逐策略生效（Bug 2） + if (config_.mesh_enabled && config_.mesh_residency) { + // expert_ids 是 int64_t*，需要转为 int* + std::vector topk_int(expert_ids, expert_ids + k); + mesh::hook::on_decode_layer(config_.mesh_residency, config_.layer_idx, tp_part_idx, + topk_int.data(), k, weights, config_.expert_num); + } + int nth = T::recommended_nth(config_.intermediate_size); pool->do_work_stealing_job( nth * activated_expert * 2, [](int _) { T::config(); }, diff --git a/kt-kernel/operators/avx2/bf16-moe.hpp b/kt-kernel/operators/avx2/bf16-moe.hpp index bdb1a8d8d..c032aab3c 100644 --- a/kt-kernel/operators/avx2/bf16-moe.hpp +++ b/kt-kernel/operators/avx2/bf16-moe.hpp @@ -15,6 +15,7 @@ #include "avx2_bf16_gemm.hpp" #include "moe_base.hpp" +#include "../mesh/mesh_hook.hpp" template class AVX2_BF16_MOE_TP : public AVX2_MOE_BASE> { @@ -64,16 +65,45 @@ class AVX2_BF16_MOE_TP : public AVX2_MOE_BASE> { void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; auto& ba = gate_up_ba_[expert_idx]; - auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; - avx2::gemm_bf16(m, config_.intermediate_size, config_.hidden_size, *ba, *bb, *bc, ith, nth); + // MESH hook + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = do_up + ? mesh::hook::get_up_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx) + : mesh::hook::get_gate_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } + + if (slot_ptr) { + auto bb = std::make_shared( + config_.intermediate_size, config_.hidden_size, slot_ptr); + avx2::gemm_bf16(m, config_.intermediate_size, config_.hidden_size, *ba, *bb, *bc, ith, nth); + } else { + auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; + avx2::gemm_bf16(m, config_.intermediate_size, config_.hidden_size, *ba, *bb, *bc, ith, nth); + } } void do_down_gemm(int expert_idx, int ith, int nth, int qlen) { int m = m_local_num_[expert_idx]; - avx2::gemm_bf16(m, config_.hidden_size, config_.intermediate_size, - *down_ba_[expert_idx], *down_bb_[expert_idx], *down_bc_[expert_idx], ith, nth); + auto& ba = down_ba_[expert_idx]; + auto& bc = down_bc_[expert_idx]; + + // MESH hook + void* slot_ptr = nullptr; + if (config_.mesh_residency) { + slot_ptr = mesh::hook::get_down_bb(config_.mesh_residency, config_.layer_idx, tp_part_idx, expert_idx); + } + + if (slot_ptr) { + auto bb = std::make_shared( + config_.hidden_size, config_.intermediate_size, slot_ptr); + avx2::gemm_bf16(m, config_.hidden_size, config_.intermediate_size, *ba, *bb, *bc, ith, nth); + } else { + avx2::gemm_bf16(m, config_.hidden_size, config_.intermediate_size, + *down_ba_[expert_idx], *down_bb_[expert_idx], *down_bc_[expert_idx], ith, nth); + } } /** @@ -81,6 +111,10 @@ class AVX2_BF16_MOE_TP : public AVX2_MOE_BASE> { * BufferB::from_mat is a trivial memcpy for AVX2 (no AMX transpose). */ void load_weights() { + // MESH 插件：权重由 MeshResidencyManager 接管，跳过原版加载路径 + if (config_.mesh_enabled && config_.mesh_residency) { + return; + } const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; auto pool = config_.pool->get_subpool(tp_part_idx); @@ -228,6 +262,14 @@ class TP_MOE> : public TP_MOEconfig.mesh_enabled && this->config.mesh_residency) { + for (auto i = 0; i < this->tp_count; i++) { + this->tps[i]->load_weights(); + } + this->weights_loaded = true; + return; + } auto& config = this->config; auto& tps = this->tps; auto& tp_count = this->tp_count; diff --git a/kt-kernel/operators/mesh/mesh_decode.hpp b/kt-kernel/operators/mesh/mesh_decode.hpp index caa54c6e3..28e038dcc 100644 --- a/kt-kernel/operators/mesh/mesh_decode.hpp +++ b/kt-kernel/operators/mesh/mesh_decode.hpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include "mesh_config.hpp" @@ -45,103 +46,92 @@ class MeshDecode { // immediate/deferred 分组结果 struct SplitResult { - std::vector immediate; // 本层立即计算 - std::vector deferred; // 推迟到下一层计算 + std::vector immediate; // 本层立即计算（CACHED + GPU + 阻塞读完成的 missing） + std::vector deferred; // 推迟到下一层计算（cached + missing） + std::vector blocking_missing; // SKILL.md 第 282 行：需要阻塞读补充 immediate 组的 missing 专家 + std::vector deferred_missing; // SKILL.md 第 284 行：deferred 组中需要提交异步 io_uring 的 missing 专家 }; /** - * @brief 对 top-k 专家进行 immediate/deferred 分组 + * @brief 对 top-k 专家进行 immediate/deferred 分组（SKILL.md 第 280-284 行） * - * B6 fix: 入口先过滤 GPU 专家（GPU 专家不参与 slot 调度） + * 分组逻辑（与 KTransformers select_deferred_experts 一致）： + * - target_immediate = original_k - defer_count（= KTransformers 的 protected_k） + * - 按 score 降序排序整个 topk，取 score 最高的 target_immediate 个作为 immediate 候选 + * （与 KTransformers 的 torch.topk(expert_scores, k=protected_k) 一致） + * - immediate 候选中：GPU 专家直接进 immediate；CPU 已 CACHED 进 immediate； + * CPU 未 CACHED 加入 blocking_missing 阻塞读后进 immediate（SKILL.md 第 282 行） + * - 其余 topk 走 deferred，其中未 CACHED 的提交异步 io_uring（SKILL.md 第 284 行） * - * @param topk 当前层 top-k 专家 ID - * @param scores 各专家的 router score（全长 expert_num） - * @param slot_pool 该层 slot 池 - * @param scheduler 调度器 - * @param layer_idx 当前层 - * @param tp_part_idx TP 分片 - * @param get_dst_ptrs 获取 slot buffer 指针的回调 - * @param is_gpu_expert 判断专家是否在 GPU 上的回调 - * @return SplitResult + * 一致性保证：MESH 的 immediate 组 = KTransformers 的 immediate 组， + * GEMM 计算的 CPU 专家都已 CACHED（未 CACHED 的会被阻塞读补充）， + * 不会走同步 SSD 加载路径。 + * + * 注意：split 只负责分组，不提交 io_uring。blocking_missing 由 on_decode_layer 阻塞读处理， + * deferred_missing 由 on_decode_layer 异步提交 io_uring。 + * + * @param topk effective_topk（含上一层层 deferred 已 CACHED 的专家） + * @param original_k 原始 top-k 数（当前 token 路由选出的，不含 prev_layer_deferred_） */ SplitResult split(const std::vector& topk, const std::vector& scores, MeshSlotPool& slot_pool, - MeshScheduler& scheduler, - int layer_idx, int tp_part_idx, - std::function(int)> get_dst_ptrs, + int layer_idx, + int original_k, std::function is_gpu_expert = nullptr) { - int k = static_cast(topk.size()); - int target_immediate = std::max(0, k - defer_count_); - - // B6 fix: 过滤 GPU 专家，GPU 专家直接进 immediate（由 GPU 计算，不需 slot） - std::vector cpu_topk; - std::vector gpu_experts; + // SKILL.md 第 280 行：target_immediate 基于原始 top-k 数，不是 effective_topk.size() + // = KTransformers 的 protected_k = num_experts_per_tok - max_deferred_experts_per_token + int target_immediate = std::max(0, original_k - defer_count_); + + // 过滤掉 -1（KTransformers select_deferred_experts 把 deferred 位置填 -1） + // on_decode_layer hook 接收 GEMM 的 expert_ids，defer 模式下含 -1 + // -1 不是有效专家 ID，提交 io_uring 会报 Bad file descriptor + std::vector valid_topk; + valid_topk.reserve(topk.size()); for (int eid : topk) { - if (is_gpu_expert && is_gpu_expert(eid)) { - gpu_experts.push_back(eid); - } else { - cpu_topk.push_back(eid); + if (eid >= 0) { + valid_topk.push_back(eid); } } - // 区分已缓存和缺失（仅 CPU 专家） - std::vector cached, missing; - for (int eid : cpu_topk) { - if (slot_pool.is_cached(eid)) { - cached.push_back(eid); - } else { - missing.push_back(eid); - } - } + // 与 KTransformers select_deferred_experts 一致：按 score 降序排序整个 topk， + // 取 score 最高的 target_immediate 个作为 immediate 候选。 + // KTransformers 用 torch.topk(expert_scores, k=protected_k) 选 score 最高的位置， + // MESH 用 scores[expert_id] 排序（top-k 无重复时两者等价）。 + std::sort(valid_topk.begin(), valid_topk.end(), + [&scores](int a, int b) { + float sa = (a < (int)scores.size()) ? scores[a] : 0.0f; + float sb = (b < (int)scores.size()) ? scores[b] : 0.0f; + return sa > sb; // 降序 + }); + + int n_immediate = std::min(static_cast(valid_topk.size()), target_immediate); SplitResult result; - // GPU 专家直接进 immediate - result.immediate = gpu_experts; - - if (static_cast(cached.size()) > target_immediate) { - // immediate 数 > k-defer：按 router score 排序取前 target_immediate 个 - std::sort(cached.begin(), cached.end(), - [&scores](int a, int b) { - float sa = (a < (int)scores.size()) ? scores[a] : 0.0f; - float sb = (b < (int)scores.size()) ? scores[b] : 0.0f; - return sa > sb; // 降序 - }); - result.immediate.insert(result.immediate.end(), cached.begin(), cached.begin() + target_immediate); - result.deferred.assign(cached.begin() + target_immediate, cached.end()); - // 缺失专家全部进 deferred - result.deferred.insert(result.deferred.end(), missing.begin(), missing.end()); - - // 为缺失专家提交 deferred 异步读 - for (int eid : missing) { - auto ptrs = get_dst_ptrs(eid); - scheduler.submit_decode_deferred(layer_idx, eid, tp_part_idx, - ptrs[0], ptrs[1], ptrs[2]); - } - } else { - // immediate 数 < k-defer：阻塞读缺失专家直到 immediate 数量够 - result.immediate.insert(result.immediate.end(), cached.begin(), cached.end()); - int need = target_immediate - static_cast(cached.size()); - - for (int i = 0; i < need && i < static_cast(missing.size()); i++) { - int eid = missing[i]; - auto ptrs = get_dst_ptrs(eid); - // 提交 immediate 异步读（schedule_key = 当前层） - scheduler.submit_decode_immediate(layer_idx, eid, tp_part_idx, - ptrs[0], ptrs[1], ptrs[2]); - // 阻塞等待读取完成 - // 实际实现需要与 MeshIoUring 配合等待 CQE - // io_.wait_expert(eid, n_reqs); + // 遍历 immediate 候选，区分 GPU/CPU 和 CACHED/missing + for (int i = 0; i < n_immediate; i++) { + int eid = valid_topk[i]; + if (is_gpu_expert && is_gpu_expert(eid)) { + // GPU 专家直接进 immediate（由 GPU 计算，不需 slot） result.immediate.push_back(eid); + } else if (slot_pool.is_cached(eid)) { + // CPU 专家已 CACHED，进 immediate + result.immediate.push_back(eid); + } else { + // CPU 专家未 CACHED，阻塞读后进 immediate（SKILL.md 第 282 行） + result.blocking_missing.push_back(eid); + result.immediate.push_back(eid); // 阻塞读完成后会 CACHED } + } - // 剩余缺失专家走 deferred - for (int i = need; i < static_cast(missing.size()); i++) { - int eid = missing[i]; - auto ptrs = get_dst_ptrs(eid); - scheduler.submit_decode_deferred(layer_idx, eid, tp_part_idx, - ptrs[0], ptrs[1], ptrs[2]); - result.deferred.push_back(eid); + // 其余 topk 走 deferred（SKILL.md 第 284 行） + for (int i = n_immediate; i < static_cast(valid_topk.size()); i++) { + int eid = valid_topk[i]; + result.deferred.push_back(eid); + if (!is_gpu_expert || !is_gpu_expert(eid)) { + if (!slot_pool.is_cached(eid)) { + result.deferred_missing.push_back(eid); + } } } @@ -156,6 +146,10 @@ class MeshDecode { * 2. 查 victim 的 slot_idx * 3. 调用 overwrite(slot_idx, new_expert_id) * + * 引用计数保护：overwrite 内部 CV 等待 active_readers 归零， + * GEMM 时 acquire_reader 持有的专家不会被驱逐（SKILL.md 第 62/125 行）。 + * immediate 组是 score 最高的 CACHED 专家，evict 选 score 最低的，不会选 immediate 组。 + * * @param slot_pool 该层 slot 池 * @param scorer 驱逐评分器 * @param layer_idx 当前层 @@ -165,11 +159,11 @@ class MeshDecode { int evict_for_new_expert(MeshSlotPool& slot_pool, const EvictionScorer& scorer, int layer_idx, int new_expert_id) { - // 前五层满配检查 + // 前五层满配检查：用实际 slot 数量判断 + // 当 slot_pool.cap() < expert_num_ 时，前 5 层也无法容纳所有专家，必须允许驱逐 if (is_front_layer(layer_idx)) { - int cap = get_effective_cap(layer_idx, slot_pool.cap()); - if (cap >= expert_num_) { - return -1; // 满配，无需驱逐 + if (slot_pool.cap() >= expert_num_) { + return -1; // 实际 slot 数 >= 专家总数，满配无需驱逐 } } diff --git a/kt-kernel/operators/mesh/mesh_eviction.hpp b/kt-kernel/operators/mesh/mesh_eviction.hpp index 1f6868905..1bc18dd92 100644 --- a/kt-kernel/operators/mesh/mesh_eviction.hpp +++ b/kt-kernel/operators/mesh/mesh_eviction.hpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -62,6 +63,29 @@ class HeatTracker { } } + // 逐层更新层内 EMA（在 on_decode_layer 中调用） + // 只更新 token_layer_heat_，不更新 global_heat_。 + // global_heat_ 的 beta_ 衰减是"跨 token"语义，每 token 只更新一次。 + // 注意：commit_token（GPU 回调路径）当前未集成（submit_gating_scores_with_cuda_stream + // 从未被调用），所以 global_heat_ 的更新由 update_global_heat() 在最后一层时触发。 + void commit_layer(const std::vector& layer_scores) { + // 层内 EMA：gamma_ 衰减，每层更新一次 + for (int e = 0; e < expert_num_ && e < (int)layer_scores.size(); e++) { + token_layer_heat_[e] = gamma_ * token_layer_heat_[e] + (1.0f - gamma_) * layer_scores[e]; + } + } + + // 跨 token 全局 EMA 更新（每 token 调用一次，由 EvictionScorer::commit_layer + // 在最后一层 layer_idx == num_layers-1 时触发） + // 之前的 bug：commit_layer 每层都应用 beta_，导致一个 token（40 层）内 + // global_heat_ 被衰减 40 次（0.5^40 ≈ 0），历史 heat 信息全部丢失， + // select_victim 退化为随机选择，高频专家被错误驱逐 → SLOT MISS → pread。 + void update_global_heat() { + for (int e = 0; e < expert_num_; e++) { + global_heat_[e] = beta_ * global_heat_[e] + (1.0f - beta_) * token_layer_heat_[e]; + } + } + // 获取某专家的全局 heat float heat(int expert_id) const { if (expert_id < 0 || expert_id >= expert_num_) return 0.0f; @@ -293,6 +317,27 @@ class EvictionScorer { } } + // 逐层更新 Heat 和 Markov（在 on_decode_layer 中调用） + // 加锁保护：多个 TP 并发调用 on_decode_layer，scorer_ 是共享对象 + // 关键：global_heat_ 的 beta_ 衰减只在最后一层（layer_idx == num_layers_-1）时更新一次， + // 而不是每层都更新。之前每层都更新导致 beta_ 被应用 40 次（0.5^40≈0），heat 信息丢失。 + // 注意：commit_token（GPU 回调路径）当前未集成，所以 global_heat_ 更新放在这里。 + void commit_layer(int layer_idx, + const std::vector& topk, + const std::vector& scores, + const std::vector& prev_topk, + const std::vector& prev_scores) { + std::lock_guard lock(mtx_); + heat_.commit_layer(scores); + if (layer_idx > 0 && !prev_topk.empty()) { + markov_.update(layer_idx - 1, prev_topk, prev_scores, topk, scores); + } + // 只在最后一层结束时更新 global_heat_（每 token 一次，正确的 beta_ 衰减语义） + if (layer_idx == num_layers_ - 1) { + heat_.update_global_heat(); + } + } + /** * @brief B2: 在第 layer_idx 层结束时，预测第 layer_idx+1 层的 cross_layer_prior * @@ -305,6 +350,7 @@ class EvictionScorer { void predict_next_layer(int layer_idx, const std::vector& topk, const std::vector& scores) { if (layer_idx < 0 || layer_idx + 1 >= num_layers_) return; + std::lock_guard lock(mtx_); std::vector predicted; markov_.predict(layer_idx, topk, scores, predicted); auto& prior = cross_layer_prior_[layer_idx + 1]; @@ -359,6 +405,9 @@ class EvictionScorer { int expert_num_; // B2: [layer][expert] 的 Markov 先验，由 predict_next_layer 累积 std::vector> cross_layer_prior_; + // 多 TP 并发保护：commit_layer/predict_next_layer 是写操作，需加锁 + // score/select_victim 是读操作，float 读写通常原子，不加锁（驱逐评分不需要精确） + mutable std::mutex mtx_; }; } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_hook.hpp b/kt-kernel/operators/mesh/mesh_hook.hpp index 746125035..37e988ada 100644 --- a/kt-kernel/operators/mesh/mesh_hook.hpp +++ b/kt-kernel/operators/mesh/mesh_hook.hpp @@ -27,6 +27,15 @@ struct HookRegistry { void* (*get_gate_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; void* (*get_up_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; void* (*get_down_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + // 同步加载版本：slot 未命中时阻塞从 SSD 加载到 slot + void* (*load_gate_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*load_up_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void* (*load_down_ptr)(void* mgr, int layer, int tp, int expert_id) = nullptr; + // 引用计数：防止 AMX 计算期间 slot 被驱逐 + void (*acquire_reader)(void* mgr, int layer, int tp, int expert_id) = nullptr; + void (*release_reader)(void* mgr, int layer, int tp, int expert_id) = nullptr; + // Bug 1 fix: decode 层级回调 — 提交异步预取 + 更新 Heat/Markov + void (*on_decode_layer)(void* mgr, int layer, int tp, const int* topk, int k, const float* weights, int expert_num) = nullptr; }; inline HookRegistry& get_registry() { @@ -64,11 +73,62 @@ inline void* get_down_bb(void* mesh_mgr, int layer, int tp, int expert_id) { return fn(mesh_mgr, layer, tp, expert_id); } +// ===== 同步加载 hook（slot 未命中时阻塞加载）===== +// MESH 模式下 gate_bb_ 为空，slot 未命中不能回退，必须同步从 SSD 加载到 slot + +inline void* load_gate_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + auto fn = get_registry().load_gate_ptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +inline void* load_up_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + auto fn = get_registry().load_up_ptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +inline void* load_down_bb(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return nullptr; + auto fn = get_registry().load_down_ptr; + if (!fn) return nullptr; + return fn(mesh_mgr, layer, tp, expert_id); +} + +// ===== 引用计数 hook ===== +// 防止 AMX 计算期间 slot 被驱逐（overwrite 会 spin-wait active_readers 归零） + +inline void acquire_reader(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return; + auto fn = get_registry().acquire_reader; + if (!fn) return; + fn(mesh_mgr, layer, tp, expert_id); +} + +inline void release_reader(void* mesh_mgr, int layer, int tp, int expert_id) { + if (!mesh_mgr) return; + auto fn = get_registry().release_reader; + if (!fn) return; + fn(mesh_mgr, layer, tp, expert_id); +} + // ===== MESH 启用判断 hook ===== inline bool is_mesh_enabled(void* mesh_mgr) { return mesh_mgr != nullptr; } +// ===== Bug 1 fix: decode 层级回调 hook ===== +// 在 forward_decode 的 GEMM 之前调用，提交异步 io_uring 预取 + 更新 Heat/Markov +inline void on_decode_layer(void* mesh_mgr, int layer, int tp, + const int* topk, int k, const float* weights, int expert_num) { + if (!mesh_mgr) return; + auto fn = get_registry().on_decode_layer; + if (!fn) return; + fn(mesh_mgr, layer, tp, topk, k, weights, expert_num); +} + } // namespace hook } // namespace mesh diff --git a/kt-kernel/operators/mesh/mesh_io_uring.hpp b/kt-kernel/operators/mesh/mesh_io_uring.hpp index 76b8facfe..bffdb39ce 100644 --- a/kt-kernel/operators/mesh/mesh_io_uring.hpp +++ b/kt-kernel/operators/mesh/mesh_io_uring.hpp @@ -19,6 +19,7 @@ #include #include #include +#include #include // liburing.h 定义了 BLOCK_SIZE 等宏，会污染全局命名空间， // 与 amx_raw_kernels.hpp / fp8-moe.hpp 中的 BLOCK_SIZE 变量冲突 @@ -97,12 +98,19 @@ class MeshIoUring { total_up_mins += layouts_tp[tp][e].up_mins_bytes; total_down_mins += layouts_tp[tp][e].down_mins_bytes; } - cache.gate_scale = numa_alloc_onnode(total_gate, numa_node); - cache.up_scale = numa_alloc_onnode(total_up, numa_node); - cache.down_scale = numa_alloc_onnode(total_down, numa_node); - cache.gate_mins = numa_alloc_onnode(total_gate_mins, numa_node); - cache.up_mins = numa_alloc_onnode(total_up_mins, numa_node); - cache.down_mins = numa_alloc_onnode(total_down_mins, numa_node); + cache.gate_scale = total_gate > 0 ? numa_alloc_onnode(total_gate, numa_node) : nullptr; + cache.up_scale = total_up > 0 ? numa_alloc_onnode(total_up, numa_node) : nullptr; + cache.down_scale = total_down > 0 ? numa_alloc_onnode(total_down, numa_node) : nullptr; + cache.gate_mins = total_gate_mins > 0 ? numa_alloc_onnode(total_gate_mins, numa_node) : nullptr; + cache.up_mins = total_up_mins > 0 ? numa_alloc_onnode(total_up_mins, numa_node) : nullptr; + cache.down_mins = total_down_mins > 0 ? numa_alloc_onnode(total_down_mins, numa_node) : nullptr; + if ((total_gate > 0 && !cache.gate_scale) || (total_up > 0 && !cache.up_scale) || + (total_down > 0 && !cache.down_scale)) { + fprintf(stderr, "[MESH] preload_scale_cache: numa_alloc_onnode FAILED! " + "layer=%d tp=%d numa=%d total_gate=%zu total_up=%zu total_down=%zu\n", + layer_idx, tp, numa_node, total_gate, total_up, total_down); + throw std::runtime_error("preload_scale_cache: numa_alloc_onnode failed (OOM?)"); + } cache.gate_scale_total = total_gate; cache.up_scale_total = total_up; cache.down_scale_total = total_down; @@ -226,51 +234,86 @@ class MeshIoUring { layer_idx, tp_part_idx, expert_id, n_reqs, 0, std::move(on_complete) }; + // SQ ring 互斥锁保护：liburing 的 io_uring_get_sqe / io_uring_submit 不是线程安全的。 + // 多 TP 线程并发 submit_load 会导致 SQ ring tail 指针竞争，SQE 丢失，CQE 永远不回来。 + // 锁范围：从 get_sqe 到 io_uring_submit，down_scattered 的 pread 在锁外执行。 + std::lock_guard sq_lock(sq_mutex_); + + // submit_load nullptr fix: 检查 io_uring_get_sqe 返回值，SQ ring 满时先 submit 再重试 + // 避免 nullptr 解引用 segfault + auto get_sqe = [this]() -> io_uring_sqe* { + io_uring_sqe* sqe = io_uring_get_sqe(&ring_); + if (sqe) return sqe; + // SQ ring 满，先 submit 释放空间 + io_uring_submit(&ring_); + sqe = io_uring_get_sqe(&ring_); + if (!sqe) { + fprintf(stderr, "[MESH] submit_load: SQ ring full after submit, aborting\n"); + std::abort(); + } + return sqe; + }; + // 提交 SQE，user_data 指向 pending - auto* sqe = io_uring_get_sqe(&ring_); + auto* sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, gate_dst, layout.gate_bytes, layout.gate_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, up_dst, layout.up_bytes, layout.up_offset); io_uring_sqe_set_data(sqe, pending); // Bug 5 fix: 先提交 gate/up SQE，再 pread down，让 io_uring 与 pread 并行 // 原代码 pread 在 io_uring_submit 之前，阻塞期间 io_uring 空闲 if (!down_scattered) { - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, down_dst, layout.down_bytes, layout.down_offset); io_uring_sqe_set_data(sqe, pending); } if (!layer_scale_loaded && has_scale) { - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, gate_scale_dst, layout.gate_scale_bytes, layout.gate_scale_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, up_scale_dst, layout.up_scale_bytes, layout.up_scale_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, down_scale_dst, layout.down_scale_bytes, layout.down_scale_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, gate_mins_dst, layout.gate_mins_bytes, layout.gate_mins_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, up_mins_dst, layout.up_mins_bytes, layout.up_mins_offset); io_uring_sqe_set_data(sqe, pending); - sqe = io_uring_get_sqe(&ring_); + sqe = get_sqe(); io_uring_prep_read(sqe, layout.fd, down_mins_dst, layout.down_mins_bytes, layout.down_mins_offset); io_uring_sqe_set_data(sqe, pending); } // Bug 5 fix: 先 submit 所有 SQE，让 io_uring 开始 DMA - io_uring_submit(&ring_); + // submit_and_wait fix: 记录 inflight SQE 数，供 submit_and_wait 等待 + // 引用计数设计（SKILL.md 第 62/122-126 行）：deferred_missing 提交的 io_uring 读取， + // 即使目标 slot 需要 evict，overwrite 会 CV 等待 active_readers 归零， + // GEMM 中 acquire_reader 持有的专家不会被驱逐。不需要阻塞等待。 + inflight_count_.fetch_add(n_reqs, std::memory_order_release); + int submit_ret = io_uring_submit(&ring_); + if (submit_ret < 0) { + // io_uring_submit 失败：SQEs 仍在 SQ ring 中（未被内核消费）， + // 后续 io_uring_submit（get_sqe 或下一个 submit_load）会重新提交它们。 + // 绝对不能 delete pending（SQEs 的 user_data 指向它，CQE 回来时会访问）， + // 也不能回滚 inflight_count_（SQEs 后续会提交，CQE 会回来，process_cqes 会递减）。 + // 否则会导致 use-after-free → segfault。 + fprintf(stderr, "[MESH] io_uring_submit failed: ret=%d errno=%d (%s), expert=%d " + "(SQEs remain in ring, will retry on next submit)\n", + submit_ret, errno, strerror(errno), expert_id); + } // Bug 5 fix: down_scattered 的 pread 在 submit 之后执行，与 io_uring 并行 if (down_scattered) { @@ -313,26 +356,41 @@ class MeshIoUring { // A6: 处理已完成的 CQE，触发 on_complete 回调 // 非阻塞：只处理当前已有的 CQE，不等待 + // 根因修复：当 cqe->res < 0 时不调用 on_complete，避免标记垃圾 slot 为 CACHED + // Bug 1 fix: 加 mutex 保护，允许多个 worker 线程在 GEMM 期间并发调用 + + // 诊断：返回当前 inflight SQE 数 + int inflight_count() const { return inflight_count_.load(std::memory_order_acquire); } + void process_cqes() { + std::lock_guard lock(cq_mutex_); unsigned head; unsigned count = 0; io_uring_cqe* cqe; io_uring_for_each_cqe(&ring_, head, cqe) { auto* pending = static_cast(io_uring_cqe_get_data(cqe)); + // submit_and_wait fix: 每个 CQE 对应一个已完成的 SQE + inflight_count_.fetch_sub(1, std::memory_order_release); if (pending) { - // 检查读取错误 if (cqe->res < 0) { - fprintf(stderr, "[MESH] io_uring read error: expert=%d tp=%d res=%d (%s)\n", + // 读取失败：不调用 on_complete，不标记 CACHED + fprintf(stderr, "[MESH] io_uring read error: expert=%d tp=%d res=%d (%s), " + "slot will NOT be marked cached\n", pending->expert_id, pending->tp_part_idx, cqe->res, strerror(-cqe->res)); - } - int completed = pending->completed_reqs.fetch_add(1) + 1; - if (completed == pending->total_reqs) { - // 所有 SQE 完成，触发回调 - if (pending->on_complete) { - pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + delete pending; + } + } else { + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + // 所有 SQE 完成，触发回调 + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; } - delete pending; } } count++; @@ -344,22 +402,34 @@ class MeshIoUring { // 阻塞等待某专家的读请求完成（通过 CQE 计数） // A6: 改为处理 CQE 并触发回调 + // 根因修复：当 cqe->res < 0 时不调用 on_complete void wait_expert(int expert_id, int n_reqs) { for (int i = 0; i < n_reqs; i++) { io_uring_cqe* cqe; - io_uring_wait_cqe(&ring_, &cqe); + // 处理 EINTR：信号中断时重试 + int wait_ret; + do { + wait_ret = io_uring_wait_cqe(&ring_, &cqe); + } while (wait_ret == -EINTR); auto* pending = static_cast(io_uring_cqe_get_data(cqe)); if (pending) { if (cqe->res < 0) { - fprintf(stderr, "[MESH] io_uring read error in wait_expert: expert=%d res=%d (%s)\n", + // 读取失败：不调用 on_complete，不标记 CACHED + fprintf(stderr, "[MESH] io_uring read error in wait_expert: expert=%d res=%d (%s), " + "slot will NOT be marked cached\n", pending->expert_id, cqe->res, strerror(-cqe->res)); - } - int completed = pending->completed_reqs.fetch_add(1) + 1; - if (completed == pending->total_reqs) { - if (pending->on_complete) { - pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + delete pending; + } + } else { + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; } - delete pending; } } io_uring_cqe_seen(&ring_, cqe); @@ -384,45 +454,75 @@ class MeshIoUring { // 提交并等待全部完成（同步接口，启动阶段用） // A6: 处理 CQE 并触发 on_complete 回调 + // 根因修复：当 cqe->res < 0（读取失败/EINTR）时，不调用 on_complete， + // 不标记 slot 为 CACHED，避免 AMX kernel 读取垃圾数据导致内存损坏。 + // submit_and_wait fix: 用 inflight_count_ 追踪已提交 SQE 数， + // 等待所有 CQE 到达后才返回。旧代码用 io_uring_sq_ready 判断， + // 但 io_uring_submit 后 sq_ready=0（SQE 已提交给内核），导致立即返回， + // on_complete 回调从未被调用，slot 永远停留在 LOADING 状态。 void submit_and_wait() { + std::lock_guard lock(cq_mutex_); io_uring_submit(&ring_); - // 等待所有 inflight 完成，处理回调 - while (true) { + // 等待所有 inflight SQE 完成 + while (inflight_count_.load(std::memory_order_acquire) > 0) { unsigned head; unsigned count = 0; io_uring_cqe* cqe; io_uring_for_each_cqe(&ring_, head, cqe) { auto* pending = static_cast(io_uring_cqe_get_data(cqe)); + // submit_and_wait fix: 每个 CQE 对应一个已完成的 SQE + inflight_count_.fetch_sub(1, std::memory_order_release); if (pending) { if (cqe->res < 0) { - fprintf(stderr, "[MESH] io_uring read error in submit_and_wait: expert=%d res=%d (%s)\n", + // 读取失败（EINTR/IO 错误等）：不调用 on_complete，不标记 CACHED + // 避免垃圾数据被 AMX kernel 消费导致 double free + fprintf(stderr, "[MESH] io_uring read error in submit_and_wait: expert=%d res=%d (%s), " + "slot will NOT be marked cached\n", pending->expert_id, cqe->res, strerror(-cqe->res)); - } - int completed = pending->completed_reqs.fetch_add(1) + 1; - if (completed == pending->total_reqs) { - if (pending->on_complete) { - pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + // 仍然计入完成数（否则 pending 永远不会被释放），但不调用 on_complete + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + // 读取失败，不调用 on_complete，直接释放 pending + delete pending; + } + } else { + int completed = pending->completed_reqs.fetch_add(1) + 1; + if (completed == pending->total_reqs) { + if (pending->on_complete) { + pending->on_complete(pending->layer_idx, pending->tp_part_idx, pending->expert_id); + } + delete pending; } - delete pending; } } count++; } - if (count == 0) { - // 没有更多 CQE，检查是否还有 inflight SQE - unsigned inflight = io_uring_sq_ready(&ring_); - if (inflight == 0) break; - // 还有 inflight，等待一个 CQE + if (count > 0) { + io_uring_cq_advance(&ring_, count); + } + // 还有 inflight，阻塞等待下一个 CQE + if (inflight_count_.load(std::memory_order_acquire) > 0) { io_uring_cqe* wait_cqe; - io_uring_wait_cqe(&ring_, &wait_cqe); - continue; + int wait_ret; + do { + wait_ret = io_uring_wait_cqe(&ring_, &wait_cqe); + } while (wait_ret == -EINTR); } - io_uring_cq_advance(&ring_, count); } } private: io_uring ring_; + // Bug 1 fix: 保护 CQ 访问的互斥锁 — process_cqes 可能在 GEMM 期间被多个 worker 线程并发调用 + std::mutex cq_mutex_; + // SQ ring 互斥锁：保护 io_uring_get_sqe + io_uring_prep_read + io_uring_submit 序列。 + // liburing 的 SQ ring 操作不是线程安全的，多 TP 线程并发 submit_load 会导致 SQ ring 损坏， + // SQE 丢失，CQE 永远不回来，blocking_load_experts 超时。 + std::mutex sq_mutex_; + // submit_and_wait fix: 跟踪已提交但未完成的 SQE 数量 + // io_uring_sq_ready() 在 io_uring_submit 后返回 0（所有 SQE 已提交给内核）， + // 但内核可能尚未完成任何读取。用此计数器追踪 inflight 请求数。 + std::atomic inflight_count_{0}; // Scale Cache（A4 fix: [layer][tp] 每层独立） struct ScaleCacheTP { @@ -447,11 +547,18 @@ class MeshIoUring { // 同步读取（启动阶段用，pread + O_DIRECT） void sync_read_to(int fd, off_t offset, size_t bytes, void* dst) { if (bytes == 0) return; + if (!dst) { + fprintf(stderr, "[MESH] sync_read_to: dst is NULL! fd=%d offset=%ld bytes=%zu — aborting\n", + fd, (long)offset, bytes); + std::abort(); + } size_t done = 0; while (done < bytes) { ssize_t n = pread(fd, (char*)dst + done, bytes - done, offset + done); if (n < 0) { - throw std::runtime_error("MeshIoUring: pread failed, errno=" + std::to_string(errno)); + fprintf(stderr, "[MESH] sync_read_to: pread failed errno=%d (%s) fd=%d offset=%ld bytes=%zu done=%zu — aborting\n", + errno, strerror(errno), fd, (long)offset, bytes, done); + std::abort(); } done += n; } diff --git a/kt-kernel/operators/mesh/mesh_residency.hpp b/kt-kernel/operators/mesh/mesh_residency.hpp index 680986022..8dd3e1589 100644 --- a/kt-kernel/operators/mesh/mesh_residency.hpp +++ b/kt-kernel/operators/mesh/mesh_residency.hpp @@ -15,9 +15,14 @@ #pragma once #include +#include +#include #include #include +#include #include +#include +#include #include #include "mesh_config.hpp" @@ -100,11 +105,26 @@ class MeshResidencyManager { decode_ = std::make_unique(config); handoff_ = std::make_unique(); + // 初始化 per-(layer,tp) mutex 数组 + load_mutexes_.resize(config.total_layers); + for (int l = 0; l < config.total_layers; l++) { + load_mutexes_[l].reserve(config.tp_count); + for (int tp = 0; tp < config.tp_count; tp++) { + load_mutexes_[l].push_back(std::make_unique()); + } + } + // 初始化 temporal 双缓冲 prefill_->init_temporal(); // B4 fix: 注入 gate_up_bytes 用于 temporal_ptr 偏移计算 prefill_->set_gate_up_bytes(compute_gate_up_bytes(config)); + // per-tp defer 队列状态初始化（不同 TP 的 forward_decode 并发调用 on_decode_layer） + prev_layer_deferred_.resize(config.tp_count); + total_deferred_.assign(config.tp_count, 0); + prev_layer_topk_.resize(config.tp_count); + prev_layer_scores_.resize(config.tp_count); + fprintf(stderr, "[MESH init DIAG] init complete, pools_ size=%zu\n", pools_.size()); } @@ -156,73 +176,126 @@ class MeshResidencyManager { // 预加载 Scale Cache（AMXINT4 专用）— 每层独立预加载 if (config_.weight_type == WeightType::AMXINT4) { + fprintf(stderr, "[MESH bootstrap] preload_scale_cache start, layers=%d\n", config_.total_layers); for (int l = 0; l < config_.total_layers; l++) { io_->preload_scale_cache(config_.expert_num, config_.tp_count, numa_nodes_, layouts_[l], l); } + fprintf(stderr, "[MESH bootstrap] preload_scale_cache done\n"); } - // 每层每 TP 读前 cap 个专家 + // pread 重试 lambda（带 EINTR 重试） + auto pread_retry = [](int fd, void* dst, size_t bytes, off_t offset) -> bool { + if (bytes == 0) return true; + size_t done = 0; + while (done < bytes) { + ssize_t n = pread(fd, (char*)dst + done, bytes - done, offset + done); + if (n < 0) { + if (errno == EINTR) continue; + fprintf(stderr, "[MESH] bootstrap pread failed: errno=%d (%s), bytes=%zu offset=%ld\n", + errno, strerror(errno), bytes, (long)offset); + return false; + } + if (n == 0) { + fprintf(stderr, "[MESH] bootstrap pread EOF: bytes=%zu done=%zu offset=%ld\n", + bytes, done, (long)offset); + return false; + } + done += (size_t)n; + } + return true; + }; + + // 每层每 TP 读前 cap 个 CPU 专家 + // SKILL.md 第 39 行：读取所有专家的前 cap 个进入 slot + // Bug fix: 原代码 `for (e = 0; e < cap; e++)` 只遍历 e=0~cap-1，跳过 GPU 专家后 + // 实际加载的 CPU 专家 < cap，ID >= cap 的 CPU 专家从未加载 → hit_rate < 100% + // 修复：遍历所有专家，加载前 cap 个 CPU 专家，slot_idx 从 0 递增 + // slot 存 TP 切分后的权重（intermediate_size / tp_count）， + // 与 AMX_MOE_BASE::config_.intermediate_size（被 TP_MOE_Common 切分）匹配。 + // 每个 pool 只读取自己 tp 对应的文件切片，不合并。 for (int l = 0; l < config_.total_layers; l++) { for (int tp = 0; tp < config_.tp_count; tp++) { MeshSlotPool& pool = pools_[l][tp]; - for (int e = 0; e < config_.cap && e < config_.expert_num; e++) { + int slot_idx = 0; + for (int e = 0; e < config_.expert_num && slot_idx < config_.cap; e++) { if (is_gpu_expert(e)) continue; // GPU expert 跳过 + // 绑定 slot（slot_idx 从 0 递增，expert_id = e） + pool.bind(slot_idx, e); + const auto& layout = layouts_[l][tp][e]; - void* gate_dst = pool.gate_ptr(e); - void* up_dst = pool.up_ptr(e); - void* down_dst = pool.down_ptr(e); - - // 绑定 slot - pool.bind(e, e); - - // 提交 io_uring 读 - // A6: 传入 layer_idx 和 on_complete 回调，完成后 mark_cached - // 辅因1 fix: on_complete 中调用 copy_scale_from_cache 把 scale 写入 slot - io_->submit_load(e, tp, layout, gate_dst, up_dst, down_dst, - nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, - ReadPriority::Demand, - /*layer_idx=*/l, - /*on_complete=*/[this, l, tp, e](int, int, int) { - // Bug 4 fix: 先 copy scale，再 mark_cached - // mark_cached 后状态为 CACHED，AMX kernel 可能立即读取 - // 若先 mark_cached 后 copy scale，AMX 会读到零 scale → NaN - if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(l)) { - // Bug 3 fix: bootstrap 中 bind(e, e) 使 slot_idx == e - int slot_idx = pools_[l][tp].expert_to_slot_idx(e); - if (slot_idx >= 0) { - void* gs = pools_[l][tp].gate_scale_ptr(slot_idx); - void* us = pools_[l][tp].up_scale_ptr(slot_idx); - void* ds = pools_[l][tp].down_scale_ptr(slot_idx); - io_->copy_scale_from_cache(l, tp, e, gs, us, ds, - nullptr, nullptr, nullptr, - layouts_[l][tp]); - } - } - // Bug 3 fix: mark_cached 参数是 slot_idx，bootstrap 中 slot_idx == e - pools_[l][tp].mark_cached(e); - }); - } - } - } - // 等待所有读完成 - io_->submit_and_wait(); + // 读取 gate 权重 + if (!pread_retry(layout.fd, pool.gate_ptr(slot_idx), layout.gate_bytes, layout.gate_offset)) { + fprintf(stderr, "[MESH] bootstrap: gate pread failed expert=%d layer=%d tp=%d\n", e, l, tp); + pool.unbind(slot_idx); // 加载失败，释放 slot + slot_idx--; + continue; + } + // 读取 up 权重 + if (!pread_retry(layout.fd, pool.up_ptr(slot_idx), layout.up_bytes, layout.up_offset)) { + fprintf(stderr, "[MESH] bootstrap: up pread failed expert=%d layer=%d tp=%d\n", e, l, tp); + pool.unbind(slot_idx); + slot_idx--; + continue; + } + // 读取 down 权重 + if (layout.down_stride > 0) { + // BF16 down_proj [E,H,I] 行主序，TP 沿 I 切不连续，逐行读取 + size_t row_bytes = layout.down_bytes / layout.down_rows; + off_t src_off = layout.down_offset; + bool down_ok = true; + for (int r = 0; r < layout.down_rows; r++) { + if (!pread_retry(layout.fd, static_cast(pool.down_ptr(slot_idx)) + r * row_bytes, row_bytes, src_off)) { + fprintf(stderr, "[MESH] bootstrap: down row %d pread failed expert=%d\n", r, e); + down_ok = false; + break; + } + src_off += layout.down_stride; + } + if (!down_ok) { + pool.unbind(slot_idx); + slot_idx--; + continue; + } + } else { + if (!pread_retry(layout.fd, pool.down_ptr(slot_idx), layout.down_bytes, layout.down_offset)) { + fprintf(stderr, "[MESH] bootstrap: down pread failed expert=%d layer=%d tp=%d\n", e, l, tp); + pool.unbind(slot_idx); + slot_idx--; + continue; + } + } + + // AMXINT4: 从 scale cache 拷贝 scale 数据到 slot + // 每个 tp 的 scale 对应自己的权重切片，直接复制 + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(l)) { + io_->copy_scale_from_cache(l, tp, e, + pool.gate_scale_ptr(slot_idx), + pool.up_scale_ptr(slot_idx), + pool.down_scale_ptr(slot_idx), + nullptr, nullptr, nullptr, + layouts_[l][tp]); + } - // 标记所有 slot 为 CACHED - for (int l = 0; l < config_.total_layers; l++) { - for (int tp = 0; tp < config_.tp_count; tp++) { - MeshSlotPool& pool = pools_[l][tp]; - for (int e = 0; e < config_.cap && e < config_.expert_num; e++) { - if (is_gpu_expert(e)) continue; - pool.mark_cached(e); + // 标记 slot 为 CACHED + pool.mark_cached(slot_idx); + + stats_.io_uring_read_count++; + stats_.io_uring_read_bytes += layout.gate_bytes + layout.up_bytes + layout.down_bytes; + slot_idx++; } } } + + fprintf(stderr, "[MESH] bootstrap: preloaded experts per layer per tp (cap=%d)\n", + config_.cap); } // ===== 权重指针查询（KT 计算用）===== + // 以下方法均内联 acquire_reader：返回非 nullptr 时 reader 已持有， + // 调用者 GEMM 完后必须调用 release_reader。 void* get_gate_ptr(int layer, int tp, int expert_id) { if (layer < 0 || layer >= (int)pools_.size()) return nullptr; @@ -242,6 +315,96 @@ class MeshResidencyManager { return pools_[layer][tp].expert_down_ptr(expert_id); } + // 带 acquire_reader 的版本（hook 用）：返回非 nullptr 时 reader 已持有 + // 原子化：lookup + state check + reader + pointer 在同一把 slot_mtx_ 锁内完成， + // 消除 TOCTOU 竞争（指针和 reader 保证在同一个 slot 上） + void* get_gate_ptr_with_reader(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].acquire_gate_ptr(expert_id); + } + + void* get_up_ptr_with_reader(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].acquire_up_ptr(expert_id); + } + + void* get_down_ptr_with_reader(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + return pools_[layer][tp].acquire_down_ptr(expert_id); + } + + // ===== 同步加载（slot 未命中时阻塞加载）===== + // MESH 模式下 gate_bb_ 为空，slot 未命中不能回退，必须同步从 SSD 加载到 slot + // 返回非 nullptr 时 reader 已持有，调用者 GEMM 完后必须调用 release_reader。 + + void* get_or_load_gate_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + + // 快速路径：原子化 acquire（lookup + state + reader + pointer 在同一锁内） + void* ptr = pools_[layer][tp].acquire_gate_ptr(expert_id); + if (ptr) return ptr; + + // 未缓存或 acquire 失败（被驱逐）：同步加载 + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + + // 加载完成后再原子化 acquire（可能又被其他线程驱逐，重试一次） + ptr = pools_[layer][tp].acquire_gate_ptr(expert_id); + if (ptr) return ptr; + + // 极端竞争下被驱逐，再加载一次 + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + ptr = pools_[layer][tp].acquire_gate_ptr(expert_id); + return ptr; + } + + void* get_or_load_up_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + + void* ptr = pools_[layer][tp].acquire_up_ptr(expert_id); + if (ptr) return ptr; + + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + + ptr = pools_[layer][tp].acquire_up_ptr(expert_id); + if (ptr) return ptr; + + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + ptr = pools_[layer][tp].acquire_up_ptr(expert_id); + return ptr; + } + + void* get_or_load_down_ptr(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return nullptr; + if (tp < 0 || tp >= (int)pools_[layer].size()) return nullptr; + + void* ptr = pools_[layer][tp].acquire_down_ptr(expert_id); + if (ptr) return ptr; + + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + + ptr = pools_[layer][tp].acquire_down_ptr(expert_id); + if (ptr) return ptr; + + if (!load_expert_sync(layer, tp, expert_id)) return nullptr; + ptr = pools_[layer][tp].acquire_down_ptr(expert_id); + return ptr; + } + + // ===== 引用计数（防止 AMX 计算期间 slot 被驱逐）===== + // acquire_reader 已内联到 get/get_or_load 方法中，外部无需调用。 + // release_reader 仍由 do_gate_up_gemm/do_down_gemm 在 GEMM 完后调用。 + + void release_reader(int layer, int tp, int expert_id) { + if (layer < 0 || layer >= (int)pools_.size()) return; + if (tp < 0 || tp >= (int)pools_[layer].size()) return; + pools_[layer][tp].release_reader(expert_id); + } + // ===== Prefill 阶段回调 ===== void on_prefill_layer_start(int layer_idx, int qlen, @@ -274,9 +437,12 @@ class MeshResidencyManager { void on_decode_token_start() { scheduler_->inc_timeline_step(); - // B5: 每 token 开始时清空跨层 defer 队列 - prev_layer_deferred_.clear(); - total_deferred_ = 0; + // B5: 每 token 开始时清空所有 tp 的跨层 defer 队列 + // on_decode_token_start 由 Python 侧调用（非并发），安全清空所有 tp + for (int tp = 0; tp < (int)prev_layer_deferred_.size(); tp++) { + prev_layer_deferred_[tp].clear(); + total_deferred_[tp] = 0; + } // B9: 统计 stats_.decode_token_count++; } @@ -298,51 +464,124 @@ class MeshResidencyManager { const std::vector& topk, const std::vector& scores, int tp_part_idx) { + // 检测 double call：KTransformers defer 机制每层调用 GEMM 两次 + // （experts_base.py:427-454 两次 submit_with_cuda_stream） + // 第一次（immediate GEMM）：topk 中 -1 数量 = defer_count_（deferred 位置填 -1） + // 第二次（deferred GEMM）：topk 中 -1 数量 = num_experts_per_tok - defer_count_（immediate 位置填 -1） + // + // SKILL.md 第 286 行：当前层计算只执行 immediate 部分，合并上一层 deferred 输出由 KT 原版负责。 + // 第二次 GEMM 计算 deferred 组专家。 + // + // 根因修复：第一次 on_decode_layer 的 topk 中 deferred 位置是 -1，看不到真实 expert_id， + // 无法为 deferred 专家提交 io_uring 预取。第二次调用时必须 blocking_load_experts， + // 否则 GEMM 会走 load_expert_sync 的同步 pread 路径（SLOT MISS + 串行阻塞读）。 + // blocking_load_experts 会提交 io_uring（并行 I/O）并 spin-wait 直到 CACHED， + // 比让 GEMM 逐个 pread（串行同步 I/O）快得多。 + if (decode_->defer_count() > 0) { + int neg_one_count = 0; + for (int eid : topk) { + if (eid < 0) neg_one_count++; + } + if (neg_one_count > decode_->defer_count()) { + // 第二次调用（deferred GEMM）：处理 deferred 组专家 + io_->process_cqes(); + + // 提取 deferred 专家（topk 中非 -1 的） + std::vector deferred_experts; + deferred_experts.reserve(topk.size() - neg_one_count); + for (int eid : topk) { + if (eid >= 0) deferred_experts.push_back(eid); + } + + // 对未 CACHED 的 deferred 专家做 blocking_load（提交 io_uring + spin-wait） + // 避免后续 GEMM 走 load_expert_sync 的同步 pread 路径 + if (!deferred_experts.empty()) { + blocking_load_experts(layer_idx, tp_part_idx, deferred_experts); + } + + return MeshDecode::SplitResult(); + } + } + // A7: 先处理已完成的 io_uring CQE，触发 mark_cached 回调 + auto _t0 = std::chrono::steady_clock::now(); io_->process_cqes(); + auto _t1 = std::chrono::steady_clock::now(); + + // on_decode_token_start 未被 Python 侧调用，在 layer_idx==0 时自动执行 + // CAS 确保每 token 只有一个 TP 执行 inc_timeline_step + if (layer_idx == 0) { + int expected = decode_token_seq_.load(std::memory_order_acquire); + if (decode_token_seq_.compare_exchange_strong(expected, expected + 1, + std::memory_order_acq_rel)) { + scheduler_->inc_timeline_step(); + stats_.decode_token_count++; + } + // 每 token 第一层清空自己的 total_deferred_（无竞争，每个 TP 只操作自己的） + total_deferred_[tp_part_idx] = 0; + } - // B5: 把上一层的 deferred 专家加入当前层的 immediate 候选 - // 它们应该已经通过 io_uring 读取完成 - std::vector effective_topk = topk; - if (!prev_layer_deferred_.empty()) { - // 检查 deferred 专家是否已 CACHED - MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; - for (int eid : prev_layer_deferred_) { - if (pool.is_cached(eid)) { - // 已缓存，加入 immediate 候选 - if (std::find(effective_topk.begin(), effective_topk.end(), eid) == effective_topk.end()) { - effective_topk.push_back(eid); - } + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; + + // B5: 上一层的 deferred 专家在当前层必须被计算处理完毕（SKILL.md 第 100 行） + // 已 cached 的：KTransformers 原版 defer 机制会在当前层合并它们的输出，MESH 不需要管 + // 未 cached 的：需要阻塞读，确保它们在当前层被计算 + std::vector prev_deferred_blocking; + if (!prev_layer_deferred_[tp_part_idx].empty()) { + for (int eid : prev_layer_deferred_[tp_part_idx]) { + if (!pool.is_cached(eid) && !is_gpu_expert(eid)) { + prev_deferred_blocking.push_back(eid); } - // 未缓存的 deferred 专家在本层继续等待（会在 split 中被当作 missing 处理） } - // 清空上一层的 defer 队列 - prev_layer_deferred_.clear(); + prev_layer_deferred_[tp_part_idx].clear(); } - MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; - auto result = decode_->split(effective_topk, scores, pool, *scheduler_, layer_idx, tp_part_idx, - [this, &pool, layer_idx, tp_part_idx](int eid) { - // A5 fix: 用 expert_*_ptr 按 expert_id 查 slot， - // 不能把 expert_id 当 slot_idx 传给 gate_ptr（eid>=cap 时越界） - return std::vector{ - pool.expert_gate_ptr(eid), - pool.expert_up_ptr(eid), - pool.expert_down_ptr(eid)}; - }, - // B6: 传入 is_gpu_expert 回调，过滤 GPU 专家 + // SKILL.md 第 280 行：split 只对原始 topk 分组 + // target_immediate = k - defer_count，deferred 最多 defer_count 个 + // 不把上一层的 deferred 加入 effective_topk（否则 deferred 指数级增长） + auto result = decode_->split(topk, scores, pool, layer_idx, + static_cast(topk.size()), [this](int eid) { return is_gpu_expert(eid); }); + auto _t2 = std::chrono::steady_clock::now(); + + // 上一层的 deferred 未 cached 的加入 blocking_missing（阻塞读） + result.blocking_missing.insert(result.blocking_missing.end(), + prev_deferred_blocking.begin(), + prev_deferred_blocking.end()); + + // SKILL.md 第 282 行：blocking_missing 阻塞读补充 immediate 组 + // 为 blocking_missing 分配 slot + 提交 io_uring + 阻塞等待完成 + blocking_load_experts(layer_idx, tp_part_idx, result.blocking_missing); + auto _t3 = std::chrono::steady_clock::now(); + // 阻塞读完成后，只把 prev_deferred_blocking 加入 immediate 组 + // （split 返回的 blocking_missing 已在 immediate 中，不重复加入） + result.immediate.insert(result.immediate.end(), + prev_deferred_blocking.begin(), + prev_deferred_blocking.end()); + + // SKILL.md 第 284 行：deferred_missing 异步 io_uring 读取，推迟到下一层 + for (int eid : result.deferred_missing) { + // 获取 slot 指针（可能为 nullptr，drain_and_submit 会分配 slot） + auto ptrs = std::vector{ + pool.expert_gate_ptr(eid), + pool.expert_up_ptr(eid), + pool.expert_down_ptr(eid)}; + scheduler_->submit_decode_deferred(layer_idx, eid, tp_part_idx, + ptrs[0], ptrs[1], ptrs[2]); + } - // B5: 更新跨层 defer 队列 - prev_layer_deferred_ = result.deferred; - total_deferred_ += static_cast(result.deferred.size()); + // B5: 更新跨层 defer 队列（只包含当前层的 deferred，最多 defer_count_ 个） + prev_layer_deferred_[tp_part_idx] = result.deferred; + total_deferred_[tp_part_idx] += static_cast(result.deferred.size()); // B9: 统计 stats_.decode_immediate_count += static_cast(result.immediate.size()); stats_.decode_deferred_count += static_cast(result.deferred.size()); stats_.defer_count += static_cast(result.deferred.size()); - // hit/miss 统计：effective_topk 中已缓存的是 hit，未缓存的是 miss - for (int eid : effective_topk) { + // hit/miss 统计：只对原始 topk（不含上一层 deferred） + // 跳过 -1（KTransformers defer 机制产生的无效位置，不是真实 miss） + for (int eid : topk) { + if (eid < 0) continue; // 跳过 -1 if (is_gpu_expert(eid)) continue; // GPU 专家不计入 hit/miss if (pool.is_cached(eid)) { stats_.cache_hit_count++; @@ -351,29 +590,338 @@ class MeshResidencyManager { } } - // B5: overflow 检查 — 如果 defer 累积超过阈值，阻塞等待 io_uring 完成 - if (total_deferred_ > config_.max_deferred_per_token * config_.total_layers) { - // B9: 统计 + // B5: overflow 检查 — 如果 defer 累积超过阈值，处理已完成的 CQE 并重置计数 + // 不用 submit_and_wait（持锁阻塞等 CQE，大量 inflight 时死锁）。 + // 引用计数保护下，deferred_missing 提交的 io_uring 会在后续层 process_cqes 时完成， + // 不需要在这里阻塞等待。 + if (total_deferred_[tp_part_idx] > config_.max_deferred_per_token * config_.total_layers) { stats_.defer_overflow_count++; - // 阻塞等待所有 inflight io_uring 完成 - io_->submit_and_wait(); - io_->process_cqes(); - total_deferred_ = 0; // 重置计数 + io_->process_cqes(); // 只处理已完成的 CQE，不阻塞等待 + total_deferred_[tp_part_idx] = 0; // 重置计数 } - // A7: 排空调度器队列，提交 io_uring 读取 - drain_and_submit(); - - // B2: 用当前层的实际路由（原始 topk + scores）预测下一层的 cross_layer_prior - // 放在 drain_and_submit 之后：当前层 IO 已提交，为下一层驱逐评分做准备 - // 注意：用原始 topk/scores，不用 effective_topk（后者混入了上一层 deferred，不是真实路由） - scorer_->predict_next_layer(layer_idx, topk, scores); + // A7: 排空当前 TP 的调度器队列，提交 io_uring 读取 + // SKILL.md 第 54 行：只处理本 TP 的请求，不跨 NUMA 操作 + drain_and_submit(tp_part_idx); + auto _t4 = std::chrono::steady_clock::now(); + + // SKILL.md 第 135 行：满驻留自动跳过 + // 当 slot 容量 ≥ CPU 专家数时，驱逐不可能发生，Heat/Markov 信号无价值 + int cpu_expert_num = config_.expert_num - config_.num_gpu_experts; + if (!scorer_->should_skip_layer(pool.cap(), cpu_expert_num)) { + // 方案 B: 只 tp_part_idx==0 更新 scorer，避免 NUMA 间 mtx_ 锁竞争 + // 安全性: TP_MOE_Common 按 intermediate_size 切分，各 NUMA 看到的 topk/scores 相同， + // 所以 NUMA 1 重复提交是冗余的，跳过后结果完全一致。 + if (tp_part_idx == 0) { + scorer_->commit_layer(layer_idx, topk, scores, prev_layer_topk_[tp_part_idx], prev_layer_scores_[tp_part_idx]); + // B2: 用当前层的实际路由预测下一层的 cross_layer_prior + scorer_->predict_next_layer(layer_idx, topk, scores); + } + prev_layer_topk_[tp_part_idx] = topk; + prev_layer_scores_[tp_part_idx] = scores; + } + auto _t5 = std::chrono::steady_clock::now(); + + // 时间统计（每 200 层打印一次） + { + static thread_local uint64_t stat_n = 0; + static thread_local double stat_total = 0, stat_cqes = 0, stat_split = 0; + static thread_local double stat_block = 0, stat_drain = 0, stat_scorer = 0; + stat_n++; + stat_total += std::chrono::duration(_t5 - _t0).count(); + stat_cqes += std::chrono::duration(_t1 - _t0).count(); + stat_split += std::chrono::duration(_t2 - _t1).count(); + stat_block += std::chrono::duration(_t3 - _t2).count(); + stat_drain += std::chrono::duration(_t4 - _t3).count(); + stat_scorer+= std::chrono::duration(_t5 - _t4).count(); + if (stat_n % 200 == 0) { + fprintf(stderr, "[MESH TIME] L%d tp%d avg: total=%.3fms cqes=%.3fms split=%.3fms " + "block=%.3fms drain=%.3fms scorer=%.3fms (n=%lu)\n", + layer_idx, tp_part_idx, + stat_total / stat_n, stat_cqes / stat_n, stat_split / stat_n, + stat_block / stat_n, stat_drain / stat_n, stat_scorer / stat_n, + (unsigned long)stat_n); + } + } return result; } /** - * @brief A7: 排空调度器队列并提交 io_uring 读取 + * @brief SKILL.md 第 282 行：阻塞读 missing 专家补充 immediate 组 + * + * 为每个 missing 专家分配 slot + 提交 io_uring + 阻塞等待完成。 + * 完成后专家状态为 CACHED，可加入 immediate 组。 + * + * 引用计数保护：overwrite CV 等待 active_readers 归零， + * GEMM 时 acquire_reader 持有的专家不会被驱逐（SKILL.md 第 62/125 行）。 + */ + void blocking_load_experts(int layer_idx, int tp_part_idx, + const std::vector& experts) { + if (experts.empty()) return; + if (layouts_.empty()) return; + + MeshSlotPool& pool = pools_[layer_idx][tp_part_idx]; + + int submitted = 0; + int skipped = 0; + bool needs_wait = false; // 是否有专家需要 spin-wait（未 CACHED 的） + // fix4: 记录需要 spin-wait 的专家及其 slot，供 stale LOADING 重试使用 + std::vector> wait_list; // (eid, slot_idx) + for (int eid : experts) { + if (pool.is_cached(eid)) continue; // 已缓存，跳过 + if (is_gpu_expert(eid)) continue; // GPU 专家，跳过 + needs_wait = true; // 至少有一个专家未 CACHED，需要 spin-wait + + // 分配 slot + int slot_idx = pool.expert_to_slot_idx(eid); + bool just_bound = false; // fix4: 区分"本次刚 bind"和"之前调用遗留的 LOADING" + if (slot_idx < 0) { + slot_idx = pool.find_free_slot(); + if (slot_idx >= 0) { + // fix5: 空闲 slot，需要 bind + pool.bind(slot_idx, eid); // bind() 会设置 state=LOADING + just_bound = true; + } else { + // fix5: 没有空闲 slot，需要驱逐 + // evict_for_new_expert 调用 overwrite，overwrite 内部已持锁完成 bind + // （设置 state=LOADING, bound_expert_id, expert_to_slot_） + // 不能再调用 pool.bind，否则不持锁重复设置可能与 overwrite 的持锁操作竞争 + slot_idx = decode_->evict_for_new_expert(pool, *scorer_, layer_idx, eid); + if (slot_idx >= 0) { + stats_.eviction_count++; + just_bound = true; + } + } + } + + if (slot_idx < 0) { + fprintf(stderr, "[MESH] blocking_load: slot alloc FAILED eid=%d layer=%d tp=%d\n", + eid, layer_idx, tp_part_idx); + skipped++; + continue; // 无法分配 slot，跳过 + } + + // fix4: 防止 double-submission，但区分两种 LOADING： + // 1. just_bound=true：本次调用刚 bind，必须提交 io_uring（旧 fix3b 的 bug 就是这里误跳过） + // 2. just_bound=false：之前调用遗留的 LOADING（io_uring 可能在飞行中） + // - inflight>0：有 io_uring 在飞行中，跳过提交，spin-wait 会处理 + // - inflight==0：io_uring 丢失（提交失败/CQE 错误），重置并重新提交 + ExpertState state = pool.slot_state(slot_idx); + if (state == ExpertState::LOADING && !just_bound) { + if (io_->inflight_count() > 0) { + // io_uring 可能在飞行中，跳过提交，spin-wait 会处理 + wait_list.emplace_back(eid, slot_idx); + continue; + } + // inflight==0 但 LOADING：io_uring 丢失，重置并重新提交 + fprintf(stderr, "[MESH] blocking_load: stale LOADING (inflight=0), resubmit eid=%d slot=%d layer=%d tp=%d\n", + eid, slot_idx, layer_idx, tp_part_idx); + pool.unbind(slot_idx); + pool.bind(slot_idx, eid); + // 继续执行下面的 io_uring 提交逻辑 + } + + // 获取 slot 指针 + void* gate_dst = pool.gate_ptr(slot_idx); + void* up_dst = pool.up_ptr(slot_idx); + void* down_dst = pool.down_ptr(slot_idx); + + const auto& layout = layouts_[layer_idx][tp_part_idx][eid]; + + stats_.io_uring_read_count++; + stats_.io_uring_read_bytes += static_cast( + layout.gate_bytes + layout.up_bytes + layout.down_bytes); + + io_->submit_load( + eid, tp_part_idx, layout, + gate_dst, up_dst, down_dst, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, // scale 从 cache memcpy + ReadPriority::Demand, + /*layer_idx=*/layer_idx, + /*on_complete=*/[this, layer_idx, tp_part_idx, eid, slot_idx](int, int, int) { + // Bug 4 fix: 先 copy scale，再 mark_cached + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(layer_idx)) { + void* gs = pools_[layer_idx][tp_part_idx].gate_scale_ptr(slot_idx); + void* us = pools_[layer_idx][tp_part_idx].up_scale_ptr(slot_idx); + void* ds = pools_[layer_idx][tp_part_idx].down_scale_ptr(slot_idx); + io_->copy_scale_from_cache(layer_idx, tp_part_idx, eid, gs, us, ds, + nullptr, nullptr, nullptr, + layouts_[layer_idx][tp_part_idx]); + } + pools_[layer_idx][tp_part_idx].mark_cached(slot_idx); + }); + submitted++; + wait_list.emplace_back(eid, slot_idx); // fix4: 记录已提交的专家 + } + + if (!needs_wait) return; // 所有专家已 CACHED/GPU，无需等待 + if (submitted == 0 && skipped > 0 && wait_list.empty()) { + // 所有未 CACHED 的专家都分配 slot 失败，且没有 LOADING 中的专家 + // 无法 spin-wait，让 GEMM 的 load_expert_sync 处理 + fprintf(stderr, "[MESH] blocking_load: all %d experts skipped (no slot) layer=%d tp=%d\n", + skipped, layer_idx, tp_part_idx); + return; + } + // 注意：如果 submitted==0 但 needs_wait==true 且 wait_list 非空， + // 说明所有未 CACHED 的专家都在 LOADING 状态（io_uring 可能在飞行中）， + // 必须继续 spin-wait，不能提前返回 + + // SKILL.md 第 282 行：阻塞等待 blocking_missing 的 io_uring 完成 + // 不用 submit_and_wait（会等待所有 inflight，包括之前层 deferred_missing 提交的 io_uring， + // 如果那些读取慢会死锁）。改用 process_cqes + spin-wait 等 blocking_missing CACHED。 + // 加超时检查（120秒），超时打印警告并返回。 + // fix4: 加 stale LOADING 检测 — 等待 >5s 后，如果仍有专家 LOADING 且 inflight==0， + // 说明 io_uring 丢失，重置 slot 并重新提交 io_uring + auto t0 = std::chrono::steady_clock::now(); + int resubmit_count = 0; + while (true) { + // 处理已完成的 CQE，触发 mark_cached 回调 + io_->process_cqes(); + + // 检查所有 blocking_missing 是否已 CACHED + bool all_cached = true; + for (int eid : experts) { + if (is_gpu_expert(eid)) continue; + if (!pool.is_cached(eid)) { + all_cached = false; + break; + } + } + if (all_cached) break; + + // fix6: 检测所有未 CACHED 专家的 slot 状态（不只是 wait_list） + // fix5 只检测 wait_list，漏掉了"提交时已 CACHED 但 spin-wait 期间被其他线程 evict"的专家 + // 这些专家被 evict 后 is_cached 返回 false，all_cached=false，但 fix5 不检测它们 + // 导致 spin-wait 永远等待已不存在的专家 → 120s timeout + // fix6: 遍历 experts（输入参数），检测所有未 CACHED 专家的 slot 解绑/覆盖 + for (int eid : experts) { + if (is_gpu_expert(eid)) continue; + if (pool.is_cached(eid)) continue; + int cur_slot = pool.expert_to_slot_idx(eid); + if (cur_slot < 0) { + // slot 被解绑（被 evict_for_new_expert 的 unbind/erase 清除） + // 让 GEMM 的 load_expert_sync（持有 load_mutexes_ 锁）重新加载 + fprintf(stderr, "[MESH] blocking_load: slot unbound during spin-wait, eid=%d " + "layer=%d tp=%d inflight=%d — was cached at submit, evicted by concurrent thread, " + "returning to let GEMM load_expert_sync handle\n", + eid, layer_idx, tp_part_idx, io_->inflight_count()); + return; + } + } + + // 超时检查（120秒） + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - t0).count(); + if (elapsed > 120) { + // 诊断：打印未 CACHED 的专家和 inflight 状态 + fprintf(stderr, "[MESH] blocking_load_experts: timeout layer=%d tp=%d elapsed=%lds " + "submitted=%d skipped=%d inflight=%d\n", + layer_idx, tp_part_idx, (long)elapsed, submitted, skipped, + io_->inflight_count()); + for (int eid : experts) { + if (is_gpu_expert(eid)) continue; + if (!pool.is_cached(eid)) { + int sidx = pool.expert_to_slot_idx(eid); + fprintf(stderr, "[MESH] NOT CACHED: eid=%d slot=%d state=%d\n", + eid, sidx, sidx >= 0 ? (int)pool.slot_state(sidx) : -1); + } + } + break; + } + + // fix4: stale LOADING 检测 — 等待 >5s 后，如果仍有专家 LOADING 且 inflight==0， + // 说明 io_uring 丢失（提交失败/CQE 错误/被其他线程误处理） + // 重置 slot 状态并重新提交 io_uring + if (elapsed > 5 && io_->inflight_count() == 0 && resubmit_count < 3) { + for (auto& wl : wait_list) { + int eid = wl.first; + int sidx = wl.second; + if (is_gpu_expert(eid)) continue; + if (pool.is_cached(eid)) continue; + // fix5: 再次检查 slot 是否仍绑定到 eid（防御性，修改2已检测但中间可能有变化） + int cur_slot = pool.expert_to_slot_idx(eid); + if (cur_slot != sidx) { + fprintf(stderr, "[MESH] blocking_load: slot reassigned before stale LOADING resubmit, " + "eid=%d orig_slot=%d cur_slot=%d layer=%d tp=%d — skipping\n", + eid, sidx, cur_slot, layer_idx, tp_part_idx); + continue; + } + ExpertState st = pool.slot_state(sidx); + if (st == ExpertState::LOADING) { + fprintf(stderr, "[MESH] blocking_load: stale LOADING after %lds, resubmit eid=%d slot=%d " + "layer=%d tp=%d (resubmit #%d)\n", + (long)elapsed, eid, sidx, layer_idx, tp_part_idx, resubmit_count + 1); + pool.unbind(sidx); + pool.bind(sidx, eid); + const auto& layout = layouts_[layer_idx][tp_part_idx][eid]; + void* gate_dst = pool.gate_ptr(sidx); + void* up_dst = pool.up_ptr(sidx); + void* down_dst = pool.down_ptr(sidx); + stats_.io_uring_read_count++; + stats_.io_uring_read_bytes += static_cast( + layout.gate_bytes + layout.up_bytes + layout.down_bytes); + io_->submit_load( + eid, tp_part_idx, layout, + gate_dst, up_dst, down_dst, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + ReadPriority::Demand, + layer_idx, + [this, layer_idx, tp_part_idx, eid, sidx](int, int, int) { + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(layer_idx)) { + void* gs = pools_[layer_idx][tp_part_idx].gate_scale_ptr(sidx); + void* us = pools_[layer_idx][tp_part_idx].up_scale_ptr(sidx); + void* ds = pools_[layer_idx][tp_part_idx].down_scale_ptr(sidx); + io_->copy_scale_from_cache(layer_idx, tp_part_idx, eid, gs, us, ds, + nullptr, nullptr, nullptr, + layouts_[layer_idx][tp_part_idx]); + } + pools_[layer_idx][tp_part_idx].mark_cached(sidx); + }); + submitted++; + resubmit_count++; + } + } + } + + // 短暂 sleep，避免 100% CPU + usleep(100); // 100us + } + } + + // Bug 1+2 fix: 简化版 on_decode_layer（供 hook 调用，接收原始指针参数） + // 在 forward_decode 的 GEMM 之前调用，更新 Heat/Markov 使驱逐策略生效 + // 注意：异步 io_uring 预取暂时禁用（CQE 在 GEMM 期间无法可靠处理）， + // 改由 load_expert_sync 的并行 pread（Bug 3 fix）处理 missing 专家 + void on_decode_layer_simple(int layer_idx, int tp_part_idx, + const int* topk_ptr, int k, + const float* weights, int expert_num) { + if (!config_.enabled) return; + + // 构造 topk 和 scores（全长，top-k 位置填入 weights，其余 0） + std::vector topk(topk_ptr, topk_ptr + k); + std::vector scores(expert_num, 0.0f); + for (int i = 0; i < k; i++) { + int eid = topk_ptr[i]; + if (eid >= 0 && eid < expert_num) { + scores[eid] = weights[i]; + } + } + + // 方案 B: 只 tp_part_idx==0 更新 scorer，避免 NUMA 间 mtx_ 锁竞争 + if (tp_part_idx == 0) { + scorer_->commit_layer(layer_idx, topk, scores, prev_layer_topk_[tp_part_idx], prev_layer_scores_[tp_part_idx]); + scorer_->predict_next_layer(layer_idx, topk, scores); + } + prev_layer_topk_[tp_part_idx] = topk; + prev_layer_scores_[tp_part_idx] = scores; + } + + /** + * @brief A7: 排空当前 TP 的调度器队列并提交 io_uring 读取 + * + * SKILL.md 第 54 行：驱逐和覆盖只影响本 TP 的本地 shard，不跨 NUMA 操作 + * 只取走 tp_part_idx 对应的请求，其余 TP 的请求留在队列中由各自 TP 处理 * * 把 MeshScheduler 优先队列中的 ScheduledRequest 取出， * 为每个请求找到对应的 ExpertFileLayout 并提交给 MeshIoUring。 @@ -381,8 +929,13 @@ class MeshResidencyManager { * * B1: 如果专家未缓存且 slot 池满，先驱逐一个 victim 再 bind */ - void drain_and_submit() { - auto requests = scheduler_->drain_all(); + void drain_and_submit(int tp_part_idx = -1) { + std::vector requests; + if (tp_part_idx >= 0) { + requests = scheduler_->drain_all_for_tp(tp_part_idx); + } else { + requests = scheduler_->drain_all(); + } for (auto& req : requests) { // 边界检查 if (req.layer_idx < 0 || req.layer_idx >= (int)layouts_.size()) continue; @@ -417,6 +970,7 @@ class MeshResidencyManager { if (slot_idx < 0) { // 没有空闲 slot，需要驱逐 + // 引用计数保护：overwrite CV 等待 reader 归零，GEMM 中的专家不会被驱逐 slot_idx = decode_->evict_for_new_expert(pool, *scorer_, layer, expert); if (slot_idx >= 0) { // B9: 统计驱逐 @@ -513,9 +1067,24 @@ class MeshResidencyManager { // GPU expert mask std::vector gpu_experts_mask_; - // B5: 跨层 defer 队列状态 - std::vector prev_layer_deferred_; // 上一层的 deferred 专家 - int total_deferred_ = 0; // 跨层累积 defer 计数 + // B5: 跨层 defer 队列状态 — per-tp，因为不同 TP 的 forward_decode 并发调用 on_decode_layer + // do_numa_job 并发执行 tps[numa_id]->forward()，不同 TP 同时操作这些变量会导致 + // std::vector 的并发 clear()/operator=/iterate → use-after-free/堆损坏/段错误 + std::vector> prev_layer_deferred_; // [tp] 上一层的 deferred 专家 + std::vector total_deferred_; // [tp] 跨层累积 defer 计数 + + // Bug 2 fix: 逐层 Heat/Markov 更新所需的上一层 topk/scores — per-tp + std::vector> prev_layer_topk_; // [tp] + std::vector> prev_layer_scores_; // [tp] + + // on_decode_token_start 未被 Python 侧调用，改在 on_decode_layer layer_idx==0 时自动执行 + // 用 atomic + CAS 确保每 token 只有一个 TP 执行 inc_timeline_step + std::atomic decode_token_seq_{0}; // 已处理的 decode token 序号 + + // 并发保护：per-(layer,tp) mutex，保护 load_expert_sync 的 slot 分配和加载 + // 防止多个线程同时为同一/不同 expert 分配 slot 导致竞争 + // 用 unique_ptr 因为 std::mutex 不可移动 + std::vector>> load_mutexes_; // B9: 运行时统计 MeshStats stats_; @@ -528,13 +1097,185 @@ class MeshResidencyManager { std::unique_ptr decode_; std::unique_ptr handoff_; + // ===== 同步加载专家到 slot（slot 未命中时阻塞调用）===== + // 返回 true 表示加载成功，false 表示失败（layouts 未注入或无法分配 slot） + // 根因修复：用 pread 直接读取，替代 io_uring。 + // io_uring 的 CQE 在 EINTR 时 res<0，但旧代码仍调用 on_complete 标记 slot 为 CACHED， + // 导致 slot 内是垃圾数据，AMX kernel 读取后内存损坏 → double free。 + // Bug 3 fix: 减小锁粒度 — pread 移出锁外，允许同层同 TP 并行 I/O + // 锁内只做 slot 分配（find_free_slot / evict + bind/overwrite），锁外做 pread。 + // 同一 expert 的并发线程：发现 LOADING 状态后释放锁，spin-wait 等 CACHED。 + bool load_expert_sync(int layer, int tp, int expert_id) { + if (is_gpu_expert(expert_id)) return false; + if (layouts_.empty()) return false; + if (layer < 0 || layer >= (int)layouts_.size()) return false; + if (tp < 0 || tp >= (int)layouts_[layer].size()) return false; + if (expert_id < 0 || expert_id >= (int)layouts_[layer][tp].size()) return false; + + MeshSlotPool& pool = pools_[layer][tp]; + + // 快速路径：已缓存直接返回（无需加锁） + if (pool.is_cached(expert_id)) return true; + + // 外层 retry loop：处理 slot 被 evict 后需要重新分配的情况 + // Race condition: Thread A sets LOADING + pread, Thread B spin-waits, + // Thread C evicts the slot (overwrite) after pread completes, + // Thread B never sees CACHED → 必须检测并 retry + auto t0 = std::chrono::steady_clock::now(); + for (int outer_retry = 0; ; outer_retry++) { + if (pool.is_cached(expert_id)) return true; + + // 超时检查（120秒）— 避免超过 NCCL watchdog 的 600 秒超时 + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - t0).count(); + if (elapsed > 120) { + fprintf(stderr, "[MESH] load_expert_sync: timeout expert=%d layer=%d tp=%d elapsed=%lds\n", + expert_id, layer, tp, (long)elapsed); + return false; + } + + int slot_idx = -1; + bool wait_for_loading = false; + + // 内层 retry loop: 当所有 slot 都在 LOADING 时，evict_for_new_expert 找不到 + // CACHED slot 来驱逐。等待某个 pread 完成后 slot 变为 CACHED，再重试。 + for (int alloc_retry = 0; alloc_retry < 100000; alloc_retry++) { + if (pool.is_cached(expert_id)) return true; + { + std::lock_guard lock(*load_mutexes_[layer][tp]); + if (pool.is_cached(expert_id)) return true; + + slot_idx = pool.expert_to_slot_idx(expert_id); + if (slot_idx >= 0) { + ExpertState state = pool.slot_state(slot_idx); + if (state == ExpertState::LOADING) { + wait_for_loading = true; + } + break; // 跳出 alloc retry loop + } else { + slot_idx = pool.find_free_slot(); + if (slot_idx >= 0) { + pool.bind(slot_idx, expert_id); // 设置 LOADING + break; + } else { + slot_idx = decode_->evict_for_new_expert(pool, *scorer_, layer, expert_id); + if (slot_idx >= 0) { + break; + } + // slot_idx < 0: 所有 slot 都在 LOADING，无法驱逐 + } + } + } + // 锁已释放 — spin-wait 等待 slot 释放 + if ((alloc_retry & 63) == 0) io_->process_cqes(); + if ((alloc_retry % 1000) == 0) usleep(1); + sched_yield(); + } + + if (wait_for_loading) { + // 等待其他线程完成 pread 加载 + // 关键修复：检测 slot 是否被 evict（overwrite），如果是则 retry + while (true) { + if (pool.is_cached(expert_id)) return true; + // 检查 expert 是否仍绑定在 LOADING 状态的 slot 上 + int sidx = pool.expert_to_slot_idx(expert_id); + if (sidx < 0) break; // expert 不再绑定任何 slot — 被 evict，retry + ExpertState st = pool.slot_state(sidx); + if (st != ExpertState::LOADING) break; // 状态变了（CACHED 或被 evict），retry + io_->process_cqes(); + usleep(100); // 100us + } + continue; // 外层 retry + } + + if (slot_idx < 0) continue; // alloc 失败，外层 retry + + // 锁已释放 — 以下 pread 在锁外执行，允许同层同 TP 的其他 expert 并行 I/O + const auto& layout = layouts_[layer][tp][expert_id]; + void* gate_dst = pool.gate_ptr(slot_idx); + void* up_dst = pool.up_ptr(slot_idx); + void* down_dst = pool.down_ptr(slot_idx); + + stats_.io_uring_read_count++; + stats_.io_uring_read_bytes += layout.gate_bytes + layout.up_bytes + layout.down_bytes; + + // 用 pread 同步读取权重到 slot（带 EINTR 重试） + auto pread_retry = [](int fd, void* dst, size_t bytes, off_t offset) -> bool { + if (bytes == 0) return true; + size_t done = 0; + while (done < bytes) { + ssize_t n = pread(fd, (char*)dst + done, bytes - done, offset + done); + if (n < 0) { + if (errno == EINTR) continue; + fprintf(stderr, "[MESH] pread failed: errno=%d (%s), bytes=%zu offset=%ld\n", + errno, strerror(errno), bytes, (long)offset); + return false; + } + if (n == 0) { + fprintf(stderr, "[MESH] pread unexpected EOF: bytes=%zu done=%zu offset=%ld\n", + bytes, done, (long)offset); + return false; + } + done += (size_t)n; + } + return true; + }; + + // 读取 gate 权重（单个 tp 切片） + if (!pread_retry(layout.fd, gate_dst, layout.gate_bytes, layout.gate_offset)) { + fprintf(stderr, "[MESH] load_expert_sync: gate pread failed, expert=%d layer=%d tp=%d\n", + expert_id, layer, tp); + return false; + } + // 读取 up 权重 + if (!pread_retry(layout.fd, up_dst, layout.up_bytes, layout.up_offset)) { + fprintf(stderr, "[MESH] load_expert_sync: up pread failed, expert=%d layer=%d tp=%d\n", + expert_id, layer, tp); + return false; + } + // 读取 down 权重 + if (layout.down_stride > 0) { + size_t row_bytes = layout.down_bytes / layout.down_rows; + off_t src_off = layout.down_offset; + for (int r = 0; r < layout.down_rows; r++) { + if (!pread_retry(layout.fd, static_cast(down_dst) + r * row_bytes, row_bytes, src_off)) { + fprintf(stderr, "[MESH] load_expert_sync: down row %d pread failed, expert=%d\n", + r, expert_id); + return false; + } + src_off += layout.down_stride; + } + } else { + if (!pread_retry(layout.fd, down_dst, layout.down_bytes, layout.down_offset)) { + fprintf(stderr, "[MESH] load_expert_sync: down pread failed, expert=%d layer=%d tp=%d\n", + expert_id, layer, tp); + return false; + } + } + + // AMXINT4: 从 scale cache 拷贝 scale 数据到 slot（单个 tp 切片） + if (config_.weight_type == WeightType::AMXINT4 && io_->scale_cache_loaded(layer)) { + io_->copy_scale_from_cache(layer, tp, expert_id, + pool.gate_scale_ptr(slot_idx), + pool.up_scale_ptr(slot_idx), + pool.down_scale_ptr(slot_idx), + nullptr, nullptr, nullptr, + layouts_[layer][tp]); + } + + // 所有读取成功，标记 slot 为 CACHED（原子操作，无需锁） + pool.mark_cached(slot_idx); + return true; + } // 外层 for loop + return false; // unreachable + } + // 计算 slot 字节数（gate + up + down 三个矩阵，AMXINT4 含 scale） + // slot 存 TP 切分后的权重（intermediate_size / tp_count）， + // 与 AMX_MOE_BASE::config_.intermediate_size（被 TP_MOE_Common 切分）匹配。 + // GEMM 的 BufferB 用 n=intermediate_size/tp 计算 scale 偏移 = n*k/2， + // slot 的 gate_up_weights_bytes 必须等于这个值，否则 scale 指针错位。 size_t compute_slot_bytes(const MeshConfig& config) const { - // 根据 SKILL 第 8 节的双 NUMA 拆分： - // gate: [hidden, intermediate/tp_count] BufferB n=intermediate/tp, k=hidden - // up: [hidden, intermediate/tp_count] BufferB n=intermediate/tp, k=hidden - // down: [intermediate/tp_count, hidden] BufferB n=hidden, k=intermediate/tp - // AMXINT4 BufferB 期望: n*k/2 (权重) + n*sizeof(float) (scale) 连续存放 size_t gate_up_bytes = compute_gate_up_bytes(config); size_t down_bytes = compute_down_bytes(config); return gate_up_bytes * 2 + down_bytes; @@ -542,9 +1283,10 @@ class MeshResidencyManager { // gate 或 up 单块字节数（权重 + scale） // BufferB 构造: n=intermediate/tp, k=hidden → 权重=n*k/2, scale=n*4 + // 除以 tp_count：与 do_gate_up_gemm 中 config_.intermediate_size（已被 TP 切分）匹配 size_t compute_gate_up_bytes(const MeshConfig& config) const { int h = config.hidden_size; - int i = config.intermediate_size / config.tp_count; // TP 切分后 + int i = config.intermediate_size / config.tp_count; // TP 切分 size_t elements = static_cast(h) * i; switch (config.weight_type) { case WeightType::AMXINT4: @@ -560,7 +1302,7 @@ class MeshResidencyManager { // BufferB 构造: n=hidden, k=intermediate/tp → 权重=n*k/2, scale=n*4 size_t compute_down_bytes(const MeshConfig& config) const { int h = config.hidden_size; - int i = config.intermediate_size / config.tp_count; + int i = config.intermediate_size / config.tp_count; // TP 切分 size_t elements = static_cast(i) * h; switch (config.weight_type) { case WeightType::AMXINT4: @@ -575,14 +1317,14 @@ class MeshResidencyManager { // AMXINT4: gate/up 的纯权重大小（不含 scale），用于 slot 内 scale 指针偏移 size_t compute_gate_up_weights_bytes(const MeshConfig& config) const { int h = config.hidden_size; - int i = config.intermediate_size / config.tp_count; + int i = config.intermediate_size / config.tp_count; // TP 切分 return static_cast(h) * i / 2; } // AMXINT4: down 的纯权重大小（不含 scale） size_t compute_down_weights_bytes(const MeshConfig& config) const { int h = config.hidden_size; - int i = config.intermediate_size / config.tp_count; + int i = config.intermediate_size / config.tp_count; // TP 切分 return static_cast(i) * h / 2; } }; diff --git a/kt-kernel/operators/mesh/mesh_scheduler.hpp b/kt-kernel/operators/mesh/mesh_scheduler.hpp index b0c91a7ee..0ea477fb0 100644 --- a/kt-kernel/operators/mesh/mesh_scheduler.hpp +++ b/kt-kernel/operators/mesh/mesh_scheduler.hpp @@ -141,6 +141,33 @@ class MeshScheduler { return result; } + // SKILL.md 第 54 行：驱逐和覆盖只影响本 TP 的本地 shard，不跨 NUMA 操作 + // 只取走指定 TP 的请求，其余放回队列，避免 TP0 操作 TP1 的 pool + std::vector drain_all_for_tp(int tp_part_idx) { + std::lock_guard lock(mutex_); + std::vector result; + std::vector remaining; + while (!pq_.empty()) { + auto req = std::move(const_cast(pq_.top())); + pq_.pop(); + if (req.tp_part_idx == tp_part_idx) { + result.push_back(std::move(req)); + } else { + remaining.push_back(std::move(req)); + } + } + // 把不属于该 TP 的请求放回队列 + for (auto& req : remaining) { + pq_.push(std::move(req)); + } + // result 按 schedule_key 升序排列 + std::sort(result.begin(), result.end(), + [](const ScheduledRequest& a, const ScheduledRequest& b) { + return a.schedule_key < b.schedule_key; + }); + return result; + } + // 取出 schedule_key 最小的请求（非阻塞） bool try_pop(ScheduledRequest& out) { std::lock_guard lock(mutex_); diff --git a/kt-kernel/operators/mesh/mesh_slot_pool.hpp b/kt-kernel/operators/mesh/mesh_slot_pool.hpp index 1a88530be..5938f34d8 100644 --- a/kt-kernel/operators/mesh/mesh_slot_pool.hpp +++ b/kt-kernel/operators/mesh/mesh_slot_pool.hpp @@ -11,12 +11,16 @@ #pragma once #include +#include +#include #include #include #include #include +#include #include #include +#include #include #include @@ -93,6 +97,14 @@ class MeshSlotPool { slots_.resize(cap); slot_to_expert_.assign(cap, -1); + // per-slot mutex + CV：保护 active_readers 的 acquire/release/overwrite 同步 + // 用 unique_ptr 因为 mutex/CV 不可移动 + slot_mtx_.resize(cap); + slot_cv_.resize(cap); + for (int i = 0; i < cap; i++) { + slot_mtx_[i] = std::make_unique(); + slot_cv_[i] = std::make_unique(); + } } ~MeshSlotPool() { @@ -110,6 +122,8 @@ class MeshSlotPool { : slots_(std::move(other.slots_)), expert_to_slot_(std::move(other.expert_to_slot_)), slot_to_expert_(std::move(other.slot_to_expert_)), + slot_mtx_(std::move(other.slot_mtx_)), + slot_cv_(std::move(other.slot_cv_)), layer_idx_(other.layer_idx_), tp_part_idx_(other.tp_part_idx_), numa_node_(other.numa_node_), @@ -129,6 +143,8 @@ class MeshSlotPool { slots_ = std::move(other.slots_); expert_to_slot_ = std::move(other.expert_to_slot_); slot_to_expert_ = std::move(other.slot_to_expert_); + slot_mtx_ = std::move(other.slot_mtx_); + slot_cv_ = std::move(other.slot_cv_); layer_idx_ = other.layer_idx_; tp_part_idx_ = other.tp_part_idx_; numa_node_ = other.numa_node_; @@ -195,6 +211,40 @@ class MeshSlotPool { return down_ptr(slot_idx); } + // ===== 原子化 acquire（lookup + state check + reader + pointer 全部在同一锁内）===== + // 消除 TOCTOU 竞争：原 expert_gate_ptr + acquire_reader 分两次独立 lookup， + // 中间 expert 可能被驱逐并重新加载到另一个 slot，导致返回的指针（旧 slot）和 + // reader（新 slot）不在同一个 slot 上。GEMM 用无 reader 保护的旧 slot 指针， + // 旧 slot 被 overwrite 覆盖 → 读到空内存 → 段错误。 + // 此方法保证指针和 reader 在同一个 slot 上，overwrite 持同一把锁无法插入。 + void* acquire_gate_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + std::lock_guard lock(*slot_mtx_[slot_idx]); + if (slot_to_expert_[slot_idx] != expert_id) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + return gate_ptr(slot_idx); + } + void* acquire_up_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + std::lock_guard lock(*slot_mtx_[slot_idx]); + if (slot_to_expert_[slot_idx] != expert_id) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + return up_ptr(slot_idx); + } + void* acquire_down_ptr(int expert_id) { + int slot_idx = expert_to_slot_.lookup(expert_id); + if (slot_idx < 0) return nullptr; + std::lock_guard lock(*slot_mtx_[slot_idx]); + if (slot_to_expert_[slot_idx] != expert_id) return nullptr; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return nullptr; + slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + return down_ptr(slot_idx); + } + // ===== 状态查询 ===== bool is_cached(int expert_id) const { @@ -251,58 +301,95 @@ class MeshSlotPool { expert_to_slot_.insert(expert_id, slot_idx); } + // 解绑：将 slot 恢复为 BASELINE（加载失败时调用） + void unbind(int slot_idx) { + Slot& s = slots_[slot_idx]; + int old_expert = s.bound_expert_id; + s.set_state(ExpertState::BASELINE); + s.bound_expert_id = -1; + slot_to_expert_[slot_idx] = -1; + if (old_expert >= 0) expert_to_slot_.erase(old_expert); + } + // 标记 slot 已缓存完毕（io_uring CQE 到达后调用） void mark_cached(int slot_idx) { slots_[slot_idx].set_state(ExpertState::CACHED); } // 覆盖：解绑旧 expert，将新 expert 读入同一个 slot - // 必须等 active_readers 归零后才能执行 + // 用 condition_variable 等待 active_readers 归零（30秒超时，超时后强制覆盖） + // 持有 slot_mtx_ 期间修改绑定关系，防止 acquire_reader 的 TOCTOU 竞争 void overwrite(int slot_idx, int new_expert_id) { Slot& s = slots_[slot_idx]; - // Bug 8 fix: spin-wait 加超时（10 秒），避免 reader 泄漏导致永久死锁 - // 超时后 abort，因为继续执行会导致数据损坏 { - const int kMaxSpinIters = 100000000; // ~10s 取决于 CPU - int spin = 0; - while (s.active_readers.load(std::memory_order_acquire) > 0) { - if (++spin > kMaxSpinIters) { - fprintf(stderr, - "[MESH] overwrite: slot %d active_readers never reached 0 " - "(expert %d -> %d), aborting to avoid data corruption\n", - slot_idx, s.bound_expert_id, new_expert_id); - std::abort(); + std::unique_lock lock(*slot_mtx_[slot_idx]); + // CV 等待 active_readers 归零：GEMM 完成后 release_reader 会 notify + // 30秒超时诊断：如果 reader 未释放，说明有 reader leak 或死锁 + // 注意：wait_for 之前不能调用 fprintf，否则 stderr 缓冲区满时 fprintf 阻塞， + // 锁被永久持有，wait_for 永远不被调用，30秒超时永远不触发 → 永久死锁 + if (s.active_readers.load(std::memory_order_acquire) > 0) { + auto status = slot_cv_[slot_idx]->wait_for(lock, std::chrono::seconds(30), [&] { + return s.active_readers.load(std::memory_order_acquire) == 0; + }); + if (!status) { + // 超时：reader 未释放，强制覆盖以避免永久死锁 + // 此时 fprintf 在 wait_for 之后，锁仍持有但 wait_for 已执行过，不会阻止超时 + int still = s.active_readers.load(std::memory_order_acquire); + fprintf(stderr, "[MESH OVERWRITE TIMEOUT] layer=%d tp=%d slot=%d old_expert=%d " + "new_expert=%d active_readers=%d — FORCING overwrite after 30s timeout! " + "Reader leak detected.\n", + layer_idx_, tp_part_idx_, slot_idx, s.bound_expert_id, + new_expert_id, still); } - // Bug 8 fix: 轻量让步，减少空转功耗（不用 sched_yield 避免系统调用开销） } + // 仍在锁内：修改绑定关系，acquire_reader 会看到新状态 + int old_expert_id = s.bound_expert_id; + if (old_expert_id >= 0) { + expert_to_slot_.erase(old_expert_id); + } + s.set_state(ExpertState::LOADING); + s.bound_expert_id = new_expert_id; + slot_to_expert_[slot_idx] = new_expert_id; + expert_to_slot_.insert(new_expert_id, slot_idx); } - int old_expert_id = s.bound_expert_id; - if (old_expert_id >= 0) { - expert_to_slot_.erase(old_expert_id); - } - // Bug 8 fix: 去掉无意义的 DEMOTING 中间状态 - // 原代码 DEMOTING 立即被 LOADING 覆盖，无窗口让其他线程看到 - // reader==0 已保证安全，直接进 LOADING - s.set_state(ExpertState::LOADING); - s.bound_expert_id = new_expert_id; - slot_to_expert_[slot_idx] = new_expert_id; - expert_to_slot_.insert(new_expert_id, slot_idx); } // ===== 引用计数 ===== // AMX 计算 / GPU 搬运前递增 reader - void acquire_reader(int expert_id) { + // 返回 true 表示成功获取 reader（slot 仍绑定该 expert 且 CACHED） + // 返回 false 表示 slot 已被驱逐（调用者应走 load 路径） + // 持有 slot_mtx_ 防止与 overwrite 的 TOCTOU 竞争 + bool acquire_reader(int expert_id) { int slot_idx = expert_to_slot_.lookup(expert_id); - if (slot_idx < 0) return; + if (slot_idx < 0) return false; + std::lock_guard lock(*slot_mtx_[slot_idx]); + // 双重检查：持锁后确认 expert 仍绑定且 CACHED + if (slot_to_expert_[slot_idx] != expert_id) return false; + if (slots_[slot_idx].get_state() != ExpertState::CACHED) return false; + // 不在此处打印日志：acquire_reader 是超高频路径（每 token 每 expert 每 GEMM）， + // 持锁 fprintf 会在 stderr 缓冲区满时阻塞，导致 slot_mtx_ 被永久持有 → 死锁 slots_[slot_idx].active_readers.fetch_add(1, std::memory_order_acq_rel); + return true; } - // 计算完毕递减 reader + // 计算完毕递减 reader，并 notify overwrite 的 CV 等待 + // 死锁修复：不持有 slot_mtx_ 锁。active_readers 是 std::atomic，fetch_sub 不需要锁。 + // 原实现持锁 fetch_sub 会导致死锁：overwrite 持 slot_mtx_ CV wait 等 active_readers==0， + // release_reader 需要同一个 slot_mtx_ 才能 fetch_sub → 永久死锁。 + // 不持锁的安全性：lookup 可能和 overwrite 的 erase/insert 竞争，但 int 读写通常原子。 + // 若 lookup 返回 -1（expert 已被 evict），说明 overwrite 已 CV wait 完成，迟到 return 安全。 + // 若 lookup 返回正确 slot_idx，overwrite 还在 CV wait（active_readers>0），fetch_sub 安全。 void release_reader(int expert_id) { int slot_idx = expert_to_slot_.lookup(expert_id); - if (slot_idx < 0) return; + if (slot_idx < 0) { + // expert 已被 evict，reader 无法释放。 + // 这是潜在 reader leak 的来源之一，但此处不打印日志（超高频路径）。 + // 若需诊断，可在 overwrite 超时日志中观察到 active_readers>0。 + return; + } slots_[slot_idx].active_readers.fetch_sub(1, std::memory_order_acq_rel); + slot_cv_[slot_idx]->notify_all(); } // Bug 11 fix: 按 slot_idx 递增/递减 reader，不依赖 expert 绑定关系 @@ -314,6 +401,7 @@ class MeshSlotPool { void release_slot_reader(int slot_idx) { if (slot_idx < 0 || slot_idx >= cap_) return; slots_[slot_idx].active_readers.fetch_sub(1, std::memory_order_acq_rel); + slot_cv_[slot_idx]->notify_all(); } // ===== 驱逐扫描 ===== @@ -330,12 +418,15 @@ class MeshSlotPool { return -1; } - // 获取所有 CACHED 状态的专家 ID（驱逐评分用） + // 获取所有可驱逐的专家 ID（CACHED 且 active_readers==0） + // SKILL.md 第 62 行：不驱逐读取中/使用中的专家 + // active_readers>0 表示有 GEMM 正在读取，强制驱逐会导致读到空内存 std::vector cached_experts() const { std::vector result; for (int i = 0; i < cap_; i++) { if (slots_[i].get_state() == ExpertState::CACHED && - slots_[i].bound_expert_id >= 0) { + slots_[i].bound_expert_id >= 0 && + slots_[i].active_readers.load(std::memory_order_acquire) == 0) { result.push_back(slots_[i].bound_expert_id); } } @@ -372,6 +463,12 @@ class MeshSlotPool { std::vector slots_; // [cap_]，驱逐扫描用 + // per-slot mutex + condition_variable + // 保护 acquire_reader/release_reader/overwrite 之间的同步 + // acquire_reader 持锁检查 state + 增计数；overwrite 持锁 CV 等待计数归零 + std::vector> slot_mtx_; + std::vector> slot_cv_; + // expert_id -> slot_idx 的双向映射 // 简单实现用 vector，O(1) 查找 struct ExpertToSlot { diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 4e313ca33..3f21f287e 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -387,6 +387,35 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor): Args: physical_to_logical_map_cpu: Mapping from physical to logical expert IDs """ + # MESH 模式：跳过全量权重加载，仅创建 MoE 模块（权重由 slot pool 提供） + # C++ load_weights 会因 mesh_enabled=True 跳过，gate_bb_ 保持空分配 + # forward 时 slot 未命中通过 get_or_load_*_ptr 同步加载 + if self.mesh_enabled and self.mesh_residency_ptr: + moe_config = MOEConfig( + self.num_experts, + self.num_experts_per_tok, + self.hidden_size, + self.moe_intermediate_size, + self.gpu_experts_mask.data_ptr(), + ) + moe_config.layer_idx = self.layer_idx + moe_config.pool = self.cpu_infer.backend_ + moe_config.max_len = self.chunked_prefill_size + moe_config.mesh_enabled = self.mesh_enabled + moe_config.mesh_residency = self.mesh_residency_ptr + moe_config.load = False + + if self.method == "AMXINT4": + self.moe = AMXInt4_MOE(moe_config) + elif self.method == "AMXINT8": + self.moe = AMXInt8_MOE(moe_config) + else: + raise NotImplementedError(f"Unsupported AMX method: {self.method}") + + self.cpu_infer.submit(self.moe.load_weights_task(physical_to_logical_map_cpu.data_ptr())) + self.cpu_infer.sync() + return + gate_ptr = 0 up_ptr = 0 down_ptr = 0 @@ -552,6 +581,8 @@ def __init__( numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, swiglu_alpha: float = 0.0, + mesh_enabled: bool = False, + mesh_residency_ptr: int = 0, ): self._swiglu_alpha = float(swiglu_alpha) # Defence in depth: reject swiglu_limit on non-MXFP4/MXFP8 methods even @@ -643,6 +674,10 @@ def __init__( self.up_scales = None self.down_scales = None + # MESH 插件参数 + self.mesh_enabled = mesh_enabled + self.mesh_residency_ptr = mesh_residency_ptr + @staticmethod def _create_loader(method: str, weight_path: str): if method == "RAWINT4": @@ -693,6 +728,36 @@ def load_weights_from_tensors( def load_weights(self, physical_to_logical_map_cpu: torch.Tensor): import time + # MESH 模式：跳过磁盘权重加载，仅创建 MoE 模块 + # C++ load_weights 会因 mesh_enabled=True 跳过，forward 时由 MESH hook 提供权重 + if self.mesh_enabled and self.mesh_residency_ptr: + moe_config = MOEConfig( + self.num_experts, + self.num_experts_per_tok, + self.hidden_size, + self.moe_intermediate_size, + self.gpu_experts_mask.data_ptr(), + ) + moe_config.layer_idx = self.layer_idx + moe_config.pool = self.cpu_infer.backend_ + moe_config.max_len = self.chunked_prefill_size + moe_config.mesh_enabled = self.mesh_enabled + moe_config.mesh_residency = self.mesh_residency_ptr + + if self.method == "BF16": + if _HAS_BF16_SUPPORT: + self.moe = AMXBF16_MOE(moe_config) + else: + self.moe = AVX2BF16_MOE(moe_config) + else: + raise NotImplementedError( + f"MESH mode not supported for method={self.method!r} in NativeMoEWrapper" + ) + + self.cpu_infer.submit(self.moe.load_weights_task(physical_to_logical_map_cpu.data_ptr())) + self.cpu_infer.sync() + return + if NativeMoEWrapper._native_loader_instance is None: t_recreate_start = time.time() NativeMoEWrapper._native_loader_instance = NativeMoEWrapper._create_loader( @@ -797,6 +862,9 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor): moe_config.layer_idx = self.layer_idx moe_config.pool = self.cpu_infer.backend_ moe_config.max_len = self.chunked_prefill_size + # MESH 插件：注入 mesh 参数 + moe_config.mesh_enabled = self.mesh_enabled + moe_config.mesh_residency = self.mesh_residency_ptr # V4-Flash 2604B SwiGLU clamp; 0.0 = disabled (default for non-MXFP4 # paths). Read by `act_fn` in operators/amx/la/amx.hpp via # `apply_activation` in operators/amx/moe_base.hpp. Re-checked here diff --git a/kt-kernel/python/utils/mesh/stats.py b/kt-kernel/python/utils/mesh/stats.py index 131ff1df0..f930c9d10 100644 --- a/kt-kernel/python/utils/mesh/stats.py +++ b/kt-kernel/python/utils/mesh/stats.py @@ -72,6 +72,24 @@ def summary(self) -> str: f"{self.decode_deferred_count} deferred" ) + def to_kv(self) -> str: + """返回单行 key=value 格式，方便从日志 grep 解析。""" + return ( + f"hit_rate={self.cache_hit_rate:.4f} " + f"hit_count={self.cache_hit_count} " + f"miss_count={self.cache_miss_count} " + f"iouring_read_gib={self.io_uring_read_gib:.2f} " + f"iouring_read_count={self.io_uring_read_count} " + f"eviction_count={self.eviction_count} " + f"defer_count={self.defer_count} " + f"defer_overflow={self.defer_overflow_count} " + f"prefill_layers={self.prefill_layer_count} " + f"prefill_swaps={self.prefill_temporal_swap_count} " + f"decode_tokens={self.decode_token_count} " + f"decode_immediate={self.decode_immediate_count} " + f"decode_deferred={self.decode_deferred_count}" + ) + class MeshStatsCollector: """从 C++ 侧读取统计信息。""" diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py index 7f4128b8f..07d0a09d3 100644 --- a/kt-kernel/python/utils/mesh/wrapper.py +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -14,11 +14,14 @@ import logging import os +import threading +import time import torch from typing import List, Optional from .config import MeshConfig from .residency import MeshResidencyManager +from .stats import MeshStatsCollector logger = logging.getLogger(__name__) @@ -82,29 +85,55 @@ def __init__( # 如果未设置，需要外部补充（通常由调用者设置） logger.warning("mesh_config.total_layers not set, MESH schedule_key may be incorrect") - # 1. 创建 AMXMoEWrapper，注入 mesh_enabled=True - # AMX 内核代码完全复用，load_weights() 会走 mesh 分支（直接 return） - from ..amx import AMXMoEWrapper - self._inner = AMXMoEWrapper( - layer_idx=layer_idx, - num_experts=num_experts, - num_experts_per_tok=num_experts_per_tok, - hidden_size=hidden_size, - moe_intermediate_size=moe_intermediate_size, - gpu_experts_mask=gpu_experts_mask, - cpuinfer_threads=cpuinfer_threads, - threadpool_count=threadpool_count, - weight_path=weight_path, - chunked_prefill_size=chunked_prefill_size, - cpu_save=cpu_save, - max_deferred_experts_per_token=max_deferred_experts_per_token, - method=method, - numa_nodes=numa_nodes, - # A2: mesh_enabled 暂不开启（避免 load_weights 跳过导致 gate_bb_ 为空崩溃） - # 仅注入 mesh_residency 指针，让 do_gate_up_gemm 的 hook 可被调用 - mesh_enabled=False, - mesh_residency_ptr=0, - ) + # 1. 创建内部 Wrapper + # BF16 模式：使用 NativeMoEWrapper（支持 BF16 kernel + MESH hooks） + # AMXINT4 模式：使用 AMXMoEWrapper（AMX INT4 kernel + MESH hooks） + from ..amx import AMXMoEWrapper, NativeMoEWrapper + + is_bf16 = (method.upper() == "BF16") + if is_bf16: + # BF16 MESH 模式：NativeMoEWrapper + mesh_enabled=True + # C++ load_weights 会跳过，forward 时由 MESH hook 提供权重 + self._inner = NativeMoEWrapper( + layer_idx=layer_idx, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + gpu_experts_mask=gpu_experts_mask, + cpuinfer_threads=cpuinfer_threads, + threadpool_count=threadpool_count, + weight_path=weight_path, + chunked_prefill_size=chunked_prefill_size, + cpu_save=cpu_save, + max_deferred_experts_per_token=max_deferred_experts_per_token, + method=method, + numa_nodes=numa_nodes, + mesh_enabled=True, + mesh_residency_ptr=0, # 稍后注入 + ) + else: + # AMXINT4 MESH 模式：AMXMoEWrapper + mesh_enabled=True + # C++ load_weights 会跳过全量加载，forward 时由 MESH hook 提供权重 + # slot 未命中时通过 get_or_load_*_ptr 同步加载（而非回退到空的 gate_bb_） + self._inner = AMXMoEWrapper( + layer_idx=layer_idx, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + hidden_size=hidden_size, + moe_intermediate_size=moe_intermediate_size, + gpu_experts_mask=gpu_experts_mask, + cpuinfer_threads=cpuinfer_threads, + threadpool_count=threadpool_count, + weight_path=weight_path, + chunked_prefill_size=chunked_prefill_size, + cpu_save=cpu_save, + max_deferred_experts_per_token=max_deferred_experts_per_token, + method=method, + numa_nodes=numa_nodes, + mesh_enabled=True, + mesh_residency_ptr=0, + ) # 2. 创建/复用 ResidencyManager（进程级单例，避免 40× 内存重复） if numa_nodes is None: @@ -144,10 +173,32 @@ def __init__( self._mesh_config = mesh_config self._layer_idx = layer_idx - # A2: 注入 mesh_residency 指针到 inner AMXMoEWrapper + # A2: 注入 mesh_residency 指针到 inner Wrapper # 这样 do_gate_up_gemm/do_down_gemm 中的 mesh hook 可以调用 ResidencyManager self._inner.mesh_residency_ptr = self._residency.raw_ptr() + # Stats 收集：只在 layer 0 创建 collector，启动后台线程定期输出 + if layer_idx == 0: + self._stats_collector = MeshStatsCollector(self._residency) + self._stats_stop_event = threading.Event() + self._stats_thread = threading.Thread( + target=self._stats_loop, args=(self._stats_stop_event,), daemon=True + ) + self._stats_thread.start() + logger.info("MeshMoEWrapper: stats background thread started (interval=10s)") + else: + self._stats_collector = None + self._stats_stop_event = None + self._stats_thread = None + + def _stats_loop(self, stop_event): + """后台线程：每 10 秒输出一次 MESH stats 到日志。""" + while not stop_event.is_set(): + stop_event.wait(10) + if stop_event.is_set(): + break + self.dump_stats() + def _inject_file_layouts( self, weight_path: str, @@ -228,16 +279,17 @@ def _inject_amxint4_layouts( - blk.{L}.ffn_up_exps.{E}.numa.{N}.weight / .scale - blk.{L}.ffn_down_exps.{E}.numa.{N}.weight / .scale """ - # 收集所有需要打开的文件，用 O_DIRECT 打开 + # 收集所有需要打开的文件，用 O_DIRECT 打开（绕过 page cache） file_fds: dict[str, int] = {} def get_fd(fpath: str) -> int: if fpath not in file_fds: - # O_DIRECT 需要对齐读取，但 safetensors 偏移可能不对齐 - # 先用 O_RDONLY 打开，O_DIRECT 对齐问题在 io_uring 层处理 - fd = os.open(fpath, os.O_RDONLY) + # O_DIRECT：SSD 直接 DMA 到 NUMA buffer，不经过 page cache + # SKILL.md 第11行：MESH 核心设计原则，禁止使用 page cache + # MoE 专家权重已验证全部 512 对齐（offset 和 size） + fd = os.open(fpath, os.O_RDONLY | os.O_DIRECT) file_fds[fpath] = fd - logger.debug("_inject_amxint4_layouts: opened %s as fd=%d", fpath, fd) + logger.debug("_inject_amxint4_layouts: opened %s as fd=%d (O_DIRECT)", fpath, fd) return file_fds[fpath] injected = 0 @@ -311,7 +363,7 @@ def get_fd(fpath: str) -> int: logger.info( "_inject_amxint4_layouts: injected %d layouts (missing %d), " - "opened %d files with O_RDONLY", + "opened %d files with O_DIRECT", injected, missing, len(file_fds), ) @@ -341,19 +393,26 @@ def _inject_bf16_layouts( def get_fd(fpath: str) -> int: if fpath not in file_fds: - fd = os.open(fpath, os.O_RDONLY) + # O_DIRECT：绕过 page cache，SSD 直接 DMA 到 NUMA buffer + fd = os.open(fpath, os.O_RDONLY | os.O_DIRECT) file_fds[fpath] = fd return file_fds[fpath] injected = 0 missing = 0 + # 支持两种前缀：model.layers.{L}（标准）和 model.language_model.layers.{L}（VL 模型 TEXTONLY） for layer in range(total_layers): gate_up_key = f"model.layers.{layer}.mlp.experts.gate_up_proj" down_key = f"model.layers.{layer}.mlp.experts.down_proj" - if gate_up_key not in tensor_map or down_key not in tensor_map: - missing += 1 - continue + lm_key = f"model.language_model.layers.{layer}.mlp.experts.gate_up_proj" + lm_down = f"model.language_model.layers.{layer}.mlp.experts.down_proj" + if lm_key in tensor_map and lm_down in tensor_map: + gate_up_key = lm_key + down_key = lm_down + else: + missing += 1 + continue gu_fpath, gu_base, _ = tensor_map[gate_up_key] d_fpath, d_base, _ = tensor_map[down_key] @@ -425,20 +484,20 @@ def __getattr__(self, name: str): # ===== forward 委托 ===== def forward(self, *args, **kwargs): - """forward 委托给 AMXMoEWrapper。 + """forward 委托给内部 Wrapper（AMXMoEWrapper 或 NativeMoEWrapper）。 - MESH 通过 hook 在 AMX 内核中重定向权重指针， + MESH 通过 hook 在内核中重定向权重指针， Python 侧无需特殊处理。 """ return self._inner.forward(*args, **kwargs) def load_weights(self, physical_to_logical_map_cpu=None): - """load_weights：委托给内部 AMXMoEWrapper 加载全量权重。 + """load_weights：委托给内部 Wrapper。 - TODO: MESH 模式的最终目标是跳过全量加载，仅由 ResidencyManager 的 - slot pool 管理 cap 个专家（io_uring 按需加载）。但当前 io_uring 按需 - 加载和 AMX 内核 hook 尚未实现，self.moe 为 None 会导致 submit_forward - 崩溃。暂时先正常加载权重以保证推理可用，slot pool 作为额外开销。 + - BF16 模式：NativeMoEWrapper 的 load_weights 会跳过磁盘加载， + C++ load_weights 也跳过，forward 时由 MESH hook 提供权重。 + - AMXINT4 模式：AMXMoEWrapper 正常加载全量权重， + forward 时 MESH hook 重定向到 slot pool。 """ return self._inner.load_weights(physical_to_logical_map_cpu) @@ -478,3 +537,13 @@ def on_decode_layer(self, topk: List[int], scores: List[float], tp_part_idx: int def on_decode_token_end(self, all_layers_topk: List[List[int]], all_layers_scores: List[List[float]]) -> None: self._residency.on_decode_token_end(all_layers_topk, all_layers_scores) + + def dump_stats(self) -> None: + """输出 MESH stats 到日志。benchmark 脚本从日志 grep 获取。""" + if self._stats_collector is None: + return + try: + stats = self._stats_collector.collect() + logger.info("[MESH_STATS_KV] %s", stats.to_kv()) + except Exception as e: + logger.warning("Failed to dump MESH stats: %s", e) diff --git a/scripts/bench_bf16_py.py b/scripts/bench_bf16_py.py new file mode 100644 index 000000000..f9002c304 --- /dev/null +++ b/scripts/bench_bf16_py.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +"""通用 benchmark 脚本 — 支持 bf16 和 AMXINT4，标准和 MESH 模式。 +Usage: + 标准 kt: python3 bench_bf16_py.py kt + MESH: python3 bench_bf16_py.py mesh + +Example: + python3 bench_bf16_py.py kt 35b 10004 0,1,2,3 + python3 bench_bf16_py.py kt 397b 10004 0,1,2,3 + python3 bench_bf16_py.py mesh 35b "64 128 192 224" 10004 0,1,2,3 + python3 bench_bf16_py.py mesh 397b "128 192 256 480" 10004 0,1,2,3 +""" +import subprocess, time, json, threading, requests, sys, os + +MODE = sys.argv[1] # kt / mesh +MODEL = sys.argv[2] # 35b / 397b + +if MODE == "kt": + PORT = int(sys.argv[3]) + CUDA = sys.argv[4] + SCRIPT = f"run_kt_{MODEL}_bf16.sh" + CAP_LIST = [None] + RESULT_FILE = f"/tmp/kt_bench_{MODEL}_bf16_results.txt" +elif MODE == "mesh": + CAP_LIST = [int(x) for x in sys.argv[3].split()] + PORT = int(sys.argv[4]) + CUDA = sys.argv[5] + SCRIPT = f"run_mesh_{MODEL}_bf16.sh" + RESULT_FILE = f"/tmp/mesh_bench_{MODEL}_bf16_results.txt" +else: + print(f"Unknown mode: {MODE}") + sys.exit(1) + +with open(RESULT_FILE, "w") as f: + f.write("mode | decode_tok_s | prefill_tok_s | prefill_s | peak_sys_used_gib | peak_gpu_gib\n") + +URL = f"http://localhost:{PORT}/v1/completions" + + +def run(cmd, timeout=30): + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) + + +def wait_ready(log_file, server_pid, max_wait=900): + for i in range(max_wait // 5): + time.sleep(5) + try: + with open(log_file) as f: + content = f.read() + if "fired up and ready to roll" in content: + return True + for fatal in ["Segmentation fault", "CUDA error", "RuntimeError:", + "Traceback (most recent call last):", "core dumped", "Out of memory"]: + if fatal in content: + print(f"FATAL: {fatal}") + return False + except: + pass + try: + os.kill(server_pid, 0) + except ProcessLookupError: + print("Server process exited") + return False + return False + + +def sample_memory(stop_event, samples): + while not stop_event.is_set(): + try: + out = subprocess.check_output(["free", "-m"], text=True) + for line in out.splitlines(): + if line.startswith("Mem:"): + used = int(line.split()[2]) + samples.append(("mem", used)) + break + except: + pass + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], + text=True + ) + gpu_sum = sum(int(x) for x in out.strip().split("\n")[:4]) + samples.append(("gpu", gpu_sum)) + except: + pass + time.sleep(2) + + +def bench_prefill(): + prompt = "The quick brown fox jumps over the lazy dog. " * 260 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=300) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def bench_decode(): + data = {"model": "default", "prompt": "Write a short essay about artificial intelligence and its impact on society:", + "max_tokens": 256, "temperature": 0, "stream": True} + start = time.time() + first = None + cnt = 0 + r = requests.post(URL, json=data, stream=True, timeout=120) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + chunk = json.loads(line) + if chunk["choices"][0]["text"]: + cnt += 1 + if first is None: + first = time.time() + end = time.time() + dec = end - first if first else end - start + dtps = cnt / dec if dec > 0 else 0 + return dtps + + +for cap in CAP_LIST: + label = f"cap={cap}" if cap else "standard_kt" + print(f"\n========== Testing {MODEL} bf16 {label} ==========") + + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + + log_file = f"/tmp/sglang_bf16_{MODEL}_{label}.log" + if cap: + server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {cap} {PORT} {CUDA}" + else: + server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {PORT} {CUDA}" + proc = subprocess.Popen(server_cmd, shell=True, stdout=open(log_file, "w"), stderr=subprocess.STDOUT) + print(f"Started server PID={proc.pid}, waiting for ready...") + + if not wait_ready(log_file, proc.pid): + print(f"FAILED: Server not ready for {label}") + with open(RESULT_FILE, "a") as f: + f.write(f"{label} | FAILED | FAILED | FAILED | FAILED | FAILED\n") + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + continue + + print(f"Server ready") + time.sleep(5) + + samples = [] + stop_event = threading.Event() + mem_thread = threading.Thread(target=sample_memory, args=(stop_event, samples)) + mem_thread.start() + print("Memory sampler started") + + try: + pt, ptime, ptps = bench_prefill() + print(f"Prefill: {pt} tokens, {ptime:.2f}s, {ptps:.2f} tok/s") + except Exception as e: + print(f"Prefill failed: {e}") + pt, ptime, ptps = 0, 0, 0 + + time.sleep(3) + + try: + dtps = bench_decode() + print(f"Decode: {dtps:.2f} tok/s") + except Exception as e: + print(f"Decode failed: {e}") + dtps = 0 + + stop_event.set() + mem_thread.join() + + mem_vals = [v for k, v in samples if k == "mem"] + gpu_vals = [v for k, v in samples if k == "gpu"] + peak_mem = max(mem_vals) / 1024 if mem_vals else 0 + peak_gpu = max(gpu_vals) / 1024 if gpu_vals else 0 + print(f"Peak sys used: {peak_mem:.2f} GiB, Peak GPU: {peak_gpu:.2f} GiB, Samples: {len(mem_vals)}") + + with open(RESULT_FILE, "a") as f: + f.write(f"{label} | {dtps:.2f} | {ptps:.2f} | {ptime:.2f} | {peak_mem:.2f} | {peak_gpu:.2f}\n") + + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + +print("\n=== Results ===") +with open(RESULT_FILE) as f: + print(f.read()) diff --git a/scripts/bench_full_py.py b/scripts/bench_full_py.py new file mode 100644 index 000000000..e6ce218db --- /dev/null +++ b/scripts/bench_full_py.py @@ -0,0 +1,344 @@ +#!/usr/bin/env python3 +"""完整 benchmark 脚本 — 支持所有格式，采集所有规范指标。 + +采集指标： + 1. decode tok/s（256 tokens） + 2. 短 prompt prefill tok/s（~2601 tokens） + 3. 长 prompt prefill tok/s（~8192 tokens） + 4. peak_sys_used_gib（从 server 启动前开始采样，覆盖 bootstrap） + 5. peak_gpu_gib + 6. MESH stats（hit_rate、iouring_read_gib 等，从日志 grep） + 7. iostat 磁盘读取带宽（r_MB_s） + +Usage: + kt 模式: python3 bench_full_py.py kt + mesh 模式: python3 bench_full_py.py mesh + + MODEL: 35b / 397b + FORMAT: amxint4 / bf16 + +Example: + python3 bench_full_py.py kt 35b amxint4 10004 0,1,2,3 + python3 bench_full_py.py kt 397b bf16 10004 0,1,2,3 + python3 bench_full_py.py mesh 35b amxint4 "64 128 192 224" 10004 0,1,2,3 + python3 bench_full_py.py mesh 397b bf16 "128 192 256 480" 10004 0,1,2,3 +""" +import subprocess, time, json, threading, requests, sys, os, re + +MODE = sys.argv[1] # kt / mesh +MODEL = sys.argv[2] # 35b / 397b +FORMAT = sys.argv[3] # amxint4 / bf16 + +if MODE == "kt": + PORT = int(sys.argv[4]) + CUDA = sys.argv[5] + CAP_LIST = [None] + SCRIPT = f"run_kt_{MODEL}_{FORMAT}.sh" if FORMAT != "amxint4" else f"run_kt_{MODEL}.sh" + RESULT_FILE = f"/tmp/full_bench_{MODEL}_{FORMAT}_kt_results.txt" +elif MODE == "mesh": + CAP_LIST = [int(x) for x in sys.argv[4].split()] + PORT = int(sys.argv[5]) + CUDA = sys.argv[6] + SCRIPT = f"run_mesh_{MODEL}_cap.sh" if FORMAT != "bf16" else f"run_mesh_{MODEL}_bf16.sh" + RESULT_FILE = f"/tmp/full_bench_{MODEL}_{FORMAT}_mesh_results.txt" +else: + print(f"Unknown mode: {MODE}") + sys.exit(1) + +with open(RESULT_FILE, "w") as f: + f.write("mode | decode_tok_s | prefill_short_tok_s | prefill_long_tok_s | " + "peak_sys_gib | peak_gpu_gib | hit_rate | iouring_read_gib | " + "iostat_r_mbs | eviction_count | decode_tokens\n") + +URL = f"http://localhost:{PORT}/v1/completions" + + +def run(cmd, timeout=30): + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) + + +def wait_ready(log_file, server_pid, max_wait=900): + for i in range(max_wait // 5): + time.sleep(5) + try: + with open(log_file) as f: + content = f.read() + if "fired up and ready to roll" in content: + return True + for fatal in ["Segmentation fault", "CUDA error", "RuntimeError:", + "Traceback (most recent call last):", "core dumped", "Out of memory"]: + if fatal in content: + print(f"FATAL: {fatal}") + return False + except: + pass + try: + os.kill(server_pid, 0) + except ProcessLookupError: + print("Server process exited") + return False + return False + + +def sample_memory(stop_event, samples): + """从 server 启动前开始采样，覆盖 bootstrap + 测试全过程。""" + while not stop_event.is_set(): + try: + out = subprocess.check_output(["free", "-m"], text=True) + for line in out.splitlines(): + if line.startswith("Mem:"): + used = int(line.split()[2]) + samples.append(("mem", used)) + break + except: + pass + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], + text=True + ) + gpu_sum = sum(int(x) for x in out.strip().split("\n")[:4]) + samples.append(("gpu", gpu_sum)) + except: + pass + time.sleep(2) + + +def start_iostat(iostat_file): + """启动 iostat 后台采样，每 5 秒一次。""" + # 采样所有设备，解析时筛选数据盘 + cmd = f"iostat -x -d 5 > {iostat_file} 2>&1" + proc = subprocess.Popen(cmd, shell=True) + return proc + + +def parse_iostat(iostat_file): + """解析 iostat 输出，返回数据盘的峰值读取带宽 (MB/s)。 + + iostat -x -d 输出格式： + Device r/s rkB/s rrqm/s %rrqm r_await rareq-sz w/s wkB/s ... + dm-1 203 35097 0 0 3.77 172 220 1308 ... + + 返回 (peak_MB/s, avg_MB/s) + """ + try: + with open(iostat_file) as f: + content = f.read() + # 采样每轮所有非 loop/ram 设备的 rkB/s 总和 + values = [] + current_round = {} + for line in content.splitlines(): + parts = line.split() + if len(parts) >= 3 and parts[0] not in ("Device", "Linux", ""): + dev = parts[0] + if dev.startswith(("loop", "ram", "sr")): + continue + try: + rkbs = float(parts[2]) # rkB/s + current_round[dev] = current_round.get(dev, 0) + rkbs + except (ValueError, IndexError): + pass + elif parts and parts[0] == "Device": + if current_round: + values.append(sum(current_round.values())) + current_round = {} + if current_round: + values.append(sum(current_round.values())) + + if values: + peak_mbs = max(values) / 1024 + avg_mbs = sum(values) / len(values) / 1024 + return peak_mbs, avg_mbs + return 0.0, 0.0 + except: + return 0.0, 0.0 + + +def bench_prefill_short(): + """短 prompt prefill (~2601 tokens)。""" + prompt = "The quick brown fox jumps over the lazy dog. " * 260 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=300) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def bench_prefill_long(): + """长 prompt prefill (~8192 tokens)。""" + prompt = "The quick brown fox jumps over the lazy dog. " * 820 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=600) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def bench_decode(): + """Decode 256 tokens。""" + data = {"model": "default", "prompt": "Write a short essay about artificial intelligence and its impact on society:", + "max_tokens": 256, "temperature": 0, "stream": True} + start = time.time() + first = None + cnt = 0 + r = requests.post(URL, json=data, stream=True, timeout=120) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + chunk = json.loads(line) + if chunk["choices"][0]["text"]: + cnt += 1 + if first is None: + first = time.time() + end = time.time() + dec = end - first if first else end - start + dtps = cnt / dec if dec > 0 else 0 + return dtps + + +def grep_mesh_stats(log_file): + """从 server 日志 grep 最后一次 MESH_STATS_KV 输出。""" + try: + with open(log_file) as f: + content = f.read() + # 找最后一次 [MESH_STATS_KV] 行 + matches = re.findall(r'\[MESH_STATS_KV\]\s*(.*)', content) + if matches: + kv_line = matches[-1].strip() + # 解析 key=value 对 + stats = {} + for pair in kv_line.split(): + if "=" in pair: + k, v = pair.split("=", 1) + stats[k] = v + return stats + except: + pass + return {} + + +for cap in CAP_LIST: + label = f"cap={cap}" if cap else "standard_kt" + print(f"\n========== Testing {MODEL} {FORMAT} {label} ==========") + + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + + log_file = f"/tmp/sglang_full_{MODEL}_{FORMAT}_{label}.log" + iostat_file = f"/tmp/iostat_{MODEL}_{FORMAT}_{label}.log" + + # 1. 在 server 启动前开始内存采样（修复时序问题） + samples = [] + stop_event = threading.Event() + mem_thread = threading.Thread(target=sample_memory, args=(stop_event, samples)) + mem_thread.start() + print("Memory sampler started (pre-server)") + + # 2. 启动 iostat 采样 + iostat_proc = start_iostat(iostat_file) + print(f"iostat started (PID={iostat_proc.pid})") + + # 3. 启动 server + if cap: + server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {cap} {PORT} {CUDA}" + else: + server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {PORT} {CUDA}" + proc = subprocess.Popen(server_cmd, shell=True, stdout=open(log_file, "w"), stderr=subprocess.STDOUT) + print(f"Started server PID={proc.pid}, waiting for ready...") + + if not wait_ready(log_file, proc.pid): + print(f"FAILED: Server not ready for {label}") + with open(RESULT_FILE, "a") as f: + f.write(f"{label} | FAILED | FAILED | FAILED | FAILED | FAILED | - | - | - | - | -\n") + stop_event.set() + mem_thread.join() + iostat_proc.kill() + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + continue + + print(f"Server ready") + time.sleep(5) + + # 4. 短 prompt prefill（< 4096 threshold，走 GPU prefill） + try: + pt_s, ptime_s, ptps_s = bench_prefill_short() + print(f"Prefill short: {pt_s} tokens, {ptime_s:.2f}s, {ptps_s:.2f} tok/s") + except Exception as e: + print(f"Prefill short failed: {e}") + pt_s, ptime_s, ptps_s = 0, 0, 0 + + time.sleep(3) + + # 5. Decode（先于长 prompt prefill，确保 MESH stats 能输出） + try: + dtps = bench_decode() + print(f"Decode: {dtps:.2f} tok/s") + except Exception as e: + print(f"Decode failed: {e}") + dtps = 0 + + time.sleep(3) + + # 6. 长 prompt prefill（最后执行，可能触发 GPU prefill bug 导致 server 崩溃） + try: + pt_l, ptime_l, ptps_l = bench_prefill_long() + print(f"Prefill long: {pt_l} tokens, {ptime_l:.2f}s, {ptps_l:.2f} tok/s") + except Exception as e: + print(f"Prefill long failed: {e}") + pt_l, ptime_l, ptps_l = 0, 0, 0 + + # 7. 停止采样 + time.sleep(3) # 等待最后的 MESH stats 输出 + stop_event.set() + mem_thread.join() + iostat_proc.kill() + + # 8. 计算峰值 + mem_vals = [v for k, v in samples if k == "mem"] + gpu_vals = [v for k, v in samples if k == "gpu"] + peak_mem = max(mem_vals) / 1024 if mem_vals else 0 + peak_gpu = max(gpu_vals) / 1024 if gpu_vals else 0 + print(f"Peak sys used: {peak_mem:.2f} GiB, Peak GPU: {peak_gpu:.2f} GiB, Samples: {len(mem_vals)}") + + # 9. 解析 iostat + iostat_peak, iostat_avg = parse_iostat(iostat_file) + print(f"iostat: peak={iostat_peak:.1f} MB/s, avg={iostat_avg:.1f} MB/s") + + # 10. 从日志 grep MESH stats + mesh_stats = grep_mesh_stats(log_file) + hit_rate = float(mesh_stats.get("hit_rate", 0)) + iouring_gib = float(mesh_stats.get("iouring_read_gib", 0)) + eviction_count = int(mesh_stats.get("eviction_count", 0)) + decode_tokens = int(mesh_stats.get("decode_tokens", 0)) + if mesh_stats: + print(f"MESH stats: hit_rate={hit_rate:.4f}, iouring_read={iouring_gib:.2f} GiB, " + f"evictions={eviction_count}, decode_tokens={decode_tokens}") + else: + print("No MESH stats found in log") + + # 11. 写入结果 + with open(RESULT_FILE, "a") as f: + f.write(f"{label} | {dtps:.2f} | {ptps_s:.2f} | {ptps_l:.2f} | " + f"{peak_mem:.2f} | {peak_gpu:.2f} | " + f"{hit_rate:.4f} | {iouring_gib:.2f} | " + f"{iostat_peak:.1f} | {eviction_count} | {decode_tokens}\n") + + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + +print("\n=== Results ===") +with open(RESULT_FILE) as f: + print(f.read()) diff --git a/scripts/convert_35b_bf16_packed.py b/scripts/convert_35b_bf16_packed.py new file mode 100644 index 000000000..081f042b3 --- /dev/null +++ b/scripts/convert_35b_bf16_packed.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""将 35B bf16 deepseek 格式转为 MESH 需要的 packed 格式。 + +输入: model.layers.{L}.mlp.experts.experts.{id}.{gate,up,down}_proj.weight +输出: model.layers.{L}.mlp.experts.gate_up_proj [E, 2*I, H] + model.layers.{L}.mlp.experts.down_proj [E, H, I] + +用法: python3 convert_35b_bf16_packed.py +""" +import os +import sys +import torch +from safetensors import safe_open +from safetensors.torch import save_file +import glob + +SRC = "/mnt/data2/models/Qwen3.5-35B-A3B-Unfused" +DST = "/mnt/data2/models/Qwen3.5-35B-A3B-BF16-PACKED" +NUM_LAYERS = 40 +NUM_EXPERTS = 256 +HIDDEN = 2048 +INTER = 512 + +os.makedirs(DST, exist_ok=True) + +# 收集所有 shard 文件 +shards = sorted(glob.glob(os.path.join(SRC, "model.safetensors-*"))) +print(f"Found {len(shards)} shards") + +# 构建 tensor 索引：tensor_name -> shard_path +tensor_index = {} +for s in shards: + f = safe_open(s, "pt") + for k in f.keys(): + if "experts.experts" in k: + tensor_index[k] = s +print(f"Indexed {len(tensor_index)} expert tensors") + +for layer in range(NUM_LAYERS): + # 收集该层所有专家的 gate/up/down + gate_list = [] + up_list = [] + down_list = [] + for eid in range(NUM_EXPERTS): + g_key = f"model.layers.{layer}.mlp.experts.experts.{eid}.gate_proj.weight" + u_key = f"model.layers.{layer}.mlp.experts.experts.{eid}.up_proj.weight" + d_key = f"model.layers.{layer}.mlp.experts.experts.{eid}.down_proj.weight" + if g_key not in tensor_index: + print(f"WARNING: missing {g_key}") + continue + f = safe_open(tensor_index[g_key], "pt") + gate_list.append(f.get_tensor(g_key)) # [I, H] + f = safe_open(tensor_index[u_key], "pt") + up_list.append(f.get_tensor(u_key)) # [I, H] + f = safe_open(tensor_index[d_key], "pt") + down_list.append(f.get_tensor(d_key)) # [H, I] + + # gate_up: [E, 2*I, H] — gate 在前，up 在后 + gate_stack = torch.stack(gate_list) # [E, I, H] + up_stack = torch.stack(up_list) # [E, I, H] + gate_up = torch.cat([gate_stack, up_stack], dim=1) # [E, 2*I, H] + + # down: [E, H, I] + down = torch.stack(down_list) # [E, H, I] + + # 保存 + out_file = os.path.join(DST, f"layer_{layer:02d}.safetensors") + save_file({ + f"model.layers.{layer}.mlp.experts.gate_up_proj": gate_up.contiguous(), + f"model.layers.{layer}.mlp.experts.down_proj": down.contiguous(), + }, out_file) + print(f"Layer {layer}: saved {gate_up.shape} + {down.shape} -> {out_file}") + + # 释放内存 + del gate_list, up_list, down_list, gate_stack, up_stack, gate_up, down + +# 复制 config.json +import shutil +for fname in ["config.json", "generation_config.json"]: + src_f = os.path.join(SRC, fname) + if os.path.exists(src_f): + shutil.copy2(src_f, os.path.join(DST, fname)) +print(f"\nDone! Output: {DST}") diff --git a/scripts/kt_bench_py.py b/scripts/kt_bench_py.py new file mode 100644 index 000000000..ba704cde1 --- /dev/null +++ b/scripts/kt_bench_py.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +"""Standard kt (non-MESH) benchmark with memory sampling — Python version. +Usage: python3 kt_bench_py.py +Example: python3 kt_bench_py.py 35b 10004 0,1,2,3 + python3 kt_bench_py.py 397b 10004 0,1,2,3 +""" +import subprocess, time, json, threading, requests, sys, os + +MODEL = sys.argv[1] # 35b / 397b +PORT = int(sys.argv[2]) +CUDA = sys.argv[3] + +if MODEL == "35b": + SCRIPT = "run_kt_35b.sh" +elif MODEL == "397b": + SCRIPT = "run_kt_397b.sh" +else: + print(f"Unknown model: {MODEL}") + sys.exit(1) + +RESULT_FILE = f"/tmp/kt_bench_{MODEL}_results.txt" +with open(RESULT_FILE, "w") as f: + f.write("mode | decode_tok_s | prefill_tok_s | prefill_s | peak_sys_used_gib | peak_gpu_gib\n") + +URL = f"http://localhost:{PORT}/v1/completions" + + +def run(cmd, timeout=30): + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) + + +def wait_ready(log_file, server_pid, max_wait=900): + """Wait for server to be ready.""" + for i in range(max_wait // 5): + time.sleep(5) + try: + with open(log_file) as f: + content = f.read() + if "fired up and ready to roll" in content: + return True + for fatal in ["Segmentation fault", "CUDA error", "RuntimeError:", + "Traceback (most recent call last):", "core dumped", "Out of memory"]: + if fatal in content: + print(f"FATAL: {fatal}") + return False + except: + pass + try: + os.kill(server_pid, 0) + except ProcessLookupError: + print("Server process exited") + return False + return False + + +def sample_memory(stop_event, samples): + """Background thread to sample memory every 2 seconds.""" + while not stop_event.is_set(): + try: + out = subprocess.check_output(["free", "-m"], text=True) + for line in out.splitlines(): + if line.startswith("Mem:"): + used = int(line.split()[2]) + samples.append(("mem", used)) + break + except: + pass + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], + text=True + ) + gpu_sum = sum(int(x) for x in out.strip().split("\n")[:4]) + samples.append(("gpu", gpu_sum)) + except: + pass + time.sleep(2) + + +def bench_prefill(): + """Benchmark prefill with 2048-token prompt (CPU path, < 4096 threshold).""" + prompt = "The quick brown fox jumps over the lazy dog. " * 260 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=300) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def bench_decode(): + """Benchmark decode with 256 tokens.""" + data = {"model": "default", "prompt": "Write a short essay about artificial intelligence and its impact on society:", + "max_tokens": 256, "temperature": 0, "stream": True} + start = time.time() + first = None + cnt = 0 + r = requests.post(URL, json=data, stream=True, timeout=120) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + chunk = json.loads(line) + if chunk["choices"][0]["text"]: + cnt += 1 + if first is None: + first = time.time() + end = time.time() + dec = end - first if first else end - start + dtps = cnt / dec if dec > 0 else 0 + return dtps + + +print(f"\n========== Testing standard kt {MODEL} (no MESH) ==========") + +# Stop old server +run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) +time.sleep(5) + +# Start server (no cap argument for standard kt) +log_file = f"/tmp/sglang_kt_{MODEL}.log" +server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {PORT} {CUDA}" +proc = subprocess.Popen(server_cmd, shell=True, stdout=open(log_file, "w"), stderr=subprocess.STDOUT) +print(f"Started server PID={proc.pid}, waiting for ready...") + +if not wait_ready(log_file, proc.pid): + print(f"FAILED: Server not ready") + with open(RESULT_FILE, "a") as f: + f.write(f"standard_kt | FAILED | FAILED | FAILED | FAILED | FAILED\n") + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + sys.exit(1) + +print(f"Server ready") +time.sleep(5) + +# Start memory sampling +samples = [] +stop_event = threading.Event() +mem_thread = threading.Thread(target=sample_memory, args=(stop_event, samples)) +mem_thread.start() +print("Memory sampler started") + +# Prefill benchmark +try: + pt, ptime, ptps = bench_prefill() + print(f"Prefill: {pt} tokens, {ptime:.2f}s, {ptps:.2f} tok/s") +except Exception as e: + print(f"Prefill failed: {e}") + pt, ptime, ptps = 0, 0, 0 + +time.sleep(3) + +# Decode benchmark +try: + dtps = bench_decode() + print(f"Decode: {dtps:.2f} tok/s") +except Exception as e: + print(f"Decode failed: {e}") + dtps = 0 + +# Stop memory sampling +stop_event.set() +mem_thread.join() + +# Calculate peaks +mem_vals = [v for k, v in samples if k == "mem"] +gpu_vals = [v for k, v in samples if k == "gpu"] +peak_mem = max(mem_vals) / 1024 if mem_vals else 0 # GiB +peak_gpu = max(gpu_vals) / 1024 if gpu_vals else 0 # GiB +print(f"Peak sys used: {peak_mem:.2f} GiB, Peak GPU: {peak_gpu:.2f} GiB, Samples: {len(mem_vals)}") + +# Write result +with open(RESULT_FILE, "a") as f: + f.write(f"standard_kt | {dtps:.2f} | {ptps:.2f} | {ptime:.2f} | {peak_mem:.2f} | {peak_gpu:.2f}\n") + +# Stop server +run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) +time.sleep(5) + +print("\n=== Results ===") +with open(RESULT_FILE) as f: + print(f.read()) diff --git a/scripts/mesh_bench_py.py b/scripts/mesh_bench_py.py new file mode 100644 index 000000000..aa8e0a836 --- /dev/null +++ b/scripts/mesh_bench_py.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +"""MESH benchmark with memory sampling — Python version. +Usage: python3 mesh_bench_py.py +Example: python3 mesh_bench_py.py 35b "64 128 192 224" 10004 0,1,2,3 +""" +import subprocess, time, json, threading, requests, sys, os + +MODEL = sys.argv[1] # 35b / 397b +CAP_LIST = [int(x) for x in sys.argv[2].split()] +PORT = int(sys.argv[3]) +CUDA = sys.argv[4] + +if MODEL == "35b": + SCRIPT = "run_mesh_35b_cap.sh" +elif MODEL == "397b": + SCRIPT = "run_mesh_397b_cap.sh" +else: + print(f"Unknown model: {MODEL}") + sys.exit(1) + +RESULT_FILE = f"/tmp/mesh_bench_{MODEL}_results.txt" +with open(RESULT_FILE, "w") as f: + f.write("cap | decode_tok_s | prefill_tok_s | prefill_s | peak_sys_used_gib | peak_gpu_gib\n") + +URL = f"http://localhost:{PORT}/v1/completions" + + +def run(cmd, timeout=30): + return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) + + +def wait_ready(log_file, server_pid, max_wait=600): + """Wait for server to be ready.""" + for i in range(max_wait // 5): + time.sleep(5) + try: + with open(log_file) as f: + content = f.read() + if "fired up and ready to roll" in content: + return True + # Check fatal errors + for fatal in ["Segmentation fault", "CUDA error", "RuntimeError:", + "Traceback (most recent call last):", "core dumped", "Out of memory"]: + if fatal in content: + print(f"FATAL: {fatal}") + return False + except: + pass + # Check if process exited + try: + os.kill(server_pid, 0) + except ProcessLookupError: + print("Server process exited") + return False + return False + + +def sample_memory(stop_event, samples): + """Background thread to sample memory every 2 seconds.""" + while not stop_event.is_set(): + # System memory (free -m) + try: + out = subprocess.check_output(["free", "-m"], text=True) + for line in out.splitlines(): + if line.startswith("Mem:"): + used = int(line.split()[2]) + samples.append(("mem", used)) + break + except: + pass + # GPU memory + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], + text=True + ) + gpu_sum = sum(int(x) for x in out.strip().split("\n")[:4]) + samples.append(("gpu", gpu_sum)) + except: + pass + time.sleep(2) + + +def bench_prefill(): + """Benchmark prefill with 2048-token prompt.""" + prompt = "The quick brown fox jumps over the lazy dog. " * 260 + data = {"model": "default", "prompt": prompt, "max_tokens": 1, "temperature": 0, "stream": False} + start = time.time() + r = requests.post(URL, json=data, timeout=300) + end = time.time() + resp = r.json() + pt = resp.get("usage", {}).get("prompt_tokens", 0) + elapsed = end - start + tps = pt / elapsed if elapsed > 0 else 0 + return pt, elapsed, tps + + +def bench_decode(): + """Benchmark decode with 256 tokens.""" + data = {"model": "default", "prompt": "Write a short essay about artificial intelligence and its impact on society:", + "max_tokens": 256, "temperature": 0, "stream": True} + start = time.time() + first = None + cnt = 0 + r = requests.post(URL, json=data, stream=True, timeout=120) + for line in r.iter_lines(): + if line: + line = line.decode() + if line.startswith("data: "): + line = line[6:] + if line == "[DONE]": + break + chunk = json.loads(line) + if chunk["choices"][0]["text"]: + cnt += 1 + if first is None: + first = time.time() + end = time.time() + dec = end - first if first else end - start + dtps = cnt / dec if dec > 0 else 0 + return dtps + + +for cap in CAP_LIST: + print(f"\n========== Testing {MODEL} cap={cap} ==========") + + # Stop old server + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + + # Start server + log_file = f"/tmp/sglang_mesh_{MODEL}_cap{cap}.log" + server_cmd = f"cd /mnt/data2/tmp/qujing_mesh && bash {SCRIPT} {cap} {PORT} {CUDA}" + proc = subprocess.Popen(server_cmd, shell=True, stdout=open(log_file, "w"), stderr=subprocess.STDOUT) + print(f"Started server PID={proc.pid}, waiting for ready...") + + if not wait_ready(log_file, proc.pid): + print(f"FAILED: Server not ready for cap={cap}") + with open(RESULT_FILE, "a") as f: + f.write(f"{cap} | FAILED | FAILED | FAILED | FAILED | FAILED\n") + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + continue + + print(f"Server ready") + time.sleep(5) + + # Start memory sampling + samples = [] + stop_event = threading.Event() + mem_thread = threading.Thread(target=sample_memory, args=(stop_event, samples)) + mem_thread.start() + print("Memory sampler started") + + # Prefill benchmark + try: + pt, ptime, ptps = bench_prefill() + print(f"Prefill: {pt} tokens, {ptime:.2f}s, {ptps:.2f} tok/s") + except Exception as e: + print(f"Prefill failed: {e}") + pt, ptime, ptps = 0, 0, 0 + + time.sleep(3) + + # Decode benchmark + try: + dtps = bench_decode() + print(f"Decode: {dtps:.2f} tok/s") + except Exception as e: + print(f"Decode failed: {e}") + dtps = 0 + + # Stop memory sampling + stop_event.set() + mem_thread.join() + + # Calculate peaks + mem_vals = [v for k, v in samples if k == "mem"] + gpu_vals = [v for k, v in samples if k == "gpu"] + peak_mem = max(mem_vals) / 1024 if mem_vals else 0 # GiB + peak_gpu = max(gpu_vals) / 1024 if gpu_vals else 0 # GiB + print(f"Peak sys used: {peak_mem:.2f} GiB, Peak GPU: {peak_gpu:.2f} GiB, Samples: {len(mem_vals)}") + + # Write result + with open(RESULT_FILE, "a") as f: + f.write(f"{cap} | {dtps:.2f} | {ptps:.2f} | {ptime:.2f} | {peak_mem:.2f} | {peak_gpu:.2f}\n") + + # Stop server + run("pkill -f 'sglang.launch_server.*{}'".format(PORT)) + time.sleep(5) + +print("\n=== Results ===") +with open(RESULT_FILE) as f: + print(f.read()) diff --git a/scripts/run_full_bench.sh b/scripts/run_full_bench.sh new file mode 100644 index 000000000..7ea291c6a --- /dev/null +++ b/scripts/run_full_bench.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# 完整 benchmark 批量测试脚本 +# 等 GPU 空闲后执行所有测试组合 +# 用法: bash run_full_bench.sh +# 例: bash run_full_bench.sh 0,1,2,3 + +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +cd /mnt/data2/tmp/qujing_mesh + +CUDA=${1:-0,1,2,3} +PORT=10004 + +echo "============================================" +echo "Full Benchmark Suite" +echo "CUDA: $CUDA" +echo "Start: $(date)" +echo "============================================" + +# 35B 测试 +echo "" +echo "=== 35B AMXINT4 标准 kt ===" +python3 -u bench_full_py.py kt 35b amxint4 $PORT $CUDA + +echo "" +echo "=== 35B AMXINT4 MESH ===" +python3 -u bench_full_py.py mesh 35b amxint4 "64 128 192 224" $PORT $CUDA + +echo "" +echo "=== 35B BF16 标准 kt ===" +python3 -u bench_full_py.py kt 35b bf16 $PORT $CUDA + +echo "" +echo "=== 35B BF16 MESH ===" +python3 -u bench_full_py.py mesh 35b bf16 "64 128 192 224" $PORT $CUDA + +# 397B 测试 +echo "" +echo "=== 397B AMXINT4 标准 kt ===" +python3 -u bench_full_py.py kt 397b amxint4 $PORT $CUDA + +echo "" +echo "=== 397B AMXINT4 MESH ===" +python3 -u bench_full_py.py mesh 397b amxint4 "128 192 256 480" $PORT $CUDA + +echo "" +echo "=== 397B BF16 标准 kt ===" +python3 -u bench_full_py.py kt 397b bf16 $PORT $CUDA + +echo "" +echo "=== 397B BF16 MESH ===" +python3 -u bench_full_py.py mesh 397b bf16 "128 192 256 480" $PORT $CUDA + +echo "" +echo "============================================" +echo "All tests completed: $(date)" +echo "============================================" +echo "" +echo "=== Results Summary ===" +for f in /tmp/full_bench_*_results.txt; do + echo "--- $f ---" + cat "$f" + echo "" +done diff --git a/scripts/run_kt_35b.sh b/scripts/run_kt_35b.sh new file mode 100644 index 000000000..396512119 --- /dev/null +++ b/scripts/run_kt_35b.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# 标准 kt 35B AMXINT4 benchmark（非 MESH）— 与 MESH 对照 +# 用法: run_kt_35b.sh +# GE=32, TP=4, 开 cudagraph +PORT=$1 +CUDA=$2 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_NUM_GPU_EXPERTS=32 + +# 禁用 MESH：删除配置文件（否则 experts.py 会回退读取 /tmp/kt_mesh_config.json 激活 MESH） +rm -f /tmp/kt_mesh_config.json +unset KT_ENABLE_MESH KT_MESH_CAP KT_MESH_TOTAL_LAYERS KT_MESH_WEIGHT_TYPE + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ \ + --kt-weight-path /mnt/data2/models/Qwen3.5-35B-A3B-AMXINT4-NUMA2-MESH/ \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 \ + --kt-max-deferred-experts-per-token 3 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_kt_35b_bf16.sh b/scripts/run_kt_35b_bf16.sh new file mode 100644 index 000000000..043ec99a5 --- /dev/null +++ b/scripts/run_kt_35b_bf16.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# 标准 kt 35B BF16 benchmark（非 MESH）— 与 MESH bf16 对照 +# 用法: run_kt_35b_bf16.sh +# GE=32, TP=4, 开 cudagraph +PORT=$1 +CUDA=$2 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_NUM_GPU_EXPERTS=32 + +# 禁用 MESH：删除配置文件（否则 experts.py 会回退读取 /tmp/kt_mesh_config.json 激活 MESH） +rm -f /tmp/kt_mesh_config.json +unset KT_ENABLE_MESH KT_MESH_CAP KT_MESH_TOTAL_LAYERS KT_MESH_WEIGHT_TYPE + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ \ + --kt-weight-path /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method BF16 --kt-gpu-prefill-token-threshold 4096 \ + --kt-max-deferred-experts-per-token 3 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_kt_397b.sh b/scripts/run_kt_397b.sh new file mode 100644 index 000000000..898041659 --- /dev/null +++ b/scripts/run_kt_397b.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# 标准 kt 397B AMXINT4 benchmark（非 MESH）— 与 MESH 对照 +# 用法: run_kt_397b.sh +# GE=32, TP=4, 开 cudagraph +PORT=$1 +CUDA=$2 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_NUM_GPU_EXPERTS=32 + +# 禁用 MESH：删除配置文件（否则 experts.py 会回退读取 /tmp/kt_mesh_config.json 激活 MESH） +rm -f /tmp/kt_mesh_config.json +unset KT_ENABLE_MESH KT_MESH_CAP KT_MESH_TOTAL_LAYERS KT_MESH_WEIGHT_TYPE + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-weight-path /mnt/data2/models/Qwen3.5-397B-A17B-AMXINT4-NUMA2-MESH-FIXED \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 \ + --kt-max-deferred-experts-per-token 3 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_kt_397b_bf16.sh b/scripts/run_kt_397b_bf16.sh new file mode 100644 index 000000000..e26eb9357 --- /dev/null +++ b/scripts/run_kt_397b_bf16.sh @@ -0,0 +1,31 @@ +#!/bin/bash +# 标准 kt 397B BF16 benchmark（非 MESH）— 与 MESH bf16 对照 +# 用法: run_kt_397b_bf16.sh +# GE=32, TP=4, 开 cudagraph +PORT=$1 +CUDA=$2 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_NUM_GPU_EXPERTS=32 + +# 禁用 MESH：删除配置文件（否则 experts.py 会回退读取 /tmp/kt_mesh_config.json 激活 MESH） +rm -f /tmp/kt_mesh_config.json +unset KT_ENABLE_MESH KT_MESH_CAP KT_MESH_TOTAL_LAYERS KT_MESH_WEIGHT_TYPE + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-weight-path /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method BF16 --kt-gpu-prefill-token-threshold 4096 \ + --kt-max-deferred-experts-per-token 3 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_mesh_35b_bf16.sh b/scripts/run_mesh_35b_bf16.sh new file mode 100644 index 000000000..575e2b1b6 --- /dev/null +++ b/scripts/run_mesh_35b_bf16.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# MESH 35B BF16 benchmark — packed bf16 格式 +# 用法: run_mesh_35b_bf16.sh +# CAP: slot 容量 (64/128/192/224), 224=full +# GE=32, TP=4, 开 cudagraph +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=40 +export KT_MESH_WEIGHT_TYPE=bf16 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap) + ' weight_type=' + str(cfg.weight_type)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused/ \ + --kt-weight-path /mnt/data2/models/Qwen3.5-35B-A3B-BF16-PACKED \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method BF16 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_mesh_35b_cap.sh b/scripts/run_mesh_35b_cap.sh index 8f2f3ab7f..faed2eb9a 100644 --- a/scripts/run_mesh_35b_cap.sh +++ b/scripts/run_mesh_35b_cap.sh @@ -10,6 +10,7 @@ source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache export SGLANG_DISABLE_CUDNN_CHECK=1 +export SGLANG_DISABLE_BALANCE_CHECK=1 export CUDA_VISIBLE_DEVICES=$CUDA export KT_ENABLE_MESH=1 export KT_MESH_CAP=$CAP diff --git a/scripts/run_mesh_397b_bf16.sh b/scripts/run_mesh_397b_bf16.sh new file mode 100644 index 000000000..66af36cf2 --- /dev/null +++ b/scripts/run_mesh_397b_bf16.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# MESH 397B BF16 benchmark — packed bf16 格式（model.language_model.layers 前缀） +# 用法: run_mesh_397b_bf16.sh +# CAP: slot 容量 (128/192/256/480), 480=full +# GE=32, TP=4, 开 cudagraph +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=60 +export KT_MESH_WEIGHT_TYPE=bf16 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap) + ' weight_type=' + str(cfg.weight_type)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-weight-path /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-cpuinfer 153 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method BF16 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 4 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/scripts/run_mesh_397b_cap.sh b/scripts/run_mesh_397b_cap.sh index 94dfee3fd..9bcaf9b8d 100644 --- a/scripts/run_mesh_397b_cap.sh +++ b/scripts/run_mesh_397b_cap.sh @@ -36,4 +36,4 @@ exec python -m sglang.launch_server \ --max-running-requests 16 --max-total-tokens 20000 \ --watchdog-timeout 3000 --enable-mixed-chunk \ --tensor-parallel-size 4 --enable-p2p-check \ - --disable-shared-experts-fusion --language-only + --disable-shared-experts-fusion diff --git a/scripts/run_mesh_bench.sh b/scripts/run_mesh_bench.sh index ab3062261..5bb4c9828 100644 --- a/scripts/run_mesh_bench.sh +++ b/scripts/run_mesh_bench.sh @@ -5,6 +5,8 @@ # CAP_LIST: 空格分隔的 cap 值，如 "64 128 192 224" # PORT: 端口 # CUDA: CUDA 设备列表，如 0,1,2,3 +# +# 输出指标：decode tok/s, prefill tok/s, peak RSS, peak GPU mem, peak anon, peak file MODEL=$1 CAP_LIST=$2 @@ -13,8 +15,10 @@ CUDA=$4 if [ "$MODEL" = "35b" ]; then SCRIPT="run_mesh_35b_cap.sh" + PREFILL_TOKENS=8192 elif [ "$MODEL" = "397b" ]; then SCRIPT="run_mesh_397b_cap.sh" + PREFILL_TOKENS=8192 else echo "Unknown model: $MODEL (use 35b or 397b)" exit 1 @@ -22,7 +26,7 @@ fi RESULT_FILE="/tmp/mesh_bench_${MODEL}_results.txt" echo "=== MESH ${MODEL} benchmark results ===" > $RESULT_FILE -echo "cap | tokens | total_s | prefill_s | decode_s | decode_tok_s" >> $RESULT_FILE +echo "cap | decode_tok_s | prefill_tok_s | prefill_s | peak_rss_gib | peak_sys_used_gib | peak_gpu_gib | peak_anon_gib | peak_file_gib" >> $RESULT_FILE for CAP in $CAP_LIST; do echo "" @@ -30,7 +34,7 @@ for CAP in $CAP_LIST; do # 停旧服务 pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null - sleep 3 + sleep 5 # 启动新服务 cd /mnt/data2/tmp/qujing_mesh @@ -47,14 +51,13 @@ for CAP in $CAP_LIST; do echo "Server ready after $((i*5))s" break fi - # 只检测真正致命的错误，忽略 "import error"/"side-effect import failed" 等正常警告 + # 只检测真正致命的错误 if grep -qi "Segmentation fault\|CUDA error\|RuntimeError:\|AssertionError:\|Traceback (most recent call last):\|core dumped\|Out of memory" /tmp/sglang_mesh_${MODEL}_cap${CAP}.log 2>/dev/null; then echo "FATAL error detected in log:" grep -i "Segmentation fault\|CUDA error\|RuntimeError:\|AssertionError:\|Traceback (most recent call last):\|core dumped\|Out of memory" /tmp/sglang_mesh_${MODEL}_cap${CAP}.log | tail -5 READY=0 break fi - # 检查服务进程是否已退出 if ! kill -0 $SERVER_PID 2>/dev/null; then echo "Server process exited unexpectedly" tail -20 /tmp/sglang_mesh_${MODEL}_cap${CAP}.log @@ -65,31 +68,83 @@ for CAP in $CAP_LIST; do if [ $READY -eq 0 ]; then echo "FAILED: Server not ready for cap=${CAP}" - echo "${CAP} | FAILED" >> $RESULT_FILE + echo "${CAP} | FAILED | FAILED | FAILED | FAILED | FAILED | FAILED | FAILED | FAILED" >> $RESULT_FILE pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null - sleep 3 + sleep 5 continue fi - sleep 3 # 额外等待稳定 + sleep 5 # 额外等待稳定 - # 功能测试 - FUNC_OUT=$(curl -s http://localhost:${PORT}/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"default","prompt":"Hello, my name is","max_tokens":16,"temperature":0}' | \ - python3 -c "import json,sys; r=json.load(sys.stdin); print(r['choices'][0]['text'][:50])" 2>/dev/null) - echo "Function test output: ${FUNC_OUT}" + # 启动内存采样（后台运行，每 2 秒采样一次） + MEM_SAMPLE_FILE="/tmp/mesh_mem_samples_${MODEL}_cap${CAP}.csv" + echo "timestamp,rss_kb,sys_used_mb,sys_avail_mb,anon_kb,file_kb,gpu0_mb,gpu1_mb,gpu2_mb,gpu3_mb" > $MEM_SAMPLE_FILE + ( + while true; do + # 采样所有 sglang 相关进程的 RSS 总和（用 ps 匹配，更可靠） + RSS_TOTAL=0 + ANON_TOTAL=0 + FILE_TOTAL=0 + for pid in $(ps -eo pid,rss,args | grep -i "sglang\|launch_server\|sglang.srt" | grep -v grep | awk '{print $1}'); do + if [ -f "/proc/$pid/status" ]; then + rss=$(grep "VmRSS" /proc/$pid/status 2>/dev/null | awk '{print $2}') + RSS_TOTAL=$((RSS_TOTAL + ${rss:-0})) + fi + if [ -r "/proc/$pid/smaps_rollup" ]; then + anon=$(grep "^Anonymous:" /proc/$pid/smaps_rollup 2>/dev/null | awk '{print $2}') + ANON_TOTAL=$((ANON_TOTAL + ${anon:-0})) + # file-backed = RSS - Anonymous（smaps_rollup 无 File: 字段） + FILE_TOTAL=$((RSS_TOTAL - ANON_TOTAL)) + fi + done + # 系统总内存 + page cache（Cached 列） + SYS_MEM=$(free -m | awk '/^Mem:/ {printf "%d,%d", $3, $7}') + CACHE_MB=$(grep "^Cached:" /proc/meminfo | awk '{print $2}') + # GPU memory + GPU_LINE=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | head -4 | tr '\n' ',') + echo "$(date +%s),${RSS_TOTAL},${SYS_MEM},${ANON_TOTAL},${FILE_TOTAL},${GPU_LINE},cache_mb=${CACHE_MB}" >> $MEM_SAMPLE_FILE + sleep 2 + done + ) & + MEM_SAMPLER_PID=$! + echo "Memory sampler started PID=$MEM_SAMPLER_PID" - # 测速 source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate - SPEED_OUT=$(python3 -c " + + # ===== 1. Prefill 测速（2048 tokens，< 4096 threshold 走 CPU prefill）===== + echo "Running prefill benchmark (2048 tokens, CPU prefill path)..." + PREFILL_OUT=$(python3 -c " +import requests, time, json +# 生成约 2048 token 的 prompt（重复文本，< 4096 threshold 避免触发 GPU prefill bug） +prompt = 'The quick brown fox jumps over the lazy dog. ' * 260 # 约 2080 tokens +url = 'http://localhost:${PORT}/v1/completions' +data = {'model':'default','prompt':prompt,'max_tokens':1,'temperature':0,'stream':False} +start = time.time() +r = requests.post(url, json=data, timeout=300) +end = time.time() +resp = r.json() +# 从 usage 获取 prompt tokens +prompt_tokens = resp.get('usage',{}).get('prompt_tokens',0) +elapsed = end - start +prefill_tps = prompt_tokens / elapsed if elapsed > 0 else 0 +print(f'{prompt_tokens}|{elapsed:.2f}|{prefill_tps:.2f}') +" 2>&1) + echo "Prefill result: ${PREFILL_OUT}" + PREFILL_S=$(echo "$PREFILL_OUT" | cut -d'|' -f2) + PREFILL_TPS=$(echo "$PREFILL_OUT" | cut -d'|' -f3) + + sleep 3 + + # ===== 2. Decode 测速（streaming 256 tokens）===== + echo "Running decode benchmark (256 tokens)..." + DECODE_OUT=$(python3 -c " import requests, time, json url = 'http://localhost:${PORT}/v1/completions' data = {'model':'default','prompt':'Write a short essay about artificial intelligence and its impact on society:','max_tokens':256,'temperature':0,'stream':True} start = time.time() first_token_time = None token_count = 0 -r = requests.post(url, json=data, stream=True) +r = requests.post(url, json=data, stream=True, timeout=120) for line in r.iter_lines(): if line: line = line.decode('utf-8') @@ -107,15 +162,73 @@ total = end - start prefill = first_token_time - start if first_token_time else 0 decode = end - first_token_time if first_token_time else total dtps = token_count/decode if decode > 0 else 0 -print(f'{token_count}|{total:.2f}|{prefill:.2f}|{decode:.2f}|{dtps:.2f}') +print(f'{dtps:.2f}') +" 2>&1) + echo "Decode tok/s: ${DECODE_OUT}" + + # ===== 3. 停止内存采样，计算 peak ===== + kill $MEM_SAMPLER_PID 2>/dev/null + sleep 1 + + # 从 CSV 计算峰值 + PEAK_RSS_KB=$(python3 -c " +import csv +with open('${MEM_SAMPLE_FILE}') as f: + reader = csv.DictReader(f) + vals = [int(row['rss_kb']) for row in reader if row.get('rss_kb')] + print(max(vals) if vals else 0) +" 2>/dev/null) + PEAK_SYS_USED_MB=$(python3 -c " +import csv +with open('${MEM_SAMPLE_FILE}') as f: + reader = csv.DictReader(f) + vals = [int(row['sys_used_mb']) for row in reader if row.get('sys_used_mb')] + print(max(vals) if vals else 0) +" 2>/dev/null) + PEAK_ANON_KB=$(python3 -c " +import csv +with open('${MEM_SAMPLE_FILE}') as f: + reader = csv.DictReader(f) + vals = [int(row['anon_kb']) for row in reader if row.get('anon_kb')] + print(max(vals) if vals else 0) +" 2>/dev/null) + PEAK_FILE_KB=$(python3 -c " +import csv +with open('${MEM_SAMPLE_FILE}') as f: + reader = csv.DictReader(f) + vals = [int(row['file_kb']) for row in reader if row.get('file_kb')] + print(max(vals) if vals else 0) " 2>/dev/null) + PEAK_GPU_MB=$(python3 -c " +import csv +with open('${MEM_SAMPLE_FILE}') as f: + reader = csv.DictReader(f) + vals = [] + for row in reader: + gpu_sum = 0 + for k in ['gpu0_mb','gpu1_mb','gpu2_mb','gpu3_mb']: + v = row.get(k,'0').strip().rstrip(',') + if v: + gpu_sum += int(v) + vals.append(gpu_sum) + print(max(vals) if vals else 0) +" 2>/dev/null) + + # 转换为 GiB + PEAK_RSS_GIB=$(python3 -c "print(f'{${PEAK_RSS_KB:-0}/1024/1024:.2f}')") + PEAK_SYS_USED_GIB=$(python3 -c "print(f'{${PEAK_SYS_USED_MB:-0}/1024:.2f}')") + PEAK_ANON_GIB=$(python3 -c "print(f'{${PEAK_ANON_KB:-0}/1024/1024:.2f}')") + PEAK_FILE_GIB=$(python3 -c "print(f'{${PEAK_FILE_KB:-0}/1024/1024:.2f}')") + PEAK_GPU_GIB=$(python3 -c "print(f'{${PEAK_GPU_MB:-0}/1024:.2f}')") - echo "Speed: ${SPEED_OUT}" - echo "${CAP} | ${SPEED_OUT}" >> $RESULT_FILE + echo "Peak RSS: ${PEAK_RSS_GIB} GiB, Peak SysUsed: ${PEAK_SYS_USED_GIB} GiB, Peak GPU: ${PEAK_GPU_GIB} GiB, Peak Anon: ${PEAK_ANON_GIB} GiB, Peak File: ${PEAK_FILE_GIB} GiB" + + # 写入结果 + echo "${CAP} | ${DECODE_OUT:-0} | ${PREFILL_TPS:-0} | ${PREFILL_S:-0} | ${PEAK_RSS_GIB} | ${PEAK_SYS_USED_GIB} | ${PEAK_GPU_GIB} | ${PEAK_ANON_GIB} | ${PEAK_FILE_GIB}" >> $RESULT_FILE # 停服务 pkill -f "sglang.launch_server.*${PORT}" 2>/dev/null - sleep 3 + sleep 5 done echo "" diff --git a/scripts/wait_gpu_and_run.sh b/scripts/wait_gpu_and_run.sh new file mode 100644 index 000000000..082fd414d --- /dev/null +++ b/scripts/wait_gpu_and_run.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# 等待 GPU 空闲后自动启动完整 benchmark +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +cd /mnt/data2/tmp/qujing_mesh + +while true; do + GPU0=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0 | tr -d ' ') + GPU1=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 1 | tr -d ' ') + GPU2=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 2 | tr -d ' ') + GPU3=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 3 | tr -d ' ') + + echo "[$(date '+%H:%M:%S')] GPU0=${GPU0}MiB GPU1=${GPU1}MiB GPU2=${GPU2}MiB GPU3=${GPU3}MiB" + + if [ "$GPU0" -lt 5000 ] && [ "$GPU1" -lt 5000 ] && [ "$GPU2" -lt 5000 ] && [ "$GPU3" -lt 5000 ]; then + echo "GPU 0,1,2,3 all free! Starting benchmark..." + bash run_full_bench.sh 0,1,2,3 > /tmp/full_bench_output.log 2>&1 + echo "DONE" >> /tmp/full_bench_output.log + exit 0 + fi + + GPU4=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 4 | tr -d ' ') + GPU5=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 5 | tr -d ' ') + GPU6=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 6 | tr -d ' ') + GPU7=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 7 | tr -d ' ') + + if [ "$GPU4" -lt 5000 ] && [ "$GPU5" -lt 5000 ] && [ "$GPU6" -lt 5000 ] && [ "$GPU7" -lt 5000 ]; then + echo "GPU 4,5,6,7 all free! Starting benchmark..." + bash run_full_bench.sh 4,5,6,7 > /tmp/full_bench_output.log 2>&1 + echo "DONE" >> /tmp/full_bench_output.log + exit 0 + fi + + sleep 300 +done From 8f7e704296c2c48b51fb04e530dd366312b31b8d Mon Sep 17 00:00:00 2001 From: RaQiu <2023015520@st.cupk.edu.cn> Date: Wed, 1 Jul 2026 20:06:25 +0800 Subject: [PATCH 10/10] =?UTF-8?q?feat:=20=E6=96=B9=E6=A1=88=20A=20CPU=20ex?= =?UTF-8?q?pert=20load=20balancing=20+=2035B/397B=20BF16=20benchmark=20scr?= =?UTF-8?q?ipts?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - wrapper.py: per-process NUMA shard loading (each tensor_parallel process only loads its own NUMA shard, tp_part_idx=0) - cap_sweep.sh: add cgroup memory tracking, RAID util/read monitoring, "full" label support, full-range taskset - New scripts: cap_sweep_bf16.sh, run_mesh_35b_bf16_cap_tp2.sh, run_mesh_397b_cap_tp2.sh, align_safetensors.py (512-byte alignment) --- align_safetensors.py | 61 ++++++ cap_sweep.sh | 151 ++++++++++--- cap_sweep_bf16.sh | 281 +++++++++++++++++++++++++ kt-kernel/python/utils/mesh/wrapper.py | 135 ++++++------ run_mesh_35b_bf16_cap_tp2.sh | 37 ++++ run_mesh_397b_cap_tp2.sh | 39 ++++ 6 files changed, 607 insertions(+), 97 deletions(-) create mode 100644 align_safetensors.py create mode 100644 cap_sweep_bf16.sh create mode 100644 run_mesh_35b_bf16_cap_tp2.sh create mode 100644 run_mesh_397b_cap_tp2.sh diff --git a/align_safetensors.py b/align_safetensors.py new file mode 100644 index 000000000..8b45569b6 --- /dev/null +++ b/align_safetensors.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +"""Align safetensors header to 512-byte boundary for MESH O_DIRECT.""" +import struct +import os +import sys +import shutil + + +def align_safetensors(src, dst): + with open(src, 'rb') as f: + header_len = struct.unpack(' {size_gib:.2f} GiB, aligned={aligned}') + sys.stdout.flush() + + print('All done!') + + +if __name__ == '__main__': + main() diff --git a/cap_sweep.sh b/cap_sweep.sh index 5e830c8e9..b88e4f990 100644 --- a/cap_sweep.sh +++ b/cap_sweep.sh @@ -1,17 +1,17 @@ #!/bin/bash -# MESH cap sweep benchmark — fine-grained per-rank memory + GPU + utilization. +# MESH cap sweep benchmark — 方案 A 验证：cgroup 内存 + RAID util + GPU + NUMA. # Tests multiple cap values with decode speed + per-rank VRAM/GPU util + per-NUMA memory. # NO pkill -f: uses screen quit + ss+kill + pgrep -P for safe cleanup. # Does NOT touch gpu4 processes (port 30001/30002). # # Naming convention (per user requirement): -# - "tensor_parallel0" / "tensor_parallel1" = the two scheduler processes (was TP0/TP1) +# - "tensor_parallel0" / "tensor_parallel1" = the two scheduler processes # - GPU6 = tensor_parallel0's device, GPU7 = tensor_parallel1's device -# - threadpool0/1 = CPU expert worker pools (not tracked separately here) +# - threadpool0/1 = CPU expert worker pools # -# Usage: bash cap_sweep.sh "32 64 96 128 160 192" +# Usage: bash cap_sweep.sh "32 64 96 128 160 192 full" -CAP_LIST=${1:-"32 64 96 128 160 192"} +CAP_LIST=${1:-"32 64 96 128 160 192 full"} PORT=50052 CUDA=6,7 RANK0_GPU=6 @@ -21,9 +21,11 @@ WORKDIR=/mnt/data2/tmp/qujing_mesh SCRIPT=run_mesh_35b_cap_graph_raid0.sh RESULT=$WORKDIR/cap_sweep_results.txt VENV=$WORKDIR/.venv/bin/python +CGROUP=/sys/fs/cgroup/mesh_sweep +RAID_DEV=md397 -# Header columns (no "TP" abbreviation — use full "tensor_parallel0/1") -echo "cap | decode_tok_s | decode_tokens | total_elapsed | peak_sys_gib | tensor_parallel0_vram_gib | tensor_parallel1_vram_gib | tensor_parallel0_gpu_util_pct | tensor_parallel1_gpu_util_pct | tensor_parallel0_vmhwm_mb | tensor_parallel0_numa0_mb | tensor_parallel0_numa1_mb | tensor_parallel1_vmhwm_mb | tensor_parallel1_numa0_mb | tensor_parallel1_numa1_mb | cpu_load" > $RESULT +# Header columns +echo "cap | decode_tok_s | decode_tokens | total_elapsed | cgroup_peak_gib | peak_sys_gib | tensor_parallel0_vram_gib | tensor_parallel1_vram_gib | tensor_parallel0_gpu_util_pct | tensor_parallel1_gpu_util_pct | tensor_parallel0_vmhwm_mb | tensor_parallel0_numa0_mb | tensor_parallel0_numa1_mb | tensor_parallel1_vmhwm_mb | tensor_parallel1_numa0_mb | tensor_parallel1_numa1_mb | raid_util_pct | raid_read_mbs | cpu_load" > $RESULT stop_server() { echo " Stopping server..." @@ -55,20 +57,64 @@ stop_server() { fi } +# Setup cgroup for memory tracking (needs sudo for cgroup v2 subtree_control) +setup_cgroup() { + # 清理旧 cgroup + sudo rmdir $CGROUP 2>/dev/null + # 启用 memory controller (需要 root) + sudo bash -c 'echo +memory > /sys/fs/cgroup/cgroup.subtree_control' 2>/dev/null + sudo mkdir -p $CGROUP 2>/dev/null + # 让当前用户可读写 cgroup.procs 和 memory.current + sudo chown -R $(id -u):$(id -g) $CGROUP 2>/dev/null + if [ -f $CGROUP/memory.current ]; then + echo " cgroup $CGROUP ready" + else + echo " WARNING: cgroup $CGROUP not available, memory tracking disabled" + fi +} + +# Add process tree to cgroup +add_to_cgroup() { + local SERVER_PID=$1 + if [ ! -f $CGROUP/memory.current ]; then + return + fi + # 直接用 pgrep -f 找所有 sglang 相关进程并加入 cgroup + # 不依赖递归 pgrep -P（sglang 进程树可能较深且复杂） + local ALL_PIDS=$(pgrep -f "sglang|launch_server|scheduler_TP|kt_kernel|python.*serve" 2>/dev/null) + # 也加入 server PID 本身 + ALL_PIDS="$SERVER_PID $ALL_PIDS" + local COUNT=0 + for p in $ALL_PIDS; do + if [ -n "$p" ]; then + echo $p | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + COUNT=$((COUNT + 1)) + fi + done + # 递归加入 server PID 的所有子孙进程 + local RECURSIVE_PIDS=$(pgrep -P $SERVER_PID 2>/dev/null) + for p in $RECURSIVE_PIDS; do + echo $p | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + for GC in $(pgrep -P $p 2>/dev/null); do + echo $GC | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + for GGC in $(pgrep -P $GC 2>/dev/null); do + echo $GGC | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + done + done + done + echo " Added $COUNT processes to cgroup" +} + # Get per-rank process info: VmHWM (MB) + NUMA0 mem (MB) + NUMA1 mem (MB) -# Returns 6 space-separated integers: -# tp0_vmhwm tp0_numa0 tp0_numa1 tp1_vmhwm tp1_numa0 tp1_numa1 get_rank_info() { local SERVER_PID=$1 local TP0_VMHWM=0 TP0_NUMA0=0 TP0_NUMA1=0 local TP1_VMHWM=0 TP1_NUMA0=0 TP1_NUMA1=0 - local CHILDREN=$(pgrep -P $SERVER_PID 2>/dev/null) - for CHILD in $CHILDREN; do - # Use ps -o args= to get full command line (not truncated like /proc/PID/status Name) + # 用 pgrep -f 直接找 scheduler_TP 进程，不依赖进程树 + for CHILD in $(pgrep -f "scheduler_TP[01]" 2>/dev/null); do local ARGS=$(ps -p $CHILD -o args= 2>/dev/null) local HWM=$(grep "^VmHWM:" /proc/$CHILD/status 2>/dev/null | awk '{print $2}') [ -z "$HWM" ] && HWM=0 - # numastat -p PID: take "Total" row, Node0 and Node1 columns (MB, converted to int) local NUMA_OUT=$(numastat -p $CHILD 2>/dev/null | awk '/^Total/ {printf "%d %d", $2, $3}') local N0=$(echo "$NUMA_OUT" | awk '{print $1}') local N1=$(echo "$NUMA_OUT" | awk '{print $2}') @@ -87,30 +133,40 @@ get_rank_info() { echo "$TP0_VMHWM $TP0_NUMA0 $TP0_NUMA1 $TP1_VMHWM $TP1_NUMA0 $TP1_NUMA1" } -# Check CPU load from other processes (djh training etc.) check_cpu_load() { - local DJH_CPU=$(ps -p 1620567 -o %cpu= 2>/dev/null | awk '{printf "%.0f", $1}') local TOTAL_USED=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d% -f1) - echo "djh_cpu=${DJH_CPU}% total_user=${TOTAL_USED}%" + echo "total_user=${TOTAL_USED}%" } -# Per-second sampler: system RAM + per-rank GPU VRAM + per-rank GPU util +# Per-second sampler: cgroup memory + system RAM + per-rank GPU VRAM/util + RAID util MEM_SAMPLES=/tmp/mesh_mem_samples.txt -GPU0_VRAM_SAMPLES=/tmp/mesh_gpu0_vram.txt # GPU6 = tensor_parallel0 -GPU1_VRAM_SAMPLES=/tmp/mesh_gpu1_vram.txt # GPU7 = tensor_parallel1 +CGROUP_SAMPLES=/tmp/mesh_cgroup_samples.txt +GPU0_VRAM_SAMPLES=/tmp/mesh_gpu0_vram.txt +GPU1_VRAM_SAMPLES=/tmp/mesh_gpu1_vram.txt GPU0_UTIL_SAMPLES=/tmp/mesh_gpu0_util.txt GPU1_UTIL_SAMPLES=/tmp/mesh_gpu1_util.txt +RAID_UTIL_SAMPLES=/tmp/mesh_raid_util.txt +RAID_READ_SAMPLES=/tmp/mesh_raid_read.txt sample_loop() { > $MEM_SAMPLES + > $CGROUP_SAMPLES > $GPU0_VRAM_SAMPLES > $GPU1_VRAM_SAMPLES > $GPU0_UTIL_SAMPLES > $GPU1_UTIL_SAMPLES + > $RAID_UTIL_SAMPLES + > $RAID_READ_SAMPLES while true; do # System memory (MB used) free -m | grep "^Mem:" | awk '{print $3}' >> $MEM_SAMPLES - # Query both GPUs at once: columns = memory.used, utilization.gpu + # cgroup memory.current (bytes) + if [ -f $CGROUP/memory.current ]; then + cat $CGROUP/memory.current >> $CGROUP_SAMPLES + else + echo 0 >> $CGROUP_SAMPLES + fi + # GPU VRAM + util local OUT=$(nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader,nounits -i $RANK0_GPU,$RANK1_GPU 2>/dev/null) local L0=$(echo "$OUT" | sed -n '1p') local L1=$(echo "$OUT" | sed -n '2p') @@ -118,18 +174,35 @@ sample_loop() { echo "$L0" | awk -F', ' '{print $2}' >> $GPU0_UTIL_SAMPLES echo "$L1" | awk -F', ' '{print $1}' >> $GPU1_VRAM_SAMPLES echo "$L1" | awk -F', ' '{print $2}' >> $GPU1_UTIL_SAMPLES + # RAID md397 %util + rkB/s (iostat columns: $23=%util, $3=rkB/s) + # 注意 iostat 末尾有空行，必须用 grep 过滤 + local IOSTAT_OUT=$(iostat -dx $RAID_DEV 1 2 2>/dev/null | grep "^$RAID_DEV" | tail -1) + echo "$IOSTAT_OUT" | awk '{print $23}' >> $RAID_UTIL_SAMPLES + echo "$IOSTAT_OUT" | awk '{print $3}' >> $RAID_READ_SAMPLES sleep 1 done } -echo "Config: mesh=ON, tensor_parallel_size=2, GE=32, dual-numa(0,1), defer=3, RAID0 O_DIRECT" -echo "Cap list: $CAP_LIST" +echo "Config: 方案 A (all tensor_parallel ranks compute CPU experts), mesh=ON, tensor_parallel_size=2, GE=32, RAID0 O_DIRECT" +echo "Cap list: $CAP_LIST (full=224, i.e. 256 experts - 32 GPU experts)" echo "GPU6 = tensor_parallel0 device, GPU7 = tensor_parallel1 device" echo "Results -> $RESULT" echo "" -for CAP in $CAP_LIST; do - echo "========== Testing cap=$CAP ==========" +# Setup cgroup once +setup_cgroup + +for CAP_INPUT in $CAP_LIST; do + # "full" -> 224 + if [ "$CAP_INPUT" = "full" ]; then + CAP=224 + CAP_LABEL="full(224)" + else + CAP=$CAP_INPUT + CAP_LABEL=$CAP + fi + + echo "========== Testing cap=$CAP_LABEL ==========" # 1. Stop old server stop_server @@ -137,8 +210,8 @@ for CAP in $CAP_LIST; do # 2. Start new server LOG=/tmp/mesh_sweep_cap${CAP}.log cd $WORKDIR - screen -dmS $SCREEN_NAME bash -c "taskset -c 24-47,72-191 bash $SCRIPT $CAP $PORT $CUDA > $LOG 2>&1" - echo " Started cap=$CAP, waiting for ready..." + screen -dmS $SCREEN_NAME bash -c "taskset -c 0-191 bash $SCRIPT $CAP $PORT $CUDA > $LOG 2>&1" + echo " Started cap=$CAP_LABEL, waiting for ready..." # 3. Wait for ready (up to 10 min) READY=0 @@ -160,17 +233,18 @@ for CAP in $CAP_LIST; do done if [ $READY -eq 0 ]; then - echo " FAILED: cap=$CAP not ready" - echo "$CAP | FAILED | - | - | - | - | - | - | - | - | - | - | - | - | - | $(check_cpu_load)" >> $RESULT + echo " FAILED: cap=$CAP_LABEL not ready" + echo "$CAP_LABEL | FAILED | - | - | - | - | - | - | - | - | - | - | - | - | - | - | - | $(check_cpu_load)" >> $RESULT tail -20 $LOG 2>/dev/null continue fi sleep 5 - # 4. Get server PID (the listening process) + # 4. Get server PID and add to cgroup SERVER_PID=$(ss -tlnp 2>/dev/null | grep ":$PORT " | grep -oP 'pid=\K[0-9]+' | head -1) echo " Server PID: $SERVER_PID" + add_to_cgroup $SERVER_PID # 5. Start sampler (per-second) sample_loop & @@ -190,18 +264,27 @@ for CAP in $CAP_LIST; do kill $SAMPLER_PID 2>/dev/null wait $SAMPLER_PID 2>/dev/null - # 8. Calculate peak memory + GPU metrics + # 8. Calculate peak memory + GPU metrics + cgroup + RAID PEAK_SYS_MB=$(sort -n $MEM_SAMPLES | tail -1) PEAK_SYS_GIB=$(echo "scale=2; ${PEAK_SYS_MB:-0} / 1024" | bc) + # cgroup peak (bytes -> GiB) + PEAK_CGROUP_B=$(sort -n $CGROUP_SAMPLES | tail -1) + PEAK_CGROUP_GIB=$(echo "scale=2; ${PEAK_CGROUP_B:-0} / 1073741824" | bc) + # GPU VRAM PEAK_GPU0_VRAM_MB=$(sort -n $GPU0_VRAM_SAMPLES | tail -1) PEAK_GPU1_VRAM_MB=$(sort -n $GPU1_VRAM_SAMPLES | tail -1) PEAK_GPU0_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU0_VRAM_MB:-0} / 1024" | bc) PEAK_GPU1_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU1_VRAM_MB:-0} / 1024" | bc) - # GPU utilization: take max during decode (peak utilization) + # GPU utilization: max during decode PEAK_GPU0_UTIL=$(sort -n $GPU0_UTIL_SAMPLES | tail -1) PEAK_GPU1_UTIL=$(sort -n $GPU1_UTIL_SAMPLES | tail -1) [ -z "$PEAK_GPU0_UTIL" ] && PEAK_GPU0_UTIL=0 [ -z "$PEAK_GPU1_UTIL" ] && PEAK_GPU1_UTIL=0 + # RAID: max %util + avg read kB/s during decode + PEAK_RAID_UTIL=$(sort -n $RAID_UTIL_SAMPLES | tail -1) + [ -z "$PEAK_RAID_UTIL" ] && PEAK_RAID_UTIL=0 + AVG_RAID_READ_KBS=$(awk '{sum+=$1; n++} END {if(n>0) printf "%.0f", sum/n; else print 0}' $RAID_READ_SAMPLES) + AVG_RAID_READ_MBS=$(echo "scale=1; ${AVG_RAID_READ_KBS:-0} / 1024" | bc) # 9. Get per-rank VmHWM + NUMA breakdown RANK_INFO=$(get_rank_info $SERVER_PID) @@ -215,14 +298,16 @@ for CAP in $CAP_LIST; do # 10. CPU load CPU_LOAD=$(check_cpu_load) - echo " Peak sys: ${PEAK_SYS_GIB}GiB" + echo " Decode: ${DECODE_TPS} tok/s" + echo " cgroup peak: ${PEAK_CGROUP_GIB}GiB | sys peak: ${PEAK_SYS_GIB}GiB" echo " tensor_parallel0 (GPU${RANK0_GPU}): VRAM=${PEAK_GPU0_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU0_UTIL}%" echo " tensor_parallel1 (GPU${RANK1_GPU}): VRAM=${PEAK_GPU1_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU1_UTIL}%" echo " tensor_parallel0 proc: VmHWM=${TP0_VMHWM}MB, NUMA0=${TP0_NUMA0}MB, NUMA1=${TP0_NUMA1}MB" echo " tensor_parallel1 proc: VmHWM=${TP1_VMHWM}MB, NUMA0=${TP1_NUMA0}MB, NUMA1=${TP1_NUMA1}MB" + echo " RAID $RAID_DEV: peak_util=${PEAK_RAID_UTIL}%, avg_read=${AVG_RAID_READ_MBS}MB/s" echo " CPU load: $CPU_LOAD" - echo "$CAP | $DECODE_TPS | $DECODE_TOKENS | $TOTAL_ELAPSED | $PEAK_SYS_GIB | $PEAK_GPU0_VRAM_GIB | $PEAK_GPU1_VRAM_GIB | $PEAK_GPU0_UTIL | $PEAK_GPU1_UTIL | $TP0_VMHWM | $TP0_NUMA0 | $TP0_NUMA1 | $TP1_VMHWM | $TP1_NUMA0 | $TP1_NUMA1 | $CPU_LOAD" >> $RESULT + echo "$CAP_LABEL | $DECODE_TPS | $DECODE_TOKENS | $TOTAL_ELAPSED | $PEAK_CGROUP_GIB | $PEAK_SYS_GIB | $PEAK_GPU0_VRAM_GIB | $PEAK_GPU1_VRAM_GIB | $PEAK_GPU0_UTIL | $PEAK_GPU1_UTIL | $TP0_VMHWM | $TP0_NUMA0 | $TP0_NUMA1 | $TP1_VMHWM | $TP1_NUMA0 | $TP1_NUMA1 | $PEAK_RAID_UTIL | $AVG_RAID_READ_MBS | $CPU_LOAD" >> $RESULT echo "" done diff --git a/cap_sweep_bf16.sh b/cap_sweep_bf16.sh new file mode 100644 index 000000000..7b207af9b --- /dev/null +++ b/cap_sweep_bf16.sh @@ -0,0 +1,281 @@ +#!/bin/bash +# MESH cap sweep benchmark — 35B BF16 原版 +# Tests cap values with decode speed + per-rank VRAM/GPU util + per-NUMA memory. +# Uses md398 (3-NVMe raid0) for BF16 aligned weights. +# Naming: tensor_parallel0 (GPU6) / tensor_parallel1 (GPU7) +# +# Usage: bash cap_sweep_bf16.sh "32 64 96 128 160 192 full" + +CAP_LIST=${1:-"32 64 96 128 160 192 full"} +PORT=50052 +CUDA=6,7 +RANK0_GPU=6 +RANK1_GPU=7 +SCREEN_NAME="mesh_sweep_bf16" +WORKDIR=/mnt/data2/tmp/qujing_mesh +SCRIPT=run_mesh_35b_bf16_cap_tp2.sh +RESULT=$WORKDIR/cap_sweep_bf16_results.txt +VENV=$WORKDIR/.venv/bin/python +CGROUP=/sys/fs/cgroup/mesh_sweep_bf16 +RAID_DEV=md398 + +# Header columns +echo "cap | decode_tok_s | decode_tokens | total_elapsed | cgroup_peak_gib | peak_sys_gib | tensor_parallel0_vram_gib | tensor_parallel1_vram_gib | tensor_parallel0_gpu_util_pct | tensor_parallel1_gpu_util_pct | tensor_parallel0_vmhwm_mb | tensor_parallel0_numa0_mb | tensor_parallel0_numa1_mb | tensor_parallel1_vmhwm_mb | tensor_parallel1_numa0_mb | tensor_parallel1_numa1_mb | raid_util_pct | raid_read_mbs | cpu_load" > $RESULT + +stop_server() { + echo " Stopping server..." + screen -X -S $SCREEN_NAME quit 2>/dev/null + sleep 3 + local PID=$(ss -tlnp 2>/dev/null | grep ":$PORT " | grep -oP 'pid=\K[0-9]+' | head -1) + if [ -n "$PID" ]; then + local CHILDREN=$(pgrep -P $PID 2>/dev/null) + for CHILD in $CHILDREN; do + local GRANDCHILDREN=$(pgrep -P $CHILD 2>/dev/null) + for GC in $GRANDCHILDREN; do + kill $GC 2>/dev/null + done + kill $CHILD 2>/dev/null + done + kill $PID 2>/dev/null + fi + for i in $(seq 1 15); do + if ! ss -tln 2>/dev/null | grep -q ":$PORT "; then + break + fi + sleep 2 + done + sleep 3 + if ss -tln 2>/dev/null | grep -q ":$PORT "; then + echo " WARNING: port $PORT still in use after cleanup" + else + echo " Port $PORT free" + fi +} + +setup_cgroup() { + sudo rmdir $CGROUP 2>/dev/null + sudo bash -c 'echo +memory > /sys/fs/cgroup/cgroup.subtree_control' 2>/dev/null + sudo mkdir -p $CGROUP 2>/dev/null + sudo chown -R $(id -u):$(id -g) $CGROUP 2>/dev/null + if [ -f $CGROUP/memory.current ]; then + echo " cgroup $CGROUP ready" + else + echo " WARNING: cgroup $CGROUP not available" + fi +} + +add_to_cgroup() { + local SERVER_PID=$1 + if [ ! -f $CGROUP/memory.current ]; then + return + fi + local ALL_PIDS=$(pgrep -f "sglang|launch_server|scheduler_TP|kt_kernel|python.*serve" 2>/dev/null) + ALL_PIDS="$SERVER_PID $ALL_PIDS" + local COUNT=0 + for p in $ALL_PIDS; do + if [ -n "$p" ]; then + echo $p | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + COUNT=$((COUNT + 1)) + fi + done + local RECURSIVE_PIDS=$(pgrep -P $SERVER_PID 2>/dev/null) + for p in $RECURSIVE_PIDS; do + echo $p | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + for GC in $(pgrep -P $p 2>/dev/null); do + echo $GC | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + for GGC in $(pgrep -P $GC 2>/dev/null); do + echo $GGC | sudo tee $CGROUP/cgroup.procs > /dev/null 2>&1 + done + done + done + echo " Added $COUNT processes to cgroup" +} + +get_rank_info() { + local SERVER_PID=$1 + local TP0_VMHWM=0 TP0_NUMA0=0 TP0_NUMA1=0 + local TP1_VMHWM=0 TP1_NUMA0=0 TP1_NUMA1=0 + for CHILD in $(pgrep -f "scheduler_TP[01]" 2>/dev/null); do + local ARGS=$(ps -p $CHILD -o args= 2>/dev/null) + local HWM=$(grep "^VmHWM:" /proc/$CHILD/status 2>/dev/null | awk '{print $2}') + [ -z "$HWM" ] && HWM=0 + local NUMA_OUT=$(numastat -p $CHILD 2>/dev/null | awk '/^Total/ {printf "%d %d", $2, $3}') + local N0=$(echo "$NUMA_OUT" | awk '{print $1}') + local N1=$(echo "$NUMA_OUT" | awk '{print $2}') + [ -z "$N0" ] && N0=0 + [ -z "$N1" ] && N1=0 + if echo "$ARGS" | grep -q "scheduler_TP0"; then + TP0_VMHWM=$((HWM / 1024)) + TP0_NUMA0=$N0 + TP0_NUMA1=$N1 + elif echo "$ARGS" | grep -q "scheduler_TP1"; then + TP1_VMHWM=$((HWM / 1024)) + TP1_NUMA0=$N0 + TP1_NUMA1=$N1 + fi + done + echo "$TP0_VMHWM $TP0_NUMA0 $TP0_NUMA1 $TP1_VMHWM $TP1_NUMA0 $TP1_NUMA1" +} + +check_cpu_load() { + local TOTAL_USED=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d% -f1) + echo "total_user=${TOTAL_USED}%" +} + +MEM_SAMPLES=/tmp/mesh_bf16_mem.txt +CGROUP_SAMPLES=/tmp/mesh_bf16_cgroup.txt +GPU0_VRAM_SAMPLES=/tmp/mesh_bf16_gpu0_vram.txt +GPU1_VRAM_SAMPLES=/tmp/mesh_bf16_gpu1_vram.txt +GPU0_UTIL_SAMPLES=/tmp/mesh_bf16_gpu0_util.txt +GPU1_UTIL_SAMPLES=/tmp/mesh_bf16_gpu1_util.txt +RAID_UTIL_SAMPLES=/tmp/mesh_bf16_raid_util.txt +RAID_READ_SAMPLES=/tmp/mesh_bf16_raid_read.txt + +sample_loop() { + > $MEM_SAMPLES + > $CGROUP_SAMPLES + > $GPU0_VRAM_SAMPLES + > $GPU1_VRAM_SAMPLES + > $GPU0_UTIL_SAMPLES + > $GPU1_UTIL_SAMPLES + > $RAID_UTIL_SAMPLES + > $RAID_READ_SAMPLES + while true; do + free -m | grep "^Mem:" | awk '{print $3}' >> $MEM_SAMPLES + if [ -f $CGROUP/memory.current ]; then + cat $CGROUP/memory.current >> $CGROUP_SAMPLES + else + echo 0 >> $CGROUP_SAMPLES + fi + local OUT=$(nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader,nounits -i $RANK0_GPU,$RANK1_GPU 2>/dev/null) + local L0=$(echo "$OUT" | sed -n '1p') + local L1=$(echo "$OUT" | sed -n '2p') + echo "$L0" | awk -F', ' '{print $1}' >> $GPU0_VRAM_SAMPLES + echo "$L0" | awk -F', ' '{print $2}' >> $GPU0_UTIL_SAMPLES + echo "$L1" | awk -F', ' '{print $1}' >> $GPU1_VRAM_SAMPLES + echo "$L1" | awk -F', ' '{print $2}' >> $GPU1_UTIL_SAMPLES + local IOSTAT_OUT=$(iostat -dx $RAID_DEV 1 2 2>/dev/null | grep "^$RAID_DEV" | tail -1) + echo "$IOSTAT_OUT" | awk '{print $23}' >> $RAID_UTIL_SAMPLES + echo "$IOSTAT_OUT" | awk '{print $3}' >> $RAID_READ_SAMPLES + sleep 1 + done +} + +echo "Config: 35B BF16 原版, 方案 A, mesh=ON, tensor_parallel=2, GE=32, RAID0 (md398) O_DIRECT" +echo "Cap list: $CAP_LIST (full=224, i.e. 256 experts - 32 GPU experts)" +echo "GPU6 = tensor_parallel0 device, GPU7 = tensor_parallel1 device" +echo "Results -> $RESULT" +echo "" + +setup_cgroup + +for CAP_INPUT in $CAP_LIST; do + if [ "$CAP_INPUT" = "full" ]; then + CAP=224 + CAP_LABEL="full(224)" + else + CAP=$CAP_INPUT + CAP_LABEL=$CAP + fi + + echo "========== Testing cap=$CAP_LABEL ==========" + + stop_server + + LOG=/tmp/mesh_sweep_bf16_cap${CAP}.log + cd $WORKDIR + screen -dmS $SCREEN_NAME bash -c "taskset -c 0-191 bash $SCRIPT $CAP $PORT $CUDA > $LOG 2>&1" + echo " Started cap=$CAP_LABEL, waiting for ready..." + + READY=0 + for i in $(seq 1 120); do + sleep 5 + if grep -q "fired up and ready to roll" $LOG 2>/dev/null; then + READY=1 + echo " Server ready after $((i*5))s" + break + fi + if grep -q "scheduler is dead" $LOG 2>/dev/null; then + echo " FATAL: scheduler dead" + break + fi + if grep -q "Segmentation fault" $LOG 2>/dev/null; then + echo " FATAL: segfault" + break + fi + done + + if [ $READY -eq 0 ]; then + echo " FAILED: cap=$CAP_LABEL not ready" + echo "$CAP_LABEL | FAILED | - | - | - | - | - | - | - | - | - | - | - | - | - | - | - | - | $(check_cpu_load)" >> $RESULT + tail -20 $LOG 2>/dev/null + continue + fi + + sleep 5 + + SERVER_PID=$(ss -tlnp 2>/dev/null | grep ":$PORT " | grep -oP 'pid=\K[0-9]+' | head -1) + echo " Server PID: $SERVER_PID" + add_to_cgroup $SERVER_PID + + sample_loop & + SAMPLER_PID=$! + echo " Sampler PID: $SAMPLER_PID" + + echo " Running decode benchmark..." + BENCH_OUT=$($VENV $WORKDIR/bench_decode_stream.py $PORT 512 2>&1) + echo " $BENCH_OUT" + + DECODE_TPS=$(echo "$BENCH_OUT" | grep "decode_tok_s:" | awk '{print $2}') + DECODE_TOKENS=$(echo "$BENCH_OUT" | grep "decode_tokens:" | awk '{print $2}') + TOTAL_ELAPSED=$(echo "$BENCH_OUT" | grep "total_elapsed:" | awk '{print $2}') + + kill $SAMPLER_PID 2>/dev/null + wait $SAMPLER_PID 2>/dev/null + + PEAK_SYS_MB=$(sort -n $MEM_SAMPLES | tail -1) + PEAK_SYS_GIB=$(echo "scale=2; ${PEAK_SYS_MB:-0} / 1024" | bc) + PEAK_CGROUP_B=$(sort -n $CGROUP_SAMPLES | tail -1) + PEAK_CGROUP_GIB=$(echo "scale=2; ${PEAK_CGROUP_B:-0} / 1073741824" | bc) + PEAK_GPU0_VRAM_MB=$(sort -n $GPU0_VRAM_SAMPLES | tail -1) + PEAK_GPU1_VRAM_MB=$(sort -n $GPU1_VRAM_SAMPLES | tail -1) + PEAK_GPU0_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU0_VRAM_MB:-0} / 1024" | bc) + PEAK_GPU1_VRAM_GIB=$(echo "scale=2; ${PEAK_GPU1_VRAM_MB:-0} / 1024" | bc) + PEAK_GPU0_UTIL=$(sort -n $GPU0_UTIL_SAMPLES | tail -1) + PEAK_GPU1_UTIL=$(sort -n $GPU1_UTIL_SAMPLES | tail -1) + [ -z "$PEAK_GPU0_UTIL" ] && PEAK_GPU0_UTIL=0 + [ -z "$PEAK_GPU1_UTIL" ] && PEAK_GPU1_UTIL=0 + PEAK_RAID_UTIL=$(sort -n $RAID_UTIL_SAMPLES | tail -1) + [ -z "$PEAK_RAID_UTIL" ] && PEAK_RAID_UTIL=0 + AVG_RAID_READ_KBS=$(awk '{sum+=$1; n++} END {if(n>0) printf "%.0f", sum/n; else print 0}' $RAID_READ_SAMPLES) + AVG_RAID_READ_MBS=$(echo "scale=1; ${AVG_RAID_READ_KBS:-0} / 1024" | bc) + + RANK_INFO=$(get_rank_info $SERVER_PID) + TP0_VMHWM=$(echo "$RANK_INFO" | awk '{print $1}') + TP0_NUMA0=$(echo "$RANK_INFO" | awk '{print $2}') + TP0_NUMA1=$(echo "$RANK_INFO" | awk '{print $3}') + TP1_VMHWM=$(echo "$RANK_INFO" | awk '{print $4}') + TP1_NUMA0=$(echo "$RANK_INFO" | awk '{print $5}') + TP1_NUMA1=$(echo "$RANK_INFO" | awk '{print $6}') + + CPU_LOAD=$(check_cpu_load) + + echo " Decode: ${DECODE_TPS} tok/s" + echo " cgroup peak: ${PEAK_CGROUP_GIB}GiB | sys peak: ${PEAK_SYS_GIB}GiB" + echo " tensor_parallel0 (GPU${RANK0_GPU}): VRAM=${PEAK_GPU0_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU0_UTIL}%" + echo " tensor_parallel1 (GPU${RANK1_GPU}): VRAM=${PEAK_GPU1_VRAM_GIB}GiB, gpu_util_peak=${PEAK_GPU1_UTIL}%" + echo " tensor_parallel0 proc: VmHWM=${TP0_VMHWM}MB, NUMA0=${TP0_NUMA0}MB, NUMA1=${TP0_NUMA1}MB" + echo " tensor_parallel1 proc: VmHWM=${TP1_VMHWM}MB, NUMA0=${TP1_NUMA0}MB, NUMA1=${TP1_NUMA1}MB" + echo " RAID $RAID_DEV: peak_util=${PEAK_RAID_UTIL}%, avg_read=${AVG_RAID_READ_MBS}MB/s" + echo " CPU load: $CPU_LOAD" + + echo "$CAP_LABEL | $DECODE_TPS | $DECODE_TOKENS | $TOTAL_ELAPSED | $PEAK_CGROUP_GIB | $PEAK_SYS_GIB | $PEAK_GPU0_VRAM_GIB | $PEAK_GPU1_VRAM_GIB | $PEAK_GPU0_UTIL | $PEAK_GPU1_UTIL | $TP0_VMHWM | $TP0_NUMA0 | $TP0_NUMA1 | $TP1_VMHWM | $TP1_NUMA0 | $TP1_NUMA1 | $PEAK_RAID_UTIL | $AVG_RAID_READ_MBS | $CPU_LOAD" >> $RESULT + + echo "" +done + +stop_server + +echo "" +echo "=== Results ===" +cat $RESULT diff --git a/kt-kernel/python/utils/mesh/wrapper.py b/kt-kernel/python/utils/mesh/wrapper.py index 07d0a09d3..f39b6f471 100644 --- a/kt-kernel/python/utils/mesh/wrapper.py +++ b/kt-kernel/python/utils/mesh/wrapper.py @@ -139,6 +139,10 @@ def __init__( if numa_nodes is None: numa_nodes = list(range(threadpool_count)) + # 方案 A: 记录本进程的 NUMA 节点列表，供 _inject_file_layouts 使用 + # 每进程 numa_nodes = [tp_rank]，例如 tensor_parallel0 → [0]，tensor_parallel1 → [1] + self._numa_nodes = numa_nodes + if MeshMoEWrapper._shared_residency is None: # 第一层：创建并初始化共享 ResidencyManager self._residency = MeshResidencyManager() @@ -294,72 +298,75 @@ def get_fd(fpath: str) -> int: injected = 0 missing = 0 + # 方案 A: 每进程只加载自己 NUMA shard + # self._numa_nodes = [tp_rank]，tensor 名用 numa.{local_numa_id} + # tp_part_idx=0（进程内只有 1 个分片） + local_numa_id = self._numa_nodes[0] if self._numa_nodes else 0 for layer in range(total_layers): - for tp in range(tp_count): - for expert in range(expert_num): - # AMXINT4 NUMA-sharded tensor 名 - gate_w_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{tp}.weight" - gate_s_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{tp}.scale" - up_w_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{tp}.weight" - up_s_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{tp}.scale" - down_w_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{tp}.weight" - down_s_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{tp}.scale" - - # 检查必需的 tensor 是否存在 - required = [gate_w_key, up_w_key, down_w_key] - if not all(k in tensor_map for k in required): - missing += 1 - continue + for expert in range(expert_num): + # AMXINT4 NUMA-sharded tensor 名（用本进程 NUMA ID） + gate_w_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{local_numa_id}.weight" + gate_s_key = f"blk.{layer}.ffn_gate_exps.{expert}.numa.{local_numa_id}.scale" + up_w_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{local_numa_id}.weight" + up_s_key = f"blk.{layer}.ffn_up_exps.{expert}.numa.{local_numa_id}.scale" + down_w_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{local_numa_id}.weight" + down_s_key = f"blk.{layer}.ffn_down_exps.{expert}.numa.{local_numa_id}.scale" + + # 检查必需的 tensor 是否存在 + required = [gate_w_key, up_w_key, down_w_key] + if not all(k in tensor_map for k in required): + missing += 1 + continue - # 获取 gate 布局 - g_fpath, g_off, g_bytes = tensor_map[gate_w_key] - fd = get_fd(g_fpath) - - # 获取 up 布局 - u_fpath, u_off, u_bytes = tensor_map[up_w_key] - if u_fpath != g_fpath: - fd = get_fd(u_fpath) # 不同文件需要不同 fd - - # 获取 down 布局 - d_fpath, d_off, d_bytes = tensor_map[down_w_key] - if d_fpath not in file_fds: - fd = get_fd(d_fpath) - - # scale 偏移（AMXINT4 专用） - g_s_off = g_s_bytes = 0 - u_s_off = u_s_bytes = 0 - d_s_off = d_s_bytes = 0 - if gate_s_key in tensor_map: - _, g_s_off, g_s_bytes = tensor_map[gate_s_key] - if up_s_key in tensor_map: - _, u_s_off, u_s_bytes = tensor_map[up_s_key] - if down_s_key in tensor_map: - _, d_s_off, d_s_bytes = tensor_map[down_s_key] - - # 注意：fd 取 gate 所在文件的 fd。如果 up/down 在不同文件， - # io_uring 的 submit_load 会用 layout.fd 读取所有三个矩阵， - # 这要求三个矩阵在同一文件。safetensors 通常如此（同层同 NUMA 在同一文件）。 - # 如果跨文件，需要后续拆分为多次 submit_load。 - self._residency.set_file_layout( - layer_idx=layer, - tp_part_idx=tp, - expert_id=expert, - fd=fd, - gate_offset=g_off, - up_offset=u_off, - down_offset=d_off, - gate_bytes=g_bytes, - up_bytes=u_bytes, - down_bytes=d_bytes, - gate_scale_offset=g_s_off, - up_scale_offset=u_s_off, - down_scale_offset=d_s_off, - gate_scale_bytes=g_s_bytes, - up_scale_bytes=u_s_bytes, - down_scale_bytes=d_s_bytes, - # AMXINT4 对称量化无 mins，保持 0 - ) - injected += 1 + # 获取 gate 布局 + g_fpath, g_off, g_bytes = tensor_map[gate_w_key] + fd = get_fd(g_fpath) + + # 获取 up 布局 + u_fpath, u_off, u_bytes = tensor_map[up_w_key] + if u_fpath != g_fpath: + fd = get_fd(u_fpath) # 不同文件需要不同 fd + + # 获取 down 布局 + d_fpath, d_off, d_bytes = tensor_map[down_w_key] + if d_fpath not in file_fds: + fd = get_fd(d_fpath) + + # scale 偏移（AMXINT4 专用） + g_s_off = g_s_bytes = 0 + u_s_off = u_s_bytes = 0 + d_s_off = d_s_bytes = 0 + if gate_s_key in tensor_map: + _, g_s_off, g_s_bytes = tensor_map[gate_s_key] + if up_s_key in tensor_map: + _, u_s_off, u_s_bytes = tensor_map[up_s_key] + if down_s_key in tensor_map: + _, d_s_off, d_s_bytes = tensor_map[down_s_key] + + # 注意：fd 取 gate 所在文件的 fd。如果 up/down 在不同文件， + # io_uring 的 submit_load 会用 layout.fd 读取所有三个矩阵， + # 这要求三个矩阵在同一文件。safetensors 通常如此（同层同 NUMA 在同一文件）。 + # 如果跨文件，需要后续拆分为多次 submit_load。 + self._residency.set_file_layout( + layer_idx=layer, + tp_part_idx=0, # 方案 A: 进程内只有 1 个分片 + expert_id=expert, + fd=fd, + gate_offset=g_off, + up_offset=u_off, + down_offset=d_off, + gate_bytes=g_bytes, + up_bytes=u_bytes, + down_bytes=d_bytes, + gate_scale_offset=g_s_off, + up_scale_offset=u_s_off, + down_scale_offset=d_s_off, + gate_scale_bytes=g_s_bytes, + up_scale_bytes=u_s_bytes, + down_scale_bytes=d_s_bytes, + # AMXINT4 对称量化无 mins，保持 0 + ) + injected += 1 logger.info( "_inject_amxint4_layouts: injected %d layouts (missing %d), " diff --git a/run_mesh_35b_bf16_cap_tp2.sh b/run_mesh_35b_bf16_cap_tp2.sh new file mode 100644 index 000000000..e32838a88 --- /dev/null +++ b/run_mesh_35b_bf16_cap_tp2.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# MESH 35B BF16 benchmark — 方案 A, tensor_parallel=2, RAID0 O_DIRECT (md398) +# CAP: slot 容量 (32/64/96/128/160/192/full), full=224 (256 experts - 32 GPU experts) +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=40 +export KT_MESH_WEIGHT_TYPE=bf16 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap) + ', weight_type=' + str(cfg.weight_type)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/tmp/qujing_mesh/Qwen3.5-35B-A3B-Unfused \ + --kt-weight-path /mnt/raid_bf16/Qwen3.5-35B-A3B-BF16-MESH-ALIGNED/ \ + --kt-cpuinfer 48 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method BF16 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 2 --enable-p2p-check \ + --disable-shared-experts-fusion diff --git a/run_mesh_397b_cap_tp2.sh b/run_mesh_397b_cap_tp2.sh new file mode 100644 index 000000000..eea0d88fc --- /dev/null +++ b/run_mesh_397b_cap_tp2.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# MESH 397B AMXINT4 benchmark — 方案 A, tensor_parallel=2, RAID0 O_DIRECT +# 用法: run_mesh_397b_cap_tp2.sh +# CAP: slot 容量 (32/64/96/128/192/256/384/480), 480=full (512 - 32 GPU experts) +# GE=32, TP=2, cpuinfer=48 (每进程 48 AMX 线程匹配 48 物理核心) +CAP=$1 +PORT=$2 +CUDA=$3 +source /mnt/data2/tmp/qujing_mesh/.venv/bin/activate +export TMPDIR=/mnt/data2/tmp/qujing_mesh/.tmp +export TRITON_CACHE_DIR=/mnt/data2/tmp/qujing_mesh/.triton_cache +export SGLANG_DISABLE_CUDNN_CHECK=1 +export CUDA_VISIBLE_DEVICES=$CUDA +export KT_ENABLE_MESH=1 +export KT_MESH_CAP=$CAP +export KT_NUM_GPU_EXPERTS=32 +export KT_MESH_TOTAL_LAYERS=60 +export KT_MESH_WEIGHT_TYPE=amxint4 + +python -c " +from kt_kernel.utils.mesh.config import MeshConfig +cfg = MeshConfig.from_env() +cfg.to_file('/tmp/kt_mesh_config.json') +print('MESH config written: cap=' + str(cfg.cap)) +" + +exec python -m sglang.launch_server \ + --host 0.0.0.0 --port $PORT \ + --model /mnt/data2/models/Qwen3.5-397B-A17B-TEXTONLY \ + --kt-weight-path /mnt/raid397_amxint4/Qwen3.5-397B-A17B-AMXINT4-NUMA2-MESH-FIXED/ \ + --kt-cpuinfer 48 --kt-threadpool-count 2 --kt-num-gpu-experts 32 \ + --kt-method AMXINT4 --kt-gpu-prefill-token-threshold 4096 \ + --kt-enable-dynamic-expert-update --kt-numa-nodes 0 1 \ + --attention-backend flashinfer --trust-remote-code \ + --mem-fraction-static 0.90 --chunked-prefill-size 4096 \ + --max-running-requests 16 --max-total-tokens 20000 \ + --watchdog-timeout 3000 --enable-mixed-chunk \ + --tensor-parallel-size 2 --enable-p2p-check \ + --disable-shared-experts-fusion