# SFT + KTWrapper 功能需求

## 1. 背景与目标

### 1.1 当前问题

| 问题 | 描述 | 影响 |
|------|------|------|
| SFT 无 Python Wrapper | 直接调用 C++ 绑定 | 用户需了解底层细节 |
| 代码分离 | 推理和 SFT 各自独立 | 维护成本高 |
| 优化不共享 | 推理优化无法自动惠及 SFT | 重复工作 |
| 接口不一致 | 推理有 Wrapper，SFT 没有 | 用户体验差 |

### 1.2 目标

1. **统一入口**：`KTMoEWrapper(mode="inference"/"sft")`
2. **自动复用**：推理优化自动惠及 SFT（通过共享基类）
3. **减少 bug**：共享基础设施，减少重复代码
4. **向后兼容**：现有推理代码无需修改

---

## 2. 功能需求

### 2.1 推理功能（保持不变）

| 功能 | 方法 | 参数 | 说明 |
|------|------|------|------|
| 权重加载（预量化） | `load_weights()` | physical_to_logical_map | 从磁盘加载预量化权重 |
| 权重加载（在线量化） | `load_weights_from_tensors()` | gate, up, down, map | 从 BF16 张量在线量化 |
| 同步前向 | `forward()` | hidden_states, topk_ids, topk_weights, cuda_stream | 完整的前向推理 |
| 异步提交 | `submit_forward()` | hidden_states, topk_ids, topk_weights, cuda_stream | 非阻塞提交任务 |
| 异步同步 | `sync_forward()` | hidden_states, cuda_stream | 等待并获取结果 |
| 延迟执行 | `select_deferred_experts()` | expert_ids, scores, protected_k | 选择延迟执行的专家 |

### 2.2 SFT 功能（新增）

| 功能 | 方法 | 参数 | 说明 |
|------|------|------|------|
| 权重加载 | `load_weights()` | physical_to_logical_map | 加载基础权重 |
| LoRA 初始化 | `init_lora_weights()` | gate_a/b, up_a/b, down_a/b | 初始化 LoRA 权重 |
| SFT 前向 | `forward_sft()` | hidden_states, expert_ids, weights, save_for_backward | 带梯度缓存的前向 |
| 反向传播 | `backward()` | grad_output | 计算 LoRA 梯度 |
| 权重更新 | `update_lora_weights()` | - | 同步 LoRA 权重到 C++ |

### 2.3 共享功能

| 功能 | 说明 |
|------|------|
| CPUInfer 单例管理 | 全局共享的 CPU 推理引擎 |
| WorkerPoolConfig | NUMA 子池配置 |
| 配置验证 | 参数有效性检查 |
| 错误处理 | 统一的异常机制 |

---

## 3. 接口需求

### 3.1 工厂类接口

```python
KTMoEWrapper(
    # 基础参数（推理和 SFT 共用）
    layer_idx: int,
    num_experts: int,
    num_experts_per_tok: int,
    hidden_size: int,
    moe_intermediate_size: int,
    num_gpu_experts: int,
    cpuinfer_threads: int,
    threadpool_count: int,
    weight_path: str,
    chunked_prefill_size: int,

    # 推理特有参数
    cpu_save: bool = False,
    max_deferred_experts_per_token: Optional[int] = None,

    # 模式选择
    method: str = "AMXINT4",
    mode: str = "inference",  # "inference" 或 "sft"

    # SFT 特有参数（mode="sft" 时有效）
    lora_rank: int = 16,
    lora_alpha: float = 32.0,
    max_cache_depth: int = 1,
)
```

### 3.2 method 参数值

#### 推理模式 (mode="inference")

| method | 后端 | 说明 |
|--------|------|------|
| `AMXINT4` | AMXMoEWrapper | AMX INT4 量化 |
| `AMXINT8` | AMXMoEWrapper | AMX INT8 量化 |
| `RAWINT4` | NativeMoEWrapper | 预量化 INT4（K-Group） |
| `FP8` | NativeMoEWrapper | FP8 量化 |
| `LLAMAFILE` | LlamafileMoEWrapper | GGUF 格式 |
| `MOE_INT4` | GeneralMoEWrapper | 通用 INT4 内核 |
| `MOE_INT8` | GeneralMoEWrapper | 通用 INT8 内核 |

#### SFT 模式 (mode="sft")

| method | 后端 | 说明 |
|--------|------|------|
| `AMXBF16_SFT` | AMXSFTMoEWrapper | AMX BF16 训练 |
| `AMXINT8_SFT` | AMXSFTMoEWrapper | AMX INT8 训练 |
| `AMXINT4_SFT` | AMXSFTMoEWrapper | AMX INT4 训练 |
| `AMXINT4_KGroup_SFT` | AMXSFTMoEWrapper | AMX INT4 K-Group 训练 |

### 3.3 SFT 特有参数说明

| 参数 | 类型 | 默认值 | 说明 |
|------|------|--------|------|
| `lora_rank` | int | 16 | LoRA 低秩矩阵的秩 |
| `lora_alpha` | float | 32.0 | LoRA 缩放因子 |
| `max_cache_depth` | int | 1 | 前向缓存深度（支持多次前向后反向） |

**LoRA 缩放公式**：
```
lora_scaling = lora_alpha / lora_rank
output = base_output + lora_output * lora_scaling
```

---

## 4. 性能需求

### 4.1 推理性能

| 指标 | 要求 |
|------|------|
| 前向延迟 | 不降级（与当前实现一致） |
| 吞吐量 | 不降级 |
| 内存占用 | 不增加 |

### 4.2 SFT 性能

| 指标 | 要求 |
|------|------|
| 前向延迟 | 与直接调用 C++ 绑定一致（<5% 差异） |
| 反向延迟 | 与直接调用 C++ 绑定一致（<5% 差异） |
| 内存占用 | 与直接调用 C++ 绑定一致 |

---

## 5. 兼容性需求

### 5.1 向后兼容

```python
# 现有推理代码（无需修改）
wrapper = KTMoEWrapper(
    layer_idx=0,
    num_experts=256,
    ...,
    method="AMXINT4"
)
wrapper.load_weights(physical_map)
output = wrapper.forward(hidden_states, topk_ids, topk_weights, cuda_stream)
```

### 5.2 量化格式支持

| 格式 | 推理 | SFT |
|------|------|-----|
| BF16 | ✅ | ✅ |
| INT8 | ✅ | ✅ |
| INT4 | ✅ | ✅ |
| INT4_KGroup | ✅ | ✅ |
| FP8 | ✅ | ❌（暂不支持） |
| GGUF | ✅ | ❌（暂不支持） |

### 5.3 TP 模式支持

| 模式 | 说明 | 推理 | SFT |
|------|------|------|-----|
| TP (Tensor Parallel) | 多 NUMA 节点并行 | ✅ | ✅ |
| No-TP | 单 NUMA 节点 | ✅ | ✅ |

---

## 6. 错误处理需求

### 6.1 参数验证

| 场景 | 行为 |
|------|------|
| `mode` 无效 | `ValueError("Unknown mode: {mode}")` |
| 推理模式调用 SFT 方法 | `RuntimeError("forward_sft() not available in inference mode")` |
| SFT 模式调用推理特有方法 | `RuntimeError("submit_forward() not available in SFT mode")` |
| `method` 与 `mode` 不匹配 | `ValueError("Method {method} not supported in {mode} mode")` |

### 6.2 运行时检查

| 场景 | 行为 |
|------|------|
| 未加载权重时调用 forward | `RuntimeError("Weights not loaded")` |
| 未初始化 LoRA 时调用 backward | `RuntimeError("LoRA weights not initialized")` |
| cache_idx 超出范围 | `RuntimeError("Invalid cache index")` |

---

## 7. 使用场景示例

### 7.1 推理场景

```python
# 创建推理 Wrapper
wrapper = KTMoEWrapper(
    layer_idx=0,
    num_experts=256,
    num_experts_per_tok=8,
    hidden_size=7168,
    moe_intermediate_size=2048,
    num_gpu_experts=0,
    cpuinfer_threads=60,
    threadpool_count=4,
    weight_path="/path/to/weights",
    chunked_prefill_size=25600,
    method="AMXINT4",
    mode="inference"
)

# 加载权重
physical_map = torch.arange(256, dtype=torch.int64)
wrapper.load_weights(physical_map)

# 推理
hidden_states = torch.randn(1, 7168, dtype=torch.bfloat16).cuda()
topk_ids = torch.randint(0, 256, (1, 8)).cuda()
topk_weights = torch.rand(1, 8, dtype=torch.float32).cuda()

output = wrapper.forward(hidden_states, topk_ids, topk_weights, cuda_stream)
```

### 7.2 SFT 场景

```python
# 创建 SFT Wrapper
wrapper = KTMoEWrapper(
    layer_idx=0,
    num_experts=256,
    num_experts_per_tok=8,
    hidden_size=7168,
    moe_intermediate_size=2048,
    num_gpu_experts=0,
    cpuinfer_threads=60,
    threadpool_count=4,
    weight_path="/path/to/weights",
    chunked_prefill_size=25600,
    method="AMXBF16_SFT",
    mode="sft",
    lora_rank=16,
    lora_alpha=32.0
)

# 加载基础权重
wrapper.load_weights(physical_map)

# 初始化 LoRA 权重
gate_lora_a = torch.zeros(256, 16, 7168, dtype=torch.bfloat16)
gate_lora_b = torch.zeros(256, 2048, 16, dtype=torch.bfloat16)
# ... 其他 LoRA 权重
wrapper.init_lora_weights(gate_lora_a, gate_lora_b, ...)

# 训练循环
for batch in dataloader:
    # 前向传播
    output = wrapper.forward_sft(
        hidden_states, expert_ids, weights,
        save_for_backward=True
    )

    # 计算损失
    loss = criterion(output, target)

    # 反向传播
    grad_input, grad_loras = wrapper.backward(grad_output)

    # 更新 LoRA 权重（使用外部优化器）
    optimizer.step()

    # 同步更新后的权重到 C++
    wrapper.update_lora_weights()
```
