Demonstrate adding a cuda build + cuda custom op to torchao

[zou3519]

We copy-paste a CUDA kernel from torchvision for nms to serve as the example.

[msaroufim]

will likely need to conda install cuda -c nvidia in CI since I don't think nvcc will be available

[msaroufim]

more of a nit but should this be a recursive=True so we can have kernels grouped by folders under the cuda/ folder

[msaroufim]

ok i guess as a first useful kernel to merge, I'm gonna adapt this example to use paged attention

[msaroufim]

did you write this or was it codegened?

EDIT: My question is actually maybe can we codegen this? same as load_inline does

[zou3519]

This is not codegenned. We could codegen it, but in general codegen adds more complexity. What were you thinking for the API?

[cpuhrsch]

Thanks for sending this!

I guess this needs some version guards because it requires features from the nightlies? Or is the CI failure real for 2.2?

Processing /home/runner/work/ao/ao
  Installing build dependencies: started
  Installing build dependencies: finished with status 'done'
  Getting requirements to build wheel: started
  Getting requirements to build wheel: finished with status 'error'
  error: subprocess-exited-with-error
  
  × Getting requirements to build wheel did not run successfully.
  │ exit code: 1
  ╰─> [17 lines of output]
      Traceback (most recent call last):
        File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
          main()
        File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
          json_out['return_val'] = hook(**hook_input['kwargs'])
        File "/opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 11[8](https://github.com/pytorch-labs/ao/actions/runs/8634345339/job/23669554393?pr=130#step:5:9), in get_requires_for_build_wheel
          return hook(config_settings)
        File "/tmp/pip-build-env-4v28lou4/overlay/lib/python3.[9](https://github.com/pytorch-labs/ao/actions/runs/8634345339/job/23669554393?pr=130#step:5:10)/site-packages/setuptools/build_meta.py", line 325, in get_requires_for_build_wheel
          return self._get_build_requires(config_settings, requirements=['wheel'])
        File "/tmp/pip-build-env-4v28lou4/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 295, in _get_build_requires
          self.run_setup()
        File "/tmp/pip-build-env-4v28lou4/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 487, in run_setup
          super().run_setup(setup_script=setup_script)
        File "/tmp/pip-build-env-4v28lou4/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 3[11](https://github.com/pytorch-labs/ao/actions/runs/8634345339/job/23669554393?pr=130#step:5:12), in run_setup
          exec(code, locals())
        File "<string>", line 7, in <module>
      ModuleNotFoundError: No module named 'torch'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error

× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip.

[msaroufim]

This was upstreamed as part of #135