# MoE SFT AMX 功能需求文档

## 1. 项目背景

### 1.1 项目概述

kt-kernel 是 KTransformers 项目的高性能 CPU 推理内核，基于 Intel AMX (Advanced Matrix Extensions) 指令集实现 MoE (Mixture of Experts) 层的高效推理。本需求旨在扩展现有推理能力，支持 SFT (Supervised Fine-Tuning) 微调场景下的 LoRA 训练。

### 1.2 需求来源

在 DeepSeek-V3 等大规模 MoE 模型的微调场景中，需要在 CPU 端高效完成:
- 前向传播 (包含 LoRA 适配器计算)
- 反向传播 (计算 LoRA 权重梯度)
- 与 PyTorch 优化器的无缝集成

### 1.3 目标

在现有 AMX MoE 推理算子基础上，增加 SFT 训练支持:
1. 继承现有推理类，最大化代码复用
2. 支持 BF16 + INT8 量化模式
3. 每个 expert 独立的 LoRA 权重
4. C++ 直接访问 Python tensor 指针 (零拷贝)
5. 支持梯度检查点 (gradient checkpointing)

---

## 2. 功能需求

### 2.1 核心功能

| 功能编号 | 功能名称 | 描述 | 优先级 |
|---------|---------|------|--------|
| F-001 | LoRA 前向传播 | 在 gate/up/down 投影矩阵上应用 LoRA 适配器 | P0 |
| F-002 | LoRA 反向传播 | 计算 LoRA A/B 矩阵的梯度 | P0 |
| F-003 | 输入梯度计算 | 计算输入张量的梯度用于链式求导 | P0 |
| F-004 | 零拷贝权重访问 | C++ 直接访问 Python tensor 内存 | P0 |
| F-005 | 梯度检查点支持 | 支持多次 forward 后统一 backward | P1 |
| F-006 | INT8 量化支持 | 在量化模式下进行 SFT 训练 | P1 |

### 2.2 LoRA 适配器规格

**应用位置**: gate_proj, up_proj, down_proj 三个投影矩阵

**计算公式**:
```
output = input @ W^T + (input @ A^T @ B^T) * (alpha / rank)
```

**权重形状**:
| 权重 | 形状 | 说明 |
|------|------|------|
| gate_lora_a | [expert_num, lora_rank, hidden_size] | Gate 投影 LoRA A 矩阵 |
| gate_lora_b | [expert_num, intermediate_size, lora_rank] | Gate 投影 LoRA B 矩阵 |
| up_lora_a | [expert_num, lora_rank, hidden_size] | Up 投影 LoRA A 矩阵 |
| up_lora_b | [expert_num, intermediate_size, lora_rank] | Up 投影 LoRA B 矩阵 |
| down_lora_a | [expert_num, lora_rank, intermediate_size] | Down 投影 LoRA A 矩阵 |
| down_lora_b | [expert_num, hidden_size, lora_rank] | Down 投影 LoRA B 矩阵 |

### 2.3 数据类型

| 张量类型 | 数据类型 | 说明 |
|---------|---------|------|
| 输入 (input) | BF16 | 输入 hidden states |
| 输出 (output) | FP32 | 输出便于 loss 计算 |
| 基础权重 | BF16/INT8 | 冻结的模型权重 |
| LoRA 权重 | BF16 | 可训练的适配器权重 |
| 梯度 | BF16 | LoRA 梯度 |

---

## 3. 非功能需求

### 3.1 性能要求

| 指标 | 要求 | 说明 |
|------|------|------|
| 前向吞吐 | 与推理持平 | LoRA 额外开销 < 15% |
| 反向时延 | < 2x 前向 | 合理的训练开销 |
| 内存效率 | 共享缓冲区 | 复用 NUMA 内存池 |

### 3.2 精度要求

| 模式 | 前向误差阈值 | 反向误差阈值 |
|------|-------------|-------------|
| BF16 | < 0.05 | < 0.10 |
| INT8 | < 0.15 | < 0.25 |

**误差计算方式**:
```python
relative_diff = mean(abs(output - reference)) / mean(abs(reference))
```

### 3.3 兼容性要求

- 继承现有 AMX MoE 类接口
- 兼容 CPUInfer 异步执行模型
- 支持 pybind11 绑定调用
- 与 PyTorch 优化器无缝集成

---

## 4. 设计约束

### 4.1 架构约束

1. **类继承关系**: AMX_SFT_MOE_TP 继承自 AMX_MOE_TP
2. **TP 封装**: TP_MOE_SFT 继承自 TP_MOE
3. **CRTP 模式**: 使用 Curiously Recurring Template Pattern

### 4.2 内存约束

1. 使用 shared_mem_buffer_numa 进行 NUMA 感知分配
2. LoRA 权重通过指针直接访问 (零拷贝)
3. 中间值缓存复用推理缓冲区

### 4.3 接口约束

1. 配置使用 MOESFTConfig (继承自 GeneralMOEConfig)
2. Python 绑定使用 task 异步模型
3. 梯度缓冲区由 Python 端分配

---

## 5. 设计决策

### 5.1 关键设计选择

| 决策点 | 选择方案 | 理由 |
|--------|---------|------|
| 类继承 vs 组合 | 继承 AMX_MOE_TP | 最大化代码复用，利用现有 GEMM 优化 |
| 权重同步机制 | 零拷贝指针 | 避免每次 forward 的内存拷贝开销 |
| 梯度检查点 | 栈式缓存 | 支持多层 forward 后 backward |
| 输出类型 | FP32 | 便于 Python 端 loss 计算 |

### 5.2 取舍权衡

| 问题 | 取舍 | 原因 |
|------|------|------|
| LoRA 计算效率 | 简单循环 vs AMX | 初版使用简单循环，后续可优化为 AMX |
| 梯度累积 | 覆盖 vs 累积 | 每次 backward 覆盖，累积由 Python 端处理 |
| 基础权重梯度 | 不计算 | LoRA 微调场景基础权重冻结 |

---

## 6. 验收标准

### 6.1 功能验收

- [ ] 前向传播输出与 PyTorch 参考实现一致 (BF16 误差 < 5%)
- [ ] 反向传播梯度与 PyTorch 参考实现一致 (BF16 误差 < 10%)
- [ ] 支持完整训练循环 (forward → backward → optimizer.step)
- [ ] 零拷贝设计: 权重更新后无需手动同步

### 6.2 性能验收

- [ ] 前向吞吐保持与推理持平
- [ ] 无内存泄漏 (长时训练稳定)
- [ ] NUMA 亲和性保持

### 6.3 测试用例

| 测试名称 | 测试内容 |
|---------|---------|
| test_moe_sft_forward("bf16") | BF16 模式前向精度 |
| test_moe_sft_forward("int8") | INT8 模式前向精度 |
| test_moe_sft_backward("bf16") | BF16 模式反向精度 |
| test_moe_sft_backward("int8") | INT8 模式反向精度 |
| test_moe_sft_lora_weight_sync() | LoRA 权重同步和指针更新 |
| test_moe_sft_training_loop() | 完整训练循环 |
