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
This commit is contained in:
parent
fd5f0a1331
commit
1b4ee4d3fc
3 changed files with 52 additions and 5 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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"];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue