diff --git a/examples/disaggregated/README.md b/examples/disaggregated/README.md
index 120706dd01a..376965c5c90 100644
--- a/examples/disaggregated/README.md
+++ b/examples/disaggregated/README.md
@@ -143,3 +143,7 @@ Once the context and generation servers are launched, you can again launch the d
 ```
 trtllm-serve disaggregated -c disagg_config.yaml
 ```
+
+## Launching disaggregated serving using SLURM
+
+Please refer to the [SLURM example](./slurm/README.md) for an example of how to launch disaggregated serving using SLURM.
diff --git a/examples/disaggregated/slurm/README.md b/examples/disaggregated/slurm/README.md
new file mode 100644
index 00000000000..7549c1ae2d0
--- /dev/null
+++ b/examples/disaggregated/slurm/README.md
@@ -0,0 +1,61 @@
+# Disaggregated Serving Launcher for SLURM
+
+## Overview
+
+A tool for running disaggregated serving benchmarks with TRT-LLM on SLURM.
+
+## Usage
+
+```bash
+python3 launcher.py --account <account> --partition <partition> --time <time> --job-name <job-name> --container-image <container-image> --config-file <config-file> --experiment-path <experiment-path> --request-allocation --num-gpus <num-gpus>
+```
+
+### Configuration
+
+The configuration file should be in the following format:
+
+```yaml
+exec:
+  model_path: <model-path>
+  # Determines the disaggregated serving configuration
+  config:
+    context:
+      tp: <tp>
+      ep: <ep>
+      pp: <pp>
+      max_batch_size: <max_batch_size>
+      max_num_tokens: <max_num_tokens>
+      max_seq_len: <max_seq_len>
+      config:
+        # Determines the context server PyTorch configuration
+        print_iter_log: true
+        disable_overlap_scheduler: true
+        kv_cache_config:
+          free_gpu_memory_fraction: 0.75
+          enable_block_reuse: false
+    generation:
+      tp: <tp>
+      ep: <ep>
+      pp: <pp>
+      max_batch_size: <max_batch_size>
+      max_num_tokens: <max_num_tokens>
+      max_seq_len: <max_seq_len>
+      config:
+        # Determines the generation server PyTorch configuration
+        print_iter_log: true
+        kv_cache_config:
+          free_gpu_memory_fraction: 0.75
+          enable_block_reuse: false
+
+# Determines the profiling configuration
+profile:
+  isl: <isl>
+  osl: <osl>
+  use_benchmark_serving: true
+  concurrency:
+    - <concurrency>
+```
+
+Please refer to the [config.yaml](config.yaml) file for an example configuration.
+
+The experiment results will be saved in the experiment path.
diff --git a/examples/disaggregated/slurm/config.yaml b/examples/disaggregated/slurm/config.yaml
new file mode 100644
index 00000000000..3efb7622bd0
--- /dev/null
+++ b/examples/disaggregated/slurm/config.yaml
@@ -0,0 +1,36 @@
+exec:
+  config:
+    context:
+      tp: 4
+      ep: 4
+      pp: 1
+      max_batch_size: 4
+      max_num_tokens: 1024
+      max_seq_len: 1024
+      config:
+        kv_cache_config:
+          free_gpu_memory_fraction: 0.75
+          enable_block_reuse: false
+        print_iter_log: true
+      dp: 1
+    generation:
+      tp: 4
+      ep: 4
+      pp: 1
+      max_batch_size: 1
+      max_num_tokens: 4096
+      max_seq_len: 2048
+      config:
+        print_iter_log: true
+        kv_cache_config:
+          free_gpu_memory_fraction: 0.75
+          enable_block_reuse: false
+      dp: 1
+  model_path: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+profile:
+  isl: 1024
+  osl: 1024
+  use_benchmark_serving: true
+  concurrency:
+  - 128
+  - 256
diff --git a/examples/disaggregated/slurm/disagg_profiler/__init__.py b/examples/disaggregated/slurm/disagg_profiler/__init__.py
new file mode 100644
index 00000000000..0e7fa7a064e
--- /dev/null
+++ b/examples/disaggregated/slurm/disagg_profiler/__init__.py
@@ -0,0 +1,16 @@
+"""
+Disaggregated serving profiler package.
+
+This package contains the job management and parameter sweeping functionality
+for the TRT-LLM disaggregated serving launcher.
+"""
+
+from .job_manager import JobManager, calculate_nodes_needed, wait_for_server
+from .sweeper import (AutoSweeper, ParameterSweeper, get_slurm_allocation,
+                      run_sweep_configuration)
+
+__all__ = [
+    'JobManager', 'calculate_nodes_needed', 'wait_for_server',
+    'ParameterSweeper', 'AutoSweeper', 'get_slurm_allocation',
+    'run_sweep_configuration'
+]
diff --git a/examples/disaggregated/slurm/disagg_profiler/job_manager.py b/examples/disaggregated/slurm/disagg_profiler/job_manager.py
new file mode 100644
index 00000000000..6b36fe0bc90
--- /dev/null
+++ b/examples/disaggregated/slurm/disagg_profiler/job_manager.py
@@ -0,0 +1,1058 @@
+import json
+import os
+import subprocess
+import time
+import uuid
+
+import requests
+import yaml
+
+
+def wait_for_server(host, port, max_retries=1000, delay=5):
+    """Wait for a server to become available."""
+    for _ in range(max_retries):
+        try:
+            print(f"Checking server status at {host}:{port}")
+            state = requests.get(f"http://{host}:{port}/health")
+            if state.status_code == 200:
+                print("Server is running")
+                return True
+        except:
+            time.sleep(delay)
+
+    print(f"Server did not start after {max_retries} attempts")
+    return False
+
+
+def calculate_nodes_needed(config, num_gpus):
+    """
+    Calculate the number of nodes needed based on the configuration.
+
+    The calculation is based on:
+    1. The number of context servers and their TP size
+    2. The number of generation servers and their TP size
+    3. The number of IFB servers and their TP size
+    4. The number of GPUs per node (num_gpus)
+
+    Args:
+        config (dict): The configuration dictionary
+
+    Returns:
+        int: The number of nodes needed
+    """
+    total_gpus_needed = 0
+
+    # Calculate GPUs needed for context servers
+    if 'context' in config['exec']['config']:
+        context_config = config['exec']['config']['context']
+        context_tp = context_config['tp']
+        context_dp = context_config['dp']
+        total_gpus_needed += context_tp * context_dp
+
+    # Calculate GPUs needed for generation servers
+    if 'generation' in config['exec']['config']:
+        generation_config = config['exec']['config']['generation']
+        generation_tp = generation_config['tp']
+        generation_dp = generation_config['dp']
+        total_gpus_needed += generation_tp * generation_dp
+
+    # Calculate GPUs needed for IFB servers
+    if 'ifb' in config['exec']['config']:
+        ifb_config = config['exec']['config']['ifb']
+        ifb_tp = ifb_config['tp']
+        ifb_dp = ifb_config['dp']
+        total_gpus_needed += ifb_tp * ifb_dp
+
+    # Calculate nodes needed (round up to ensure enough nodes)
+    nodes_needed = (total_gpus_needed + num_gpus - 1) // num_gpus
+
+    # Ensure we request at least one node
+    return max(1, nodes_needed)
+
+
+class JobManager:
+    """Manages SLURM jobs for the disaggregated serving setup."""
+
+    def __init__(self,
+                 args,
+                 load_config_only=False,
+                 override_config=None,
+                 override_job_id=None):
+        """
+        Initialize the job manager.
+
+        Args:
+            args (argparse.Namespace): Command line arguments.
+            load_config_only (bool): If True, only load the configuration without launching jobs.
+            override_config (dict): Override the base config with this config if provided.
+            override_job_id (str): Override the auto-generated job ID with this ID if provided.
+        """
+        print(
+            f"DEBUG JobManager: __init__ called with load_config_only={load_config_only}, override_job_id={override_job_id}"
+        )
+
+        self.args = args
+        self.num_gpus = args.num_gpus
+        self.account = args.account
+        self.partition = args.partition
+        self.time = args.time
+        self.job_name = args.job_name
+        self.container_image = args.container_image
+        self.mounts = args.mounts
+        self.workdir = os.path.abspath(os.path.join(os.getcwd()))
+        self.context_jobs = []
+        self.generation_jobs = []
+        self.ifb_jobs = []
+        self.disagg_jobs = []
+        self.experiment_path = args.experiment_path
+
+        # Clean up any previous node info files
+        self.setup_shared_directories(override_job_id)
+
+        # Generate configuration files
+        self.base_config = self.load_config()
+
+        # If override_config is provided, use it instead of the base_config
+        if override_config:
+            self.config = override_config
+        else:
+            self.config = self.base_config
+        print(self.config)
+
+        # If we only need to load the config, return here
+        if load_config_only:
+            print(
+                "DEBUG JobManager: load_config_only=True, returning without launching jobs"
+            )
+            return
+
+        self.output_folder = self.create_output_folder()
+
+        # Check if results already exist and should be skipped
+        if hasattr(args, 'skip_existing') and args.skip_existing:
+            if self._check_existing_results():
+                print(
+                    f"Skipping execution - results already exist in {self.output_folder}"
+                )
+                return
+
+        # Run the job with the current configuration
+        print("DEBUG JobManager: About to call prepare_and_run_single_config")
+        self.prepare_and_run_single_config()
+        print("DEBUG JobManager: Completed prepare_and_run_single_config")
+
+    def create_output_folder(self):
+        # The parent directory must be model name
+        model_name = self.config['exec']['model_path'].split('/')[-1]
+        if model_name == '':
+            model_name = self.config['exec']['model_path'].split('/')[-2]
+        model_name = os.path.join(self.experiment_path, model_name)
+        os.makedirs(model_name, exist_ok=True)
+        output_folder = os.path.join(model_name, (
+            f"{self.config['profile']['isl']}_{self.config['profile']['osl']}"))
+        os.makedirs(output_folder, exist_ok=True)
+        ifb_tp = self.config['exec']['config']['ifb'][
+            'tp'] if 'ifb' in self.config['exec']['config'] else 0
+        ifb_dp = self.config['exec']['config']['ifb'][
+            'dp'] if 'ifb' in self.config['exec']['config'] else 0
+        ctx_tp = self.config['exec']['config']['context'][
+            'tp'] if 'context' in self.config['exec']['config'] else 0
+        ctx_dp = self.config['exec']['config']['context'][
+            'dp'] if 'context' in self.config['exec']['config'] else 0
+        gen_tp = self.config['exec']['config']['generation'][
+            'tp'] if 'generation' in self.config['exec']['config'] else 0
+        gen_dp = self.config['exec']['config']['generation'][
+            'dp'] if 'generation' in self.config['exec']['config'] else 0
+
+        output_folder = os.path.join(output_folder, (f"{ctx_tp}_"
+                                                     f"{ctx_dp}_"
+                                                     f"{gen_tp}_"
+                                                     f"{gen_dp}_"
+                                                     f"{ifb_tp}_"
+                                                     f"{ifb_dp}"))
+
+        for server_type in ['generation', 'ifb']:
+            if f'{server_type}' in self.config['exec']['config']:
+                config = self.config['exec']['config'][f'{server_type}'][
+                    'config']
+                pytorch_config = config[
+                    'pytorch_backend_config'] if 'pytorch_backend_config' in config else config
+                if 'disable_overlap_scheduler' not in pytorch_config or not pytorch_config[
+                        'disable_overlap_scheduler']:
+                    output_folder += '_overlap'
+                if 'use_cuda_graph' in pytorch_config:
+                    if pytorch_config['use_cuda_graph']:
+                        output_folder += '_cuda_graph'
+                if 'enable_attention_dp' in config:
+                    if config['enable_attention_dp']:
+                        output_folder += '_adp'
+                if 'speculative_config' in config:
+                    if "num_nextn_predict_layers" in config[
+                            "speculative_config"]:
+                        output_folder += f'_mtp{config["speculative_config"]["num_nextn_predict_layers"]}'
+                config = self.config['exec']['config'][f'{server_type}']
+                if 'env' in config:
+                    for key, value in config['env'].items():
+                        output_folder += f'_{key}_{value}'
+
+        if not os.path.exists(output_folder):
+            os.makedirs(output_folder)
+        return output_folder
+
+    def _check_existing_results(self):
+        """
+        Check if results already exist for this configuration.
+
+        Returns:
+            bool: True if all results exist and are valid, False otherwise
+        """
+        if not os.path.exists(self.output_folder):
+            return False
+
+        # Get concurrency levels from the configuration
+        concurrency_levels = self.config.get('profile',
+                                             {}).get('concurrency', [1])
+
+        for concurrency in concurrency_levels:
+            concurrency_path = os.path.join(self.output_folder,
+                                            str(concurrency))
+            raw_data_path = os.path.join(concurrency_path, "raw_data.json")
+
+            if not os.path.exists(raw_data_path):
+                return False
+
+            # Check if the raw_data.json is valid and contains expected results
+            try:
+                with open(raw_data_path, 'r') as f:
+                    data = json.load(f)
+
+                if 'results' not in data or 'infbench_summary' not in data[
+                        'results']:
+                    return False
+
+                if 'Requests/s' not in data['results']['infbench_summary']:
+                    return False
+            except (json.JSONDecodeError, KeyError):
+                return False
+
+        return True
+
+    def setup_shared_directories(self, override_job_id=None):
+        """Create shared directories for coordination between nodes."""
+
+        # Use a unique directory for each job, or the override_job_id if provided
+        if override_job_id:
+            job_id = override_job_id
+        else:
+            job_id = str(uuid.uuid4())
+
+        os.makedirs(f"job_{job_id}", exist_ok=True)
+        self.job_id = job_id
+
+    def load_config(self):
+        """Load the configuration file."""
+        with open(self.args.config_file, "r") as f:
+            return yaml.safe_load(f)
+
+    def create_context_config(self):
+        """Create the context server configuration file."""
+        if 'context' in self.config['exec']['config']:
+            context_config = self.config['exec']['config']['context']['config']
+
+            # Create config in temporary job directory for execution
+            config_path = os.path.join(f"job_{self.job_id}", "context.yaml")
+            with open(config_path, 'w') as f:
+                yaml.dump(context_config, f)
+
+            # Also save a copy to the output folder for reference
+            if hasattr(self, 'output_folder') and self.output_folder:
+                output_config_path = os.path.join(self.output_folder,
+                                                  "context_config.yaml")
+                try:
+                    with open(output_config_path, 'w') as f:
+                        yaml.dump(context_config,
+                                  f,
+                                  default_flow_style=False,
+                                  indent=2)
+                    print(
+                        f"Context configuration saved to {output_config_path}")
+                except Exception as e:
+                    print(
+                        f"Warning: Could not save context config to output folder: {e}"
+                    )
+
+            return os.path.abspath(config_path)
+        return None
+
+    def create_generation_config(self):
+        """Create the generation server configuration file."""
+        if 'generation' in self.config['exec']['config']:
+            generation_config = self.config['exec']['config']['generation'][
+                'config']
+
+            # Create config in temporary job directory for execution
+            config_path = os.path.join(f"job_{self.job_id}", "generation.yaml")
+            with open(config_path, 'w') as f:
+                yaml.dump(generation_config, f)
+
+            # Also save a copy to the output folder for reference
+            if hasattr(self, 'output_folder') and self.output_folder:
+                output_config_path = os.path.join(self.output_folder,
+                                                  "generation_config.yaml")
+                try:
+                    with open(output_config_path, 'w') as f:
+                        yaml.dump(generation_config,
+                                  f,
+                                  default_flow_style=False,
+                                  indent=2)
+                    print(
+                        f"Generation configuration saved to {output_config_path}"
+                    )
+                except Exception as e:
+                    print(
+                        f"Warning: Could not save generation config to output folder: {e}"
+                    )
+
+            return os.path.abspath(config_path)
+        return None
+
+    def create_ifb_config(self):
+        """Create the IFB server configuration file."""
+        if 'ifb' in self.config['exec']['config']:
+            ifb_config = self.config['exec']['config']['ifb']['config']
+
+            # Create config in temporary job directory for execution
+            config_path = os.path.join(f"job_{self.job_id}", "ifb.yaml")
+            with open(config_path, 'w') as f:
+                yaml.dump(ifb_config, f)
+
+            # Also save a copy to the output folder for reference
+            if hasattr(self, 'output_folder') and self.output_folder:
+                output_config_path = os.path.join(self.output_folder,
+                                                  "ifb_config.yaml")
+                try:
+                    with open(output_config_path, 'w') as f:
+                        yaml.dump(ifb_config,
+                                  f,
+                                  default_flow_style=False,
+                                  indent=2)
+                    print(f"IFB configuration saved to {output_config_path}")
+                except Exception as e:
+                    print(
+                        f"Warning: Could not save IFB config to output folder: {e}"
+                    )
+
+            return os.path.abspath(config_path)
+        return None
+
+    def get_node_hostnames(self):
+        """Get list of hostnames from SLURM allocation or fallback to localhost."""
+        node_hostnames = []
+
+        # Try to get hostnames from SLURM
+        try:
+            # Check if we're running in a SLURM allocation
+            if "SLURM_JOB_ID" in os.environ:
+                cmd = [
+                    "srun", "--overlap", "--oversubscribe", "scontrol", "show",
+                    "hostnames"
+                ]
+                print(f"Running command: {' '.join(cmd)}")
+                output = subprocess.check_output(cmd).decode().strip().split(
+                    '\n')
+                for hostname in output:
+                    if hostname.strip():
+                        node_hostnames.append(hostname.strip())
+
+                if not node_hostnames:
+                    raise Exception("No hostnames returned by SLURM")
+            else:
+                raise Exception(
+                    "Not running in a SLURM allocation (SLURM_JOB_ID not set)")
+        except Exception as e:
+            print(
+                f"Error getting node hostnames: {e}, falling back to localhost")
+            node_hostnames = ["localhost"]
+
+        print(f"Using node hostnames: {node_hostnames}")
+        return node_hostnames
+
+    def calculate_server_distribution(self):
+        """Calculate how to distribute context and generation servers across nodes."""
+        distribution = []
+
+        total_ctx_servers = self.config['exec']['config']['context'][
+            'dp'] if 'context' in self.config['exec']['config'] else 0
+        total_gen_servers = self.config['exec']['config']['generation'][
+            'dp'] if 'generation' in self.config['exec']['config'] else 0
+        total_ifb_servers = self.config['exec']['config']['ifb'][
+            'dp'] if 'ifb' in self.config['exec']['config'] else 0
+        port_base = 8001
+
+        gpu_index = 0
+        server_id = 0
+        j = 0
+        for ctx_server in range(total_ctx_servers):
+            gpu_indices = []
+            server_ids = []
+            server_ids.append(self.node_hostnames[server_id])
+            for i in range(self.config['exec']['config']['context']['tp']):
+                if i != 0 and i % self.num_gpus == 0:
+                    server_id += 1
+                    print(f'Adding {self.node_hostnames[server_id]}')
+                    server_ids.append(self.node_hostnames[server_id])
+                gpu_indices.append(str(gpu_index % self.num_gpus))
+                gpu_index += 1
+                gpu_index = gpu_index % self.num_gpus
+            if len(gpu_indices) > self.num_gpus:
+                gpu_indices = gpu_indices[:self.num_gpus]
+
+            distribution.append({
+                "node_id": server_id,
+                "hostnames": server_ids,
+                "gpu_indices": gpu_indices,
+                "is_generation": False,
+                "is_ifb": False,
+                "port": port_base + j
+            })
+            j += 1
+            if gpu_index % self.num_gpus == 0:
+                gpu_index = 0
+                server_id += 1
+
+        # If the last context server is not on a full node, add a generation server to the next node
+        if gpu_index % self.num_gpus != 0:
+            server_id += 1
+            gpu_index = 0
+
+        for gen_server in range(total_gen_servers):
+            gpu_indices = []
+            server_ids = []
+            server_ids.append(self.node_hostnames[server_id])
+            for i in range(self.config['exec']['config']['generation']['tp']):
+                if i != 0 and i % self.num_gpus == 0:
+                    server_id += 1
+                    server_ids.append(self.node_hostnames[server_id])
+                    print(f'Adding {self.node_hostnames[server_id]}')
+                gpu_indices.append(str(gpu_index % self.num_gpus))
+                gpu_index += 1
+                gpu_index = gpu_index % self.num_gpus
+
+            if len(gpu_indices) > self.num_gpus:
+                gpu_indices = gpu_indices[:self.num_gpus]
+
+            distribution.append({
+                "node_id": server_id,
+                "hostnames": server_ids,
+                "gpu_indices": gpu_indices,
+                "is_ifb": False,
+                "is_generation": True,
+                "port": port_base + j
+            })
+            j += 1
+            if gpu_index % self.num_gpus == 0:
+                gpu_index = 0
+                server_id += 1
+        for ifb_server in range(total_ifb_servers):
+            gpu_indices = []
+            server_ids = []
+            server_ids.append(self.node_hostnames[server_id])
+            for i in range(self.config['exec']['config']['ifb']['tp']):
+                if i != 0 and i % self.num_gpus == 0:
+                    server_id += 1
+                    server_ids.append(self.node_hostnames[server_id])
+                    print(f'Adding {self.node_hostnames[server_id]}')
+                gpu_indices.append(str(gpu_index % self.num_gpus))
+                gpu_index += 1
+                gpu_index = gpu_index % self.num_gpus
+            if len(gpu_indices) > self.num_gpus:
+                gpu_indices = gpu_indices[:self.num_gpus]
+
+            distribution.append({
+                "node_id": server_id,
+                "hostnames": server_ids,
+                "gpu_indices": gpu_indices,
+                "is_generation": False,
+                "is_ifb": True,
+                "port": port_base + j
+            })
+            j += 1
+        return distribution
+
+    def run_in_trtllm_container(self, cmd, num_tasks=1):
+        """Run a command in the TRTLLM container."""
+        cmd = [
+            "srun", "--overlap", "--oversubscribe", "-A", self.account, "-p",
+            self.partition, "-N", "1", "-n",
+            str(num_tasks), "-t", self.time, "--container-image",
+            self.container_image, "--container-mounts", self.mounts,
+            "--container-workdir", self.workdir, "--mpi", "pmix", "bash", "-c",
+            " ".join(cmd)
+        ]
+        return subprocess.run(cmd)
+
+    def spawn_in_trtllm_container(self, cmd, hostname):
+        """Spawn a command in the TRTLLM container."""
+        cmd = [
+            "srun", "--overlap", "--oversubscribe", "-A", self.account, "-p",
+            self.partition, "-N", "1", "-t", self.time, "--container-image",
+            self.container_image, "--container-mounts", self.mounts,
+            "--container-workdir", self.workdir, "--mpi", "pmix", "-w",
+            hostname, "bash", "-c", " ".join(cmd)
+        ]
+        return subprocess.Popen(cmd)
+
+    def write_node_distribution(self, distribution):
+        """Write the server distribution to a file for the disaggregated server."""
+        # Ensure the job_id directory exists
+        os.makedirs(f"job_{self.job_id}", exist_ok=True)
+
+        with open(f"job_{self.job_id}/node_distribution.yaml", "w") as f:
+            yaml.dump(distribution, f)
+
+    def launch_jobs(self):
+        """Launch SLURM jobs for all nodes."""
+        print("DEBUG JobManager: launch_jobs started")
+
+        # Check if results already exist and should be skipped
+        if hasattr(self.args, 'skip_existing') and self.args.skip_existing:
+            if self._check_existing_results():
+                print(
+                    f"Skipping job execution - results already exist in {self.output_folder}"
+                )
+                return
+
+        # Calculate server distribution
+        distribution = self.calculate_server_distribution()
+        self.write_node_distribution(distribution)
+
+        # Track all launched jobs
+        # Launch context and generation servers on each node
+        for node in distribution:
+            node_id = node["node_id"]
+            hostnames = node["hostnames"]
+            gpu_indices = node["gpu_indices"]
+            print(
+                f"Launching context server on {hostnames[0]}:{node['port']}, GPUs: {gpu_indices}"
+            )
+
+            # Launch context servers
+            if node["is_ifb"]:
+                if 'use_trtllm_bench' not in self.config[
+                        'profile'] or not self.config['profile'][
+                            'use_trtllm_bench']:
+                    log_file = f"ifb_server_{node_id}.log"
+                    ifb_jobs = self.launch_ifb_server(node_id, hostnames,
+                                                      gpu_indices, node["port"],
+                                                      log_file)
+                    self.ifb_jobs.extend(ifb_jobs)
+            elif not node["is_generation"]:
+                log_file = f"context_server_{node_id}.log"
+                ctx_jobs = self.launch_context_server(node_id, hostnames,
+                                                      gpu_indices, node["port"],
+                                                      log_file)
+                self.context_jobs.extend(ctx_jobs)
+            else:
+                log_file = f"generation_server_{node_id}.log"
+                gen_jobs = self.launch_generation_server(
+                    node_id, hostnames, gpu_indices, node["port"], log_file)
+                self.generation_jobs.extend(gen_jobs)
+            # Wait for server to start
+            time.sleep(15)
+
+        # Wait for all servers to start (give some time for them to initialize)
+        print("Waiting for all servers to start...")
+
+        for node in distribution:
+            hostnames = node["hostnames"]
+            print(f"Waiting for {hostnames[0]}:{node['port']} to start...")
+            if 'use_trtllm_bench' not in self.config[
+                    'profile'] or not self.config['profile']['use_trtllm_bench']:
+                wait_for_server(hostnames[0], node["port"])
+
+        # Collect server URLs and configure the disaggregated server
+        context_server_urls = []
+        generation_server_urls = []
+        ifb_server_urls = []
+
+        # Add all known server URLs based on the distribution
+        for node in distribution:
+            node_id = node["node_id"]
+            hostnames = node["hostnames"]
+            gpu_indices = node["gpu_indices"]
+
+            # Add context server URLs
+            node_port = node["port"]
+            if node["is_ifb"]:
+                ifb_server_urls.append(f"{hostnames[0]}:{node_port}")
+            elif not node["is_generation"]:
+                context_server_urls.append(f"{hostnames[0]}:{node_port}")
+            else:
+                generation_server_urls.append(f"{hostnames[0]}:{node_port}")
+
+        # Launch disaggregated server on the master node
+        distribution[0]["hostnames"][0]
+        port = 8000
+
+        if len(distribution) == 1 and distribution[0]["is_ifb"]:
+            distribution[0]["hostnames"][0]
+            port = distribution[0]["port"]
+        else:
+            for i in range(self.args.num_disagg_servers):
+                disagg_job = self.launch_disaggregated_server(
+                    self.node_hostnames[i], context_server_urls,
+                    generation_server_urls)
+                self.disagg_jobs.append(disagg_job)
+
+            # Wait a bit for the disaggregated server to initialize
+            print("Waiting for disaggregated server to initialize...")
+            time.sleep(30)
+
+        load_gen_jobs = []
+        # Run the loadgen tests
+        print("DEBUG JobManager: About to call run_loadgen_tests")
+        for i in range(self.args.num_disagg_servers):
+            load_gen_jobs.append(
+                self.run_loadgen_tests(self.node_hostnames[i], port,
+                                       generation_server_urls))
+        print("DEBUG JobManager: Completed run_loadgen_tests")
+
+        for loadgen_job in load_gen_jobs:
+            loadgen_job.wait()
+
+        
+        self.terminate_jobs(self.context_jobs + self.generation_jobs +
+                            self.ifb_jobs + self.disagg_jobs)
+        print("DEBUG JobManager: launch_jobs completed")
+    
+    def _build_server_launch_command(self, config, hostnames, gpu_indices,
+                                     node_port, config_path, log_file):
+        """Build the command to launch a server (context or generation).
+        
+        Returns:
+        cmd: The command to launch the server
+        """
+        if not gpu_indices or not config_path:
+            return None
+
+        # Get configuration based on server type
+        server_config = config
+        tp = server_config['tp']
+        ep = server_config['ep']
+        pp = server_config['pp']
+        max_batch_size = server_config['max_batch_size']
+        max_seq_len = server_config['max_seq_len']
+        max_num_tokens = server_config['max_num_tokens']
+        cuda_visible_devices = ",".join(gpu_indices)
+
+        print(
+            f"Launching server on {hostnames[0]}:{node_port}, GPUs: {cuda_visible_devices}"
+        )
+
+        envs = ""
+        envs += "TRTLLM_USE_UCX_KVCACHE=1 "
+        envs += " CUDA_VISIBLE_DEVICES=" + cuda_visible_devices
+        envs += " UCX_CUDA_IPC_ENABLE_MNNVL=n "
+        if 'gen_only' in config and config['gen_only']:
+            envs += " TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1"
+
+        if 'env' in config:
+            for key, value in config['env'].items():
+                envs += f" {key}={value}"
+        unset_env = ''
+        if 'unset_env' in config:
+            for key in config['unset_env']:
+                unset_env += f" -u {key}"
+        nsys = False
+        nsys_prefix = ''
+        if 'nsys' in config:
+            nsys = True
+        if nsys:
+            envs += " TLLM_PROFILE_RECORD_GC=1"
+            envs += " TLLM_NVTX_DEBUG=1"
+            if 'TLLM_PROFILE_START_STOP' in config['nsys']:
+                envs += f" TLLM_PROFILE_START_STOP={config['nsys']['TLLM_PROFILE_START_STOP']}"
+            nsys_file = os.path.join(self.output_folder, f"{log_file}.$SLURM_PROCID.nsys-rep")
+            nsys_prefix = f"nsys profile -e \"NSYS_MPI_STORE_TEAMS_PER_RANK=1\" -o {nsys_file} -f true -t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node --capture-range-end=stop --gpu-metrics-devices=none"
+        log_file = os.path.join(self.output_folder, log_file)
+
+        # Build the command
+        cmd = [
+            "srun",
+            "--overlap",
+            "--oversubscribe",
+            "-A",
+            self.account,
+            "-p",
+            self.partition,
+            "-N",
+            f"{len(hostnames)}",
+            "-n",
+            str(tp),
+            "-t",
+            self.time,
+            "--mpi",
+            "pmix",
+            "-w",
+            ",".join(hostnames),  # Specify the hostnames
+            f"--container-image={self.container_image}",
+            f"--container-mounts={self.mounts}",
+            f"--container-workdir={self.workdir}",
+            "--export=\"'CUDA_VISIBLE_DEVICES=" + cuda_visible_devices + "'\"",
+            "bash",
+            "-c",
+            envs +
+            " env " + unset_env + " " + nsys_prefix + " trtllm-llmapi-launch trtllm-serve " +
+            self.config['exec']['model_path'] + " --host 0.0.0.0 --port " +
+            str(node_port) + " --backend pytorch --extra_llm_api_options " +
+            config_path + " --tp_size " + str(tp) + " --ep_size " + str(ep) +
+            " --max_batch_size " + str(max_batch_size) + " --max_seq_len " +
+            str(max_seq_len) + " --max_num_tokens " + str(max_num_tokens) +
+            " --pp_size " + str(pp) + " &> " + log_file
+        ]
+
+        return cmd
+
+    def launch_context_server(self, node_id, hostnames, gpu_indices, node_port,
+                              log_file):
+        """Launch context server on a specific node."""
+        launched_jobs = []
+
+        cmd = self._build_server_launch_command(
+            self.config['exec']['config']['context'], hostnames, gpu_indices,
+            node_port, self.context_config_path, log_file)
+        if not cmd:
+            return launched_jobs
+
+        print(f"Running command: {' '.join(cmd)}")
+        proc = subprocess.Popen(cmd)
+        launched_jobs.append(proc)
+
+        # Small delay to avoid race conditions in resource allocation
+        time.sleep(2)
+
+        return launched_jobs
+
+    def launch_generation_server(self, node_id, hostnames, gpu_indices,
+                                 node_port, log_file):
+        """Launch generation server on a specific node."""
+        launched_jobs = []
+
+        cmd = self._build_server_launch_command(
+            self.config['exec']['config']['generation'], hostnames, gpu_indices,
+            node_port, self.generation_config_path, log_file)
+        if not cmd:
+            return launched_jobs
+
+        print(f"Running command: {' '.join(cmd)}")
+        proc = subprocess.Popen(cmd)
+        launched_jobs.append(proc)
+
+        return launched_jobs
+
+    def launch_ifb_server(self, node_id, hostnames, gpu_indices, node_port,
+                          log_file):
+        """Launch IFB server on a specific node."""
+        launched_jobs = []
+
+        cmd = self._build_server_launch_command(
+            self.config['exec']['config']['ifb'], hostnames, gpu_indices,
+            node_port, self.ifb_config_path, log_file)
+        if not cmd:
+            return launched_jobs
+
+        print(f"Running command: {' '.join(cmd)}")
+        proc = subprocess.Popen(cmd)
+        launched_jobs.append(proc)
+
+        return launched_jobs
+
+    def launch_disaggregated_server(self, hostname, context_servers,
+                                    generation_servers):
+        """Launch the disaggregated server on the master node."""
+        if not generation_servers:
+            print(
+                "No context or generation servers available. Cannot launch disaggregated server."
+            )
+            return None
+
+        config = self.config['exec']['config']
+        # Create the disaggregated server configuration
+        disagg_config = {
+            "hostname": hostname,
+            "port": 8000,
+            "backend": "pytorch",
+            "model": self.config['exec']['model_path'],
+            "context_servers": {
+                "num_instances":
+                len(context_servers),
+                "urls":
+                context_servers,
+                "tensor_parallel_size":
+                config['context']['tp'] if 'context' in config else 1,
+            },
+            "generation_servers": {
+                "num_instances":
+                len(generation_servers),
+                "urls":
+                generation_servers,
+                "tensor_parallel_size":
+                self.config['exec']['config']['generation']['tp'],
+            }
+        }
+
+        # Write the configuration to a file
+        disagg_config_path = os.path.join(f"job_{self.job_id}",
+                                          f"disagg_config_{hostname}.yaml")
+        with open(disagg_config_path, 'w') as f:
+            yaml.dump(disagg_config, f)
+        disagg_config_path = os.path.abspath(disagg_config_path)
+
+        # Also save a copy to the output folder for reference
+        if hasattr(self, 'output_folder') and self.output_folder:
+            output_config_path = os.path.join(self.output_folder,
+                                              f"disagg_config_{hostname}.yaml")
+            try:
+                with open(output_config_path, 'w') as f:
+                    yaml.dump(disagg_config,
+                              f,
+                              default_flow_style=False,
+                              indent=2)
+                print(
+                    f"Disaggregated server configuration saved to {output_config_path}"
+                )
+            except Exception as e:
+                print(
+                    f"Warning: Could not save disaggregated config to output folder: {e}"
+                )
+
+        print(f"Launching disaggregated server on {hostname}:8000")
+        envs = ""
+        gen_only = self.config['exec']['config']['generation'].get(
+            'gen_only', False)
+        if gen_only:
+            envs += " TRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1"
+
+        # Store the log to the output folder
+        log_file = os.path.join(self.output_folder,
+                                f"disagg_server_{hostname}.log")
+        # Build the command
+        cmd = [
+            "srun",
+            "--overlap",
+            "--oversubscribe",
+            "-A",
+            self.account,
+            "-p",
+            self.partition,
+            "-N",
+            "1",
+            "-n",
+            "1",
+            "-t",
+            self.time,
+            "-w",
+            hostname,  # Specify the hostname
+            "-J",
+            f"{self.job_name}_disagg",
+            "--mpi",
+            "pmix",
+            f"--container-image={self.container_image}",
+            f"--container-mounts={self.mounts}",
+            f"--container-workdir={self.workdir}",
+            "--export=OMPI_ALLOW_RUN_AS_ROOT=1",
+            "bash",
+            "-c",
+            f" {envs} trtllm-serve disaggregated -c " + disagg_config_path +
+            " --request_timeout " + " 1800000" + " &> " + log_file
+        ]
+
+        print(f"Running command: {' '.join(cmd)}")
+        return subprocess.Popen(cmd)
+
+    def run_loadgen_tests(self, hostname, port, generation_server_urls):
+        """Run loadgen tests on the master node."""
+        print(
+            f"DEBUG JobManager: run_loadgen_tests called with hostname={hostname}, port={port}"
+        )
+        output_folder = self.output_folder
+
+        # Run loadgen tests for each concurrency level
+        if 'use_trtllm_bench' in self.config['profile'] and self.config[
+                'profile']['use_trtllm_bench']:
+            self.prepare_dataset(output_folder)
+
+        concurrency_levels = self.config['profile']['concurrency']
+        print(
+            f"DEBUG JobManager: Will run loadgen for concurrency levels: {concurrency_levels}"
+        )
+
+        for concurrency in concurrency_levels:
+            print(
+                f"DEBUG JobManager: Starting loadgen for concurrency {concurrency}"
+            )
+            concurrency_output_folder = f"{output_folder}/{concurrency}"
+            if not os.path.exists(concurrency_output_folder):
+                os.makedirs(concurrency_output_folder)
+
+            # Build and run the loadgen command
+            if 'use_trtllm_bench' in self.config['profile'] and self.config[
+                    'profile']['use_trtllm_bench']:
+                self.run_trtllm_bench_for_concurrency(
+                    hostname, port, concurrency, concurrency_output_folder)
+            elif 'use_benchmark_serving' in self.config[
+                    'profile'] and self.config['profile'][
+                        'use_benchmark_serving']:
+                pid = self.run_benchmark_serving_for_concurrency(
+                    hostname, port, concurrency, concurrency_output_folder)
+                pid.wait()
+            print(
+                f"DEBUG JobManager: Completed loadgen for concurrency {concurrency}"
+            )
+
+        print("DEBUG JobManager: run_loadgen_tests completed")
+        return pid
+
+    def record_kv_cache_stats(self, generation_server_urls, output_folder):
+        """Record the kv cache stats for a specific concurrency level."""
+        print(
+            f"DEBUG JobManager: Recording kv cache stats for generation servers"
+        )
+        # Get the kv cache stats
+        # Send a request to "/metrics" and store the output in a file
+        for url in generation_server_urls:
+            response = requests.get(f"http://{url}/metrics")
+            with open(
+                    os.path.join(output_folder,
+                                 f"{url.split(':')[0]}_kv_cache_stats.txt"),
+                    "w") as f:
+                f.write(response.text)
+        print(
+            f"DEBUG JobManager: Recorded kv cache stats for generation servers")
+
+    def prepare_dataset(self, output_folder):
+        """Prepare the dataset for the loadgen tests."""
+        input_tokens = str(self.config['profile']['isl'])
+        output_tokens = str(self.config['profile']['osl'])
+        num_prompts = str(
+            8 * self.config['exec']['config']['ifb']['max_batch_size'])
+
+        # If the dataset file already exists, skip the preparation
+        dataset_file = os.path.join(output_folder, "dataset.json")
+        if os.path.exists(dataset_file):
+            return
+
+        cmd = [
+            "python3", self.args.prepare_dataset_script, "--stdout",
+            "--tokenizer", self.config['exec']['model_path'], "token-norm-dist",
+            "--input-mean", input_tokens, "--output-mean", output_tokens,
+            "--input-stdev", "0", "--output-stdev", "0", "--num-requests",
+            num_prompts, ">", dataset_file
+        ]
+        print(f"Running command: {' '.join(cmd)}")
+        self.run_in_trtllm_container(cmd)
+
+    def run_trtllm_bench_for_concurrency(self, hostname, port, concurrency,
+                                         output_folder):
+        """Run the trtllm bench for a specific concurrency level."""
+        config = self.config['exec']['config']['ifb']
+        tp = config['tp'] if 'tp' in config else 0
+        ep = config['ep'] if 'ep' in config else 0
+        max_batch_size = config[
+            'max_batch_size'] if 'max_batch_size' in config else 0
+        max_num_tokens = config[
+            'max_num_tokens'] if 'max_num_tokens' in config else 0
+        gpu_fraction = config['config']['kv_cache_config'][
+            'free_gpu_memory_fraction'] if 'config' in config and 'kv_cache_config' in config[
+                'config'] else 0
+        num_requests = concurrency * 8
+        # Parent of output_folder is the dataset file
+        dataset_file = os.path.join(os.path.dirname(output_folder),
+                                    "dataset.json")
+        # Redirect the output to a file
+        cmd = [
+            "trtllm-llmapi-launch", "trtllm-bench", "-m",
+            "deepseek-ai/DeepSeek-R1", "--model_path",
+            f"{self.config['exec']['model_path']}", "throughput", "--tp",
+            f"{tp}", "--ep", f"{ep}", "--warmup", "0", "--dataset",
+            f"{dataset_file}", "--backend", "pytorch", "--max_batch_size",
+            f"{max_batch_size}", "--max_num_tokens", f"{max_num_tokens}",
+            "--kv_cache_free_gpu_mem_fraction", f"{gpu_fraction}",
+            "--extra_llm_api_options", f"{self.ifb_config_path}",
+            "--num_requests", f"{num_requests}", "--concurrency",
+            f"{concurrency}", "--report_json",
+            f"{output_folder}/trtllm_bench_throughput.json", "--streaming",
+            "&>" + os.path.join(output_folder, "output.log")
+        ]
+        print(f"Running command: {' '.join(cmd)}")
+        self.run_in_trtllm_container(cmd, tp)
+
+    def run_benchmark_serving_for_concurrency(self, hostname, port, concurrency,
+                                              output_folder):
+        """Run the benchmark serving test for a specific concurrency level."""
+        model_name = self.config['exec']['model_path']
+        dataset_path = self.config['profile']['dataset_path']
+        num_prompts = self.config['profile']['num_prompts']
+        cmd = [
+            "/usr/bin/python3", "-m",
+            "tensorrt_llm.serve.scripts.benchmark_serving", "--model",
+            model_name, "--tokenizer", model_name, "--dataset-name",
+            "trtllm_custom", "--dataset-path", dataset_path, "--num-prompts",
+            str(num_prompts), "--host", hostname, "--port",
+            str(port), "--max-concurrency",
+            str(concurrency)
+        ]
+        if 'ignore_eos' in self.config['profile'] and self.config['profile'][
+                'ignore_eos']:
+            cmd += ["--ignore-eos"]
+        cmd += ["&> " + os.path.join(output_folder, "benchmark_serving.log")]
+        print(f"Running command: {' '.join(cmd)}")
+        return self.spawn_in_trtllm_container(cmd, hostname)
+
+    def terminate_jobs(self, jobs):
+        """Terminate all running jobs."""
+        
+        # Then terminate the jobs
+        for job in jobs:
+            if job and job.poll() is None:
+                try:
+                    job.terminate()
+                    job.wait(timeout=5)
+                except:
+                    job.kill()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.terminate_jobs(self.context_jobs + self.generation_jobs +
+                            self.ifb_jobs + self.disagg_jobs)
+
+    def save_config_to_output_folder(self):
+        """Save the complete configuration used for this run to the output folder."""
+        if not hasattr(self, 'output_folder') or not self.output_folder:
+            return
+
+        config_file_path = os.path.join(self.output_folder, "config.yaml")
+        try:
+            with open(config_file_path, 'w') as f:
+                yaml.dump(self.config, f, default_flow_style=False, indent=2)
+            print(f"Configuration saved to {config_file_path}")
+        except Exception as e:
+            print(
+                f"Warning: Could not save configuration to output folder: {e}")
+
+    def prepare_and_run_single_config(self):
+        """Prepare and run a single configuration."""
+        print("DEBUG JobManager: prepare_and_run_single_config started")
+
+        # Create the configuration files
+        self.context_config_path = self.create_context_config()
+        self.generation_config_path = self.create_generation_config()
+        self.ifb_config_path = self.create_ifb_config()
+
+        # Save the complete configuration to the output folder
+        self.save_config_to_output_folder()
+
+        # Get hostnames of available nodes
+        self.node_hostnames = self.get_node_hostnames()
diff --git a/examples/disaggregated/slurm/disagg_profiler/sweeper.py b/examples/disaggregated/slurm/disagg_profiler/sweeper.py
new file mode 100644
index 00000000000..9469ff7284e
--- /dev/null
+++ b/examples/disaggregated/slurm/disagg_profiler/sweeper.py
@@ -0,0 +1,1327 @@
+import argparse
+import copy
+import itertools
+import json
+import math
+import os
+import re
+import subprocess
+import sys
+import time
+from functools import partial
+from multiprocessing import Pool
+
+import yaml
+
+from .job_manager import JobManager, calculate_nodes_needed
+
+
+def get_slurm_allocation(account, partition, time_limit, job_name, num_nodes):
+    """
+    Request a SLURM allocation and wait for it to be granted.
+
+    Args:
+        account (str): SLURM account
+        partition (str): SLURM partition
+        time_limit (str): Time limit for the allocation
+        job_name (str): Job name
+        num_nodes (int): Number of nodes to request
+
+    Returns:
+        str: Job ID of the SLURM allocation, or None if allocation failed
+    """
+    print(
+        f"Requesting SLURM allocation for {num_nodes} nodes on partition {partition}..."
+    )
+
+    # capture SLURM_NODELIST inside allocation
+    cmd = [
+        "salloc", "-A", account, "-p", partition, "-N",
+        str(num_nodes), "-t", time_limit, "-J", job_name, "--no-shell"
+    ]
+
+    try:
+        # Run the command and stream the output in real-time
+        process = subprocess.Popen(cmd,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.STDOUT,
+                                   text=True,
+                                   bufsize=1,
+                                   universal_newlines=True)
+
+        output_lines = []
+        job_id = None
+
+        # Process the output in real-time
+        for line in process.stdout:
+            print(line.strip())
+            output_lines.append(line)
+
+            # Look for job ID in the output (e.g., "salloc: Granted job allocation 12345")
+            job_id_match = re.search(r"Granted job allocation (\d+)", line)
+            if job_id_match:
+                job_id = job_id_match.group(1)
+                print(
+                    f"Successfully obtained SLURM allocation with job ID: {job_id}"
+                )
+                break
+
+        # Wait for the process to complete
+        process.wait()
+
+        if job_id:
+            return job_id
+        else:
+            print(
+                "Failed to obtain SLURM allocation. Check the output above for details."
+            )
+            return None
+    except Exception as e:
+        print(f"Error requesting SLURM allocation: {e}")
+        return None
+
+
+def run_sweep_configuration(config_params, args_dict):
+    """
+    Run a single sweep configuration in a separate process.
+
+    Args:
+        config_params (tuple): (config, description, config_num, total_configs, timestamp_base)
+        args_dict (dict): Dictionary of arguments needed to run the configuration
+
+    Returns:
+        bool: True if successful, False otherwise
+    """
+    config, description, config_num, total_configs, timestamp_base = config_params
+
+    print(f"Running configuration {config_num}/{total_configs}: {description}")
+
+    # Reconstruct args from dict
+    args = argparse.Namespace()
+    for key, value in args_dict.items():
+        setattr(args, key, value)
+
+    # Calculate nodes needed for this specific configuration
+    nodes_needed = calculate_nodes_needed(config, args.num_gpus)
+    print(f"Configuration requires {nodes_needed} nodes")
+
+    # Set a unique job ID with timestamp, index, and sweep parameters for better tracking
+    job_id = f"sweep_{timestamp_base}_{config_num}of{total_configs}_{description}"
+
+    # If allocation requested and not using a single allocation, allocate resources for just this configuration
+    slurm_job_id = None
+    if args.request_allocation and not args.use_single_allocation:
+        print(
+            f"Requesting allocation for configuration {config_num}/{total_configs}"
+        )
+        slurm_job_id = get_slurm_allocation(
+            args.account, args.partition, args.time,
+            f"{args.job_name}_sweep{config_num}", nodes_needed)
+
+        if not slurm_job_id:
+            print(
+                f"Failed to get allocation for configuration {config_num}/{total_configs}, skipping"
+            )
+            return False
+
+        # Set SLURM_JOB_ID for this run
+        os.environ["SLURM_JOB_ID"] = slurm_job_id
+
+    try:
+        # Create and run the job manager with this configuration
+        # JobManager will automatically check for existing results if --skip-existing is set
+        with JobManager(args,
+                        load_config_only=False,
+                        override_config=config,
+                        override_job_id=job_id) as job_manager:
+            job_manager.launch_jobs()
+        return True
+    except Exception as e:
+        print(f"Error running configuration {config_num}: {e}")
+        return False
+    finally:
+        # Release the allocation if we requested one for this specific configuration
+        if slurm_job_id:
+            print(
+                f"Releasing SLURM allocation (job ID: {slurm_job_id}) for configuration {config_num}"
+            )
+            subprocess.run(["scancel", slurm_job_id])
+            print(f"SLURM allocation released for configuration {config_num}")
+
+
+class AutoSweeper:
+    """
+    Automatically determines optimal sweep configurations by profiling context
+    and generation servers separately to balance their request processing rates.
+
+    This class:
+    1. Profiles context-only servers
+    2. Profiles generation-only servers
+    3. Analyzes the performance ratios
+    4. Generates sweep configurations to balance the rates
+    """
+
+    def __init__(self, args, max_gpus=None):
+        """
+        Initialize the auto sweeper.
+
+        Args:
+            args (argparse.Namespace): Command line arguments
+            max_gpus (int): Maximum number of GPUs to use across all servers
+        """
+        self.args = args
+        self.max_gpus = max_gpus or 32  # Default limit
+
+        # Load base configuration
+        temp_job_manager = JobManager(args, load_config_only=True)
+        self.base_config = temp_job_manager.base_config
+
+        # Storage for profiling results
+        self.context_results = {}
+        self.generation_results = {}
+
+    def create_context_only_config(self, base_config, dp=1, tp=1):
+        """
+        Create a context-only configuration from a disaggregated config.
+
+        Args:
+            base_config (dict): Base disaggregated configuration
+            dp (int): Data parallelism for context server
+            tp (int): Tensor parallelism for context server
+
+        Returns:
+            dict: Context-only configuration
+        """
+        config = copy.deepcopy(base_config)
+
+        # Remove generation and ifb servers
+        if 'generation' in config['exec']['config']:
+            del config['exec']['config']['generation']
+        if 'ifb' in config['exec']['config']:
+            del config['exec']['config']['ifb']
+
+        # Configure context server - get base params from existing config or defaults
+        base_context = config['exec']['config'].get('context', {})
+
+        config['exec']['config']['ifb'] = {
+            'tp': tp,
+            'dp': dp,
+            'ep': base_context.get('ep', 1),
+            'pp': base_context.get('pp', 1),
+            'max_batch_size': base_context.get('max_batch_size', 512),
+            'max_seq_len': base_context.get('max_seq_len', 8192),
+            'max_num_tokens': base_context.get('max_num_tokens', 8192),
+            'config': base_context.get('config')
+        }
+        if 'context' in config['exec']['config']:
+            del config['exec']['config']['context']
+
+        # Set OSL to 1 for context profiling (we only care about context processing, not generation)
+        if 'profile' in config:
+            config['profile']['osl'] = 1
+            print(
+                f"Set OSL=1 for context profiling (was {config['profile'].get('osl', 'unknown')})"
+            )
+
+        return config
+
+    def create_generation_only_config(self, base_config, dp=1, tp=1):
+        """
+        Create a generation-only configuration from a disaggregated config.
+
+        Args:
+            base_config (dict): Base disaggregated configuration
+            dp (int): Data parallelism for generation server
+            tp (int): Tensor parallelism for generation server
+
+        Returns:
+            dict: Generation-only configuration
+        """
+        config = copy.deepcopy(base_config)
+
+        # Remove context servers
+        if 'context' in config['exec']['config']:
+            del config['exec']['config']['context']
+
+        # Get base params from existing generation config or use ifb config as template
+        base_gen = config['exec']['config'].get('generation', {})
+        base_ifb = config['exec']['config'].get('ifb', {})
+
+        # Use generation config as base, fallback to ifb, then defaults
+        template_config = base_gen if base_gen else base_ifb
+
+        # Configure as IFB (generation-only) server
+        config['exec']['config']['generation'] = {
+            'tp': tp,
+            'dp': dp,
+            'ep': template_config.get('ep', 1),
+            'pp': template_config.get('pp', 1),
+            'max_batch_size': template_config.get('max_batch_size', 512),
+            'max_seq_len': template_config.get('max_seq_len', 8192),
+            'max_num_tokens': template_config.get('max_num_tokens', 8192),
+            'gen_only': True,  # Set generation-only mode
+            'config': template_config.get('config')
+        }
+
+        # Remove the generation config since we're using IFB for gen-only
+        if 'ifb' in config['exec']['config']:
+            del config['exec']['config']['ifb']
+
+        return config
+
+    def parse_loadgen_results(self, output_folder):
+        """
+        Parse loadgen results to extract Requests/s for each concurrency level.
+
+        Args:
+            output_folder (str): Path to output folder containing results
+
+        Returns:
+            dict: Mapping of concurrency -> requests_per_second
+        """
+        results = {}
+
+        # Look for concurrency subdirectories
+        if not os.path.exists(output_folder):
+            print(f"Output folder {output_folder} does not exist")
+            return results
+
+        for item in os.listdir(output_folder):
+            concurrency_path = os.path.join(output_folder, item)
+            if os.path.isdir(concurrency_path) and item.isdigit():
+                concurrency = int(item)
+
+                benchmark_serving_path = os.path.join(concurrency_path,
+                                                      "benchmark_serving.log")
+                if os.path.exists(benchmark_serving_path):
+                    with open(benchmark_serving_path, 'r') as f:
+                        for line in f:
+                            if "Request throughput (req/s):" in line:
+                                rps = float(
+                                    line.split("Request throughput (req/s):")
+                                    [1].strip())
+                                results[concurrency] = rps
+                                print(
+                                    f"Concurrency {concurrency}: {rps:.3f} Requests/s"
+                                )
+                else:
+                    # Look for raw_data.json file
+                    raw_data_path = os.path.join(concurrency_path,
+                                                 "raw_data.json")
+                    if os.path.exists(raw_data_path):
+                        try:
+                            with open(raw_data_path, 'r') as f:
+                                data = json.load(f)
+
+                            if 'results' in data and 'infbench_summary' in data[
+                                    'results']:
+                                rps = data['results']['infbench_summary'][
+                                    'Requests/s']
+                                results[concurrency] = rps
+                                print(
+                                    f"Concurrency {concurrency}: {rps:.3f} Requests/s"
+                                )
+                        except Exception as e:
+                            print(f"Error parsing {raw_data_path}: {e}")
+
+        return results
+
+    def profile_context_servers(self, tp_values=None, dp_values=None):
+        """
+        Profile context-only servers with different TP/DP configurations.
+
+        Args:
+            tp_values (list): List of TP values to test (if None, uses default from config)
+            dp_values (list): List of DP values to test (if None, uses default from config)
+        """
+        if tp_values is None:
+            # Extract default TP value from base configuration
+            default_tp = 1
+            if 'context' in self.base_config['exec']['config']:
+                default_tp = self.base_config['exec']['config']['context'].get(
+                    'tp', 1)
+            elif 'ifb' in self.base_config['exec']['config']:
+                default_tp = self.base_config['exec']['config']['ifb'].get(
+                    'tp', 1)
+            tp_values = [default_tp]
+            print(f"Using default context TP value: {default_tp}")
+        if dp_values is None:
+            # Extract default DP value from base configuration
+            default_dp = 1
+            if 'context' in self.base_config['exec']['config']:
+                default_dp = self.base_config['exec']['config']['context'].get(
+                    'dp', 1)
+            elif 'ifb' in self.base_config['exec']['config']:
+                default_dp = self.base_config['exec']['config']['ifb'].get(
+                    'dp', 1)
+            dp_values = [default_dp]
+            print(f"Using default context DP value: {default_dp}")
+
+        print("=== Profiling Context Servers ===")
+        print(f"DEBUG: tp_values = {tp_values}, dp_values = {dp_values}")
+        print(
+            f"DEBUG: Will run {len(tp_values)} x {len(dp_values)} = {len(tp_values) * len(dp_values)} iterations"
+        )
+
+        iteration_count = 0
+        for tp in tp_values:
+            for dp in dp_values:
+                iteration_count += 1
+                print(
+                    f"DEBUG: Starting iteration {iteration_count} with TP={tp}, DP={dp}"
+                )
+
+                gpus_needed = tp * dp
+                if gpus_needed > self.max_gpus:
+                    print(
+                        f"Skipping context TP={tp}, DP={dp} (needs {gpus_needed} GPUs > {self.max_gpus})"
+                    )
+                    continue
+
+                print(f"Profiling context server: TP={tp}, DP={dp}")
+
+                # Create context-only config
+                config = self.create_context_only_config(self.base_config,
+                                                         dp=dp,
+                                                         tp=tp)
+
+                # Calculate nodes needed for this configuration
+                nodes_needed = calculate_nodes_needed(config,
+                                                      self.args.num_gpus)
+
+                # Request allocation for this specific configuration
+                slurm_job_id = None
+                if hasattr(
+                        self.args,
+                        'request_allocation') and self.args.request_allocation:
+                    print(
+                        f"Requesting allocation for context TP={tp}, DP={dp} ({nodes_needed} nodes)"
+                    )
+                    slurm_job_id = get_slurm_allocation(
+                        self.args.account, self.args.partition, self.args.time,
+                        f"{self.args.job_name}_ctx_tp{tp}_dp{dp}", nodes_needed)
+
+                    if not slurm_job_id:
+                        print(
+                            f"Failed to get allocation for context TP={tp}, DP={dp}, skipping"
+                        )
+                        self.context_results[(tp, dp)] = {}
+                        continue
+
+                    # Set SLURM_JOB_ID for this run
+                    os.environ["SLURM_JOB_ID"] = slurm_job_id
+                    print(
+                        f"Using allocation {slurm_job_id} for context TP={tp}, DP={dp}"
+                    )
+
+                # Set job ID for this configuration
+                job_id = f"ctx_profile_tp{tp}_dp{dp}"
+
+                try:
+                    # Run the profiling
+                    # JobManager will automatically check for existing results if --skip-existing is set
+                    with JobManager(self.args,
+                                    load_config_only=False,
+                                    override_config=config,
+                                    override_job_id=job_id) as job_manager:
+                        job_manager.launch_jobs()
+
+                    # Parse results
+                    results = self.parse_loadgen_results(
+                        job_manager.output_folder)
+                    self.context_results[(tp, dp)] = results
+
+                    print(f"Context TP={tp}, DP={dp} results: {results}")
+
+                except Exception as e:
+                    print(f"Error profiling context TP={tp}, DP={dp}: {e}")
+                    self.context_results[(tp, dp)] = {}
+                finally:
+                    # Release the allocation if we requested one
+                    if slurm_job_id:
+                        print(
+                            f"Releasing allocation {slurm_job_id} for context TP={tp}, DP={dp}"
+                        )
+                        subprocess.run(["scancel", slurm_job_id])
+                        # Clear the environment variable
+                        if "SLURM_JOB_ID" in os.environ:
+                            del os.environ["SLURM_JOB_ID"]
+
+        print(
+            f"DEBUG: Completed context profiling with {iteration_count} total iterations"
+        )
+
+    def profile_generation_servers(self, tp_values=None, dp_values=None):
+        """
+        Profile generation-only servers with different TP/DP configurations.
+
+        Args:
+            tp_values (list): List of TP values to test (if None, uses default from config)
+            dp_values (list): List of DP values to test (if None, uses default from config)
+        """
+        if tp_values is None:
+            # Extract default TP value from base configuration
+            default_tp = 1
+            if 'generation' in self.base_config['exec']['config']:
+                default_tp = self.base_config['exec']['config'][
+                    'generation'].get('tp', 1)
+            elif 'ifb' in self.base_config['exec']['config']:
+                default_tp = self.base_config['exec']['config']['ifb'].get(
+                    'tp', 1)
+            tp_values = [default_tp]
+            print(f"Using default generation TP value: {default_tp}")
+        if dp_values is None:
+            # Extract default DP value from base configuration
+            default_dp = 1
+            if 'generation' in self.base_config['exec']['config']:
+                default_dp = self.base_config['exec']['config'][
+                    'generation'].get('dp', 1)
+            elif 'ifb' in self.base_config['exec']['config']:
+                default_dp = self.base_config['exec']['config']['ifb'].get(
+                    'dp', 1)
+            dp_values = [default_dp]
+            print(f"Using default generation DP value: {default_dp}")
+
+        print("=== Profiling Generation Servers ===")
+        print(f"DEBUG: tp_values = {tp_values}, dp_values = {dp_values}")
+        print(
+            f"DEBUG: Will run {len(tp_values)} x {len(dp_values)} = {len(tp_values) * len(dp_values)} iterations"
+        )
+
+        iteration_count = 0
+        for tp in tp_values:
+            for dp in dp_values:
+                iteration_count += 1
+                print(
+                    f"DEBUG: Starting iteration {iteration_count} with TP={tp}, DP={dp}"
+                )
+
+                gpus_needed = tp * dp
+                if gpus_needed > self.max_gpus:
+                    print(
+                        f"Skipping generation TP={tp}, DP={dp} (needs {gpus_needed} GPUs > {self.max_gpus})"
+                    )
+                    continue
+
+                print(f"Profiling generation server: TP={tp}, DP={dp}")
+
+                # Create generation-only config
+                config = self.create_generation_only_config(self.base_config,
+                                                            dp=dp,
+                                                            tp=tp)
+
+                # Calculate nodes needed for this configuration
+                nodes_needed = calculate_nodes_needed(config,
+                                                      self.args.num_gpus)
+
+                # Request allocation for this specific configuration
+                slurm_job_id = None
+                if hasattr(
+                        self.args,
+                        'request_allocation') and self.args.request_allocation:
+                    print(
+                        f"Requesting allocation for generation TP={tp}, DP={dp} ({nodes_needed} nodes)"
+                    )
+                    slurm_job_id = get_slurm_allocation(
+                        self.args.account, self.args.partition, self.args.time,
+                        f"{self.args.job_name}_gen_tp{tp}_dp{dp}", nodes_needed)
+
+                    if not slurm_job_id:
+                        print(
+                            f"Failed to get allocation for generation TP={tp}, DP={dp}, skipping"
+                        )
+                        self.generation_results[(tp, dp)] = {}
+                        continue
+
+                    # Set SLURM_JOB_ID for this run
+                    os.environ["SLURM_JOB_ID"] = slurm_job_id
+                    print(
+                        f"Using allocation {slurm_job_id} for generation TP={tp}, DP={dp}"
+                    )
+
+                # Set job ID for this configuration
+                job_id = f"gen_profile_tp{tp}_dp{dp}"
+
+                try:
+                    # Run the profiling
+                    # JobManager will automatically check for existing results if --skip-existing is set
+                    with JobManager(self.args,
+                                    load_config_only=False,
+                                    override_config=config,
+                                    override_job_id=job_id) as job_manager:
+                        job_manager.launch_jobs()
+
+                    # Parse results
+                    results = self.parse_loadgen_results(
+                        job_manager.output_folder)
+                    self.generation_results[(tp, dp)] = results
+
+                    print(f"Generation TP={tp}, DP={dp} results: {results}")
+
+                except Exception as e:
+                    print(f"Error profiling generation TP={tp}, DP={dp}: {e}")
+                    self.generation_results[(tp, dp)] = {}
+                finally:
+                    # Release the allocation if we requested one
+                    if slurm_job_id:
+                        print(
+                            f"Releasing allocation {slurm_job_id} for generation TP={tp}, DP={dp}"
+                        )
+                        subprocess.run(["scancel", slurm_job_id])
+                        # Clear the environment variable
+                        if "SLURM_JOB_ID" in os.environ:
+                            del os.environ["SLURM_JOB_ID"]
+
+        print(
+            f"DEBUG: Completed generation profiling with {iteration_count} total iterations"
+        )
+
+    def calculate_optimal_ratios(self):
+        """
+        Calculate optimal server ratios based on profiling results.
+
+        Returns:
+            list: List of optimal configurations with ratios (deduplicated)
+        """
+        print("=== Calculating Optimal Ratios ===")
+
+        all_configs = []
+
+        # For each concurrency level, find the best combinations
+        all_concurrencies = set()
+
+        # Collect all concurrency levels from both results
+        for results in self.context_results.values():
+            all_concurrencies.update(results.keys())
+        for results in self.generation_results.values():
+            all_concurrencies.update(results.keys())
+
+        for concurrency in sorted(all_concurrencies):
+            print(f"\nAnalyzing concurrency level: {concurrency}")
+
+            # Compare each context config with each generation config
+            for (ctx_tp, ctx_dp), ctx_results in self.context_results.items():
+                if concurrency not in ctx_results:
+                    continue
+
+                ctx_rps = ctx_results[concurrency]
+
+                for (gen_tp,
+                     gen_dp), gen_results in self.generation_results.items():
+                    if concurrency not in gen_results:
+                        continue
+
+                    gen_rps = gen_results[concurrency]
+
+                    if gen_rps <= 0:
+                        continue
+
+                    # Calculate the ratio of context to generation throughput
+                    ratio = ctx_rps / gen_rps
+
+                    # Determine optimal server counts
+                    if ratio < 1:
+                        # Context is slower (bottleneck) - need more context servers
+                        ctx_servers_needed = math.ceil(1 / ratio)
+                        gen_servers_needed = 1
+                    else:
+                        # Generation is slower (bottleneck) - need more generation servers
+                        gen_servers_needed = math.ceil(ratio)
+                        ctx_servers_needed = 1
+
+                    total_gpus = (ctx_servers_needed * ctx_tp * ctx_dp +
+                                  gen_servers_needed * gen_tp * gen_dp)
+
+                    if total_gpus <= self.max_gpus:
+                        config = {
+                            'concurrency':
+                            concurrency,
+                            'ctx_tp':
+                            ctx_tp,
+                            'ctx_dp':
+                            ctx_dp,
+                            'ctx_servers':
+                            ctx_servers_needed,
+                            'ctx_rps':
+                            ctx_rps,
+                            'gen_tp':
+                            gen_tp,
+                            'gen_dp':
+                            gen_dp,
+                            'gen_servers':
+                            gen_servers_needed,
+                            'gen_rps':
+                            gen_rps,
+                            'ratio':
+                            ratio,
+                            'total_gpus':
+                            total_gpus,
+                            'balanced_rps':
+                            min(ctx_rps * ctx_servers_needed,
+                                gen_rps * gen_servers_needed)
+                        }
+                        all_configs.append(config)
+
+        # Deduplicate configurations based on actual server configuration
+        # Key: (ctx_tp, ctx_dp, ctx_servers, gen_tp, gen_dp, gen_servers)
+        unique_configs = {}
+
+        for config in all_configs:
+            # Create a key based on the actual configuration parameters that matter
+            config_key = (config['ctx_tp'], config['ctx_dp'],
+                          config['ctx_servers'], config['gen_tp'],
+                          config['gen_dp'], config['gen_servers'])
+
+            # Keep the configuration with the highest balanced RPS
+            if config_key not in unique_configs or config[
+                    'balanced_rps'] > unique_configs[config_key]['balanced_rps']:
+                unique_configs[config_key] = config
+
+        # Convert back to list and sort by balanced throughput
+        optimal_configs = list(unique_configs.values())
+        optimal_configs.sort(key=lambda x: x['balanced_rps'], reverse=True)
+
+        print(
+            f"\nFound {len(all_configs)} total configurations across all concurrency levels"
+        )
+        print(
+            f"After deduplication: {len(optimal_configs)} unique configurations"
+        )
+
+        for config in optimal_configs:
+            print(
+                f"  Config: Ctx({config['ctx_tp']},{config['ctx_dp']})x{config['ctx_servers']} "
+                f"+ Gen({config['gen_tp']},{config['gen_dp']})x{config['gen_servers']} "
+                f"= {config['balanced_rps']:.2f} RPS using {config['total_gpus']} GPUs "
+                f"(from concurrency {config['concurrency']})")
+
+        return optimal_configs
+
+    def generate_sweep_configurations(self, optimal_configs):
+        """
+        Generate sweep configurations based on optimal ratios.
+
+        Args:
+            optimal_configs (list): List of optimal configurations
+
+        Returns:
+            list: List of sweep configurations
+        """
+        print("=== Generating Sweep Configurations ===")
+
+        sweep_configs = []
+
+        for config in optimal_configs:
+            # Create the disaggregated configuration
+            disagg_config = copy.deepcopy(self.base_config)
+
+            # Get base configs for proper parameter inheritance
+            base_context = self.base_config['exec']['config'].get('context', {})
+            base_generation = self.base_config['exec']['config'].get(
+                'generation', {})
+
+            # Configure context servers
+            disagg_config['exec']['config']['context'] = {
+                'tp': config['ctx_tp'],
+                'dp': config['ctx_servers'],  # Use server count as DP
+                'ep': base_context.get('ep', 1),
+                'pp': base_context.get('pp', 1),
+                'max_batch_size': base_context.get('max_batch_size', 512),
+                'max_seq_len': base_context.get('max_seq_len', 8192),
+                'max_num_tokens': base_context.get('max_num_tokens', 8192),
+                'config': base_context.get('config', {})
+            }
+
+            # Configure generation servers
+            disagg_config['exec']['config']['generation'] = {
+                'tp': config['gen_tp'],
+                'dp': config['gen_servers'],  # Use server count as DP
+                'ep': base_generation.get('ep', 1),
+                'pp': base_generation.get('pp', 1),
+                'max_batch_size': base_generation.get('max_batch_size', 512),
+                'max_seq_len': base_generation.get('max_seq_len', 8192),
+                'max_num_tokens': base_generation.get('max_num_tokens', 8192),
+                'gen_only': False,
+                'config': base_generation.get('config', {})
+            }
+
+            # Remove ifb config if it exists (we're using disaggregated mode)
+            if 'ifb' in disagg_config['exec']['config']:
+                del disagg_config['exec']['config']['ifb']
+
+            description = (
+                f"c{config['concurrency']}_ctx{config['ctx_tp']}x{config['ctx_servers']}_"
+                f"gen{config['gen_tp']}x{config['gen_servers']}")
+
+            sweep_configs.append((disagg_config, description))
+
+        print(f"Generated {len(sweep_configs)} sweep configurations")
+        return sweep_configs
+
+    def run_auto_sweep(self,
+                       ctx_tp_values=None,
+                       ctx_dp_values=None,
+                       gen_tp_values=None,
+                       gen_dp_values=None):
+        """
+        Run the complete auto-sweep process.
+
+        Args:
+            ctx_tp_values (list): TP values to test for context servers (if None, uses default from config)
+            ctx_dp_values (list): DP values to test for context servers (if None, uses default from config)
+            gen_tp_values (list): TP values to test for generation servers (if None, uses default from config)
+            gen_dp_values (list): DP values to test for generation servers (if None, uses default from config)
+        """
+        print(f"Starting auto-sweep with max {self.max_gpus} GPUs")
+        print(
+            f"DEBUG: Called with ctx_tp_values={ctx_tp_values}, ctx_dp_values={ctx_dp_values}"
+        )
+        print(
+            f"DEBUG: Called with gen_tp_values={gen_tp_values}, gen_dp_values={gen_dp_values}"
+        )
+
+        # Step 1: Profile context servers
+        print("DEBUG: Step 1 - About to call profile_context_servers")
+        self.profile_context_servers(ctx_tp_values, ctx_dp_values)
+        print("DEBUG: Step 1 - Completed profile_context_servers")
+
+        # Step 2: Profile generation servers
+        print("DEBUG: Step 2 - About to call profile_generation_servers")
+        self.profile_generation_servers(gen_tp_values, gen_dp_values)
+        print("DEBUG: Step 2 - Completed profile_generation_servers")
+
+        # Step 3: Calculate optimal ratios
+        print("DEBUG: Step 3 - About to calculate optimal ratios")
+        optimal_configs = self.calculate_optimal_ratios()
+
+        if not optimal_configs:
+            print("No optimal configurations found!")
+            return
+
+        # Step 4: Generate sweep configurations
+        print("DEBUG: Step 4 - About to generate sweep configurations")
+        sweep_configs = self.generate_sweep_configurations(optimal_configs)
+
+        # Step 5: Run the sweep configurations
+        if sweep_configs:
+            print("=== Running Optimal Sweep Configurations ===")
+
+            # Create a parameter sweeper to run the generated configurations
+            # Ensure we use separate allocations for each configuration
+            temp_args = copy.deepcopy(self.args)
+            temp_args.use_single_allocation = False  # Force separate allocations
+
+            sweeper = ParameterSweeper(
+                "", temp_args)  # Empty config since we provide directly
+            sweeper.run_generated_sweeps(sweep_configs)
+        else:
+            print("No sweep configurations generated!")
+
+
+class ParameterSweeper:
+    """
+    Manages parameter sweeps for the disaggregated serving setup.
+
+    This class handles the parsing of sweep configurations, generates all combinations
+    of parameter values to be tested, and runs the corresponding jobs using JobManager.
+    It abstracts away the details of parameter manipulation and provides a clean interface
+    for running multiple configurations.
+
+    Example usage:
+        sweeper = ParameterSweeper('sweep_config.yaml', args)
+        sweeper.run_sweeps()
+    """
+
+    def __init__(self, sweep_config_path, args):
+        """
+        Initialize the parameter sweeper.
+
+        Args:
+            sweep_config_path (str): Path to the sweep configuration file or a string
+                                    representation of the sweep configuration.
+            args (argparse.Namespace): Command line arguments.
+        """
+        self.sweep_config_path = sweep_config_path
+        self.args = args
+        self.job_manager = None
+        self.sweep_params = self.parse_sweep_config()
+
+    def parse_sweep_config(self):
+        """
+        Parse the sweep configuration string into a list of parameter dictionaries.
+
+        The sweep configuration can be either:
+        - A path to a YAML file containing the sweep definition
+        - A string representation of a Python list of dictionaries
+        - Empty string (for generated configs)
+
+        Returns:
+            List of dictionaries, where each dictionary defines a sweep.
+        """
+        # If no sweep config provided (for generated configs), return empty list
+        if not self.sweep_config_path:
+            return []
+
+        # Load the sweep configuration from the file or string
+        if os.path.exists(self.sweep_config_path):
+            with open(self.sweep_config_path, 'r') as f:
+                sweep_params = yaml.safe_load(f)
+        else:
+            # Try to parse as a Python dictionary
+            try:
+                sweep_params = eval(self.sweep_config_path)
+            except Exception as e:
+                print(
+                    f"Error parsing sweep configuration: {self.sweep_config_path}"
+                )
+                print(f"Exception: {e}")
+                sys.exit(1)
+
+        return sweep_params
+
+    def run_sweeps(self):
+        """
+        Generate all sweep configurations and run jobs for each one.
+
+        This is the main method that orchestrates the sweeping process:
+        1. Load the base configuration
+        2. Generate all configurations from sweep parameters
+        3. Run a job for each configuration with appropriate allocation
+        """
+        print(
+            f"Running parameter sweep with configuration: {self.sweep_config_path}"
+        )
+
+        # Create a temporary job manager to load the base config
+        temp_args = copy.deepcopy(self.args)
+        temp_args.sweep_config = None  # Disable sweep for the temporary job manager
+
+        temp_job_manager = JobManager(temp_args, load_config_only=True)
+        base_config = temp_job_manager.base_config
+
+        # Generate all configurations
+        all_configs = self.generate_sweep_configurations(base_config)
+
+        # Run the configurations
+        self._run_sweep_configurations(all_configs)
+
+    def run_generated_sweeps(self, sweep_configs):
+        """
+        Run sweep configurations that are already generated.
+
+        Args:
+            sweep_configs (list): List of (config, description) tuples
+        """
+        print(f"Running {len(sweep_configs)} generated sweep configurations")
+
+        # Convert to the expected format
+        all_configs = sweep_configs
+
+        # Run the configurations
+        self._run_sweep_configurations(all_configs)
+
+    def _run_sweep_configurations(self, all_configs):
+        """
+        Internal method to run a list of sweep configurations.
+
+        Args:
+            all_configs (list): List of (config, description) tuples
+        """
+        # Check if any sweep configuration needs more nodes than the base configuration
+        base_config = all_configs[0][0] if all_configs else {}
+        base_nodes = calculate_nodes_needed(base_config)
+        max_nodes = base_nodes
+
+        for config, _ in all_configs:
+            config_nodes = calculate_nodes_needed(config)
+            if config_nodes > max_nodes:
+                max_nodes = config_nodes
+
+        if max_nodes > base_nodes:
+            print(
+                f"Warning: Some sweep configurations require more nodes ({max_nodes}) than the base configuration ({base_nodes})."
+            )
+
+        total_configs = len(all_configs)
+        timestamp_base = time.strftime("%Y%m%d-%H%M%S")
+
+        # If using a single allocation for all sweeps, get it now
+        single_allocation_job_id = None
+        if self.args.request_allocation and self.args.use_single_allocation:
+            print(
+                f"Requesting a single allocation for all sweep configurations")
+            print(f"Using maximum node count: {max_nodes}")
+            single_allocation_job_id = get_slurm_allocation(
+                self.args.account, self.args.partition, self.args.time,
+                f"{self.args.job_name}_sweep_all", max_nodes)
+
+            if not single_allocation_job_id:
+                print(
+                    "Failed to get allocation for all configurations, exiting")
+                return
+
+            # Set SLURM_JOB_ID for this run
+            os.environ["SLURM_JOB_ID"] = single_allocation_job_id
+            print(
+                f"Using single allocation with job ID: {single_allocation_job_id} for all configurations"
+            )
+
+        # Determine number of parallel processes to use
+        num_parallel = min(self.args.parallel_sweeps, total_configs)
+        print(
+            f"Running {total_configs} configurations with {num_parallel} parallel processes"
+        )
+
+        # Prepare config parameters for parallel execution
+        config_params = []
+        for idx, (config, description) in enumerate(all_configs):
+            config_num = idx + 1
+            config_params.append((config, description, config_num,
+                                  total_configs, timestamp_base))
+
+        # Convert args to dictionary for pickling
+        args_dict = vars(copy.deepcopy(self.args))
+        args_dict[
+            'sweep_config'] = None  # Disable sweep for the individual jobs
+
+        try:
+            # Run configurations in parallel if requested
+            if num_parallel > 1:
+                print(
+                    f"Starting parallel execution with {num_parallel} processes"
+                )
+                with Pool(processes=num_parallel) as pool:
+                    results = pool.map(
+                        partial(run_sweep_configuration, args_dict=args_dict),
+                        config_params)
+
+                # Report results
+                successful = sum(1 for r in results if r)
+                print(
+                    f"Completed {successful}/{total_configs} configurations successfully"
+                )
+            else:
+                # Run configurations sequentially
+                print("Starting sequential execution")
+                for params in config_params:
+                    run_sweep_configuration(params, args_dict)
+        finally:
+            # Release the single allocation if we requested one
+            if single_allocation_job_id:
+                print(
+                    f"Releasing single SLURM allocation (job ID: {single_allocation_job_id})"
+                )
+                subprocess.run(["scancel", single_allocation_job_id])
+                print("SLURM allocation released.")
+
+    def generate_sweep_configurations(self, base_config):
+        """
+        Generate all configurations based on the sweep parameters.
+
+        Args:
+            base_config (dict): Base configuration to apply parameter changes to.
+
+        Returns:
+            List of (config, description) tuples, where config is the modified configuration
+            and description is a string describing the parameter values.
+        """
+        all_configs = []
+
+        # Process each sweep
+        for sweep_idx, sweep in enumerate(self.sweep_params):
+            param_names = []
+            param_values = []
+
+            for param_path, values in sweep.items():
+                param_names.append(param_path)
+                param_values.append(values)
+
+            print(
+                f"Generating configurations for sweep {sweep_idx+1}/{len(self.sweep_params)}"
+            )
+
+            # Generate all combinations for this sweep
+            for idx, combination in enumerate(itertools.product(*param_values)):
+                # Create a new config for this combination
+                config = copy.deepcopy(base_config)
+
+                # Apply the parameter values to the config
+                param_description = []
+                for i, param_path in enumerate(param_names):
+                    value = combination[i]
+                    self.set_config_value(config, param_path, value)
+                    param_description.append(
+                        f"{param_path.split('.')[-1]}={value}")
+
+                combination_desc = "_".join(param_description)
+                all_configs.append((config, combination_desc))
+
+        return all_configs
+
+    def set_config_value(self, config, param_path, value):
+        """
+        Set a configuration value based on a dot-notation path.
+
+        Supports synchronized parameters using the "sync" prefix:
+        - sync.tp: Sets both context.tp and generation.tp to the same value
+        - sync.dp: Sets both context.dp and generation.dp to the same value
+        - sync.max_batch_size: Sets both context.max_batch_size and generation.max_batch_size
+
+        Args:
+            config (dict): Configuration to modify
+            param_path (str): Parameter path in dot notation
+            value: Value to set
+        """
+        parts = param_path.split('.')
+
+        # Handle synchronized parameters (sync prefix)
+        if parts[0] == 'sync' and len(parts) > 1:
+            sync_param = '.'.join(
+                parts[1:])  # Join all parts after 'sync' to handle nested paths
+
+            # Apply to both context and generation servers
+            context_path = f"context.{sync_param}"
+            generation_path = f"generation.{sync_param}"
+
+            print(
+                f"Synchronizing parameter {sync_param}={value} across context and generation servers"
+            )
+
+            # Set context server parameter
+            if 'context' in config.get('exec', {}).get('config', {}):
+                self.set_config_value(config, context_path, value)
+
+            # Set generation server parameter
+            if 'generation' in config.get('exec', {}).get('config', {}):
+                self.set_config_value(config, generation_path, value)
+
+            return
+
+        # Special parameter name mappings for commonly used shortened names
+        if parts[0] == 'context' and len(parts) > 1:
+            parts = ['exec', 'config', 'context'] + parts[1:]
+        elif parts[0] == 'generation' and len(parts) > 1:
+            parts = ['exec', 'config', 'generation'] + parts[1:]
+        elif parts[0] == 'gen' and len(parts) > 1:
+            parts = ['exec', 'config', 'generation'] + parts[1:]
+        elif parts[0] == 'ifb' and len(parts) > 1:
+            parts = ['exec', 'config', 'ifb'] + parts[1:]
+
+        # Navigate to the correct position in the config
+        current = config
+        for part in parts[:-1]:
+            if part not in current:
+                current[part] = {}
+            current = current[part]
+
+        # Set the value
+        current[parts[-1]] = value
+        print(f"Set {param_path} to {value}")
+
+
+class MultiConfigSweeper:
+    """
+    Handles multiple configurations for context and generation servers, testing all combinations.
+
+    This class supports configurations where context and generation servers can each have
+    multiple predefined configurations, and all combinations of these configurations are tested.
+
+    The configuration format allows either:
+    1. Single configuration (existing format)
+    2. List of configurations for exhaustive combination testing
+
+    Example multi-config format:
+    exec:
+      config:
+        context:
+          - tp: 2
+            dp: 1
+            max_batch_size: 128
+            config: {...}
+          - tp: 4
+            dp: 1
+            max_batch_size: 256
+            config: {...}
+        generation:
+          - tp: 2
+            dp: 2
+            max_batch_size: 64
+            config: {...}
+          - tp: 4
+            dp: 1
+            max_batch_size: 128
+            config: {...}
+    """
+
+    def __init__(self, args):
+        """
+        Initialize the multi-config sweeper.
+
+        Args:
+            args (argparse.Namespace): Command line arguments
+        """
+        self.args = args
+
+        # Load base configuration
+        temp_job_manager = JobManager(args, load_config_only=True)
+        self.base_config = temp_job_manager.base_config
+
+    def is_multi_config(self, config):
+        """
+        Check if the configuration uses multiple configurations format.
+
+        Args:
+            config (dict): Configuration to check
+
+        Returns:
+            bool: True if multi-config format is detected
+        """
+        exec_config = config.get('exec', {}).get('config', {})
+
+        # Check if context or generation sections are lists
+        for section in ['context', 'generation', 'ifb']:
+            if section in exec_config and isinstance(exec_config[section],
+                                                     list):
+                return True
+
+        return False
+
+    def extract_multi_configs(self, config):
+        """
+        Extract multiple configurations from the config.
+
+        Args:
+            config (dict): Configuration with potential multi-config format
+
+        Returns:
+            tuple: (context_configs, generation_configs, ifb_configs) where each is a list
+        """
+        exec_config = config.get('exec', {}).get('config', {})
+
+        # Extract context configurations
+        context_configs = []
+        if 'context' in exec_config:
+            if isinstance(exec_config['context'], list):
+                context_configs = exec_config['context']
+            else:
+                context_configs = [exec_config['context']]
+
+        # Extract generation configurations
+        generation_configs = []
+        if 'generation' in exec_config:
+            if isinstance(exec_config['generation'], list):
+                generation_configs = exec_config['generation']
+            else:
+                generation_configs = [exec_config['generation']]
+
+        # Extract ifb configurations
+        ifb_configs = []
+        if 'ifb' in exec_config:
+            if isinstance(exec_config['ifb'], list):
+                ifb_configs = exec_config['ifb']
+            else:
+                ifb_configs = [exec_config['ifb']]
+
+        return context_configs, generation_configs, ifb_configs
+
+    def generate_config_combinations(self, base_config):
+        """
+        Generate all combinations of configurations.
+
+        Args:
+            base_config (dict): Base configuration
+
+        Returns:
+            list: List of (config, description) tuples for all combinations
+        """
+        if not self.is_multi_config(base_config):
+            print(
+                "Configuration is not in multi-config format, returning single config"
+            )
+            return [(base_config, "single_config")]
+
+        context_configs, generation_configs, ifb_configs = self.extract_multi_configs(
+            base_config)
+
+        combinations = []
+        combination_id = 0
+
+        # Generate combinations based on available server types
+        if context_configs and generation_configs:
+            # Disaggregated mode: context + generation combinations
+            print(
+                f"Generating combinations for disaggregated mode: {len(context_configs)} context × {len(generation_configs)} generation"
+            )
+
+            for ctx_idx, ctx_config in enumerate(context_configs):
+                for gen_idx, gen_config in enumerate(generation_configs):
+                    combination_id += 1
+
+                    # Create combined configuration
+                    combined_config = copy.deepcopy(base_config)
+                    combined_config['exec']['config']['context'] = ctx_config
+                    combined_config['exec']['config']['generation'] = gen_config
+
+                    # Remove ifb if present (disaggregated mode)
+                    if 'ifb' in combined_config['exec']['config']:
+                        del combined_config['exec']['config']['ifb']
+
+                    # Generate description
+                    description = f"ctx{ctx_idx+1}_gen{gen_idx+1}"
+
+                    combinations.append((combined_config, description))
+
+        elif ifb_configs:
+            # IFB mode: single server combinations
+            print(
+                f"Generating combinations for IFB mode: {len(ifb_configs)} configurations"
+            )
+
+            for ifb_idx, ifb_config in enumerate(ifb_configs):
+                combination_id += 1
+
+                # Create combined configuration
+                combined_config = copy.deepcopy(base_config)
+                combined_config['exec']['config']['ifb'] = ifb_config
+
+                # Remove context and generation if present (IFB mode)
+                for section in ['context', 'generation']:
+                    if section in combined_config['exec']['config']:
+                        del combined_config['exec']['config'][section]
+
+                # Generate description
+                description = f"ifb{ifb_idx+1}"
+
+                combinations.append((combined_config, description))
+
+        else:
+            print("Warning: No valid server configurations found")
+            return [(base_config, "fallback_config")]
+
+        print(f"Generated {len(combinations)} configuration combinations")
+        return combinations
+
+    def run_multi_config_sweep(self):
+        """
+        Run sweep with all configuration combinations.
+        """
+        print("=== Multi-Configuration Sweep Mode ===")
+        print(
+            "Generating all combinations of context and generation server configurations"
+        )
+
+        # Generate all configuration combinations
+        all_configs = self.generate_config_combinations(self.base_config)
+
+        if not all_configs:
+            print("No configurations generated, exiting")
+            return
+
+        print(f"Will run {len(all_configs)} configuration combinations:")
+        for i, (config, description) in enumerate(all_configs, 1):
+            print(f"  {i}: {description}")
+
+        # Use existing sweep infrastructure to run all combinations
+        parameter_sweeper = ParameterSweeper(
+            "", self.args)  # Empty sweep config path
+        parameter_sweeper.run_generated_sweeps(all_configs)
diff --git a/examples/disaggregated/slurm/launcher.py b/examples/disaggregated/slurm/launcher.py
new file mode 100644
index 00000000000..777d49319c8
--- /dev/null
+++ b/examples/disaggregated/slurm/launcher.py
@@ -0,0 +1,266 @@
+import argparse
+import os
+import subprocess
+import sys
+
+from disagg_profiler.job_manager import JobManager, calculate_nodes_needed
+from disagg_profiler.sweeper import (AutoSweeper, MultiConfigSweeper,
+                                     ParameterSweeper, get_slurm_allocation)
+
+
+def setup_argparse():
+    """Set up command line argument parsing."""
+    parser = argparse.ArgumentParser(
+        description="TRT-LLM Disaggregated Serving Launcher")
+
+    # SLURM job arguments
+    parser.add_argument("--account",
+                        type=str,
+                        required=True,
+                        help="SLURM account")
+    parser.add_argument("--partition",
+                        type=str,
+                        required=True,
+                        help="SLURM partition")
+    parser.add_argument("--time",
+                        type=str,
+                        default="04:00:00",
+                        help="SLURM time limit")
+    parser.add_argument("--job-name",
+                        type=str,
+                        required=True,
+                        help="SLURM job name")
+    parser.add_argument("--container-image",
+                        type=str,
+                        required=True,
+                        help="Container image")
+    parser.add_argument("--mounts",
+                        type=str,
+                        required=True,
+                        help="Container mount points")
+    parser.add_argument("--num-gpus",
+                        type=int,
+                        required=True,
+                        help="Number of GPUs on each node")
+    parser.add_argument("--config-file",
+                        type=str,
+                        default="./config.yaml",
+                        help="Path to the configuration file")
+    parser.add_argument("--experiment-path",
+                        type=str,
+                        default=".",
+                        help="Path to the configuration file")
+    parser.add_argument(
+        "--prepare-dataset-script",
+        type=str,
+        default=None,
+        help=
+        "Path to the prepare dataset script. Required if trtllm_bench is used.")
+    parser.add_argument(
+        "--sweep-config",
+        type=str,
+        help=
+        """Parameter sweep configuration as a Python string or a path to a YAML file.
+
+                        Format:
+                        - A list of dictionaries, where each dictionary defines a sweep
+                        - Each dictionary maps parameter paths to lists of values to try
+
+                        Example: '[{"context.dp": [1, 2, 3, 4]}, {"gen.dp": [1, 2, 3]}]'
+
+                        This will first sweep context.dp with values 1-4 (keeping gen.dp=1),
+                        then sweep gen.dp with values 1-3 (keeping context.dp=1).
+
+                        Parameter paths can use shortened names:
+                        - context.dp (for exec.config.context.dp)
+                        - gen.dp (for exec.config.generation.dp)
+                        - ifb.dp (for exec.config.ifb.dp)
+
+                        Or full paths like exec.config.context.dp
+
+                        For multiple parameters, use:
+                        '[{"context.tp": [2, 4, 8], "context.dp": [1, 2, 4]}]'
+                        This will try all combinations of tp and dp (9 combinations total).
+                        """)
+    parser.add_argument(
+        "--auto-sweep",
+        action="store_true",
+        help=
+        """Automatically determine optimal sweep configurations by profiling context and generation servers.
+
+                        This mode will:
+                        1. Profile context-only servers with different TP/DP configurations
+                        2. Profile generation-only servers with different TP/DP configurations
+                        3. Analyze the performance ratios to determine optimal server counts
+                        4. Generate and run sweep configurations that balance context and generation throughput
+
+                        Auto-sweep uses separate SLURM allocations for each configuration, providing better
+                        resource utilization and fault tolerance compared to a single large allocation.
+
+                        The goal is to ensure the request processing rates of context and generation servers are balanced.
+                        If ctx_rps / gen_rps < 1, more generation servers are needed to match context throughput.
+                        If ctx_rps / gen_rps > 1, more context servers are needed to match generation throughput.
+
+                        Both the profiling phases and final sweep execution can be parallelized using --parallel-sweeps.
+                        """)
+    parser.add_argument(
+        "--max-gpus",
+        type=int,
+        default=32,
+        help=
+        "Maximum number of GPUs to use across all servers in auto-sweep mode (default: 32)"
+    )
+    parser.add_argument(
+        "--request-allocation",
+        action="store_true",
+        help=
+        "Request a SLURM allocation before running jobs. This is useful for interactive testing."
+    )
+    parser.add_argument(
+        "--use-single-allocation",
+        action="store_true",
+        help=
+        "When running sweeps, use a single allocation for all configurations instead of separate allocations for each."
+    )
+    parser.add_argument(
+        "--parallel-sweeps",
+        type=int,
+        default=1,
+        help="""Number of sweep configurations to run in parallel (default: 1).
+
+                        This applies to:
+                        - Manual sweep mode: final sweep configurations
+                        - Auto-sweep mode: profiling phases AND final sweep configurations
+
+                        For auto-sweep, this parallelizes both context server profiling, generation server profiling,
+                        and the final optimal configuration sweep. This can significantly reduce total runtime.
+                        """)
+    parser.add_argument("--skip-existing",
+                        action="store_true",
+                        help="""Skip configurations that already have results.
+
+                        This is useful for resuming interrupted sweeps or re-running experiments.
+                        The launcher will check for existing output folders with results and skip
+                        those configurations, only running new or incomplete ones.
+
+                        This applies to:
+                        - Manual sweep mode: skip sweep configurations with existing results
+                        - Auto-sweep mode: skip profiling runs AND final sweep configurations with existing results
+                        """)
+    parser.add_argument("--num-disagg-servers",
+                        type=int,
+                        default=1,
+                        help="Number of disaggregated servers to use.")
+
+    return parser
+
+
+def main():
+    parser = setup_argparse()
+    args = parser.parse_args()
+
+    # Load base configuration to check for multi-config format
+    temp_job_manager = JobManager(args, load_config_only=True)
+    base_config = temp_job_manager.base_config
+
+    # Auto-detect multi-config format
+    multi_sweeper = MultiConfigSweeper(args)
+    is_multi_config = multi_sweeper.is_multi_config(base_config)
+
+    # Validate arguments - can't use multiple sweep modes at once
+    sweep_modes = [args.auto_sweep, bool(args.sweep_config), is_multi_config]
+    active_modes = []
+    if args.auto_sweep:
+        active_modes.append("--auto-sweep")
+    if args.sweep_config:
+        active_modes.append("--sweep-config")
+    if is_multi_config:
+        active_modes.append("multi-config (auto-detected)")
+
+    if sum(sweep_modes) > 1:
+        print("Error: Cannot use multiple sweep modes at the same time")
+        print(f"Detected modes: {', '.join(active_modes)}")
+        print(
+            "Choose one of: --auto-sweep, --sweep-config, or define single configurations (not lists) in your config file"
+        )
+        sys.exit(1)
+
+    # Handle allocation requests for different modes
+    job_id = None
+    if args.request_allocation:
+        if args.auto_sweep:
+            # For auto-sweep, use separate allocations for each configuration
+            # This provides better resource utilization and fault tolerance
+            print(
+                "Auto-sweep mode: will request separate allocations for each configuration"
+            )
+            num_nodes = None  # No upfront allocation
+        elif args.sweep_config and args.use_single_allocation:
+            # For manual sweep with single allocation, calculate based on base config
+            # This is a conservative estimate - actual sweep configs might need more
+            base_nodes = calculate_nodes_needed(base_config, args.num_gpus)
+            print(
+                f"Manual sweep mode: using {base_nodes} nodes based on base configuration"
+            )
+            num_nodes = base_nodes
+        elif not args.sweep_config and not is_multi_config:
+            # For single configuration mode
+            num_nodes = calculate_nodes_needed(base_config, args.num_gpus)
+            print(
+                f"Single config mode: calculated {num_nodes} nodes needed based on configuration"
+            )
+        else:
+            # For manual sweep without single allocation or multi-config, no pre-allocation needed
+            print(
+                "Multi-config or manual sweep mode without single allocation - allocations will be requested per configuration"
+            )
+            num_nodes = None
+
+        if num_nodes is not None:
+            # Request the allocation
+            job_id = get_slurm_allocation(args.account, args.partition,
+                                          args.time, args.job_name, num_nodes)
+
+            # Exit if allocation failed
+            if not job_id:
+                print("Exiting due to allocation failure.")
+                sys.exit(1)
+
+            # Set the SLURM_JOB_ID environment variable to ensure the jobs use this allocation
+            os.environ["SLURM_JOB_ID"] = job_id
+            print(f"Using SLURM allocation with job ID: {job_id}")
+
+    try:
+        # Auto-sweep mode: automatically determine optimal configurations
+        if args.auto_sweep:
+            print("=== Auto-Sweep Mode ===")
+            auto_sweeper = AutoSweeper(args, max_gpus=args.max_gpus)
+            auto_sweeper.run_auto_sweep()
+        # Multi-config mode: test all combinations of predefined configurations (auto-detected)
+        elif is_multi_config:
+            print("=== Multi-Configuration Mode (Auto-Detected) ===")
+            print(
+                "Detected lists in configuration file - will test all combinations"
+            )
+            multi_sweeper.run_multi_config_sweep()
+        # Manual sweep mode: use provided sweep config
+        elif args.sweep_config:
+            print("=== Manual Sweep Mode ===")
+            sweeper = ParameterSweeper(args.sweep_config, args)
+            sweeper.run_sweeps()
+        else:
+            # Standard single-run mode
+            print("=== Single Configuration Mode ===")
+            with JobManager(args) as job_manager:
+                job_manager.launch_jobs()
+    finally:
+        # Release the allocation if we requested one and it's not a manual sweep (which handles its own allocations)
+        if job_id and (not args.sweep_config
+                       or args.use_single_allocation) and not is_multi_config:
+            print(f"Releasing SLURM allocation (job ID: {job_id})...")
+            subprocess.run(["scancel", job_id])
+            print("SLURM allocation released.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/disaggregated/slurm/slurm_populate_urls.py b/examples/disaggregated/slurm/slurm_populate_urls.py
deleted file mode 100644
index abe8122dbe5..00000000000
--- a/examples/disaggregated/slurm/slurm_populate_urls.py
+++ /dev/null
@@ -1,164 +0,0 @@
-import argparse
-import os
-import re
-
-import yaml
-
-# Parse command line arguments
-parser = argparse.ArgumentParser(
-    description='Update YAML configuration with SLURM node information.')
-parser.add_argument(
-    '--nodelist_env_var',
-    type=str,
-    default='SLURM_JOB_NODELIST',
-    help=
-    'Name of the env var that provides the list of nodes as dev[7-8,11,13] for example'
-)
-parser.add_argument(
-    '--tasks_per_node_env_var',
-    type=str,
-    default='SLURM_TASKS_PER_NODE',
-    help=
-    'Name of the env var that provides the tasks per node as 8(x3),2 for example'
-)
-parser.add_argument('--disagg_server_port',
-                    type=int,
-                    default=8000,
-                    help='The port to use for disagg server')
-parser.add_argument('--worker_start_port',
-                    type=int,
-                    default=8001,
-                    help='The starting port to use for workers')
-parser.add_argument('--input_yaml',
-                    type=str,
-                    default='config.yaml',
-                    help='Path to the input YAML file')
-parser.add_argument('--output_yaml',
-                    type=str,
-                    default='output_config.yaml',
-                    help='Path to the output YAML file')
-args = parser.parse_args()
-
-# Parse SLURM_JOB_NODELIST and SLURM_TASKS_PER_NODE from environment variables
-print("---")
-slurm_job_nodelist = os.getenv(args.nodelist_env_var, '')
-if not slurm_job_nodelist:
-    raise ValueError(f"Environment variable {args.nodelist_env_var} not found.")
-print(f"{args.nodelist_env_var}: {slurm_job_nodelist}")
-slurm_tasks_per_node = os.getenv(args.tasks_per_node_env_var, '')
-if not slurm_tasks_per_node:
-    raise ValueError(
-        f"Environment variable {args.tasks_per_node_env_var} not found.")
-print(f"{args.tasks_per_node_env_var}: {slurm_tasks_per_node}")
-print("---")
-
-# Generate list of nodes
-node_prefix = re.match(r'^[a-zA-Z]+', slurm_job_nodelist).group(0)
-node_range = re.search(r'\[(.*?)\]', slurm_job_nodelist).group(1)
-nodes = []
-for part in node_range.split(','):
-    if '-' in part:
-        start, end = map(int, part.split('-'))
-        nodes.extend([f"{node_prefix}{i}" for i in range(start, end + 1)])
-    else:
-        nodes.append(f"{node_prefix}{part}")
-print(f"Nodes: {nodes}")
-
-# Generate tasks per node
-tasks_per_node = []
-for part in slurm_tasks_per_node.split(','):
-    if '(x' in part:
-        count, repeat = map(int, re.findall(r'\d+', part))
-        tasks_per_node.extend([count] * repeat)
-    else:
-        tasks_per_node.append(int(part))
-print(f"Tasks_per_node: {tasks_per_node}")
-
-if (len(tasks_per_node) != len(nodes)):
-    raise ValueError(
-        f"Number of nodes and tasks per node do not match. Number of nodes: {len(nodes)}, Number of tasks per node: {len(tasks_per_node)}"
-    )
-
-max_tasks_per_node = max(tasks_per_node)
-task_nodes = []
-for node, tasks in zip(nodes, tasks_per_node):
-    task_nodes.extend([node] * tasks)
-
-print(f"Task nodes: {task_nodes}")
-print("---")
-
-
-# Function to generate URLs
-def generate_urls(ctx_or_gen,
-                  num_instances,
-                  tensor_parallel_size,
-                  pipeline_parallel_size,
-                  max_task_per_node,
-                  nodes,
-                  task_nodes,
-                  node_to_port,
-                  task_nodes_offset=0):
-    urls = []
-
-    for instance in range(num_instances):
-        tasks_needed = tensor_parallel_size * pipeline_parallel_size
-
-        if (task_nodes_offset + tasks_needed) > len(task_nodes):
-            print(f"{ctx_or_gen} urls so far: {urls}")
-            raise ValueError(
-                f"For {ctx_or_gen} instance {instance}, there are not enough tasks available. task_nodes_offset: {task_nodes_offset}, tasks_needed: {tasks_needed}, len(task_nodes): {len(task_nodes)}"
-            )
-
-        # Minimum number of nodes needed for that instance
-        min_node = (tasks_needed + max_tasks_per_node - 1) / max_tasks_per_node
-        instance_nodes = set(task_nodes[task_nodes_offset:task_nodes_offset +
-                                        tasks_needed])
-        if len(instance_nodes) > min_node:
-            raise ValueError(
-                f"Tasks for a instance {instance} of {ctx_or_gen} instances use more node than expected. Nodes used: {instance_nodes}, number of nodes expected: {min_node}, max_tasks_per_node: {max_tasks_per_node}"
-            )
-
-        node = task_nodes[task_nodes_offset]
-        port = node_to_port[node]
-        node_to_port[node] += 1
-        task_nodes_offset += tasks_needed
-
-        urls.append(f"{node}:{port}")
-
-    print(f"{ctx_or_gen} urls: {urls}")
-    return urls, task_nodes_offset
-
-
-# Load the YAML file
-with open(args.input_yaml, 'r') as file:
-    config = yaml.safe_load(file)
-
-# Keep track of the port number for each node
-node_ports = {}
-for node in nodes:
-    node_ports[node] = args.worker_start_port
-
-# Generate URLs for context_servers and generation_servers
-context_urls, task_node_offset = generate_urls(
-    "ctx", config['context_servers']['num_instances'],
-    config['context_servers']['tensor_parallel_size'],
-    config['context_servers']['pipeline_parallel_size'], max_tasks_per_node,
-    nodes, task_nodes, node_ports)
-
-generation_urls, _ = generate_urls(
-    "gen", config['generation_servers']['num_instances'],
-    config['generation_servers']['tensor_parallel_size'],
-    config['generation_servers']['pipeline_parallel_size'], max_tasks_per_node,
-    nodes, task_nodes, node_ports, task_node_offset)
-
-# Update the YAML configuration
-config['hostname'] = nodes[0]
-config['port'] = args.disagg_server_port
-config['context_servers']['urls'] = context_urls
-config['generation_servers']['urls'] = generation_urls
-
-# Save the updated YAML file
-with open(args.output_yaml, 'w') as file:
-    yaml.safe_dump(config, file, sort_keys=False)
-
-print("YAML file updated successfully.")