Mirroristas/whisper
2
0
Fork 0
mirror of https://github.com/openai/whisper.git synced 2025-11-24 06:26:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactored changes to update triton kernel only once

Browse Source
This commit is contained in:
ExtReMLapin 2025-05-07 10:08:12 +02:00 committed by GitHub
parent 6335afea82
commit 285c21b32e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
whisper/triton_ops.py
View File

@ -60,7 +60,6 @@ def median_kernel(filter_width: int):
tl.store(y_ptr + offsets, MIDDLE_ROW_HERE, mask=mask) # noqa: F821 tl.store(y_ptr + offsets, MIDDLE_ROW_HERE, mask=mask) # noqa: F821
kernel = triton.JITFunction(kernel.fn) kernel = triton.JITFunction(kernel.fn)
triton_3_kernel_update = hasattr(kernel, "_unsafe_update_src")
new_kernel = kernel.src.replace( new_kernel = kernel.src.replace(
" LOAD_ALL_ROWS_HERE", " LOAD_ALL_ROWS_HERE",
@ -71,12 +70,8 @@ def median_kernel(filter_width: int):
] ]
), ),
) )
if triton_3_kernel_update is True:
kernel._unsafe_update_src(new_kernel)
else:
kernel.src = new_kernel
new_kernel = kernel.src.replace( new_kernel = new_kernel.replace(
" BUBBLESORT_HERE", " BUBBLESORT_HERE",
"\n\n".join( "\n\n".join(
[ [
@ -97,14 +92,10 @@ def median_kernel(filter_width: int):
] ]
), ),
) )
if triton_3_kernel_update is True:
kernel._unsafe_update_src(new_kernel)
else:
kernel.src = new_kernel
new_kernel = kernel.src.replace("MIDDLE_ROW_HERE", f"row{filter_width // 2}") new_kernel = new_kernel.replace("MIDDLE_ROW_HERE", f"row{filter_width // 2}")
if triton_3_kernel_update is True: if hasattr(kernel, "_unsafe_update_src") is True:
kernel._unsafe_update_src(new_kernel) kernel._unsafe_update_src(new_kernel)
kernel.hash = None kernel.hash = None
else: else: