diff --git a/.ci/test_infra/oss_ci_benchmark_v3.py b/.ci/test_infra/oss_ci_benchmark_v3.py index 12e34e176..27121b1fc 100644 --- a/.ci/test_infra/oss_ci_benchmark_v3.py +++ b/.ci/test_infra/oss_ci_benchmark_v3.py @@ -24,6 +24,11 @@ "gpu_count": 1, "avail_gpu_mem_in_gb": 288, }, + "linux.dgx.b200": { + "name": "linux.dgx.b200", + "gpu_count": 1, + "avail_gpu_mem_in_gb": 192, + } } @@ -52,7 +57,7 @@ def setup_tritonbench_cwd(): def parse_runners( runner_name: str, runner_type: str, envs: Dict[str, str] ) -> List[Dict[str, Any]]: - runner_mapping = RUNNER_TYPE_MAPPING[runner_type].copy() + runner_mapping = RUNNER_TYPE_MAPPING.get(runner_type, {}).copy() runner_mapping["name"] = runner_name runner_mapping["gpu_info"] = envs["device"] runner_mapping["extra_info"] = {}