Skip to content

update kernel generator example to be more clear #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

update kernel generator example to be more clear #101

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
075ae7a
update kernel generator example to be more clear
zanderjiang Oct 24, 2025
e8d6eca
Merge branch 'main' of https://github.com/flashinfer-ai/flashinfer-be…
zanderjiang Oct 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 16 additions & 8 deletions examples/kernel_generator/kernel_generator_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,29 @@ def main():
"""
Generate optimized solutions for all definitions in the traceset.
"""
model_name = "gpt-5-2025-08-07" # Choose model here
language = "triton"
target_gpu = "B200"
# TODO: select model, language, target gpu, definition
model_name = "gpt-5-2025-08-07" # Choose author-model
language = "triton" # Target solution language
target_gpu = "B200" # Choose solution target GPU
definition = "" # Leave empty to generate solutions for all definitions
Comment on lines +24 to +26
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Oct 24, 2025

Choose a reason for hiding this comment

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

medium

For better readability and adherence to Python's PEP 8 style guide, inline comments should be separated by at least two spaces from the statement. I've also aligned the comments for better visual structure.1

Suggested change
language = "triton" # Target solution language
target_gpu = "B200" # Choose solution target GPU
definition = "" # Leave empty to generate solutions for all definitions
language = "triton" # Target solution language
target_gpu = "B200" # Choose solution target GPU
definition = "" # Leave empty to generate solutions for all definitions

Style Guide References

Footnotes

  1. PEP 8 E261 states that inline comments should be separated by at least two spaces from the statement. ↩


# TODO: adjust local path to traceset
traceset_path = "/home/akj2/flashinfer-trace"
traceset_path = "/path/to/flashinfer-trace"
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Oct 24, 2025

Choose a reason for hiding this comment

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

medium

While using a placeholder path is an improvement, the script doesn't handle the case where the user forgets to change it. If run with the default path, it will create an empty directory, find no definitions, and do nothing without a clear explanation. This can be confusing for new users. It would be more robust to check if any definitions were loaded from the provided path and exit with a helpful error message if not.

For example, you could add this check after loading the traceset:

if not traceset.definitions:
    print(f"Error: No definitions found in traceset at '{traceset_path}'.")
    print("Please ensure `traceset_path` points to a valid flashinfer-trace directory.")
    return


print(f"Loading TraceSet from: {traceset_path}")
traceset = TraceSet.from_path(traceset_path)

# all_definitions = list(traceset.definitions.keys())
# Filter for rmsnorm definitions only
all_definitions = [name for name in traceset.definitions.keys() if "rmsnorm" in name.lower()]
all_definitions = list(traceset.definitions.keys())

print(f"All definitions found: {len(all_definitions)}")
if definition:
if definition in all_definitions:
all_definitions = [definition]
print(f"Generating solution {definition}")
else:
print(f"Definition '{definition}' not found in traceset")
return

print(f"Found {len(all_definitions)} definitions to generate solutions")

api_key = os.getenv("LLM_API_KEY")
base_url = os.getenv("BASE_URL")
Expand Down
Loading