NVFP4 recipe with GEMM via BF16 dequant#518
Merged
matthiasdiener merged 101 commits intodevfrom Apr 29, 2026
Merged
Conversation
Remove TODO regarding userbuffers
Userbuffer Enablement for ROCm
* Update Dockerfile to use ROCm TheRock * Update wheels building script to work with ROCm TheRock and the latest Manylinux image * Support default ROCm location /opt/rocm/core * Fix UB code build on TheRock * Support comma separated list of target GPU architectures * Guess ROCm build from HIP_PLATFORM
ipanfilo
reviewed
Apr 23, 2026
aris134
reviewed
Apr 27, 2026
aris134
reviewed
Apr 27, 2026
Comment on lines
+112
to
+113
| const float fp8_max = te_fp8_fnuz() ? 240.0f : 448.0f; | ||
| const float factor_inv = 1.0f / (6.0f * fp8_max); |
Contributor
There was a problem hiding this comment.
Same comment as above regarding using Numeric_Traits_fp8e4m3 here
aris134
reviewed
Apr 27, 2026
aris134
reviewed
Apr 27, 2026
Comment on lines
+438
to
+470
| if (is_fp4_dtype(param.Atype)) { | ||
| hip_bfloat16* a_bf16 = reinterpret_cast<hip_bfloat16*>(ws_ptr); | ||
| ws_ptr += a_bf16_bytes; | ||
| const int64_t total_a = static_cast<int64_t>(m) * k; | ||
| const auto& a_sinv = (transa == CUBLAS_OP_T) ? inputA.scale_inv | ||
| : inputA.columnwise_scale_inv; | ||
| const int64_t a_num_cols = (transa == CUBLAS_OP_T) | ||
| ? inputA.data.shape.back() | ||
| : inputA.columnwise_data.shape.back(); | ||
| const int64_t a_scale_stride = (a_sinv.shape.size() >= 2) ? a_sinv.shape[1] : (a_num_cols / 16); | ||
| launch_dequant_fp4_to_bf16(param.A, param.A_scale_inv, a_bf16, total_a, | ||
| a_num_cols, a_scale_stride, stream); | ||
| param.A = a_bf16; | ||
| param.Atype = DType::kBFloat16; | ||
| param.A_scale_inv = nullptr; | ||
| } | ||
|
|
||
| if (is_fp4_dtype(param.Btype)) { | ||
| hip_bfloat16* b_bf16 = reinterpret_cast<hip_bfloat16*>(ws_ptr); | ||
| ws_ptr += b_bf16_bytes; | ||
| const int64_t total_b = static_cast<int64_t>(k) * n; | ||
| const auto& b_sinv = (transb == CUBLAS_OP_N) ? inputB.scale_inv | ||
| : inputB.columnwise_scale_inv; | ||
| const int64_t b_num_cols = (transb == CUBLAS_OP_N) | ||
| ? inputB.data.shape.back() | ||
| : inputB.columnwise_data.shape.back(); | ||
| const int64_t b_scale_stride = (b_sinv.shape.size() >= 2) ? b_sinv.shape[1] : (b_num_cols / 16); | ||
| launch_dequant_fp4_to_bf16(param.B, param.B_scale_inv, b_bf16, total_b, | ||
| b_num_cols, b_scale_stride, stream); | ||
| param.B = b_bf16; | ||
| param.Btype = DType::kBFloat16; | ||
| param.B_scale_inv = nullptr; | ||
| } |
Contributor
There was a problem hiding this comment.
Minor comment: would it make sense to factor the repeated FP4→BF16 staging logic for A/B into a small helper? The two blocks look structurally similar, aside from the operand-specific shape/layout details.
Contributor
Author
There was a problem hiding this comment.
Thanks, I factored this out into a lambda function in fae76d3
aris134
approved these changes
Apr 27, 2026
Contributor
aris134
left a comment
aris134
left a comment
There was a problem hiding this comment.
LGTM. I left one minor non-blocking suggestion, but this looks good to me overall.
ipanfilo
requested changes
Apr 28, 2026
matthiasdiener
force-pushed
the
mdiener/nvfp4-gemm
branch
7 times, most recently
from
0f240ad to
9ed88ff
Compare
matthiasdiener
force-pushed
the
mdiener/nvfp4-gemm
branch
from
9ed88ff to
a6f4787
Compare
ipanfilo
approved these changes
Apr 29, 2026
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
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.
Description
Part of https://github.com/ROCm/frameworks-internal/issues/15682
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: