Skip to content

[Cache] Add cachedPrefixes for caching repeated system prompts #664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+353 −16
Open

[Cache] Add cachedPrefixes for caching repeated system prompts #664

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2af2b16
Added cachedPrefixes to cache long system prompts when creating engine.
YiyanZhai Feb 4, 2025
92e00b9
Cleaned console logs
YiyanZhai Feb 4, 2025
dfd7ba8
Addressed review comments
YiyanZhai Feb 17, 2025
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
14 changes: 14 additions & 0 deletions examples/prefix-caching/README.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
# WebLLM App for Prefix Caching Demo

This example demonstrates the use of `cachedPrefixes` in WebLLM.
To try it out, you can do the following steps under this folder

```bash
npm install
npm start
```

Note if you would like to hack WebLLM core package.
You can change web-llm dependencies as `"file:../.."`, and follow the build from source
instruction in the project to build webllm locally. This option is only recommended
if you would like to hack WebLLM core package.
20 changes: 20 additions & 0 deletions examples/prefix-caching/package.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
{
"name": "prefix-caching-example",
"version": "0.1.0",
"private": true,
"scripts": {
"start": "parcel src/prefix-caching.html --port 8888",
"build": "parcel build src/prefix-caching.html --dist-dir lib"
},
"devDependencies": {
"buffer": "^5.7.1",
"parcel": "^2.8.3",
"process": "^0.11.10",
"tslib": "^2.3.1",
"typescript": "^4.9.5",
"url": "^0.11.3"
},
"dependencies": {
"@mlc-ai/web-llm": "^0.2.78"
}
}
23 changes: 23 additions & 0 deletions examples/prefix-caching/src/prefix-caching.html
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
<!doctype html>
<html>
<script>
webLLMGlobal = {};
</script>
<body>
<h2>WebLLM Prefix Caching Test Page</h2>
Open console to see output
<br />
<br />
<label id="init-label"> </label>

<h3>Prompt</h3>
<label id="prompt-label"> </label>

<h3>Response</h3>
<label id="generate-label"> </label>
<br />
<label id="stats-label"> </label>

<script type="module" src="./prefix-caching.ts"></script>
</body>
</html>
142 changes: 142 additions & 0 deletions examples/prefix-caching/src/prefix-caching.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
import * as webllm from "@mlc-ai/web-llm";

const SYSTEM_PROMPT_PREFIX =
"You are a helpful assistant running in the user's browser, responsible for answering questions.";

function setLabel(id: string, text: string) {
const label = document.getElementById(id);
if (label == null) {
throw Error("Cannot find label " + id);
}
label.innerText = text;
}

async function testPrefix() {
const initProgressCallback = (report: webllm.InitProgressReport) => {
setLabel("init-label", report.text);
};

const selectedModel = "Llama-3.1-8B-Instruct-q4f32_1-MLC";
const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine(
selectedModel,
{
initProgressCallback: initProgressCallback,
logLevel: "INFO",
// Prefilling KV cache for efficiency
cachedPrefixes: [[{ role: "system", content: SYSTEM_PROMPT_PREFIX }]],
},
{
context_window_size: 2048,
},
);

const reply_using_prefix = await engine.chat.completions.create({
messages: [
{ role: "system", content: SYSTEM_PROMPT_PREFIX },
{ role: "user", content: "List three US states." },
],
// below configurations are all optional
n: 1,
temperature: 1.5,
max_tokens: 64,
logprobs: true,
top_logprobs: 2,
});
console.log(reply_using_prefix);
console.log(reply_using_prefix.usage);
}

async function testWithoutPrefix() {
const initProgressCallback = (report: webllm.InitProgressReport) => {
setLabel("init-label", report.text);
};

const selectedModel = "Llama-3.1-8B-Instruct-q4f32_1-MLC";
// Engine Initialization without cachedPrefixes
const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine(
selectedModel,
{
initProgressCallback: initProgressCallback,
logLevel: "INFO",
},
{
context_window_size: 2048,
},
);

const reply_without_prefix = await engine.chat.completions.create({
messages: [
{ role: "system", content: SYSTEM_PROMPT_PREFIX },
{ role: "user", content: "List three US states." },
],
// below configurations are all optional
n: 1,
temperature: 1.5,
max_tokens: 64,
logprobs: true,
top_logprobs: 2,
});
console.log(reply_without_prefix);
console.log(reply_without_prefix.usage);
}

async function testMultiRound() {
const initProgressCallback = (report: webllm.InitProgressReport) => {
setLabel("init-label", report.text);
};

const selectedModel = "Llama-3.1-8B-Instruct-q4f32_1-MLC";
const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine(
selectedModel,
{
initProgressCallback: initProgressCallback,
logLevel: "INFO",
cachedPrefixes: [[{ role: "system", content: SYSTEM_PROMPT_PREFIX }]], // Prefilling KV cache for efficiency
},
{
context_window_size: 2048,
},
);

// First Completion with cachedPrefixes
const reply0 = await engine.chat.completions.create({
messages: [
{ role: "system", content: SYSTEM_PROMPT_PREFIX },
{ role: "user", content: "List three US states." },
],
// below configurations are all optional
n: 1,
temperature: 1.5,
max_tokens: 64,
logprobs: true,
top_logprobs: 2,
});
console.log(reply0);
console.log(reply0.usage);

// Second Completion with cachedPrefixes
const reply1 = await engine.chat.completions.create({
messages: [
{ role: "system", content: SYSTEM_PROMPT_PREFIX },
{ role: "user", content: "Where is the US capital?" },
],
// below configurations are all optional
n: 1,
temperature: 1.5,
max_tokens: 64,
logprobs: true,
top_logprobs: 2,
});
console.log(reply1);
console.log(reply1.usage);
}

async function main() {
await testPrefix();

await testWithoutPrefix();

await testMultiRound();
}

main();
6 changes: 6 additions & 0 deletions src/config.ts
View file
Original file line number Diff line number Diff line change
@@ -9,6 +9,7 @@ import {
NonNegativeError,
RangeError,
} from "./error";
import { ChatCompletionMessageParam } from "./openai_api_protocols/chat_completion";

/**
* Conversation template config
@@ -105,15 +106,20 @@ export interface ChatOptions extends Partial<ChatConfig> {}
* appConfig: Configure the app, including the list of models and whether to use IndexedDB cache.
* initProgressCallback: A callback for showing the progress of loading the model.
* logitProcessorRegistry: A register for stateful logit processors, see `webllm.LogitProcessor`.
* cachedPrefixes: Specifies a system prompt (prefix) that will be prefilled when loading the engine
* to create their corresponding KV cache and store them for reuse. These cached kv pairs persist
* until the engine is reloaded.
*
* @note All fields are optional, and `logitProcessorRegistry` is only used for `MLCEngine` and not
* other `MLCEngine`s.
* @note cachedPrefixes is experimental. It may change in future versions.
*/
export interface MLCEngineConfig {
appConfig?: AppConfig;
initProgressCallback?: InitProgressCallback;
logitProcessorRegistry?: Map<string, LogitProcessor>;
logLevel?: LogLevel;
cachedPrefixes?: ChatCompletionMessageParam[][];
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's add docs to MLCEngineConfig, specifying the behavior of cachedPrefixes (e.g. will prefill when loading the engine to create the prefixes' KV, will only dispose these KV when reloading the engine). Perhaps we can also mark this as experimental to signify potential future API/behavior change

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also add an examples/cached_prefixes? Where we can demonstrate the prefill time difference between using cachedPrefixes and not using it. We should also test whether the behavior is expected in multi-turn conversation.

}

/**
12 changes: 12 additions & 0 deletions src/engine.ts
View file
Original file line number Diff line number Diff line change
@@ -131,6 +131,7 @@ export class MLCEngine implements MLCEngineInterface {
private logitProcessorRegistry?: Map<string, LogitProcessor>;
private initProgressCallback?: InitProgressCallback;
private appConfig: AppConfig;
private cachedPrefixes: ChatCompletionMessageParam[][];

// Signals and flags
private interruptSignal = false;
@@ -149,6 +150,7 @@ export class MLCEngine implements MLCEngineInterface {
this.setLogLevel(engineConfig?.logLevel || DefaultLogLevel);
this.setInitProgressCallback(engineConfig?.initProgressCallback);
this.setLogitProcessorRegistry(engineConfig?.logitProcessorRegistry);
this.cachedPrefixes = engineConfig?.cachedPrefixes || [];

this.chat = new API.Chat(this);
this.completions = new API.Completions(this);
@@ -392,6 +394,16 @@ export class MLCEngine implements MLCEngineInterface {
this.loadedModelIdToPipeline.set(modelId, newPipeline);
this.loadedModelIdToLock.set(modelId, new CustomLock());

// Call prefillConvSequence() if cachedPrefixes is specified
if (
newPipeline instanceof LLMChatPipeline &&
this.cachedPrefixes.length > 0
) {
for (let i = 0; i < this.cachedPrefixes.length; i++) {
await newPipeline.prefillConvSequence(this.cachedPrefixes[i]);
}
}

// Clean up
const tend = performance.now();
if (this.initProgressCallback !== undefined) {
Loading