🔗 Helpful Links

🧪 See artifacts and rendered test results athud.pytorch.org/pr/148890

• 📄 PreviewPython docs built from this PR
• 📄 PreviewC++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit thebot commands wiki or ouroffice hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit260d3eb with merge base1bf443e ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@jamesjwu has imported this pull request. If you are a Meta employee, you can view this diffon Phabricator.

@jamesjwu has imported this pull request. If you are a Meta employee, you can view this diffon Phabricator.

Can we add support for larger number of args (fall back to dynamic alloc), and the special stuff like hooks, shared memory, custom ops? I would really like to have only one code path since I think that will be easier to maintain. I suspect stuff like args>50 will have poor coverage in our test suite and is likey to get broken without us noticing.

I'm happy to implement these additions so that our coverage is as complete as possible, but I am very worried about trying to ship this without a config flag to turn it off initially or fall back on (especially internally, where rollouts can take weeks). So I think we should still be able to support the fallback to triton for now. I do agree directionally that we should try to get to the case where all kernels are statically launchable.

Re: falling back to dynamic alloc: what do you think about falling back to something likealloca, to just allocate it directly on the stack? I know it might not be super portable, but maybe it's possible to support.

I fixed a bug just now (when we have a global scratch, MAX_ARGS on C++ side actually needs to be MAX_ARGS + 1 lol), will kick off a new benchmark run.

I realize there's a big list of features that we technically need to support to be "feature parity" with triton's launcher capabilities, so I wrote down them all in this issue:#149440, and will fix them one by one:#149442.

The goal is that withconfig.static_cuda_launch=True, all triton kernels are statically launched (or else we hard error). This should decrease the complexity of maintaining the code base, because then either all of the CompileResults are triton launched, or all of them are statically launched, but there's no in between.

@jansel I do have a concern, though, with user defined triton kernels, because they can essentially containanything. So there's a bunch of features that PT2 generated kernels will never encounter (i.e.num_ctas>1,CUTensorMaps, etc), but would need to be implemented in the static launcher just to support user defined ones. We could fix the finite list of these today, but the bigger worry is that until we can fully upstream it to triton side, there's always a risk of triton supporting a new feature that StaticCudaLauncher can't support, and then having user defined triton kernels break. What do you think our options are here? Will we need to keep a fallback option around for those?

• Fixed a bunch of bugs and rebased onto main
• Ran a new benchmark:
  https://hud.pytorch.org/benchmark/huggingface/inductor_with_cudagraphs?dashboard=torchinductor&startTime=Wed%2C%2012%20Mar%202025%2015%3A18%3A30%20GMT&stopTime=Wed%2C%2019%20Mar%202025%2015%3A18%3A30%20GMT&granularity=hour&mode=training&dtype=amp&deviceName=cuda%20(a100)&lBranch=gh/jamesjwu/119/head&lCommit=be6d18d1307e269d4dac6594f25dae1941787e69&rBranch=gh/jamesjwu/119/base&rCommit=f461654bc928f8af1580b6ce1a13635af99e3ab3

(For posterity, base revision isf461654 and head revision isbe6d18d)

Overall looks better, no more accuracy failures or fail to runs. One model, XGLMForCausalLM, seemed to have regressed runtime perf, but I couldn't repro it locally. I checked on various other revisions (i.e. main, other base revisions I've run), and I think this particular base revision might have just gotten a high perf speedup due to variance (2.5x), though I will rerun the head revision one more time to confirm.

Ugh there's one more fail_accuracy on timm_efficientnet on torchbench. But I checked on main revisions, and it seems to have been failing accuracy until today?

Re: falling back to dynamic alloc: what do you think about falling back to something like alloca, to just allocate it directly on the stack? I know it might not be super portable, but maybe it's possible to support.

This seems fine to me.

For user-defined Triton kernels maybe we just always fallback to the upstream launcher? That way the codepath will at least be well tested.

For user-defined Triton kernels maybe we just always fallback to the upstream launcher? That way the codepath will at least be well tested.

This seems good to me, I had thought your point was to try to get it so we don't have to keep both codepaths. But we can make it so only user defined triton kernels use the old launcher system and everything else uses static launch.

One last round of benchmarks, showing stuff looks good (this is before I flipped the config to default off):

https://hud.pytorch.org/benchmark/compilers?dashboard=torchinductor&startTime=Thu%2C%2013%20Mar%202025%2015%3A39%3A38%20GMT&stopTime=Thu%2C%2020%20Mar%202025%2015%3A39%3A38%20GMT&granularity=hour&mode=inference&dtype=bfloat16&deviceName=cuda%20(a100)&lBranch=gh/jamesjwu/119/orig&lCommit=4673097e21a35158bb28402e35af65f620846f9f&rBranch=main&rCommit=e8a35eb7da1f7cbd62827943c1e00849a8592298

Will flip the config to default on once I land#149442 and some other in progress PRs to reach closer feature parity with triton's launcher.

@pytorchbot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in thewiki.

Questions? Feedback? Please reach out to thePyTorch DevX Team