[None][test] Fix mpi-type issue and add wideep acc test to dev's l0 local flow

jenkins/scripts/perf/local/submit.py

@@ -1,7 +1,10 @@
 #!/usr/bin/env python3
 import argparse
+import copy
+import json
 import os
 import re
+import shutil
 from datetime import datetime
 
 import yaml
@@ -312,7 +315,9 @@ def generate_srun_args(args, runtime_mode, timestamp):
 
     lines.append("--container-env=NVIDIA_IMEX_CHANNELS")
 
-    if is_aggr:
+    if args.mpi_type:
+        lines.append(f"--mpi={args.mpi_type}")
+    elif is_aggr:
         lines.append("--mpi=pmi2")
 
     return lines
@@ -352,6 +357,26 @@ def generate_pytest_command(
     return pytest_command, test_list_content, test_list_path
 
 
+def replace_env_in_file(work_dir: str, file_path: str, env_vars: dict) -> str:
+    """Read a file, replace env var placeholders, write to work_dir/lm_eval_configs/.
+
+    Returns the lm_eval_configs directory path (for use as --include_path).
+    Also copies utils.py from the same directory if present (needed for GPQA task).
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        content = f.read()
+    for key, value in env_vars.items():
+        content = content.replace(key, value)
+    tmp_dir = os.path.join(work_dir, "lm_eval_configs")
+    os.makedirs(tmp_dir, exist_ok=True)
+    with open(os.path.join(tmp_dir, os.path.basename(file_path)), "w", encoding="utf-8") as f:
+        f.write(content)
+    utils_py = os.path.join(os.path.dirname(file_path), "utils.py")
+    if os.path.exists(utils_py):
+        shutil.copy(utils_py, tmp_dir)
+    return tmp_dir
+
+
 def remove_whitespace_lines(lines):
     """Remove empty lines and strip whitespace."""
     return [line for line in lines if line.strip()]
@@ -422,6 +447,12 @@ def main():
         help="Nsys start-stop range for generation workers in disaggregated mode (default: 1-100)",
     )
     parser.add_argument("--test-prefix", default="", help="Test prefix")
+    parser.add_argument(
+        "--mpi-type",
+        default="",
+        help="MPI type for srun (e.g. pmix, pmi2). If not set, aggregated runs default to"
+        " --mpi=pmi2; non-aggregated runs omit --mpi entirely.",
+    )
 
     args = parser.parse_args()
 
@@ -692,6 +723,22 @@ def main():
             ]
         )
 
+    # Export accuracy config to BENCHMARK node (disagg only)
+    if runtime_mode == "disaggregated":
+        acc_cfg = config.get("accuracy", {})
+        if acc_cfg.get("enable_accuracy_test"):
+            env_sub = {"LLM_MODELS_ROOT": args.llm_models_root}
+            processed = copy.deepcopy(acc_cfg)
+            for task_cfg in processed.get("tasks", {}).values():
+                extra = task_cfg.get("extra_kwargs", {})
+                if "custom_config" in extra:
+                    cfg_path = extra.pop("custom_config")
+                    if not os.path.isabs(cfg_path):
+                        cfg_path = os.path.join(llm_src, cfg_path)
+                    extra["include_path"] = replace_env_in_file(work_dir, cfg_path, env_sub)
+            script_prefix_lines.append(f"export ACCURACY_CONFIG_JSON='{json.dumps(processed)}'")
+            srun_args_lines.append("--container-env=ACCURACY_CONFIG_JSON")
+
     # Remove whitespace lines
     script_prefix_lines = remove_whitespace_lines(script_prefix_lines)
 

tests/integration/defs/perf/test_perf_sanity.py

@@ -611,6 +611,7 @@ class DisaggTestCmds(NamedTuple):
     num_gen_servers: int
     output_dir: str
     test_output_dir: str
+    model_name: str = ""
 
     def _generate_hostname_file(self, server_idx: int, port: int):
         """Create hostname file for coordination."""
@@ -825,6 +826,22 @@ def run_cmd(self, server_idx: int) -> List[str]:
                         benchmark_ctx.write(output)
                     outputs.append(output)
 
+                # Run accuracy tests after benchmark (if configured)
+                acc_cfg_json = os.environ.get("ACCURACY_CONFIG_JSON")
+                if acc_cfg_json:
+                    import json as _json
+
+                    acc_cfg = _json.loads(acc_cfg_json)
+                    if acc_cfg.get("enable_accuracy_test"):
+                        _run_accuracy_tests(
+                            acc_cfg,
+                            self.model_name,
+                            disagg_server_hostname,
+                            disagg_server_port,
+                            self.test_output_dir,
+                            server_idx,
+                        )
+
             finally:
                 with open(benchmark_status_file, "w") as status_file:
                     status_file.write("Done")
@@ -835,6 +852,60 @@ def get_cmd_str(self, server_idx: int) -> List[str]:
         return ["multi-node disaggregated server tests, please check config files"]
 
 
+def _run_accuracy_tests(
+    accuracy_cfg: dict,
+    model_name: str,
+    server_hostname: str,
+    server_port: int,
+    output_dir: str,
+    server_idx: int,
+) -> None:
+    """Run lm_eval against the running disagg server. Saves results only — no validation."""
+    endpoint_map = {
+        "local-completions": "v1/completions",
+        "local-chat-completions": "v1/chat/completions",
+    }
+    env_var = accuracy_cfg.get("env_var") or {}
+    model_path = get_model_dir(model_name)
+
+    for task_name, task_cfg in accuracy_cfg.get("tasks", {}).items():
+        model_type = task_cfg.get("model", "local-completions")
+        model_args_extra = task_cfg.get("model_args_extra", "")
+        extra_kwargs = task_cfg.get("extra_kwargs", {})
+        base_url = f"http://{server_hostname}:{server_port}/{endpoint_map.get(model_type, 'v1/completions')}"
+        model_args = f"model={model_path},base_url={base_url},{model_args_extra}"
+
+        acc_output_dir = os.path.join(output_dir, f"accuracy_eval_{task_name}.{server_idx}")
+        log_file = os.path.join(output_dir, f"accuracy_eval_{task_name}.{server_idx}.log")
+        os.makedirs(acc_output_dir, exist_ok=True)
+
+        cmd = [
+            "lm_eval",
+            "--model",
+            model_type,
+            "--tasks",
+            task_name,
+            "--model_args",
+            model_args,
+            "--log_samples",
+            "--output_path",
+            acc_output_dir,
+        ]
+        if "include_path" in extra_kwargs:
+            cmd += ["--include_path", extra_kwargs["include_path"]]
+        for k, v in extra_kwargs.items():
+            if k == "include_path":
+                continue
+            cmd += [f"--{k}"] if isinstance(v, bool) and v else [f"--{k}", str(v)]
+
+        run_env = copy.deepcopy(os.environ)
+        run_env.update({k: str(v) for k, v in env_var.items()})
+        print_info(f"[Accuracy] Running {task_name}, output: {log_file}")
+        with open(log_file, "w") as lf:
+            ret = subprocess.run(cmd, env=run_env, stdout=lf, stderr=subprocess.STDOUT)
+        print_info(f"[Accuracy] {task_name} done, exit_code={ret.returncode}")
+
+
 def parse_select_pattern(select_pattern: str) -> list:
     """Parse select pattern (server config names).
 
@@ -1288,6 +1359,7 @@ def _get_disagg_commands(self, output_dir: str, test_output_dir: str):
             num_gen_servers=disagg_config.num_gen_servers,
             output_dir=output_dir,
             test_output_dir=test_output_dir,
+            model_name=disagg_config.model_name,
         )
 
     def _check_benchmark_errors(self, output: str) -> None:

tests/integration/lm_eval_configs/gpqa_diamond_local.yaml

@@ -1,6 +1,10 @@
 # Modified from tensorrt_llm/evaluate/lm_eval_tasks/gpqa/cot_zeroshot_aa/gpqa_diamond_cot_zeroshot_aa.yaml
 task: gpqa_diamond_local
-dataset_path: HF_HOME/datasets/Idavidrein___gpqa
+dataset_path: csv
+dataset_name: null
+dataset_kwargs:
+  data_files:
+    train: LLM_MODELS_ROOT/datasets/gpqa/gpqa_diamond.csv
 tag: gpqa
 output_type: generate_until
 process_docs: !function utils.process_gpqa_docs

tests/integration/test_lists/qa/llm_perf_multinode.txt

@@ -158,5 +158,14 @@ perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_deepseek-v32-fp4_8k1k_
 perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_deepseek-v32-fp4_8k1k_ctx8_gen1_dep32_bs16_eplb288_mtp3_con512_ccb-NIXL]
 perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_kimi-k2-thinking-fp4_1k1k_ctx3_gen1_dep32_bs1024_eplb384_mtp0_con16384_ccb-NIXL]
 perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_kimi-k2-thinking-fp4_8k1k_ctx8_gen1_dep32_bs256_eplb416_mtp0_con8192_ccb-NIXL]
+
+# accuracy cases
+perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_accuracy-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL]
+perf/test_perf_sanity.py::test_e2e[disagg-gen_only-wideep_accuracy-deepseek-r1-fp4_gpqa_diamond_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL]
+perf/test_perf_sanity.py::test_e2e[disagg-e2e-wideep_accuracy-kimi-k2-thinking-fp4_1k1k_ctx3_gen1_dep32_bs1024_eplb384_mtp0_ccb-NIXL]
+
+# stress cases
+perf/test_perf_sanity.py::test_e2e[disagg-e2e-wideep_stress-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL]
+
 # GB200 supported cases
 # GB300 supported cases

tests/scripts/perf/disaggregated/wideep_accuracy-deepseek-r1-fp4_1k1k_ctx2_gen1_dep16_bs128_eplb288_mtp3_ccb-NIXL.yaml

@@ -0,0 +1,127 @@
+metadata:
+  model_name: deepseek_r1_0528_fp4_v2
+  precision: fp4
+  model_dir_name: DeepSeek-R1-0528-FP4-v2
+  supported_gpus:
+  - GB200
+  - GB300
+  script_file: disaggr_torch.slurm
+  benchmark_type: 1k1k
+  accuracy:
+    datasets:
+    - dataset_name: gsm8k_local
+      expected_value: 0.9454
+      threshold_type: hypothesis_test
+      filter_type: flexible-extract
+slurm:
+  script_file: disaggr_torch.slurm
+  partition: <partition>
+  account: <account>
+  job_time: 03:00:00
+  job_name: unified-benchmark
+  extra_args: "--gres=gpu:4"
+  numa_bind: true
+benchmark:
+  mode: gen_only
+  use_nv_sa_benchmark: false
+  multi_round: 1
+  benchmark_ratio: 0.8
+  streaming: true
+  concurrency_list: '2048'
+  input_length: 1024
+  output_length: 1024
+  dataset_file: datasets/perf-ci/deepseek_r1-1k1k-20480-ratio-1_for_serve.json
+hardware:
+  gpus_per_node: 4
+  num_ctx_servers: 2
+  num_gen_servers: 1
+environment:
+  container_mount: <container_mount>
+  container_image: <container_image>
+  model_path: <model_path>
+  trtllm_repo: ''
+  build_wheel: false
+  work_dir: <full_path_to_work_dir>
+  worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1
+    TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes
+  server_env_var: TRTLLM_SERVER_DISABLE_GC=1
+profiling:
+  nsys_on: false
+accuracy:
+  enable_accuracy_test: true
+  env_var:
+    HF_HOME: <hf_home_path>
+  tasks:
+    gsm8k_local:
+      model: "local-completions"
+      model_args_extra: "num_concurrent=512,max_retries=3,tokenized_requests=false,timeout=7200,max_gen_toks=16384"
+      extra_kwargs:
+        trust_remote_code: true
+        custom_config: tests/integration/lm_eval_configs/gsm8k_local.yaml
+worker_config:
+  gen:
+    enable_layerwise_nvtx_marker: true
+    tensor_parallel_size: 16
+    moe_expert_parallel_size: 16
+    enable_attention_dp: true
+    enable_lm_head_tp_in_adp: true
+    pipeline_parallel_size: 1
+    max_batch_size: 128
+    max_num_tokens: 512
+    max_seq_len: 2251
+    cuda_graph_config:
+      enable_padding: true
+      batch_sizes:
+      - 1
+      - 2
+      - 4
+      - 8
+      - 16
+      - 32
+      - 64
+      - 128
+      - 256
+      - 512
+      - 768
+      - 1024
+      - 2048
+    print_iter_log: true
+    kv_cache_config:
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.9
+      dtype: fp8
+    moe_config:
+      backend: WIDEEP
+      load_balancer:
+        num_slots: 288
+        layer_updates_per_iter: 1
+    cache_transceiver_config:
+      max_tokens_in_buffer: 4608
+      backend: NIXL
+    stream_interval: 20
+    num_postprocess_workers: 4
+    speculative_config:
+      decoding_type: MTP
+      num_nextn_predict_layers: 3
+  ctx:
+    enable_layerwise_nvtx_marker: true
+    max_batch_size: 4
+    max_num_tokens: 4608
+    max_seq_len: 2251
+    tensor_parallel_size: 4
+    moe_expert_parallel_size: 4
+    enable_attention_dp: true
+    pipeline_parallel_size: 1
+    print_iter_log: true
+    cuda_graph_config: null
+    disable_overlap_scheduler: true
+    kv_cache_config:
+      enable_block_reuse: false
+      free_gpu_memory_fraction: 0.85
+      dtype: fp8
+    cache_transceiver_config:
+      max_tokens_in_buffer: 4608
+      backend: NIXL
+    speculative_config:
+      decoding_type: MTP
+      num_nextn_predict_layers: 3