Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 0 additions & 5 deletions examples/mem_os/simple_vllm_memos.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@
"""
Simple example demonstrating how to use VLLMLLM with an existing vLLM server.
Requires a vLLM server to be running.
Simple example demonstrating how to use VLLMLLM with an existing vLLM server.
Requires a vLLM server to be running.
"""

from memos.configs.llm import VLLMLLMConfig
Expand All @@ -15,11 +13,9 @@ def main():

# Configuration for connecting to existing vLLM server
config = VLLMLLMConfig(
model_name_or_path="/mnt/afs/models/hf_models/Qwen2.5-7B", # MUST MATCH the --model arg of vLLM server
model_name_or_path="/mnt/afs/models/hf_models/Qwen2.5-7B", # MUST MATCH the --model arg of vLLM server
api_key="", # Not needed for local server
api_base="http://localhost:8088/v1", # vLLM server address with /v1
api_base="http://localhost:8088/v1", # vLLM server address with /v1
temperature=0.7,
max_tokens=512,
top_p=0.9,
Expand All @@ -39,7 +35,6 @@ def main():
try:
prompt = llm.build_vllm_kv_cache(system_messages)
print(f"✓ KV cache built successfully for prompt: '{prompt[:100]}...'")
print(f"✓ KV cache built successfully for prompt: '{prompt[:100]}...'")
except Exception as e:
print(f"✗ Failed to build KV cache: {e}")

Expand Down
29 changes: 0 additions & 29 deletions src/memos/llms/vllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,7 @@ def __init__(self, config: VLLMLLMConfig):
api_key=api_key,
base_url=getattr(self.config, "api_base", "http://localhost:8088/v1")
)
api_key = getattr(self.config, "api_key", "dummy")
if not api_key:
api_key = "dummy"

import openai
self.client = openai.Client(
api_key=api_key,
base_url=getattr(self.config, "api_base", "http://localhost:8088/v1")
)

def build_vllm_kv_cache(self, messages: Any) -> str:
def build_vllm_kv_cache(self, messages: Any) -> str:
"""
Build a KV cache from chat messages via one vLLM request.
Expand All @@ -67,21 +57,12 @@ def build_vllm_kv_cache(self, messages: Any) -> str:

if not prompt.strip():
raise ValueError("Prompt is empty, cannot build KV cache.")
raise ValueError("Prompt is empty, cannot build KV cache.")

# 3. Send request to vLLM server to preload the KV cache
if self.client:
try:
# Use the processed messages for the API call
# 3. Send request to vLLM server to preload the KV cache
if self.client:
try:
# Use the processed messages for the API call
prefill_kwargs = {
"model": self.config.model_name_or_path,
"messages": processed_messages,
"max_tokens": 2,
"temperature": 0.0,
"model": self.config.model_name_or_path,
"messages": processed_messages,
"max_tokens": 2,
Expand All @@ -90,8 +71,6 @@ def build_vllm_kv_cache(self, messages: Any) -> str:
}
self.client.chat.completions.create(**prefill_kwargs)
logger.info(f"vLLM KV cache prefill completed for prompt: '{prompt[:100]}...'")
self.client.chat.completions.create(**prefill_kwargs)
logger.info(f"vLLM KV cache prefill completed for prompt: '{prompt[:100]}...'")
except Exception as e:
logger.warning(f"Failed to prefill vLLM KV cache: {e}")

Expand All @@ -101,7 +80,6 @@ def generate(self, messages: list[MessageDict]) -> str:
"""
Generate a response from the model.
"""
if self.client:
if self.client:
return self._generate_with_api_client(messages)
else:
Expand All @@ -111,11 +89,8 @@ def _generate_with_api_client(self, messages: list[MessageDict]) -> str:
"""
Generate response using vLLM API client.
"""
if self.client:
if self.client:
completion_kwargs = {
"model": self.config.model_name_or_path,
"messages": messages,
"model": self.config.model_name_or_path,
"messages": messages,
"temperature": float(getattr(self.config, "temperature", 0.8)),
Expand All @@ -127,9 +102,6 @@ def _generate_with_api_client(self, messages: list[MessageDict]) -> str:
response_text = response.choices[0].message.content or ""
logger.info(f"VLLM API response: {response_text}")
return remove_thinking_tags(response_text) if getattr(self.config, "remove_think_prefix", False) else response_text
response_text = response.choices[0].message.content or ""
logger.info(f"VLLM API response: {response_text}")
return remove_thinking_tags(response_text) if getattr(self.config, "remove_think_prefix", False) else response_text
else:
raise RuntimeError("API client is not available")

Expand All @@ -142,7 +114,6 @@ def _messages_to_prompt(self, messages: list[MessageDict]) -> str:
role = msg["role"]
content = msg["content"]
prompt_parts.append(f"{role.capitalize()}: {content}")
prompt_parts.append(f"{role.capitalize()}: {content}")
return "\n".join(prompt_parts)

def generate_stream(self, messages: list[MessageDict]):
Expand Down
Loading