Ingest FP8 attn scales and use them in ROCm FlashAttention #338
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
mawong-amd
force-pushed
the
ingest_fp8_attn_scales
branch
from
4e42946 to
9ba2fab
Compare
gshtras
approved these changes
Dec 19, 2024
gshtras
reviewed
Dec 19, 2024
| @@ -428,7 +434,9 @@ def load_weights(self, weights: Iterable[Tuple[str, | |||
| param = params_dict[scale_name] | |||
| weight_loader = getattr(param, "weight_loader", | |||
| default_weight_loader) | |||
| loaded_weight = loaded_weight[0] | |||
| if loaded_weight.shape: | |||
There was a problem hiding this comment.
I targeted only the models which unconditionally do this logic
gshtras
requested changes
Dec 19, 2024
There was a problem hiding this comment.
Without explicitly disabling VLLM_USE_ROCM_FP8_ATTN now quark quantized models (amd/Meta-Llama-3.1-70B-Instruct-FP8-KV) fail with a triton exception:
python: /root/.triton/llvm/llvm-c08c6a71-ubuntu-x64/include/llvm/Support/Casting.h:566: decltype(auto) llvm::cast(const From&) [with To = mlir::detail::TypedValue<mlir::RankedTensorType>; From = mlir::OpResult]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.
…ntion for dynamic quantization
EDIT: @ilia-cher identified the issue and provided a simple fix that works on older Triton. Still recommended to upgrade to latest. |
mawong-amd
force-pushed
the
ingest_fp8_attn_scales
branch
from
9639307 to
1ed1389
Compare
mawong-amd
changed the title
gshtras
previously approved these changes
Dec 19, 2024
gshtras
approved these changes
Dec 19, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Thanks to the work of @ilia-cher in #301, Triton FA supports per-tensor quantized FP8 almost-everything (quantized first and second GEMMs and attention output <- needs per-tensor quantized Q, K, V, softmax(QK^T) and corresponding scales).
This PR enables the aforementioned quantization routines in Triton FA and ROCm PA if a quantized (text-only) Llama model contains attention output scales and an appropriate environment variable is set (
VLLM_USE_ROCM_FP8_FLASH_ATTN={True/1}, off by default). Extending this to other model architectures is straightforward but not done for now. Accuracy might dip for Triton FA if not all scales are present in the model.