# MoE SFT AMX 算子接口文档

## 1. 概述

`moe_sft_amx` 是用于 MoE (Mixture of Experts) 层 LoRA 微调的高性能算子，基于 Intel AMX (Advanced Matrix Extensions) 加速。该算子支持 BF16 和 INT8 量化模式，提供前向传播和反向传播功能。

### 1.1 主要特性

- **LoRA 微调**: 支持在 gate/up/down 三个投影矩阵上应用 LoRA 适配器
- **量化模式**: 支持 BF16 和 INT8 两种精度
- **AMX 加速**: 利用 Intel AMX 指令集进行高效矩阵运算
- **异步执行**: 通过 CPUInfer 实现异步任务提交和执行
- **零拷贝/复制设计**: 部分 LoRA 权重零拷贝，部分需要分区复制（详见下文）
- **梯度检查点**: 支持 forward 保存中间值用于 backward

---

## 2. 数据流

```
Training Step (no-TP 模式 - 完全零拷贝):
┌─────────────────────────────────────────────────────────────────────┐
│  Python                           │           C++                   │
├───────────────────────────────────┼─────────────────────────────────┤
│  config.gate_lora_a = ptr         │  直接访问 Python tensor 内存    │
│  (零拷贝, 在初始化时设置)          │                                 │
│                                   │                                 │
│  1. forward_sft_task()        ────>  前向传播 (保存中间值)          │
│     output <──────────────────────   返回输出 (float32)             │
│                                   │                                 │
│  2. backward_task()           ────>  反向传播                       │
│     grad_lora_* <─────────────────   写入 LoRA 梯度到指定 buffer    │
│                                   │                                 │
│  3. optimizer.step()              │                                 │
│     原地更新 LoRA 权重             │  下次 forward 自动看到更新      │
│     (零拷贝, 无需同步)             │                                 │
│                                   │                                 │
│  4. 下一个 step, 回到 1           │                                 │
└───────────────────────────────────┴─────────────────────────────────┘

Training Step (TP 模式 - 部分需要同步):
┌─────────────────────────────────────────────────────────────────────┐
│  Python                           │           C++                   │
├───────────────────────────────────┼─────────────────────────────────┤
│  config.gate_lora_a = ptr         │  分区权重被复制到 NUMA 节点     │
│  (初始化时设置)                    │                                 │
│                                   │                                 │
│  1. forward_sft_task()        ────>  前向传播 (使用分区权重)        │
│     output <──────────────────────   返回输出 (float32)             │
│                                   │                                 │
│  2. backward_task()           ────>  反向传播 (梯度分区+合并)       │
│     grad_lora_* <─────────────────   写入完整梯度                   │
│                                   │                                 │
│  3. optimizer.step()              │                                 │
│     原地更新 LoRA 权重             │                                 │
│                                   │                                 │
│  4. update_lora_weights_task()────>  ★ 必须！重新复制分区权重 ★     │
│     (TP 模式必需)                  │                                 │
│                                   │                                 │
│  5. 下一个 step, 回到 1           │                                 │
└───────────────────────────────────┴─────────────────────────────────┘
```

---

## 3. 配置参数

### 3.1 MOESFTConfig

| 参数 | 类型 | 说明 |
|------|------|------|
| `expert_num` | int | 专家总数 |
| `num_experts_per_tok` | int | 每个 token 激活的专家数 (top-k) |
| `hidden_size` | int | 隐藏层维度 |
| `intermediate_size` | int | MLP 中间层维度 |
| `lora_rank` | int | LoRA 秩 (r) |
| `lora_alpha` | float | LoRA 缩放因子 (alpha) |
| `layer_idx` | int | 层索引 |
| `max_len` | int | 最大序列长度 |
| `max_cache_depth` | int | 最大缓存深度 (用于梯度检查点) |
| `gate_proj` | int64 | gate 投影权重指针 |
| `up_proj` | int64 | up 投影权重指针 |
| `down_proj` | int64 | down 投影权重指针 |
| `gate_lora_a` | int64 | gate LoRA A 权重指针 (零拷贝) |
| `gate_lora_b` | int64 | gate LoRA B 权重指针 (零拷贝) |
| `up_lora_a` | int64 | up LoRA A 权重指针 (零拷贝) |
| `up_lora_b` | int64 | up LoRA B 权重指针 (零拷贝) |
| `down_lora_a` | int64 | down LoRA A 权重指针 (零拷贝) |
| `down_lora_b` | int64 | down LoRA B 权重指针 (零拷贝) |
| `pool` | WorkerPool* | CPUInfer 后端线程池 |

---

## 4. 权重格式

### 4.1 基础权重 (冻结)

```python
gate_proj: Tensor  # [expert_num, intermediate_size, hidden_size], bf16
up_proj: Tensor    # [expert_num, intermediate_size, hidden_size], bf16
down_proj: Tensor  # [expert_num, hidden_size, intermediate_size], bf16
```

### 4.2 LoRA 适配器权重 (可训练)

每个投影矩阵有两个 LoRA 矩阵 A 和 B:

```python
# Gate 投影 LoRA
gate_lora_a: Tensor  # [expert_num, lora_rank, hidden_size], bf16
gate_lora_b: Tensor  # [expert_num, intermediate_size, lora_rank], bf16

# Up 投影 LoRA
up_lora_a: Tensor    # [expert_num, lora_rank, hidden_size], bf16
up_lora_b: Tensor    # [expert_num, intermediate_size, lora_rank], bf16

# Down 投影 LoRA
down_lora_a: Tensor  # [expert_num, lora_rank, intermediate_size], bf16
down_lora_b: Tensor  # [expert_num, hidden_size, lora_rank], bf16
```

### 4.3 LoRA 计算公式

```
output = input @ W^T + (input @ A^T @ B^T) * (alpha / rank)
```

其中:
- `W` 是基础权重 (冻结)
- `A` 和 `B` 是 LoRA 适配器矩阵 (可训练)
- `alpha / rank` 是缩放因子

---

## 5. Python API 接口

### 5.1 创建实例

```python
import kt_kernel
kt_kernel_ext = kt_kernel.kt_kernel_ext

# 创建 CPUInfer 实例
CPUInfer = kt_kernel_ext.CPUInfer(num_threads)

# 创建配置 (使用属性设置)
config = kt_kernel_ext.moe.MOESFTConfig()
config.expert_num = expert_num
config.num_experts_per_tok = num_experts_per_tok
config.hidden_size = hidden_size
config.intermediate_size = intermediate_size
config.lora_rank = lora_rank
config.lora_alpha = lora_alpha
config.max_cache_depth = 1  # 梯度检查点缓存深度
config.max_len = max_len
config.layer_idx = 0

# 设置基础权重指针
config.gate_proj = gate_proj.data_ptr()
config.up_proj = up_proj.data_ptr()
config.down_proj = down_proj.data_ptr()

# 设置 LoRA 权重指针 (零拷贝 - 直接指向 Python tensor)
config.gate_lora_a = gate_lora_a.data_ptr()
config.gate_lora_b = gate_lora_b.data_ptr()
config.up_lora_a = up_lora_a.data_ptr()
config.up_lora_b = up_lora_b.data_ptr()
config.down_lora_a = down_lora_a.data_ptr()
config.down_lora_b = down_lora_b.data_ptr()

config.pool = CPUInfer.backend_

# 创建 MOE SFT 实例
# BF16 模式:
moe = kt_kernel_ext.moe.AMXBF16_SFT_MOE(config)
# 或 INT8 模式:
moe = kt_kernel_ext.moe.AMXInt8_SFT_MOE(config)
```

### 5.2 加载基础权重

```python
# 加载并量化基础权重
CPUInfer.submit(moe.load_weights_task())
CPUInfer.sync()
```

### 5.3 预热 (可选)

```python
CPUInfer.submit(moe.warm_up_task())
CPUInfer.sync()
```

### 5.4 前向传播

```python
# 输入张量
bsz_tensor = torch.tensor([qlen], device="cpu")           # 批大小
expert_ids = torch.tensor(..., dtype=torch.int64)         # [qlen, k]
weights = torch.tensor(..., dtype=torch.float32)          # [qlen, k]
input_data = torch.tensor(..., dtype=torch.bfloat16)      # [qlen, hidden_size]
output = torch.zeros((qlen, hidden_size), dtype=torch.float32)  # 输出为 float32

CPUInfer.submit(moe.forward_sft_task(
    bsz_tensor.data_ptr(),
    num_experts_per_tok,
    expert_ids.data_ptr(),
    weights.data_ptr(),
    input_data.data_ptr(),
    output.data_ptr(),
    save_for_backward=True  # 是否保存中间值用于反向传播
))
CPUInfer.sync()
```

**参数说明**:

| 参数 | 类型 | 说明 |
|------|------|------|
| `bsz_ptr` | int64 | 批大小张量指针 |
| `num_experts_per_tok` | int | 每 token 专家数 |
| `expert_ids_ptr` | int64 | 专家 ID 张量指针 [qlen, k] |
| `weights_ptr` | int64 | 路由权重张量指针 [qlen, k] |
| `input_ptr` | int64 | 输入张量指针 [qlen, hidden_size] |
| `output_ptr` | int64 | 输出张量指针 [qlen, hidden_size], float32 |
| `save_for_backward` | bool | 是否保存中间值 |

### 5.5 反向传播

```python
# 分配梯度缓冲区
grad_output = torch.tensor(..., dtype=torch.bfloat16)     # [qlen, hidden_size]
grad_input = torch.zeros((qlen, hidden_size), dtype=torch.bfloat16)

grad_gate_lora_a = torch.zeros_like(gate_lora_a)
grad_gate_lora_b = torch.zeros_like(gate_lora_b)
grad_up_lora_a = torch.zeros_like(up_lora_a)
grad_up_lora_b = torch.zeros_like(up_lora_b)
grad_down_lora_a = torch.zeros_like(down_lora_a)
grad_down_lora_b = torch.zeros_like(down_lora_b)

CPUInfer.submit(moe.backward_task(
    grad_output.data_ptr(),
    grad_input.data_ptr(),
    grad_gate_lora_a.data_ptr(),
    grad_gate_lora_b.data_ptr(),
    grad_up_lora_a.data_ptr(),
    grad_up_lora_b.data_ptr(),
    grad_down_lora_a.data_ptr(),
    grad_down_lora_b.data_ptr()
))
CPUInfer.sync()
```

**参数说明**:

| 参数 | 类型 | 说明 |
|------|------|------|
| `grad_output_ptr` | int64 | 上游梯度指针 [qlen, hidden_size] |
| `grad_input_ptr` | int64 | 输入梯度输出指针 [qlen, hidden_size] |
| `grad_gate_lora_a/b_ptr` | int64 | gate LoRA 梯度输出指针 |
| `grad_up_lora_a/b_ptr` | int64 | up LoRA 梯度输出指针 |
| `grad_down_lora_a/b_ptr` | int64 | down LoRA 梯度输出指针 |

### 5.6 更新 LoRA 权重指针

#### 5.6.1 何时需要调用

| 场景 | no-TP 模式 | TP 模式 |
|------|-----------|---------|
| tensor 被重新分配（非原地操作） | ✓ 需要 | ✓ 需要 |
| 原地更新（optimizer.step()） | ❌ 不需要 | ✓ **需要** |

**TP 模式特殊说明**：

在 TP 模式下，以下 LoRA 权重会被**复制**到各 NUMA 节点（非零拷贝）：

| 权重 | 分区方式 | 修改后是否需要同步 |
|------|---------|-------------------|
| `gate_lora_a` | 零拷贝 | ❌ 不需要 |
| `gate_lora_b` | 复制（连续块） | ✓ **需要** |
| `up_lora_a` | 零拷贝 | ❌ 不需要 |
| `up_lora_b` | 复制（连续块） | ✓ **需要** |
| `down_lora_a` | 复制（逐行） | ✓ **需要** |
| `down_lora_b` | 零拷贝 | ❌ 不需要 |

因此，**在 TP 模式下，每次 optimizer.step() 后都必须调用 update_lora_weights_task()**。

#### 5.6.2 使用示例

```python
# 例如: 当 tensor 被重新分配后
new_gate_lora_a = some_operation_that_creates_new_tensor(gate_lora_a)

# 更新 C++ 端的指针
CPUInfer.submit(moe.update_lora_weights_task(
    new_gate_lora_a.data_ptr(),
    new_gate_lora_b.data_ptr(),
    new_up_lora_a.data_ptr(),
    new_up_lora_b.data_ptr(),
    new_down_lora_a.data_ptr(),
    new_down_lora_b.data_ptr()
))
CPUInfer.sync()
```

#### 5.6.3 TP 模式训练循环示例

```python
# TP 模式下的完整训练循环
for step in range(num_steps):
    # Forward
    CPUInfer.submit(moe.forward_sft_task(..., save_for_backward=True))
    CPUInfer.sync()

    # Compute loss
    loss = compute_loss(output)
    grad_output = compute_grad(loss)

    # Backward
    CPUInfer.submit(moe.backward_task(grad_output.data_ptr(), ...))
    CPUInfer.sync()

    # Update weights
    optimizer.step()

    # ★ TP 模式必需：同步分区权重 ★
    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()
```

**注意**: 在 no-TP 模式下，如果使用原地操作 (如 `tensor.add_()`, `optimizer.step()`), 则不需要调用此接口, 因为零拷贝设计会自动看到更新。

---

## 6. Forward Cache 机制

### 6.1 概述

SFT MoE 使用 ForwardCache 保存前向传播中间值，用于反向传播计算 LoRA 梯度。

### 6.2 缓存内容

| 缓存字段 | 内容 | 保存时机 | Backward 用途 |
|---------|------|----------|--------------|
| `input_cache` | 原始输入 (token order) | `save_to_cache()` | `backward_gate_up`: 计算 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 梯度 |

### 6.3 保存时序

```
Forward 执行流程:
┌─────────────────────────────────────────────────────────────────────┐
│ Step 1-4: Routing, gather input                                     │
│                                                                     │
│ Step 5: Gate/Up GEMM + LoRA                                         │
│         m_local_gate_output_ = gate projection 输出                 │
│         m_local_up_output_ = up projection 输出                     │
│                                                                     │
│ ★ save_to_cache() ★                                                 │
│   - 保存 input (原始 token order)                                   │
│   - 保存 gate_output_cache (激活前)                                  │
│   - 保存 up_output_cache (激活前)                                    │
│                                                                     │
│ Step 6: apply_activation()                                          │
│         m_local_gate_output_ = silu(gate) × up                      │
│                                                                     │
│ ★ save_intermediate_to_cache() ★                                    │
│   - 保存 intermediate_cache (激活后)                                 │
│                                                                     │
│ Step 7-8: Down GEMM + LoRA, scatter output                          │
└─────────────────────────────────────────────────────────────────────┘
```

### 6.4 内存估算

**计算公式**:

```
单个 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

梯度缓冲区:
  grad_intermediate_: max_len × k × intermediate_size × 2 bytes
  grad_gate_output_:  max_len × k × intermediate_size × 2 bytes
  grad_up_output_:    max_len × k × intermediate_size × 2 bytes
```

**DeepSeek-V3 示例**:

| 参数 | 值 |
|------|-----|
| max_len | 25600 |
| k (num_experts_per_tok) | 8 |
| hidden_size | 7168 |
| intermediate_size | 2048 |

| 缓冲区 | 大小 |
|--------|------|
| input_cache | 350 MB |
| gate_output_cache | 800 MB |
| up_output_cache | 800 MB |
| intermediate_cache | 800 MB |
| **单个 cache slot** | **2.75 GB** |
| 梯度缓冲区 (3个) | 2.4 GB |
| **总计 (depth=1)** | **≈ 5.15 GB** |

**注意**: 如果 `max_cache_depth > 1`，cache 内存按倍数增长。

---

## 7. 精度要求

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

精度计算方式:
```python
relative_diff = mean(abs(output - reference)) / mean(abs(reference))
```

---

## 8. 注意事项

### 8.1 通用注意事项

1. **内存对齐**: 所有张量必须是 contiguous 的
2. **异步执行**: 使用 `CPUInfer.submit()` 提交任务后需要调用 `CPUInfer.sync()` 等待完成
3. **梯度缓冲区**: 反向传播会覆盖梯度缓冲区，不会累积
4. **基础权重冻结**: `load_weights_task()` 只需调用一次，基础权重在训练过程中不变
5. **输出格式**: `forward_sft_task()` 输出为 float32, 便于后续 loss 计算

### 8.2 no-TP 模式注意事项

1. **零拷贝设计**: 所有 LoRA 权重通过指针直接访问 Python tensor, 无需每次 forward 前同步
2. **指针更新**: 仅当 LoRA tensor 被重新分配 (非原地操作) 时, 需要调用 `update_lora_weights_task()`

### 8.3 TP 模式注意事项

1. **部分零拷贝**: 只有不含 `intermediate_size` 维度的权重是零拷贝的
   - 零拷贝: `gate_lora_a`, `up_lora_a`, `down_lora_b`
   - 分区复制: `gate_lora_b`, `up_lora_b`, `down_lora_a`
2. **必须同步**: 每次 `optimizer.step()` 后必须调用 `update_lora_weights_task()` 同步分区权重
3. **梯度分区合并**: backward 会自动处理梯度的分区和合并，用户无需额外操作
4. **内存开销**: 分区权重会增加内存开销（每个 NUMA 节点保存分区副本）

---

## 9. API 变更记录

### v2.1 (当前版本) - TP 模式支持

- **修复**: Bug #19-22 - 完整支持 TP 模式 forward/backward
- **变更**: `update_lora_weights_task()` - TP 模式下每次 optimizer.step() 后必须调用
- **新增**: TP 模式权重分区（基础权重 + LoRA 权重）
- **新增**: TP 模式梯度分区与合并

### v2.0 - 零拷贝设计

- **新增**: `MOESFTConfig` 支持直接设置 LoRA 权重指针
- **新增**: `forward_sft_task()` - SFT 专用前向传播
- **新增**: `update_lora_weights_task()` - 更新 LoRA 权重指针
- **移除**: `sync_lora_weights_task()` - 不再需要每次同步
- **变更**: `load_weights_task()` - 替代 `load_base_weights_task()`, 无需 mapping 参数
- **变更**: `backward_task()` - 简化参数, 使用缓存的路由信息
- **变更**: 输出格式从 bf16 改为 float32

### v1.0 (旧版本)

- `sync_lora_weights_task()` - 每次 forward 前同步 LoRA 权重
- `forward_task()` - 通用前向传播
- `backward_task()` - 需要传入完整的路由信息
