diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index 337c5184f20..ca2f9908966 100644 --- a/examples/specdec_bench/run.py +++ b/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.). 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. 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" ) diff --git a/examples/specdec_bench/specdec_bench/models/base.py b/examples/specdec_bench/specdec_bench/models/base.py index 43d3a1337d3..c14f7010027 100644 --- a/examples/specdec_bench/specdec_bench/models/base.py +++ b/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 diff --git a/examples/specdec_bench/specdec_bench/models/sglang.py b/examples/specdec_bench/specdec_bench/models/sglang.py index 99a66b0647e..f95b9eb0ed8 100644 --- a/examples/specdec_bench/specdec_bench/models/sglang.py +++ b/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 diff --git a/examples/specdec_bench/specdec_bench/models/trtllm_torch_api.py b/examples/specdec_bench/specdec_bench/models/trtllm_torch_api.py index 25a2aed6323..0bbefc02cd5 100644 --- a/examples/specdec_bench/specdec_bench/models/trtllm_torch_api.py +++ b/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, diff --git a/examples/specdec_bench/specdec_bench/models/vllm.py b/examples/specdec_bench/specdec_bench/models/vllm.py index 52bf35f1a0d..a344842ebf2 100644 --- a/examples/specdec_bench/specdec_bench/models/vllm.py +++ b/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() diff --git a/examples/specdec_bench/specdec_bench/utils.py b/examples/specdec_bench/specdec_bench/utils.py index 9a52d0ceac2..5de76f90c5a 100644 --- a/examples/specdec_bench/specdec_bench/utils.py +++ b/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 diff --git a/examples/specdec_bench/upload_to_s3.py b/examples/specdec_bench/upload_to_s3.py index a0868101082..067ea25bce3 100644 --- a/examples/specdec_bench/upload_to_s3.py +++ b/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", diff --git a/tools/launcher/common/specdec_bench/upload_to_s3.sh b/tools/launcher/common/specdec_bench/upload_to_s3.sh new file mode 100755 index 00000000000..59db5b00fc5 --- /dev/null +++ b/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" diff --git a/tools/launcher/core.py b/tools/launcher/core.py index 8382e0de4f4..7d53cca0db4 100644 --- a/tools/launcher/core.py +++ b/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 diff --git a/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench.yaml b/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench.yaml index 4bd925f2436..d6872339f8c 100644 --- a/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench.yaml +++ b/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench.yaml @@ -1,9 +1,19 @@ -# SPEED-bench smoke run for Qwen3.5-4B via vLLM (autoregressive baseline). +# SPEED-bench run for Qwen3.5-4B via vLLM (autoregressive baseline). # -# Reads nvidia/SPEED-Bench-Internal/qualitative through Qwen/Qwen3.5-4B with -# --speculative_algorithm NONE (no draft model) and writes timing.json + -# aa_timing.json + acceptance_rate.json + specbench_responses.jsonl + -# specbench_results.json to /scratchspace/specdec_bench/. +# Two-task pipeline: +# task_0 Quantitative quality split (nvidia/SPEED-Bench-Internal/qualitative) +# task_1 Long-context throughput split (nvidia/SPEED-Bench-Internal/throughput_32k) +# +# Both use --speculative_algorithm NONE (no draft model) — this is the +# autoregressive baseline that the MTP variant in specdec_bench_mtp.yaml is +# compared against. +# +# Results write to /scratchspace/qwen35_4b_none_vllm//. The +# pensieve-intern `specdec_bench` workflow's wrap_up stage owns publishing +# these to s3://team-specdec-workgroup/results/qwen35_4b_none_vllm// +# with provenance stamps (jira_ticket + huggingface_model_id). Sweep-name +# convention: __ so multi-model / multi-engine +# records don't collide in S3. # # The qwen3_5 model_type needs transformers >= 4.58, which is NOT in # vllm/vllm-openai:latest yet — use the qwen3_5-cu130 tag instead. @@ -21,6 +31,10 @@ pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3.5-4B + # Step 1: qualitative split — quality / acceptance-rate numbers. + # tp_size=2 + concurrency=32 trades aa_timing fidelity for ~30x wall-clock + # speedup; acceptance-length (AL) is concurrency-independent and is the + # primary metric we care about for this split. task_0: script: common/specdec_bench/run.sh args: @@ -28,14 +42,44 @@ pipeline: - --dataset_path /hf-local/nvidia/SPEED-Bench-Internal/qualitative - --engine VLLM - --speculative_algorithm NONE - - --tp_size 1 + - --tp_size 2 - --ep_size 1 - - --concurrency 1 + - --concurrency 32 + - --output_length 4096 + - --aa_timing + - --show_progress + - --save_dir /scratchspace/qwen35_4b_none_vllm/qualitative + environment: + - HF_MODEL_CKPT: <> + - HF_LOCAL: /hf-local + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:qwen3_5-cu130 + + # Step 2: throughput_32k split — long-context throughput. + # `--num_requests 80` caps the run at 80 samples (split has 1,536) so it fits + # in the 4h Slurm time-limit; each 32K-input sample takes ~60-90s. + # tp_size=2 doubles the KV-cache budget across 2 GPUs; concurrency=8 keeps + # 8 * 32K = 256K tokens of in-flight KV under that doubled budget. + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench-Internal/throughput_32k + - --engine VLLM + - --speculative_algorithm NONE + - --tp_size 2 + - --ep_size 1 + - --concurrency 8 - --num_requests 80 - --output_length 4096 + - --max_seq_len 40960 - --aa_timing - --show_progress - - --save_dir /scratchspace/specdec_bench + - --save_dir /scratchspace/qwen35_4b_none_vllm/throughput_32k environment: - HF_MODEL_CKPT: <> - HF_LOCAL: /hf-local @@ -43,5 +87,11 @@ pipeline: _factory_: "slurm_factory" nodes: 1 ntasks_per_node: 1 - gpus_per_node: 1 + gpus_per_node: 2 container: vllm/vllm-openai:qwen3_5-cu130 + +# S3 upload is intentionally not a task in this YAML — the bench pipeline only +# writes results to /scratchspace/qwen35_4b_none_vllm//. The +# pensieve-intern specdec_bench workflow's wrap_up stage owns harvesting these +# from lustre and publishing them to the team S3 vault with provenance stamps +# (jira_ticket + huggingface_model_id) for the "official record" tracking. diff --git a/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml b/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml index 7a6720cd8ec..b1f5590b827 100644 --- a/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml +++ b/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml @@ -5,6 +5,17 @@ # with draft_length=3 to produce real acceptance-rate numbers instead of the # trivial AR=1 that NONE yields. # +# Two-task pipeline: +# task_0 Quantitative quality split (nvidia/SPEED-Bench-Internal/qualitative) +# task_1 Long-context throughput split (nvidia/SPEED-Bench-Internal/throughput_32k) +# +# Results write to /scratchspace/qwen35_4b_mtp_vllm//. The +# pensieve-intern `specdec_bench` workflow's wrap_up stage owns publishing +# these to s3://team-specdec-workgroup/results/qwen35_4b_mtp_vllm// +# with provenance stamps (jira_ticket + huggingface_model_id). Sweep-name +# convention: __ so multi-model / multi-engine +# records don't collide in S3. +# # Slurm run on cw_dfw: # uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml --yes @@ -14,6 +25,10 @@ pipeline: global_vars: hf_model: /hf-local/Qwen/Qwen3.5-4B + # Step 1: qualitative split — quality / acceptance-rate numbers with MTP draft=3. + # tp_size=2 + concurrency=32 trades aa_timing fidelity for ~30x wall-clock + # speedup; acceptance-length (AL) is concurrency-independent and is the + # primary metric we care about for this split. task_0: script: common/specdec_bench/run.sh args: @@ -22,14 +37,13 @@ pipeline: - --engine VLLM - --speculative_algorithm MTP - --draft_length 3 - - --tp_size 1 + - --tp_size 2 - --ep_size 1 - - --concurrency 1 - - --num_requests 80 + - --concurrency 32 - --output_length 4096 - --aa_timing - --show_progress - - --save_dir /scratchspace/specdec_bench_mtp + - --save_dir /scratchspace/qwen35_4b_mtp_vllm/qualitative environment: - HF_MODEL_CKPT: <> - HF_LOCAL: /hf-local @@ -52,5 +66,44 @@ pipeline: _factory_: "slurm_factory" nodes: 1 ntasks_per_node: 1 - gpus_per_node: 1 + gpus_per_node: 2 container: vllm/vllm-openai:qwen3_5-cu130 + + # Step 2: throughput_32k split — long-context throughput with MTP draft=3. + # `--num_requests 80` caps the run at 80 samples (split has 1,536) so it fits + # in the 4h Slurm time-limit; each 32K-input sample takes ~60-90s. + # tp_size=2 doubles the KV-cache budget across 2 GPUs; concurrency=8 keeps + # 8 * 32K = 256K tokens of in-flight KV under that doubled budget. + task_1: + script: common/specdec_bench/run.sh + args: + - --dataset speed + - --dataset_path /hf-local/nvidia/SPEED-Bench-Internal/throughput_32k + - --engine VLLM + - --speculative_algorithm MTP + - --draft_length 3 + - --tp_size 2 + - --ep_size 1 + - --concurrency 8 + - --num_requests 80 + - --output_length 4096 + - --max_seq_len 40960 + - --aa_timing + - --show_progress + - --save_dir /scratchspace/qwen35_4b_mtp_vllm/throughput_32k + environment: + - HF_MODEL_CKPT: <> + - HF_LOCAL: /hf-local + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:qwen3_5-cu130 + + +# S3 upload is intentionally not a task in this YAML — the bench pipeline only +# writes results to /scratchspace/qwen35_4b_mtp_vllm//. The +# pensieve-intern specdec_bench workflow's wrap_up stage owns harvesting these +# from lustre and publishing them to the team S3 vault with provenance stamps +# (jira_ticket + huggingface_model_id) for the "official record" tracking.