finish

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

Author: Ubospica

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(

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
@@ -175,15 +175,15 @@ def dispatch(
         if sol_name:
             sol = self._trace_set.get_solution(sol_name)
             if sol:
-                runnable = get_builder_registry().build(defn, sol)
+                runnable = BuilderRegistry.get_instance().build(defn, sol)
 
         # 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)
+                    runnable = BuilderRegistry.get_instance().build(defn, sol)
                 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,7 +115,7 @@ 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)
@@ -169,7 +169,7 @@ 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] = {}
@@ -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:

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.

flashinfer_bench/bench/evaluators/default.py

@@ -16,8 +16,7 @@
     make_eval,
     normalize_outputs,
 )
-from flashinfer_bench.compile.registry import get_builder_registry
-from flashinfer_bench.compile.runnable import Runnable
+from flashinfer_bench.compile import BuilderRegistry, Runnable
 from flashinfer_bench.data.definition import Definition
 from flashinfer_bench.data.trace import (
     Correctness,
@@ -46,7 +45,7 @@ def build_baseline(
         traceset_root: Optional[Path] = None,
     ) -> DeviceBaseline:
         output_dtypes = {k: dtype_str_to_torch_dtype(v.dtype) for k, v in defn.outputs.items()}
-        ref_runnable = get_builder_registry().build_reference(defn)
+        ref_runnable = BuilderRegistry.get_instance().build_reference(defn)
         loaded_stensors = (
             load_safetensors(defn, workload, traceset_root)
             if any(d.type == "safetensors" for d in workload.inputs.values())

flashinfer_bench/bench/evaluators/sampling.py

@@ -18,8 +18,7 @@
     make_eval,
     normalize_outputs,
 )
-from flashinfer_bench.compile.registry import get_builder_registry
-from flashinfer_bench.compile.runnable import Runnable
+from flashinfer_bench.compile import BuilderRegistry, Runnable
 from flashinfer_bench.data.definition import Definition
 from flashinfer_bench.data.trace import Correctness, Evaluation, EvaluationStatus, Workload
 from flashinfer_bench.utils import dtype_str_to_torch_dtype
@@ -42,7 +41,7 @@ def build_baseline(
         device: str,
         traceset_root: Optional[Path] = None,
     ) -> DeviceBaseline:
-        ref_runnable = get_builder_registry().build_reference(defn)
+        ref_runnable = BuilderRegistry.get_instance().build_reference(defn)
         loaded_stensors = (
             load_safetensors(defn, workload, traceset_root)
             if any(d.type == "safetensors" for d in workload.inputs.values())

flashinfer_bench/bench/runner/isolated_runner.py

@@ -14,7 +14,7 @@
 from flashinfer_bench.bench.config import BenchmarkConfig
 from flashinfer_bench.bench.evaluators import resolve_evaluator
 from flashinfer_bench.bench.utils import make_eval
-from flashinfer_bench.compile import Runnable, get_builder_registry
+from flashinfer_bench.compile import BuilderRegistry, Runnable
 from flashinfer_bench.data import Definition, Evaluation, EvaluationStatus, Solution, Workload
 from flashinfer_bench.logging import get_logger
 from flashinfer_bench.utils import redirect_stdio_to_file
@@ -230,7 +230,7 @@ def _solution_worker_main(
     redirect_stdio_to_file(log_path)
     try:
         torch.cuda.set_device(int(device.split(":")[1]))
-        registry = get_builder_registry()
+        registry = BuilderRegistry.get_instance()
 
         # Handshake
         conn.send({"cmd": "READY"})

flashinfer_bench/bench/runner/persistent_runner.py

@@ -16,7 +16,7 @@
 from flashinfer_bench.bench.config import BenchmarkConfig
 from flashinfer_bench.bench.evaluators import resolve_evaluator
 from flashinfer_bench.bench.utils import make_eval
-from flashinfer_bench.compile import BuildError, get_builder_registry
+from flashinfer_bench.compile import BuilderRegistry, BuildError
 from flashinfer_bench.data import Definition, Evaluation, EvaluationStatus, Solution, Workload
 from flashinfer_bench.logging import get_logger
 from flashinfer_bench.utils import redirect_stdio_to_file
@@ -65,7 +65,7 @@ def __init__(self, device: str, log_dir: str = "/tmp/flashinfer_bench") -> None:
         self._device = device
         self._log_dir = log_dir
         self._baselines: Dict[BaselineHandle, DeviceBaseline] = {}
-        self._registry = get_builder_registry()
+        self._registry = BuilderRegistry.get_instance()
 
         # Solution failure tracking
         self._failure_records: Dict[str, SolutionFailureRecord] = {}
@@ -648,7 +648,7 @@ def _persistent_worker_main(conn: mp.connection.Connection, device: str, log_dir
     """
     try:
         torch.cuda.set_device(int(device.split(":")[1]))
-        registry = get_builder_registry()
+        registry = BuilderRegistry.get_instance()
 
         conn.send({"cmd": WorkerResponse.READY.value})
 

flashinfer_bench/compile/__init__.py

@@ -1,10 +1,20 @@
 """Compiler subsystem package.
 
-Exports common builder types for convenience.
+This package provides the infrastructure for building solutions into executable runnables.
+It includes:
+- Builder: Abstract base class for different language/build system implementations
+- BuilderRegistry: Central registry for managing and dispatching builders
+- Runnable: Executable wrapper around compiled solutions
+- RunnableMetadata: Metadata about build process and source
+
+The typical workflow is:
+1. Get the singleton registry: registry = BuilderRegistry.get_instance()
+2. Build a solution: runnable = registry.build(definition, solution)
+3. Execute: result = runnable(**inputs)
 """
 
 from .builder import Builder, BuildError
-from .registry import BuilderRegistry, get_builder_registry
-from .runnable import Runnable
+from .registry import BuilderRegistry
+from .runnable import Runnable, RunnableMetadata
 
-__all__ = ["Builder", "BuildError", "BuilderRegistry", "Runnable", "get_builder_registry"]
+__all__ = ["Builder", "BuildError", "BuilderRegistry", "Runnable", "RunnableMetadata"]