diff --git a/examples/disaggregated/slurm/benchmark/submit.py b/examples/disaggregated/slurm/benchmark/submit.py index f1cf53f87996..60c51c81626f 100644 --- a/examples/disaggregated/slurm/benchmark/submit.py +++ b/examples/disaggregated/slurm/benchmark/submit.py @@ -291,15 +291,17 @@ def build_worker_environment(worker_config, env_config, role, benchmark_mode, upsert_env_config(env_config, 'worker_env_var', 'TRTLLM_DISAGG_BENCHMARK_GEN_ONLY', 'TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1') - if benchmark_mode == "gen_only": - upsert_env_config(env_config, 'worker_env_var', + if benchmark_mode == "gen_only" and role == "GEN": + # GEN worker only: skipping transfer-state polling helps the generation + # worker, but the same flag on the CTX worker has been seen to hang + # gen_only runs with KV blocks never released. + upsert_env_config(env_config, 'gen_worker_env_var', 'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP', 'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1') - if role == "GEN": - concurrency = _parse_positive_concurrency(concurrency) - upsert_env_config(env_config, 'gen_worker_env_var', - 'TLLM_BENCHMARK_REQ_QUEUES_SIZE', - f'TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency}') + concurrency = _parse_positive_concurrency(concurrency) + upsert_env_config(env_config, 'gen_worker_env_var', + 'TLLM_BENCHMARK_REQ_QUEUES_SIZE', + f'TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency}') # 2. Add profiling env vars to env_config (conditional) if nsys_on: diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 8e571eb2f890..e359b81bfff7 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -901,9 +901,8 @@ def main(): srun_args_lines.append("--container-env=TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") elif "gen_only" in bm_config.get("mode", ""): concurrency = bm_config.get("concurrency", 1) - ctx_worker_env_vars = ( - f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 {ctx_worker_env_vars}" - ) + # GEN worker only: the same flag on the CTX worker has been seen to + # hang gen_only runs with KV blocks never released. gen_worker_env_vars = ( f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 " f"TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency} {gen_worker_env_vars}" diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index f1c54fb13690..dad770df9654 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -796,9 +796,8 @@ def main(): srun_args_lines.append("--container-env=TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") elif benchmark_mode == "gen_only": concurrency = benchmark_config.get("concurrency", 1) - ctx_worker_env_vars = ( - f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 {ctx_worker_env_vars}" - ) + # GEN worker only: the same flag on the CTX worker has been seen to + # hang gen_only runs with KV blocks never released. gen_worker_env_vars = ( f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 " f"TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency} {gen_worker_env_vars}"