# SFT + KTWrapper 功能具体实现

## 1. 文件变更清单

| 文件 | 操作 | 说明 |
|------|------|------|
| `python/experts.py` | 修改 | 添加 mode 参数和 SFT 分支 |
| `python/experts_base.py` | 修改 | 提取 `_MoEBase` 共享基类 |
| `python/experts_sft.py` | 新增 | `BaseSFTMoEWrapper` 和 `KExpertsSFTBuffer` |
| `python/utils/amx_sft.py` | 新增 | `AMXSFTMoEWrapper` 实现 |
| `examples/test_moe_sft_wrapper.py` | 新增 | Wrapper 版 SFT 测试 |

---

## 2. _MoEBase 共享基类实现

### 2.1 从 BaseMoEWrapper 提取

**修改文件**: `python/experts_base.py`

```python
# 在文件开头添加
import threading
from typing import ClassVar, Optional

class _MoEBase:
    """推理和 SFT 共享的基类

    职责:
    1. 管理 CPUInfer 单例
    2. 提供 WorkerPoolConfig 构建逻辑
    3. 基础参数验证
    """

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

    @classmethod
    def _get_cpu_infer(
        cls,
        cpuinfer_threads: int,
        threadpool_count: int
    ) -> "kt_kernel_ext.CPUInfer":
        """获取或创建 CPUInfer 单例

        Args:
            cpuinfer_threads: 每个 NUMA 子池的线程数
            threadpool_count: NUMA 子池数量（TP 数量）

        Returns:
            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,
        moe_intermediate_size: int,
        num_experts_per_tok: int,
    ) -> None:
        """验证基础配置参数

        Raises:
            ValueError: 参数无效时抛出
        """
        if num_experts <= 0:
            raise ValueError(f"num_experts must be positive, got {num_experts}")
        if hidden_size <= 0:
            raise ValueError(f"hidden_size must be positive, got {hidden_size}")
        if moe_intermediate_size <= 0:
            raise ValueError(f"moe_intermediate_size must be positive, got {moe_intermediate_size}")
        if num_experts_per_tok <= 0:
            raise ValueError(f"num_experts_per_tok must be positive, got {num_experts_per_tok}")
        if num_experts_per_tok > num_experts:
            raise ValueError(
                f"num_experts_per_tok ({num_experts_per_tok}) cannot exceed "
                f"num_experts ({num_experts})"
            )
```

### 2.2 修改 BaseMoEWrapper 继承关系

```python
# 修改 BaseMoEWrapper 的定义
class BaseMoEWrapper(_MoEBase, ABC):
    """推理 MoE 的基类

    继承自 _MoEBase 以共享 CPUInfer 单例管理
    """

    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)

        # 验证基础配置
        self._validate_base_config(
            num_experts=num_experts,
            hidden_size=hidden_size,
            moe_intermediate_size=moe_intermediate_size,
            num_experts_per_tok=num_experts_per_tok,
        )

        # ... 其余现有初始化逻辑保持不变 ...

    # 添加 SFT 方法的错误提示
    def forward_sft(self, *args, **kwargs):
        raise RuntimeError(
            "forward_sft() is not available in inference mode. "
            "Use forward() instead, or create wrapper with mode='sft'."
        )

    def backward(self, *args, **kwargs):
        raise RuntimeError(
            "backward() is not available in inference mode. "
            "Create wrapper with mode='sft' to use SFT features."
        )

    def init_lora_weights(self, *args, **kwargs):
        raise RuntimeError(
            "init_lora_weights() is not available in inference mode. "
            "Create wrapper with mode='sft' to use SFT features."
        )

    def update_lora_weights(self, *args, **kwargs):
        raise RuntimeError(
            "update_lora_weights() is not available in inference mode. "
            "Create wrapper with mode='sft' to use SFT features."
        )
```

---

## 3. BaseSFTMoEWrapper 实现

### 3.1 新建文件

**新建文件**: `python/experts_sft.py`

```python
"""SFT MoE Wrapper 基类和缓冲区管理

提供 SFT（Supervised Fine-Tuning）模式的 MoE Wrapper 基类，
支持 LoRA 微调的前向传播、反向传播和权重更新。
"""

import threading
from abc import ABC, abstractmethod
from typing import ClassVar, Dict, Optional, Tuple

import torch

from .experts_base import _MoEBase

try:
    import kt_kernel_ext
except ImportError:
    kt_kernel_ext = None


class KExpertsSFTBuffer:
    """SFT 模式的 CPU 缓冲区管理

    与推理模式的 KExpertsCPUBuffer 不同：
    - 不需要双缓冲（SFT 是同步执行）
    - 需要额外的梯度缓冲区
    - 需要 LoRA 梯度缓冲区

    Attributes:
        capture_buffers: 缓冲区缓存字典
    """

    _buffer_lock: ClassVar[threading.Lock] = threading.Lock()
    capture_buffers: ClassVar[Dict[tuple, "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,
    ):
        """初始化 SFT 缓冲区

        Args:
            qlen: 序列长度
            hidden_size: 隐藏层维度
            moe_intermediate_size: MoE 中间层维度
            num_experts: 专家总数
            num_experts_per_tok: 每 token 激活的专家数
            lora_rank: LoRA 低秩矩阵的秩
            dtype: 数据类型
        """
        self.qlen = qlen
        self.hidden_size = hidden_size
        self.moe_intermediate_size = moe_intermediate_size
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.lora_rank = lora_rank
        self.dtype = dtype

        # ========== 前向缓冲 ==========
        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 个）==========
        # Gate LoRA 梯度
        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"
        )

        # Up LoRA 梯度
        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"
        )

        # Down LoRA 梯度
        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 缓冲区（线程安全）

        使用参数组合作为缓存键，复用已创建的缓冲区。

        Args:
            qlen: 序列长度
            hidden_size: 隐藏层维度
            moe_intermediate_size: MoE 中间层维度
            num_experts: 专家总数
            num_experts_per_tok: 每 token 激活的专家数
            lora_rank: LoRA 秩
            dtype: 数据类型

        Returns:
            KExpertsSFTBuffer 实例
        """
        key = (
            qlen, hidden_size, moe_intermediate_size,
            num_experts, num_experts_per_tok, lora_rank, dtype
        )

        with cls._buffer_lock:
            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]

    @classmethod
    def clear_cache(cls) -> None:
        """清除所有缓存的缓冲区"""
        with cls._buffer_lock:
            cls.capture_buffers.clear()

    def get_lora_grads(self) -> Dict[str, torch.Tensor]:
        """获取所有 LoRA 梯度的字典

        Returns:
            包含 6 个 LoRA 梯度张量的字典
        """
        return {
            "grad_gate_lora_a": self.grad_gate_lora_a,
            "grad_gate_lora_b": self.grad_gate_lora_b,
            "grad_up_lora_a": self.grad_up_lora_a,
            "grad_up_lora_b": self.grad_up_lora_b,
            "grad_down_lora_a": self.grad_down_lora_a,
            "grad_down_lora_b": self.grad_down_lora_b,
        }


class BaseSFTMoEWrapper(_MoEBase, ABC):
    """SFT MoE 的基类

    提供 LoRA 微调所需的前向传播、反向传播和权重更新功能。
    与推理基类 BaseMoEWrapper 的主要区别：

    1. 支持 forward_sft() 带梯度缓存的前向传播
    2. 支持 backward() 反向传播计算 LoRA 梯度
    3. 支持 update_lora_weights() 同步 LoRA 权重到 C++ 后端
    4. 使用 KExpertsSFTBuffer 而非 KExpertsCPUBuffer

    Attributes:
        lora_rank: LoRA 低秩矩阵的秩
        lora_alpha: LoRA 缩放因子
        lora_scaling: 实际缩放值 (lora_alpha / lora_rank)
        max_cache_depth: 前向缓存深度
    """

    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,
    ):
        """初始化 SFT MoE Wrapper

        Args:
            layer_idx: 层索引
            num_experts: 专家总数
            num_experts_per_tok: 每 token 激活的专家数
            hidden_size: 隐藏层维度
            moe_intermediate_size: MoE 中间层维度
            num_gpu_experts: GPU 上的专家数（SFT 通常为 0）
            cpuinfer_threads: CPU 推理线程数
            threadpool_count: NUMA 子池数量
            weight_path: 权重路径
            chunked_prefill_size: 分块预填充大小
            lora_rank: LoRA 秩
            lora_alpha: LoRA 缩放因子
            max_cache_depth: 前向缓存深度
        """
        # 获取共享的 CPUInfer 实例
        self.cpu_infer = self._get_cpu_infer(cpuinfer_threads, threadpool_count)

        # 验证基础配置
        self._validate_base_config(
            num_experts=num_experts,
            hidden_size=hidden_size,
            moe_intermediate_size=moe_intermediate_size,
            num_experts_per_tok=num_experts_per_tok,
        )

        # 验证 SFT 特有参数
        self._validate_sft_config(lora_rank, lora_alpha, max_cache_depth)

        # 保存配置
        self.layer_idx = layer_idx
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.hidden_size = hidden_size
        self.moe_intermediate_size = moe_intermediate_size
        self.num_gpu_experts = num_gpu_experts
        self.weight_path = weight_path
        self.chunked_prefill_size = chunked_prefill_size
        self.threadpool_count = 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
        self._cache_depth: int = 0

    @staticmethod
    def _validate_sft_config(
        lora_rank: int,
        lora_alpha: float,
        max_cache_depth: int
    ) -> None:
        """验证 SFT 特有参数

        Raises:
            ValueError: 参数无效时抛出
        """
        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}")

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

        Args:
            physical_to_logical_map: 物理到逻辑专家的映射
        """
        ...

    @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 权重

        Args:
            gate_lora_a: Gate LoRA A 矩阵 [num_experts, lora_rank, hidden_size]
            gate_lora_b: Gate LoRA B 矩阵 [num_experts, intermediate_size, lora_rank]
            up_lora_a: Up LoRA A 矩阵 [num_experts, lora_rank, hidden_size]
            up_lora_b: Up LoRA B 矩阵 [num_experts, intermediate_size, lora_rank]
            down_lora_a: Down LoRA A 矩阵 [num_experts, lora_rank, intermediate_size]
            down_lora_b: Down LoRA B 矩阵 [num_experts, hidden_size, lora_rank]
        """
        ...

    @abstractmethod
    def forward_sft(
        self,
        hidden_states: torch.Tensor,
        expert_ids: torch.Tensor,
        weights: torch.Tensor,
        save_for_backward: bool = True,
    ) -> torch.Tensor:
        """SFT 前向传播

        Args:
            hidden_states: 输入隐藏状态 [qlen, hidden_size]
            expert_ids: 专家 ID [qlen, num_experts_per_tok]
            weights: 专家权重 [qlen, num_experts_per_tok]
            save_for_backward: 是否保存激活值用于反向传播

        Returns:
            输出隐藏状态 [qlen, hidden_size]
        """
        ...

    @abstractmethod
    def backward(
        self,
        grad_output: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """反向传播

        Args:
            grad_output: 输出梯度 [qlen, hidden_size]

        Returns:
            grad_input: 输入梯度 [qlen, hidden_size]
            grad_loras: LoRA 梯度字典
        """
        ...

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

        在使用外部优化器更新 LoRA 权重后调用此方法，
        将更新后的权重同步到 C++ 后端。
        """
        ...

    # 推理方法的错误提示
    def forward(self, *args, **kwargs):
        raise RuntimeError(
            "forward() is not available in SFT mode. "
            "Use forward_sft() instead."
        )

    def submit_forward(self, *args, **kwargs):
        raise RuntimeError(
            "submit_forward() is not available in SFT mode. "
            "SFT mode uses synchronous forward_sft()."
        )

    def sync_forward(self, *args, **kwargs):
        raise RuntimeError(
            "sync_forward() is not available in SFT mode."
        )

    def select_deferred_experts(self, *args, **kwargs):
        raise RuntimeError(
            "select_deferred_experts() is not available in SFT mode."
        )
```

---

## 4. AMXSFTMoEWrapper 实现

### 4.1 新建文件

**新建文件**: `python/utils/amx_sft.py`

```python
"""AMX SFT MoE Wrapper 实现

基于 AMX 指令集的 SFT MoE Wrapper，支持 BF16、INT8、INT4 和 INT4_KGroup 量化。
"""

from typing import Dict, Optional, Tuple

import torch

from ..experts_sft import BaseSFTMoEWrapper, KExpertsSFTBuffer

try:
    import kt_kernel_ext
except ImportError:
    kt_kernel_ext = None


class AMXSFTMoEWrapper(BaseSFTMoEWrapper):
    """AMX SFT MoE Wrapper

    使用 Intel AMX 指令集加速的 SFT MoE 实现。
    支持的量化方法：
    - AMXBF16_SFT: BF16 精度
    - AMXINT8_SFT: INT8 量化
    - AMXINT4_SFT: INT4 量化
    - AMXINT4_KGroup_SFT: INT4 K-Group 量化（AWQ/K2）

    Attributes:
        method: 量化方法
        moe: C++ MoE 实例
    """

    SUPPORTED_METHODS = {
        "AMXBF16_SFT",
        "AMXINT8_SFT",
        "AMXINT4_SFT",
        "AMXINT4_KGroup_SFT",
    }

    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,
        method: str = "AMXBF16_SFT",
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        max_cache_depth: int = 1,
    ):
        """初始化 AMX SFT MoE Wrapper

        Args:
            method: 量化方法，必须是 SUPPORTED_METHODS 之一
            其他参数见 BaseSFTMoEWrapper
        """
        # 验证 method
        if method not in self.SUPPORTED_METHODS:
            raise ValueError(
                f"Unsupported method: {method}. "
                f"Supported methods: {self.SUPPORTED_METHODS}"
            )

        # 调用父类初始化
        super().__init__(
            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,
            lora_rank=lora_rank,
            lora_alpha=lora_alpha,
            max_cache_depth=max_cache_depth,
        )

        self.method = method

        # 创建 MOESFTConfig
        config = kt_kernel_ext.moe.MOESFTConfig()
        config.expert_num = num_experts
        config.hidden_size = hidden_size
        config.intermediate_size = moe_intermediate_size
        config.experts_per_token = num_experts_per_tok
        config.weight_path = weight_path
        config.layer_idx = layer_idx
        config.lora_rank = lora_rank
        config.lora_alpha = lora_alpha
        config.max_cache_depth = max_cache_depth
        config.tp_size = threadpool_count

        # 根据 method 创建对应的 MoE 实例
        self.moe = self._create_moe_instance(config, method)

        # 预热
        self._warm_up()

    def _create_moe_instance(
        self,
        config: "kt_kernel_ext.moe.MOESFTConfig",
        method: str
    ):
        """根据 method 创建 MoE 实例"""
        if method == "AMXBF16_SFT":
            return kt_kernel_ext.moe.AMXBF16_SFT_MOE(config)
        elif method == "AMXINT8_SFT":
            return kt_kernel_ext.moe.AMXInt8_SFT_MOE(config)
        elif method == "AMXINT4_SFT":
            return kt_kernel_ext.moe.AMXInt4_SFT_MOE(config)
        elif method == "AMXINT4_KGroup_SFT":
            return kt_kernel_ext.moe.AMXInt4KGroup_SFT_MOE(config)
        else:
            raise ValueError(f"Unknown method: {method}")

    def _warm_up(self) -> None:
        """预热 MoE 实例"""
        self.cpu_infer.submit(self.moe.warm_up_task())
        self.cpu_infer.sync()

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

        Args:
            physical_to_logical_map: 物理到逻辑专家的映射 [num_experts]
        """
        if physical_to_logical_map.dtype != torch.int64:
            physical_to_logical_map = physical_to_logical_map.to(torch.int64)

        # 确保在 CPU 上且连续
        if physical_to_logical_map.device.type != "cpu":
            physical_to_logical_map = physical_to_logical_map.cpu()
        physical_to_logical_map = physical_to_logical_map.contiguous()

        # 提交加载任务
        self.cpu_infer.submit(
            self.moe.load_weights_task(physical_to_logical_map.data_ptr())
        )
        self.cpu_infer.sync()

        self._weights_loaded = True

    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 权重

        所有权重必须是 BF16 格式，形状：
        - gate_lora_a: [num_experts, lora_rank, hidden_size]
        - gate_lora_b: [num_experts, intermediate_size, lora_rank]
        - up_lora_a: [num_experts, lora_rank, hidden_size]
        - up_lora_b: [num_experts, intermediate_size, lora_rank]
        - down_lora_a: [num_experts, lora_rank, intermediate_size]
        - down_lora_b: [num_experts, hidden_size, lora_rank]
        """
        # 验证并保存权重引用
        self.gate_lora_a = self._validate_lora_weight(
            gate_lora_a, "gate_lora_a",
            (self.num_experts, self.lora_rank, self.hidden_size)
        )
        self.gate_lora_b = self._validate_lora_weight(
            gate_lora_b, "gate_lora_b",
            (self.num_experts, self.moe_intermediate_size, self.lora_rank)
        )
        self.up_lora_a = self._validate_lora_weight(
            up_lora_a, "up_lora_a",
            (self.num_experts, self.lora_rank, self.hidden_size)
        )
        self.up_lora_b = self._validate_lora_weight(
            up_lora_b, "up_lora_b",
            (self.num_experts, self.moe_intermediate_size, self.lora_rank)
        )
        self.down_lora_a = self._validate_lora_weight(
            down_lora_a, "down_lora_a",
            (self.num_experts, self.lora_rank, self.moe_intermediate_size)
        )
        self.down_lora_b = self._validate_lora_weight(
            down_lora_b, "down_lora_b",
            (self.num_experts, self.hidden_size, self.lora_rank)
        )

        # 同步到 C++ 后端
        self._sync_lora_weights_to_cpp()

        self._lora_initialized = True

    def _validate_lora_weight(
        self,
        weight: torch.Tensor,
        name: str,
        expected_shape: tuple
    ) -> torch.Tensor:
        """验证 LoRA 权重格式"""
        # 检查形状
        if weight.shape != expected_shape:
            raise ValueError(
                f"{name} shape mismatch: expected {expected_shape}, "
                f"got {weight.shape}"
            )

        # 确保 BF16、CPU、连续
        if weight.dtype != torch.bfloat16:
            weight = weight.to(torch.bfloat16)
        if weight.device.type != "cpu":
            weight = weight.cpu()
        if not weight.is_contiguous():
            weight = weight.contiguous()

        return weight

    def _sync_lora_weights_to_cpp(self) -> None:
        """同步 LoRA 权重指针到 C++ 后端"""
        self.cpu_infer.submit(
            self.moe.update_lora_weights_task(
                self.gate_lora_a.data_ptr(),
                self.gate_lora_b.data_ptr(),
                self.up_lora_a.data_ptr(),
                self.up_lora_b.data_ptr(),
                self.down_lora_a.data_ptr(),
                self.down_lora_b.data_ptr(),
            )
        )
        self.cpu_infer.sync()

    def forward_sft(
        self,
        hidden_states: torch.Tensor,
        expert_ids: torch.Tensor,
        weights: torch.Tensor,
        save_for_backward: bool = True,
    ) -> torch.Tensor:
        """SFT 前向传播

        Args:
            hidden_states: 输入 [qlen, hidden_size]，可以是 GPU 或 CPU 张量
            expert_ids: 专家 ID [qlen, num_experts_per_tok]
            weights: 专家权重 [qlen, num_experts_per_tok]
            save_for_backward: 是否保存用于反向传播

        Returns:
            输出 [qlen, hidden_size]，与输入同设备
        """
        # 状态检查
        if not self._weights_loaded:
            raise RuntimeError("Weights not loaded. Call load_weights() first.")
        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."
            )

        # 检查缓存深度
        if save_for_backward and self._cache_depth >= self.max_cache_depth:
            raise RuntimeError(
                f"Forward cache full (depth={self._cache_depth}, "
                f"max={self.max_cache_depth}). "
                "Call backward() to free cache slots."
            )

        qlen = hidden_states.shape[0]
        input_device = hidden_states.device

        # 获取 SFT 缓冲区
        buffer = KExpertsSFTBuffer.get_buffer(
            qlen=qlen,
            hidden_size=self.hidden_size,
            moe_intermediate_size=self.moe_intermediate_size,
            num_experts=self.num_experts,
            num_experts_per_tok=self.num_experts_per_tok,
            lora_rank=self.lora_rank,
            dtype=hidden_states.dtype,
        )

        # 复制输入到 CPU
        buffer.input_cpu[:qlen].copy_(hidden_states, non_blocking=True)
        buffer.expert_ids_cpu[:qlen].copy_(expert_ids, non_blocking=True)
        buffer.weights_cpu[:qlen].copy_(weights, non_blocking=True)

        # 同步 GPU->CPU 传输
        if input_device.type == "cuda":
            torch.cuda.current_stream().synchronize()

        # 提交前向任务
        self.cpu_infer.submit(
            self.moe.forward_sft_task(
                qlen,
                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()

        # 更新缓存深度
        if save_for_backward:
            self._cache_depth += 1

        # 返回输出（复制回原设备）
        if input_device.type == "cuda":
            output = torch.empty_like(hidden_states)
            output.copy_(buffer.output_cpu[:qlen], non_blocking=True)
            return output
        else:
            return buffer.output_cpu[:qlen].clone()

    def backward(
        self,
        grad_output: torch.Tensor,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """反向传播

        Args:
            grad_output: 输出梯度 [qlen, hidden_size]

        Returns:
            grad_input: 输入梯度 [qlen, hidden_size]
            grad_loras: LoRA 梯度字典
        """
        # 状态检查
        if self._cache_depth <= 0:
            raise RuntimeError(
                "No forward cache available. "
                "Call forward_sft() with save_for_backward=True first."
            )

        qlen = grad_output.shape[0]
        input_device = grad_output.device

        # 获取缓冲区
        buffer = KExpertsSFTBuffer.get_buffer(
            qlen=qlen,
            hidden_size=self.hidden_size,
            moe_intermediate_size=self.moe_intermediate_size,
            num_experts=self.num_experts,
            num_experts_per_tok=self.num_experts_per_tok,
            lora_rank=self.lora_rank,
            dtype=grad_output.dtype,
        )

        # 复制梯度到 CPU
        buffer.grad_output_cpu[:qlen].copy_(grad_output, non_blocking=True)

        # 同步 GPU->CPU 传输
        if input_device.type == "cuda":
            torch.cuda.current_stream().synchronize()

        # 提交反向任务
        self.cpu_infer.submit(
            self.moe.backward_task(
                qlen,
                self.num_experts_per_tok,
                buffer.grad_output_cpu.data_ptr(),
                buffer.grad_input_cpu.data_ptr(),
                buffer.grad_gate_lora_a.data_ptr(),
                buffer.grad_gate_lora_b.data_ptr(),
                buffer.grad_up_lora_a.data_ptr(),
                buffer.grad_up_lora_b.data_ptr(),
                buffer.grad_down_lora_a.data_ptr(),
                buffer.grad_down_lora_b.data_ptr(),
            )
        )
        self.cpu_infer.sync()

        # 更新缓存深度
        self._cache_depth -= 1

        # 准备返回值
        if input_device.type == "cuda":
            grad_input = torch.empty_like(grad_output)
            grad_input.copy_(buffer.grad_input_cpu[:qlen], non_blocking=True)
        else:
            grad_input = buffer.grad_input_cpu[:qlen].clone()

        grad_loras = {
            "grad_gate_lora_a": buffer.grad_gate_lora_a.clone(),
            "grad_gate_lora_b": buffer.grad_gate_lora_b.clone(),
            "grad_up_lora_a": buffer.grad_up_lora_a.clone(),
            "grad_up_lora_b": buffer.grad_up_lora_b.clone(),
            "grad_down_lora_a": buffer.grad_down_lora_a.clone(),
            "grad_down_lora_b": buffer.grad_down_lora_b.clone(),
        }

        return grad_input, grad_loras

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

        在使用外部优化器更新 LoRA 权重后调用。
        """
        if not self._lora_initialized:
            raise RuntimeError(
                "LoRA weights not initialized. "
                "Call init_lora_weights() first."
            )

        self._sync_lora_weights_to_cpp()
```

---

## 5. KTMoEWrapper 工厂类修改

### 5.1 修改文件

**修改文件**: `python/experts.py`

```python
"""KTMoEWrapper 工厂类

提供统一的 MoE Wrapper 创建入口，支持推理和 SFT 两种模式。
"""

from typing import Optional

try:
    import kt_kernel_ext
except ImportError:
    kt_kernel_ext = None


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

    根据 mode 参数创建推理或 SFT 模式的 Wrapper。

    Usage:
        # 推理模式（默认）
        wrapper = KTMoEWrapper(..., method="AMXINT4", mode="inference")

        # SFT 模式
        wrapper = KTMoEWrapper(..., method="AMXBF16_SFT", mode="sft",
                              lora_rank=16, lora_alpha=32.0)
    """

    # 推理模式支持的 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,
    ):
        """创建 MoE Wrapper 实例

        Args:
            layer_idx: 层索引
            num_experts: 专家总数
            num_experts_per_tok: 每 token 激活的专家数
            hidden_size: 隐藏层维度
            moe_intermediate_size: MoE 中间层维度
            num_gpu_experts: GPU 上的专家数
            cpuinfer_threads: CPU 推理线程数
            threadpool_count: NUMA 子池数量
            weight_path: 权重路径
            chunked_prefill_size: 分块预填充大小
            cpu_save: 是否保存到 CPU 内存（推理模式）
            max_deferred_experts_per_token: 延迟专家数（推理模式）
            method: 后端方法
            mode: 模式 ("inference" 或 "sft")
            lora_rank: LoRA 秩（SFT 模式）
            lora_alpha: LoRA 缩放因子（SFT 模式）
            max_cache_depth: 前向缓存深度（SFT 模式）

        Returns:
            BaseMoEWrapper 或 BaseSFTMoEWrapper 的子类实例

        Raises:
            ValueError: mode 或 method 无效时抛出
        """
        # 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}' is not supported in inference mode. "
                f"Supported methods: {sorted(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: {sorted(cls.SFT_METHODS)}"
            )

        # 3. 准备基础参数
        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,
        }

        # 4. 根据 mode 创建对应的 Wrapper
        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)
        else:
            raise ValueError(f"Unknown inference method: {method}")

    @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)
        else:
            raise ValueError(f"Unknown SFT method: {method}")
```

---

## 6. 实现检查清单

### 6.1 Phase 1: 基础设施

- [ ] 从 `BaseMoEWrapper` 提取 `_MoEBase` 共享基类
  - [ ] 移动 `_cpu_infer_instance` 类变量
  - [ ] 移动 `_get_cpu_infer()` 类方法
  - [ ] 添加线程锁保护
  - [ ] 添加 `_validate_base_config()` 方法

- [ ] 修改 `BaseMoEWrapper` 继承 `_MoEBase`
  - [ ] 更新初始化逻辑
  - [ ] 添加 SFT 方法的错误提示

### 6.2 Phase 2: SFT 实现

- [ ] 创建 `experts_sft.py`
  - [ ] 实现 `KExpertsSFTBuffer` 类
  - [ ] 实现 `BaseSFTMoEWrapper` 抽象基类

- [ ] 创建 `utils/amx_sft.py`
  - [ ] 实现 `AMXSFTMoEWrapper` 类
  - [ ] 实现 `load_weights()` 方法
  - [ ] 实现 `init_lora_weights()` 方法
  - [ ] 实现 `forward_sft()` 方法
  - [ ] 实现 `backward()` 方法
  - [ ] 实现 `update_lora_weights()` 方法

### 6.3 Phase 3: 工厂类

- [ ] 修改 `experts.py` 中的 `KTMoEWrapper`
  - [ ] 添加 `mode` 参数
  - [ ] 添加 SFT 特有参数
  - [ ] 添加 `SFT_METHODS` 集合
  - [ ] 实现 `_create_sft_wrapper()` 方法
  - [ ] 添加 mode/method 验证逻辑

### 6.4 Phase 4: 测试

- [ ] 创建 `test_moe_sft_wrapper.py`
  - [ ] 前向精度测试
  - [ ] 反向精度测试
  - [ ] 训练循环测试
  - [ ] 性能测试
  - [ ] 与直接 C++ 调用的对比测试
