[OMNIML-4788] specdec_bench/Qwen3.5-4B: throughput_32k benchmark + S3 upload step

examples/specdec_bench/run.py

@@ -35,6 +35,21 @@
     "AUTO_DEPLOY": models.AutoDeployModel,
     "SPECBENCH_MEDUSA": models.SpecBenchMedusaModel,
 }
+
+# Translation table for --max_seq_len. Each engine spells the same
+# concept (max input + output sequence the engine should reserve)
+# differently:
+#   VLLM   → max_model_len   (AsyncEngineArgs)
+#   TRTLLM → max_seq_len     (LLM(...))
+#   SGLANG → context_length  (sgl.Engine)
+# Mapping applied in run_simple() so cell YAMLs use one CLI flag
+# regardless of --engine. New engines: add an entry + a comment in
+# the wrapper's __init__ pointing back here.
+_MAX_SEQ_LEN_KEY = {
+    "VLLM": "max_model_len",
+    "TRTLLM": "max_seq_len",
+    "SGLANG": "context_length",
+}
 datasets_available = {
     "mtbench": datasets.MTBench,
     "random": datasets.RandomToken,
@@ -145,8 +160,27 @@ def run_simple(args):
         dataset = datasets.RandomToken(tokenizer, args.random_isl, **dataset_kwargs)
     elif args.specbench is not None:
         dataset = datasets.SpecBench(args.specbench, **dataset_kwargs)
+    # CLI overrides take precedence over --runtime_params; supplying neither
+    # leaves engine_args empty (engine auto-derives sequence length) and
+    # sampling_kwargs defaulting to greedy (temperature=0).
+    #
+    # --max_seq_len is the generic sequence-length cap; _MAX_SEQ_LEN_KEY
+    # (module scope) maps it to the engine-specific kwarg so cell / variant
+    # YAMLs can use one flag regardless of --engine. Engines outside the
+    # table fall back to --runtime_params (engine_args.<their-key>).
     engine_args = args.runtime_params.get("engine_args", {})
+    if args.max_seq_len is not None:
+        key = _MAX_SEQ_LEN_KEY.get(args.engine)
+        if key is None:
+            raise ValueError(
+                f"--max_seq_len is not wired for --engine {args.engine}. "
+                f"Use --runtime_params with engine_args.<key> for this engine, "
+                f"or extend _MAX_SEQ_LEN_KEY in run.py."
+            )
+        engine_args[key] = args.max_seq_len
     sampling_kwargs = args.runtime_params.get("sampling_kwargs", {"temperature": 0})
+    if args.temperature is not None:
+        sampling_kwargs["temperature"] = args.temperature
     model_class = engines_available[args.engine]
     model = model_class(
         args.model_dir,
@@ -155,6 +189,7 @@ def run_simple(args):
         speculative_algorithm=args.speculative_algorithm,
         draft_model_dir=args.draft_model_dir,
         speculative_num_steps=args.draft_length,
+        speculative_num_draft_tokens=args.block_size,
         tensor_parallel_size=args.tp_size,
         moe_expert_parallel_size=args.ep_size,
         trust_remote_code=args.trust_remote_code,
@@ -288,10 +323,46 @@ def run_simple(args):
         default=None,
         help="Path to the runtime params yaml file",
     )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        required=False,
+        default=None,
+        help=(
+            "Sampling temperature. Overrides sampling_kwargs.temperature from "
+            "--runtime_params if both set. Default when neither is set: 0 (greedy)."
+        ),
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        required=False,
+        default=None,
+        help=(
+            "Max sequence length the engine should reserve (input + output). "
+            "Maps to the engine-specific kwarg at the model-wrapper seam: "
+            "VLLM → max_model_len, TRTLLM → max_seq_len, SGLANG → context_length. "
+            "Overrides the same key in --runtime_params engine_args if both "
+            "are set. When neither is set, the engine auto-derives from the "
+            "model config + memory budget, which can cap below the input "
+            "length on tight GPUs. Set to 40960 for the SPEED-Bench "
+            "throughput_32k split (32K input + 4K output + 4K headroom)."
+        ),
+    )
     parser.add_argument(
         "--output_length", type=int, required=False, default=4096, help="Output length"
     )
     parser.add_argument("--draft_length", type=int, required=False, default=3, help="Draft length")
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        required=False,
+        default=None,
+        help=(
+            "DFlash block size (num_speculative_tokens). Use instead of --draft_length "
+            "for DFLASH: block_size = draft_length + 1."
+        ),
+    )
     parser.add_argument(
         "--tp_size", type=int, required=False, default=4, help="Tensor parallel size"
     )

examples/specdec_bench/specdec_bench/models/base.py

@@ -15,6 +15,27 @@
 
 
 class Model:
+    """Base class for inference-engine wrappers.
+
+    Cross-engine kwarg conventions (read by run.py, set on **kwargs):
+
+    - ``sampling_kwargs``: dict-shaped sampling config (``temperature``,
+      etc.). Universal; every engine consumes it.
+    - ``max_model_len`` / ``max_seq_len`` / ``context_length``: max
+      input+output sequence length the engine should reserve. The CLI
+      flag ``--max_seq_len`` in run.py is generic; it is translated to
+      one of these three engine-specific kwargs at the run_simple()
+      seam based on ``--engine``. New engine wrappers should read one
+      of these names and add the mapping in run.py's
+      ``_MAX_SEQ_LEN_KEY``.
+
+    Engine-specific kwargs (``mem_fraction_static`` for SGLang,
+    ``enable_chunked_prefill`` for TRT-LLM, etc.) are passed through
+    ``**kwargs`` from ``--runtime_params engine_args`` without
+    translation — those are the engine's own surface, not part of the
+    cross-engine contract.
+    """
+
     def __init__(self, model_dir, tokenizer, max_draft_length):
         raise NotImplementedError
 

examples/specdec_bench/specdec_bench/models/sglang.py

@@ -26,6 +26,12 @@
 
 
 class SGLANGModel(Model):
+    # Cross-engine ``--max_seq_len`` (run.py) lands in kwargs under the
+    # SGLang-native name ``context_length`` (see run.py's
+    # ``_MAX_SEQ_LEN_KEY``) and is forwarded into ``sgl.Engine(...)``
+    # via ``engine_kwargs["context_length"]`` below. ``None`` lets
+    # SGLang auto-derive from the model config.
+
     def __init__(
         self,
         model_dir,
@@ -58,6 +64,11 @@ def __init__(
             "enable_torch_compile": kwargs.get("enable_torch_compile", False),
             "cuda_graph_max_bs": max_concurrent_requests,
             "disable_cuda_graph": False,
+            # Cross-engine `--max_seq_len` from run.py lands here as
+            # `context_length` (sgl.Engine's spelling). None lets SGLang
+            # auto-derive from the model config — same auto-default
+            # behavior as vLLM's max_model_len=None.
+            "context_length": kwargs.get("context_length"),
         }
         if speculative_algorithm is not None:
             # https://github.com/sgl-project/sglang/pull/3582

examples/specdec_bench/specdec_bench/models/trtllm_torch_api.py

@@ -37,6 +37,11 @@
 
 
 class TRTLLMPYTModel(Model):
+    # Cross-engine ``--max_seq_len`` (run.py) lands in kwargs under the
+    # TRT-LLM-native name ``max_seq_len`` (passthrough — same word, see
+    # run.py's ``_MAX_SEQ_LEN_KEY``) and is read by ``create_executor``
+    # below into ``LLM(max_seq_len=…)``.
+
     def __init__(
         self,
         model_path,

examples/specdec_bench/specdec_bench/models/vllm.py

@@ -29,6 +29,10 @@
 
 
 class VLLMModel(Model):
+    # Cross-engine ``--max_seq_len`` (run.py) lands in kwargs under the
+    # vLLM-native name ``max_model_len`` (see run.py's ``_MAX_SEQ_LEN_KEY``)
+    # and is read at line ~92 below into AsyncEngineArgs.
+
     def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs):
         specdec = None
         if kwargs.get("speculative_algorithm") == "EAGLE3":
@@ -76,6 +80,7 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs
             num_speculative_tokens = 1
         else:
             num_speculative_tokens = specdec.get("num_speculative_tokens", 3)
+
         engine_args = AsyncEngineArgs(
             model=model_dir,
             tokenizer=kwargs.get("tokenizer_path"),
@@ -88,6 +93,7 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs
             skip_tokenizer_init=False,
             async_scheduling=kwargs.get("async_scheduling", True),
             enforce_eager=False,
+            max_model_len=kwargs.get("max_model_len"),
         )
         self.engine_args = engine_args
         self.model = AsyncLLM.from_engine_args(engine_args)
@@ -102,7 +108,7 @@ def __init__(self, model_dir, max_concurrent_requests, sampling_kwargs, **kwargs
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)
 
-    async def run(self, prompt_ids, max_length, end_id, request_id, turn_id):
+    async def run(self, prompt_ids, max_length, end_id, request_id, turn_id):  # pragma: no cover
         output_dict = {}
         self.sampling_config.max_tokens = max_length
         self.sampling_config.stop_token_ids = [end_id]
@@ -134,7 +140,7 @@ async def run(self, prompt_ids, max_length, end_id, request_id, turn_id):
         ]
         return output_dict
 
-    async def generate(self, prompt_ids, request_id, turn_id):
+    async def generate(self, prompt_ids, request_id, turn_id):  # pragma: no cover
         timing = []
         timing.append(time.perf_counter())
         outputs = []
@@ -152,7 +158,7 @@ async def generate(self, prompt_ids, request_id, turn_id):
                 break
         return outputs, timing, full_tokens
 
-    def get_serving_config(self):
+    def get_serving_config(self):  # pragma: no cover
         """Dump the AsyncEngineArgs dataclass plus the runtime vllm_config when available."""
         try:
             import dataclasses
@@ -171,7 +177,7 @@ def get_serving_config(self):
             pass
         return cfg
 
-    def stop(self):
+    def stop(self):  # pragma: no cover
         try:
             self.loop.run_until_complete(self.model.shutdown())
             self.loop.close()

examples/specdec_bench/specdec_bench/utils.py

@@ -196,6 +196,8 @@ def _checkpoint_provenance(model_dir):
 
 
 def _is_sensitive_key(key):
+    if not isinstance(key, str):
+        return False
     klow = key.lower()
     if klow in _SENSITIVE_KEY_ALLOWLIST:
         return False

examples/specdec_bench/upload_to_s3.py

@@ -72,8 +72,10 @@ def _check_provenance(run_dir: Path) -> list[str]:
 
 # ── S3 helpers ────────────────────────────────────────────────────────────────
 # Endpoint, key id, and secret default to empty and are taken from --endpoint /
-# --key-id / --secret (or the corresponding S3_ENDPOINT / S3_KEY_ID / S3_SECRET
-# env vars).
+# --key-id / --secret (or the corresponding SPECDEC_BENCH_S3_ENDPOINT /
+# SPECDEC_BENCH_S3_KEY_ID / SPECDEC_BENCH_S3_SECRET env vars). The prefix
+# disambiguates from any other S3 credentials a CI runner or user shell might
+# carry — only specdec_bench result uploads use this set.
 
 
 def parse_s3_path(path: str) -> tuple[str, str]:
@@ -175,19 +177,19 @@ def main():
     )
     parser.add_argument(
         "--endpoint",
-        default=os.environ.get("S3_ENDPOINT", ""),
-        help="S3 endpoint URL",
+        default=os.environ.get("SPECDEC_BENCH_S3_ENDPOINT", ""),
+        help="S3 endpoint URL (default: $SPECDEC_BENCH_S3_ENDPOINT)",
     )
     parser.add_argument(
         "--key-id",
-        default=os.environ.get("S3_KEY_ID", ""),
+        default=os.environ.get("SPECDEC_BENCH_S3_KEY_ID", ""),
         dest="key_id",
-        help="S3 access key ID",
+        help="S3 access key ID (default: $SPECDEC_BENCH_S3_KEY_ID)",
     )
     parser.add_argument(
         "--secret",
-        default=os.environ.get("S3_SECRET", ""),
-        help="S3 secret access key",
+        default=os.environ.get("SPECDEC_BENCH_S3_SECRET", ""),
+        help="S3 secret access key (default: $SPECDEC_BENCH_S3_SECRET)",
     )
     parser.add_argument(
         "--skip-existing",

tools/launcher/common/specdec_bench/upload_to_s3.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
+source ${SCRIPT_DIR}/../service_utils.sh
+
+trap 'error_handler $0 $LINENO' ERR
+trap 'exit_handler' EXIT
+
+###################################################################################################
+# Upload a specdec_bench results directory to S3. Thin wrapper around
+# examples/specdec_bench/upload_to_s3.py.
+#
+# YAML usage:
+#   task_2:
+#     script: common/specdec_bench/upload_to_s3.sh
+#     args:
+#       - /scratchspace/specdec_bench
+#       - s3://team-specdec-workgroup/results
+#       - --skip-existing                  # optional
+#       - --allow-incomplete-provenance    # optional, for runs without CONTAINER_IMAGE set
+#
+# Required env (or pass via --endpoint / --key-id / --secret to the underlying script):
+#   S3_ENDPOINT, S3_KEY_ID, S3_SECRET
+
+# Install boto3 if not already in the container. Warm pipelines where an
+# earlier specdec_bench task ran will already have it from run.sh.
+if ! pip show boto3 >/dev/null 2>&1; then
+    if ! pip install -r modules/Model-Optimizer/examples/specdec_bench/requirements.txt; then
+        report_result "FAIL: upload_to_s3: pip install requirements.txt failed"
+        exit 1
+    fi
+fi
+
+if ! python3 modules/Model-Optimizer/examples/specdec_bench/upload_to_s3.py "${@}"; then
+    report_result "FAIL: upload_to_s3: upload_to_s3.py exited non-zero"
+    exit 1
+fi
+
+report_result "PASS: upload_to_s3 completed"

tools/launcher/core.py

@@ -38,18 +38,29 @@
 def get_default_env(experiment_title=None):
     """Return (slurm_env, local_env) dicts for the given experiment title."""
     title = experiment_title or DEFAULT_EXPERIMENT_TITLE
+    # specdec_bench upload credentials — forwarded so that the YAML pipeline
+    # step `common/specdec_bench/upload_to_s3.sh` can publish to the team
+    # S3 bucket without baking secrets into committed YAMLs. The prefix
+    # disambiguates from any other S3 creds a CI runner might carry.
+    specdec_s3 = {
+        "SPECDEC_BENCH_S3_ENDPOINT": os.getenv("SPECDEC_BENCH_S3_ENDPOINT", ""),
+        "SPECDEC_BENCH_S3_KEY_ID": os.getenv("SPECDEC_BENCH_S3_KEY_ID", ""),
+        "SPECDEC_BENCH_S3_SECRET": os.getenv("SPECDEC_BENCH_S3_SECRET", ""),
+    }
     slurm_env = {
         "TRITON_CACHE_DIR": f"/{title}/triton-cache",
         "HF_HOME": f"/{title}/hf-cache",
         "HF_TOKEN": os.getenv("HF_TOKEN", ""),
         "MLM_SKIP_INSTALL": "1",
         "LAUNCH_SCRIPT": "python",
+        **specdec_s3,
     }
     local_env = {
         "TRITON_CACHE_DIR": f"/{title}/triton-cache",
         "HF_HOME": f"/{title}/hf-cache",
         "HF_TOKEN": os.getenv("HF_TOKEN", ""),
         "MLM_SKIP_INSTALL": "1",
+        **specdec_s3,
     }
     return slurm_env, local_env