# SFT + KTWrapper 功能架构设计

## 1. 总体架构图

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                            Python API 层                                     │
│  ┌───────────────────────────────────────────────────────────────────────┐  │
│  │                  KTMoEWrapper (统一工厂入口)                            │  │
│  │         mode="inference" → 推理      mode="sft" → SFT                 │  │
│  └───────────────────────────────────────────────────────────────────────┘  │
│                                    │                                         │
│  ┌───────────────────────────────────────────────────────────────────────┐  │
│  │                        _MoEBase (共享基类)                              │  │
│  │    - CPUInfer 单例管理 (_cpu_infer_instance)                           │  │
│  │    - WorkerPoolConfig 构建                                             │  │
│  │    - 基础配置验证                                                       │  │
│  └───────────────────────────────────────────────────────────────────────┘  │
│               │                                      │                       │
│  ┌─────────────────────────┐          ┌─────────────────────────┐           │
│  │    BaseMoEWrapper       │          │   BaseSFTMoEWrapper     │           │
│  │    (推理基类-不变)       │          │     (SFT基类-新增)       │           │
│  │  ───────────────────    │          │  ───────────────────    │           │
│  │  - forward()            │          │  - forward_sft()        │           │
│  │  - submit_forward()     │          │  - backward()           │           │
│  │  - sync_forward()       │          │  - update_lora_weights()│           │
│  │  - load_weights()       │          │  - init_lora_weights()  │           │
│  │  - KExpertsCPUBuffer    │          │  - KExpertsSFTBuffer    │           │
│  └─────────────────────────┘          └─────────────────────────┘           │
│               │                                      │                       │
│  ┌─────────────────────────┐          ┌─────────────────────────┐           │
│  │  AMXMoEWrapper          │          │  AMXSFTMoEWrapper       │           │
│  │  NativeMoEWrapper       │          │  (新增)                  │           │
│  │  LlamafileMoEWrapper    │          │                         │           │
│  │  GeneralMoEWrapper      │          │                         │           │
│  └─────────────────────────┘          └─────────────────────────┘           │
└─────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                            C++ 后端层                                        │
│  ┌───────────────────────────────────────────────────────────────────────┐  │
│  │                     CPUInfer (推理引擎单例)                             │  │
│  │              WorkerPool → TaskQueue → Worker Threads                  │  │
│  └───────────────────────────────────────────────────────────────────────┘  │
│               │                                      │                       │
│  ┌─────────────────────────┐          ┌─────────────────────────┐           │
│  │      TP_MOE<T>          │          │     TP_MOE_SFT<T>       │           │
│  │    (推理 MoE 基类)       │ ◄─继承── │    (SFT MoE 基类)        │           │
│  └─────────────────────────┘          └─────────────────────────┘           │
│               │                                      │                       │
│  ┌─────────────────────────┐          ┌─────────────────────────┐           │
│  │    AMX_MOE_TP<T>        │          │   AMX_SFT_MOE_TP<T>     │           │
│  │   (AMX 推理实现)         │ ◄─继承── │    (AMX SFT 实现)        │           │
│  └─────────────────────────┘          └─────────────────────────┘           │
└─────────────────────────────────────────────────────────────────────────────┘
```

---

## 2. 类设计

### 2.1 _MoEBase 共享基类

```python
class _MoEBase:
    """推理和 SFT 共享的基类，管理 CPUInfer 单例"""

    _cpu_infer_instance: ClassVar[Optional[CPUInfer]] = None
    _cpu_infer_lock: ClassVar[threading.Lock] = threading.Lock()

    @classmethod
    def _get_cpu_infer(
        cls,
        cpuinfer_threads: int,
        threadpool_count: int
    ) -> CPUInfer:
        """获取或创建 CPUInfer 单例"""
        with cls._cpu_infer_lock:
            if cls._cpu_infer_instance is None:
                worker_config = kt_kernel_ext.WorkerPoolConfig()
                worker_config.max_threads_per_subpool = cpuinfer_threads
                worker_config.subpool_count = threadpool_count
                cls._cpu_infer_instance = kt_kernel_ext.CPUInfer(worker_config)
            return cls._cpu_infer_instance

    @classmethod
    def _validate_base_config(cls, num_experts: int, hidden_size: int, ...):
        """基础参数验证"""
        if num_experts <= 0:
            raise ValueError("num_experts must be positive")
        if hidden_size <= 0:
            raise ValueError("hidden_size must be positive")
        # ...
```

### 2.2 BaseMoEWrapper（推理基类-不变）

```python
class BaseMoEWrapper(_MoEBase, ABC):
    """推理 MoE 的基类（保持现有实现不变）"""

    def __init__(
        self,
        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,
    ):
        # 获取共享的 CPUInfer 实例
        self.cpu_infer = self._get_cpu_infer(cpuinfer_threads, threadpool_count)
        # ... 现有初始化逻辑 ...

    @abstractmethod
    def load_weights(self, physical_to_logical_map: torch.Tensor) -> None: ...

    @abstractmethod
    def forward(self, hidden_states, topk_ids, topk_weights, cuda_stream) -> torch.Tensor: ...

    def submit_forward(self, hidden_states, topk_ids, topk_weights, cuda_stream) -> None: ...

    def sync_forward(self, hidden_states, cuda_stream) -> torch.Tensor: ...
```

### 2.3 BaseSFTMoEWrapper（SFT 基类-新增）

```python
class BaseSFTMoEWrapper(_MoEBase, ABC):
    """SFT MoE 的基类（新增）"""

    def __init__(
        self,
        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,
        # SFT 特有参数
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        max_cache_depth: int = 1,
    ):
        # 获取共享的 CPUInfer 实例
        self.cpu_infer = self._get_cpu_infer(cpuinfer_threads, threadpool_count)

        # SFT 特有配置
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.lora_scaling = lora_alpha / lora_rank
        self.max_cache_depth = max_cache_depth

        # LoRA 权重占位符
        self.gate_lora_a: Optional[torch.Tensor] = None
        self.gate_lora_b: Optional[torch.Tensor] = None
        self.up_lora_a: Optional[torch.Tensor] = None
        self.up_lora_b: Optional[torch.Tensor] = None
        self.down_lora_a: Optional[torch.Tensor] = None
        self.down_lora_b: Optional[torch.Tensor] = None

        # 权重加载状态
        self._weights_loaded: bool = False
        self._lora_initialized: bool = False

    @abstractmethod
    def load_weights(self, physical_to_logical_map: torch.Tensor) -> None:
        """加载基础权重"""
        ...

    @abstractmethod
    def init_lora_weights(
        self,
        gate_lora_a: torch.Tensor,
        gate_lora_b: torch.Tensor,
        up_lora_a: torch.Tensor,
        up_lora_b: torch.Tensor,
        down_lora_a: torch.Tensor,
        down_lora_b: torch.Tensor,
    ) -> None:
        """初始化 LoRA 权重"""
        ...

    @abstractmethod
    def forward_sft(
        self,
        hidden_states: torch.Tensor,
        expert_ids: torch.Tensor,
        weights: torch.Tensor,
        save_for_backward: bool = True,
    ) -> torch.Tensor:
        """SFT 前向传播（带梯度缓存）"""
        ...

    @abstractmethod
    def backward(
        self,
        grad_output: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """反向传播，返回输入梯度和 LoRA 梯度"""
        ...

    @abstractmethod
    def update_lora_weights(self) -> None:
        """同步 LoRA 权重到 C++ 后端"""
        ...
```

### 2.4 KTMoEWrapper 工厂类修改

```python
class KTMoEWrapper:
    """统一的 MoE Wrapper 工厂类"""

    # 推理模式支持的 method
    INFERENCE_METHODS = {
        "AMXINT4", "AMXINT8",  # AMX 量化
        "RAWINT4", "FP8",      # Native 量化
        "LLAMAFILE",           # GGUF 格式
        "MOE_INT4", "MOE_INT8" # 通用内核
    }

    # SFT 模式支持的 method
    SFT_METHODS = {
        "AMXBF16_SFT",         # AMX BF16
        "AMXINT8_SFT",         # AMX INT8
        "AMXINT4_SFT",         # AMX INT4
        "AMXINT4_KGroup_SFT",  # AMX INT4 K-Group
    }

    def __new__(
        cls,
        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",
        # SFT 特有参数
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        max_cache_depth: int = 1,
    ):
        # 1. 验证 mode 参数
        if mode not in ("inference", "sft"):
            raise ValueError(f"Unknown mode: {mode}. Must be 'inference' or 'sft'")

        # 2. 验证 method 与 mode 的匹配
        if mode == "inference" and method not in cls.INFERENCE_METHODS:
            raise ValueError(f"Method '{method}' not supported in inference mode")
        if mode == "sft" and method not in cls.SFT_METHODS:
            raise ValueError(f"Method '{method}' not supported in SFT mode")

        # 3. 根据 mode 创建对应的 Wrapper
        base_kwargs = {
            "layer_idx": layer_idx,
            "num_experts": num_experts,
            "num_experts_per_tok": num_experts_per_tok,
            "hidden_size": hidden_size,
            "moe_intermediate_size": moe_intermediate_size,
            "num_gpu_experts": num_gpu_experts,
            "cpuinfer_threads": cpuinfer_threads,
            "threadpool_count": threadpool_count,
            "weight_path": weight_path,
            "chunked_prefill_size": chunked_prefill_size,
        }

        if mode == "inference":
            return cls._create_inference_wrapper(
                method=method,
                cpu_save=cpu_save,
                max_deferred_experts_per_token=max_deferred_experts_per_token,
                **base_kwargs
            )
        else:  # mode == "sft"
            return cls._create_sft_wrapper(
                method=method,
                lora_rank=lora_rank,
                lora_alpha=lora_alpha,
                max_cache_depth=max_cache_depth,
                **base_kwargs
            )

    @classmethod
    def _create_inference_wrapper(cls, method: str, **kwargs):
        """创建推理模式的 Wrapper（现有逻辑）"""
        if method in ("AMXINT4", "AMXINT8"):
            from .utils.amx import AMXMoEWrapper
            return AMXMoEWrapper(method=method, **kwargs)
        elif method in ("RAWINT4", "FP8"):
            from .utils.native import NativeMoEWrapper
            return NativeMoEWrapper(method=method, **kwargs)
        elif method == "LLAMAFILE":
            from .utils.llamafile import LlamafileMoEWrapper
            return LlamafileMoEWrapper(**kwargs)
        elif method in ("MOE_INT4", "MOE_INT8"):
            from .utils.general import GeneralMoEWrapper
            return GeneralMoEWrapper(method=method, **kwargs)

    @classmethod
    def _create_sft_wrapper(cls, method: str, **kwargs):
        """创建 SFT 模式的 Wrapper（新增）"""
        if method in ("AMXBF16_SFT", "AMXINT8_SFT", "AMXINT4_SFT", "AMXINT4_KGroup_SFT"):
            from .utils.amx_sft import AMXSFTMoEWrapper
            return AMXSFTMoEWrapper(method=method, **kwargs)
```

---

## 3. 缓冲区设计

### 3.1 KExpertsCPUBuffer（推理-不变）

```python
class KExpertsCPUBuffer:
    """推理模式的 CPU 缓冲区管理（保持现有实现）"""

    # 双缓冲结构：7 元组
    # (input, immediate_ids, deferred_ids, weights, output, bsz, output_gpu)

    capture_buffers: ClassVar[Dict[int, Any]] = {}

    @classmethod
    def get_buffer(
        cls,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
        num_deferred_experts: int,
        capture_slot: int,
        # ... 其他参数
    ) -> Tuple:
        """获取或创建缓冲区"""
        # ... 现有实现 ...
```

### 3.2 KExpertsSFTBuffer（SFT-新增）

```python
class KExpertsSFTBuffer:
    """SFT 模式的 CPU 缓冲区管理（新增）"""

    capture_buffers: ClassVar[Dict[int, "KExpertsSFTBuffer"]] = {}

    def __init__(
        self,
        qlen: int,
        hidden_size: int,
        moe_intermediate_size: int,
        num_experts: int,
        num_experts_per_tok: int,
        lora_rank: int,
        dtype: torch.dtype = torch.bfloat16,
    ):
        # 前向缓冲
        self.input_cpu = torch.empty(
            (qlen, hidden_size), dtype=dtype, device="cpu", pin_memory=True
        )
        self.expert_ids_cpu = torch.empty(
            (qlen, num_experts_per_tok), dtype=torch.int64, device="cpu", pin_memory=True
        )
        self.weights_cpu = torch.empty(
            (qlen, num_experts_per_tok), dtype=torch.float32, device="cpu", pin_memory=True
        )
        self.output_cpu = torch.empty(
            (qlen, hidden_size), dtype=dtype, device="cpu", pin_memory=True
        )

        # 反向缓冲
        self.grad_output_cpu = torch.empty(
            (qlen, hidden_size), dtype=dtype, device="cpu", pin_memory=True
        )
        self.grad_input_cpu = torch.empty(
            (qlen, hidden_size), dtype=dtype, device="cpu", pin_memory=True
        )

        # LoRA 梯度缓冲（6 个）
        self.grad_gate_lora_a = torch.empty(
            (num_experts, lora_rank, hidden_size), dtype=dtype, device="cpu"
        )
        self.grad_gate_lora_b = torch.empty(
            (num_experts, moe_intermediate_size, lora_rank), dtype=dtype, device="cpu"
        )
        self.grad_up_lora_a = torch.empty(
            (num_experts, lora_rank, hidden_size), dtype=dtype, device="cpu"
        )
        self.grad_up_lora_b = torch.empty(
            (num_experts, moe_intermediate_size, lora_rank), dtype=dtype, device="cpu"
        )
        self.grad_down_lora_a = torch.empty(
            (num_experts, lora_rank, moe_intermediate_size), dtype=dtype, device="cpu"
        )
        self.grad_down_lora_b = torch.empty(
            (num_experts, hidden_size, lora_rank), dtype=dtype, device="cpu"
        )

    @classmethod
    def get_buffer(
        cls,
        qlen: int,
        hidden_size: int,
        moe_intermediate_size: int,
        num_experts: int,
        num_experts_per_tok: int,
        lora_rank: int,
        dtype: torch.dtype = torch.bfloat16,
    ) -> "KExpertsSFTBuffer":
        """获取或创建 SFT 缓冲区"""
        key = (qlen, hidden_size, moe_intermediate_size, num_experts, num_experts_per_tok, lora_rank, dtype)

        if key not in cls.capture_buffers:
            cls.capture_buffers[key] = cls(
                qlen=qlen,
                hidden_size=hidden_size,
                moe_intermediate_size=moe_intermediate_size,
                num_experts=num_experts,
                num_experts_per_tok=num_experts_per_tok,
                lora_rank=lora_rank,
                dtype=dtype,
            )

        return cls.capture_buffers[key]
```

### 3.3 缓冲区对比

| 特性 | KExpertsCPUBuffer (推理) | KExpertsSFTBuffer (SFT) |
|------|-------------------------|-------------------------|
| 用途 | GPU-CPU 异步传输 | 前向/反向数据存储 |
| 双缓冲 | ✅ 支持 | ❌ 不需要 |
| 梯度缓冲 | ❌ 无 | ✅ 6 个 LoRA 梯度 |
| Pin Memory | ✅ 部分 | ✅ 全部 |
| 缓存策略 | capture_slot 索引 | 参数组合哈希 |

---

## 4. 状态管理

### 4.1 推理模式（无状态）

```
┌─────────────────────────────────────────┐
│           推理状态管理                   │
│                                         │
│  每次 forward() 调用：                   │
│  1. 获取缓冲槽位 (capture_slot)          │
│  2. 复制输入数据                         │
│  3. 提交任务到 CPUInfer                  │
│  4. 同步等待结果                         │
│  5. 返回输出（缓冲槽位可复用）            │
│                                         │
│  特点：                                  │
│  - 无状态，每次调用独立                   │
│  - 双缓冲槽位管理                        │
│  - 无需保存中间结果                       │
└─────────────────────────────────────────┘
```

### 4.2 SFT 模式（有状态）

```
┌─────────────────────────────────────────┐
│            SFT 状态管理                  │
│                                         │
│  forward_sft() 调用：                    │
│  1. 保存输入到 ForwardCache              │
│  2. 执行前向计算                         │
│  3. 保存激活值到 ForwardCache            │
│  4. cache_idx 自增                       │
│                                         │
│  backward() 调用：                       │
│  1. 从 ForwardCache 恢复激活值           │
│  2. 计算 LoRA 梯度                       │
│  3. cache_idx 自减                       │
│  4. 返回梯度                             │
│                                         │
│  状态变量：                              │
│  - cache_idx: int                       │
│  - ForwardCache: 栈结构                  │
│  - LoRA 权重: 6 个张量                   │
│  - _weights_loaded: bool                │
│  - _lora_initialized: bool              │
└─────────────────────────────────────────┘
```

### 4.3 ForwardCache 结构（C++ 层）

```cpp
// 位于 operators/amx/sft_moe.hpp
struct ForwardCache {
    // 输入缓存
    std::vector<std::vector<T>> input_cache;      // [cache_depth][qlen * hidden]
    std::vector<std::vector<int64_t>> expert_ids_cache;
    std::vector<std::vector<float>> weights_cache;

    // 激活值缓存（用于反向传播）
    std::vector<std::vector<T>> gate_output_cache;
    std::vector<std::vector<T>> up_output_cache;
    std::vector<std::vector<T>> act_output_cache;

    // 缓存深度管理
    int max_depth;
    int current_depth;

    void push(/* 前向数据 */);
    void pop(/* 恢复数据 */);
    void clear();
};
```

---

## 5. 数据流设计

### 5.1 推理数据流

```
┌─────────────────────────────────────────────────────────────────────┐
│                        推理数据流                                    │
│                                                                     │
│  GPU                          CPU                                   │
│  ┌─────────┐                  ┌─────────────────────────────────┐   │
│  │ hidden  │ ───D2H Copy───► │ KExpertsCPUBuffer.input          │   │
│  │ states  │                  │            │                    │   │
│  └─────────┘                  │            ▼                    │   │
│                               │    ┌─────────────┐              │   │
│                               │    │  CPUInfer   │              │   │
│                               │    │  submit()   │              │   │
│                               │    └──────┬──────┘              │   │
│                               │           ▼                     │   │
│                               │    ┌─────────────┐              │   │
│                               │    │  TP_MOE     │              │   │
│                               │    │  forward()  │              │   │
│                               │    └──────┬──────┘              │   │
│                               │           ▼                     │   │
│  ┌─────────┐                  │ KExpertsCPUBuffer.output        │   │
│  │ output  │ ◄───H2D Copy─── │                                  │   │
│  └─────────┘                  └─────────────────────────────────┘   │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘
```

### 5.2 SFT 数据流

```
┌─────────────────────────────────────────────────────────────────────┐
│                         SFT 数据流                                   │
│                                                                     │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                      前向传播                                  │  │
│  │  GPU                          CPU                             │  │
│  │  ┌─────────┐                  ┌─────────────────────────────┐ │  │
│  │  │ hidden  │ ───D2H Copy───► │ KExpertsSFTBuffer.input      │ │  │
│  │  │ states  │                  │            │                │ │  │
│  │  └─────────┘                  │            ▼                │ │  │
│  │                               │    ┌─────────────┐          │ │  │
│  │                               │    │  forward_   │          │ │  │
│  │                               │    │  sft_task() │          │ │  │
│  │                               │    └──────┬──────┘          │ │  │
│  │                               │           │                 │ │  │
│  │                               │           ▼                 │ │  │
│  │                               │    ┌─────────────┐          │ │  │
│  │                               │    │ ForwardCache│ (保存)    │ │  │
│  │                               │    └──────┬──────┘          │ │  │
│  │                               │           ▼                 │ │  │
│  │  ┌─────────┐                  │ KExpertsSFTBuffer.output    │ │  │
│  │  │ output  │ ◄───H2D Copy─── │                              │ │  │
│  │  └─────────┘                  └─────────────────────────────┘ │  │
│  └───────────────────────────────────────────────────────────────┘  │
│                                                                     │
│  ┌───────────────────────────────────────────────────────────────┐  │
│  │                      反向传播                                  │  │
│  │  GPU                          CPU                             │  │
│  │  ┌─────────┐                  ┌─────────────────────────────┐ │  │
│  │  │ grad_   │ ───D2H Copy───► │ KExpertsSFTBuffer.grad_out   │ │  │
│  │  │ output  │                  │            │                │ │  │
│  │  └─────────┘                  │            ▼                │ │  │
│  │                               │    ┌─────────────┐          │ │  │
│  │                               │    │ ForwardCache│ (恢复)    │ │  │
│  │                               │    └──────┬──────┘          │ │  │
│  │                               │           ▼                 │ │  │
│  │                               │    ┌─────────────┐          │ │  │
│  │                               │    │  backward_  │          │ │  │
│  │                               │    │  task()     │          │ │  │
│  │                               │    └──────┬──────┘          │ │  │
│  │                               │           │                 │ │  │
│  │                               │     ┌─────┴─────┐           │ │  │
│  │                               │     ▼           ▼           │ │  │
│  │  ┌─────────┐                  │ grad_input  grad_loras      │ │  │
│  │  │ grad_   │ ◄───H2D Copy─── │                              │ │  │
│  │  │ input   │                  └─────────────────────────────┘ │  │
│  │  └─────────┘                                                  │  │
│  └───────────────────────────────────────────────────────────────┘  │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘
```

---

## 6. 错误处理设计

### 6.1 模式检查

```python
class BaseSFTMoEWrapper(_MoEBase, ABC):
    """SFT 模式特有方法的错误处理"""

    def forward(self, *args, **kwargs):
        """推理模式的 forward 在 SFT 中不可用"""
        raise RuntimeError(
            "forward() is not available in SFT mode. "
            "Use forward_sft() instead."
        )

    def submit_forward(self, *args, **kwargs):
        """异步前向在 SFT 中不可用"""
        raise RuntimeError(
            "submit_forward() is not available in SFT mode. "
            "SFT mode uses synchronous forward_sft()."
        )

    def sync_forward(self, *args, **kwargs):
        """异步同步在 SFT 中不可用"""
        raise RuntimeError(
            "sync_forward() is not available in SFT mode."
        )


class BaseMoEWrapper(_MoEBase, ABC):
    """推理模式特有方法的错误处理"""

    def forward_sft(self, *args, **kwargs):
        """SFT 前向在推理模式中不可用"""
        raise RuntimeError(
            "forward_sft() is not available in inference mode. "
            "Use forward() instead."
        )

    def backward(self, *args, **kwargs):
        """反向传播在推理模式中不可用"""
        raise RuntimeError(
            "backward() is not available in inference mode."
        )

    def init_lora_weights(self, *args, **kwargs):
        """LoRA 初始化在推理模式中不可用"""
        raise RuntimeError(
            "init_lora_weights() is not available in inference mode."
        )
```

### 6.2 状态检查

```python
class BaseSFTMoEWrapper(_MoEBase, ABC):
    """运行时状态检查"""

    def forward_sft(self, hidden_states, expert_ids, weights, save_for_backward=True):
        # 检查权重是否已加载
        if not self._weights_loaded:
            raise RuntimeError(
                "Weights not loaded. Call load_weights() first."
            )

        # 检查 LoRA 是否已初始化（如果需要训练）
        if save_for_backward and not self._lora_initialized:
            raise RuntimeError(
                "LoRA weights not initialized. "
                "Call init_lora_weights() first, or set save_for_backward=False."
            )

        # ... 前向逻辑 ...

    def backward(self, grad_output):
        # 检查是否有缓存的前向数据
        if self._cache_depth <= 0:
            raise RuntimeError(
                "No forward cache available. "
                "Call forward_sft() with save_for_backward=True first."
            )

        # ... 反向逻辑 ...
```

### 6.3 参数验证

```python
def __new__(cls, ..., mode: str = "inference", method: str = "AMXINT4", ...):
    # 模式验证
    if mode not in ("inference", "sft"):
        raise ValueError(
            f"Unknown mode: '{mode}'. Must be 'inference' or 'sft'."
        )

    # method 与 mode 匹配验证
    if mode == "inference" and method not in cls.INFERENCE_METHODS:
        raise ValueError(
            f"Method '{method}' is not supported in inference mode. "
            f"Supported methods: {cls.INFERENCE_METHODS}"
        )

    if mode == "sft" and method not in cls.SFT_METHODS:
        raise ValueError(
            f"Method '{method}' is not supported in SFT mode. "
            f"Supported methods: {cls.SFT_METHODS}"
        )

    # SFT 参数验证
    if mode == "sft":
        if lora_rank <= 0:
            raise ValueError(f"lora_rank must be positive, got {lora_rank}")
        if lora_alpha <= 0:
            raise ValueError(f"lora_alpha must be positive, got {lora_alpha}")
        if max_cache_depth <= 0:
            raise ValueError(f"max_cache_depth must be positive, got {max_cache_depth}")
```

---

## 7. 线程安全设计

### 7.1 CPUInfer 单例保护

```python
class _MoEBase:
    _cpu_infer_instance: ClassVar[Optional[CPUInfer]] = None
    _cpu_infer_lock: ClassVar[threading.Lock] = threading.Lock()

    @classmethod
    def _get_cpu_infer(cls, cpuinfer_threads: int, threadpool_count: int) -> CPUInfer:
        """线程安全的单例获取"""
        with cls._cpu_infer_lock:
            if cls._cpu_infer_instance is None:
                # 双重检查锁定
                worker_config = kt_kernel_ext.WorkerPoolConfig()
                worker_config.max_threads_per_subpool = cpuinfer_threads
                worker_config.subpool_count = threadpool_count
                cls._cpu_infer_instance = kt_kernel_ext.CPUInfer(worker_config)
            return cls._cpu_infer_instance
```

### 7.2 缓冲区访问保护

```python
class KExpertsSFTBuffer:
    _buffer_lock: ClassVar[threading.Lock] = threading.Lock()
    capture_buffers: ClassVar[Dict[tuple, "KExpertsSFTBuffer"]] = {}

    @classmethod
    def get_buffer(cls, ...) -> "KExpertsSFTBuffer":
        """线程安全的缓冲区获取"""
        key = (qlen, hidden_size, ...)

        with cls._buffer_lock:
            if key not in cls.capture_buffers:
                cls.capture_buffers[key] = cls(...)
            return cls.capture_buffers[key]
```

---

## 8. 文件结构

### 8.1 修改后的文件结构

```
kt-kernel/python/
├── experts.py              # 修改：添加 mode 参数和 SFT 分支
├── experts_base.py         # 修改：提取 _MoEBase 共享基类
├── experts_sft.py          # 新增：BaseSFTMoEWrapper 和 KExpertsSFTBuffer
└── utils/
    ├── amx.py              # 不变：AMXMoEWrapper
    ├── amx_sft.py          # 新增：AMXSFTMoEWrapper
    ├── native.py           # 不变
    ├── llamafile.py        # 不变
    └── general.py          # 不变
```

### 8.2 导入关系

```python
# experts.py
from .experts_base import BaseMoEWrapper, _MoEBase
from .experts_sft import BaseSFTMoEWrapper

# experts_base.py
from .experts_base import _MoEBase  # 共享基类

# experts_sft.py
from .experts_base import _MoEBase

# utils/amx_sft.py
from ..experts_sft import BaseSFTMoEWrapper
```

---

## 9. forward_sft 与 forward 关系设计决策

### 9.1 背景问题

推理模式的 `forward()` 方法中包含多个优化分支（延迟专家执行、跨层增量执行、双缓冲、异步同步等）。
如果 `forward_sft()` 完全独立实现，这些优化无法自动复用。

### 9.2 两种设计方案

#### 方案 A：调用复用
```python
def forward_sft(self, hidden_states, expert_ids, weights, save_for_backward=True):
    # 调用推理的 forward，禁用不兼容的优化
    output = self.forward(
        hidden_states, expert_ids, weights,
        use_deferred_experts=False,  # 禁用延迟专家
        use_async=False,             # 禁用异步
        use_double_buffer=False,     # 禁用双缓冲
    )

    if save_for_backward:
        # 保存激活值到 ForwardCache
        self._save_for_backward(hidden_states, expert_ids, weights, output)

    return output
```

**优点**：
- 自动复用推理优化
- 代码更少

**缺点**：
- 推理优化可能意外影响 SFT 梯度正确性
- forward() 接口需要添加大量禁用参数
- 维护成本高（两边耦合）

#### 方案 B：独立实现（复制粘贴）
```python
def forward_sft(self, hidden_states, expert_ids, weights, save_for_backward=True):
    # 完全独立的前向实现
    # 手动复制有用的优化逻辑

    buffer = KExpertsSFTBuffer.get_buffer(...)
    buffer.input_cpu.copy_(hidden_states)
    buffer.expert_ids_cpu.copy_(expert_ids)
    buffer.weights_cpu.copy_(weights)

    self.cpu_infer.submit(
        self.moe.forward_sft_task(
            buffer.bsz_tensor.data_ptr(),
            self.num_experts_per_tok,
            buffer.expert_ids_cpu.data_ptr(),
            buffer.weights_cpu.data_ptr(),
            buffer.input_cpu.data_ptr(),
            buffer.output_cpu.data_ptr(),
            save_for_backward,
        )
    )
    self.cpu_infer.sync()

    return buffer.output_cpu.clone()
```

**优点**：
- 推理优化变更不会意外破坏 SFT 梯度
- SFT 可以针对性优化
- 代码独立，更安全

**缺点**：
- 无法自动复用推理优化
- 需要手动同步有用的优化

### 9.3 Python 层推理优化分析

| 优化功能 | 推理目的 | SFT 可用？ | 原因 |
|---------|---------|----------|------|
| 延迟专家执行 | 减少内存峰值 | ❌ | 反向传播需要**所有**专家的激活值 |
| 跨层增量执行 | 隐藏 CPU 计算延迟 | ❌ | 训练需要精确梯度，不能跨层合并 |
| 双缓冲 | GPU-CPU 异步流水线 | ❌ | SFT 必须同步（保存激活值） |
| 异步同步 | 提高并行度 | ❌ | 同上 |
| prefill/decode 分支 | 优化不同场景 | ⚠️ | 部分可用，需评估 |

### 9.4 C++ 层复用分析

| 优化功能 | 复用方式 | 说明 |
|---------|---------|------|
| AMX 内核优化 | 自动继承 | TP_MOE_SFT 继承 TP_MOE |
| 线程池调度 | 自动共享 | 共享 CPUInfer 单例 |
| 量化算法 | 自动继承 | 通过模板参数复用 |
| NUMA 优化 | 自动继承 | WorkerPool 配置复用 |

### 9.5 最终决策：方案 B（独立实现）

**选择理由**：

1. **需求本质不同**
   - 推理追求**低延迟**，可以牺牲一定精度
   - SFT 追求**梯度正确性**，不能有任何精度损失

2. **更安全**
   - 推理的激进优化（延迟专家、异步执行）不会意外破坏 SFT 梯度
   - 每次推理优化更新不需要验证对 SFT 的影响

3. **实际复用有限**
   - 能复用的优化大部分在 C++ 层（已通过继承自动复用）
   - Python 层的优化对 SFT 几乎不适用

4. **维护成本可控**
   - 真正对 SFT 有用的优化很少
   - 手动同步工作量不大

### 9.6 后续优化指南

如果未来有推理优化对 SFT 也有用，应该：

1. **评估安全性**：确认该优化不会影响梯度计算
2. **手动复制**：将优化逻辑复制到 `forward_sft()`
3. **独立测试**：验证 SFT 梯度精度未受影响

```python
# 示例：将 prefill/decode 分支优化复制到 forward_sft

def forward_sft(self, hidden_states, expert_ids, weights, save_for_backward=True):
    qlen = hidden_states.shape[0]

    # 从推理中复制的优化：区分 prefill 和 decode 场景
    if qlen > 32:  # prefill 场景
        # 使用批处理优化
        return self._forward_sft_prefill(hidden_states, expert_ids, weights, save_for_backward)
    else:  # decode 场景
        # 使用低延迟路径
        return self._forward_sft_decode(hidden_states, expert_ids, weights, save_for_backward)
```

---

## 10. 总结

### 10.1 设计原则

1. **统一入口**：KTMoEWrapper 作为唯一工厂接口
2. **内部分离**：推理和 SFT 使用独立的基类和缓冲区
3. **共享基础设施**：CPUInfer 单例、基础验证逻辑共享
4. **C++ 层复用**：通过继承自动复用 AMX 内核优化
5. **Python 层独立**：forward_sft 独立实现，确保梯度正确性

### 10.2 代码复用总结

| 层级 | 复用内容 | 复用方式 |
|------|---------|---------|
| Python 工厂层 | KTMoEWrapper 入口 | 统一接口 |
| Python 基类层 | CPUInfer 单例 | _MoEBase 共享 |
| Python 基类层 | 参数验证 | _MoEBase 共享 |
| Python 实现层 | forward 逻辑 | **独立实现** |
| C++ 基类层 | TP_MOE 优化 | 继承复用 |
| C++ 实现层 | AMX 内核优化 | 继承复用 |
| C++ 实现层 | 量化算法 | 模板参数复用 |
