Reducing the cuDNN Download and Installation sizes - cuDNN JIT

[jhalabi-nv]

Hello,

The cuDNN team is working on a new cuDNN package to reduce our download and installation size.

There are two key pieces:

1. A new “library configuration” that tells cuDNN that only the “graph JIT” subset of the library is needed. This is set with an environment variable at runtime.
2. A set of corresponding packages (e.g. PyPI, apt, etc.) that only include that minimal subset.

1 is already available, but 2 is still a WIP. We’re posting here because we understand that the download and installation size of cuDNN is a pain point for this project, and we want to keep you in the loop about our roadmap. Also, you can try out 1 and emulate 2 by removing the unnecessary sub-libraries, this can serve as a proof of concept.

More specifically about 1:

In cuDNN 9.2, we released the GRAPH_JIT_ONLY configuration mode. The cuDNN flash attention kernels are runtime compiled, and thus included in this GRAPH_JIT_ONLY configuration. This means that you could enable GRAPH_JIT_ONLY and remove the unnecessary cuDNN sub-libraries, which would reduce the cuDNN installation size from 835MB to 14MB.

The cuDNN Library Configuration section illustrates how this can be achieved. You’ll need to set an environment variable and remove any non-required cuDNN sub-libraries.

Then once we have 2 available which should reduce the download size from 435MB to 6MB on Ubuntu, we'll create a PR to update the llm.c documentation. That's coming soon.

In the meantime, we'd be excited to see if anyone can repro our setup with GRAPH_JIT_ONLY, , demonstrating the significantly reduced cuDNN footprint.