[Feature] Add float4_e2m1_unpacked dtype (#2271)

* Add float4_e2m1_unpacked dtype

* Fix TIR cast for float4_e2m1_unpacked dtype.

Route custom[float4_e2m1_unpacked] through Float4E2M1Unpacked FFI
so T.float4_e2m1_unpacked(0.0) works in prim_func builders.

* Limit float4_e2m1_unpacked vector lanes and fix TMA transaction bytes.

Cap unpacked FP4 vector dtypes at x32 and count TMA/mbarrier
transactions using 4-bit payload size (0.5 byte per element).

* lint

Author: Rachmanino

src/backend/cuda/op/copy.cc

@@ -42,9 +42,17 @@ PrimExpr MakeTmaLeaderCondition(PrimExpr thread_extent) {
   return Call(DataType::Bool(), tl_shuffle_elect(), {std::move(thread_extent)});
 }
 
-PrimExpr TMABytesFromElements(PrimExpr elements, DataType dtype) {
+int TMATransactionElementBits(DataType dtype) {
+  // float4_e2m1_unpacked stores one logical FP4 value per 8-bit SMEM slot,
+  // but TMA/mbarrier transaction counts reflect the moved FP4 payload (4b).
+  if (dtype.is_float4_e2m1_unpacked()) {
+    return 4;
+  }
+  return dtype.bits();
+}
+
+PrimExpr TMABytesFromElements(PrimExpr elements, DataType dtype, int bits) {
   PrimExpr elements_i64 = cast(DataType::Int(64), elements);
-  int bits = dtype.bits();
   if (bits % 8 == 0) {
     return elements_i64 * IntImm(DataType::Int(64), bits / 8);
   }
@@ -53,8 +61,26 @@ PrimExpr TMABytesFromElements(PrimExpr elements, DataType dtype) {
                   IntImm(DataType::Int(64), 8));
 }
 
+int64_t TMABytesFromElements(int64_t elements, DataType dtype, int bits) {
+  return (elements * bits + 7) / 8;
+}
+
+PrimExpr TMABytesFromElements(PrimExpr elements, DataType dtype) {
+  return TMABytesFromElements(elements, dtype, dtype.bits());
+}
+
 int64_t TMABytesFromElements(int64_t elements, DataType dtype) {
-  return (elements * dtype.bits() + 7) / 8;
+  return TMABytesFromElements(elements, dtype, dtype.bits());
+}
+
+PrimExpr TMATransactionBytesFromElements(PrimExpr elements, DataType dtype) {
+  return TMABytesFromElements(elements, dtype,
+                              TMATransactionElementBits(dtype));
+}
+
+int64_t TMATransactionBytesFromElements(int64_t elements, DataType dtype) {
+  return TMABytesFromElements(elements, dtype,
+                              TMATransactionElementBits(dtype));
 }
 
 int64_t TMAElementsForBytes(int64_t bytes, DataType dtype) {
@@ -1656,10 +1682,11 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T,
     PrimExpr total_bytes;
     if ((*inner_box_dim) != instruction_dim) {
       int loop_extent = (*inner_box_dim) / instruction_dim;
-      total_bytes = TMABytesFromElements(total_elements * loop_extent,
-                                         shared_tensor->dtype);
+      total_bytes = TMATransactionBytesFromElements(
+          total_elements * loop_extent, shared_tensor->dtype);
     } else {
-      total_bytes = TMABytesFromElements(total_elements, shared_tensor->dtype);
+      total_bytes =
+          TMATransactionBytesFromElements(total_elements, shared_tensor->dtype);
     }
 
     Stmt barrier_before_tma_stmt;
@@ -1959,7 +1986,8 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &T,
   }
 
   Stmt tma_copy;
-  PrimExpr total_bytes = TMABytesFromElements(elements, shared_tensor->dtype);
+  PrimExpr total_bytes =
+      TMATransactionBytesFromElements(elements, shared_tensor->dtype);
   if (is_load) {
     PrimExpr mbar_arg = barrier_base_id >= 0 ? mbar_handle : PrimExpr(0);
     tma_copy = Evaluate(Call(DataType::Handle(), tma_load(),

testing/python/language/test_tilelang_language_float4_e2m1_unpacked_dtype.py

@@ -0,0 +1,66 @@
+"""Tests for float4_e2m1_unpacked (CUTLASS float_e2m1_unpacksmem_t)."""
+
+import tilelang.language as T
+from tilelang import tvm as tvm
+
+
+def test_float4_e2m1_unpacked_dtype_properties():
+    dt = T.float4_e2m1_unpacked
+    assert dt.bits == 8
+    assert int(dt.type_code) == 131
+    assert str(dt) == "custom[float4_e2m1_unpacked]8"
+    assert dt.lanes == 1
+
+
+def test_float4_e2m1_unpacked_vector_lanes():
+    dt2 = T.float4_e2m1_unpackedx2
+    assert dt2.bits == 8
+    assert int(dt2.type_code) == 131
+    assert dt2.lanes == 2
+
+
+def test_float4_e2m1_unpacked_tir_cast():
+    from tilelang import language as T_lang
+
+    @T_lang.prim_func
+    def main():
+        T_lang.evaluate(T_lang.float4_e2m1_unpacked(0.0))
+
+    assert main.body.value.dtype.is_float4_e2m1_unpacked()
+
+
+def test_float4_e2m1_unpacked_distinct_from_packed():
+    packed = T.float4_e2m1fn
+    unpacked = T.float4_e2m1_unpacked
+    assert packed.bits == 4
+    assert unpacked.bits == 8
+    assert packed != unpacked
+
+
+def test_float4_e2m1_unpacked_dtype_helpers():
+    from tilelang.language.dtypes import is_float4, is_float4_e2m1fn, is_float4_e2m1_unpacked
+
+    packed = T.float4_e2m1fn
+    unpacked = T.float4_e2m1_unpacked
+    assert packed.is_float4_e2m1fn()
+    assert not packed.is_float4_e2m1_unpacked()
+    assert packed.is_float4()
+    assert unpacked.is_float4_e2m1_unpacked()
+    assert not unpacked.is_float4_e2m1fn()
+    assert unpacked.is_float4()
+    assert T.float4_e2m1fnx2.is_float4()
+    assert T.float4_e2m1_unpackedx4.is_float4()
+    assert is_float4_e2m1fn(packed)
+    assert is_float4_e2m1_unpacked(unpacked)
+    assert is_float4("float4_e2m1fn")
+    assert is_float4("custom[float4_e2m1_unpacked]8")
+    assert not is_float4(T.float16)
+
+
+if __name__ == "__main__":
+    test_float4_e2m1_unpacked_dtype_properties()
+    test_float4_e2m1_unpacked_vector_lanes()
+    test_float4_e2m1_unpacked_tir_cast()
+    test_float4_e2m1_unpacked_distinct_from_packed()
+    test_float4_e2m1_unpacked_dtype_helpers()
+    print("ok")

tilelang/language/copy_op.py

@@ -402,8 +402,12 @@ def tma_gather4_bytes(K_box, dtype: str) -> int:
     """
     if dtype not in _TMA_SUPPORTED_DTYPES:
         raise ValueError(f"Unsupported dtype: {dtype}")
-    elem_bytes = tvm.DataType(dtype).bits // 8
-    return 4 * K_box * elem_bytes
+    dt = tvm.DataType(dtype)
+    if dt.is_float4_e2m1_unpacked():
+        elem_bits = 4
+    else:
+        elem_bits = dt.bits
+    return (4 * K_box * elem_bits + 7) // 8
 
 
 def tma_scatter4(

tilelang/language/dtypes.py

@@ -17,6 +17,9 @@ def bits(self) -> int: ...
         @property
         def bytes(self) -> int: ...
         def as_torch(self) -> torch.dtype: ...
+        def is_float4_e2m1fn(self) -> bool: ...
+        def is_float4_e2m1_unpacked(self) -> bool: ...
+        def is_float4(self) -> bool: ...
 else:
     dtype = tvm.DataType
 
@@ -163,6 +166,10 @@ def __dtype_call__(self: dtype, *args, is_size_var: bool = False) -> tirx.Var:
         val = "BFloat" + self[6:]
     elif self.startswith("tfloat"):
         val = "TensorFloat" + self[6:]
+    elif self.is_float4_e2m1_unpacked():
+        val = "Float4E2M1Unpacked"
+        if self.lanes > 1:
+            val += f"x{self.lanes}"
     else:
         raise TypeError(f"Invalid type {self}")
     if "_" in val:
@@ -246,10 +253,61 @@ def __dtype_bytes__(self: dtype) -> int:
     return self.itemsize
 
 
+_FLOAT4_E2M1FN_BITS = 4
+_FLOAT4_E2M1FN_TYPE_CODE = 17
+_FLOAT4_E2M1_UNPACKED_BITS = 8
+_FLOAT4_E2M1_UNPACKED_TYPE_CODE = 131
+
+
+def __dtype_is_float4_e2m1fn__(self: dtype) -> bool:
+    """Packed 4-bit FP4 E2M1 (CUTLASS float_e2m1_t).
+
+    Use for pure FP4 workloads on tcgen05 mxf4 / mxf4nvf4 (packed SMEM, half
+    the footprint). Global tensors and packed SMEM roundtrips use this dtype.
+    TMA: ``CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B``.
+    """
+    return self.bits == _FLOAT4_E2M1FN_BITS and int(self.type_code) == _FLOAT4_E2M1FN_TYPE_CODE
+
+
+def __dtype_is_float4_e2m1_unpacked__(self: dtype) -> bool:
+    """8-bit FP4 E2M1 unpacked storage (CUTLASS float_e2m1_unpacksmem_t).
+
+    Only for tcgen05 ``kind::f8f6f4`` and block-scaled ``kind::mxf8f6f4``
+    mixed-precision paths where sub-byte types share 16-byte SMEM/TMEM padding
+    (one byte per FP4 element). Not for mxf4 / mxf4nvf4 packed FP4 kernels.
+    Pair with ``T.tma_copy`` from packed global FP4; TMA:
+    ``CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B``.
+    """
+    return self.bits == _FLOAT4_E2M1_UNPACKED_BITS and int(self.type_code) == _FLOAT4_E2M1_UNPACKED_TYPE_CODE
+
+
+def __dtype_is_float4__(self: dtype) -> bool:
+    """Whether this is any FP4 E2M1 logical variant (packed or unpacked storage)."""
+    return __dtype_is_float4_e2m1fn__(self) or __dtype_is_float4_e2m1_unpacked__(self)
+
+
+def is_float4_e2m1fn(value: AnyDType) -> bool:
+    """Whether *value* is the packed 4-bit FP4 E2M1 variant."""
+    return get_tvm_dtype(value).is_float4_e2m1fn()
+
+
+def is_float4_e2m1_unpacked(value: AnyDType) -> bool:
+    """Whether *value* is the 8-bit FP4 E2M1 unpacked shared-memory storage variant."""
+    return get_tvm_dtype(value).is_float4_e2m1_unpacked()
+
+
+def is_float4(value: AnyDType) -> bool:
+    """Whether *value* is any FP4 E2M1 logical variant (packed or unpacked)."""
+    return get_tvm_dtype(value).is_float4()
+
+
 dtype.__call__ = __dtype_call__
 dtype.__new__ = __dtype_new__
 dtype.as_torch = __dtype_as_torch__
 dtype.bytes = property(__dtype_bytes__)
+dtype.is_float4_e2m1fn = __dtype_is_float4_e2m1fn__
+dtype.is_float4_e2m1_unpacked = __dtype_is_float4_e2m1_unpacked__
+dtype.is_float4 = __dtype_is_float4__
 
 
 def get_tvm_dtype(value: AnyDType) -> dtype:
@@ -423,6 +481,12 @@ class float4_e2m1fnx8(dtype): ...
     class float4_e2m1fnx16(dtype): ...
     class float4_e2m1fnx32(dtype): ...
     class float4_e2m1fnx64(dtype): ...
+    class float4_e2m1_unpacked(dtype): ...
+    class float4_e2m1_unpackedx2(dtype): ...
+    class float4_e2m1_unpackedx4(dtype): ...
+    class float4_e2m1_unpackedx8(dtype): ...
+    class float4_e2m1_unpackedx16(dtype): ...
+    class float4_e2m1_unpackedx32(dtype): ...
     class bfloat16(dtype): ...
     class bfloat16x2(dtype): ...
     class tfloat32(dtype): ...
@@ -599,6 +663,12 @@ class tfloat32x64(dtype): ...
     float4_e2m1fnx16 = dtype("float4_e2m1fnx16")
     float4_e2m1fnx32 = dtype("float4_e2m1fnx32")
     float4_e2m1fnx64 = dtype("float4_e2m1fnx64")
+    float4_e2m1_unpacked = dtype("custom[float4_e2m1_unpacked]8")
+    float4_e2m1_unpackedx2 = dtype("custom[float4_e2m1_unpacked]8x2")
+    float4_e2m1_unpackedx4 = dtype("custom[float4_e2m1_unpacked]8x4")
+    float4_e2m1_unpackedx8 = dtype("custom[float4_e2m1_unpacked]8x8")
+    float4_e2m1_unpackedx16 = dtype("custom[float4_e2m1_unpacked]8x16")
+    float4_e2m1_unpackedx32 = dtype("custom[float4_e2m1_unpacked]8x32")
     bfloat16 = dtype("bfloat16")
     bfloat16x2 = dtype("bfloat16x2")
     tfloat32 = dtype("custom[tfloat32]")
@@ -773,6 +843,12 @@ class tfloat32x64(dtype): ...
     "float4_e2m1fnx16",
     "float4_e2m1fnx32",
     "float4_e2m1fnx64",
+    "float4_e2m1_unpacked",
+    "float4_e2m1_unpackedx2",
+    "float4_e2m1_unpackedx4",
+    "float4_e2m1_unpackedx8",
+    "float4_e2m1_unpackedx16",
+    "float4_e2m1_unpackedx32",
     "bfloat16",
     "bfloat16x2",
     "tfloat32",
@@ -788,4 +864,7 @@ class tfloat32x64(dtype): ...
     "dtype",
     "AnyDType",
     "get_tvm_dtype",
+    "is_float4_e2m1fn",
+    "is_float4_e2m1_unpacked",
+    "is_float4",
 ]