Lightning-AI / Lightning-AI/lightning-thunder
FP8 Linear and conv with cudnn
@vedaanta is already working on this.
Since Jun 3, 2024.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
CuDNN provides flexible support for performant gemm/conv with fp8 quantization. Thunder introducing fp8 casts in its traces can benefit from cudnn fusions.
### Motivation
Today, thunder uses TE's fp8 linear which delegates quantization strategies to TE, making it opaque to thunder. If thunder plans to handle fp8 casts itself, performant and flexible kernels from cudnn can help.
Cudnn's support is described here: [cuDNN's runtime fusion engine](https://docs.nvidia.com/deeplearning/cudnn/latest/developer/graph-api.html#generic-runtime-fusion-engines)
For fp8 specifically, cudnn can provide the following graph as one fused kernel:
The graph is flexible, meaning:
- final output, C, can be in fp16/bf16/fp32.
- amax operation can be skipped for static scaling strategies
The corresponding backward graphs are also supported. (Though they require offline transpose on Hopper)
### Pitch
Have cudnn executor claim gemm/conv along with the fp8 casts around them.
CC @IvanYashchuk @kshitij12345 @Anerudhan
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.