From d783032f8dd6f9e6e7942af2003cbaa4d5018af5 Mon Sep 17 00:00:00 2001 From: Darkroom4364 Date: Sat, 23 May 2026 10:46:26 +0200 Subject: [PATCH] feat: generate kernel-aware NAS candidates --- docs/research/kernel-aware-nas.md | 12 ++-- scripts/nas_experiment.py | 61 +++++++++++++++++- src/research_engine/arch_cost_model.py | 89 ++++++++++++++++++++++++++ tests/test_arch_cost_model.py | 60 +++++++++++++++++ 4 files changed, 216 insertions(+), 6 deletions(-) diff --git a/docs/research/kernel-aware-nas.md b/docs/research/kernel-aware-nas.md index 6734817..5accf9b 100644 --- a/docs/research/kernel-aware-nas.md +++ b/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. diff --git a/scripts/nas_experiment.py b/scripts/nas_experiment.py index 7501c42..173ebfd 100644 --- a/scripts/nas_experiment.py +++ b/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,7 +216,11 @@ 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"], }, } @@ -191,7 +228,10 @@ def build_report(model: ArchitectureCostModel) -> dict: 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}") diff --git a/src/research_engine/arch_cost_model.py b/src/research_engine/arch_cost_model.py index dee95c4..80bd558 100644 --- a/src/research_engine/arch_cost_model.py +++ b/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.""" diff --git a/tests/test_arch_cost_model.py b/tests/test_arch_cost_model.py index 348fbe4..b8b544b 100644 --- a/tests/test_arch_cost_model.py +++ b/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()