Add disagg bs==1 test

Signed-off-by: Yifei Zhang <[email protected]>

Author: yifeizhang-c

tests/integration/defs/disaggregated/test_configs/disagg_config_genbs1.yaml

@@ -0,0 +1,45 @@
+model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+hostname: localhost
+port: 8000
+backend: pytorch
+cuda_graph_config: null
+context_servers:
+  num_instances: 1
+  max_batch_size: 1
+  max_num_tokens: 3000
+  max_seq_len: 4096
+  free_gpu_memory_fraction: 0.85
+  tensor_parallel_size: 2
+  moe_expert_parallel_size: 2
+  enable_attention_dp: true
+  pipeline_parallel_size: 1
+  disable_overlap_scheduler: true
+  kv_cache_config:
+    enable_block_reuse: false
+    free_gpu_memory_fraction: 0.85
+    dtype: fp8
+  cache_transceiver_config:
+    backend: default
+  urls:
+    - "localhost:8001"
+generation_servers:
+  num_instances: 1
+  tensor_parallel_size: 2
+  moe_expert_parallel_size: 2
+  enable_attention_dp: true
+  pipeline_parallel_size: 1
+  max_batch_size: 1
+  max_num_tokens: 4
+  max_seq_len: 2251
+  free_gpu_memory_fraction: 0.7
+  kv_cache_config:
+    enable_block_reuse: false
+    free_gpu_memory_fraction: 0.7
+    dtype: fp8
+  moe_config:
+    backend: CUTLASS
+  cache_transceiver_config:
+    backend: default
+  stream_interval: 20
+  urls:
+    - "localhost:8002"

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -44,6 +44,7 @@ def get_test_config(test_desc, example_dir, test_root):
         "gen_only": (2, f"{test_configs_root}/disagg_config_gen_only.yaml"),
         "gen_only_trt_backend":
         (2, f"{test_configs_root}/disagg_config_gen_only_trt_backend.yaml"),
+        "genbs1": (4, f"{test_configs_root}/disagg_config_genbs1.yaml"),
         "4_ranks": (4, f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml"),
         "4_ranks_trt_backend":
         (4,
@@ -384,6 +385,28 @@ def test_disaggregated_benchmark_gen_only_trt_backend(
                            cwd=llm_venv.get_working_directory())
 
 
+@pytest.mark.skip_less_device(4)
+@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
+                         indirect=True)
+def test_disaggregated_benchmark_genbs1(disaggregated_test_root,
+                                        disaggregated_example_root, llm_venv,
+                                        llama_model_root):
+    src_dst_dict = {
+        llama_model_root:
+        f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    }
+    for src, dst in src_dst_dict.items():
+        if not os.path.islink(dst):
+            os.makedirs(os.path.dirname(dst), exist_ok=True)
+            os.symlink(src, dst, target_is_directory=True)
+
+    run_disaggregated_test(disaggregated_example_root,
+                           "genbs1",
+                           env=llm_venv._new_env,
+                           cwd=llm_venv.get_working_directory(),
+                           prompt_file="long_prompts.json")
+
+
 @pytest.mark.skip_less_device(2)
 @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
                          indirect=True)