[Quantization] Add cutlass kernel for FP8 #43304
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
|
||
| _quantization_kernel = get_kernel("RedHatAI/quantization") | ||
| except Exception as e: | ||
| logger.warning_once(f"Failed to load CUTLASS quantization kernel: {e}. Falling back to Triton.") |
There was a problem hiding this comment.
Do we want to also log that we're using the Redhat kernel in case it was successfully loaded?
There was a problem hiding this comment.
Yes we can do that I think
| Check if CUTLASS blockwise FP8 matmul is supported for the given block size. | ||
| CUTLASS blockwise kernels require: | ||
| - SM90+ (Hopper or newer) |
There was a problem hiding this comment.
This is fine IMO. When hardware is available, users should be able to max them out!
| kernel = _get_quantization_kernel() | ||
| if kernel is None: | ||
| return False |
There was a problem hiding this comment.
I think we also need to check that kernels is installed no?
| # Global for the CUTLASS quantization kernel (lazily loaded) | ||
| _quantization_kernel = None | ||
|
|
||
|
|
||
| def _get_quantization_kernel(): | ||
| """Lazily load the CUTLASS quantization kernel from HuggingFace Hub.""" | ||
| global _quantization_kernel | ||
| if _quantization_kernel is None: | ||
| try: | ||
| from .hub_kernels import get_kernel | ||
|
|
||
| _quantization_kernel = get_kernel("RedHatAI/quantization") | ||
| except Exception as e: | ||
| logger.warning_once(f"Failed to load CUTLASS quantization kernel: {e}. Falling back to Triton.") |
There was a problem hiding this comment.
instead of a single kernel, can we try to create a dict where we store multiple kernels ? We can leave this for a follow-up PR but the idea would be to also move the triton kernels in kernels.
|
View the CircleCI Test Summary for this PR: https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43304&sha=cfd4b9 |
What does this PR do?
adds the cutlass kernel for scaled matmul, the performance is much better than triton for the specific block size : (128, 128):
All FP8 tests passing !