From 1b4ee4d3fc003701f63ed19fe6f51219f7caf700 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 12 Jun 2026 06:33:58 +0000
Subject: [PATCH] studio: surface DiffusionGemma generation stats and auto-size
 the context

- backend: launch the visual decoder with auto-sized context (--maxtok 0) and read
  the resolved budget back into the reported context length
- frontend: show generation time, speed, tokens, denoising steps and blocks in the
  message timing tooltip for diffusion models
---
 studio/backend/core/inference/llama_cpp.py    | 29 +++++++++++++++----
 .../assistant-ui/message-timing.tsx           | 24 +++++++++++++++
 .../src/features/chat/api/chat-adapter.ts     |  4 +++
 3 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index fd2182bd91..74a7afd49a 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -2328,9 +2328,11 @@ class LlamaCppBackend:
 
         self._kill_process()
         self._port = self._find_free_port()
-        # The whole [prompt | 256-canvas] must fit one ubatch; the model's native
-        # context (262144) is irrelevant here, so cap to a sane per-turn budget.
-        maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 16384) else 8192
+        # The whole [prompt | 256-canvas] must fit one non-causal ubatch. Default to auto-size (0): the
+        # visual server probes the largest context that actually fits this GPU's VRAM (capped at the
+        # model's training context), which is far better than the old fixed 8192. Honor an explicit,
+        # in-range user n_ctx as an override.
+        maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0
         gpu = os.environ.get("DG_GPU", "0")
 
         cmd = list(shim_cmd) + [
@@ -2406,8 +2408,9 @@ class LlamaCppBackend:
                 self._hf_variant = None
         else:
             self._hf_variant = None
-        self._effective_context_length = maxtok
-        self._max_context_length = self._context_length or maxtok
+        # Provisional until the server reports the budget it resolved (auto-size picks it from VRAM).
+        self._effective_context_length = maxtok or self._context_length
+        self._max_context_length = self._context_length or maxtok or None
 
         healthy = self._wait_for_health(timeout = 600.0)
         if healthy:
@@ -2416,6 +2419,22 @@ class LlamaCppBackend:
             if extra_args is not None:
                 self._extra_args = list(extra_args)
                 self._extra_args_source = (model_identifier, hf_variant)
+            # The visual server logs "... MAXTOK=<N> ..." with the per-turn context budget it actually
+            # resolved (auto-sized to fit VRAM when launched with --maxtok 0). Read it back so the UI
+            # context bar reflects the real budget rather than the requested value.
+            chosen = maxtok
+            try:
+                import re as _re
+                for _ln in reversed(self._stdout_lines):
+                    _m = _re.search(r"MAXTOK=(\d+)", _ln)
+                    if _m:
+                        chosen = int(_m.group(1))
+                        break
+            except Exception:
+                pass
+            if chosen and chosen > 0:
+                self._effective_context_length = chosen
+                self._max_context_length = chosen
             self._requested_n_ctx = int(n_ctx)
         else:
             self._healthy = False
diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx
index 447fd763e3..e175bb6426 100644
--- a/studio/frontend/src/components/assistant-ui/message-timing.tsx
+++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx
@@ -135,6 +135,30 @@ export const MessageTiming: FC<{
                   </span>
                 </div>
               )}
+              {timing.firstTokenTime !== undefined && (
+                <div className="flex items-center justify-between gap-4">
+                  <span className="text-muted-foreground">First token</span>
+                  <span className="font-mono tabular-nums">
+                    {formatTimingMs(timing.firstTokenTime)}
+                  </span>
+                </div>
+              )}
+              {st?.diffusion_steps != null && (
+                <div className="flex items-center justify-between gap-4">
+                  <span className="text-muted-foreground">Denoising steps</span>
+                  <span className="font-mono tabular-nums">
+                    {formatNumber(st.diffusion_steps)}
+                  </span>
+                </div>
+              )}
+              {st?.diffusion_blocks != null && (
+                <div className="flex items-center justify-between gap-4">
+                  <span className="text-muted-foreground">Blocks</span>
+                  <span className="font-mono tabular-nums">
+                    {formatNumber(st.diffusion_blocks)}
+                  </span>
+                </div>
+              )}
               {cacheHits > 0 && (
                 <div className="flex items-center justify-between gap-4">
                   <span className="text-muted-foreground">Cache hits</span>
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index d944c7ab74..d25417af86 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -113,6 +113,10 @@ interface ServerTimings {
   predicted_ms: number;
   predicted_per_token_ms: number;
   predicted_per_second: number;
+  // DiffusionGemma-only extras (present when serving a diffusion model; ignored otherwise).
+  diffusion_blocks?: number;
+  diffusion_steps?: number;
+  diffusion_canvas?: number;
 }
 
 type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];