# SFT-MOE-AMX 算子架构文档

## 1. 概述

SFT-MOE-AMX 是用于监督微调 (Supervised Fine-Tuning) 场景的 Mixture of Experts 算子，使用 Intel AMX 指令集加速 BF16 矩阵乘法，并支持 LoRA 适配器。

### 1.1 核心文件

| 文件 | 描述 |
|------|------|
| `sft_moe.hpp` | AMX SFT MOE 核心实现 |
| `moe-sft-tp.hpp` | Tensor Parallelism 包装器 |
| `amx_raw_buffers.hpp` | BufferA/B/C 定义 |

### 1.2 类继承关系

```
MOE_Base<T>                    // 基础 MOE 类
    ↓
AMXBF16_SFT_MOE               // SFT MOE 实现
    ↓
TP_MOE_SFT                     // TP 包装器
```

---

## 2. 核心数据结构

### 2.1 AMX Buffer 系统

AMX 指令集要求特定的内存布局，使用三种 Buffer:

| Buffer | 用途 | 特点 |
|--------|------|------|
| BufferA | 存储左矩阵 (activations) | 行主序，M_STEP 分块 |
| BufferB | 存储右矩阵 (weights) | 列主序，预转置 |
| BufferC | 存储输出结果 | 累积器格式 |

### 2.2 关键对齐参数

```cpp
constexpr int M_STEP = 64;      // M 维度分块大小
constexpr int N_BLOCK = 64;     // N 维度分块大小
constexpr int K_STEP = 32;      // K 维度分块大小 (lora_rank 对齐)
```

### 2.3 padded_lora_rank

为了对齐 K_STEP，lora_rank 会被 padding:

```cpp
padded_lora_rank_ = (lora_rank + K_STEP - 1) / K_STEP * K_STEP;
// 例如: lora_rank=8 -> padded_lora_rank=32
```

---

## 3. Forward 计算流程

### 3.1 流程图

```
输入: input [qlen, hidden_size]
      expert_ids [qlen, num_experts_per_tok]
      weights [qlen, num_experts_per_tok]

Step 1: Expert Routing
    → m_local_num_[expert_id]: 每个 expert 处理的 token 数量
    → m_expert_id_map_[task_id]: 激活的 expert ID 映射

Step 2: Buffer 分配
    → gate_up_ba_[expert_id]: 输入 BufferA
    → gate_bc_[expert_id], up_bc_[expert_id]: 输出 BufferC
    → down_ba_[expert_id], down_bc_[expert_id]: Down 投影用

Step 3: 输入复制
    → m_local_input_ptr_[expert_id]: 按 expert 分组的输入

Step 4: 输入量化
    → gate_up_ba_[expert_id]->from_mat(): BF16 → BufferA 格式

Step 5: Gate + Up 基础 GEMM
    → gate_bc_ = input @ gate_proj.T
    → up_bc_ = input @ up_proj.T
    → 输出: m_local_gate_output_ptr_, m_local_up_output_ptr_

Step 5.5: Gate + Up LoRA (可选)
    → lora_out = (input @ lora_A.T) @ lora_B.T
    → gate/up_output += lora_out * lora_scaling

Step 6: 激活函数
    → intermediate = silu(gate) * up
    → 结果存储在 m_local_gate_output_ptr_

Step 7: Intermediate 量化
    → down_ba_[expert_id]->from_mat(): 准备 Down 投影输入

Step 8: Down 基础 GEMM
    → down_bc_ = intermediate @ down_proj.T
    → 输出: m_local_down_output_ptr_

Step 8.5: Down LoRA (可选)
    → lora_out = (intermediate @ lora_A.T) @ lora_B.T
    → down_output += lora_out * lora_scaling

Step 9: 加权合并
    → output[i] = Σ(weights[i,j] * expert_output[expert_ids[i,j]])

输出: output [qlen, hidden_size]
```

### 3.2 关键函数

| 函数 | 位置 | 描述 |
|------|------|------|
| `forward_sft` | sft_moe.hpp:399 | Forward 主入口 |
| `do_gate_up_gemm` | moe.hpp | Gate/Up 基础 GEMM |
| `do_down_gemm` | moe.hpp | Down 基础 GEMM |
| `compute_lora_gate_up_amx` | sft_moe.hpp:1557 | Gate/Up LoRA (AMX) |
| `compute_lora_down_amx` | sft_moe.hpp:1654 | Down LoRA (AMX) |

---

## 4. LoRA 计算架构

### 4.1 LoRA 公式

```
output = base_output + (input @ lora_A.T @ lora_B.T) * lora_scaling
```

其中:
- `lora_scaling = lora_alpha / lora_rank`

### 4.2 LoRA Buffer 结构

```cpp
// Gate LoRA
gate_lora_a_bb_[expert_id]          // lora_A 转为 BufferB 格式
gate_lora_b_bb_[expert_id]          // lora_B 转为 BufferB 格式
lora_gate_intermediate_bc_[expert_id]  // 中间结果 BufferC
lora_gate_intermediate_ba_[expert_id]  // 中间结果 BufferA (step2)
lora_gate_intermediate_ptr_[expert_id] // 中间结果 BF16 指针
lora_gate_out_bc_[expert_id]        // LoRA 输出 BufferC

// Up LoRA (结构相同)
up_lora_a_bb_[expert_id]
up_lora_b_bb_[expert_id]
lora_up_intermediate_bc_[expert_id]
lora_up_intermediate_ba_[expert_id]
lora_up_intermediate_ptr_[expert_id]
lora_up_out_bc_[expert_id]
```

### 4.3 LoRA 计算步骤 (compute_lora_gate_up_amx)

```
Step 1: input @ lora_A.T
    Input:  gate_up_ba_[expert_id]           // [m, hidden_size]
    Weight: gate_lora_a_bb_[expert_id]       // [padded_lora_rank, hidden_size]
    Output: lora_gate_intermediate_bc_       // [m, padded_lora_rank]
    Then:   bc->to_mat() → lora_gate_intermediate_ptr_

Step 2: Quantize intermediate
    Input:  lora_gate_intermediate_ptr_      // BF16 [m, padded_lora_rank]
    Output: lora_gate_intermediate_ba_       // BufferA 格式

Step 3: intermediate @ lora_B.T + add to main
    Input:  lora_gate_intermediate_ba_       // [m, padded_lora_rank]
    Weight: gate_lora_b_bb_[expert_id]       // [intermediate_size, padded_lora_rank]
    Output: lora_gate_out_bc_                // [m, intermediate_size]
    Then:   add_lora_output_to_main() → m_local_gate_output_ptr_
```

---

## 5. 权重准备流程

### 5.1 权重加载 (load_weights_task)

从 Python 传入的原始 BF16 权重指针:
- `gate_proj_`, `up_proj_`, `down_proj_`: 基础权重
- `gate_lora_a_`, `gate_lora_b_`: Gate LoRA 权重
- `up_lora_a_`, `up_lora_b_`: Up LoRA 权重
- `down_lora_a_`, `down_lora_b_`: Down LoRA 权重

### 5.2 BufferB 预转换

基础权重在 `load_weights_task` 中转换:
```cpp
gate_bb_[expert_id]->from_mat(src_ptr, ...)  // BF16 → BufferB
```

LoRA 权重在 `prepare_lora_weights` 中延迟转换:
```cpp
convert_lora_a_to_buffer_b()  // lora_A → BufferB
convert_lora_b_to_buffer_b()  // lora_B → BufferB (需要 padding)
```

### 5.3 convert_lora_b_to_buffer_b 函数

关键的 BufferB 转换函数 (sft_moe.hpp:1491):

```cpp
void convert_lora_b_to_buffer_b(
    const ggml_bf16_t* src,     // 原始 BF16 权重 [expert_num, n, k]
    std::vector<std::shared_ptr<BufferB>>& dst_vec,  // 目标 BufferB 数组
    int src_n,                   // intermediate_size (1408)
    int src_k                    // lora_rank (8)
) {
    for (int expert_idx = 0; expert_idx < config_.expert_num; expert_idx++) {
        // 源数据偏移计算
        const ggml_bf16_t* expert_src = src + expert_idx * src_n * src_k;

        // 转换为 BufferB 格式 (会 padding 到 padded_lora_rank)
        dst_vec[expert_idx]->from_mat(expert_src, 0, 1, 0, src_n, 0, src_k);
    }
}
```

### 5.4 BufferB::from_mat 函数

实际的数据复制发生在 `amx_raw_buffers.hpp` 中:

```cpp
void from_mat(const ggml_bf16_t* src, ...) {
    // 使用 AVX-512 指令复制数据
    // 偏移计算: src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin
    avx512_copy_32xbf16(dst, src + offset);
}
```

---

## 6. NaN 调试点

### 6.1 调试宏定义

```cpp
// NaN 检查助手
NaNCheckResult check_bf16_buffer_for_nan(buf, size, label);
NaNCheckResult check_fp32_buffer_for_nan(buf, size, label);
```

### 6.2 Forward 中的检查点

| Step | 检查点 | 描述 |
|------|--------|------|
| 5 | GATE/UP 基础 GEMM 后 | 检查 m_local_gate_output_ptr_ |
| 5.5 | GATE/UP LoRA 后 | 检查 LoRA 加法后的结果 |
| 6 | 激活函数后 | 检查 silu * up 结果 |
| 8 | DOWN 基础 GEMM 后 | 检查 m_local_down_output_ptr_ |
| 8.5 | DOWN LoRA 后 | 检查 LoRA 加法后的结果 |

---

## 7. 内存布局

### 7.1 Expert 权重内存布局

```
原始 BF16 权重 (连续存储):
┌─────────────────────────────────────────────────────┐
│ Expert 0 │ Expert 1 │ ... │ Expert 63             │
└─────────────────────────────────────────────────────┘
     ↓ offset = expert_id * n * k

Expert 内部布局 (row-major):
┌─────────────────────────────────────────────────────┐
│ row 0: [col 0, col 1, ..., col k-1]               │
│ row 1: [col 0, col 1, ..., col k-1]               │
│ ...                                                 │
│ row n-1: [col 0, col 1, ..., col k-1]             │
└─────────────────────────────────────────────────────┘
```

### 7.2 BufferB 内存布局

```
BufferB (AMX 优化布局):
┌─────────────────────────────────────────────────────┐
│ N_BLOCK 0:                                          │
│   K_STEP 0: [n0, n1, ..., n63] x [k0..k31]        │
│   K_STEP 1: [n0, n1, ..., n63] x [k32..k63]       │
│   ...                                               │
│ N_BLOCK 1:                                          │
│   ...                                               │
└─────────────────────────────────────────────────────┘
```

---

## 8. 内存分配架构

### 8.1 内存分配策略

SFT-MOE 使用两种内存分配策略：

| 策略 | 函数 | 适用场景 | 特点 |
|------|------|----------|------|
| `aligned_alloc` | 独立分配 | LoRA 缓冲区 | 防止内存污染 |
| `shared_mem_buffer` | 共享池 | Backward/Cache 缓冲区 | 内存复用 |

### 8.2 关键变量定义

```cpp
// 配置参数
E = expert_num           // 专家数量
H = hidden_size          // 隐藏维度
I = intermediate_size    // MLP 中间维度
L = max_len              // 最大序列长度
K = num_experts_per_tok  // 每 token 激活专家数
R = lora_rank            // LoRA rank
R' = padded_lora_rank    // 对齐后的 LoRA rank (32 的倍数)

// 计算变量
max_m = align64(L)                    // 单个 expert 最大 token 数
max_total_tokens = L × K              // 所有激活 expert 的 token 总数
```

### 8.3 内存计算公式

#### LoRA 缓冲区 (aligned_alloc)

```cpp
// BufferB 用于存储 LoRA 权重 (每个 expert 独立)
lora_bb_pool = E × (
    BufferB(R', H) × 2 +        // gate_lora_a, up_lora_a
    BufferB(I, R') × 2 +        // gate_lora_b, up_lora_b
    BufferB(H, R') × 2 +        // gate_lora_a_t, up_lora_a_t
    BufferB(R', I) × 2 +        // gate_lora_b_t, up_lora_b_t
    BufferB(R', I) +            // down_lora_a
    BufferB(H, R')              // down_lora_b
)

// BufferA/C 用于计算中间结果 (共享池)
lora_ba_pool = BufferA(max_total_tokens, R') × 2      // gate + up
lora_bc_inter_pool = BufferC(max_total_tokens, R') × 2
lora_bc_out_pool = BufferC(max_total_tokens, I) × 2 + BufferC(max_total_tokens, H)
lora_intermediate_bf16_pool = max_total_tokens × R' × sizeof(bf16) × 2
```

#### Backward 缓冲区 (shared_mem_buffer)

```cpp
backward_ba_pool = BufferA(max_total_tokens, H)
backward_bc_pool = BufferC(max_total_tokens, I) + BufferC(max_total_tokens, H)
grad_output_bf16_pool = max_total_tokens × H × sizeof(bf16)
backward_bb_pool = E × (BufferB(H, I) × 2 + BufferB(I, H))
```

#### 其他缓冲区

```cpp
grad_buffer = L × K × I × sizeof(bf16) × 3
cache_total = (L × H × sizeof(bf16) + L × K × I × sizeof(bf16) × 3) × cache_depth
```

### 8.4 实测内存数据 (accuracy 模式)

配置: E=256, H=7168, I=2048, L=25600, K=8, R=16, R'=32

| 缓冲区 | 大小 | 备注 |
|--------|------|------|
| lora_bb_pool | 720 MB | LoRA 权重 BufferB |
| lora_ba_pool | 25 MB | LoRA 中间 BufferA |
| lora_bc_inter_pool | 50 MB | LoRA 中间 BufferC |
| lora_bc_out_pool | 8.59 GB | LoRA 输出 BufferC |
| lora_intermediate_bf16_pool | 25 MB | LoRA BF16 中间结果 |
| backward_ba_pool | 2.73 GB | Backward BufferA |
| backward_bc_pool | 7.03 GB | Backward BufferC |
| grad_output_bf16_pool | 2.73 GB | 梯度 BF16 缓冲 |
| backward_bb_pool | 21.00 GB | 转置基础权重 BufferB |
| grad_buffer ×3 | 2.34 GB | 梯度缓冲区 |
| cache_total | 2.69 GB | 缓存 (depth=1) |
| **总计** | **47.93 GB** | |

### 8.5 内存分配位置

| 函数 | 位置 | 分配内容 |
|------|------|----------|
| `init_all_buffers` | sft_moe.hpp:906 | 计算大小，调用 aligned_alloc 和 shared_mem_buffer |
| `init_lora_amx_buffers` | sft_moe.hpp:1170 | 初始化 Buffer 对象（数据指针为 nullptr） |
| `forward_sft` | sft_moe.hpp:399 | 动态分配 LoRA BufferA/C 的数据指针 |
| `backward_*_amx` | sft_moe.hpp | 动态分配 backward BufferA/C 的数据指针 |

---

## 9. 已知问题

### 9.1 Expert 17-24 NaN 问题 (Bug-A) ✅ 已解决

**症状**:
- 只有 Expert 17-24 (连续 8 个) 在 Step 5.5 产生 NaN
- PyTorch 参考实现正常

**根本原因**:
`shared_mem_buffer.alloc` 内存共享机制导致 LoRA BufferB 内存池与其他缓冲区共享物理内存，其他代码写入时污染了 Expert 17-24 的 BufferB 数据。

**修复方案**:
将 LoRA 相关内存池从 `mem_requests.append_pointer()` 改为 `aligned_alloc()` 独立分配。

**修复状态**: ✅ 已解决 (2026-01-11)

### 9.2 accuracy 模式内存过大问题 (Bug-C) ✅ 已解决

**症状**:
- 运行 accuracy 模式时内存需求过大（理论 4+ TB）

**根本原因**:
1. `max_m` 计算错误：使用 `max_len × num_experts_per_tok` 而非 `max_len`
2. 为每个 expert 独立分配 max_m 大小的缓冲区

**修复方案**:
1. 修正 `max_m = align64(max_len)`
2. 使用共享缓冲区池，forward/backward 时动态分配

**修复效果**: 内存从 ~4 TB 降到 ~48 GB

**修复状态**: ✅ 已解决 (2026-01-11)

详见: [bug记录文档.md](./bug记录文档.md)
