# MoE SFT AMX 功能最终实现文档

## 1. 文件变更清单

### 1.1 新增文件

| 文件路径 | 说明 |
|---------|------|
| `operators/amx/sft_moe.hpp` | AMX_SFT_MOE_TP 类实现 |
| `operators/moe-sft-tp.hpp` | TP_MOE_SFT 封装类 |
| `examples/test_moe_sft_amx.py` | 测试文件 |
| `docs/sft_moe_amx/` | 文档目录 |

### 1.2 修改文件

| 文件路径 | 修改内容 |
|---------|---------|
| `operators/common.hpp` | 新增 MOESFTConfig 配置类 |
| `ext_bindings.cpp` | 新增 SFT MOE Python 绑定 |

---

## 2. 类实现详情

### 2.1 MOESFTConfig (operators/common.hpp)

```cpp
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;
    void* gate_lora_b = nullptr;
    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;

    // 构造函数
    MOESFTConfig() : GeneralMOEConfig() {}
    MOESFTConfig(int expert_num, int routed_expert_num,
                 int hidden_size, int intermediate_size);
};
```

### 2.2 ForwardCache (operators/amx/sft_moe.hpp)

用于梯度检查点的缓存结构:

```cpp
struct ForwardCache {
    // 中间值缓存
    ggml_bf16_t* input_cache = nullptr;
    ggml_bf16_t* gate_output_cache = nullptr;
    ggml_bf16_t* up_output_cache = nullptr;
    ggml_bf16_t* intermediate_cache = nullptr;

    // 路由信息缓存
    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.2.1 缓存字段详解

| 缓存字段 | 内容 | 保存时机 | 用途 |
|---------|------|----------|-----|
| `input_cache` | 原始输入 (token order) | `save_to_cache()` | `backward_gate_up` 计算 LoRA 梯度 |
| `gate_output_cache` | gate 输出 (激活前) | `save_to_cache()` | `backward_activation` 计算 SiLU 梯度 |
| `up_output_cache` | up 输出 (激活前) | `save_to_cache()` | `backward_activation` 计算 SiLU 梯度 |
| `intermediate_cache` | silu(gate) × up (激活后) | `save_intermediate_to_cache()` | `backward_down` 计算 down LoRA 梯度 |

**保存时机说明**:

1. **`save_to_cache()`** 在 `apply_activation` **之前**调用
   - 保存 `gate_output_cache` (激活前的 gate projection 输出)
   - 保存 `up_output_cache` (激活前的 up projection 输出)
   - 保存 `input_cache` (原始 token order 的输入，用于 LoRA 梯度计算)

2. **`save_intermediate_to_cache()`** 在 `apply_activation` **之后**调用
   - 保存 `intermediate_cache` = silu(gate) × up (激活后的中间值)
   - 这是 Bug #17c 的修复：down LoRA 梯度需要激活后的 intermediate，而非激活前的 gate

**Backward 使用**:
- `backward_down`: 使用 `intermediate_cache` (激活后) 计算 down LoRA 梯度
- `backward_activation`: 使用 `gate_output_cache` + `up_output_cache` (激活前) 计算 SiLU 梯度
- `backward_gate_up`: 使用 `input_cache` (原始 token order) 计算 gate/up LoRA 梯度

#### 2.2.2 内存估算

**单个 cache slot 大小计算公式**:

```
input_cache:        max_len × hidden_size × 2 bytes
gate_output_cache:  max_len × k × intermediate_size × 2 bytes
up_output_cache:    max_len × k × intermediate_size × 2 bytes
intermediate_cache: max_len × k × intermediate_size × 2 bytes
```

**示例 (DeepSeek-V3 参数)**:
```
参数:
  max_len = 25600
  k = 8 (num_experts_per_tok)
  hidden_size = 7168
  intermediate_size = 2048

计算:
  input_cache:        25600 × 7168 × 2 = 350 MB
  gate_output_cache:  25600 × 8 × 2048 × 2 = 800 MB
  up_output_cache:    25600 × 8 × 2048 × 2 = 800 MB
  intermediate_cache: 25600 × 8 × 2048 × 2 = 800 MB

  单个 cache slot 总计 ≈ 2.75 GB

  如果 max_cache_depth = 2，则总缓存需求 ≈ 5.5 GB
```

**梯度缓冲区大小**:
```
grad_intermediate_: max_len × k × intermediate_size × 2 bytes ≈ 800 MB
grad_gate_output_:  max_len × k × intermediate_size × 2 bytes ≈ 800 MB
grad_up_output_:    max_len × k × intermediate_size × 2 bytes ≈ 800 MB

梯度缓冲区总计 ≈ 2.4 GB
```

### 2.3 AMX_SFT_MOE_TP (operators/amx/sft_moe.hpp)

继承自 AMX_MOE_TP，添加 SFT 训练支持:

**公开方法**:

| 方法 | 签名 | 说明 |
|------|------|------|
| 构造函数 | `AMX_SFT_MOE_TP(MOESFTConfig config, int tp_part_idx = 0)` | 初始化 |
| forward_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)` | SFT 前向传播 |
| backward | `void backward(const void* grad_output, void* grad_input, void* grad_gate_lora_a, ...)` | 反向传播 |
| update_lora_weights | `void update_lora_weights(void* gate_lora_a, ...)` | 更新权重指针 |

**私有方法**:

| 方法 | 说明 |
|------|------|
| `init_lora_buffers()` | 初始化 LoRA 中间缓冲区 |
| `init_cache_buffers()` | 初始化缓存缓冲区 |
| `init_grad_buffers()` | 初始化梯度缓冲区 |
| `compute_lora_gate_up()` | 计算 gate/up LoRA |
| `compute_lora_down()` | 计算 down LoRA |
| `push_cache()` / `pop_cache()` | 缓存栈管理 |
| `save_to_cache()` | 保存中间值到缓存 |
| `backward_down()` | down 投影反向传播 |
| `backward_activation()` | 激活函数反向传播 |
| `backward_gate_up()` | gate/up 投影反向传播 |

### 2.4 TP_MOE_SFT (operators/moe-sft-tp.hpp)

多 NUMA 节点封装，实现权重分区和梯度合并:

```cpp
template <class T>
class TP_MOE_SFT : public TP_MOE<T> {
public:
    MOESFTConfig sft_config;

    // Bug #19 fix: 分区后的 LoRA 权重指针（含 intermediate_size 维度的需要分区）
    std::vector<ggml_bf16_t*> partitioned_gate_lora_b_;  // 连续块分片
    std::vector<ggml_bf16_t*> partitioned_up_lora_b_;    // 连续块分片
    std::vector<ggml_bf16_t*> partitioned_down_lora_a_;  // 逐行分片

    // Bug #20 fix: 分区后的基础权重指针（backward 需要 BF16 权重）
    std::vector<ggml_bf16_t*> partitioned_gate_proj_;
    std::vector<ggml_bf16_t*> partitioned_up_proj_;
    std::vector<ggml_bf16_t*> partitioned_down_proj_;

    TP_MOE_SFT(MOESFTConfig config);
    ~TP_MOE_SFT();  // 释放分区权重

    // 主要接口
    void load_weights() override;  // Bug #19 fix: 添加基础权重分区
    void forward_sft(int* qlen_ptr, int k, ...);
    void backward(const void* grad_output, ...);  // Bug #21 fix: 添加梯度分区合并
    void update_lora_weights(void* gate_lora_a, ...);  // 添加 LoRA 分区

    // 内存管理
    void free_partitioned_lora_weights();
    void free_partitioned_base_weights();

    // Python 绑定
    void forward_sft_binding(intptr_t qlen_ptr, ...);
    void backward_binding(intptr_t grad_output, ...);
    void update_lora_weights_binding(intptr_t gate_lora_a, ...);
};
```

#### 2.4.1 TP 权重分区策略

在 TP 模式下，`intermediate_size` 维度被分区到多个 NUMA 节点。权重分区规则：

| 权重类型 | 原始形状 | 分区方式 | 分区后形状 |
|---------|---------|---------|-----------|
| gate_proj / up_proj | `[E, I, H]` | 连续块 | `[E, I/N, H]` |
| down_proj | `[E, H, I]` | 逐行 | `[E, H, I/N]` |
| gate_lora_b / up_lora_b | `[E, I, R]` | 连续块 | `[E, I/N, R]` |
| down_lora_a | `[E, R, I]` | 逐行 | `[E, R, I/N]` |
| gate_lora_a / up_lora_a | `[E, R, H]` | 不分区（零拷贝） | `[E, R, H]` |
| down_lora_b | `[E, H, R]` | 不分区（零拷贝） | `[E, H, R]` |

其中：E = expert_num, I = intermediate_size, H = hidden_size, R = lora_rank, N = tp_count

#### 2.4.2 Backward 梯度合并

Bug #21 修复：梯度需要与权重对称处理：

```cpp
void backward(...) {
    // 1. 为每个 NUMA 分配分区梯度 buffer
    std::vector<ggml_bf16_t*> part_grad_gate_lora_b(tp_count);
    // ...

    // 2. 每个 NUMA 计算分区梯度
    pool->do_numa_job([...](int numa_id) {
        tps[numa_id]->backward(grad_output, grad_input,
                               grad_gate_lora_a,                  // 不分区
                               part_grad_gate_lora_b[numa_id],    // 分区
                               ...);
    });

    // 3. 合并分区梯度到完整梯度
    for (int i = 0; i < tp_count; i++) {
        // 连续块合并 / 逐行合并
    }

    // 4. 清理临时 buffer
}
```

---

## 3. Python 绑定

### 3.1 MOESFTConfig 绑定 (ext_bindings.cpp)

```cpp
py::class_<MOESFTConfig, GeneralMOEConfig>(moe_module, "MOESFTConfig")
    .def(py::init<>())
    .def(py::init<int, int, int, int>())
    .def_readwrite("lora_rank", &MOESFTConfig::lora_rank)
    .def_readwrite("lora_alpha", &MOESFTConfig::lora_alpha)
    .def_readwrite("max_cache_depth", &MOESFTConfig::max_cache_depth)
    .DEF_PTR_PROPERTY(MOESFTConfig, gate_lora_a)
    .DEF_PTR_PROPERTY(MOESFTConfig, gate_lora_b)
    .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_a)
    .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_b)
    .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_a)
    .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b);
```

### 3.2 SFT MOE 类绑定

使用模板函数绑定 BF16 和 INT8 两种模式:

```cpp
template <class T>
void bind_moe_sft_module(py::module_& moe_module, const char* name) {
    using SFT_MOE = TP_MOE_SFT<T>;
    py::class_<SFT_MOE>(moe_module, name)
        .def(py::init<MOESFTConfig>())
        .def("load_weights_task", ...)
        .def("warm_up_task", ...)
        .def("forward_sft_task", [](SFT_MOE& self, intptr_t qlen_ptr, int k,
                                     intptr_t expert_ids, intptr_t weights,
                                     intptr_t input, intptr_t output,
                                     bool save_for_backward) {
            return create_job([&self, ...] {
                self.forward_sft_binding(...);
            });
        })
        .def("backward_task", [](SFT_MOE& self, intptr_t grad_output, ...) {
            return create_job([&self, ...] {
                self.backward_binding(...);
            });
        })
        .def("update_lora_weights_task", [](SFT_MOE& self, ...) {
            return create_job([&self, ...] {
                self.update_lora_weights_binding(...);
            });
        });
}

// 实例化
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224BF>>(moe_module, "AMXBF16_SFT_MOE");
bind_moe_sft_module<AMX_SFT_MOE_TP<amx::GemmKernel224Int8>>(moe_module, "AMXInt8_SFT_MOE");
```

---

## 4. 内存布局

### 4.1 LoRA 权重布局

```
gate_lora_a: [expert_num, lora_rank, hidden_size]
             连续存储，按 expert_idx 索引

gate_lora_b: [expert_num, intermediate_size, lora_rank]
             连续存储，按 expert_idx 索引

偏移计算:
  expert_lora_a = gate_lora_a + expert_idx * lora_rank * hidden_size
  expert_lora_b = gate_lora_b + expert_idx * intermediate_size * lora_rank
```

### 4.2 缓存布局

```
cache_stack_: std::vector<ForwardCache>  [max_cache_depth]
  ├── cache_stack_[0]
  │   ├── input_cache:        [max_len, hidden_size]
  │   ├── gate_output_cache:  [max_len * k, intermediate_size]
  │   ├── up_output_cache:    [max_len * k, intermediate_size]
  │   └── intermediate_cache: [max_len * k, intermediate_size]
  ├── cache_stack_[1]
  │   └── ...
  └── ...
```

### 4.3 梯度缓冲区布局

```
grad_intermediate_: [max_len * k, intermediate_size]
grad_gate_output_:  [max_len * k, intermediate_size]
grad_up_output_:    [max_len * k, intermediate_size]
```

---

## 5. 数据流

### 5.1 训练循环

```
Python                                 C++
  │                                     │
  ├─ config.lora_a = tensor.ptr() ─────>│ 零拷贝指针初始化
  │                                     │
  ├─ forward_sft_task() ───────────────>│ 前向传播
  │   save_for_backward=True            │   ├─ Gate + Up GEMM
  │                                     │   ├─ Gate + Up LoRA
  │<───────────── output (float32) ─────│   ├─ save_to_cache()
  │                                     │   ├─ Activation
  │                                     │   ├─ Down GEMM
  │                                     │   └─ Down LoRA
  │                                     │
  ├─ loss = compute_loss(output)        │
  ├─ grad_output = d_loss/d_output      │
  │                                     │
  ├─ backward_task() ──────────────────>│ 反向传播
  │                                     │   ├─ pop_cache()
  │<───────────── grad_lora_* ──────────│   ├─ backward_down()
  │                                     │   ├─ backward_activation()
  │                                     │   └─ backward_gate_up()
  │                                     │
  ├─ param.grad = grad_lora_*           │
  ├─ optimizer.step()                   │ 原地更新权重
  │   (in-place update)                 │
  │                                     │
  ├─ update_lora_weights_task() ───────>│ Bug #22: TP 模式必须同步分区权重
  │   (TP 模式必需)                       │   ↓ 重新复制分区权重
  │                                     │
  └─ 下一个 step                         │
```

### 5.2 输入输出类型

| 张量 | 数据类型 | 说明 |
|------|---------|------|
| input | bf16 | 输入 hidden states |
| output | float32 | 输出便于 loss 计算 |
| expert_ids | int64 | 专家路由索引 |
| weights | float32 | 专家路由权重 |
| grad_output | bf16 | 上游梯度 |
| grad_input | bf16 | 输入梯度 |
| grad_lora_* | bf16 | LoRA 梯度 |

---

## 6. API 变更记录

### 6.1 从 v1.0 到 v2.0 的变更

| 变更类型 | v1.0 | v2.0 | 原因 |
|---------|------|------|------|
| 新增 | - | MOESFTConfig | 统一 SFT 配置 |
| 新增 | - | forward_sft_task() | SFT 专用前向 |
| 新增 | - | update_lora_weights_task() | 指针更新 |
| 移除 | sync_lora_weights_task() | - | 零拷贝设计不需要 |
| 变更 | load_base_weights_task(mapping) | load_weights_task() | 简化接口 |
| 变更 | backward_task(routing_info, ...) | backward_task(grad, ...) | 使用缓存路由 |
| 变更 | output: bf16 | output: float32 | 便于 loss 计算 |

### 6.2 Python API 示例对比

**v1.0 (旧)**:
```python
# 每次 forward 前同步
CPUInfer.submit(moe.sync_lora_weights_task(
    gate_lora_a.data_ptr(), ...
))
CPUInfer.sync()

# forward
CPUInfer.submit(moe.forward_task(...))

# backward 需要传入路由信息
CPUInfer.submit(moe.backward_task(
    qlen, k, expert_ids.data_ptr(), weights.data_ptr(),
    grad_output.data_ptr(), ...
))
```

**v2.0 (新)**:
```python
# 初始化时设置指针 (零拷贝)
config.gate_lora_a = gate_lora_a.data_ptr()

# forward (无需同步)
CPUInfer.submit(moe.forward_sft_task(
    bsz_tensor.data_ptr(), k, expert_ids.data_ptr(), weights.data_ptr(),
    input.data_ptr(), output.data_ptr(), True  # save_for_backward
))

# backward (使用缓存的路由信息)
CPUInfer.submit(moe.backward_task(
    grad_output.data_ptr(), grad_input.data_ptr(),
    grad_gate_lora_a.data_ptr(), ...
))

# optimizer.step() 原地更新
optimizer.step()

# TP 模式：必须同步分区权重 (Bug #22)
CPUInfer.submit(moe.update_lora_weights_task(
    gate_lora_a.data_ptr(), gate_lora_b.data_ptr(),
    up_lora_a.data_ptr(), up_lora_b.data_ptr(),
    down_lora_a.data_ptr(), down_lora_b.data_ptr(),
))
CPUInfer.sync()
```

---

## 7. 代码位置索引

> **注**: 行号基于 2026-01-04 Bug #22 修复后的代码版本

| 功能 | 文件 | 行号范围 |
|------|------|---------|
| MOESFTConfig 定义 | operators/common.hpp | 293-316 |
| ForwardCache 定义 | operators/amx/sft_moe.hpp | 23-41 |
| AMX_SFT_MOE_TP 类 | operators/amx/sft_moe.hpp | 54-1197 |
| forward_sft 实现 | operators/amx/sft_moe.hpp | 164-356 |
| backward 实现 | operators/amx/sft_moe.hpp | 372-428 |
| compute_lora_gate_up | operators/amx/sft_moe.hpp | 552-614 |
| compute_lora_down | operators/amx/sft_moe.hpp | 619-676 |
| TP_MOE_SFT 类 | operators/moe-sft-tp.hpp | 30-462 |
| TP_MOE_SFT::load_weights | operators/moe-sft-tp.hpp | 73-148 |
| TP_MOE_SFT::backward | operators/moe-sft-tp.hpp | 254-319 |
| TP_MOE_SFT::update_lora_weights | operators/moe-sft-tp.hpp | 347-408 |
| Python 绑定 | ext_bindings.cpp | 749-781 |
| 测试代码 (TP) | examples/test_moe_sft_amx.py | 1-1100+ |
| 测试代码 (no-TP) | examples/test_moe_sft_amx_no_tp.py | 1-1100+ |
