diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte
index 1b1d2d07..c24548a5 100644
--- a/dashboard/src/lib/components/ChatMessages.svelte
+++ b/dashboard/src/lib/components/ChatMessages.svelte
@@ -12,6 +12,7 @@
} from "$lib/stores/app.svelte";
import type { MessageAttachment } from "$lib/stores/app.svelte";
import MarkdownContent from "./MarkdownContent.svelte";
+ import PrefillProgressBar from "./PrefillProgressBar.svelte";
import TokenHeatmap from "./TokenHeatmap.svelte";
import PrefillProgressBar from "./PrefillProgressBar.svelte";
import ImageLightbox from "./ImageLightbox.svelte";
@@ -625,7 +626,9 @@
- {#if loading && !message.content}
+ {#if loading && !message.content && prefill}
+
+ {:else if loading && !message.content}
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index d13e7f2b..3c2c9ae5 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -276,6 +276,8 @@ export interface TokenData {
export interface PrefillProgress {
processed: number;
total: number;
+ /** Timestamp (performance.now()) when prefill started. */
+ startedAt: number;
}
export interface Message {
@@ -520,12 +522,12 @@ class AppStore {
messages = $state([]);
currentResponse = $state("");
isLoading = $state(false);
+ prefillProgress = $state(null);
// Performance metrics
ttftMs = $state(null); // Time to first token in ms
tps = $state(null); // Tokens per second
totalTokens = $state(0); // Total tokens in current response
- prefillProgress = $state(null);
// Abort controller for stopping generation
private currentAbortController: AbortController | null = null;
@@ -2018,6 +2020,7 @@ class AppStore {
): Promise {
const decoder = new TextDecoder();
let buffer = "";
+ let currentEventType = "";
while (true) {
const { done, value } = await reader.read();
@@ -2033,7 +2036,15 @@ class AppStore {
for (const line of lines) {
const trimmed = line.trim();
- if (!trimmed) continue;
+ if (!trimmed) {
+ currentEventType = "";
+ continue;
+ }
+
+ if (trimmed.startsWith("event: ")) {
+ currentEventType = trimmed.slice(7);
+ continue;
+ }
// Handle SSE comments (": key json") for prefill progress etc.
if (trimmed.startsWith(": ") && onEvent) {
@@ -2055,14 +2066,22 @@ class AppStore {
if (trimmed.startsWith("data: ")) {
const data = trimmed.slice(6);
- if (data === "[DONE]") continue;
+ if (data === "[DONE]") {
+ currentEventType = "";
+ continue;
+ }
try {
- const parsed = JSON.parse(data) as T;
- onChunk(parsed);
+ const parsed = JSON.parse(data);
+ if (currentEventType && onEvent?.[currentEventType]) {
+ onEvent[currentEventType](parsed);
+ } else {
+ onChunk(parsed as T);
+ }
} catch {
// Skip malformed JSON
}
+ currentEventType = "";
}
}
}
@@ -2163,6 +2182,7 @@ class AppStore {
this.isLoading = true;
this.currentResponse = "";
+ this.prefillProgress = null;
this.ttftMs = null;
this.tps = null;
this.totalTokens = 0;
@@ -2367,6 +2387,11 @@ class AppStore {
}
if (tokenContent) {
+ // Clear prefill progress once tokens start arriving
+ if (this.prefillProgress !== null) {
+ this.prefillProgress = null;
+ }
+
// Track first token for TTFT
if (firstTokenTime === null) {
firstTokenTime = performance.now();
@@ -2420,6 +2445,7 @@ class AppStore {
this.prefillProgress = {
processed: inner.processed_tokens,
total: inner.total_tokens,
+ startedAt: this.prefillProgress?.startedAt ?? performance.now(),
};
},
},
@@ -2474,6 +2500,7 @@ class AppStore {
this.prefillProgress = null;
this.isLoading = false;
this.currentResponse = "";
+ this.prefillProgress = null;
this.saveConversationsToStorage();
}
}
@@ -3106,10 +3133,10 @@ export const hasStartedChat = () => appStore.hasStartedChat;
export const messages = () => appStore.messages;
export const currentResponse = () => appStore.currentResponse;
export const isLoading = () => appStore.isLoading;
+export const prefillProgress = () => appStore.prefillProgress;
export const ttftMs = () => appStore.ttftMs;
export const tps = () => appStore.tps;
export const totalTokens = () => appStore.totalTokens;
-export const prefillProgress = () => appStore.prefillProgress;
export const topologyData = () => appStore.topologyData;
export const instances = () => appStore.instances;
export const runners = () => appStore.runners;
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index 6e926205..bbb762e6 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -90,6 +90,8 @@ def prefill(
)
if has_ssm:
snapshots.append(snapshot_ssm_states(cache))
+ if on_prefill_progress is not None:
+ on_prefill_progress(processed, total)
if on_prefill_progress is not None:
on_prefill_progress(processed, total)