pd_engine_optimized.py

#!/usr/bin/env python3
"""PD-Engine 优化版 - 使用 MLX 内置 generate_step

优化点：
1. 用 MLX generate_step 替代手动 decode 循环
2. 动态路由：短 prompt 用 GPU，长 prompt 用 ANE
3. 完整的性能测试框架
"""

import sys
import time
import argparse
from typing import Optional, List, Tuple

import mlx.core as mx
from mlx_lm import load
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
import coremltools as ct
import numpy as np

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

# ============================================================
# 配置
# ============================================================

class Config:
    """全局配置"""
    # 模型配置
    MODEL_PATH = "tinyllama-mlx"
    TOKENIZER_PATH = "tinyllama/tinyllama-1.1b-chat-v1.0"
    ANE_MODEL_PATH = "model_ane_full.mlpackage"
    
    # 架构配置
    NUM_LAYERS = 22
    NUM_HEADS = 32
    NUM_KV_HEADS = 4
    HEAD_DIM = 64
    ROPE_THETA = 10000.0
    
    # 性能配置
    ANE_THRESHOLD = 64  # tokens，超过此值用 ANE
    ANE_FIXED_LEN = 128  # ANE 固定处理长度


# ============================================================
# RoPE 实现
# ============================================================

def apply_rope_mlx(x: mx.array, positions: mx.array, theta: float = 10000.0) -> mx.array:
    """Apply Rotary Positional Embedding.
    
    Args:
        x: (batch, num_heads, seq_len, head_dim)
        positions: (batch, seq_len) or (1, seq_len)
        theta: RoPE base frequency
    
    Returns:
        Rotated tensor of same shape
    """
    batch, num_heads, seq_len, head_dim = x.shape
    
    # Compute frequencies
    freqs = 1.0 / (theta ** (mx.arange(0, head_dim, 2)[:head_dim // 2] / head_dim))
    
    # Compute angles
    angles = positions[:, :, None] * freqs[None, None, :]
    
    cos = mx.cos(angles)
    sin = mx.sin(angles)
    
    # Split and rotate
    x1 = x[..., :head_dim // 2]
    x2 = x[..., head_dim // 2:]
    
    # Broadcast
    cos = cos[:, None, :, :]  # (batch, 1, seq, dim//2)
    sin = sin[:, None, :, :]
    
    return mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)


# ============================================================
# ANE Prefill
# ============================================================

class ANEPrefiller:
    """ANE Prefill 引擎"""
    
    def __init__(self, model_path: str, config: Config):
        self.config = config
        print(f"📦 加载 ANE 模型: {model_path}")
        self.model = ct.models.MLModel(model_path)
        print("   ✅ ANE 模型加载完成")
    
    def prefill(self, tokens: List[int]) -> Tuple[List[dict], float]:
        """运行 ANE Prefill，返回 KV cache"""
        seq_len = len(tokens)
        
        # Pad to fixed length
        padded_len = self.config.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 = self.model.predict({"input_ids": input_ids_np.astype(np.int32)})
        ane_time = (time.time() - start) * 1000
        
        k_raw = output['var_3891']  # (1, 2816, 4, 64)
        v_raw = output['var_3894']
        
        # Extract KV for each layer
        kv_layers = []
        positions = mx.arange(seq_len, dtype=mx.int64).reshape(1, seq_len)
        
        for layer_idx in range(self.config.NUM_LAYERS):
            start_idx = layer_idx * padded_len
            end_idx = start_idx + seq_len
            k_layer = k_raw[:, start_idx:end_idx, :, :]
            v_layer = v_raw[:, start_idx:end_idx, :, :]
            
            # Convert to MLX format
            k_mx = mx.array(k_layer.transpose(0, 2, 1, 3))  # (1, heads, seq, dim)
            v_mx = mx.array(v_layer.transpose(0, 2, 1, 3))
            
            # Apply RoPE to K
            k_mx = apply_rope_mlx(k_mx, positions, self.config.ROPE_THETA)
            
            kv_layers.append({'k': k_mx, 'v': v_mx})
        
        return kv_layers, ane_time


# ============================================================
# GPU Prefill
# ============================================================

class GPUPrefiller:
    """GPU Prefill 引擎（用于短 prompt）"""
    
    def __init__(self, model, config: Config):
        self.model = model
        self.config = config
    
    def prefill(self, tokens: List[int]) -> Tuple[list, float]:
        """运行 GPU Prefill，返回 KV cache"""
        start = time.time()
        
        # Create KV cache
        kv_cache = make_prompt_cache(self.model)
        
        # Run forward pass
        input_ids = mx.array(tokens).reshape(1, -1)
        _ = self.model(input_ids, cache=kv_cache)
        
        gpu_time = (time.time() - start) * 1000
        
        return kv_cache, gpu_time


# ============================================================
# 混合推理引擎
# ============================================================

class HybridEngine:
    """ANE + GPU 混合推理引擎"""
    
    def __init__(self, config: Optional[Config] = None):
        self.config = config or Config()
        
        # Load MLX model
        print(f"\n📦 加载 MLX 模型: {self.config.MODEL_PATH}")
        self.model, _ = load(self.config.MODEL_PATH)
        self.tokenizer = AutoTokenizer.from_pretrained(self.config.TOKENIZER_PATH)
        print("   ✅ MLX 模型加载完成")
        
        # Load ANE model
        self.ane_prefiller = ANEPrefiller(self.config.ANE_MODEL_PATH, self.config)
        self.gpu_prefiller = GPUPrefiller(self.model, self.config)
    
    def generate(
        self,
        prompt: str,
        max_tokens: int = 50,
        temp: float = 0.0,
        verbose: bool = True
    ) -> dict:
        """生成文本，自动选择最优路径
        
        Returns:
            {
                'output': str,
                'prefill_time_ms': float,
                'decode_time_ms': float,
                'total_time_ms': float,
                'tok_per_sec': float,
                'route': 'ane' | 'gpu'
            }
        """
        tokens = self.tokenizer.encode(prompt)
        seq_len = len(tokens)
        
        if verbose:
            print(f"\n📝 Prompt: '{prompt}' ({seq_len} tokens)")
        
        # 动态路由
        use_ane = seq_len > self.config.ANE_THRESHOLD
        route = 'ane' if use_ane else 'gpu'
        
        if verbose:
            print(f"   路由: {route.upper()} (threshold={self.config.ANE_THRESHOLD})")
        
        # Prefill
        if use_ane:
            kv_cache, prefill_time = self.ane_prefiller.prefill(tokens)
            # TODO: 将 ANE KV 转换为 MLX KVCache 格式
            # 暂时用 GPU prefill
            kv_cache, prefill_time = self.gpu_prefiller.prefill(tokens)
        else:
            kv_cache, prefill_time = self.gpu_prefiller.prefill(tokens)
        
        if verbose:
            print(f"   Prefill: {prefill_time:.2f}ms")
        
        # Decode (使用 MLX generate_step)
        start_decode = time.time()
        
        sampler = make_sampler(temp)
        generated_tokens = []
        
        for token, _ in generate_step(
            prompt=mx.array(tokens),
            model=self.model,
            max_tokens=max_tokens,
            sampler=sampler,
            prompt_cache=kv_cache,
        ):
            generated_tokens.append(int(token.item()))
        
        decode_time = (time.time() - start_decode) * 1000
        output = self.tokenizer.decode(generated_tokens)
        
        total_time = prefill_time + decode_time
        tok_per_sec = max_tokens / (decode_time / 1000) if decode_time > 0 else 0
        
        if verbose:
            print(f"   Decode: {decode_time:.2f}ms ({tok_per_sec:.1f} tok/s)")
            print(f"   Total: {total_time:.2f}ms")
            print(f"\n💬 Output: '{output}'")
        
        return {
            'output': output,
            'prefill_time_ms': prefill_time,
            'decode_time_ms': decode_time,
            'total_time_ms': total_time,
            'tok_per_sec': tok_per_sec,
            'route': route
        }


# ============================================================
# Benchmark 框架
# ============================================================

class Benchmark:
    """性能测试框架"""
    
    def __init__(self, engine: HybridEngine):
        self.engine = engine
    
    def run_suite(self, test_cases: List[dict]):
        """运行测试套件
        
        Args:
            test_cases: [{'prompt': str, 'max_tokens': int, 'expected': str}]
        """
        print("\n" + "="*60)
        print("📊 性能测试")
        print("="*60)
        
        results = []
        
        for i, case in enumerate(test_cases, 1):
            print(f"\n[Test {i}/{len(test_cases)}]")
            
            result = self.engine.generate(
                prompt=case['prompt'],
                max_tokens=case.get('max_tokens', 20),
                verbose=False
            )
            
            result['test_case'] = case
            result['has_expected'] = case.get('expected', '').lower() in result['output'].lower()
            
            print(f"   Prompt: '{case['prompt'][:50]}...'")
            print(f"   Route: {result['route'].upper()}")
            print(f"   Prefill: {result['prefill_time_ms']:.2f}ms")
            print(f"   Decode: {result['decode_time_ms']:.2f}ms ({result['tok_per_sec']:.1f} tok/s)")
            print(f"   Output: '{result['output'][:50]}...'")
            print(f"   Expected: {'✅' if result['has_expected'] else '⚠️'}")
            
            results.append(result)
        
        # 统计
        self._print_summary(results)
        
        return results
    
    def _print_summary(self, results: List[dict]):
        """打印统计摘要"""
        print("\n" + "="*60)
        print("📈 测试总结")
        print("="*60)
        
        # 按路由分组
        ane_results = [r for r in results if r['route'] == 'ane']
        gpu_results = [r for r in results if r['route'] == 'gpu']
        
        if ane_results:
            avg_prefill = sum(r['prefill_time_ms'] for r in ane_results) / len(ane_results)
            avg_decode = sum(r['tok_per_sec'] for r in ane_results) / len(ane_results)
            print(f"\nANE 路由 ({len(ane_results)} tests):")
            print(f"   平均 Prefill: {avg_prefill:.2f}ms")
            print(f"   平均 Decode: {avg_decode:.1f} tok/s")
        
        if gpu_results:
            avg_prefill = sum(r['prefill_time_ms'] for r in gpu_results) / len(gpu_results)
            avg_decode = sum(r['tok_per_sec'] for r in gpu_results) / len(gpu_results)
            print(f"\nGPU 路由 ({len(gpu_results)} tests):")
            print(f"   平均 Prefill: {avg_prefill:.2f}ms")
            print(f"   平均 Decode: {avg_decode:.1f} tok/s")
        
        # 准确率
        accuracy = sum(1 for r in results if r['has_expected']) / len(results) * 100
        print(f"\n准确率: {accuracy:.1f}%")


# ============================================================
# 主程序
# ============================================================

def main():
    parser = argparse.ArgumentParser(description='PD-Engine 优化版')
    parser.add_argument('--prompt', type=str, help='输入 prompt')
    parser.add_argument('--max-tokens', type=int, default=30, help='最大生成 tokens')
    parser.add_argument('--benchmark', action='store_true', help='运行性能测试')
    args = parser.parse_args()
    
    # 初始化引擎
    engine = HybridEngine()
    
    if args.benchmark:
        # 运行 benchmark
        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': ''},
            {'prompt': 'The quick brown fox jumps over the lazy dog. ' * 3, 'max_tokens': 30, 'expected': ''},
            {'prompt': 'In a galaxy far far away, there lived a', 'max_tokens': 25, 'expected': ''},
        ]
        
        benchmark = Benchmark(engine)
        benchmark.run_suite(test_cases)
    
    elif args.prompt:
        # 单次生成
        engine.generate(args.prompt, max_tokens=args.max_tokens)
    
    else:
        # 交互模式
        print("\n💬 交互模式 (输入 'quit' 退出)")
        print("-"*60)
        
        while True:
            try:
                prompt = input("\n你: ").strip()
                if prompt.lower() == 'quit':
                    break
                if not prompt:
                    continue
                
                engine.generate(prompt, max_tokens=50)
            
            except (EOFError, KeyboardInterrupt):
                break
        
        print("\n\n👋 再见！")


if __name__ == "__main__":
    main()