-
Notifications
You must be signed in to change notification settings - Fork 3.7k
[Dtype] Low-precision Blackwell Datatype Support #18027
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
c9d49a3
to
c17241a
Compare
Thanks @DerrickYLJ for coauthoring it! Would you like to add yourself as coauthor in the PR? |
9f0210f
to
98596ff
Compare
cc @Hzfengsy |
Thanks for reminding me, I will take a close look tomorrow. also cc @LeiWang1999 |
2890d8e
to
52365da
Compare
Co-authored-by: DerrickYLJ <[email protected]>
52365da
to
1c36432
Compare
Co-authored-by: DerrickYLJ <[email protected]>
fef9930
to
481b6c8
Compare
@tvm-bot rerun |
1 similar comment
@tvm-bot rerun |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM.
This PR updates the datatype names in our NCCL integration to align with the latest changes in apache#18027. These changes were missed in the previous PR.
// E5M2 format, consistent with IEEE-754 | ||
case DataType::kFloat8_e4m3fnuz: | ||
// UE4M3 format, not consistent with IEEE-754 | ||
return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@Kathryn-cat , hi, why are the FloatConfig of e4m3, e4m3b11fnuz, e4m3fn and e4m3fnuz exactly the same?
This PR focuses on supporting FP4/FP8 data types introduced in Blackwell architectures (sm_100).
TVM nd array stores subbyte data types in compact format, thus two FP4 would be stored in 1 byte. The size calculator for array allocator is modified accordingly.
Subtype arithmetic
The type
__nv_fp4_e2m1
from<cuda_fp4.h>
is a tag type and does not support pointer arithmetic. Accordingly, the compiler does not support index operations on an array declared with__nv_fp4_e2m1
directly. If any index operations likearr[0] + arr[1]
is desired, user should declare the array as vector type like__nv_fp4x2_e2m1
.For example, suppose user creates an array A of type
__nv_fp4_e2m1
with values[-1 2 0.5 -6 -6 -2 2 3 4 1 -3 4 -2 2...]
Printing out values of A[0], A[1], ... will show
This is because
__nv_fp4_e2m1
is only a tag type. When it advances pointer, it advance by 1-byte at a time, yielding the upper 4 bits in the packed memory buffer. As a result, we should avoid directly doing indexing on__nv_fp4_e2m1
for arithmetic operations.If user passes in
__nv_fp4_e2m1
nd array and perform indexing, we can convert it to__nv_fp4x2_e2m1
and recalculate the indices if possible, but this requires more careful handling in the lowering process.Thus, the original corresponding test case in
test_target_codegen_cuda_fp4.py
is removed.