Generate kernel-aware NAS candidates

docs/research/kernel-aware-nas.md

@@ -15,8 +15,12 @@ to choose model dimensions that the kernels execute efficiently.
   total layer latency or latency normalized by `hidden_dim * ffn_dim`.
 - `ArchitectureCostModel.sweep_dimension()` exposes tile-efficiency cliffs for
   dimensions that do not land on the 128-wide matmul tile boundary.
+- `generate_nas_candidates()` expands a base config across hidden width, head
+  dimension, GQA ratio, FFN ratio, sliding window, and QK-norm placement while
+  skipping invalid head/KV combinations.
 - `scripts/nas_experiment.py` compares known model shapes against novel
-  candidates and prints a fastest-first NAS ranking.
+  candidates, prints a fastest-first NAS ranking, and ranks the generated
+  candidate space.
 
 Run:
 
@@ -34,6 +38,7 @@ The experiment can answer hardware-facing architecture questions such as:
 
 - which candidate has the lowest predicted layer latency on a target GPU;
 - which kernel dominates the layer latency for that candidate;
+- which generated head-dim / GQA / FFN-ratio / QK-norm candidates rank fastest;
 - whether a hidden size, head size, or FFN size falls off the 128-wide tile
   boundary; and
 - whether the same ranking holds across A100, T4, and H100 profiles.
@@ -52,8 +57,7 @@ constraint before selecting an architecture.
 
 1. Calibrate the hardcoded effective throughput profiles from the persisted
    `.noeris` kernel performance database instead of manual constants.
-2. Add candidate generation over head dimension, GQA ratio, FFN ratio, sliding
-   window size, and QK-norm placement.
-3. Add a multi-hardware comparison command that writes A100, T4, and H100
+2. Add a multi-hardware comparison command that writes A100, T4, and H100
    reports in one invocation.
+3. Add candidate-quality constraints before selecting architectures.
 4. Validate the top candidates with real Triton layer benchmarks.

scripts/nas_experiment.py

@@ -28,6 +28,7 @@
 _mod = importlib.util.module_from_spec(_spec)
 _spec.loader.exec_module(_mod)
 ArchitectureCostModel = _mod.ArchitectureCostModel
+generate_nas_candidates = _mod.generate_nas_candidates
 _tile_efficiency = _mod._tile_efficiency
 
 
@@ -113,6 +114,30 @@ def _ranking_rows(ranked: list[dict]) -> list[dict]:
     ]
 
 
+def generated_candidate_configs() -> list[dict]:
+    """Generate a compact 2B-class candidate space for kernel-aware NAS."""
+    base = {
+        "hidden_dim": 2048,
+        "num_heads": 16,
+        "num_kv_heads": 2,
+        "head_dim": 128,
+        "ffn_dim": 8192,
+        "seq_len": 2048,
+        "batch_size": 1,
+        "use_qk_norm": True,
+        "window_size": 1024,
+    }
+    return generate_nas_candidates(
+        base,
+        hidden_dims=[1536, 2048, 2560],
+        head_dims=[64, 128, 256],
+        ffn_ratios=[3.0, 4.0, 5.333],
+        kv_head_counts=[1, 2, 4, 8],
+        window_sizes=[512, 1024, None],
+        qk_norm_options=[True, False],
+    )
+
+
 def build_report(model: ArchitectureCostModel) -> dict:
     """Build a machine-readable NAS report for one hardware profile."""
     all_configs = KNOWN_CONFIGS + NOVEL_CONFIGS
@@ -158,17 +183,25 @@ def build_report(model: ArchitectureCostModel) -> dict:
 
     fastest = model.rank_configs(all_configs, metric="total_ms")
     most_efficient = model.rank_configs(all_configs, metric="ms_per_mparam_proxy")
+    generated = generated_candidate_configs()
+    generated_fastest = model.rank_configs(generated, metric="total_ms")
+    generated_efficient = model.rank_configs(generated, metric="ms_per_mparam_proxy")
 
     return {
         "schema_version": 1,
         "experiment": "kernel_aware_nas",
         "hardware": model.hardware,
         "candidate_count": len(all_configs),
+        "generated_candidate_count": len(generated),
         "comparison": comparisons,
         "rankings": {
             "total_ms": _ranking_rows(fastest),
             "ms_per_mparam_proxy": _ranking_rows(most_efficient),
         },
+        "generated_search": {
+            "total_ms_top": _ranking_rows(generated_fastest[:25]),
+            "ms_per_mparam_proxy_top": _ranking_rows(generated_efficient[:25]),
+        },
         "kernel_cliffs": {
             "hidden_dim": {
                 "base_config": hidden_base,
@@ -183,15 +216,22 @@ def build_report(model: ArchitectureCostModel) -> dict:
             "fastest_config": fastest[0]["name"],
             "fastest_total_ms": fastest[0]["total_ms"],
             "most_efficient_config": most_efficient[0]["name"],
-            "most_efficient_ms_per_mparam_proxy": most_efficient[0]["ms_per_mparam_proxy"],
+            "most_efficient_ms_per_mparam_proxy": (
+                most_efficient[0]["ms_per_mparam_proxy"]
+            ),
+            "fastest_generated_config": generated_fastest[0]["name"],
+            "fastest_generated_total_ms": generated_fastest[0]["total_ms"],
         },
     }
 
 
 def write_report(report: dict, output_path: Path) -> None:
     """Write a deterministic JSON artifact."""
     output_path.parent.mkdir(parents=True, exist_ok=True)
-    output_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    output_path.write_text(
+        json.dumps(report, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
 
 
 def run_comparison(model: ArchitectureCostModel) -> None:
@@ -240,6 +280,22 @@ def run_comparison(model: ArchitectureCostModel) -> None:
               f"bottleneck={row['bottleneck']}")
 
 
+def run_generated_search(model: ArchitectureCostModel, top_n: int = 10) -> None:
+    """Generate and rank a broader kernel-aware architecture candidate set."""
+    generated = generated_candidate_configs()
+    ranked = model.rank_configs(generated)
+
+    print(f"\n{'='*90}")
+    print(f"  Generated NAS candidates: top {top_n} of {len(generated)}")
+    print(f"{'='*90}")
+    print(f"  {'Rank':>4s} {'Config':38s} {'Total ms':>9s} {'Eff':>8s} {'Bottleneck':>14s}")
+    print("  " + "-" * 82)
+    for row in ranked[:top_n]:
+        print(f"  #{row['rank']:02d}  {row['name'][:38]:38s} "
+              f"{row['total_ms']:9.3f} {row['ms_per_mparam_proxy']:8.4f} "
+              f"{row['bottleneck']:>14s}")
+
+
 def run_kernel_cliff_test(model: ArchitectureCostModel) -> None:
     """Test whether tile-unaligned dimensions cause performance cliffs."""
     print(f"\n{'='*90}")
@@ -298,6 +354,7 @@ def main() -> None:
     report = build_report(model)
 
     run_comparison(model)
+    run_generated_search(model)
     run_kernel_cliff_test(model)
 
     print(f"\n{'='*90}")

src/research_engine/arch_cost_model.py

@@ -86,6 +86,14 @@
 MATMUL_TILE_SIZE = 128  # typical Triton BLOCK_M/N
 
 
+def _round_up_to_multiple(value: float, multiple: int) -> int:
+    return int(math.ceil(value / multiple) * multiple)
+
+
+def _format_ratio(value: float) -> str:
+    return f"{value:g}".replace(".", "p")
+
+
 def _is_tile_aligned(dim: int, tile: int = MATMUL_TILE_SIZE) -> bool:
     return dim % tile == 0
 
@@ -111,6 +119,87 @@ def _tile_penalty(dim: int, tile: int = MATMUL_TILE_SIZE) -> dict[str, Any]:
     }
 
 
+def generate_nas_candidates(
+    base_config: dict[str, Any],
+    *,
+    hidden_dims: list[int] | None = None,
+    head_dims: list[int] | None = None,
+    ffn_ratios: list[float] | None = None,
+    kv_head_counts: list[int] | None = None,
+    window_sizes: list[int | None] | None = None,
+    qk_norm_options: list[bool] | None = None,
+    tile_multiple: int = MATMUL_TILE_SIZE,
+    max_candidates: int | None = None,
+) -> list[dict[str, Any]]:
+    """Generate tile-aligned transformer architecture candidates.
+
+    The search varies dimensions that are natural handles for kernel-aware NAS:
+    hidden width, attention head width, GQA ratio, FFN expansion, sliding-window
+    attention, and QK-norm placement. Invalid combinations are skipped.
+    """
+    hidden_dims = hidden_dims or [base_config["hidden_dim"]]
+    head_dims = head_dims or [64, 128, 256]
+    ffn_ratios = ffn_ratios or [3.0, 4.0, 5.333]
+    kv_head_counts = kv_head_counts or [1, 2, 4, 8]
+    window_sizes = window_sizes or [base_config.get("window_size"), None]
+    qk_norm_options = qk_norm_options or [base_config.get("use_qk_norm", True), False]
+
+    candidates: list[dict[str, Any]] = []
+    seen_names: set[str] = set()
+    seen_shapes: set[tuple[Any, ...]] = set()
+
+    for hidden_dim in hidden_dims:
+        for head_dim in head_dims:
+            if hidden_dim % head_dim != 0:
+                continue
+            num_heads = hidden_dim // head_dim
+            for num_kv_heads in kv_head_counts:
+                if num_kv_heads > num_heads or num_heads % num_kv_heads != 0:
+                    continue
+                for ffn_ratio in ffn_ratios:
+                    ffn_dim = _round_up_to_multiple(hidden_dim * ffn_ratio, tile_multiple)
+                    for window_size in window_sizes:
+                        for use_qk_norm in qk_norm_options:
+                            name = (
+                                f"gen_h{hidden_dim}_hd{head_dim}_kv{num_kv_heads}_"
+                                f"ffn{_format_ratio(ffn_ratio)}_"
+                                f"win{window_size or 'full'}_"
+                                f"{'qknorm' if use_qk_norm else 'rope'}"
+                            )
+                            shape_key = (
+                                hidden_dim,
+                                num_heads,
+                                num_kv_heads,
+                                head_dim,
+                                ffn_dim,
+                                window_size,
+                                use_qk_norm,
+                            )
+                            if name in seen_names or shape_key in seen_shapes:
+                                continue
+                            cfg = {
+                                **base_config,
+                                "name": name,
+                                "hidden_dim": hidden_dim,
+                                "num_heads": num_heads,
+                                "num_kv_heads": num_kv_heads,
+                                "head_dim": head_dim,
+                                "ffn_dim": ffn_dim,
+                                "window_size": window_size,
+                                "use_qk_norm": use_qk_norm,
+                            }
+                            candidates.append(cfg)
+                            seen_names.add(name)
+                            seen_shapes.add(shape_key)
+                            if (
+                                max_candidates is not None
+                                and len(candidates) >= max_candidates
+                            ):
+                                return candidates
+
+    return candidates
+
+
 class ArchitectureCostModel:
     """Predicts layer latency from architecture config using real kernel data."""
 

tests/test_arch_cost_model.py

@@ -15,6 +15,7 @@
 from research_engine.arch_cost_model import (
     HARDWARE_PROFILES,
     ArchitectureCostModel,
+    generate_nas_candidates,
     _is_tile_aligned,
     _tile_efficiency,
 )
@@ -156,6 +157,47 @@ def test_sweep_dimension_includes_tile_efficiency(self) -> None:
         self.assertFalse(results[1]["aligned_128"])
         self.assertLess(results[1]["tile_efficiency"], 1.0)
 
+    def test_generate_nas_candidates_varies_architecture_knobs(self) -> None:
+        candidates = generate_nas_candidates(
+            BASE_CONFIG,
+            hidden_dims=[2048],
+            head_dims=[64, 128],
+            ffn_ratios=[3.0, 4.0],
+            kv_head_counts=[1, 2, 64],
+            window_sizes=[512, None],
+            qk_norm_options=[True, False],
+        )
+
+        names = [cfg["name"] for cfg in candidates]
+        self.assertEqual(len(names), len(set(names)))
+        self.assertGreater(len(candidates), 1)
+        self.assertTrue(all(cfg["ffn_dim"] % 128 == 0 for cfg in candidates))
+        self.assertTrue(
+            all(cfg["num_heads"] % cfg["num_kv_heads"] == 0 for cfg in candidates)
+        )
+        self.assertTrue(
+            all(cfg["num_kv_heads"] <= cfg["num_heads"] for cfg in candidates)
+        )
+        self.assertEqual({cfg["head_dim"] for cfg in candidates}, {64, 128})
+        self.assertEqual({cfg["use_qk_norm"] for cfg in candidates}, {True, False})
+        self.assertNotIn(64, {cfg["num_kv_heads"] for cfg in candidates})
+
+    def test_generated_candidates_are_rankable(self) -> None:
+        candidates = generate_nas_candidates(
+            BASE_CONFIG,
+            hidden_dims=[1536, 2048],
+            head_dims=[128, 256],
+            ffn_ratios=[3.0],
+            kv_head_counts=[1, 2],
+            window_sizes=[1024],
+            qk_norm_options=[True],
+        )
+        ranked = ArchitectureCostModel("a100").rank_configs(candidates)
+
+        self.assertEqual(len(ranked), len(candidates))
+        self.assertEqual(ranked[0]["rank"], 1)
+        self.assertLessEqual(ranked[0]["total_ms"], ranked[-1]["total_ms"])
+
 
 class NasExperimentTests(unittest.TestCase):
     def test_build_report_has_expected_schema_for_each_hardware(self) -> None:
@@ -182,6 +224,12 @@ def test_build_report_has_expected_schema_for_each_hardware(self) -> None:
                 report["summary"]["fastest_config"],
                 report["rankings"]["total_ms"][0]["name"],
             )
+            self.assertGreater(report["generated_candidate_count"], 0)
+            self.assertEqual(len(report["generated_search"]["total_ms_top"]), 25)
+            self.assertEqual(
+                report["summary"]["fastest_generated_config"],
+                report["generated_search"]["total_ms_top"][0]["name"],
+            )
             self.assertIn("hidden_dim", report["kernel_cliffs"])
             self.assertIn("ffn_dim", report["kernel_cliffs"])
 
@@ -246,6 +294,18 @@ def test_run_comparison_prints_nas_ranking(self) -> None:
 
         self.assertIn("Fastest-first NAS ranking", stdout.getvalue())
 
+    def test_generated_search_prints_top_candidates(self) -> None:
+        module = _load_nas_experiment_module()
+        stdout = io.StringIO()
+
+        with redirect_stdout(stdout):
+            module.run_generated_search(ArchitectureCostModel("a100"), top_n=3)
+
+        output = stdout.getvalue()
+        self.assertIn("Generated NAS candidates", output)
+        self.assertIn("top 3", output)
+        self.assertIn("#01", output)
+
 
 if __name__ == "__main__":
     unittest.main()