# SFT + KTWrapper 架构分析

## 1. 现有推理架构

### 1.1 Python API 层

#### 类继承关系
```
KTMoEWrapper (工厂类)
    │
    └── __new__() 根据 method 参数返回具体实现
            │
            ├── AMXMoEWrapper (AMXINT4/AMXINT8)
            │
            ├── NativeMoEWrapper (RAWINT4/FP8)
            │
            ├── LlamafileMoEWrapper (LLAMAFILE)
            │
            └── GeneralMoEWrapper (MOE_INT4/MOE_INT8)
                    │
                    └── 继承自 BaseMoEWrapper (基类)
```

#### BaseMoEWrapper 核心功能
- **CPUInfer 单例管理**：全局共享的 CPU 推理引擎
- **KExpertsCPUBuffer 缓冲区管理**：双缓冲机制，支持 GPU-CPU 异步传输
- **异步执行**：`submit_forward()` / `sync_forward()` 分离提交和同步

#### 关键文件
| 文件 | 内容 |
|------|------|
| `python/experts.py` | KTMoEWrapper 工厂类 |
| `python/experts_base.py` | BaseMoEWrapper 基类 + KExpertsCPUBuffer |
| `python/utils/amx.py` | AMXMoEWrapper / NativeMoEWrapper |
| `python/utils/llamafile.py` | LlamafileMoEWrapper |
| `python/utils/moe_kernel.py` | GeneralMoEWrapper |

### 1.2 C++ 后端层

#### 类继承关系
```
CPUInfer
    │
    ├── WorkerPool (NUMA 感知线程池)
    │       │
    │       ├── InNumaPool (NUMA 子池)
    │       │
    │       └── Work-Stealing 调度
    │
    └── TaskQueue (无锁任务队列)

MoE_Interface (接口类)
    │
    └── TP_MOE<T> (Tensor Parallel 包装)
            │
            └── T = AMX_MOE_TP<GemmKernel>
                    │
                    └── AMX_MOE_BASE<T, Derived> (CRTP 基类)
```

#### 关键文件
| 文件 | 内容 |
|------|------|
| `cpu_backend/cpuinfer.h` | CPUInfer 类 |
| `cpu_backend/worker_pool.h` | WorkerPool / InNumaPool |
| `cpu_backend/task_queue.h` | TaskQueue |
| `operators/moe-tp.hpp` | TP_MOE / TP_MOE_Common |
| `operators/amx/moe.hpp` | AMX_MOE_TP |
| `operators/amx/moe_base.hpp` | AMX_MOE_BASE |

---

## 2. 现有 SFT 架构

### 2.1 C++ 层（已实现）

#### 类继承关系
```
TP_MOE<T> (推理)
    │
    └── TP_MOE_SFT<T> (SFT 扩展)
            │
            ├── forward_sft_binding() - 带梯度缓存的前向
            │
            ├── backward_binding() - 反向传播
            │
            └── update_lora_weights_binding() - LoRA 权重更新

AMX_MOE_TP<GemmKernel> (推理)
    │
    └── AMX_SFT_MOE_TP<GemmKernel, BaseMOE> (SFT 扩展)
            │
            ├── ForwardCache - 梯度检查点
            │
            ├── LoRA 权重存储
            │
            └── 反向传播实现
```

#### MOESFTConfig 配置结构
```cpp
struct MOESFTConfig : public GeneralMOEConfig {
    // LoRA 配置
    int lora_rank = 16;
    float lora_alpha = 32.0f;

    // LoRA 权重指针
    void* gate_lora_a;   // [expert_num, lora_rank, hidden_size]
    void* gate_lora_b;   // [expert_num, intermediate_size, lora_rank]
    void* up_lora_a, *up_lora_b;
    void* down_lora_a, *down_lora_b;

    // 梯度检查点
    int max_cache_depth = 1;
};
```

#### ForwardCache 结构
```cpp
struct ForwardCache {
    ggml_bf16_t* input_cache;        // 原始输入
    ggml_bf16_t* gate_output_cache;  // Gate 投影输出
    ggml_bf16_t* up_output_cache;    // Up 投影输出
    ggml_bf16_t* intermediate_cache; // 激活后的中间值

    // 路由信息
    std::vector<int64_t> expert_ids_cache;
    std::vector<float> weights_cache;
    std::vector<int> m_local_num_cache;
};
```

#### 关键文件
| 文件 | 内容 |
|------|------|
| `operators/moe-sft-tp.hpp` | TP_MOE_SFT 类 |
| `operators/amx/sft_moe.hpp` | AMX_SFT_MOE_TP 类 |
| `operators/common.hpp` | MOESFTConfig / ForwardCache |
| `ext_bindings.cpp` | Python 绑定 |

### 2.2 Python 层（当前状态）

**问题**：SFT 目前无 Python Wrapper，直接调用 C++ 绑定

```python
# 当前使用方式（直接调用 C++ 绑定）
config = kt_kernel_ext.moe.MOESFTConfig(...)
config.lora_rank = 16
config.lora_alpha = 32.0
moe = kt_kernel_ext.moe.AMXBF16_SFT_MOE(config)

CPUInfer.submit(moe.load_weights_task(...))
CPUInfer.sync()
CPUInfer.submit(moe.forward_sft_task(...))
CPUInfer.sync()
```

**缺点**：
- 用户需要了解 C++ 绑定细节
- 无法复用推理层的缓冲区管理
- 无法复用 CPUInfer 单例管理
- 接口不一致

---

## 3. 继承关系图

### 3.1 C++ 层完整继承图

```
                                    MoE_Interface
                                         │
                    ┌────────────────────┴────────────────────┐
                    │                                          │
              TP_MOE_Common<T>                                 │
                    │                                          │
              TP_MOE<T>                                        │
                    │                                          │
          ┌────────┴────────┐                                  │
          │                 │                                  │
    (推理实例化)       TP_MOE_SFT<T>                            │
          │                 │                                  │
          │            (SFT 扩展)                              │
          │                                                    │
          │                                                    │
    AMX_MOE_BASE<T, Derived>  ◄────── CRTP ──────┐            │
          │                                       │            │
    AMX_MOE_TP<T>                                 │            │
          │                                       │            │
          │                    AMX_SFT_MOE_TP<T, BaseMOE>      │
          │                           │                        │
          │                    (继承自 BaseMOE<T>)             │
          │                           │                        │
          └───────────────────────────┘                        │
                                                               │
    GemmKernel224BF / GemmKernel224Int8 / GemmKernel224Int4    │
                    │                                          │
                    └──────────────► 模板参数 T ◄──────────────┘
```

### 3.2 Python 层目标继承图（设计目标）

```
                    _MoEBase (共享基类)
                         │
                         ├── _cpu_infer_instance (单例)
                         │
                         └── _get_cpu_infer() (类方法)
                              │
          ┌──────────────────┴──────────────────┐
          │                                      │
   BaseMoEWrapper                        BaseSFTMoEWrapper
   (推理基类-不变)                         (SFT 基类-新增)
          │                                      │
          ├── forward()                          ├── forward_sft()
          ├── submit_forward()                   ├── backward()
          ├── sync_forward()                     └── update_lora_weights()
          │                                      │
          │                                      │
  ┌───────┼───────┐                      ┌───────┴───────┐
  │       │       │                      │               │
AMXMoE  Native  Llamafile              AMXSFTMoE     (其他SFT)
Wrapper Wrapper  Wrapper               Wrapper        Wrapper
```

---

## 4. 代码复用分析

### 4.1 C++ 层复用（自动）

| 组件 | 推理 | SFT | 复用方式 |
|------|------|-----|---------|
| AMX GemmKernel | ✅ | ✅ | 继承（AMX_SFT_MOE_TP 继承自 AMX_MOE_TP） |
| WorkerPool | ✅ | ✅ | 共享（通过 config.pool） |
| TaskQueue | ✅ | ✅ | 共享（通过 CPUInfer） |
| NUMA 调度 | ✅ | ✅ | 继承 |
| 量化算法 | ✅ | ✅ | 继承 |
| TP 分区 | ✅ | ✅ | 继承（TP_MOE_SFT 继承自 TP_MOE） |

**结论**：C++ 层 SFT 已经通过继承自动复用推理优化，无需额外工作。

### 4.2 Python 层复用（需设计）

| 组件 | 推理 | SFT | 复用情况 |
|------|------|-----|---------|
| CPUInfer 单例 | ✅ | ❌ | 需要提取到共享基类 |
| WorkerPoolConfig | ✅ | ❌ | 需要提取到共享基类 |
| 缓冲区管理 | KExpertsCPUBuffer | 需新建 | 需求不同，无法直接复用 |
| forward 逻辑 | forward() | forward_sft() | 签名不同，无法复用 |
| 权重加载 | load_weights() | load_weights() | 可部分复用 |

**结论**：Python 层主要复用 CPUInfer 单例管理，其他需要独立实现。

---

## 5. 方案对比

### 5.1 方案 A：完全合并

```python
class KTMoEWrapper:
    def __new__(cls, ..., mode="inference"/"sft"):
        # 返回同一个类的不同配置
```

| 优点 | 缺点 |
|------|------|
| 统一入口 | 接口污染（推理用户看到 SFT 参数） |
| 代码复用最大化 | 基类膨胀（推理需要处理 SFT 方法） |
| | 缓冲区管理复杂 |
| | 状态管理混乱 |
| | Bug 风险高 |

### 5.2 方案 B：完全分离

```python
class KTMoEWrapper:      # 推理
class KTMoESFTWrapper:   # SFT（完全独立）
```

| 优点 | 缺点 |
|------|------|
| 职责分离清晰 | 代码重复（CPUInfer 管理） |
| 各自演进不干扰 | 推理优化不能自动惠及 SFT |
| 接口精简 | 需要维护两套代码 |

### 5.3 方案 C：轻度合并（推荐）

```python
class _MoEBase:           # 共享基类
class BaseMoEWrapper:     # 推理基类（继承 _MoEBase）
class BaseSFTMoEWrapper:  # SFT 基类（继承 _MoEBase）
class KTMoEWrapper:       # 统一工厂入口
```

| 优点 | 缺点 |
|------|------|
| 统一入口 | 轻微的接口参数增加 |
| 共享 CPUInfer 管理 | |
| 各自独立的缓冲区 | |
| 推理代码不变，零风险 | |
| SFT 独立实现 | |

---

## 6. 推荐方案：轻度合并

### 6.1 设计原则

1. **共享基础设施**：CPUInfer 单例、WorkerPoolConfig
2. **分离业务逻辑**：推理和 SFT 各自独立的 forward/backward
3. **独立缓冲区**：推理用 KExpertsCPUBuffer，SFT 用 KExpertsSFTBuffer
4. **统一入口**：KTMoEWrapper 工厂类通过 mode 参数区分

### 6.2 架构图

```
┌─────────────────────────────────────────────────────────────────┐
│                        Python API 层                             │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │             KTMoEWrapper (统一工厂入口)                       │ │
│  │    mode="inference" → 推理    mode="sft" → SFT              │ │
│  └─────────────────────────────────────────────────────────────┘ │
│                              │                                   │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │                     _MoEBase (共享基类)                       │ │
│  │  - CPUInfer 单例管理                                          │ │
│  │  - WorkerPoolConfig 构建                                      │ │
│  └─────────────────────────────────────────────────────────────┘ │
│           │                              │                       │
│  ┌────────────────────┐      ┌────────────────────┐             │
│  │  BaseMoEWrapper    │      │ BaseSFTMoEWrapper  │             │
│  │  (推理基类-不变)    │      │   (SFT基类-新增)   │             │
│  │  - forward()       │      │ - forward_sft()    │             │
│  │  - submit_forward()│      │ - backward()       │             │
│  │  - sync_forward()  │      │ - update_lora()    │             │
│  │  - KExpertsCPU     │      │ - KExpertsSFT      │             │
│  │    Buffer          │      │   Buffer           │             │
│  └────────────────────┘      └────────────────────┘             │
│           │                              │                       │
│  ┌────────────────────┐      ┌────────────────────┐             │
│  │  AMXMoEWrapper     │      │  AMXSFTMoEWrapper  │             │
│  │  NativeMoEWrapper  │      │  (新增)            │             │
│  │  LlamafileMoEWrapper│      │                    │             │
│  │  GeneralMoEWrapper │      │                    │             │
│  └────────────────────┘      └────────────────────┘             │
└─────────────────────────────────────────────────────────────────┘
                              │
                    pybind11 (kt_kernel_ext)
                              │
┌─────────────────────────────────────────────────────────────────┐
│                        C++ 后端层                                │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │                      CPUInfer                                 │ │
│  │  - submit() / sync()                                         │ │
│  │  - submit_with_cuda_stream()                                 │ │
│  └─────────────────────────────────────────────────────────────┘ │
│                              │                                   │
│  ┌──────────────────────┬──────────────────────┐                │
│  │     WorkerPool       │     TaskQueue        │                │
│  │  - NUMA 感知线程池     │  - 无锁任务队列       │                │
│  └──────────────────────┴──────────────────────┘                │
│                              │                                   │
│  ┌────────────────────────────────────────────────────────────┐ │
│  │                    算子层                                    │ │
│  │  ┌─────────────────────┐  ┌─────────────────────┐          │ │
│  │  │ TP_MOE<T> (推理)     │  │ TP_MOE_SFT<T> (SFT) │          │ │
│  │  │  └─AMX_MOE_TP<T>    │  │  └─AMX_SFT_MOE_TP<T>│          │ │
│  │  └─────────────────────┘  └─────────────────────┘          │ │
│  └────────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
```

### 6.3 复用效果

| 优化类型 | 推理 | SFT | 复用方式 |
|---------|------|-----|---------|
| AMX 内核优化 | ✅ | ✅ | C++ 继承（自动） |
| 线程池优化 | ✅ | ✅ | 共享 CPUInfer 单例 |
| NUMA 调度优化 | ✅ | ✅ | 共享 WorkerPool |
| 量化算法优化 | ✅ | ✅ | C++ 继承（自动） |
| TP 分区优化 | ✅ | ✅ | C++ 继承（自动） |
| Python CPUInfer | ✅ | ✅ | 共享 _MoEBase |
| 缓冲区管理 | ✅ | ✅ | 各自独立（需求不同） |
