refactor: Builder System (#120)

This PR aims to simplify the logic of all builders and provide a unified
abstraction. It also provides a simpler abstraction for Runnable. It
also:

1. Renames CUDABuilder to TorchBuilder and adds CPP support to it
2. Enhances the CI unit test to reduce execution time
3. Removes get_builder_registry and provides class-level singleton
BuilderRegistry.get_instance() instead
4. Provides a hash function for Solution.
5. Provides utilities for tests

Signed-off-by: Ubospica <[email protected]>

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added TorchBuilder for compiling CUDA solutions via PyTorch
extensions.
* Introduced Solution hashing for deterministic caching and entry-point
extraction utilities.

* **Bug Fixes**
* Improved runtime rank validation and error handling with better error
messages.
  * Removed legacy CUDA builder; replaced with TorchBuilder.

* **Improvements**
  * Simplified builder registry access pattern.
  * Enhanced metadata tracking and cleanup workflows.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Ubospica <[email protected]>

Author: Ubospica

examples/fi_gqa_e2e_example.py

@@ -1,3 +1,6 @@
+import flashinfer
+import torch
+
 from flashinfer_bench.apply import ApplyConfig, enable_apply
 from flashinfer_bench.data import (
     BuildSpec,
@@ -16,11 +19,11 @@
     Workload,
 )
 
-ts = TraceSet.from_path("../flashinfer-trace")  # path to your flashinfer-trace
+trace_set = TraceSet.from_path("../flashinfer-trace")  # path to your flashinfer-trace
 
 # Add a pseudo solution with an unrealistically large speedup so it gets selected
 def_name = "gqa_paged_prefill_causal_h32_kv8_d128_ps1"
-if def_name in ts.definitions:
+if def_name in trace_set.definitions:
     pseudo_code = (
         "import torch\n\n"
         "def run(q, k_cache, v_cache, qo_indptr, kv_indptr, kv_indices, sm_scale):\n"
@@ -36,7 +39,7 @@
         definition=def_name,
         author="pseudo",
         spec=BuildSpec(
-            language=SupportedLanguages.PYTHON, target_hardware=["gpu"], entry_point="main.py::run"
+            language=SupportedLanguages.PYTHON, target_hardware=["cuda"], entry_point="main.py::run"
         ),
         sources=[SourceFile(path="main.py", content=pseudo_code)],
         description="Pseudo solution that prints greeting and returns zero tensors.",
@@ -77,17 +80,13 @@
         ),
     )
 
-    ts.solutions.setdefault(def_name, []).append(pseudo_sol)
-    ts.traces.setdefault(def_name, []).append(pseudo_trace)
+    trace_set.solutions.setdefault(def_name, []).append(pseudo_sol)
+    trace_set.traces.setdefault(def_name, []).append(pseudo_trace)
 
 # Enable apply against the in-memory augmented trace set
-enable_apply(ts, ApplyConfig(on_miss_policy="use_def_best"))
-
-import flashinfer
+enable_apply(trace_set, ApplyConfig(on_miss_policy="use_def_best"))
 
 # FlashInfer official example
-import torch
-
 num_layers = 32
 num_qo_heads = 32
 num_kv_heads = 8

examples/win_at_p.py

@@ -26,7 +26,6 @@
 Notes:
 - If multiple runs exist for the same author within a group, we take the MIN latency for that author in that group.
 - By default, the baseline author ('flashinfer') is EXCLUDED from output curves; use --include-baseline to include it.
-
 """
 
 import argparse

flashinfer_bench/apply/key.py

@@ -12,7 +12,7 @@
 class ApplyKey:
     axes: Tuple[Tuple[str, int], ...] = field(default_factory=tuple)
     # Features extracted from input tensors
-    feats: Tuple[Tuple[str, Union[int, Union[float, bool]]], ...] = field(default_factory=tuple)
+    feats: Tuple[Tuple[str, Union[int, float, bool]], ...] = field(default_factory=tuple)
 
     def encode(self) -> str:
         return json.dumps(
@@ -39,10 +39,10 @@ def __eq__(self, other: Any) -> bool:
 
 
 class ApplyKeyBuilder(ABC):
-    def __init__(self, defn: Definition) -> None:
-        self.defn = defn
+    def __init__(self, definition: Definition) -> None:
+        self.definition = definition
         # axis -> (input, dim_idx)
-        self._axis_proj: Dict[str, Tuple[str, int]] = self._collect_var_axis_projections(defn)
+        self._axis_proj: Dict[str, Tuple[str, int]] = self._collect_var_axis_projections(definition)
 
     @abstractmethod
     def build_from_runtime(self, runtime_kwargs: Dict[str, Any]) -> ApplyKey:
@@ -59,14 +59,14 @@ def features(self, runtime_kwargs: Dict[str, Any]) -> Tuple[Tuple[str, Any], ...
         """Lightweight feature extraction"""
         ...
 
-    def _collect_var_axis_projections(self, defn: Definition) -> Dict[str, Tuple[str, int]]:
+    def _collect_var_axis_projections(self, definition: Definition) -> Dict[str, Tuple[str, int]]:
         """
         Iterate over the shape of inputs, find the first occurrence of each var axis:
           axis_name -> (input_name, dim_idx)
         """
         proj: Dict[str, Tuple[str, int]] = {}
-        axis_defs = defn.axes
-        inputs = defn.inputs
+        axis_defs = definition.axes
+        inputs = definition.inputs
 
         for inp_name, spec in inputs.items():
             shape = spec.shape
@@ -82,15 +82,15 @@ def _collect_var_axis_projections(self, defn: Definition) -> Dict[str, Tuple[str
         var_axes = [k for k, v in axis_defs.items() if isinstance(v, AxisVar)]
         missing = [ax for ax in var_axes if ax not in proj]
         if missing:
-            raise ValueError(f"Cannot locate var axes {missing} from inputs of '{defn.name}'")
+            raise ValueError(f"Cannot locate var axes {missing} from inputs of '{definition.name}'")
         return proj
 
     def _materialize_axes(self, runtime_kwargs: Dict[str, Any]) -> Dict[str, int]:
         axes: Dict[str, int] = {}
         for axis, (inp, dim_idx) in self._axis_proj.items():
             if inp not in runtime_kwargs:
                 raise KeyError(
-                    f"Missing runtime input '{inp}' for axis '{axis}' in '{self.defn.name}'"
+                    f"Missing runtime input '{inp}' for axis '{axis}' in '{self.definition.name}'"
                 )
             val = runtime_kwargs[inp]
             shape = val.shape
@@ -144,9 +144,9 @@ def for_type(cls, type_name: str) -> Type[ApplyKeyBuilder]:
         return cls._REGISTRY.get(type_name, AxesOnlyKeyBuilder)
 
     @classmethod
-    def specialize(cls, defn: Definition) -> ApplyKeyBuilder:
-        builder_cls = cls.for_type(defn.op_type)
-        return builder_cls(defn)
+    def specialize(cls, definition: Definition) -> ApplyKeyBuilder:
+        builder_cls = cls.for_type(definition.op_type)
+        return builder_cls(definition)
 
 
 ApplyKeyFactory.register("gemm", GEMMKeyBuilder)

flashinfer_bench/apply/runtime.py

@@ -10,7 +10,7 @@
 
 from typing import Any, Callable, Dict, Mapping, Optional
 
-from flashinfer_bench.compile import get_builder_registry
+from flashinfer_bench.compile import BuilderRegistry
 from flashinfer_bench.data import TraceSet
 from flashinfer_bench.env import get_fib_dataset_path, get_fib_enable_apply
 from flashinfer_bench.logging import get_logger
@@ -156,34 +156,34 @@ def dispatch(
                 pass
 
         # Then try to run apply logic
-        defn = self._trace_set.definitions.get(def_name)
-        if defn is None:
+        definition = self._trace_set.definitions.get(def_name)
+        if definition is None:
             if fallback is None:
                 raise RuntimeError(f"Definition '{def_name}' not found and no fallback provided")
             return fallback(**runtime_kwargs)
 
         # Build key
-        builder = self._key_builders.get(defn.name)
+        builder = self._key_builders.get(definition.name)
         if builder is None:
-            builder = ApplyKeyFactory.specialize(defn)
-            self._key_builders[defn.name] = builder
+            builder = ApplyKeyFactory.specialize(definition)
+            self._key_builders[definition.name] = builder
         key = builder.build_from_runtime(runtime_kwargs)
 
         # Lookup solution
         sol_name = self._table.match_solution(def_name, key)
         runnable = None
         if sol_name:
-            sol = self._trace_set.get_solution(sol_name)
-            if sol:
-                runnable = get_builder_registry().build(defn, sol)
+            solution = self._trace_set.get_solution(sol_name)
+            if solution:
+                runnable = BuilderRegistry.get_instance().build(definition, solution)
 
         # Miss policy
         if runnable is None:
             if self._apply_config.on_miss_policy == "use_def_best":
                 best_sol_name = self._table.def_best.get(def_name)
-                sol = self._trace_set.get_solution(best_sol_name)
-                if defn and sol:
-                    runnable = get_builder_registry().build(defn, sol)
+                solution = self._trace_set.get_solution(best_sol_name)
+                if definition and solution:
+                    runnable = BuilderRegistry.get_instance().build(definition, solution)
                 if runnable is not None:
                     return runnable(**runtime_kwargs)
             if fallback is None:

flashinfer_bench/apply/table.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
-from flashinfer_bench.compile import Runnable, get_builder_registry
+from flashinfer_bench.compile import BuilderRegistry, Runnable
 from flashinfer_bench.data import Trace, TraceSet
 from flashinfer_bench.env import get_fib_cache_path
 
@@ -115,13 +115,13 @@ def load_or_build(cls, trace_set: TraceSet, apply_config: ApplyConfig) -> "Apply
                 index[def_name] = bucket
 
             def_best: Dict[str, str] = {}
-            reg = get_builder_registry()
+            reg = BuilderRegistry.get_instance()
 
             for def_name, sol_name in raw["def_best"].items():
-                defn = trace_set.definitions.get(def_name)
-                sol = trace_set.get_solution(sol_name)
-                if defn and sol:
-                    reg.build(defn, sol)
+                definition = trace_set.definitions.get(def_name)
+                solution = trace_set.get_solution(sol_name)
+                if definition and solution:
+                    reg.build(definition, solution)
                     def_best[def_name] = sol_name
 
             table = cls(digest=digest, index=index, def_best=def_best)
@@ -169,12 +169,12 @@ def _build(cls, trace_set: TraceSet, apply_config: ApplyConfig) -> "ApplyTable":
             The newly built apply table.
         """
         digest = cls._digest(trace_set, apply_config)
-        reg = get_builder_registry()
+        reg = BuilderRegistry.get_instance()
 
         index: Dict[str, Dict[ApplyKey, str]] = {}
         def_best: Dict[str, Runnable] = {}
 
-        for def_name, defn in trace_set.definitions.items():
+        for def_name, definition in trace_set.definitions.items():
             per_key, ranked = cls._sweep_def(
                 trace_set, def_name, apply_config.max_atol, apply_config.max_rtol
             )
@@ -189,11 +189,11 @@ def _build(cls, trace_set: TraceSet, apply_config: ApplyConfig) -> "ApplyTable":
             # Build def_best
             if ranked:
                 best_sol_name = ranked[0][0]
-                sol = trace_set.get_solution(best_sol_name)
-                if sol:
+                solution = trace_set.get_solution(best_sol_name)
+                if solution:
                     if apply_config.on_miss_policy == "use_def_best":
                         # Only AOT if on_miss_policy is use_def_best
-                        reg.build(defn, sol)
+                        reg.build(definition, solution)
                     def_best[def_name] = best_sol_name
 
         return cls(digest=digest, index=index, def_best=def_best)
@@ -267,7 +267,7 @@ def _prewarm_aot(cls, trace_set: TraceSet, config: ApplyConfig, table: "ApplyTab
         """
         if not (config.aot_ratio and config.aot_ratio > 0.0):
             return
-        reg = get_builder_registry()
+        reg = BuilderRegistry.get_instance()
 
         for def_name, bucket in table.index.items():
             if not bucket:
@@ -277,20 +277,20 @@ def _prewarm_aot(cls, trace_set: TraceSet, config: ApplyConfig, table: "ApplyTab
             ranked = sorted(win_counts.items(), key=lambda kv: kv[1], reverse=True)
             cutoff = max(1, int(len(ranked) * config.aot_ratio))
 
-            defn = trace_set.definitions.get(def_name)
-            if not defn:
+            definition = trace_set.definitions.get(def_name)
+            if not definition:
                 continue
             for sol_name, _ in ranked[:cutoff]:
-                sol = trace_set.get_solution(sol_name)
-                if sol:
-                    reg.build(defn, sol)
+                solution = trace_set.get_solution(sol_name)
+                if solution:
+                    reg.build(definition, solution)
 
         if config.on_miss_policy == "use_def_best":
             for def_name, sol_name in table.def_best.items():
-                defn = trace_set.definitions.get(def_name)
-                sol = trace_set.get_solution(sol_name)
-                if defn and sol:
-                    reg.build(defn, sol)
+                definition = trace_set.definitions.get(def_name)
+                solution = trace_set.get_solution(sol_name)
+                if definition and solution:
+                    reg.build(definition, solution)
 
     @classmethod
     def _digest(cls, trace_set: TraceSet, config: ApplyConfig) -> str:
@@ -313,23 +313,23 @@ def _digest(cls, trace_set: TraceSet, config: ApplyConfig) -> str:
             SHA256 hash digest as a hexadecimal string.
         """
         d = trace_set.to_dict()
-        for defn in d["definitions"].values():
+        for definition in d["definitions"].values():
             for drop in ("description", "tags", "reference", "constraints"):
-                defn.pop(drop, None)
+                definition.pop(drop, None)
         for sol_list in d["solutions"].values():
-            for sol in sol_list:
-                spec = sol.get("spec", {}) or {}
+            for solution in sol_list:
+                spec = solution.get("spec", {}) or {}
                 deps = spec.get("dependencies") or []
                 spec["dependencies"] = sorted(deps)
                 new_sources = []
-                for sf in sol.get("sources") or []:
+                for sf in solution.get("sources") or []:
                     new_sources.append(
                         {
                             "path": sf["path"],
                             "sha1": hashlib.sha1(sf["content"].encode("utf-8")).hexdigest(),
                         }
                     )
-                sol["sources"] = new_sources
+                solution["sources"] = new_sources
         kept_traces: List[Dict[str, Any]] = []
         for traces in d["traces"].values():
             for trace in traces:

flashinfer_bench/bench/benchmark.py

@@ -3,7 +3,7 @@
 from collections import defaultdict
 from typing import List
 
-from flashinfer_bench.compile import get_builder_registry
+from flashinfer_bench.compile import BuilderRegistry
 from flashinfer_bench.data import EvaluationStatus, Trace, TraceSet
 from flashinfer_bench.logging import get_logger
 
@@ -46,7 +46,7 @@ def __init__(self, trace_set: TraceSet, config: BenchmarkConfig = None) -> None:
             self._runner = PersistentRunner(logger, self._config.log_dir)
 
         # Setup registry
-        self._registry = get_builder_registry()
+        self._registry = BuilderRegistry.get_instance()
 
     def get_trace_set(self) -> TraceSet:
         """Get the TraceSet associated with this benchmark.
@@ -79,8 +79,8 @@ def run_all(self, dump_traces: bool = True, resume: bool = False) -> TraceSet:
         definitions_to_run = self._trace_set.definitions.items()
         if self._config.definitions is not None:
             definitions_to_run = [
-                (name, defn)
-                for name, defn in definitions_to_run
+                (name, definition)
+                for name, definition in definitions_to_run
                 if name in self._config.definitions
             ]
             provided_defs = set(self._config.definitions)
@@ -89,7 +89,7 @@ def run_all(self, dump_traces: bool = True, resume: bool = False) -> TraceSet:
             if missing_defs:
                 logger.warning(f"Definitions not found in trace set: {sorted(missing_defs)}")
 
-        for def_name, defn in definitions_to_run:
+        for def_name, definition in definitions_to_run:
             sols = self._trace_set.solutions.get(def_name, [])
             if not sols:
                 logger.warning(f"No solutions found for def={def_name}, skipping definition")
@@ -128,7 +128,7 @@ def run_all(self, dump_traces: bool = True, resume: bool = False) -> TraceSet:
 
                 try:
                     results = self._runner.run_workload(
-                        defn, wl, sols_to_run, self._config, self._trace_set.root
+                        definition, wl, sols_to_run, self._config, self._trace_set.root
                     )
                 except RuntimeError as e:
                     logger.error(f"Failed to run workload {wl.uuid}: {e}")