From c8cb505ecd198c63ea44a98d23d53498f8eb1933 Mon Sep 17 00:00:00 2001
From: edwingao28 <105094497+edwingao28@users.noreply.github.com>
Date: Fri, 16 Oct 2026 23:57:26 +0000
Subject: [PATCH] feat(matrix): filter test-config rows by TP

Add --tp to test-config so a run selects single-node rows by TP from the master config instead of a duplicated selector file. Absent TP values and multi-node rows fail.
---
 inferencex-e2e/infx/matrix/generate.py        | 22 ++++++++++
 .../matrix/test_generate_sweep_configs.py     | 40 +++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/inferencex-e2e/infx/matrix/generate.py b/inferencex-e2e/infx/matrix/generate.py
index 15816a534e..c828976cbb 100644
--- a/inferencex-e2e/infx/matrix/generate.py
+++ b/inferencex-e2e/infx/matrix/generate.py
@@ -1441,6 +1441,15 @@ def filter_exp_names(entries: list[dict], exp_names: list[str]) -> list[dict]:
     return [entry for entry in entries if entry.get(Fields.EXP_NAME.value) in requested]
 
 
+def filter_tp(entries: list[dict], tps: list[int]) -> list[dict]:
+    if any(Fields.TP.value not in entry for entry in entries):
+        raise ValueError("--tp filters single-node rows only")
+    missing = sorted(set(tps) - {entry[Fields.TP.value] for entry in entries})
+    if missing:
+        raise ValueError("TP value(s) not found: " + ", ".join(map(str, missing)))
+    return [entry for entry in entries if entry[Fields.TP.value] in tps]
+
+
 def apply_node_type_defaults(args: argparse.Namespace) -> argparse.Namespace:
     if (
         hasattr(args, "single_node")
@@ -1666,6 +1675,13 @@ def main() -> list[dict]:
         required=False,
         help="Only include these concurrency values. Values must exist in the config conc-range/list.",
     )
+    test_config_keys_parser.add_argument(
+        "--tp",
+        nargs="+",
+        type=int,
+        required=False,
+        help="Only include single-node rows with these TP sizes.",
+    )
     test_config_keys_parser.add_argument(
         "--exp-names",
         nargs="+",
@@ -1711,6 +1727,12 @@ def main() -> list[dict]:
     else:
         parser.error(f"Unknown command: {args.command}")
 
+    if args.command == "test-config" and args.tp:
+        try:
+            matrix_values = filter_tp(matrix_values, args.tp)
+        except ValueError as error:
+            parser.error(str(error))
+
     if args.command == "test-config" and args.exp_names:
         try:
             matrix_values = filter_exp_names(matrix_values, args.exp_names)
diff --git a/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py b/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py
index 698d3940ff..6c0fa059ac 100644
--- a/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py
+++ b/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py
@@ -17,6 +17,7 @@
     apply_node_type_defaults,
     expand_config_keys,
     filter_exp_names,
+    filter_tp,
     generate_full_sweep,
     generate_test_config_sweep,
     mark_all_eval_entries,
@@ -3222,6 +3223,45 @@ def test_rejects_missing_ambiguous_or_duplicate_names(
 
 
 
+class TestFilterTp:
+    @pytest.mark.parametrize(
+        ("entries", "message"),
+        (
+            ([{"tp": 8}], "not found: 4"),
+            ([{"prefill": {"tp": 4}, "decode": {"tp": 4}}], "single-node rows only"),
+        ),
+    )
+    def test_rejects_absent_tp_or_multinode(self, entries, message):
+        with pytest.raises(ValueError, match=message):
+            filter_tp(entries, [4])
+
+    def test_test_config_keeps_only_requested_tp(
+        self, monkeypatch, sample_single_node_config, sample_runner_config
+    ):
+        config = copy.deepcopy(sample_single_node_config)
+        sequence = config["dsr1-fp8-mi300x-sglang"]["scenarios"]["fixed-seq-len"][1]
+        sequence["search-space"] = [
+            {"tp": 8, "conc-start": 1, "conc-end": 4},
+            {"tp": 4, "conc-start": 1, "conc-end": 64},
+        ]
+        monkeypatch.setattr(generate_sweep_configs, "load_config_files", lambda _: config)
+        monkeypatch.setattr(
+            generate_sweep_configs, "load_runner_file", lambda _: sample_runner_config
+        )
+        monkeypatch.setattr(sys, "argv", [
+            "generate_sweep_configs.py", "test-config",
+            "--config-files", "dummy.yaml",
+            "--config-keys", "dsr1-fp8-mi300x-sglang",
+            "--seq-lens", "8k1k", "--tp", "4", "--conc", "1", "2", "4", "--no-evals",
+        ])
+
+        result = generate_sweep_configs.main()
+
+        assert sorted((entry["tp"], entry["conc"]) for entry in result) == [
+            (4, 1), (4, 2), (4, 4),
+        ]
+
+
 class TestExpandConfigKeys:
 
     AVAILABLE = [