Cialtion
/

SimpleTool

@@ -1,423 +0,0 @@
-"""
-vLLM Prefill/Decode 分离性能测试脚本
-=============================================
-核心监控指标：
-- Prefill tokens 数量 (实际需要计算的)
-- KV Cache 命中 tokens 数量
-- Prefill 耗时 (ms)
-- Decode 耗时 (ms per token)
-- 每步 overhead 分析
-"""
-import time
-import uuid
-import numpy as np
-from typing import List, Dict, Any, Optional
-from dataclasses import dataclass, field
-from vllm import LLM, SamplingParams
-@dataclass
-class StepMetrics:
-    """单步指标"""
-    step_idx: int
-    step_type: str  # "prefill" or "decode"
-    duration_ms: float
-    tokens_processed: int = 0  # prefill时是处理的token数，decode时是1
-@dataclass
-class RequestMetrics:
-    """单请求完整指标"""
-    request_id: str
-    tag: str
-    # Token 统计
-    total_prompt_tokens: int = 0
-    cached_tokens: int = 0           # KV cache 命中的
-    computed_tokens: int = 0         # 实际需要 prefill 的
-    output_tokens: int = 0
-    # 时间指标 (ms)
-    prefill_ms: float = 0.0
-    decode_total_ms: float = 0.0
-    decode_per_token_ms: float = 0.0
-    total_ms: float = 0.0
-    # 每步详情
-    steps: List[StepMetrics] = field(default_factory=list)
-    # 输出
-    output_text: str = ""
-    stop_reason: str = ""
-def extract_cache_metrics(output) -> Dict[str, int]:
-    """
-    从 vLLM RequestOutput 提取 cache 相关指标
-    兼容不同 vLLM 版本
-    """
-    result = {
-        'num_cached_tokens': 0,
-        'num_computed_tokens': 0,
-        'num_prompt_tokens': 0,
-    }
-    try:
-        # vLLM 0.5+ metrics
-        if hasattr(output, 'metrics') and output.metrics:
-            m = output.metrics
-            result['num_cached_tokens'] = getattr(m, 'num_cached_tokens', 0) or 0
-            result['num_computed_tokens'] = getattr(m, 'num_computed_tokens', 0) or 0
-            result['num_prompt_tokens'] = getattr(m, 'num_prompt_tokens', 0) or 0
-        # 备用：直接从 prompt_token_ids 获取
-        if result['num_prompt_tokens'] == 0 and hasattr(output, 'prompt_token_ids'):
-            result['num_prompt_tokens'] = len(output.prompt_token_ids)
-    except Exception:
-        pass
-    return result
-class PDProfiler:
-    """Prefill/Decode 性能分析器"""
-    def __init__(self, llm: LLM):
-        self.llm = llm
-        self.engine = llm.llm_engine
-        self.tokenizer = llm.get_tokenizer()
-    def profile_request(
-        self,
-        prompt: str,
-        tag: str = "default",
-        sampling_params: Optional[SamplingParams] = None
-    ) -> RequestMetrics:
-        """
-        分析单个请求的 P/D 性能
-        """
-        if sampling_params is None:
-            sampling_params = SamplingParams(temperature=0.0, max_tokens=32)
-        metrics = RequestMetrics(
-            request_id=str(uuid.uuid4()),
-            tag=tag
-        )
-        # 计算 prompt tokens
-        prompt_tokens = self.tokenizer.encode(prompt, add_special_tokens=False)
-        metrics.total_prompt_tokens = len(prompt_tokens)
-        # 提交请求
-        self.engine.add_request(metrics.request_id, prompt, sampling_params)
-        # Step 循环，记录每步耗时
-        step_idx = 0
-        start_time = time.perf_counter()
-        prev_output_len = 0
-        while self.engine.has_unfinished_requests():
-            step_start = time.perf_counter()
-            outputs = self.engine.step()
-            step_end = time.perf_counter()
-            step_ms = (step_end - step_start) * 1000
-            for out in outputs:
-                if out.request_id != metrics.request_id:
-                    continue
-                # 判断是 prefill 还是 decode
-                current_output_len = len(out.outputs[0].token_ids) if out.outputs else 0
-                if step_idx == 0:
-                    # 第一步是 prefill
-                    step_type = "prefill"
-                    tokens_in_step = metrics.total_prompt_tokens
-                else:
-                    step_type = "decode"
-                    tokens_in_step = current_output_len - prev_output_len
-                metrics.steps.append(StepMetrics(
-                    step_idx=step_idx,
-                    step_type=step_type,
-                    duration_ms=step_ms,
-                    tokens_processed=tokens_in_step
-                ))
-                prev_output_len = current_output_len
-                # 请求完成时提取最终指标
-                if out.finished:
-                    cache_info = extract_cache_metrics(out)
-                    metrics.cached_tokens = cache_info['num_cached_tokens']
-                    metrics.computed_tokens = cache_info['num_computed_tokens']
-                    # 如果 vLLM 没返回 computed_tokens，手动计算
-                    if metrics.computed_tokens == 0 and metrics.cached_tokens > 0:
-                        metrics.computed_tokens = metrics.total_prompt_tokens - metrics.cached_tokens
-                    elif metrics.computed_tokens == 0 and metrics.cached_tokens == 0:
-                        metrics.computed_tokens = metrics.total_prompt_tokens
-                    metrics.output_tokens = current_output_len
-                    metrics.output_text = out.outputs[0].text if out.outputs else ""
-                    metrics.stop_reason = str(getattr(out.outputs[0], 'finish_reason', '')) if out.outputs else ""
-            step_idx += 1
-        metrics.total_ms = (time.perf_counter() - start_time) * 1000
-        # 汇总时间指标
-        prefill_steps = [s for s in metrics.steps if s.step_type == "prefill"]
-        decode_steps = [s for s in metrics.steps if s.step_type == "decode"]
-        metrics.prefill_ms = sum(s.duration_ms for s in prefill_steps)
-        metrics.decode_total_ms = sum(s.duration_ms for s in decode_steps)
-        if metrics.output_tokens > 0:
-            metrics.decode_per_token_ms = metrics.decode_total_ms / metrics.output_tokens
-        return metrics
-    def warmup(self, prompt: str):
-        """预热，确保 KV cache 被填充"""
-        self.engine.add_request(
-            str(uuid.uuid4()),
-            prompt,
-            SamplingParams(max_tokens=1)
-        )
-        while self.engine.has_unfinished_requests():
-            self.engine.step()
-def print_metrics_table(metrics_list: List[RequestMetrics], title: str = ""):
-    """打印性能指标表格"""
-    print(f"\n{'='*120}")
-    if title:
-        print(f" {title}")
-    print(f"{'='*120}")
-    # 表头
-    headers = [
-        "Tag", "PromptTok", "Cached", "Computed", "OutTok",
-        "Prefill(ms)", "Decode(ms)", "Dec/Tok(ms)", "Total(ms)", "Output"
-    ]
-    widths = [12, 10, 8, 10, 8, 12, 12, 12, 10, 30]
-    header_line = " | ".join(f"{h:<{w}}" for h, w in zip(headers, widths))
-    print(header_line)
-    print("-" * 120)
-    for m in metrics_list:
-        output_preview = m.output_text[:28].replace('\n', '\\n') + "..." if len(m.output_text) > 28 else m.output_text.replace('\n', '\\n')
-        row = [
-            m.tag[:12],
-            str(m.total_prompt_tokens),
-            str(m.cached_tokens),
-            str(m.computed_tokens),
-            str(m.output_tokens),
-            f"{m.prefill_ms:.2f}",
-            f"{m.decode_total_ms:.2f}",
-            f"{m.decode_per_token_ms:.2f}",
-            f"{m.total_ms:.2f}",
-            output_preview
-        ]
-        row_line = " | ".join(f"{v:<{w}}" for v, w in zip(row, widths))
-        print(row_line)
-    print("-" * 120)
-    # 汇总统计
-    if len(metrics_list) > 1:
-        avg_prefill = np.mean([m.prefill_ms for m in metrics_list])
-        avg_decode_per_tok = np.mean([m.decode_per_token_ms for m in metrics_list if m.decode_per_token_ms > 0])
-        total_cached = sum(m.cached_tokens for m in metrics_list)
-        total_computed = sum(m.computed_tokens for m in metrics_list)
-        cache_hit_rate = total_cached / (total_cached + total_computed) * 100 if (total_cached + total_computed) > 0 else 0
-        print(f"[Summary] Avg Prefill: {avg_prefill:.2f}ms | Avg Decode/Tok: {avg_decode_per_tok:.2f}ms | Cache Hit Rate: {cache_hit_rate:.1f}%")
-def print_step_details(metrics: RequestMetrics):
-    """打印单请求的每步详情"""
-    print(f"\n[Step Details for '{metrics.tag}']")
-    print(f"  {'Step':<6} {'Type':<10} {'Duration(ms)':<14} {'Tokens':<8}")
-    print(f"  {'-'*40}")
-    for s in metrics.steps:
-        print(f"  {s.step_idx:<6} {s.step_type:<10} {s.duration_ms:<14.2f} {s.tokens_processed:<8}")
-# ==============================================================================
-# 主测试
-# ==============================================================================
-def main():
-    # ==================== 配置 ====================
-    MODEL_PATH = "./RT-Qwen3-4B-AWQ"
-    print("="*60)
-    print(" vLLM Prefill/Decode 分离性能测试")
-    print("="*60)
-    # 初始化 LLM
-    print("\n[Init] Loading model...")
-    llm = LLM(
-        model=MODEL_PATH,
-        trust_remote_code=True,
-        enable_prefix_caching=True,
-        tensor_parallel_size=1,
-        max_num_seqs=16,
-        gpu_memory_utilization=0.8,
-        enforce_eager=False,
-        block_size=16,
-        max_model_len=8192,
-        # 关键：增大 chunk size 减少 chunked prefill 开销
-        # max_num_batched_tokens=8192,  # 可选：设置更大的值
-    )
-    profiler = PDProfiler(llm)
-    tokenizer = llm.get_tokenizer()
-    # ==================== System Prompt ====================
-    system_prompt = (
-        "<|im_start|>system\n"
-        "You are a multi-head parallel function calling model. \n"
-        "## Output Heads\n\n"
-        "**Head 0 - <content>**: Natural language response\n"
-        "- Format: <content>response text</content>\n"
-        "- Answer what you want to say while you are calling a function\n\n"
-        "**Head 1 - <function>**: Function names to call\n"
-        "- Format: <function>name</function>\n"
-        "- Name: must match tool defined name\n\n"
-        "**Head 2-7 - <arg1>、<arg2>、<arg3>、<arg4>、<arg5>、<arg6>**: Function arguments by position\n"
-        "- Format: <argN>value</argN> \n"
-        "- Strictly fill in according to the parameter order of the tool you intend to call\n"
-        "- Note the special restrictions of parameter definitions for corresponding positions\n"
-        "- If the corresponding tool definition has required parameters, these must be filled in\n"
-        "- Infer the user's actual needs.\n"
-        "- If Unnecessary: <argN><|null|></argN>\n\n"
-        "**Environment - The information you have.\n**History - The tools you have called.\n\n"
-        "## Available Tools:\n\n"
-        '{"type": "function", "function": {"name": "open_wifi_settings", "description": "Opens the Wi-Fi settings.", "parameters": {"type": "object", "properties": {}}}}\n'
-        '{"type": "function", "function": {"name": "create_contact", "description": "Creates a contact in the phone\'s contact list.", "parameters": {"type": "object", "properties": {"first_name": {"type": "string", "description": "The first name of the contact."}, "last_name": {"type": "string", "description": "The last name of the contact."}, "email": {"type": "string", "description": "The email address of the contact.", "optional": true}, "phone_number": {"type": "string", "description": "The phone number of the contact.", "optional": true}}, "required": ["first_name", "last_name"]}}}\n'
-        '{"type": "function", "function": {"name": "show_map", "description": "Shows a location on the map.", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "The location to search for. May be the name of a place, a business, or an address."}}, "required": ["query"]}}}\n'
-        '{"type": "function", "function": {"name": "create_calendar_event", "description": "Creates a new calendar event.", "parameters": {"type": "object", "properties": {"title": {"type": "string", "description": "The title of the event."}, "datetime": {"type": "string", "description": "The date and time of the event in the format YYYY-MM-DDTHH:MM:SS."}}, "required": ["title", "datetime"]}}}\n'
-        '{"type": "function", "function": {"name": "send_email", "description": "Sends an email.", "parameters": {"type": "object", "properties": {"to": {"type": "string", "description": "The email address of the recipient."}, "subject": {"type": "string", "description": "The subject of the email."}, "body": {"type": "string", "description": "The body of the email.", "optional": true}}, "required": ["to", "subject"]}}}\n'
-        '{"type": "function", "function": {"name": "turn_off_flashlight", "description": "Turns the flashlight off.", "parameters": {"type": "object", "properties": {}}}}\n'
-        '{"type": "function", "function": {"name": "turn_on_flashlight", "description": "Turns the flashlight on.", "parameters": {"type": "object", "properties": {}}}}\n'
-        "<|im_end|>\n"
-    )
-    system_tokens = len(tokenizer.encode(system_prompt, add_special_tokens=False))
-    print(f"[System Prompt] {system_tokens} tokens")
-    # ==================== 测试用例 ====================
-    test_queries = [
-        ("create_contact", "<|im_start|>user\nenvironment: [\"No develop information provided\"]\nhistory: []\n\nCan you please save a new contact for me? The name is Lena Petrova, the phone number is +359 888 123 456, and the email is lena.petrova.design@webmail.com.<|im_end|>\n<|im_start|>assistant\n"),
-        ("send_email", "<|im_start|>user\nenvironment: [\"No develop information provided\"]\nhistory: []\n\nPlease send an email to javier.ortega@ecotradeintl.com with the subject 'Update on Q4 Report' and the body 'I've uploaded the revised figures to the shared drive.'<|im_end|>\n<|im_start|>assistant\n"),
-        ("calendar", "<|im_start|>user\nenvironment: [\"No develop information provided\"]\nhistory: []\n\nPlease set up a new calendar event for 'Team Lunch with Marketing' on May 13, 2025 at 1:30 PM.<|im_end|>\n<|im_start|>assistant\n"),
-        ("flashlight_map", "<|im_start|>user\nenvironment: [\"No develop information provided\"]\nhistory: []\n\nTurn on the flashlight and show me the location of the Sunnyvale Library on the map.<|im_end|>\n<|im_start|>assistant\n"),
-        ("flashlight_map_w_history", "<|im_start|>user\nenvironment: [\"No develop information provided\"]\nhistory: [turn_on_flashlight()]\n\nTurn on the flashlight and show me the location of the Sunnyvale Library on the map.<|im_end|>\n<|im_start|>assistant\n"),
-    ]
-    head_tags = ["<function>", "<arg1>", "<arg2>", "<arg3>", "<arg4>", "<arg5>"]
-    stop_tokens = [
-        "<|null|>", "</content>", "</function>",
-        "</arg1>", "</arg2>", "</arg3>", "</arg4>", "</arg5>", "</arg6>"
-    ]
-    sampling_params = SamplingParams(
-        temperature=0.0,
-        max_tokens=16,
-        stop=stop_tokens,
-        include_stop_str_in_output=True
-    )
-    # ==================== 测试循环 ====================
-    for query_idx, (query_name, query) in enumerate(test_queries):
-        print(f"\n{'#'*80}")
-        print(f"# ROUND {query_idx + 1}: {query_name}")
-        print(f"{'#'*80}")
-        full_prefix = system_prompt + query
-        prefix_tokens = len(tokenizer.encode(full_prefix, add_special_tokens=False))
-        query_tokens = prefix_tokens - system_tokens
-        print(f"[Prefix] System: {system_tokens} + Query: {query_tokens} = Total: {prefix_tokens} tokens")
-        # ---------------------------------------------------------
-        # 1. 冷启动 Warmup (填充 KV cache)
-        # ---------------------------------------------------------
-        print(f"\n--- Phase 1: Cold Start Warmup ---")
-        warmup_metrics = profiler.profile_request(
-            full_prefix,
-            tag="warmup",
-            sampling_params=SamplingParams(max_tokens=1)
-        )
-        print(f"[Warmup] Prefill {warmup_metrics.computed_tokens} tokens in {warmup_metrics.prefill_ms:.2f}ms")
-        print(f"         Tokens/sec: {warmup_metrics.computed_tokens / warmup_metrics.prefill_ms * 1000:.0f}")
-        # ---------------------------------------------------------
-        # 2. 热启动测试 (KV cache 应该命中)
-        # ---------------------------------------------------------
-        print(f"\n--- Phase 2: Hot Start (Cache Hit Expected) ---")
-        all_metrics = []
-        for head_tag in head_tags:
-            head_prompt = full_prefix + head_tag
-            metrics = profiler.profile_request(
-                head_prompt,
-                tag=head_tag,
-                sampling_params=sampling_params
-            )
-            all_metrics.append(metrics)
-        # 打印表格
-        print_metrics_table(all_metrics, f"Round {query_idx + 1}: {query_name}")
-        # 打印第一个 head 的步骤详情
-        if all_metrics:
-            print_step_details(all_metrics[0])
-    # ==================== 额外测试：不同序列长度的冷启动性能 ====================
-    print(f"\n{'='*80}")
-    print(" BONUS: Cold Start Prefill Performance vs Sequence Length")
-    print(f"{'='*80}")
-    # 构造不同长度的 prompt
-    base_prompt = system_prompt
-    padding_text = "This is padding text to test prefill performance. " * 10
-    length_tests = []
-    for target_len in [256, 512, 1024, 2048]:
-        # 构造指定长度的 prompt
-        test_prompt = base_prompt
-        current_len = len(tokenizer.encode(test_prompt, add_special_tokens=False))
-        while current_len < target_len:
-            test_prompt += padding_text
-            current_len = len(tokenizer.encode(test_prompt, add_special_tokens=False))
-        # 冷启动测试 (新 prompt，无 cache)
-        metrics = profiler.profile_request(
-            test_prompt + f"<unique_{uuid.uuid4().hex[:8]}>",  # 确保无 cache
-            tag=f"len_{target_len}",
-            sampling_params=SamplingParams(max_tokens=1)
-        )
-        length_tests.append(metrics)
-        throughput = metrics.computed_tokens / metrics.prefill_ms * 1000 if metrics.prefill_ms > 0 else 0
-        print(f"[Seq {target_len:4d}] Prefill: {metrics.prefill_ms:7.2f}ms | Computed: {metrics.computed_tokens:4d} | Throughput: {throughput:8.0f} tok/s")
-    print("\n[Analysis] If Prefill time grows faster than linearly, chunked prefill overhead is significant.")
-    print("[Tip] Try increasing max_num_batched_tokens or disabling chunked prefill for latency-critical workloads.")
-if __name__ == "__main__":
-    main()