Skip to content

Add autotune agent and tools#31

Merged
NinaCai merged 8 commits intomainfrom
nina-autotune
Apr 30, 2026
Merged

Add autotune agent and tools#31
NinaCai merged 8 commits intomainfrom
nina-autotune

Conversation

@NinaCai
Copy link
Copy Markdown
Collaborator

@NinaCai NinaCai commented Apr 22, 2026

Add autotune agent
Add autotune tool
Integrate with the system

@google-cla
Copy link
Copy Markdown

google-cla Bot commented Apr 22, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@NinaCai NinaCai changed the title [DO NOT MERGE]add autotune agent and tools Add autotune agent and tools Apr 27, 2026
@shangkunwang01 shangkunwang01 self-requested a review April 28, 2026 19:50
Comment thread MaxKernel/hitl_agent/subagents/autotuning/prompts/autotune_prompt.py Outdated
Comment thread MaxKernel/hitl_agent/subagents/autotuning/prompts/autotune_prompt.py Outdated
Comment thread MaxKernel/hitl_agent/subagents/autotuning/prompts/autotune_prompt.py Outdated
Comment thread MaxKernel/hitl_agent/subagents/autotuning/agent.py Outdated
Comment thread MaxKernel/hitl_agent/subagents/autotuning/autotune_tool.py
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The autotune will not be run on CPU. I would suggest removing this end point in cpu_server.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd say let's keep cpu_server in case tpu server is not available for customers or for infra_verification if they need.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The current implementation will not fallback to cpu backend (the default backend in autotune_kernel is tpu and when it is used in the AutotuneRunner, the backend is not specified. As a result it will still only use tpu). A good guidance for this kind of situation would be go/tott/737.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, I added cpu fallback logic so when there is no tpu resources, users are still able to use cpu for infra verification.

Comment thread MaxKernel/hitl_agent/server_utils/tpu_server.py Outdated
@NinaCai NinaCai requested a review from shangkunwang01 April 29, 2026 17:58
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would suggest letting the eval_server to handle the backend selection logic if we want to use both. See examples like in accelerator-agents/MaxKernel/hitl_agent/subagents/profiling/kernel_profile.py and accelerator-agents/MaxKernel/hitl_agent/subagents/kernel_writing/kernel_compilation.py.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

@NinaCai NinaCai merged commit db2a5b6 into main Apr 30, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants