[NV] [PR1] feat(synthetic-al): generic synthetic-acceptance injection scaffold

runners/inject_synthetic_acceptance.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Inject synthetic acceptance parameters into an srt-slurm recipe (generic driver).
+
+This is the framework-agnostic half of the synthetic-acceptance mechanism. It
+decides *whether* to inject (the ``SYNTHETIC_ACCEPTANCE`` flag) and *what* mean
+acceptance length to inject, then delegates the actual recipe rewrite to a
+framework-specific backend (see ``runners/synthetic_injectors/``).
+
+The script is a no-op (exit 0, file untouched) when:
+  - SYNTHETIC_ACCEPTANCE is unset/false,
+so existing callers that do not opt in get exactly the previous behavior. When
+enabled it requires a backend registered for the given framework; the vLLM
+backend is added in a follow-up framework-support change.
+
+Environment variables:
+  SYNTHETIC_ACCEPTANCE         "true" to enable (default: "false")
+  SYNTHETIC_ACCEPTANCE_LENGTH  target mean acceptance length; if unset, it is
+                               auto-resolved from the reference AL YAML using
+                               MODEL_PREFIX (+ NUM_SPEC_TOKENS / THINKING_MODE)
+  NUM_SPEC_TOKENS              number of speculative tokens (for auto-lookup;
+                               falls back to the value parsed from the recipe)
+  MODEL_PREFIX                 model prefix key in the reference YAML (e.g. "dsv4")
+  THINKING_MODE                "thinking_on" / "thinking_off" — only used when the
+                               reference YAML is in the thinking matrix form
+                               (default: "thinking_on")
+  FRAMEWORK                    framework key selecting the backend (e.g.
+                               "dynamo-vllm"); may also be passed as argv[2].
+
+Usage (from a runner; use an absolute path since runners cd into the srt-slurm
+clone before invoking this):
+  python3 "$GITHUB_WORKSPACE/runners/inject_synthetic_acceptance.py" "${CONFIG_FILE%%:*}" "$FRAMEWORK"
+"""
+
+import os
+import sys
+
+from synthetic_injectors import get_injector
+
+# MODEL_PREFIX -> top-level key in speedbench-reference-al.yaml.
+MODEL_PREFIX_TO_YAML_KEY = {
+    "dsv4": "deepseek-v4-pro",
+    "dsr1": "deepseek-r1",
+}
+
+
+def _log(msg):
+    print(f"[Synthetic AR] {msg}")
+
+
+def _yaml_key(model_prefix):
+    return MODEL_PREFIX_TO_YAML_KEY.get(model_prefix, model_prefix)
+
+
+def _lookup_al(model_block, num_spec_tokens):
+    """Resolve AL for num_spec_tokens from either reference-YAML shape.
+
+    Flat list form:   [ {1: 1.90}, {2: 2.60}, ... ]
+    Thinking matrix:  { thinking_on: {1: ...}, thinking_off: {1: ...} }
+    """
+    # Flat list form (each item is a single-key {level: al} mapping).
+    if isinstance(model_block, list):
+        for item in model_block:
+            if num_spec_tokens in item:
+                return item[num_spec_tokens]
+        return None
+
+    if isinstance(model_block, dict):
+        # Thinking matrix form: pick the requested mode, then index by level.
+        if any(str(k).startswith("thinking") for k in model_block):
+            mode = os.environ.get("THINKING_MODE", "thinking_on").strip() or "thinking_on"
+            mode_block = model_block.get(mode)
+            if mode_block is None:
+                sys.exit(
+                    f"ERROR: THINKING_MODE='{mode}' not found in reference YAML "
+                    f"(available: {sorted(model_block)})"
+                )
+            return mode_block.get(num_spec_tokens)
+        # Plain {level: al} mapping.
+        return model_block.get(num_spec_tokens)
+
+    return None
+
+
+def _resolve_al(config_text, injector, ref_yaml):
+    explicit = os.environ.get("SYNTHETIC_ACCEPTANCE_LENGTH", "").strip()
+    if explicit:
+        return float(explicit)
+
+    if not os.path.isfile(ref_yaml):
+        sys.exit(
+            "ERROR: SYNTHETIC_ACCEPTANCE_LENGTH not set and reference YAML not "
+            f"found: {ref_yaml}"
+        )
+
+    import yaml  # local import: only needed on the auto-lookup path
+
+    with open(ref_yaml) as f:
+        data = yaml.safe_load(f)
+
+    key = _yaml_key(os.environ.get("MODEL_PREFIX", ""))
+    model_block = data.get(key)
+    if model_block is None:
+        sys.exit(f'ERROR: model key "{key}" not found in {ref_yaml}')
+
+    nst_env = os.environ.get("NUM_SPEC_TOKENS", "").strip()
+    if nst_env:
+        num_spec_tokens = int(nst_env)
+    else:
+        num_spec_tokens = injector.spec_tokens_from_recipe(config_text) or 2
+
+    al = _lookup_al(model_block, num_spec_tokens)
+    if al is None:
+        sys.exit(f"ERROR: num_spec_tokens={num_spec_tokens} not found for {key} in {ref_yaml}")
+
+    _log(
+        f"Auto-resolved AL={al} from {ref_yaml} "
+        f"(model={key}, num_spec_tokens={num_spec_tokens})"
+    )
+    return float(al)
+
+
+def inject(config_file, framework):
+    if os.environ.get("SYNTHETIC_ACCEPTANCE", "false") != "true":
+        return 0
+
+    injector = get_injector(framework)
+    if injector is None:
+        sys.exit(
+            "ERROR: SYNTHETIC_ACCEPTANCE=true but no synthetic-acceptance "
+            f"injector is registered for FRAMEWORK='{framework}'"
+        )
+
+    with open(config_file) as f:
+        content = f.read()
+
+    al = _resolve_al(
+        content,
+        injector,
+        os.path.join(os.path.dirname(__file__), "..", "benchmarks", "speedbench-reference-al.yaml"),
+    )
+
+    _log(f"Injecting synthetic acceptance (length={al}) into {config_file}")
+
+    new_content, count = injector.rewrite(content, al, _log)
+
+    if count == 0:
+        _log("WARNING: No speculative-config entries found to modify; leaving recipe unchanged")
+        return 0
+
+    with open(config_file, "w") as f:
+        f.write(new_content)
+    _log(f"Modified {count} speculative-config entries")
+    return 0
+
+
+def main(argv):
+    if len(argv) not in (2, 3):
+        sys.exit("Usage: inject_synthetic_acceptance.py CONFIG_FILE [FRAMEWORK]")
+    framework = argv[2] if len(argv) == 3 else os.environ.get("FRAMEWORK", "")
+    return inject(argv[1], framework)
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv))

runners/launch_gb200-nv.sh

@@ -52,11 +52,15 @@ elif [[ $FRAMEWORK == "dynamo-vllm" ]]; then
         export MODEL_PATH="/mnt/lustre01/models/kimi-k2.5-nvfp4"
         export SRT_SLURM_MODEL_PREFIX="kimi-k2.5-nvfp4"
     elif [[ $MODEL_PREFIX == "dsv4" && $PRECISION == "fp4" ]]; then
-        # The FP4 checkpoint is staged on compute-visible Lustre. The former
-        # /mnt/numa1 path is no longer present on watchtower compute nodes;
-        # the lowercase Lustre sibling is the FP8 checkpoint, so keep the
-        # NVFP4 path explicit here.
-        export MODEL_PATH="/mnt/lustre01/models/DeepSeek-V4-Pro-NVFP4/"
+        # FP4 checkpoint on compute-visible Lustre (the /mnt/numa1 path is gone
+        # on watchtower compute nodes). Use the base DeepSeek-V4-Pro checkpoint,
+        # NOT the -NVFP4 re-quant: the recipe's served identity is plain
+        # deepseek-ai/DeepSeek-V4-Pro and the pinned v0.20.1 container's
+        # deepseek_v4 loader doesn't define the NVFP4 export's extra quant
+        # params (e.g. ffn.experts.w13_input_scale), which KeyErrors at load.
+        # The lowercase Lustre sibling is the FP8 checkpoint, so name the
+        # CamelCase FP4 path explicitly (Linux is case-sensitive).
+        export MODEL_PATH="/mnt/lustre01/models/DeepSeek-V4-Pro"
         export SRT_SLURM_MODEL_PREFIX="deepseek-v4-pro"
     elif [[ $MODEL_PREFIX == "minimaxm2.5" && $PRECISION == "fp4" ]]; then
         export MODEL_PATH="/mnt/lustre01/models/MiniMax-M2.5-NVFP4"
@@ -84,8 +88,12 @@ NGINX_IMAGE="nginx:1.27.4"
 uses_watchtower_shared_fs() {
     case "$MODEL_PREFIX" in
         minimaxm2.5|minimaxm3|kimik2.5) return 0 ;;
-        *) return 1 ;;
     esac
+    # dsv4 multinode runs only under dynamo-vllm on watchtower, which likewise
+    # needs the srt-slurm workspace/outputs on a compute-visible shared FS
+    # (the runner home is not cross-mounted to compute nodes).
+    [[ "$FRAMEWORK" == "dynamo-vllm" && "$MODEL_PREFIX" == "dsv4" ]] && return 0
+    return 1
 }
 
 # === Cluster diagnostic probe for watchtower-hosted sweeps ===
@@ -462,6 +470,11 @@ fi
 # Keep the Slurm job name aligned with the GitHub runner name.
 sed -i "s/^name:.*/name: \"${RUNNER_NAME}\"/" "$CONFIG_PATH"
 
+# Optionally inject synthetic acceptance into the recipe's speculative-config
+# when SYNTHETIC_ACCEPTANCE=true (no-op otherwise). Must run after the name
+# override and before srtctl apply so the rendered job picks it up.
+python3 "$GITHUB_WORKSPACE/runners/inject_synthetic_acceptance.py" "$CONFIG_PATH" "$FRAMEWORK"
+
 # Don't leak the login-node venv to the compute-node orchestrator. sbatch's
 # default --export=ALL propagates VIRTUAL_ENV (set by `source
 # .venv/bin/activate` above) into job_script_minimal.j2, whose

runners/synthetic_injectors/__init__.py

@@ -0,0 +1,34 @@
+"""Framework-specific synthetic-acceptance injectors.
+
+The generic driver (``inject_synthetic_acceptance.py``) resolves *whether* to
+inject and *what* acceptance length to use, then hands the recipe text to a
+framework-specific injector registered here. This keeps the driver
+framework-agnostic: adding a new backend (e.g. sglang, trtllm) means dropping a
+module in this package and registering it, without touching the driver.
+
+A backend is any object exposing:
+
+    rewrite(content: str, al: float, log) -> tuple[str, int]
+        Return the recipe text with the synthetic acceptance parameters applied
+        and the number of entries modified. ``log`` is a ``callable(str)`` used
+        for human-readable progress output.
+
+    spec_tokens_from_recipe(content: str) -> int | None
+        Best-effort extraction of the speculative-token count from the recipe,
+        used only for reference-AL auto-lookup. May return ``None``.
+"""
+
+# framework value passed by the runner (e.g. "dynamo-vllm") -> backend module.
+# Populated by backend modules registering themselves (see e.g. vllm.py, added
+# in the framework-support PR). Empty here: the generic layer alone injects
+# nothing, so merging it changes no behavior.
+_INJECTORS = {}
+
+
+def register(framework, backend):
+    _INJECTORS[framework] = backend
+
+
+def get_injector(framework):
+    """Return the backend for ``framework`` (as passed by the runner), or None."""
+    return _INJECTORS.get(framework)