# MoE SFT AMX 功能详细设计文档

## 1. 现有架构分析

### 1.1 推理类继承关系

```
AMX_MOE_TP<T>                    # 单 NUMA 节点 MoE 实现
    ↓ 继承
TP_MOE<T>                        # 多 NUMA 节点封装
```

### 1.2 关键类分析

#### AMX_MOE_TP 核心成员

```cpp
template <class T>
class AMX_MOE_TP {
protected:
    // 配置
    GeneralMOEConfig config_;
    int tp_part_idx;  // NUMA 节点 ID

    // 专家路由信息
    std::vector<int> m_local_num_;           // [expert_num] 每个专家的 token 数
    std::vector<std::vector<int>> m_local_pos_;  // [qlen][k] 位置映射
    std::vector<int> m_expert_id_map_;       // 激活的专家 ID 列表

    // 中间缓冲区
    ggml_bf16_t* m_local_input_;             // [max_len * k, hidden_size]
    ggml_bf16_t* m_local_gate_output_;       // [max_len * k, intermediate_size]
    ggml_bf16_t* m_local_up_output_;         // [max_len * k, intermediate_size]
    ggml_bf16_t* m_local_down_output_;       // [max_len * k, hidden_size]

    // GEMM Buffer 对象
    std::vector<T::BufferA*> gate_up_ba_;    // gate/up 输入 buffer
    std::vector<T::BufferB*> gate_bb_, up_bb_;  // gate/up 权重 buffer
    std::vector<T::BufferC*> gate_bc_, up_bc_;  // gate/up 输出 buffer
    std::vector<T::BufferA*> down_ba_;       // down 输入 buffer
    std::vector<T::BufferC*> down_bc_;       // down 输出 buffer
};
```

#### TP_MOE 核心成员

```cpp
template <class T>
class TP_MOE {
protected:
    GeneralMOEConfig config;
    int tp_count;                            // NUMA 节点数
    std::vector<T*> tps;                     // 各 NUMA 节点的实例
    std::vector<float*> local_output_numa;   // 各 NUMA 节点的输出

    void merge_results(int qlen, void* output);  // 合并各节点结果
};
```

---

## 2. 新增类结构设计

### 2.1 类继承关系

```
AMX_MOE_TP<T>                          # 现有推理类
    ↓ 继承
AMX_SFT_MOE_TP<T>                      # SFT 扩展类
    ↓ 封装
TP_MOE<AMX_SFT_MOE_TP<GemmKernel>>
    ↓ 继承
TP_MOE_SFT<AMX_SFT_MOE_TP<GemmKernel>> # TP 封装
```

### 2.2 MOESFTConfig 配置类

```cpp
// operators/common.hpp
struct MOESFTConfig : public GeneralMOEConfig {
    // LoRA 配置
    int lora_rank = 16;
    float lora_alpha = 32.0f;
    float lora_scaling() const { return lora_alpha / lora_rank; }

    // LoRA 权重指针 (零拷贝)
    void* gate_lora_a = nullptr;  // [expert_num, lora_rank, hidden_size]
    void* gate_lora_b = nullptr;  // [expert_num, intermediate_size, lora_rank]
    void* up_lora_a = nullptr;
    void* up_lora_b = nullptr;
    void* down_lora_a = nullptr;
    void* down_lora_b = nullptr;

    // 梯度检查点配置
    int max_cache_depth = 1;
};
```

### 2.3 ForwardCache 缓存结构

```cpp
// operators/amx/sft_moe.hpp
struct ForwardCache {
    // 中间值缓存 (用于 backward)
    ggml_bf16_t* input_cache = nullptr;         // [qlen, hidden_size]
    ggml_bf16_t* gate_output_cache = nullptr;   // [tokens_total, intermediate_size]
    ggml_bf16_t* up_output_cache = nullptr;     // [tokens_total, intermediate_size]
    ggml_bf16_t* intermediate_cache = nullptr;  // [tokens_total, intermediate_size]

    // 路由信息缓存
    std::vector<int64_t> expert_ids_cache;
    std::vector<float> weights_cache;
    std::vector<int> m_local_num_cache;
    std::vector<std::vector<int>> m_local_pos_cache;
    std::vector<int> m_expert_id_map_cache;
    int qlen_cache = 0;
    int k_cache = 0;
    int activated_expert_cache = 0;

    bool valid = false;
};
```

### 2.4 AMX_SFT_MOE_TP 类定义

```cpp
// operators/amx/sft_moe.hpp
template <class T>
class AMX_SFT_MOE_TP : public AMX_MOE_TP<T> {
private:
    using Base = AMX_MOE_TP<T>;
    // 继承 Base 的 protected 成员...

    // SFT 配置
    MOESFTConfig sft_config_;

    // LoRA 配置
    int lora_rank_;
    float lora_scaling_;

    // LoRA 权重指针 (直接指向 Python tensor)
    ggml_bf16_t* gate_lora_a_;  // [expert_num, lora_rank, hidden_size]
    ggml_bf16_t* gate_lora_b_;  // [expert_num, intermediate_size, lora_rank]
    ggml_bf16_t* up_lora_a_;
    ggml_bf16_t* up_lora_b_;
    ggml_bf16_t* down_lora_a_;
    ggml_bf16_t* down_lora_b_;

    // LoRA 中间缓冲区 (原始设计，已废弃)
    // 注: Bug #18 修复后，compute_lora_gate_up 和 compute_lora_down 改用
    // 线程本地 std::vector<float> local_intermediate，避免并行任务间的 race condition
    ggml_bf16_t* lora_intermediate_;  // [max_len * k, lora_rank] (保留但不再使用)

    // Forward cache 栈
    std::vector<ForwardCache> cache_stack_;
    int cache_stack_top_ = 0;
    int max_cache_depth_;

    // 梯度中间缓冲区
    ggml_bf16_t* grad_intermediate_;  // [max_len * k, intermediate_size]
    ggml_bf16_t* grad_gate_output_;   // [max_len * k, intermediate_size]
    ggml_bf16_t* grad_up_output_;     // [max_len * k, intermediate_size]

public:
    AMX_SFT_MOE_TP(MOESFTConfig config, int tp_part_idx = 0);

    // SFT 专用接口
    void forward_sft(int qlen, int k, const int64_t* expert_ids,
                     const float* weights, const void* input, void* output,
                     bool save_for_backward);

    void backward(const void* grad_output, void* grad_input,
                  void* grad_gate_lora_a, void* grad_gate_lora_b,
                  void* grad_up_lora_a, void* grad_up_lora_b,
                  void* grad_down_lora_a, void* grad_down_lora_b);

    void update_lora_weights(void* gate_lora_a, void* gate_lora_b,
                             void* up_lora_a, void* up_lora_b,
                             void* down_lora_a, void* down_lora_b);

private:
    void init_lora_buffers();
    void init_cache_buffers();
    void init_grad_buffers();

    void compute_lora_gate_up(int qlen, int activated_expert);
    void compute_lora_down(int qlen, int activated_expert);

    ForwardCache& push_cache();
    ForwardCache pop_cache();
    void save_to_cache(ForwardCache& cache, int qlen, int k,
                       const int64_t* expert_ids, const float* weights,
                       int activated_expert);

    void backward_down(const ForwardCache& cache, const void* grad_output,
                       void* grad_down_lora_a, void* grad_down_lora_b);
    void backward_activation(const ForwardCache& cache);
    void backward_gate_up(const ForwardCache& cache, void* grad_input,
                          void* grad_gate_lora_a, void* grad_gate_lora_b,
                          void* grad_up_lora_a, void* grad_up_lora_b);
};
```

### 2.5 TP_MOE_SFT 类定义

```cpp
// operators/moe-sft-tp.hpp
template <class T>
class TP_MOE_SFT : public TP_MOE<T> {
public:
    using Base = TP_MOE<T>;
    MOESFTConfig sft_config;

    TP_MOE_SFT(MOESFTConfig config);

    void forward_sft(int* qlen_ptr, int k, const int64_t* expert_ids,
                     const float* weights, const void* input, void* output,
                     bool save_for_backward);

    void backward(const void* grad_output, void* grad_input,
                  void* grad_gate_lora_a, void* grad_gate_lora_b,
                  void* grad_up_lora_a, void* grad_up_lora_b,
                  void* grad_down_lora_a, void* grad_down_lora_b);

    void update_lora_weights(void* gate_lora_a, void* gate_lora_b,
                             void* up_lora_a, void* up_lora_b,
                             void* down_lora_a, void* down_lora_b);

    // Python bindings
    void forward_sft_binding(...);
    void backward_binding(...);
    void update_lora_weights_binding(...);
};
```

---

## 3. 核心算法实现

### 3.1 forward_sft 算法流程

```
forward_sft(qlen, k, expert_ids, weights, input, output, save_for_backward):
    1. 专家路由 (复用基类逻辑)
       - 计算 m_local_num_[expert] = 每个专家的 token 数
       - 计算 m_local_pos_[i][j] = token i 在专家 j 中的位置
       - 构建 m_expert_id_map_ = 激活的专家列表

    2. 缓冲区分配 (复用基类逻辑)
       - 分配 gate/up/down 的 input/output buffer

    3. 输入复制
       - 将 input 按 expert_ids 路由到各专家的 buffer

    4. 输入量化
       - gate_up_ba_[expert]->from_mat(input)

    5. Gate + Up GEMM (基类方法)
       - do_gate_up_gemm()

    6. Gate + Up LoRA (新增)
       - compute_lora_gate_up()
       - gate_output += (input @ gate_lora_A^T @ gate_lora_B^T) * scaling
       - up_output += (input @ up_lora_A^T @ up_lora_B^T) * scaling

    7. 保存缓存 (如果 save_for_backward)
       - push_cache()
       - save_to_cache(input, gate_output, up_output, routing_info)

    8. 激活函数
       - intermediate = silu(gate_output) * up_output

    9. 中间值量化
       - down_ba_[expert]->from_mat(intermediate)

    10. Down GEMM (基类方法)
        - do_down_gemm()

    11. Down LoRA (新增)
        - compute_lora_down()
        - down_output += (intermediate @ down_lora_A^T @ down_lora_B^T) * scaling

    12. 加权合并
        - output[i] = Σ weights[i][j] * down_output[expert_ids[i][j]]
```

### 3.2 LoRA 计算细节

```cpp
void compute_lora_gate_up(int qlen, int activated_expert) {
    // 对每个激活的专家并行处理
    pool->do_work_stealing_job(activated_expert * 2,
        [this](int task_id) {
            bool do_up = task_id % 2;
            int expert_idx = m_expert_id_map_[task_id / 2];

            // 获取当前专家的 LoRA 权重
            ggml_bf16_t* lora_a = do_up ? up_lora_a_ : gate_lora_a_;
            ggml_bf16_t* lora_b = do_up ? up_lora_b_ : gate_lora_b_;

            // 偏移到当前专家的权重
            lora_a += expert_idx * lora_rank_ * hidden_size;
            lora_b += expert_idx * intermediate_size * lora_rank_;

            // 注: Bug #18 修复后，使用线程本地 local_intermediate 代替共享 lora_intermediate_
            // 原因: activated_expert * 2 个任务并行，同一 expert 的 gate/up 任务会冲突
            std::vector<float> local_intermediate(num_tokens * lora_rank_);

            // Step 1: intermediate = input @ lora_A^T
            // [num_tokens, hidden_size] @ [lora_rank, hidden_size]^T
            // → [num_tokens, lora_rank]
            for (t in num_tokens):
                for (r in lora_rank):
                    sum = 0
                    for (h in hidden_size):
                        sum += input[t, h] * lora_a[r, h]
                    local_intermediate[t * lora_rank_ + r] = sum

            // Step 2: output += intermediate @ lora_B^T * scaling
            // [num_tokens, lora_rank] @ [intermediate_size, lora_rank]^T
            // → [num_tokens, intermediate_size]
            for (t in num_tokens):
                for (i in intermediate_size):
                    sum = 0
                    for (r in lora_rank):
                        sum += local_intermediate[t * lora_rank_ + r] * lora_b[i, r]
                    output[t, i] += sum * lora_scaling_
        });
}
```

### 3.3 backward 算法流程

```
backward(grad_output, grad_input, grad_gate_lora_a/b, grad_up_lora_a/b, grad_down_lora_a/b):
    1. 弹出缓存
       - cache = pop_cache()
       - 恢复 qlen, k, routing_info

    2. Down 投影反向
       - backward_down()
       - grad_intermediate = grad_output @ down_proj^T + LoRA 反向
       - grad_down_lora_a/b = 梯度累积

    3. 激活函数反向
       - backward_activation()
       - y = silu(gate) * up
       - grad_gate = grad_intermediate * up * d_silu(gate)
       - grad_up = grad_intermediate * silu(gate)

    4. Gate/Up 投影反向
       - backward_gate_up()
       - grad_input = grad_gate @ gate_proj^T + grad_up @ up_proj^T + LoRA 反向
       - grad_gate_lora_a/b = 梯度累积
       - grad_up_lora_a/b = 梯度累积
```

### 3.4 LoRA 梯度计算

对于 `y = x @ W^T + (x @ A^T @ B^T) * scaling`:

```
grad_x = grad_y @ W + (grad_y @ B @ A) * scaling

grad_B = (grad_y^T @ (x @ A^T)) * scaling
       = grad_y^T @ intermediate * scaling
       其中 intermediate = x @ A^T

grad_A = ((B^T @ grad_y^T) @ x) * scaling
       = (grad_y @ B)^T @ x * scaling
```

---

## 4. 内存管理策略

### 4.1 内存分配

使用 `shared_mem_buffer_numa` 进行 NUMA 感知分配:

```cpp
void init_lora_buffers() {
    // LoRA 中间缓冲区 (原始设计，Bug #18 修复后不再使用)
    // 注: 现在 compute_lora_gate_up/down 使用线程本地 std::vector<float>
    // 此分配代码保留但缓冲区已废弃
    lora_intermediate_pool_bytes_ =
        sizeof(ggml_bf16_t) * max_len * num_experts_per_tok * lora_rank_;

    MemoryRequest mem_requests;
    mem_requests.append_pointer(&lora_intermediate_pool_, lora_intermediate_pool_bytes_);
    shared_mem_buffer_numa.alloc(tp_part_idx, this, mem_requests);
}

void init_cache_buffers() {
    // 每个缓存槽的大小
    cache_slot_bytes_input_ = max_len * hidden_size * sizeof(ggml_bf16_t);
    cache_slot_bytes_intermediate_ =
        max_len * num_experts_per_tok * intermediate_size * sizeof(ggml_bf16_t);

    MemoryRequest mem_requests;
    mem_requests.append_pointer(&cache_input_pool_,
                                cache_slot_bytes_input_ * max_cache_depth_);
    mem_requests.append_pointer(&cache_gate_output_pool_,
                                cache_slot_bytes_intermediate_ * max_cache_depth_);
    // ... 其他缓冲区
    shared_mem_buffer_numa.alloc(tp_part_idx, this, mem_requests);
}
```

### 4.2 缓存栈管理

```cpp
ForwardCache& push_cache() {
    if (cache_stack_top_ >= max_cache_depth_) {
        throw std::runtime_error("Forward cache stack overflow");
    }
    return cache_stack_[cache_stack_top_++];
}

ForwardCache pop_cache() {
    if (cache_stack_top_ <= 0) {
        throw std::runtime_error("Forward cache stack underflow");
    }
    return cache_stack_[--cache_stack_top_];
}
```

### 4.3 零拷贝设计

LoRA 权重通过指针直接访问 Python tensor 内存:

```cpp
// 构造时从配置获取指针
AMX_SFT_MOE_TP(MOESFTConfig config, int tp_part_idx)
    : Base(static_cast<GeneralMOEConfig>(config), tp_part_idx) {
    gate_lora_a_ = (ggml_bf16_t*)config.gate_lora_a;
    gate_lora_b_ = (ggml_bf16_t*)config.gate_lora_b;
    // ...
}

// 当 Python tensor 重新分配时更新指针
void update_lora_weights(void* gate_lora_a, ...) {
    gate_lora_a_ = (ggml_bf16_t*)gate_lora_a;
    // ...
}
```

---

## 5. 并行化设计

### 5.1 NUMA 并行

TP_MOE_SFT 在多个 NUMA 节点上并行执行:

```cpp
void forward_sft(...) {
    // 在每个 NUMA 节点上并行执行
    pool->dispense_backend()->do_numa_job([this, ...](int numa_id) {
        tps[numa_id]->forward_sft(qlen, k, expert_ids, weights,
                                   input, local_output_numa[numa_id],
                                   save_for_backward);
    });

    // 合并各节点结果
    merge_results(qlen, output);
}
```

### 5.2 专家级并行

在单个 NUMA 节点内，对激活的专家并行处理:

```cpp
// Gate + Up LoRA 并行处理 2 * activated_expert 个任务
pool->do_work_stealing_job(activated_expert * 2,
    [this](int task_id) {
        bool do_up = task_id % 2;
        int expert_idx = m_expert_id_map_[task_id / 2];
        // 处理单个专家的 gate 或 up LoRA
    });
```

### 5.3 Token 级并行

对于输入复制和结果合并，按 token 并行:

```cpp
pool->do_work_stealing_job(qlen, [&](int i) {
    // 处理第 i 个 token
    for (int j = 0; j < k; j++) {
        // 路由到对应专家
    }
});
```

---

## 6. Python 绑定设计

### 6.1 pybind11 绑定

```cpp
// ext_bindings.cpp

// 绑定 MOESFTConfig
py::class_<MOESFTConfig>(moe_module, "MOESFTConfig")
    .def(py::init<>())
    .def_readwrite("expert_num", &MOESFTConfig::expert_num)
    .def_readwrite("num_experts_per_tok", &MOESFTConfig::num_experts_per_tok)
    .def_readwrite("hidden_size", &MOESFTConfig::hidden_size)
    .def_readwrite("intermediate_size", &MOESFTConfig::intermediate_size)
    .def_readwrite("lora_rank", &MOESFTConfig::lora_rank)
    .def_readwrite("lora_alpha", &MOESFTConfig::lora_alpha)
    .def_readwrite("max_cache_depth", &MOESFTConfig::max_cache_depth)
    .def_readwrite("max_len", &MOESFTConfig::max_len)
    .def_readwrite("layer_idx", &MOESFTConfig::layer_idx)
    .def_readwrite("gate_proj", &MOESFTConfig::gate_proj)
    .def_readwrite("up_proj", &MOESFTConfig::up_proj)
    .def_readwrite("down_proj", &MOESFTConfig::down_proj)
    .def_readwrite("gate_lora_a", &MOESFTConfig::gate_lora_a)
    .def_readwrite("gate_lora_b", &MOESFTConfig::gate_lora_b)
    .def_readwrite("up_lora_a", &MOESFTConfig::up_lora_a)
    .def_readwrite("up_lora_b", &MOESFTConfig::up_lora_b)
    .def_readwrite("down_lora_a", &MOESFTConfig::down_lora_a)
    .def_readwrite("down_lora_b", &MOESFTConfig::down_lora_b)
    .def_readwrite("pool", &MOESFTConfig::pool);

// 绑定 SFT MOE 类
using AMXBF16_SFT_MOE = TP_MOE_SFT<AMX_SFT_MOE_TP<GemmKernel224BF>>;
py::class_<AMXBF16_SFT_MOE>(moe_module, "AMXBF16_SFT_MOE")
    .def(py::init<MOESFTConfig>())
    .def("load_weights_task", ...)
    .def("warm_up_task", ...)
    .def("forward_sft_task", ...)
    .def("backward_task", ...)
    .def("update_lora_weights_task", ...);
```

### 6.2 异步任务模式

所有操作通过 CPUInfer 的 task 模式异步执行:

```cpp
.def("forward_sft_task", [](AMXBF16_SFT_MOE& self, intptr_t qlen_ptr, ...) {
    return create_job([&self, qlen_ptr, ...] {
        self.forward_sft_binding(qlen_ptr, ...);
    });
})
```
