import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

# 探针1: 使用小模型（GPT-2）
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 探针2: 有重复模式的文本（便于观察cache）
tracer_text = "TRACER_TOKEN " * 5  # 重复5次

def trace_kvcache_with_profiling(text, max_new_tokens=5):
    """完整追踪KV Cache的生成和使用"""

    # 编码
    inputs = tokenizer(text, return_tensors="pt")
    print(f"\n[实验设置]")
    print(f"输入文本：\"{text}\"")
    print(f"Token数量：{len(tokenizer.encode(text))}")

    # ============================================================
    # 阶段1：Prefill（处理输入的prompt）
    # ============================================================
    print(f"\n{'='*60}")
    print(f"[阶段1] Prefill阶段 - 处理prompt")
    print(f"{'='*60}")

    start = time.time()
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1,  # 只生成1个token
            output_hidden_states=True,
            return_dict_in_generate=True
        )
    prefill_time = time.time() - start

    first_token_id = outputs['sequences'][0, -1].item()
    print(f"✓ 生成第1个token：ID={first_token_id}")
    print(f"✓ 耗时：{prefill_time:.3f}秒")

    # 检查KV Cache
    past_kv = outputs.past_key_values
    if past_kv is not None:
        print(f"\n[KV Cache详情]")
        print(f"  层数：{len(past_kv)}")
        print(f"  第0层KV shape: {past_kv[0][0].shape}")
        print(f"  说明：[batch, num_heads, seq_len, head_dim]")

    # ============================================================
    # 阶段2：Decode（逐个生成后续token）
    # ============================================================
    print(f"\n{'='*60}")
    print(f"[阶段2] Decode阶段 - 逐个生成token")
    print(f"{'='*60}")

    past_kv = outputs.past_key_values
    current_input = outputs['sequences'][:, -1:]  # 最后一个token

    decode_times = []
    memory_usage = []

    for i in range(max_new_tokens - 1):
        start = time.time()
        with torch.no_grad():
            outputs = model(
                current_input,
                past_key_values=past_kv,
                use_cache=True  # 关键：使用cache
            )
        token_time = time.time() - start
        decode_times.append(token_time)

        # 获取下一个token
        next_token_logits = outputs.logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

        # 更新输入和cache
        current_input = next_token
        past_kv = outputs.past_key_values

        memory_mb = torch.cuda.memory_allocated() / 1024**2
        memory_usage.append(memory_mb)

        print(f"[生成token {i+2}] 耗时：{token_time:.3f}秒，显存：{memory_mb:.2f} MB")

    # ============================================================
    # 总结
    # ============================================================
    print(f"\n{'='*60}")
    print(f"[总结] 性能分析")
    print(f"{'='*60}")
    print(f"Prefill时间（第1个token）：{prefill_time:.3f}秒")
    print(f"Decode平均时间（后续token）：{sum(decode_times)/len(decode_times):.3f}秒")
    print(f"Prefill是Decode的{prefill_time / (sum(decode_times)/len(decode_times)):.1f}倍")

# 执行探针
trace_kvcache_with_profiling(tracer_text, max_new_tokens=5)

# 步骤1: 运行探针
python kvcache_tracer.py

# 步骤2: 在另一个终端监控GPU
watch -n 0.5 nvidia-smi

# 步骤3: 对比使用cache vs 不使用cache
python -c "
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained('gpt2')
tokenizer = AutoTokenizer.from_pretrained('gpt2')
text = 'Hello world'
inputs = tokenizer(text, return_tensors='pt')

# 使用cache
start = time.time()
_ = model.generate(**inputs, max_new_tokens=10, use_cache=True)
time_with_cache = time.time() - start

# 不使用cache
start = time.time()
_ = model.generate(**inputs, max_new_tokens=10, use_cache=False)
time_without_cache = time.time() - start

print(f'使用cache: {time_with_cache:.3f}s')
print(f'不使用cache: {time_without_cache:.3f}s')
print(f'加速比: {time_without_cache/time_with_cache:.2f}x')
"