test_ane_vs_gpu.py

#!/usr/bin/env python3
"""对比测试 ANE Prefill vs GPU Prefill"""

import sys
sys.path.insert(0, '/Users/xiphis/pd-engine')

import time
import mlx.core as mx
from mlx_lm import load, generate
from mlx_lm.models.cache import make_prompt_cache
from transformers import AutoTokenizer
import coremltools as ct
import numpy as np

print("="*60)
print("📊 ANE Prefill vs GPU Prefill 详细对比测试")
print("="*60)

# 加载模型
print("\n📦 加载模型...")
model, _ = load('tinyllama-mlx')
tokenizer = AutoTokenizer.from_pretrained('tinyllama/tinyllama-1.1b-chat-v1.0')
ane_model = ct.models.MLModel('model_ane_full.mlpackage')
print("   ✅ 模型加载完成")

# 配置
NUM_LAYERS = 22
NUM_KV_HEADS = 4
HEAD_DIM = 64
ROPE_THETA = 10000.0
ANE_FIXED_LEN = 128

def apply_rope_mlx(x, positions, theta=10000.0):
    """Apply RoPE"""
    batch, num_heads, seq_len, head_dim = x.shape
    freqs = 1.0 / (theta ** (mx.arange(0, head_dim, 2)[:head_dim // 2] / head_dim))
    angles = positions[:, :, None] * freqs[None, None, :]
    cos = mx.cos(angles)[:, None, :, :]
    sin = mx.sin(angles)[:, None, :, :]
    x1 = x[..., :head_dim // 2]
    x2 = x[..., head_dim // 2:]
    return mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)

def ane_prefill(tokens):
    """ANE Prefill"""
    seq_len = len(tokens)
    
    # Pad to fixed length
    # 如果正好是 128 的倍数，不需要 padding
    if seq_len % ANE_FIXED_LEN == 0:
        padded_len = seq_len
    else:
        padded_len = ((seq_len // ANE_FIXED_LEN) + 1) * ANE_FIXED_LEN
    
    input_ids_np = np.pad(tokens, (0, padded_len - seq_len), mode='constant').reshape(1, -1)
    
    # Run ANE
    start = time.time()
    output = ane_model.predict({"input_ids": input_ids_np.astype(np.int32)})
    ane_time = (time.time() - start) * 1000
    
    # Extract KV
    k_raw = output['var_3891']
    v_raw = output['var_3894']
    
    # Prepare KV cache
    positions = mx.arange(seq_len, dtype=mx.int64).reshape(1, seq_len)
    kv_layers = []
    
    for layer_idx in range(NUM_LAYERS):
        start_idx = layer_idx * ANE_FIXED_LEN
        end_idx = start_idx + seq_len
        k_layer = k_raw[:, start_idx:end_idx, :, :]
        v_layer = v_raw[:, start_idx:end_idx, :, :]
        
        k_mx = mx.array(k_layer.transpose(0, 2, 1, 3))
        v_mx = mx.array(v_layer.transpose(0, 2, 1, 3))
        k_mx = apply_rope_mlx(k_mx, positions, ROPE_THETA)
        
        kv_layers.append({'k': k_mx, 'v': v_mx})
    
    return kv_layers, ane_time, input_ids_np.shape[1]  # 返回实际 padding 后的长度

def gpu_prefill(tokens):
    """GPU Prefill"""
    start = time.time()
    
    kv_cache = make_prompt_cache(model)
    input_ids = mx.array(tokens).reshape(1, -1)
    _ = model(input_ids, cache=kv_cache)
    
    gpu_time = (time.time() - start) * 1000
    return kv_cache, gpu_time

# 测试不同长度
test_configs = [
    {'tokens': 10, 'desc': '极短'},
    {'tokens': 32, 'desc': '短'},
    {'tokens': 64, 'desc': '中等'},
    {'tokens': 100, 'desc': '较长'},
    {'tokens': 128, 'desc': 'ANE 固定长度'},
    {'tokens': 200, 'desc': '超过 ANE 长度'},
    {'tokens': 256, 'desc': '2x ANE 长度'},
    {'tokens': 512, 'desc': '4x ANE 长度'},
]

base_text = "The quick brown fox jumps over the lazy dog. This is a test sentence. "

print("\n" + "="*60)
print("📈 详细测试结果")
print("="*60)

results = []

for config in test_configs:
    target_tokens = config['tokens']
    desc = config['desc']
    
    print(f"\n{'='*60}")
    print(f"📝 {desc} ({target_tokens} tokens)")
    print(f"{'='*60}")
    
    # 生成 prompt
    prompt = base_text
    while len(tokenizer.encode(prompt)) < target_tokens:
        prompt = prompt + " " + base_text
    tokens = tokenizer.encode(prompt)[:target_tokens]
    actual_tokens = len(tokens)
    
    print(f"   实际 tokens: {actual_tokens}")
    
    # ========== ANE Prefill 测试 ==========
    print(f"\n[🔵 ANE Prefill]")
    
    # 对于超过 128 的 prompt，需要多次 ANE 调用
    if actual_tokens > ANE_FIXED_LEN:
        print(f"   ⚠️ Prompt 超过 {ANE_FIXED_LEN} tokens，需要多次 ANE 调用")
        num_calls = (actual_tokens + ANE_FIXED_LEN - 1) // ANE_FIXED_LEN
        print(f"   需要 {num_calls} 次 ANE 调用")
        
        # 模拟多次调用
        total_ane_time = 0
        for i in range(num_calls):
            start_idx = i * ANE_FIXED_LEN
            end_idx = min((i + 1) * ANE_FIXED_LEN, actual_tokens)
            chunk_tokens = tokens[start_idx:end_idx]
            
            kv, chunk_time, padded_len = ane_prefill(chunk_tokens)
            total_ane_time += chunk_time
            
            print(f"   Chunk {i+1}: {len(chunk_tokens)} tokens (padded to {padded_len}), {chunk_time:.2f}ms")
        
        ane_time = total_ane_time
    else:
        kv_ane, ane_time, padded_len = ane_prefill(tokens)
        print(f"   Padding: {actual_tokens} → {padded_len} tokens")
    
    print(f"   ✅ ANE Prefill 总时间: {ane_time:.2f}ms")
    ane_tok_per_sec = actual_tokens / (ane_time / 1000) if ane_time > 0 else 0
    print(f"   速度: {ane_tok_per_sec:.1f} tok/s")
    
    # ========== GPU Prefill 测试 ==========
    print(f"\n[🟢 GPU Prefill]")
    
    kv_gpu, gpu_time = gpu_prefill(tokens)
    print(f"   ✅ GPU Prefill: {gpu_time:.2f}ms")
    gpu_tok_per_sec = actual_tokens / (gpu_time / 1000) if gpu_time > 0 else 0
    print(f"   速度: {gpu_tok_per_sec:.1f} tok/s")
    
    # ========== 对比 ==========
    print(f"\n[📊 对比]")
    if gpu_time < ane_time:
        speedup = ane_time / gpu_time
        print(f"   🏆 GPU 快 {speedup:.1f}x")
        winner = "GPU"
    else:
        speedup = gpu_time / ane_time
        print(f"   🏆 ANE 快 {speedup:.1f}x")
        winner = "ANE"
    
    results.append({
        'tokens': actual_tokens,
        'desc': desc,
        'ane_time_ms': ane_time,
        'ane_tok_per_sec': ane_tok_per_sec,
        'gpu_time_ms': gpu_time,
        'gpu_tok_per_sec': gpu_tok_per_sec,
        'winner': winner,
    })

# 总结
print("\n" + "="*60)
print("📊 性能总结")
print("="*60)

print("\n| Tokens | ANE (ms) | ANE tok/s | GPU (ms) | GPU tok/s | 胜者 |")
print("|--------|----------|-----------|----------|-----------|------|")

for r in results:
    print(f"| {r['tokens']:6d} | {r['ane_time_ms']:8.2f} | {r['ane_tok_per_sec']:9.1f} | {r['gpu_time_ms']:8.2f} | {r['gpu_tok_per_sec']:9.1f} | {r['winner']:4s} |")

# 统计
gpu_wins = sum(1 for r in results if r['winner'] == 'GPU')
ane_wins = sum(1 for r in results if r['winner'] == 'ANE')

print(f"\nGPU 胜出: {gpu_wins}/{len(results)} 次")
print(f"ANE 胜出: {ane_wins}/{len(results)} 次")

# 平均性能
avg_ane_time = sum(r['ane_time_ms'] for r in results) / len(results)
avg_gpu_time = sum(r['gpu_time_ms'] for r in results) / len(results)

print(f"\n平均 ANE Prefill: {avg_ane_time:.2f}ms")
print(f"平均 GPU Prefill: {avg_gpu_time:.2f}ms")

print("\n" + "="*60)
print("💡 结论")
print("="*60)

if gpu_wins > ane_wins:
    print("\n✅ GPU Prefill 在大多数场景下更快")
    print("   建议：短/中 prompt 使用 GPU Prefill")
else:
    print("\n✅ ANE Prefill 在某些场景下更快")

print("\n" + "="*60)