# SFT + KTWrapper 算子接口

## 1. Python API 接口

### 1.1 KTMoEWrapper 工厂类

#### 构造函数签名

```python
class KTMoEWrapper:
    def __new__(
        cls,
        # ========== 基础参数（推理和 SFT 共用）==========
        layer_idx: int,                                    # 层索引
        num_experts: int,                                  # 专家总数
        num_experts_per_tok: int,                          # 每 token 激活的专家数 (top-k)
        hidden_size: int,                                  # 隐藏层维度
        moe_intermediate_size: int,                        # MoE 中间层维度
        num_gpu_experts: int,                              # GPU 上的专家数（SFT 通常为 0）
        cpuinfer_threads: int,                             # CPU 推理线程数
        threadpool_count: int,                             # NUMA 子池数量（TP 数量）
        weight_path: str,                                  # 权重路径
        chunked_prefill_size: int,                         # 分块预填充大小

        # ========== 推理特有参数 ==========
        cpu_save: bool = False,                            # 是否保存到 CPU 内存
        max_deferred_experts_per_token: Optional[int] = None,  # 延迟执行的专家数

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

        # ========== SFT 特有参数（mode="sft" 时有效）==========
        lora_rank: int = 16,                               # LoRA 低秩矩阵的秩
        lora_alpha: float = 32.0,                          # LoRA 缩放因子
        max_cache_depth: int = 1,                          # 前向缓存深度
    ) -> Union[BaseMoEWrapper, BaseSFTMoEWrapper]:
        ...
```

#### 参数说明

| 参数 | 类型 | 默认值 | 模式 | 说明 |
|------|------|--------|------|------|
| `layer_idx` | int | - | 共用 | 层索引，用于加载对应层的权重 |
| `num_experts` | int | - | 共用 | 专家总数（如 DeepSeek-V3 为 256） |
| `num_experts_per_tok` | int | - | 共用 | 每 token 激活的专家数（top-k 值） |
| `hidden_size` | int | - | 共用 | 隐藏层维度（如 7168） |
| `moe_intermediate_size` | int | - | 共用 | MoE 中间层维度（如 2048） |
| `num_gpu_experts` | int | - | 共用 | GPU 上的专家数（SFT 通常为 0） |
| `cpuinfer_threads` | int | - | 共用 | CPU 推理线程数 |
| `threadpool_count` | int | - | 共用 | NUMA 子池数量（用于 TP 并行） |
| `weight_path` | str | - | 共用 | 权重文件所在目录 |
| `chunked_prefill_size` | int | - | 共用 | 分块预填充大小 |
| `cpu_save` | bool | False | 推理 | 是否将结果保存到 CPU 内存 |
| `max_deferred_experts_per_token` | int | None | 推理 | 每 token 延迟执行的最大专家数 |
| `method` | str | "AMXINT4" | 共用 | 后端方法（见下表） |
| `mode` | str | "inference" | 共用 | 模式选择 |
| `lora_rank` | int | 16 | SFT | LoRA 低秩矩阵的秩 |
| `lora_alpha` | float | 32.0 | SFT | LoRA 缩放因子 |
| `max_cache_depth` | int | 1 | SFT | 前向缓存深度 |

---

### 1.2 method 参数值

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

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

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

| method | 后端类 | 量化类型 | 说明 |
|--------|--------|----------|------|
| `AMXBF16_SFT` | AMXSFTMoEWrapper | BF16 | AMX BF16 精度训练 |
| `AMXINT8_SFT` | AMXSFTMoEWrapper | INT8 | AMX INT8 量化训练 |
| `AMXINT4_SFT` | AMXSFTMoEWrapper | INT4 | AMX INT4 量化训练 |
| `AMXINT4_KGroup_SFT` | AMXSFTMoEWrapper | INT4_KGroup | AMX INT4 K-Group 训练 |

---

### 1.3 BaseMoEWrapper 推理接口

| 方法 | 签名 | 说明 |
|------|------|------|
| `load_weights` | `(physical_to_logical_map: Tensor) -> None` | 加载预量化权重 |
| `load_weights_from_tensors` | `(gate, up, down, map) -> None` | 在线量化加载 |
| `forward` | `(hidden_states, topk_ids, topk_weights, cuda_stream) -> Tensor` | 同步前向 |
| `submit_forward` | `(hidden_states, topk_ids, topk_weights, cuda_stream) -> None` | 异步提交 |
| `sync_forward` | `(hidden_states, cuda_stream) -> Tensor` | 同步获取结果 |
| `select_deferred_experts` | `(expert_ids, scores, protected_k) -> Tuple` | 选择延迟专家 |

### 1.4 BaseSFTMoEWrapper SFT 接口

| 方法 | 签名 | 说明 |
|------|------|------|
| `load_weights` | `(physical_to_logical_map: Tensor) -> None` | 加载权重 |
| `init_lora_weights` | `(gate_a, gate_b, up_a, up_b, down_a, down_b) -> None` | 初始化 LoRA |
| `forward_sft` | `(hidden_states, expert_ids, weights, save_for_backward) -> Tensor` | SFT 前向 |
| `backward` | `(grad_output) -> Tuple[Tensor, Dict]` | 反向传播 |
| `update_lora_weights` | `() -> None` | 同步 LoRA 权重 |

---

## 2. C++ 绑定接口

### 2.1 MOESFTConfig 结构

| 字段 | 类型 | 默认值 | 说明 |
|------|------|--------|------|
| `expert_num` | int | 0 | 专家总数 |
| `hidden_size` | int | 0 | 隐藏层维度 |
| `intermediate_size` | int | 0 | 中间层维度 |
| `experts_per_token` | int | 0 | 每 token 专家数 |
| `weight_path` | string | "" | 权重路径 |
| `layer_idx` | int | 0 | 层索引 |
| `tp_size` | int | 1 | TP 并行数 |
| `lora_rank` | int | 16 | LoRA 秩 |
| `lora_alpha` | float | 32.0 | LoRA alpha |
| `max_cache_depth` | int | 1 | 缓存深度 |

### 2.2 TP_MOE_SFT 方法

| 方法 | 参数 | 说明 |
|------|------|------|
| `warm_up_task` | `()` | 预热任务 |
| `load_weights_task` | `(physical_map_ptr)` | 加载权重 |
| `forward_sft_task` | `(qlen, k, expert_ids, weights, input, output, save)` | 前向任务 |
| `backward_task` | `(qlen, k, grad_out, grad_in, grad_lora_a/b...)` | 反向任务 |
| `update_lora_weights_task` | `(gate_a, gate_b, up_a, up_b, down_a, down_b)` | 更新 LoRA |

---

## 3. 使用示例

### 3.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(1024, 7168, dtype=torch.bfloat16).cuda()
topk_ids = torch.randint(0, 256, (1024, 8), dtype=torch.int64).cuda()
topk_weights = torch.rand(1024, 8, dtype=torch.float32).cuda()
cuda_stream = torch.cuda.current_stream().cuda_stream

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

### 3.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)
up_lora_a = torch.zeros(256, 16, 7168, dtype=torch.bfloat16)
up_lora_b = torch.zeros(256, 2048, 16, dtype=torch.bfloat16)
down_lora_a = torch.zeros(256, 16, 2048, dtype=torch.bfloat16)
down_lora_b = torch.zeros(256, 7168, 16, dtype=torch.bfloat16)

wrapper.init_lora_weights(
    gate_lora_a, gate_lora_b,
    up_lora_a, up_lora_b,
    down_lora_a, down_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()
```

---

## 4. 错误码和异常

### 4.1 ValueError（参数错误）

| 错误消息 | 原因 | 解决方法 |
|----------|------|----------|
| `Unknown mode: '{mode}'` | mode 参数无效 | 使用 "inference" 或 "sft" |
| `Method '{method}' not supported` | method 与 mode 不匹配 | 参考 method 表格 |
| `num_experts must be positive` | 参数 <= 0 | 使用正整数 |
| `{name} shape mismatch` | LoRA 权重形状错误 | 检查维度 |

### 4.2 RuntimeError（运行时错误）

| 错误消息 | 原因 | 解决方法 |
|----------|------|----------|
| `Weights not loaded` | 未调用 load_weights() | 先加载权重 |
| `LoRA weights not initialized` | 未调用 init_lora_weights() | 先初始化 LoRA |
| `forward_sft() not available in inference mode` | 模式不匹配 | 使用 mode="sft" |
| `Forward cache full` | 缓存超过 max_cache_depth | 调用 backward() 释放 |
| `No forward cache available` | 无缓存数据 | 先调用 forward_sft() |
