[None][test] Revert the gen_only warmup probe and fix the per-iter device step time metric

examples/disaggregated/slurm/benchmark/submit.py

@@ -304,17 +304,30 @@ 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" 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')
-        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}')
+    if benchmark_mode == "gen_only":
+        # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP is empirically load-bearing
+        # on the CTX worker for gen_only throughput on NIXL disagg workloads
+        # (removing it on CTX regressed deepseek-r1-fp4 8k1k dep4/tep8 by
+        # -27.9 %; see NVBug 6627789 / bisect PR #17535). Default: apply on
+        # both workers, matching the pre-PR contract. Workloads that observed
+        # the CTX-side hang the PR was scoped for can opt out by exporting
+        # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX=0 in the submitting
+        # shell; only the GEN worker receives the flag in that case.
+        _apply_ctx_kv_overlap = os.environ.get(
+            'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX', '1') != '0'
+        if _apply_ctx_kv_overlap:
+            upsert_env_config(env_config, 'worker_env_var',
+                              'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP',
+                              'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1')
+        elif role == "GEN":
+            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}')
 
     # 2. Add profiling env vars to env_config (conditional)
     if nsys_on:

jenkins/scripts/perf/local/submit.py

@@ -901,8 +901,18 @@ 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)
-            # GEN worker only: the same flag on the CTX worker has been seen to
-            # hang gen_only runs with KV blocks never released.
+            # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP is empirically
+            # load-bearing on the CTX worker for gen_only throughput on
+            # NIXL disagg (removing it on CTX regressed deepseek-r1-fp4
+            # 8k1k dep4/tep8 by -27.9 %; see NVBug 6627789 / PR #17535).
+            # Default: apply on both workers, matching the pre-PR
+            # contract. Workloads that observed the CTX-side hang the PR
+            # was scoped for can opt out by exporting
+            # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX=0.
+            if os.environ.get("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX", "1") != "0":
+                ctx_worker_env_vars = (
+                    f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 {ctx_worker_env_vars}"
+                )
             gen_worker_env_vars = (
                 f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 "
                 f"TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency} {gen_worker_env_vars}"

jenkins/scripts/perf/submit.py

@@ -896,8 +896,18 @@ def main():
             srun_args_lines.append("--container-env=TRTLLM_DISAGG_BENCHMARK_GEN_ONLY")
         elif benchmark_mode == "gen_only":
             concurrency = benchmark_config.get("concurrency", 1)
-            # GEN worker only: the same flag on the CTX worker has been seen to
-            # hang gen_only runs with KV blocks never released.
+            # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP is empirically
+            # load-bearing on the CTX worker for gen_only throughput on
+            # NIXL disagg (removing it on CTX regressed deepseek-r1-fp4
+            # 8k1k dep4/tep8 by -27.9 %; see NVBug 6627789 / PR #17535).
+            # Default: apply on both workers, matching the pre-PR
+            # contract. Workloads that observed the CTX-side hang the PR
+            # was scoped for can opt out by exporting
+            # TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX=0.
+            if os.environ.get("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP_CTX", "1") != "0":
+                ctx_worker_env_vars = (
+                    f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 {ctx_worker_env_vars}"
+                )
             gen_worker_env_vars = (
                 f"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 "
                 f"TLLM_BENCHMARK_REQ_QUEUES_SIZE={concurrency} {gen_worker_env_vars}"