test_optimized.py

#!/usr/bin/env python3
"""测试优化版 PD-Engine"""

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.generate import generate_step
from mlx_lm.sample_utils import make_sampler
from mlx_lm.models.cache import make_prompt_cache
from transformers import AutoTokenizer

print("="*60)
print("🚀 PD-Engine 优化版测试")
print("="*60)

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

# 测试用例
test_cases = [
    {'prompt': 'The capital of France is', 'max_tokens': 20, 'expected': 'Paris'},
    {'prompt': '2 + 2 equals', 'max_tokens': 10, 'expected': '4'},
    {'prompt': 'Hello, my name is', 'max_tokens': 15, 'expected': ''},
]

print("\n" + "="*60)
print("📊 测试结果")
print("="*60)

results = []

for case in test_cases:
    print(f"\n📝 Prompt: '{case['prompt']}'")
    
    tokens = tokenizer.encode(case['prompt'])
    
    # 方法1: 使用 generate_step (优化版)
    print("\n[方法1: generate_step]")
    start = time.time()
    
    kv_cache = make_prompt_cache(model)
    sampler = make_sampler(0.0)
    generated = []
    
    for token, _ in generate_step(
        prompt=mx.array(tokens),
        model=model,
        max_tokens=case['max_tokens'],
        sampler=sampler,
        prompt_cache=kv_cache,
    ):
        if isinstance(token, mx.array):
            generated.append(int(token.item()))
        else:
            generated.append(int(token))
    
    time1 = (time.time() - start) * 1000
    output1 = tokenizer.decode(generated)
    tok_per_sec1 = case['max_tokens'] / (time1 / 1000)
    
    print(f"   Time: {time1:.2f}ms ({tok_per_sec1:.1f} tok/s)")
    print(f"   Output: '{output1}'")
    
    # 方法2: 使用 generate (内置)
    print("\n[方法2: generate (内置)]")
    start = time.time()
    
    output2 = generate(model, tokenizer, prompt=case['prompt'], max_tokens=case['max_tokens'], verbose=False)
    
    time2 = (time.time() - start) * 1000
    tok_per_sec2 = case['max_tokens'] / (time2 / 1000)
    
    print(f"   Time: {time2:.2f}ms ({tok_per_sec2:.1f} tok/s)")
    print(f"   Output: '{output2}'")
    
    # 对比
    speedup = time2 / time1 if time1 > 0 else 0
    print(f"\n   ⚡ Speedup: {speedup:.2f}x")
    
    results.append({
        'prompt': case['prompt'],
        'time_generate_step': time1,
        'time_generate': time2,
        'tok_per_sec_generate_step': tok_per_sec1,
        'tok_per_sec_generate': tok_per_sec2,
    })

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

avg_tok_per_sec_gs = sum(r['tok_per_sec_generate_step'] for r in results) / len(results)
avg_tok_per_sec_g = sum(r['tok_per_sec_generate'] for r in results) / len(results)

print(f"\n| 方法 | 平均 tok/s |")
print(f"|------|-----------|")
print(f"| generate_step | {avg_tok_per_sec_gs:.1f} |")
print(f"| generate (内置) | {avg_tok_per_sec_g:.1f} |")

if avg_tok_per_sec_gs > avg_tok_per_sec_g:
    speedup = avg_tok_per_sec_gs / avg_tok_per_sec_g
    print(f"\n✅ generate_step 快 {speedup:.2f}x")
else:
    slowdown = avg_tok_per_sec_g / avg_tok_per_sec_gs
    print(f"\n⚠️ generate_step 慢 {slowdown:.2f}x")

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