Skip to content

Commit ba9f174

Browse files
authored
[3rdparty] Bump DLPack to v1.1 for float8/6/4 dtype supports (#17831)
This PR bumps the 3rdparty dlpack version to v1.1, which brings the support of float8, float6 and float4.
1 parent 982b46c commit ba9f174

File tree

3 files changed

+23
-7
lines changed

3 files changed

+23
-7
lines changed

include/tvm/runtime/data_type.h

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,17 @@ class DataType {
5656
kFloat = kDLFloat,
5757
kHandle = TVMArgTypeCode::kTVMOpaqueHandle,
5858
kBFloat = kDLBfloat,
59-
kFloat8_e4m3fn = 6U,
60-
kFloat8_e5m2 = 7U,
61-
kFloat4_e2m1fn = 8U,
59+
kFloat8_e3m4 = kDLFloat8_e3m4,
60+
kFloat8_e4m3 = kDLFloat8_e4m3,
61+
kFloat8_e4m3b11fnuz = kDLFloat8_e4m3b11fnuz,
62+
kFloat8_e4m3fn = kDLFloat8_e4m3fn,
63+
kFloat8_e4m3fnuz = kDLFloat8_e4m3fnuz,
64+
kFloat8_e5m2 = kDLFloat8_e5m2,
65+
kFloat8_e5m2fnuz = kDLFloat8_e5m2fnuz,
66+
kFloat8_e8m0fnu = kDLFloat8_e8m0fnu,
67+
kFloat6_e2m3fn = kDLFloat6_e2m3fn,
68+
kFloat6_e3m2fn = kDLFloat6_e3m2fn,
69+
kFloat4_e2m1fn = kDLFloat4_e2m1fn,
6270
kCustomBegin = 129
6371
};
6472
/*! \brief default constructor */

python/tvm/_ffi/runtime_ctypes.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,17 @@ class DataTypeCode(object):
6666
FLOAT = 2
6767
HANDLE = 3
6868
BFLOAT = 4
69-
Float8E4M3FN = 6
70-
Float8E5M2 = 7
71-
Float4E2M1FN = 8
69+
Float8E3M4 = 7
70+
Float8E4M3 = 8
71+
Float8E4M3B11FNUZ = 9
72+
Float8E4M3FN = 10
73+
Float8E4M3FNUZ = 11
74+
Float8E5M2 = 12
75+
Float8E5M2FNUZ = 13
76+
Float8E8M0FNU = 14
77+
Float6E2M3FN = 15
78+
Float6E3M2FN = 16
79+
Float4E2M1FN = 17
7280

7381

7482
class DataType(ctypes.Structure):

0 commit comments

Comments
 (0)