test_e2e.py

#!/usr/bin/env python3
"""PD-Engine 端到端验证

目标：验证 ANE Prefill + GPU Decode 完整流程
"""

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

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

print("="*60)
print("🎯 PD-Engine 端到端验证")
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

def apply_rope_mlx(x: mx.array, positions: mx.array, theta: float = 10000.0) -> mx.array:
    """Apply RoPE to K/V cache."""
    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)
    sin = mx.sin(angles)
    x1 = x[..., :head_dim // 2]
    x2 = x[..., head_dim // 2:]
    cos = cos[:, None, :, :]
    sin = sin[:, None, :, :]
    return mx.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)


# ========== 测试用例 ==========
test_cases = [
    ("The capital of France is", "Paris"),
    ("Hello, my name is", "John"),
    ("2 + 2 equals", "4"),
]

print("\n" + "="*60)
print("📊 端到端测试")
print("="*60)

for prompt, expected_keyword in test_cases:
    print(f"\n{'='*60}")
    print(f"📝 Prompt: '{prompt}'")
    print(f"   期望关键词: '{expected_keyword}'")
    print(f"{'='*60}")
    
    tokens = tokenizer.encode(prompt)
    seq_len = len(tokens)
    
    # ===== Step 1: ANE Prefill =====
    print("\n[Step 1: ANE Prefill]")
    start = time.time()
    
    # Pad to 128
    input_ids_np = np.pad(tokens, (0, 128 - seq_len), mode='constant').reshape(1, -1)
    ane_output = ane_model.predict({"input_ids": input_ids_np.astype(np.int32)})
    
    k_raw = ane_output['var_3891']  # (1, 2816, 4, 64)
    v_raw = ane_output['var_3894']
    
    ane_time = (time.time() - start) * 1000
    print(f"   ✅ ANE Prefill: {ane_time:.2f}ms")
    print(f"   K shape: {k_raw.shape}, V shape: {v_raw.shape}")
    
    # ===== Step 2: 准备 KV Cache =====
    print("\n[Step 2: 准备 KV Cache]")
    
    # 提取每层的 KV
    kv_layers = []
    positions = mx.arange(seq_len, dtype=mx.int64).reshape(1, seq_len)
    
    for layer_idx in range(num_layers):
        start_idx = layer_idx * 128
        end_idx = start_idx + seq_len
        k_layer = k_raw[:, start_idx:end_idx, :, :]  # (1, seq, 4, 64)
        v_layer = v_raw[:, start_idx:end_idx, :, :]
        
        # 转换为 (batch, heads, seq, dim)
        k_mx = mx.array(k_layer.transpose(0, 2, 1, 3))
        v_mx = mx.array(v_layer.transpose(0, 2, 1, 3))
        
        # 应用 RoPE 到 K
        k_mx = apply_rope_mlx(k_mx, positions, rope_theta)
        
        kv_layers.append({'k': k_mx, 'v': v_mx})
    
    print(f"   ✅ KV Cache 准备完成: {len(kv_layers)} layers")
    
    # ===== Step 3: GPU Decode =====
    print("\n[Step 3: GPU Decode]")
    start = time.time()
    
    # 简化版 decode：直接用 embedding + forward
    generated_tokens = []
    last_token = mx.array([[tokens[-1]]])
    
    for step in range(10):
        # Embedding
        hidden = model.model.embed_tokens(last_token)
        
        # Forward through layers
        for layer_idx, layer in enumerate(model.model.layers):
            # Pre-LN
            hidden_normed = layer.input_layernorm(hidden)
            
            # Q projection
            q = layer.self_attn.q_proj(hidden_normed)
            q = q.reshape(1, 1, 32, 64).transpose(0, 2, 1, 3)  # (1, 32, 1, 64)
            
            # Apply RoPE to Q
            pos = mx.array([[seq_len + step]], dtype=mx.int64)
            q = apply_rope_mlx(q, pos, rope_theta)
            
            # Get K, V from cache
            k_cache = kv_layers[layer_idx]['k']  # (1, 4, seq, 64)
            v_cache = kv_layers[layer_idx]['v']
            
            # Compute new K, V for this position
            new_k = layer.self_attn.k_proj(hidden_normed)
            new_v = layer.self_attn.v_proj(hidden_normed)
            new_k = new_k.reshape(1, 1, 4, 64).transpose(0, 2, 1, 3)
            new_v = new_v.reshape(1, 1, 4, 64).transpose(0, 2, 1, 3)
            
            # Apply RoPE to new K
            new_k = apply_rope_mlx(new_k, pos, rope_theta)
            
            # Extend cache
            k_cache = mx.concatenate([k_cache, new_k], axis=2)
            v_cache = mx.concatenate([v_cache, new_v], axis=2)
            
            # Update cache
            kv_layers[layer_idx]['k'] = k_cache
            kv_layers[layer_idx]['v'] = v_cache
            
            # GQA: expand KV
            k_expanded = mx.repeat(k_cache, 8, axis=1)  # 32 heads
            v_expanded = mx.repeat(v_cache, 8, axis=1)
            
            # Attention
            scale = 1.0 / mx.sqrt(mx.array(64.0))
            scores = mx.matmul(q, k_expanded.transpose(0, 1, 3, 2)) * scale
            
            # Softmax
            attn_weights = mx.softmax(scores.astype(mx.float32), axis=-1).astype(mx.float16)
            attn_output = mx.matmul(attn_weights, v_expanded)
            
            # Output projection
            attn_output = attn_output.transpose(0, 2, 1, 3).reshape(1, 1, -1)
            attn_output = layer.self_attn.o_proj(attn_output)
            
            # Residual
            hidden = hidden + attn_output
            
            # MLP
            residual = hidden
            hidden = layer.post_attention_layernorm(hidden)
            hidden = residual + layer.mlp(hidden)
        
        # LM Head
        hidden = model.model.norm(hidden)
        logits = model.lm_head(hidden)
        
        # Sample
        next_token = mx.argmax(logits[0, 0, :])
        generated_tokens.append(int(next_token.tolist()))
        last_token = mx.array([[int(next_token.tolist())]])
    
    decode_time = (time.time() - start) * 1000
    output_text = tokenizer.decode(generated_tokens)
    
    print(f"   ✅ GPU Decode: {decode_time:.2f}ms")
    print(f"   生成速度: {10 / (decode_time/1000):.1f} tok/s")
    
    # ===== 结果验证 =====
    print("\n[结果验证]")
    print(f"   输出: '{output_text}'")
    
    if expected_keyword.lower() in output_text.lower():
        print(f"   ✅ 包含期望关键词: '{expected_keyword}'")
    else:
        print(f"   ⚠️ 未包含期望关键词: '{expected_keyword}'")
    
    # ===== 对比纯 GPU =====
    print("\n[对比: 纯 GPU]")
    gpu_output = generate(model, tokenizer, prompt=prompt, max_tokens=10, verbose=False)
    print(f"   输出: '{gpu_output}'")

print("\n" + "="*60)
print("✅ 端到端验证完成")
print("="*60)