Add support for more batch sizes in cpu_graph_runner (#13881)

2026-03-19 09:50:56 -07:00
parent 4c52b7fcc6
commit 274581fb77
5 changed files with 135 additions and 25 deletions
@@ -31,9 +31,11 @@ class TestCPUGraph(CustomTestCase):
            "0.05",
            "--enable-torch-compile",
            "--torch-compile-max-bs",
-            "1",
+            "2",
+            "--cuda-graph-bs",
+            "2",
        ],
-        min_throughput=10,
+        min_throughput=7,
    )
    def test_latency_torch_compile_cpu(self):
        return DEFAULT_MLA_MODEL_NAME_FOR_TEST
@@ -58,8 +60,8 @@ class TestCPUGraph(CustomTestCase):
                "--trust-remote-code",
                "--disable-overlap-schedule",
                "--enable-torch-compile",
-                "--torch-compile-max-bs",
-                "1",
+                "--cuda-graph-bs",
+                "2",
                "--tp",
                f"{n_numa_node}",
            ],