[None] [fix] Minor fixes to slurm and benchmark scripts (#7453)

Signed-off-by: Kaiyu Xie <[email protected]>

Author: kaiyux

examples/disaggregated/slurm/benchmark/gen_worker_config.py

@@ -48,6 +48,11 @@ def gen_config_file(work_dir: str,
         server_port: Server port
     """
     ctx_config = {
+        'build_config': {
+            'max_batch_size': ctx_batch_size,
+            'max_num_tokens': ctx_max_num_tokens,
+            'max_seq_len': ctx_max_seq_len,
+        },
         'max_batch_size': ctx_batch_size,
         'max_num_tokens': ctx_max_num_tokens,
         'max_seq_len': ctx_max_seq_len,
@@ -79,6 +84,11 @@ def gen_config_file(work_dir: str,
         gen_moe_backend = "TRTLLM"
 
     gen_config = {
+        'build_config': {
+            'max_batch_size': gen_batch_size,
+            'max_num_tokens': gen_max_num_tokens,
+            'max_seq_len': gen_max_seq_len,
+        },
         'tensor_parallel_size': gen_tp_size,
         'moe_expert_parallel_size': gen_tp_size,
         'enable_attention_dp': True if gen_enable_attention_dp else False,

tensorrt_llm/serve/scripts/benchmark_dataset.py

@@ -494,11 +494,14 @@ def sample(
 
             # Filter out sequences that are too long or too short
             requests = []
-            for prompt, initial_prompt_len, cached_token_ids in zip(
-                    dataset, prompt_lengths, prompt_token_ids):
-                i = len(requests)
-                if i == num_requests:
-                    break
+            dataset_len = len(dataset)
+
+            for i in range(num_requests):
+                # Use modulo to cycle through the dataset when num_requests > dataset_len
+                dataset_idx = i % dataset_len
+                prompt = dataset[dataset_idx]
+                initial_prompt_len = prompt_lengths[dataset_idx]
+                cached_token_ids = prompt_token_ids[dataset_idx]
 
                 # Skip empty prompt
                 if initial_prompt_len == 0:
@@ -534,9 +537,6 @@ def sample(
                         prompt_len=total_input_len,
                         expected_output_len=int(output_lens[i]),
                     ))
-            assert len(requests) == num_requests, (
-                f"Only {len(requests)} requests sampled from sharegpt dataset, {num_requests} requests are needed"
-            )
         else:
             for i in range(num_requests):
                 inner_seq = ((offsets[i] + i + np.arange(input_lens[i])) %
@@ -1131,6 +1131,7 @@ def sample(
         if parser_fn is None:
             raise ValueError(f"Unsupported dataset path: {self.dataset_path}")
 
+        sampled_requests = []
         for item in self.data:
             if len(prompts) >= num_requests:
                 break