Skip to content

Commit 0ea0592

Browse files
authored
[rocm7.1_internal_testing] Add triton_kernels wheel generation (#2566)
Fixes #ISSUE_NUMBER
1 parent 93fa949 commit 0ea0592

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

.github/scripts/build_triton_wheel.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,13 @@ def build_triton(
162162
cwd=triton_basedir,
163163
)
164164

165+
# For gpt-oss models, triton requires this extra triton_kernels wheel
166+
# triton_kernels came after pytorch release/2.8
167+
triton_kernels_dir = Path(f"{triton_basedir}/python/triton_kernels")
168+
check_call([sys.executable, "-m", "build", "--wheel"], cwd=triton_kernels_dir, env=env)
169+
kernels_whl_path = next(iter((triton_kernels_dir / "dist").glob("*.whl")))
170+
shutil.copy(kernels_whl_path, Path.cwd())
171+
165172
return Path.cwd() / whl_path.name
166173

167174

0 commit comments

Comments
 (0)