Skip to content

[autotuner] Introduce BenchmarkProvider abstraction for kernel benchmarking#1928

Open
hinriksnaer wants to merge 8 commits intopytorch:mainfrom
hinriksnaer:benchmark-provider
Open

[autotuner] Introduce BenchmarkProvider abstraction for kernel benchmarking#1928
hinriksnaer wants to merge 8 commits intopytorch:mainfrom
hinriksnaer:benchmark-provider

Conversation

@hinriksnaer
Copy link
Copy Markdown
Collaborator

sets up foundation for #1803

Motivation

Establishes a substitution point for alternative benchmarking strategies (e.g. cross-node precompilation, overlapped precompile+benchmark). Search algorithms call self.benchmark_provider.benchmark_function() instead of self.benchmark_function(). The provider owns the single-config benchmarking pipeline while _benchmark, benchmark_batch, and create_precompile_future remain on BaseSearch for follow-up migration. Once the changes are finalized, overall maintainability and effort required to introduce more fine grained stability measurements should be improved.

Summary

  • Introduce BenchmarkProvider abstract class and LocalBenchmarkProvider implementation
  • Move benchmark_function and its dependencies (baseline computation, accuracy validation, precompile management) from BaseSearch into the provider
  • BaseSearch accepts benchmark_provider_cls parameter to allow alternative benchmarking strategies
  • No behavioral changes, same code runs through the provider instead of directly on BaseSearch

Follow-up work

  • Migrate _benchmark and create_precompile_future to the provider
  • Migrate set_adaptive_compile_timeout as post_initial_benchmark
  • Decouple accuracy validation from the benchmarking provider

Discussion

This is a step towards decoupling benchmarking strategies from search algorithms. The goal is to enable alternative benchmarking strategies without them being tied to individual search algorithm implementations, and to improve maintainability by giving benchmarking a clear boundary.

This PR moves benchmark_function and its direct dependencies into the provider as the smallest self-contained unit that establishes the abstraction without doing a comprehensive rewrite. The remaining methods on BaseSearch are marked with TODOs for incremental migration.

Given the scope of these changes, any feedback is very much appreciated.

Replace the back-reference to the full BaseSearch object with a narrow
5-field PrecompileContext dataclass (settings, log, kernel, args, jobs).
This makes PrecompileFuture's dependencies explicit and testable in
isolation.

This change is part of a larger effort to introduce a BenchmarkProvider
abstraction that encapsulates the benchmarking pipeline. PrecompileFuture
currently holds a reference to BaseSearch to access 5 fields — when the
BenchmarkProvider owns precompilation, it needs to pass its own context
into PrecompileFuture without involving the search object. PrecompileContext
makes that possible.
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Apr 2, 2026
@hinriksnaer hinriksnaer marked this pull request as draft April 2, 2026 17:55
@hinriksnaer hinriksnaer marked this pull request as ready for review April 2, 2026 23:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant