Add dsr1 and gpt-oss test cases

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

Author: chenfeiz0326

tests/integration/defs/perf/open_search_db_utils.py

@@ -40,11 +40,13 @@
     "l_ep",
     "l_pp",
     "l_max_num_tokens",
+    "l_cuda_graph_max_batch_size",
     "b_enable_chunked_prefill",
     "b_disable_overlap_scheduler",
     "s_attention_backend",
     "s_moe_backend",
     "l_moe_max_num_tokens",
+    "l_num_postprocess_workers",
     "l_stream_interval",
     "b_enable_attention_dp",
     "b_attention_dp_balance",
@@ -55,6 +57,11 @@
     "d_free_gpu_memory_fraction",
     "l_max_batch_size",
     "b_enable_padding",
+    "s_spec_decoding_type",
+    "l_num_nextn_predict_layers",
+    "l_eagle3_layers_to_capture",
+    "l_max_draft_len",
+    "s_speculative_model_dir",
 ]
 
 # Client config fields to compare
@@ -64,6 +71,8 @@
     "l_isl",
     "l_osl",
     "d_random_range_ratio",
+    "s_backend",
+    "b_use_chat_template",
 ]
 
 # Metrics where larger is better

tests/integration/defs/perf/test_perf.py

@@ -103,6 +103,7 @@
     "deepseek_r1_nvfp4": "DeepSeek-R1/DeepSeek-R1-FP4",
     "deepseek_r1_0528_fp8": "DeepSeek-R1/DeepSeek-R1-0528/",
     "deepseek_r1_0528_fp4": "DeepSeek-R1/DeepSeek-R1-0528-FP4/",
+    "deepseek_r1_0528_fp4_v2": "DeepSeek-R1/DeepSeek-R1-0528-FP4-v2/",
     "deepseek_v3_lite_fp8": "DeepSeek-V3-Lite/fp8",
     "deepseek_v3_lite_nvfp4": "DeepSeek-V3-Lite/nvfp4_moe_only",
     "qwen2_7b_instruct": "Qwen2-7B-Instruct",
@@ -513,11 +514,13 @@ def __init__(
         max_num_tokens: int,
         attention_backend: str,
         max_batch_size: int,
+        cuda_graph_max_batch_size: int = 0,
         pp: int = 1,
         enable_chunked_prefill: bool = False,
         disable_overlap_scheduler: bool = False,
         moe_backend: str = "",
         moe_max_num_tokens: int = 0,
+        num_postprocess_workers: int = 0,
         stream_interval: int = 10,
         enable_attention_dp: bool = False,
         attention_dp_balance: bool = False,
@@ -527,6 +530,11 @@ def __init__(
         enable_block_reuse: bool = False,
         free_gpu_memory_fraction: float = 0.8,
         enable_padding: bool = True,
+        spec_decoding_type: str = "",
+        num_nextn_predict_layers: int = 0,
+        eagle3_layers_to_capture: int = 0,
+        max_draft_len: int = 0,
+        speculative_model_dir: str = "",
     ):
         self.name = name
         self.model_name = model_name
@@ -540,6 +548,7 @@ def __init__(
         self.attention_backend = attention_backend
         self.moe_backend = moe_backend
         self.moe_max_num_tokens = moe_max_num_tokens
+        self.num_postprocess_workers = num_postprocess_workers
         self.stream_interval = stream_interval
         self.enable_attention_dp = enable_attention_dp
         self.attention_dp_balance = attention_dp_balance
@@ -549,7 +558,13 @@ def __init__(
         self.enable_block_reuse = enable_block_reuse
         self.free_gpu_memory_fraction = free_gpu_memory_fraction
         self.max_batch_size = max_batch_size
+        self.cuda_graph_max_batch_size = max_batch_size if cuda_graph_max_batch_size == 0 else cuda_graph_max_batch_size
         self.enable_padding = enable_padding
+        self.spec_decoding_type = spec_decoding_type
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.eagle3_layers_to_capture = eagle3_layers_to_capture
+        self.max_draft_len = max_draft_len
+        self.speculative_model_dir = speculative_model_dir
 
         self.model_path = ""
 
@@ -567,7 +582,7 @@ def to_cmd(self, working_dir: str) -> List[str]:
 
     def to_db_data(self) -> dict:
         """Convert ServerConfig to Database data"""
-        return {
+        db_data = {
             "s_model_name": self.model_name.lower(),
             "l_gpus": self.gpus,
             "l_tp": self.tp,
@@ -588,9 +603,30 @@ def to_db_data(self) -> dict:
             "b_enable_block_reuse": self.enable_block_reuse,
             "d_free_gpu_memory_fraction": self.free_gpu_memory_fraction,
             "l_max_batch_size": self.max_batch_size,
+            "l_cuda_graph_max_batch_size": self.cuda_graph_max_batch_size,
             "b_enable_padding": self.enable_padding,
+            "s_spec_decoding_type": self.spec_decoding_type,
+            "l_num_nextn_predict_layers": self.num_nextn_predict_layers,
+            "l_eagle3_layers_to_capture": self.eagle3_layers_to_capture,
+            "l_max_draft_len": self.max_draft_len,
+            "s_speculative_model_dir": self.speculative_model_dir,
             "s_server_log_link": "",
         }
+        if self.num_postprocess_workers > 0:
+            db_data["l_num_postprocess_workers"] = self.num_postprocess_workers
+        if self.spec_decoding_type:
+            db_data["s_spec_decoding_type"] = self.spec_decoding_type
+            if self.num_nextn_predict_layers > 0:
+                db_data[
+                    "l_num_nextn_predict_layers"] = self.num_nextn_predict_layers
+            if self.eagle3_layers_to_capture > 0:
+                db_data[
+                    "l_eagle3_layers_to_capture"] = self.eagle3_layers_to_capture
+            if self.max_draft_len > 0:
+                db_data["l_max_draft_len"] = self.max_draft_len
+            if self.speculative_model_dir:
+                db_data["s_speculative_model_dir"] = self.speculative_model_dir
+        return db_data
 
     def generate_extra_llm_api_config(self) -> str:
         """Generate extra-llm-api-config.yml content"""
@@ -599,21 +635,28 @@ def generate_extra_llm_api_config(self) -> str:
             f"moe_expert_parallel_size: {self.ep}",
             f"pipeline_parallel_size: {self.pp}",
             f"max_num_tokens: {self.max_num_tokens}",
+            f"max_batch_size: {self.max_batch_size}",
             f"enable_attention_dp: {str(self.enable_attention_dp).lower()}",
             f"disable_overlap_scheduler: {str(self.disable_overlap_scheduler).lower()}",
-            f"stream_interval: {self.stream_interval}",
             f"attn_backend: {self.attention_backend}",
             f"enable_chunked_prefill: {str(self.enable_chunked_prefill).lower()}",
             "cuda_graph_config:",
             f"  enable_padding: {str(self.enable_padding).lower()}",
-            f"  max_batch_size: {self.max_batch_size}",
+            f"  max_batch_size: {self.cuda_graph_max_batch_size}",
             "kv_cache_config:",
             f"  dtype: {self.kv_cache_dtype}",
             f"  free_gpu_memory_fraction: {self.free_gpu_memory_fraction}",
             f"  enable_block_reuse: {str(self.enable_block_reuse).lower()}",
             "print_iter_log: false",
         ]
 
+        if self.stream_interval > 0:
+            config_lines.append(f"stream_interval: {self.stream_interval}")
+
+        if self.num_postprocess_workers > 0:
+            config_lines.append(
+                f"num_postprocess_workers: {self.num_postprocess_workers}")
+
         # Add moe_config if moe_backend is specified
         if self.moe_backend:
             config_lines.append("moe_config:")
@@ -629,6 +672,23 @@ def generate_extra_llm_api_config(self) -> str:
                 f"  batching_wait_iters: {self.batching_wait_iters}")
             config_lines.append(f"  timeout_iters: {self.timeout_iters}")
 
+        if self.spec_decoding_type:
+            config_lines.append("speculative_config:")
+            config_lines.append(f"  decoding_type: {self.spec_decoding_type}")
+            if self.num_nextn_predict_layers > 0:
+                config_lines.append(
+                    f"  num_nextn_predict_layers: {self.num_nextn_predict_layers}"
+                )
+            if self.eagle3_layers_to_capture > 0:
+                config_lines.append(
+                    f"  eagle3_layers_to_capture: {self.eagle3_layers_to_capture}"
+                )
+            if self.max_draft_len > 0:
+                config_lines.append(f"  max_draft_len: {self.max_draft_len}")
+            if self.speculative_model_dir:
+                config_lines.append(
+                    f"  speculative_model_dir: {speculative_model_dir}")
+
         return "\n".join(config_lines)
 
 
@@ -644,22 +704,26 @@ def __init__(self,
                  iterations: int,
                  isl: int,
                  osl: int,
-                 random_range_ratio: float = 0.0):
+                 random_range_ratio: float = 0.0,
+                 backend: str = "",
+                 use_chat_template: bool = False):
         self.name = name
         self.model_name = model_name
         self.concurrency = concurrency
         self.iterations = iterations
         self.isl = isl
         self.osl = osl
         self.random_range_ratio = random_range_ratio
+        self.backend = backend
+        self.use_chat_template = use_chat_template
 
         self.model_path = ""
 
     def to_cmd(self, working_dir: str) -> List[str]:
         model_dir = get_model_dir(self.model_name)
         self.model_path = model_dir if os.path.exists(
             model_dir) else self.model_name
-        return [
+        benchmark_cmd = [
             "python", "-m", "tensorrt_llm.serve.scripts.benchmark_serving",
             "--model", self.model_path, "--dataset-name", "random",
             "--random-ids", "--num-prompts",
@@ -670,17 +734,30 @@ def to_cmd(self, working_dir: str) -> List[str]:
             "--percentile-metrics", "ttft,tpot,itl,e2el", "--max-concurrency",
             str(self.concurrency)
         ]
+        if self.backend:
+            benchmark_cmd.append("--backend")
+            benchmark_cmd.append(self.backend)
+        if self.use_chat_template:
+            benchmark_cmd.append("--use-chat-template")
+        return benchmark_cmd
 
     def to_db_data(self) -> dict:
         """Convert ClientConfig to Database data"""
-        return {
+        db_data = {
             "l_concurrency": self.concurrency,
             "l_iterations": self.iterations,
             "l_isl": self.isl,
             "l_osl": self.osl,
             "d_random_range_ratio": self.random_range_ratio,
+            "s_backend": self.backend,
+            "b_use_chat_template": self.use_chat_template,
             "s_client_log_link": "",
         }
+        if self.backend:
+            db_data["s_backend"] = self.backend
+        if self.use_chat_template:
+            db_data["b_use_chat_template"] = self.use_chat_template
+        return db_data
 
 
 def parse_select_pattern(select_pattern: str):