Merge branch 'main' into leo/dgx-spark-integrations

bench/exo_bench.py

@@ -501,20 +501,28 @@ def main() -> int:
                                 all_rows.append(row)
 
                             if batch_results:
-                                agg_gen_tps = sum(
+                                valid_gen_tps = [
                                     x["stats"]["generation_tps"]
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
+                                ]
+                                agg_gen_tps = (
+                                    mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
+                                gen_tps = agg_gen_tps / concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
+                                    f"gen_tps={gen_tps:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        gen_tps = mean(
+                            x["stats"]["generation_tps"] / x["concurrency"]
+                            for x in runs
+                        )
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                         peak = mean(

dashboard/src/lib/components/ChatForm.svelte

@@ -1,9 +1,6 @@
 <script lang="ts">
   import {
     isLoading,
-    sendMessage,
-    generateImage,
-    editImage,
     editingImage,
     clearEditingImage,
     selectedChatModel,
@@ -28,7 +25,7 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
+    onAutoSend: (
       content: string,
       files?: {
         id: string;
@@ -216,49 +213,10 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
-    // Use image editing if in edit mode
-    if (isEditMode && currentEditingImage && content) {
-      editImage(content, currentEditingImage.imageDataUrl);
-    }
-    // If user attached an image with an ImageToImage model, use edit endpoint
-    else if (
-      currentModel &&
-      modelSupportsImageEditing(currentModel) &&
-      files.length > 0 &&
-      content
-    ) {
-      // Use the first attached image for editing
-      const imageFile = files[0];
-      if (imageFile.preview) {
-        editImage(content, imageFile.preview);
-      }
-    } else if (
-      currentModel &&
-      modelSupportsTextToImage(currentModel) &&
-      content
-    ) {
-      // Use image generation for text-to-image models
-      generateImage(content);
-    } else {
-      sendMessage(
-        content,
-        files,
-        modelSupportsThinking() ? thinkingEnabled : null,
-      );
-    }
-
+    // Parent controls all send logic (including image routing,
+    // launching non-running models before sending, etc.)
+    onAutoSend(content, files);
     onSend?.();
-
-    // Refocus the textarea after sending
     setTimeout(() => textareaRef?.focus(), 10);
   }
 

dashboard/src/lib/stores/app.svelte.ts

@@ -1587,6 +1587,7 @@ class AppStore {
     // Remove messages after user message (including the user message for image requests
     // since generateImage/editImage will re-add it)
     this.messages = this.messages.slice(0, lastUserIndex);
+    this.updateActiveConversation();
 
     switch (requestType) {
       case "image-generation":

dashboard/src/routes/+page.svelte

@@ -42,6 +42,9 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    generateImage,
+    editImage,
+    editingImage,
     messages,
     debugMode,
     toggleDebugMode,
@@ -838,6 +841,52 @@
     if (!model?.tasks) return false;
     return model.tasks.includes("ImageToImage");
   }
+
+  // Route a message to the correct endpoint based on model capabilities.
+  // Image models go to generateImage/editImage; text models go to sendMessage.
+  function routeMessage(
+    content: string,
+    files?: {
+      id: string;
+      name: string;
+      type: string;
+      textContent?: string;
+      preview?: string;
+    }[],
+  ) {
+    const model = selectedChatModel();
+    if (!model) {
+      sendMessage(content, files, null);
+      return;
+    }
+
+    const currentEditImage = editingImage();
+
+    // Image editing mode (explicit edit or attached image with ImageToImage model)
+    if (currentEditImage && content && modelSupportsImageEditing(model)) {
+      editImage(content, currentEditImage.imageDataUrl);
+      return;
+    }
+    if (
+      modelSupportsImageEditing(model) &&
+      files?.length &&
+      files[0].preview &&
+      content
+    ) {
+      editImage(content, files[0].preview);
+      return;
+    }
+
+    // Text-to-image generation
+    if (modelSupportsImageGeneration(model) && content) {
+      generateImage(content);
+      return;
+    }
+
+    // Default: text chat
+    sendMessage(content, files, null);
+  }
+
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
   type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
 
@@ -2858,7 +2907,7 @@
         // Running model is same or better tier — use it directly
         setSelectedChatModel(bestRunning.id);
         if (!chatStarted) createConversation();
-        sendMessage(content, files);
+        routeMessage(content, files);
         return;
       }
     }
@@ -2875,7 +2924,7 @@
     if (hasRunningInstance(autoModel.id)) {
       setSelectedChatModel(autoModel.id);
       if (!chatStarted) createConversation();
-      sendMessage(content, files);
+      routeMessage(content, files);
       return;
     }
 
@@ -3028,7 +3077,7 @@
       if (pendingAutoMessage) {
         const msg = pendingAutoMessage;
         pendingAutoMessage = null;
-        sendMessage(msg.content, msg.files);
+        routeMessage(msg.content, msg.files);
       }
       return;
     }
@@ -3107,7 +3156,7 @@
     // Model is selected and running — send directly
     if (model && hasRunningInstance(model)) {
       chatLaunchState = "ready";
-      sendMessage(content, files, null);
+      routeMessage(content, files);
       return;
     }
 

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -58,8 +58,6 @@ class _EngineTask:
     matched_index: int | None
     cache_snapshots: list[CacheSnapshot] | None
     detokenizer: StreamingDetokenizer
-    prefill_tps: float = 0.0
-    prompt_token_count: int = 0
     on_generation_token: Callable[[], None] | None = None
     generated_text_parts: list[str] = field(default_factory=list)
     potential_stop_sequence_text: str = ""
@@ -67,6 +65,7 @@ class _EngineTask:
     generation_start_time: float = 0.0
     in_thinking: bool = False
     reasoning_tokens: int = 0
+    prefill_tps: float = 0.0
 
 
 @dataclass(eq=False)
@@ -141,7 +140,7 @@ class ExoBatchGenerator:
             top_k=task_params.top_k if task_params.top_k is not None else 0,
         )
 
-        prefill_tps, prefill_tokens, cache_snapshots = prefill(
+        _prefill_tps, _prefill_tokens, cache_snapshots = prefill(
             self.model,
             self.tokenizer,
             sampler,
@@ -208,11 +207,10 @@ class ExoBatchGenerator:
             prefix_hit_length=prefix_hit_length,
             matched_index=matched_index,
             cache_snapshots=cache_snapshots or None,
-            prefill_tps=prefill_tps,
-            prompt_token_count=prefill_tokens,
             detokenizer=self.tokenizer.detokenizer,
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
+            prefill_tps=_prefill_tps,
         )
 
         return uid
@@ -289,15 +287,22 @@ class ExoBatchGenerator:
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                generation_elapsed = time.perf_counter() - state.generation_start_time
-                generation_tps = (
-                    state.completion_tokens / generation_elapsed
-                    if generation_elapsed > 0
-                    else 0.0
-                )
+                try:
+                    mlx_stats = self._exo_gen.stats()
+                    generation_tps = mlx_stats.generation_tps
+                except ZeroDivisionError:
+                    generation_elapsed = (
+                        time.perf_counter() - state.generation_start_time
+                    )
+                    generation_tps = (
+                        state.completion_tokens / generation_elapsed
+                        if generation_elapsed > 0
+                        else 0.0
+                    )
+
                 stats = GenerationStats(
                     prompt_tps=state.prefill_tps,
-                    generation_tps=float(generation_tps),
+                    generation_tps=generation_tps,
                     prompt_tokens=len(state.all_prompt_tokens),
                     generation_tokens=state.completion_tokens,
                     peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),