# SFT-MOE-AMX Real Data 调试 - 功能需求文档

## 1. 背景

### 1.1 项目背景

kt-kernel 是 KTransformers 的高性能算子库，其中 SFT-MOE-AMX 是用于监督微调 (SFT) 场景的 Mixture of Experts 算子，使用 Intel AMX 指令集加速 BF16 矩阵乘法。

### 1.2 问题背景

在使用真实训练数据进行测试时，SFT-MOE-AMX 算子产生大量 NaN，而使用随机生成的测试数据则正常通过。

---

## 2. 测试环境

### 2.1 硬件

- CPU: 支持 AMX 指令集的 Intel Xeon (Sapphire Rapids 或更新)
- 内存: 足够运行 64 expert 的 MoE 模型

### 2.2 软件

- Python 环境: `ref-llama` (conda)
- 测试框架: PyTorch
- 编译环境: C++17, CMake

---

## 3. 功能需求

### 3.1 核心功能: SFT-MOE Forward

**输入:**
| 参数 | 类型 | 形状 | 描述 |
|------|------|------|------|
| input_data | bf16 | [qlen, hidden_size] | 输入隐藏状态 |
| expert_ids | int64 | [qlen, num_experts_per_tok] | 每个 token 选择的 expert ID |
| weights | fp32 | [qlen, num_experts_per_tok] | 每个 expert 的权重 |

**权重:**
| 参数 | 类型 | 形状 |
|------|------|------|
| gate_proj | bf16 | [expert_num, intermediate_size, hidden_size] |
| up_proj | bf16 | [expert_num, intermediate_size, hidden_size] |
| down_proj | bf16 | [expert_num, hidden_size, intermediate_size] |
| gate_lora_a | bf16 | [expert_num, lora_rank, hidden_size] |
| gate_lora_b | bf16 | [expert_num, intermediate_size, lora_rank] |
| up_lora_a | bf16 | [expert_num, lora_rank, hidden_size] |
| up_lora_b | bf16 | [expert_num, intermediate_size, lora_rank] |
| down_lora_a | bf16 | [expert_num, lora_rank, intermediate_size] |
| down_lora_b | bf16 | [expert_num, hidden_size, lora_rank] |

**输出:**
| 参数 | 类型 | 形状 | 描述 |
|------|------|------|------|
| output | bf16 | [qlen, hidden_size] | 输出隐藏状态 |

**计算公式:**
```
对于每个 token i 和选中的 expert e:
  x = input_data[i]

  # Gate 计算
  gate_base = x @ gate_proj[e].T
  gate_lora = (x @ gate_lora_a[e].T) @ gate_lora_b[e].T
  gate = gate_base + gate_lora * lora_scaling

  # Up 计算
  up_base = x @ up_proj[e].T
  up_lora = (x @ up_lora_a[e].T) @ up_lora_b[e].T
  up = up_base + up_lora * lora_scaling

  # 激活
  intermediate = silu(gate) * up

  # Down 计算
  down_base = intermediate @ down_proj[e].T
  down_lora = (intermediate @ down_lora_a[e].T) @ down_lora_b[e].T
  expert_output = down_base + down_lora * lora_scaling

  output[i] += weights[i, j] * expert_output
```

### 3.2 测试模式

#### 3.2.1 accuracy 模式

使用随机生成的数据测试数值精度:
- 生成随机输入和权重
- 对比 AMX 输出与 PyTorch 参考实现
- 验证最大误差在可接受范围内

#### 3.2.2 real_data 模式

使用真实训练数据测试:
- 从 `/mnt/data/lpl/kt_nan_debug_data.pt` 加载数据
- 数据来自 LlamaFactory 训练过程
- 验证输出不包含 NaN/Inf

#### 3.2.3 perf 模式

测试性能:
- 多次迭代测量执行时间
- 计算吞吐量

---

## 4. 配置参数

### 4.1 accuracy 模式默认配置

```python
expert_num = 256
hidden_size = 7168
intermediate_size = 2048
num_experts_per_tok = 6
lora_rank = 8
lora_alpha = 16.0
qlen = 1000
```

### 4.2 real_data 模式配置 (从 pt 文件读取)

```python
expert_num = 64
hidden_size = 2048
intermediate_size = 1408
num_experts_per_tok = 6
lora_rank = 8
lora_alpha = 16.0
qlen = 48
layer_idx = 1
```

---

## 5. 验收标准

### 5.1 功能验收

| 测试 | 验收标准 |
|------|----------|
| accuracy forward | 最大误差 < 0.1 |
| accuracy backward | 最大误差 < 0.5 (梯度累积) |
| real_data forward | NaN 数量 = 0 |
| real_data backward | NaN 数量 = 0 |

### 5.2 性能验收

| 指标 | 目标 |
|------|------|
| Forward 时间 | < PyTorch 参考实现的 50% |
| 内存占用 | 合理范围内 |

---

## 6. 测试数据格式

### 6.1 pt 文件结构

```python
{
    'input_data': tensor[bf16, qlen x hidden_size],
    'expert_ids': tensor[int64, qlen x num_experts_per_tok],
    'weights': tensor[fp32, qlen x num_experts_per_tok],
    'gate_proj': tensor[bf16, expert_num x intermediate_size x hidden_size],
    'up_proj': tensor[bf16, expert_num x intermediate_size x hidden_size],
    'down_proj': tensor[bf16, expert_num x hidden_size x intermediate_size],
    'gate_lora_a': tensor[bf16, expert_num x lora_rank x hidden_size],
    'gate_lora_b': tensor[bf16, expert_num x intermediate_size x lora_rank],
    'up_lora_a': tensor[bf16, expert_num x lora_rank x hidden_size],
    'up_lora_b': tensor[bf16, expert_num x intermediate_size x lora_rank],
    'down_lora_a': tensor[bf16, expert_num x lora_rank x intermediate_size],
    'down_lora_b': tensor[bf16, expert_num x hidden_size x lora_rank],
    'expert_num': int,
    'hidden_size': int,
    'intermediate_size': int,
    'num_experts_per_tok': int,
    'layer_idx': int,
}
```

### 6.2 数据验证

加载 pt 文件后应验证:
1. 所有张量不含 NaN/Inf
2. 张量形状与配置一致
3. expert_ids 值域在 [0, expert_num) 内
4. weights 非负且归一化

---

## 7. 相关文件

| 文件 | 描述 |
|------|------|
| `test_moe_sft_amx_no_tp.py` | 主测试文件 |
| `sft_moe.hpp` | AMX SFT MOE 实现 |
| `moe-sft-tp.hpp` | TP 包装器 |
| `kt_nan_debug_data.pt` | 真实训练数据 |
