From e46558bc284d0c49678a71bd358b38895a557cee Mon Sep 17 00:00:00 2001 From: SZCJW <792430652@qq.com> Date: Sat, 11 Apr 2026 01:03:03 +0800 Subject: [PATCH] x --- chat_with_speed.py | 99 ++++++++++++++++++++++++++++++++++++++++++++++ config.json | 6 ++- 2 files changed, 103 insertions(+), 2 deletions(-) create mode 100644 chat_with_speed.py diff --git a/chat_with_speed.py b/chat_with_speed.py new file mode 100644 index 0000000..9012294 --- /dev/null +++ b/chat_with_speed.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +""" +交互式对话脚本,用于与vLLM模型进行对话并计算token生成速度 + +使用方法: +1. 进入Docker容器:docker exec -it rocm-vllm-openai bash +2. 运行:python chat_with_speed.py +3. 输入提示词与模型对话 +4. 输入 'exit' 退出 +""" + +import json +import time +import httpx + +# 模型服务地址 +API_URL = "http://localhost:8001/v1/chat/completions" +# API密钥 +API_KEY = "sk-szcjw" +# 模型名称 +MODEL_NAME = "Qwen_local_model" + +def chat_with_model(): + """交互式对话函数""" + print("=== vLLM 交互式对话工具 ===") + print("输入提示词与模型对话,输入 'exit' 退出") + print("=" * 50) + + # 对话历史 + messages = [] + + while True: + # 获取用户输入 + user_input = input("用户: ").strip() + + if user_input.lower() == "exit": + print("退出对话...") + break + + if not user_input: + continue + + # 添加用户消息到对话历史 + messages.append({"role": "user", "content": user_input}) + + # 准备请求数据 + payload = { + "model": MODEL_NAME, + "messages": messages, + "max_tokens": 1000, + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20 + } + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {API_KEY}" + } + + print("模型: ", end="", flush=True) + + # 记录开始时间 + start_time = time.time() + + try: + # 发送请求 + response = httpx.post(API_URL, json=payload, headers=headers, timeout=300.0) + response.raise_for_status() + + # 解析响应 + result = response.json() + + # 获取模型回复 + assistant_message = result["choices"][0]["message"]["content"] + print(assistant_message) + + # 添加模型回复到对话历史 + messages.append({"role": "assistant", "content": assistant_message}) + + # 计算token速度 + usage = result.get("usage", {}) + completion_tokens = usage.get("completion_tokens", 0) + end_time = time.time() + elapsed_time = end_time - start_time + + if completion_tokens > 0 and elapsed_time > 0: + tokens_per_second = completion_tokens / elapsed_time + print(f"\n[速度统计] 生成 {completion_tokens} tokens,用时 {elapsed_time:.2f} 秒,速度: {tokens_per_second:.2f} tokens/s") + else: + print("\n[速度统计] 无法计算速度") + + except Exception as e: + print(f"\n错误: {e}") + + print("=" * 50) + +if __name__ == "__main__": + chat_with_model() diff --git a/config.json b/config.json index b7fee89..f0f7677 100644 --- a/config.json +++ b/config.json @@ -27,7 +27,7 @@ "ctx": "32768", "trust_remote": true, "valid_tp": [1, 2], - "max_num_seqs": "8", + "max_num_seqs": "16", "max_tokens": "32768", "gpu_util": "0.98", "tool_call_parser": "qwen3_coder", @@ -102,7 +102,9 @@ "valid_tp": [2], "max_num_seqs": "12", "max_tokens": "65536", - "gpu_util": "0.94", + "max_num_batched_tokens": 32768, + "kv_cache_dtype": "fp8", + "gpu_util": "0.92", "enforce_eager": false, "enable_prefix_caching": true, "tool_call_parser": "qwen3_coder",