Upload vllm_rt_qwen_mobile_actions.py with huggingface_hub

vllm_rt_qwen_mobile_actions.py

@@ -0,0 +1,423 @@
+"""
+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()