Restore custom live prompts on benchmark config

engines/live_request.py

@@ -115,6 +115,9 @@ class LiveRuntimeSettings:
     learned_page_selector_prompt_variant: str | None
     learned_page_selector_path: str | None
     compact_task_name: str | None
+    prompt_text: str
+    use_exact_length_prompt: bool
+    is_custom_prompt: bool
 
 
 def canonicalize_benchmark_payload(request: dict[str, Any]) -> dict[str, Any]:
@@ -159,7 +162,6 @@ def resolve_live_runtime_settings(
     decode_steps: int,
     max_live_context: int,
 ) -> LiveRuntimeSettings:
-    del decode_steps
     request = canonicalize_benchmark_payload(request)
     model_key = str(request.get("model") or "")
     if model_key not in MODEL_BY_KEY:
@@ -191,16 +193,10 @@ def resolve_live_runtime_settings(
         raise ValueError("Only the Qwen benchmark-backed live lane is supported in this Space.")
 
     custom_prompt = str(request.get("custom_prompt") or "").strip()
-    if custom_prompt:
-        raise ValueError(
-            "Live Qwen compare now replays the benchmark rows directly. Custom prompts are disabled until a "
-            "benchmark-faithful custom lane exists."
-        )
-
     preset_key = str(request.get("preset") or "")
     if preset_key not in PRESET_BY_KEY:
         raise ValueError("Live Qwen compare requires a benchmark-backed preset selection.")
-    if preset_key == "longbench_mini":
+    if not custom_prompt and preset_key == "longbench_mini":
         raise ValueError(
             "Live LongBench replay is not wired in this Space yet. Preset-backed compare remains the valid benchmark view."
         )
@@ -215,10 +211,15 @@ def resolve_live_runtime_settings(
         )
 
     compact_task_name = None
-    benchmark_decode_steps = 8
-    if preset_key == "compact_task":
+    prompt_text = custom_prompt
+    use_exact_length_prompt = False
+    benchmark_decode_steps = int(decode_steps or DEFAULT_LIVE_DECODE_STEPS)
+    if not custom_prompt and preset_key == "compact_task":
         compact_task_name = _compact_task_name(model_key, context_length)
         benchmark_decode_steps = COMPACT_TASK_DECODE_STEPS[compact_task_name]
+    elif not custom_prompt and preset_key == "backend_truth":
+        prompt_text = "Cache locality matters for fast decoding."
+        use_exact_length_prompt = True
 
     return LiveRuntimeSettings(
         model_key=model.key,
@@ -241,6 +242,9 @@ def resolve_live_runtime_settings(
         learned_page_selector_prompt_variant=variant_config["learned_page_selector_prompt_variant"],
         learned_page_selector_path=selector_artifact_path,
         compact_task_name=compact_task_name,
+        prompt_text=prompt_text,
+        use_exact_length_prompt=use_exact_length_prompt,
+        is_custom_prompt=bool(custom_prompt),
     )
 
 

scripts/space_dotcache_runner.py

@@ -12,7 +12,7 @@ if str(REPO_ROOT) not in sys.path:
 
 from benchmarks.bench_qwen35_attention_subset_dotcache_serving import _build_dotcache_config  # noqa: E402
 from dotcache.integrations.qwen35 import Qwen35AttentionSubsetDotCacheHarness  # noqa: E402
-from engines.live_request import resolve_live_runtime_settings  # noqa: E402
+from engines.live_request import DEFAULT_LIVE_DECODE_STEPS, resolve_live_runtime_settings  # noqa: E402
 from scripts.run_qwen35_task_selector_compare import _set_selector_prompt_tags, _task_specs  # noqa: E402
 from scripts.space_runner_common import (  # noqa: E402
     configure_model_cache_env,
@@ -176,7 +176,7 @@ def main() -> int:
     request = load_request_from_stdin()
     settings = resolve_live_runtime_settings(
         request,
-        decode_steps=0,
+        decode_steps=int(os.getenv("DOTCACHE_SPACE_LIVE_DECODE_STEPS", str(DEFAULT_LIVE_DECODE_STEPS))),
         max_live_context=int(os.getenv("DOTCACHE_SPACE_MAX_LIVE_CONTEXT", "4096")),
     )
 
@@ -197,7 +197,24 @@ def main() -> int:
         weight_quantization=os.getenv("DOTCACHE_SPACE_WEIGHT_QUANTIZATION", "none"),
     )
 
-    if settings.benchmark_suite == "compact_task":
+    if settings.is_custom_prompt:
+        input_ids, attention_mask = harness.tokenize_prompt(settings.prompt_text)
+        prompt_length = int(input_ids.shape[1])
+        if prompt_length > settings.context_length:
+            raise ValueError(
+                f"Custom prompt tokenized to {prompt_length} tokens, which exceeds the selected context limit "
+                f"of {settings.context_length}. Increase Context length or shorten the prompt."
+            )
+        decode_steps = int(settings.decode_steps)
+        record = harness.run_attention_subset_dotcache_serving(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decode_steps=decode_steps,
+            profile_backend=True,
+        )
+        generated_ids = list(record.get("dotcache_generated_ids") or [])
+        text = decode_generated_text(harness.tokenizer, generated_ids, limit=decode_steps)
+    elif settings.benchmark_suite == "compact_task":
         task_spec = _task_prompt_inputs(harness, settings)
         input_ids = task_spec["input_ids"]
         attention_mask = task_spec["attention_mask"]
@@ -212,10 +229,10 @@ def main() -> int:
         decoded_text = decode_generated_text(harness.tokenizer, generated_ids, limit=decode_steps)
         scored = dict(task_spec["score_fn"](decoded_text))
         text = str(scored.get("task_generated_text_cleaned") or scored.get("task_generated_value") or decoded_text)
-    elif settings.benchmark_suite == "backend_truth":
+    elif settings.use_exact_length_prompt:
         input_ids, attention_mask = _build_exact_length_inputs(
             harness,
-            prompt_unit=BACKEND_TRUTH_PROMPT_UNIT,
+            prompt_unit=settings.prompt_text or BACKEND_TRUTH_PROMPT_UNIT,
             prompt_length=settings.context_length,
         )
         decode_steps = int(settings.decode_steps)

space_app.py

@@ -51,7 +51,7 @@ MAX_UI_LOG_LINES = 80
 MAX_UI_LOG_CHARS = 12000
 MAX_UI_TRACE_ITEMS = 8
 MAX_ARENA_MATCHES = 64
-CUSTOM_LIVE_PROMPTS_ENABLED = False
+CUSTOM_LIVE_PROMPTS_ENABLED = True
 LIVE_EXAMPLE_PROMPTS = {
     "retrieval": (
         "Context:\n"
@@ -1000,10 +1000,9 @@ def _build_request(
     shortlist_policy: str,
     compare_against_dense: bool,
 ) -> DemoRequest:
-    normalized_custom_prompt = custom_prompt.strip()
     return DemoRequest(
         model=model,
-        preset=None if normalized_custom_prompt else (preset or None),
+        preset=preset or None,
         custom_prompt=custom_prompt,
         context_length=int(context_length),
         mode=mode,
@@ -1096,7 +1095,7 @@ def _live_context_guard_copy(mode: str) -> str:
     if mode == "custom":
         return (
             "<div class='live-context-note active'>"
-            "<strong>⚡ Live mode (ZeroGPU)</strong><span>Your prompt runs live on shared GPU. Context stays limited to <code>512</code> or <code>1024</code> for responsiveness. Expect a short queue delay, and try prompts that depend on earlier text to see differences.</span>"
+            "<strong>⚡ Live mode (ZeroGPU)</strong><span>Your prompt runs on the real Space system using the selected preset's benchmark configuration. Use this to sanity-check that live behavior stays in the same ballpark as the cached paper row.</span>"
             "</div>"
         )
     if mode == "benchmark":
@@ -1112,23 +1111,22 @@ def _live_context_guard_copy(mode: str) -> str:
     )
 
 
+def _available_live_context_choices() -> list[int]:
+    max_live_context = int(os.getenv("DOTCACHE_SPACE_MAX_LIVE_CONTEXT", "4096"))
+    choices = [value for value in FULL_CONTEXT_CHOICES if int(value) <= max_live_context]
+    return choices or [min(max_live_context, 1024)]
+
+
 def _guard_live_context(prompt_mode: str, custom_prompt: str, current_context_length: int) -> tuple[Any, str]:
     if _is_live_prompt_mode(prompt_mode):
-        if not CUSTOM_LIVE_PROMPTS_ENABLED:
-            guarded_value = int(current_context_length)
-            if guarded_value not in BENCHMARK_LIVE_CONTEXT_CHOICES:
-                guarded_value = 2048
-            return (
-                gr.update(choices=BENCHMARK_LIVE_CONTEXT_CHOICES, value=guarded_value),
-                _live_context_guard_copy("benchmark"),
-            )
         has_custom_prompt = bool(custom_prompt.strip())
         if has_custom_prompt:
+            available_choices = _available_live_context_choices()
             guarded_value = int(current_context_length)
-            if guarded_value not in LIVE_SAFE_CONTEXT_CHOICES:
-                guarded_value = 1024
+            if guarded_value not in available_choices:
+                guarded_value = available_choices[min(len(available_choices) - 1, 1)]
             return (
-                gr.update(choices=LIVE_SAFE_CONTEXT_CHOICES, value=guarded_value),
+                gr.update(choices=available_choices, value=guarded_value),
                 _live_context_guard_copy("custom"),
             )
         guarded_value = int(current_context_length)
@@ -1153,8 +1151,6 @@ def _is_live_prompt_mode(prompt_mode: str) -> bool:
 
 
 def _effective_custom_prompt(prompt_mode: str, custom_prompt: str) -> str:
-    if not CUSTOM_LIVE_PROMPTS_ENABLED:
-        return ""
     if _is_live_prompt_mode(prompt_mode):
         return custom_prompt
     return ""