# SFT + KTWrapper 测试使用

## 1. 环境准备

### 1.1 依赖安装

```bash
# 激活 conda 环境
conda activate ref

# 确认 PyTorch 版本
python -c "import torch; print(torch.__version__)"
```

### 1.2 编译 kt-kernel

```bash
cd /home/lpl/ktransformers-llama/kt-kernel

# 清理旧构建（可选）
rm -rf build

# 创建构建目录
mkdir build && cd build

# 配置 CMake
cmake .. -DCMAKE_BUILD_TYPE=Release

# 编译
make -j$(nproc)
```

### 1.3 环境变量

```bash
# 添加到 PYTHONPATH
export PYTHONPATH=$PYTHONPATH:/home/lpl/ktransformers-llama/kt-kernel/build

# 验证导入
python -c "import kt_kernel_ext; print('OK')"
```

---

## 2. 测试用例列表

### 2.1 推理测试（保持现有）

| 测试文件 | 功能 | 命令 |
|----------|------|------|
| `test_moe_amx.py` | BF16/INT8 推理精度 | `python examples/test_moe_amx.py --mode accuracy` |
| `test_moe_amx.py` | 推理性能 | `python examples/test_moe_amx.py --mode perf` |

### 2.2 SFT 测试（Wrapper 版本-新增）

| 测试文件 | 功能 | 命令 |
|----------|------|------|
| `test_moe_sft_wrapper.py` | SFT 前向精度 | `python examples/test_moe_sft_wrapper.py --mode forward` |
| `test_moe_sft_wrapper.py` | SFT 反向精度 | `python examples/test_moe_sft_wrapper.py --mode backward` |
| `test_moe_sft_wrapper.py` | 训练循环 | `python examples/test_moe_sft_wrapper.py --mode training` |
| `test_moe_sft_wrapper.py` | SFT 性能 | `python examples/test_moe_sft_wrapper.py --mode perf` |
| `test_moe_sft_wrapper.py` | 全部测试 | `python examples/test_moe_sft_wrapper.py --mode all` |

---

## 3. 测试配置

### 3.1 模型参数（DeepSeek-V3）

```python
# 模型配置
expert_num = 256
hidden_size = 7168
intermediate_size = 2048
num_experts_per_tok = 8

# LoRA 配置
lora_rank = 16
lora_alpha = 32.0
lora_scaling = lora_alpha / lora_rank  # = 2.0
```

### 3.2 测试参数

```python
# 精度测试
accuracy_qlen = 128        # 精度测试序列长度
accuracy_iter = 10         # 精度测试迭代次数

# 性能测试
perf_qlen = 128            # 性能测试序列长度
perf_warmup_iter = 5       # 预热迭代次数
perf_test_iter = 20        # 测试迭代次数
```

### 3.3 精度阈值

| 模式 | 前向阈值 | 反向阈值 | 说明 |
|------|----------|----------|------|
| BF16 | 5% | 10% | 高精度模式 |
| INT8 | 10% | 15% | 中等量化 |
| INT4 | 35% | 40% | 低精度量化 |
| INT4_KGroup | 20% | 25% | K-Group 量化 |

---

## 4. 精度验证方法

### 4.1 相对误差计算

```python
def compute_relative_error(amx_output, torch_output):
    """计算相对误差"""
    diff = torch.abs(amx_output - torch_output).mean()
    base = torch.abs(torch_output).mean()
    return (diff / base).item()

# 使用
error = compute_relative_error(amx_output, torch_output)
assert error < threshold, f"Error {error:.2%} exceeds threshold {threshold:.2%}"
```

### 4.2 PyTorch 参考实现

```python
def moe_sft_torch_forward(
    hidden_states,     # [qlen, hidden_size]
    expert_ids,        # [qlen, k]
    weights,           # [qlen, k]
    gate_weight,       # [num_experts, intermediate_size, hidden_size]
    up_weight,         # [num_experts, intermediate_size, hidden_size]
    down_weight,       # [num_experts, hidden_size, intermediate_size]
    gate_lora_a,       # [num_experts, lora_rank, hidden_size]
    gate_lora_b,       # [num_experts, intermediate_size, lora_rank]
    up_lora_a,
    up_lora_b,
    down_lora_a,
    down_lora_b,
    lora_scaling,
):
    """PyTorch 参考实现（用于精度验证）"""
    qlen, hidden_size = hidden_states.shape
    k = expert_ids.shape[1]

    output = torch.zeros_like(hidden_states)

    for i in range(qlen):
        for j in range(k):
            expert_id = expert_ids[i, j].item()
            weight = weights[i, j].item()

            x = hidden_states[i:i+1]  # [1, hidden_size]

            # Gate: base + LoRA
            gate_out = x @ gate_weight[expert_id].T
            gate_lora = (x @ gate_lora_a[expert_id].T) @ gate_lora_b[expert_id].T
            gate_out = gate_out + gate_lora * lora_scaling

            # Up: base + LoRA
            up_out = x @ up_weight[expert_id].T
            up_lora = (x @ up_lora_a[expert_id].T) @ up_lora_b[expert_id].T
            up_out = up_out + up_lora * lora_scaling

            # Activation (SiLU)
            act_out = torch.nn.functional.silu(gate_out) * up_out

            # Down: base + LoRA
            down_out = act_out @ down_weight[expert_id].T
            down_lora = (act_out @ down_lora_a[expert_id].T) @ down_lora_b[expert_id].T
            down_out = down_out + down_lora * lora_scaling

            output[i] += down_out.squeeze(0) * weight

    return output
```

### 4.3 反向传播验证

```python
def verify_backward(wrapper, torch_ref):
    """验证反向传播梯度"""
    # 准备输入
    hidden_states = torch.randn(qlen, hidden_size, dtype=torch.bfloat16)
    hidden_states.requires_grad = True

    # PyTorch 参考前向
    torch_output = torch_ref(hidden_states)
    torch_loss = torch_output.sum()
    torch_loss.backward()
    torch_grad_input = hidden_states.grad.clone()

    # Wrapper 前向 + 反向
    hidden_states_copy = hidden_states.detach().clone()
    wrapper_output = wrapper.forward_sft(hidden_states_copy, expert_ids, weights)
    grad_output = torch.ones_like(wrapper_output)
    wrapper_grad_input, grad_loras = wrapper.backward(grad_output)

    # 比较梯度
    error = compute_relative_error(wrapper_grad_input, torch_grad_input)
    print(f"Backward gradient error: {error:.2%}")
    assert error < backward_threshold
```

---

## 5. 性能测试方法

### 5.1 测试代码

```python
import time
import torch

def benchmark_forward(wrapper, hidden_states, expert_ids, weights, warmup=5, repeat=20):
    """前向性能测试"""
    # 预热
    for _ in range(warmup):
        _ = wrapper.forward_sft(hidden_states, expert_ids, weights, save_for_backward=False)

    # 计时
    torch.cuda.synchronize() if hidden_states.is_cuda else None
    start = time.perf_counter()

    for _ in range(repeat):
        _ = wrapper.forward_sft(hidden_states, expert_ids, weights, save_for_backward=False)

    torch.cuda.synchronize() if hidden_states.is_cuda else None
    end = time.perf_counter()

    avg_time = (end - start) / repeat * 1000  # ms
    return avg_time


def benchmark_backward(wrapper, hidden_states, expert_ids, weights, warmup=5, repeat=20):
    """反向性能测试"""
    grad_output = torch.randn_like(hidden_states)

    # 预热
    for _ in range(warmup):
        _ = wrapper.forward_sft(hidden_states, expert_ids, weights, save_for_backward=True)
        _ = wrapper.backward(grad_output)

    # 计时
    start = time.perf_counter()

    for _ in range(repeat):
        _ = wrapper.forward_sft(hidden_states, expert_ids, weights, save_for_backward=True)
        _, _ = wrapper.backward(grad_output)

    end = time.perf_counter()

    avg_time = (end - start) / repeat * 1000  # ms
    return avg_time
```

### 5.2 性能指标

| 指标 | 计算方式 | 单位 |
|------|----------|------|
| 前向延迟 | 单次前向时间 | ms |
| 反向延迟 | 单次反向时间 | ms |
| 前向+反向延迟 | 前向 + 反向总时间 | ms |
| 吞吐量 | qlen / 延迟 | tokens/s |

### 5.3 性能对比

```python
def compare_with_native(wrapper, native_moe, hidden_states, expert_ids, weights):
    """与直接 C++ 调用对比"""
    # Wrapper 性能
    wrapper_time = benchmark_forward(wrapper, hidden_states, expert_ids, weights)

    # Native C++ 性能（直接调用）
    # ... native 测试代码 ...

    overhead = (wrapper_time - native_time) / native_time * 100
    print(f"Wrapper overhead: {overhead:.1f}%")
    assert overhead < 5, "Wrapper overhead exceeds 5%"
```

---

## 6. 运行示例

### 6.1 快速验证

```bash
conda activate ref
cd /home/lpl/ktransformers-llama/kt-kernel

# 推理模式（验证 Wrapper 不破坏现有功能）
python examples/test_moe_amx.py --mode accuracy

# SFT 模式（新增功能验证）
python examples/test_moe_sft_wrapper.py --mode all
```

### 6.2 单项测试

```bash
# 仅测试前向精度
python examples/test_moe_sft_wrapper.py --mode forward --method AMXBF16_SFT

# 仅测试反向精度
python examples/test_moe_sft_wrapper.py --mode backward --method AMXBF16_SFT

# 仅测试性能
python examples/test_moe_sft_wrapper.py --mode perf --method AMXBF16_SFT
```

### 6.3 批量测试所有量化方法

```bash
# 测试所有 SFT 量化方法
for method in AMXBF16_SFT AMXINT8_SFT AMXINT4_SFT AMXINT4_KGroup_SFT; do
    echo "Testing $method..."
    python examples/test_moe_sft_wrapper.py --mode all --method $method
done
```

### 6.4 性能对比测试

```bash
# 对比 Wrapper 与直接 C++ 调用的性能差异
python examples/test_moe_sft_wrapper.py --mode perf --compare-native
```

---

## 7. 测试文件模板

### 7.1 test_moe_sft_wrapper.py 结构

```python
#!/usr/bin/env python
"""SFT Wrapper 测试脚本"""

import argparse
import torch
import sys
sys.path.insert(0, "/home/lpl/ktransformers-llama/kt-kernel/build")

from python.experts import KTMoEWrapper


def test_forward_accuracy(method: str):
    """测试前向精度"""
    print(f"\n=== Testing forward accuracy ({method}) ===")
    # ... 测试代码 ...


def test_backward_accuracy(method: str):
    """测试反向精度"""
    print(f"\n=== Testing backward accuracy ({method}) ===")
    # ... 测试代码 ...


def test_training_loop(method: str):
    """测试训练循环"""
    print(f"\n=== Testing training loop ({method}) ===")
    # ... 测试代码 ...


def test_performance(method: str, compare_native: bool = False):
    """测试性能"""
    print(f"\n=== Testing performance ({method}) ===")
    # ... 测试代码 ...


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["forward", "backward", "training", "perf", "all"],
                        default="all")
    parser.add_argument("--method", default="AMXBF16_SFT")
    parser.add_argument("--compare-native", action="store_true")
    args = parser.parse_args()

    if args.mode in ("forward", "all"):
        test_forward_accuracy(args.method)

    if args.mode in ("backward", "all"):
        test_backward_accuracy(args.method)

    if args.mode in ("training", "all"):
        test_training_loop(args.method)

    if args.mode in ("perf", "all"):
        test_performance(args.method, args.compare_native)

    print("\n=== All tests passed! ===")


if __name__ == "__main__":
    main()
```

---

## 8. 常见问题排查

### 8.1 导入错误

**问题**: `ImportError: No module named 'kt_kernel_ext'`

**解决**:
```bash
# 检查 PYTHONPATH
echo $PYTHONPATH

# 添加正确路径
export PYTHONPATH=$PYTHONPATH:/home/lpl/ktransformers-llama/kt-kernel/build

# 或在代码中添加
import sys
sys.path.insert(0, "/home/lpl/ktransformers-llama/kt-kernel/build")
```

### 8.2 精度超标

**问题**: `AssertionError: diff > threshold`

**排查步骤**:
1. 检查权重是否正确加载
   ```python
   print(f"Weights loaded: {wrapper._weights_loaded}")
   ```

2. 检查 LoRA 缩放因子
   ```python
   print(f"LoRA scaling: {wrapper.lora_scaling}")
   # 应该等于 lora_alpha / lora_rank
   ```

3. 对比中间值找出发散点
   ```python
   # 分别检查 gate、up、down 的输出
   ```

4. 检查数据类型
   ```python
   print(f"Input dtype: {hidden_states.dtype}")
   print(f"Weight dtype: {gate_lora_a.dtype}")
   ```

### 8.3 内存不足

**问题**: `RuntimeError: CUDA out of memory`

**解决**:
```python
# 减小序列长度
qlen = 64  # 从 128 减小

# 或使用 CPU 测试
hidden_states = torch.randn(qlen, hidden_size, dtype=torch.bfloat16)  # 不加 .cuda()
```

### 8.4 缓存溢出

**问题**: `RuntimeError: Forward cache full`

**解决**:
```python
# 方法 1: 增大 max_cache_depth
wrapper = KTMoEWrapper(..., max_cache_depth=2)

# 方法 2: 及时调用 backward 释放缓存
output = wrapper.forward_sft(..., save_for_backward=True)
grad_input, grad_loras = wrapper.backward(grad_output)  # 释放缓存
```

### 8.5 TP 模式问题

**问题**: TP 模式输出不正确

**排查步骤**:
1. 检查 threadpool_count 是否正确
   ```python
   print(f"TP size: {wrapper.threadpool_count}")
   ```

2. 检查 intermediate_size 是否可被 TP 数量整除
   ```python
   assert intermediate_size % threadpool_count == 0
   ```

3. 检查权重分片是否正确
   ```python
   # 每个 TP 分片的 intermediate_size
   tp_intermediate = intermediate_size // threadpool_count
   ```

---

## 9. CI/CD 集成

### 9.1 GitHub Actions 配置

```yaml
# .github/workflows/test_sft_wrapper.yml
name: SFT Wrapper Tests

on:
  push:
    paths:
      - 'python/experts*.py'
      - 'python/utils/amx_sft.py'
      - 'operators/moe-sft-tp.hpp'
      - 'operators/amx/sft_moe.hpp'

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: '3.10'

      - name: Install dependencies
        run: |
          pip install torch numpy

      - name: Build kt-kernel
        run: |
          mkdir build && cd build
          cmake .. -DCMAKE_BUILD_TYPE=Release
          make -j$(nproc)

      - name: Run tests
        run: |
          export PYTHONPATH=$PYTHONPATH:$(pwd)/build
          python examples/test_moe_sft_wrapper.py --mode all
```

### 9.2 本地测试脚本

```bash
#!/bin/bash
# scripts/test_sft_wrapper.sh

set -e

echo "Activating conda environment..."
source /path/to/anaconda/bin/activate ref

echo "Building kt-kernel..."
cd /home/lpl/ktransformers-llama/kt-kernel
mkdir -p build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(nproc)

echo "Running tests..."
export PYTHONPATH=$PYTHONPATH:$(pwd)
cd ..

# 运行所有测试
python examples/test_moe_sft_wrapper.py --mode all

echo "All tests passed!"
```
