Repository navigation
esmfold2-trimul: support torch.compile, upload to biohub - #1214
github-actions[bot] merged 2 commits into
Conversation
transformers is turning Hub kernels on by default and only maps kernels that work under torch.compile (huggingface/transformers#47773). Make the TriMul kernel compile-safe: - Drop @torch._dynamo.disable from the entry point. It forced a graph break at every TriMul call, and an error under fullgraph=True. - Declare can_torch_compile on the layer so it can be registered for Mode.TORCH_COMPILE. - Pass STRIDE_B/N_COL as runtime args instead of tl.constexpr. Both scale with sequence length, so they recompiled the Triton kernel, and the compiled graph, for every new length. - Use tl.cast instead of M.to in the LayerNorm kernels. M.to fails when M is a plain int (Triton's M == 1 specialization, and torch.compile's kernel mutation analysis). This also lifts the B == L == 1 restriction. Also point the Hub repo-id at biohub/esmfold2-trimul, the repo transformers loads this kernel from, so release builds are uploaded there as a Hub PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
This PR changes source code for the following kernel(s): If you have the kernel-bot build permissions, then trigger |
|
/kernel-bot build esmfold2-trimul |
|
Build request processed. Command: Dispatched (1):
|
There was a problem hiding this comment.
looks good I think we just need to use the kernels library in the tests too
**prob something like the following (may need minor adjustments)
import kernels
esmfold2_trimul = kernels.get_kernel("kernels-community/esmfold2-trimul", version=1)
|
@drbh one question, actually - I think I got a bit confused here, I'm not sure what the procedure is when a kernel is under an external trusted org (in this case, |
|
@Rocketknight1 good question, updating the **ps the kernel bot also supports which will build and push to https://huggingface.co/kernels-staging/kernels and can be used to test the final build of a kernel which could be helpful for testing in the future. also let us know if we should remove the https://huggingface.co/kernels/kernels-community/esmfold2-trimul repo in the case the kernel is always intended to be published under |
|
/kernel-bot build-and-stage esmfold2-trimul |
|
Build request processed. Command: Dispatched (1):
Hub uploads: |
|
@drbh done! And yes, this will live in |
|
/kernel-bot merge-and-upload esmfold2-trimul |
|
Build request failed. Command: Failure: Failed to merge PR before build/upload. Check mergeability and required checks. |
|
/kernel-bot merge-and-upload esmfold2-trimul |
|
Build request processed. Command: Merge result: Pull Request successfully merged Dispatched (1):
|
|
/kernel-bot merge-and-upload esmfold2-trimul |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
Hub uploads:
|
|
@drbh why did it emit 6 update PRs to the same kernel instead of a single one? |
2 participants
torch.compile() fixes for the ESMFold2 kernels!