float4_e2m1fn and float8_e8m0fnu

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

cext/tile_kernel.cpp

@@ -46,8 +46,10 @@ static PyObject* g_default_tile_context;
     X(float16, 16, 1, kDLFloat) \
     X(float32, 32, 1, kDLFloat) \
     X(float64, 64, 1, kDLFloat) \
+    X(bfloat16, 16, 1, kDLBfloat) \
     X(float8_e4m3fn, 8, 1, kDLFloat8_e4m3fn) \
-    X(float8_e5m2, 8, 1, kDLFloat8_e5m2)
+    X(float8_e5m2, 8, 1, kDLFloat8_e5m2) \
+    X(float8_e8m0fnu, 8, 1, kDLFloat8_e8m0fnu)
 
 
 #define DECLARE_TORCH_DTYPE_GLOBAL(name, bitwidth, lanes, typecode) \

changelog.d/future/f8e8m0fnu_f4e2m1fn.md

@@ -0,0 +1,6 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Add `ct.float8_e8m0fnu` dtype (8-bit, unsigned, 8 exponent bits, 0 mantissa bits). A restricted float type.
+- Add `ct.float4_e2m1fn` dtype (4-bit, 1 sign bit, 2 exponent bits, 1 mantissa bit). A restricted float type.
+- Compiling float8_e8m0fnu, float4_e2m1fn operations for SM80 family, or SM90 family will raise `TileUnsupportedFeatureError`

src/cuda/tile/_bytecode/float.py

@@ -2,60 +2,113 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from enum import Enum, auto
 import math
 import struct
 from typing import NamedTuple
 
+from cuda.tile._exception import TileInternalError
+
 from .type import SimpleType
 
 
+class NonFiniteBehavior(Enum):
+    IEEE = auto()
+    NanOnlyAllOnes = auto()
+    FiniteOnly = auto()
+
+
 class _FloatSpec(NamedTuple):
     bitwidth: int
     emin: int
     emax: int
     exp_bits: int
     precision: int
-    finite_only: bool = False
+    non_finite_behavior: NonFiniteBehavior = NonFiniteBehavior.IEEE
 
 
 _specs = {
     SimpleType.F16: _FloatSpec(16, -14, 15, 5, 10),
     SimpleType.BF16: _FloatSpec(16, -126, 127, 8, 7),
     SimpleType.F32: _FloatSpec(32, -126, 127, 8, 23),
     SimpleType.TF32: _FloatSpec(19, -126, 127, 8, 10),
-    SimpleType.F8E4M3FN: _FloatSpec(8, -6, 8, 4, 3, finite_only=True),
+    SimpleType.F8E4M3FN: _FloatSpec(8, -6, 8, 4, 3, NonFiniteBehavior.NanOnlyAllOnes),
     SimpleType.F8E5M2: _FloatSpec(8, -14, 15, 5, 2),
+    SimpleType.F8E8M0FNU: _FloatSpec(8, -127, 127, 8, 0, NonFiniteBehavior.NanOnlyAllOnes),
+    SimpleType.F4E2M1FN: _FloatSpec(4, 0, 2, 2, 1, NonFiniteBehavior.FiniteOnly),
 }
 
 
-def float_max_value(ty: SimpleType) -> float:
-    return _specs[ty].max_value()
-
-
 def float_bit_size(ty: SimpleType) -> int:
     return 64 if ty == SimpleType.F64 else _specs[ty].bitwidth
 
 
 def float_to_bits(val: float, ty: SimpleType) -> int:
     if ty == SimpleType.F64:
         return struct.unpack("<Q", struct.pack("<d", val))[0]
+    elif ty == SimpleType.F8E8M0FNU:
+        spec = _specs[ty]
+        return _convert_f8e8m0fnu(val, *spec)
     else:
         spec = _specs[ty]
         return _convert_float(val, *spec)
 
 
+def _convert_f8e8m0fnu(val: float,
+                       bitwidth: int,
+                       emin: int,
+                       emax: int,
+                       _exp_bits: int,
+                       _precision: int,
+                       _non_finite_behavior: NonFiniteBehavior) -> int:
+    nan = (1 << bitwidth) - 1  # NaN is encoded as all ones
+    smallest_representable = 0  # 2^(-127)
+    bias = -emin  # no subnormals
+
+    if math.copysign(1.0, val) < 0:
+        raise TileInternalError("negative values cannot be represented in an unsigned float format")
+
+    if val == 0.0:
+        return smallest_representable
+    if not math.isfinite(val):
+        return nan
+
+    m, e = math.frexp(val)
+    m *= 2   # [1.0, 2.0)
+    e -= 1   # val = m * 2^e
+
+    if e > emax:
+        return nan
+    if e < emin:
+        return smallest_representable
+
+    m -= 1.0
+    e += bias
+
+    # With 0 mantissa bits, the only implicit significand value is 1 (odd),
+    # so round to nearest even (RNE) ties always round up.
+    round_up = (m >= 0.5)
+    if round_up:
+        m = 0
+        e += 1
+        if e > emax + bias:
+            return nan
+
+    return e
+
+
 def _convert_float(val: float,
                    bitwidth: int,
                    emin: int,
                    emax: int,
                    exp_bits: int,
                    precision: int,
-                   finite_only: bool) -> int:
+                   non_finite_behavior: NonFiniteBehavior) -> int:
     if val == 0.0:
         sign = math.copysign(1.0, val) < 0.0
         return sign << (bitwidth - 1)
     elif not math.isfinite(val):
-        return _convert_nonfinite(val, bitwidth, exp_bits, precision, finite_only)
+        return _convert_nonfinite(val, bitwidth, exp_bits, precision, non_finite_behavior)
 
     sign, val = (1, -val) if (val < 0) else (0, val)
     m, e = math.frexp(val)
@@ -64,7 +117,7 @@ def _convert_float(val: float,
 
     if e > emax:
         return _convert_nonfinite(-math.inf if sign else math.inf,
-                                  bitwidth, exp_bits, precision, finite_only)
+                                  bitwidth, exp_bits, precision, non_finite_behavior)
 
     if e < emin:
         m = math.ldexp(m, e - emin)
@@ -73,26 +126,37 @@ def _convert_float(val: float,
         m -= 1.0
         e += -emin + 1
 
-    # Round to nearest, ties to even
+    # Round to nearest, ties to even (RNE)
+    # The following RNE implementation breaks when precision is 0
+    assert precision > 0
     m = round(m * (1 << precision))
     if m == (1 << precision):
         m = 0
         e += 1
         if e > emax - emin + 1:
             return _convert_nonfinite(-math.inf if sign else math.inf,
-                                      bitwidth, exp_bits, precision, finite_only)
+                                      bitwidth, exp_bits, precision, non_finite_behavior)
     bits = (sign << (bitwidth - 1)) | (e << precision) | m
     return bits
 
 
-def _convert_nonfinite(val, bitwidth, exp_bits, precision, finite_only) -> int:
-    if finite_only:
-        # NaN is encoded as all ones
-        sign = math.copysign(1.0, val) < 0.0
-        return (sign << (bitwidth - 1)) | ((1 << (bitwidth - 1)) - 1)
-    else:
-        # Exponent is all ones. Truncate the low bits, preserve the rest of the payload
-        float64_bits, = struct.unpack("<Q", struct.pack("<d", val))
-        payload = (float64_bits >> (52 - precision)) & ((1 << precision) - 1)
-        hi_bits = (float64_bits >> (63 - exp_bits)) << precision
-        return hi_bits | payload
+def _convert_nonfinite(val, bitwidth, exp_bits, precision, non_finite_behavior) -> int:
+    match non_finite_behavior:
+        case NonFiniteBehavior.NanOnlyAllOnes:
+            # NaN is encoded as all ones
+            sign = math.copysign(1.0, val) < 0.0
+            return (sign << (bitwidth - 1)) | ((1 << (bitwidth - 1)) - 1)
+        case NonFiniteBehavior.FiniteOnly:
+            if math.isnan(val):
+                raise TileInternalError("NaN cannot be represented in a finite-only float format")
+            # Clamp to max representable magnitude, preserve sign
+            sign = math.copysign(1.0, val) < 0.0
+            return (sign << (bitwidth - 1)) | ((1 << (bitwidth - 1)) - 1)
+        case NonFiniteBehavior.IEEE:
+            # Exponent is all ones. Truncate the low bits, preserve the rest of the payload
+            float64_bits, = struct.unpack("<Q", struct.pack("<d", val))
+            payload = (float64_bits >> (52 - precision)) & ((1 << precision) - 1)
+            hi_bits = (float64_bits >> (63 - exp_bits)) << precision
+            return hi_bits | payload
+        case _:
+            assert False

src/cuda/tile/_bytecode/type.py

@@ -44,7 +44,8 @@ class SimpleType(enum.Enum):
     F8E4M3FN = b"\x0a"
     F8E5M2 = b"\x0b"
     Token = b"\x11"
-    Unknown = b"\x12"
+    F8E8M0FNU = b"\x12"  # since 13.2
+    F4E2M1FN = b"\x13"  # since 13.3
 
 
 class _CompositeType(enum.Enum):

src/cuda/tile/_compile.py

@@ -50,7 +50,7 @@
 )
 
 from cuda.tile._passes.alias_analysis import alias_analysis_pass
-from cuda.tile._passes.check_ampere_fp8 import check_ampere_fp8
+from cuda.tile._passes.check_dtype_support import check_dtype_support
 from cuda.tile._passes.dce import dead_code_elimination_pass
 from cuda.tile._passes.token_order import token_order_pass
 from cuda.tile._cache import cache_key, cache_lookup, cache_store, evict_lru
@@ -219,7 +219,7 @@ def compile_tile(pyfunc,
         print(f'\n{code}', file=sys.stderr)
 
     sm_arch = get_sm_arch()
-    check_ampere_fp8(func_ir.body, sm_arch)
+    check_dtype_support(func_ir.body, sm_arch, bytecode_version)
 
     bytecode_generator = functools.partial(generate_bytecode_for_kernel,
                                            func_ir, compiler_options, sm_arch)

src/cuda/tile/_datatype.py

@@ -17,6 +17,7 @@
            "int8", "int16", "int32", "int64",
            "float16", "float32", "float64",
            "bfloat16", "tfloat32", "float8_e4m3fn", "float8_e5m2",
+           "float8_e8m0fnu", "float4_e2m1fn",
            "DType", "NumericDType", "ArithmeticDType",
            "NumericDTypeCategories"]
 
@@ -139,17 +140,25 @@ class ArithmeticDType(NumericDType):
 and 7 mantissa bits."""
 
 tfloat32 = NumericDType("tfloat32", 32, float, bc.SimpleType.TF32)
-tfloat32.__doc__ = """A 32-bit tensor floating-point |arithmetic dtype| with 1 sign \
+tfloat32.__doc__ = """A 32-bit tensor floating-point |numeric dtype| with 1 sign \
 bit, 8 exponent bits, and 10 mantissa bits (19-bit representation stored in 32-bit container)."""
 
 float8_e4m3fn = NumericDType("float8_e4m3fn", 8, float, bc.SimpleType.F8E4M3FN)
-float8_e4m3fn.__doc__ = """A 8-bit floating-point |arithmetic dtype| with 1 sign bit, \
+float8_e4m3fn.__doc__ = """An 8-bit floating-point |numeric dtype| with 1 sign bit, \
 4 exponent bits, and 3 mantissa bits."""
 
 float8_e5m2 = NumericDType("float8_e5m2", 8, float, bc.SimpleType.F8E5M2)
-float8_e5m2.__doc__ = """A 8-bit floating-point |arithmetic dtype| with 1 sign bit, \
+float8_e5m2.__doc__ = """An 8-bit floating-point |numeric dtype| with 1 sign bit, \
 5 exponent bits, and 2 mantissa bits."""
 
+float8_e8m0fnu = NumericDType("float8_e8m0fnu", 8, float, bc.SimpleType.F8E8M0FNU)
+float8_e8m0fnu.__doc__ = """An 8-bit floating-point |numeric dtype| with no sign bit, \
+8 exponent bits, and 0 mantissa bits."""
+
+float4_e2m1fn = NumericDType("float4_e2m1fn", 4, float, bc.SimpleType.F4E2M1FN)
+float4_e2m1fn.__doc__ = """A 4-bit floating-point |numeric dtype| with 1 sign bit, \
+2 exponent bits, and 1 mantissa bit."""
+
 
 class DTypeEnum(IntEnum):
     B1 = 0
@@ -168,6 +177,8 @@ class DTypeEnum(IntEnum):
     TF32 = 13
     F8E4M3FN = 14
     F8E5M2 = 15
+    F8E8M0FNU = 16
+    F4E2M1FN = 17
 
 
 dtype_to_enum = {
@@ -187,6 +198,8 @@ class DTypeEnum(IntEnum):
     tfloat32: DTypeEnum.TF32,
     float8_e4m3fn: DTypeEnum.F8E4M3FN,
     float8_e5m2: DTypeEnum.F8E5M2,
+    float8_e8m0fnu: DTypeEnum.F8E8M0FNU,
+    float4_e2m1fn: DTypeEnum.F4E2M1FN,
 }
 _enum_to_dtype = dict((i, t) for t, i in dtype_to_enum.items())
 
@@ -209,7 +222,7 @@ class NumericDTypeCategories:
     Boolean = [bool_]
     Integral = [uint8, uint16, uint32, uint64, int8, int16, int32, int64]
     Float = [float16, float32, float64, bfloat16]
-    RestrictedFloat = [tfloat32, float8_e4m3fn, float8_e5m2]
+    RestrictedFloat = [tfloat32, float8_e4m3fn, float8_e5m2, float8_e8m0fnu, float4_e2m1fn]
 
     @classmethod
     def get_category(cls, t: DType) -> NumericDTypeCategory:
@@ -323,6 +336,8 @@ class _DTypePromotionImpl:
     tf = DTypeEnum.TF32
     f8e4m3fn = DTypeEnum.F8E4M3FN
     f8e5m2 = DTypeEnum.F8E5M2
+    f8e8m0fnu = DTypeEnum.F8E8M0FNU
+    f4e2m1fn = DTypeEnum.F4E2M1FN
     na = None
 
     # Entries for restricted arithmetic dtypes will never be reached, but we need to keep them
@@ -337,23 +352,25 @@ class _DTypePromotionImpl:
     # Restricted floats requires explicit type cast
     # Float16 and BFloat 16 requires explicit type cast
     _common_dtype_table = [
-     # b1, u8, u16, u32, u64, i8,  i16, i32, i64, f16, f32, f64, bf,  tf, f8e4m3fn, f8e5m2
-     [b1,  u8,  u16, u32, u64, i8,  i16, i32, i64, f16, f32, f64, bf,  na,  na,       na],    # b1
-     [u8,  u8,  u16, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na],    # u8
-     [u16, u16, u16, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na],    # u16
-     [u32, u32, u32, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na],    # u32
-     [u64, u64, u64, u64, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na],    # u64
-     [i8,  na,  na,  na,  na,  i8,  i16, i32, i64, f16, f32, f64, bf,  na,  na,       na],    # i8
-     [i16, na,  na,  na,  na,  i16, i16, i32, i64, f16, f32, f64, bf,  na,  na,       na],    # i16
-     [i32, na,  na,  na,  na,  i32, i32, i32, i64, f16, f32, f64, bf,  na,  na,       na],    # i32
-     [i64, na,  na,  na,  na,  i64, i64, i64, i64, f16, f32, f64, bf,  na,  na,       na],    # i64
-     [f16, f16, f16, f16, f16, f16, f16, f16, f16, f16, f32, f64, na,  na,  na,       na],    # f16
-     [f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f64, f32, na,  na,       na],    # f32
-     [f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, na,  na,       na],    # f64
-     [bf,  bf,  bf,  bf,  bf,  bf,  bf,  bf,  bf,  na,  f32, f64, bf,  na,  na,       na],    # bf
-     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  tf,  na,       na],    # tf
-     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  f8e4m3fn, na],    # f8e4m3fn # noqa
-     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,   f8e5m2],    # f8e5m2  # noqa
+     # b1, u8, u16, u32, u64, i8,  i16, i32, i64, f16, f32, f64, bf,  tf, f8e4m3fn, f8e5m2, f8e8m0fnu, f4e2m1fn    # noqa
+     [b1,  u8,  u16, u32, u64, i8,  i16, i32, i64, f16, f32, f64, bf,  na,  na,       na,        na,        na],        # b1  # noqa
+     [u8,  u8,  u16, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na,        na,        na],        # u8  # noqa
+     [u16, u16, u16, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na,        na,        na],        # u16  # noqa
+     [u32, u32, u32, u32, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na,        na,        na],        # u32  # noqa
+     [u64, u64, u64, u64, u64, na,  na,  na,  na,  f16, f32, f64, bf,  na,  na,       na,        na,        na],        # u64  # noqa
+     [i8,  na,  na,  na,  na,  i8,  i16, i32, i64, f16, f32, f64, bf,  na,  na,       na,        na,        na],        # i8  # noqa
+     [i16, na,  na,  na,  na,  i16, i16, i32, i64, f16, f32, f64, bf,  na,  na,       na,        na,        na],        # i16  # noqa
+     [i32, na,  na,  na,  na,  i32, i32, i32, i64, f16, f32, f64, bf,  na,  na,       na,        na,        na],        # i32  # noqa
+     [i64, na,  na,  na,  na,  i64, i64, i64, i64, f16, f32, f64, bf,  na,  na,       na,        na,        na],        # i64  # noqa
+     [f16, f16, f16, f16, f16, f16, f16, f16, f16, f16, f32, f64, na,  na,  na,       na,        na,        na],        # f16  # noqa
+     [f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f64, f32, na,  na,       na,        na,        na],        # f32  # noqa
+     [f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, f64, na,  na,       na,        na,        na],        # f64  # noqa
+     [bf,  bf,  bf,  bf,  bf,  bf,  bf,  bf,  bf,  na,  f32, f64, bf,  na,  na,       na,        na,        na],        # bf  # noqa
+     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  tf,  na,       na,        na,        na],        # tf  # noqa
+     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  f8e4m3fn, na,        na,        na],        # f8e4m3fn  # noqa
+     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,       f8e5m2,    na,        na],        # f8e5m2  # noqa
+     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,       na,        f8e8m0fnu, na],        # f8e8m0fnu  # noqa
+     [na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,  na,       na,        na,        f4e2m1fn],  # f4e2m1fn  # noqa
     ]
 
     @classmethod
@@ -423,14 +440,22 @@ def _resolve_mma_supported_dtype(x_dtype: DType,
 
 def _generate_rst_dtype_promotion_table() -> str:
     """Generate an RST table representation of the dtype promotion rules."""
-    return _generate_rst_table(_DTypePromotionImpl._common_dtype_table)
+    import cuda.tile
+    # Skip dtypes not exposed in cuda.tile yet. Promomotion table is append only.
+    n = sum(1 for dtype in _enum_to_dtype.values() if hasattr(cuda.tile, dtype.name))
+    table = _DTypePromotionImpl._common_dtype_table
+    return _generate_rst_table([row[:n] for row in table[:n]])
 
 
 def _generate_rst_numeric_dtypes() -> str:
     """Generate RST documentation for numeric datatypes."""
+    import cuda.tile
     content = []
 
     for dtype in numeric_dtypes:
+        # Skip dtypes not exposed in cuda.tile yet
+        if not hasattr(cuda.tile, dtype.name):
+            continue
         content.append(f".. autodata:: cuda.tile.{dtype.name}")
         content.append("   :annotation:")
         content.append("")  # Empty line between types

src/cuda/tile/_ir/ir.py

@@ -16,7 +16,8 @@
 from typing import (
     List, Optional, Dict, Tuple, Any, TYPE_CHECKING, Sequence, Iterator
 )
-from .type import Type, InvalidType
+
+from cuda.tile._ir.type import Type, InvalidType
 from cuda.tile._exception import (
     TileTypeError, Loc, TileInternalError
 )

src/cuda/tile/_ir/typing_support.py

@@ -249,7 +249,9 @@ def get_constant_value(val: Any) -> Any:
 # =====CuTile native support ===========
 # register cuTile native dtype types
 for dtype in datatype.dtype_to_enum:
-    register_dtypes({dtype: dtype}, usable_as_constructor=True)
+    # only allow byte aligned dtypes as constructors
+    usable_as_constructor = (dtype.bitwidth % 8 == 0)
+    register_dtypes({dtype: dtype}, usable_as_constructor)
 
 
 # ========= Numpy support ===========
@@ -322,7 +324,8 @@ def get_constant_value(val: Any) -> Any:
         torch.bool: datatype.bool_,
         torch.bfloat16: datatype.bfloat16,
         torch.float8_e4m3fn: datatype.float8_e4m3fn,
-        torch.float8_e5m2: datatype.float8_e5m2
+        torch.float8_e5m2: datatype.float8_e5m2,
+        torch.float8_e8m0fnu: datatype.float8_e8m0fnu,
     })
 
     @register_type_handler(torch.Tensor, allow_subtypes=True)