Style-align SM120 FP4 MMA support

Author: TerminusAkivili

src/tl_templates/cuda/cuda_fp4.h

@@ -5,6 +5,7 @@
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) ||                      \
     (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ >= 1200))
 #include <cuda_fp4.h>
+#include <cute/numeric/numeric_types.hpp>
 
 // Wrapper for __nv_fp4_e2m1 with implicit conversions
 struct fp4_e2_t {
@@ -45,6 +46,12 @@ struct fp4_e2_t {
   TL_DEVICE operator __half() const { return __half(float(*this)); }
 };
 
+namespace tl {
+template <> struct to_cute_type<::fp4_e2_t> {
+  using type = cute::float_e2m1_t;
+};
+} // namespace tl
+
 class fp4_e2_2_t {
 public:
   __nv_fp4x2_storage_t __x;
@@ -163,26 +170,27 @@ TL_DEVICE fp4_e2_32_t make_fp4_e2_32_t(
 }
 
 // Pack sixty-four fp4_e2_t values.
-template <typename... Args>
-TL_DEVICE fp4_e2_64_t make_fp4_e2_64_t(Args... args) {
-  static_assert(sizeof...(Args) == 64,
-                "make_fp4_e2_64_t expects 64 fp4 values");
-  fp4_e2_t values[64] = {fp4_e2_t(args)...};
+TL_DEVICE fp4_e2_64_t make_fp4_e2_64_t(
+    fp4_e2_t x0, fp4_e2_t x1, fp4_e2_t x2, fp4_e2_t x3, fp4_e2_t x4,
+    fp4_e2_t x5, fp4_e2_t x6, fp4_e2_t x7, fp4_e2_t x8, fp4_e2_t x9,
+    fp4_e2_t x10, fp4_e2_t x11, fp4_e2_t x12, fp4_e2_t x13, fp4_e2_t x14,
+    fp4_e2_t x15, fp4_e2_t x16, fp4_e2_t x17, fp4_e2_t x18, fp4_e2_t x19,
+    fp4_e2_t x20, fp4_e2_t x21, fp4_e2_t x22, fp4_e2_t x23, fp4_e2_t x24,
+    fp4_e2_t x25, fp4_e2_t x26, fp4_e2_t x27, fp4_e2_t x28, fp4_e2_t x29,
+    fp4_e2_t x30, fp4_e2_t x31, fp4_e2_t y0, fp4_e2_t y1, fp4_e2_t y2,
+    fp4_e2_t y3, fp4_e2_t y4, fp4_e2_t y5, fp4_e2_t y6, fp4_e2_t y7,
+    fp4_e2_t y8, fp4_e2_t y9, fp4_e2_t y10, fp4_e2_t y11, fp4_e2_t y12,
+    fp4_e2_t y13, fp4_e2_t y14, fp4_e2_t y15, fp4_e2_t y16, fp4_e2_t y17,
+    fp4_e2_t y18, fp4_e2_t y19, fp4_e2_t y20, fp4_e2_t y21, fp4_e2_t y22,
+    fp4_e2_t y23, fp4_e2_t y24, fp4_e2_t y25, fp4_e2_t y26, fp4_e2_t y27,
+    fp4_e2_t y28, fp4_e2_t y29, fp4_e2_t y30, fp4_e2_t y31) {
   fp4_e2_64_t result;
-  result.x = make_fp4_e2_32_t(
-      values[0], values[1], values[2], values[3], values[4], values[5],
-      values[6], values[7], values[8], values[9], values[10], values[11],
-      values[12], values[13], values[14], values[15], values[16], values[17],
-      values[18], values[19], values[20], values[21], values[22], values[23],
-      values[24], values[25], values[26], values[27], values[28], values[29],
-      values[30], values[31]);
-  result.y = make_fp4_e2_32_t(
-      values[32], values[33], values[34], values[35], values[36], values[37],
-      values[38], values[39], values[40], values[41], values[42], values[43],
-      values[44], values[45], values[46], values[47], values[48], values[49],
-      values[50], values[51], values[52], values[53], values[54], values[55],
-      values[56], values[57], values[58], values[59], values[60], values[61],
-      values[62], values[63]);
+  result.x = make_fp4_e2_32_t(x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11,
+                              x12, x13, x14, x15, x16, x17, x18, x19, x20, x21,
+                              x22, x23, x24, x25, x26, x27, x28, x29, x30, x31);
+  result.y = make_fp4_e2_32_t(y0, y1, y2, y3, y4, y5, y6, y7, y8, y9, y10, y11,
+                              y12, y13, y14, y15, y16, y17, y18, y19, y20, y21,
+                              y22, y23, y24, y25, y26, y27, y28, y29, y30, y31);
   return result;
 }
 

src/tl_templates/cuda/gemm_mma.h

@@ -46,11 +46,6 @@ using _X = Underscore;
 #include "cuda_fp8.h"
 #include <cute/arch/mma_sm120.hpp>
 #include <cute/arch/mma_sm80.hpp>
-namespace tl {
-template <> struct to_cute_type<fp4_e2_t> {
-  using type = cute::float_e2m1_t;
-};
-} // namespace tl
 TL_DISPATCH_MMA_TEMPLATE(fp4_e2_t, fp4_e2_t, float, SM120_16x8x32_TN)
 TL_DISPATCH_MMA_TEMPLATE(fp8_e4_t, fp4_e2_t, float, SM120_16x8x32_TN)
 TL_DISPATCH_MMA_TEMPLATE(fp4_e2_t, fp8_e4_t, float, SM120_16x8x32_TN)

src/tl_templates/cuda/instruction/mma.h

@@ -69,9 +69,9 @@ struct MmaDispatcher {
   }
 };
 
-#define TL_DEFINE_MMA_DISPATCHER(ATypeEnum, BTypeEnum, CTypeEnum, MValue,      \
-                                 NValue, KValue, TransAValue, TransBValue,     \
-                                 SaturateValue, ImplType)                      \
+#define TL_DEFINE_MMA_DISPATCHER_IMPL(                                         \
+    ATypeEnum, BTypeEnum, CTypeEnum, MValue, NValue, KValue, TransAValue,      \
+    TransBValue, SaturateValue, ShiftAValue, ShiftBValue, ImplType)            \
   template <>                                                                  \
   struct MmaDispatcher<DataType::ATypeEnum, DataType::BTypeEnum,               \
                        DataType::CTypeEnum, MValue, NValue, KValue,            \
@@ -84,12 +84,46 @@ struct MmaDispatcher {
     static_assert(                                                             \
         std::is_same_v<typename Traits::DReg, typename Traits::CReg>,          \
         "tl::mma_sync requires matching accumulator/output regs");             \
+    template <bool Shift, class Reg>                                           \
+    static TL_DEVICE Reg maybe_shift_fp4_reg(Reg reg) {                        \
+      if constexpr (Shift) {                                                   \
+        return reg << 2;                                                       \
+      } else {                                                                 \
+        return reg;                                                            \
+      }                                                                        \
+    }                                                                          \
     static TL_DEVICE void exec(CRegType *d, const ARegType *a,                 \
                                const BRegType *b, const CRegType *c) {         \
-      call_fma<Impl>(d, a, b, c);                                              \
+      if constexpr (ShiftAValue || ShiftBValue) {                              \
+        ARegType as[Traits::kARegs];                                           \
+        BRegType bs[Traits::kBRegs];                                           \
+        _Pragma("unroll") for (int i = 0; i < Traits::kARegs; ++i) {           \
+          as[i] = maybe_shift_fp4_reg<ShiftAValue>(a[i]);                      \
+        }                                                                      \
+        _Pragma("unroll") for (int i = 0; i < Traits::kBRegs; ++i) {           \
+          bs[i] = maybe_shift_fp4_reg<ShiftBValue>(b[i]);                      \
+        }                                                                      \
+        call_fma<Impl>(d, as, bs, c);                                          \
+      } else {                                                                 \
+        call_fma<Impl>(d, a, b, c);                                            \
+      }                                                                        \
     }                                                                          \
   };
 
+#define TL_DEFINE_MMA_DISPATCHER(ATypeEnum, BTypeEnum, CTypeEnum, MValue,      \
+                                 NValue, KValue, TransAValue, TransBValue,     \
+                                 SaturateValue, ImplType)                      \
+  TL_DEFINE_MMA_DISPATCHER_IMPL(ATypeEnum, BTypeEnum, CTypeEnum, MValue,       \
+                                NValue, KValue, TransAValue, TransBValue,      \
+                                SaturateValue, false, false, ImplType)
+
+#define TL_DEFINE_MMA_DISPATCHER_WITH_FP4_SHIFT(                               \
+    ATypeEnum, BTypeEnum, CTypeEnum, MValue, NValue, KValue, TransAValue,      \
+    TransBValue, SaturateValue, ShiftAValue, ShiftBValue, ImplType)            \
+  TL_DEFINE_MMA_DISPATCHER_IMPL(                                               \
+      ATypeEnum, BTypeEnum, CTypeEnum, MValue, NValue, KValue, TransAValue,    \
+      TransBValue, SaturateValue, ShiftAValue, ShiftBValue, ImplType)
+
 // FP16 inputs (TN layout: A row-major, B column-major)
 TL_DEFINE_MMA_DISPATCHER(kFloat16, kFloat16, kFloat16, 16, 8, 16, false, true,
                          false, cute::SM80_16x8x16_F16F16F16F16_TN)
@@ -154,14 +188,19 @@ using SM120_FP8_FP4_F32_TN =
     cute::SM120_16x8x32_TN<cute::float_e4m3_t, cute::float_e2m1_t, float>;
 using SM120_FP4_FP8_F32_TN =
     cute::SM120_16x8x32_TN<cute::float_e2m1_t, cute::float_e4m3_t, float>;
-TL_DEFINE_MMA_DISPATCHER(kFloat4_e2m1fn, kFloat4_e2m1fn, kFloat32, 16, 8, 32,
-                         false, true, false, SM120_FP4_FP4_F32_TN)
-TL_DEFINE_MMA_DISPATCHER(kFloat8_e4m3, kFloat4_e2m1fn, kFloat32, 16, 8, 32,
-                         false, true, false, SM120_FP8_FP4_F32_TN)
-TL_DEFINE_MMA_DISPATCHER(kFloat4_e2m1fn, kFloat8_e4m3, kFloat32, 16, 8, 32,
-                         false, true, false, SM120_FP4_FP8_F32_TN)
-
+TL_DEFINE_MMA_DISPATCHER_WITH_FP4_SHIFT(kFloat4_e2m1fn, kFloat4_e2m1fn,
+                                        kFloat32, 16, 8, 32, false, true, false,
+                                        true, true, SM120_FP4_FP4_F32_TN)
+TL_DEFINE_MMA_DISPATCHER_WITH_FP4_SHIFT(kFloat8_e4m3, kFloat4_e2m1fn, kFloat32,
+                                        16, 8, 32, false, true, false, false,
+                                        true, SM120_FP8_FP4_F32_TN)
+TL_DEFINE_MMA_DISPATCHER_WITH_FP4_SHIFT(kFloat4_e2m1fn, kFloat8_e4m3, kFloat32,
+                                        16, 8, 32, false, true, false, true,
+                                        false, SM120_FP4_FP8_F32_TN)
+
+#undef TL_DEFINE_MMA_DISPATCHER_WITH_FP4_SHIFT
 #undef TL_DEFINE_MMA_DISPATCHER
+#undef TL_DEFINE_MMA_DISPATCHER_IMPL
 
 } // namespace detail
 
@@ -178,37 +217,7 @@ TL_DEVICE void mma_sync(
                                            TransB, Saturate>;
   static_assert(!std::is_void_v<typename Dispatcher::CRegType>,
                 "tl::mma_sync: unsupported configuration");
-  if constexpr (AType == DataType::kFloat4_e2m1fn ||
-                BType == DataType::kFloat4_e2m1fn) {
-    // SM120 f8f6f4 MMA expects FP4 operands in the same register placement as
-    // CuTe's b4x16 load path. Shift only FP4 operands; mixed FP8 operands keep
-    // their native register bits.
-    using AReg = typename Dispatcher::ARegType;
-    using BReg = typename Dispatcher::BRegType;
-    constexpr int nA = detail::MmaImplTraits<typename Dispatcher::Impl>::kARegs;
-    constexpr int nB = detail::MmaImplTraits<typename Dispatcher::Impl>::kBRegs;
-    AReg as[nA];
-    BReg bs[nB];
-#pragma unroll
-    for (int i = 0; i < nA; ++i) {
-      if constexpr (AType == DataType::kFloat4_e2m1fn) {
-        as[i] = a[i] << 2;
-      } else {
-        as[i] = a[i];
-      }
-    }
-#pragma unroll
-    for (int i = 0; i < nB; ++i) {
-      if constexpr (BType == DataType::kFloat4_e2m1fn) {
-        bs[i] = b[i] << 2;
-      } else {
-        bs[i] = b[i];
-      }
-    }
-    Dispatcher::exec(c, as, bs, c);
-  } else {
-    Dispatcher::exec(c, a, b, c);
-  }
+  Dispatcher::exec(c, a, b, c);
 }
 
 } // namespace tl

tilelang/cuda/intrinsics/layout/utils.py

@@ -28,9 +28,9 @@ def get_ldmatrix_offset(
     transposed: bool = False,
 ):
     assert matrix in ["A", "B"], "matrix should be either A or B"
-    dtype = DataType(dtype)
-    dtype_bits = dtype.bits
-    is_fp4_e2m1fn = dtype_bits == 4 and str(dtype) == "float4_e2m1fn"
+    dtype_obj = DataType(dtype)
+    dtype_bits = dtype_obj.bits
+    is_fp4_e2m1fn = dtype_bits == 4 and str(dtype_obj) == "float4_e2m1fn"
     if dtype_bits == 32:
         if matrix == "B" and transposed:
             transform_func = ldmatrix_32x4_to_shared_16x8_layout_b
@@ -78,7 +78,7 @@ def get_ldmatrix_offset(
         else:
             raise ValueError("ldmatrix only supports B transposed and A non-transposed for int8")
     else:
-        raise ValueError(f"Unsupported dtype {dtype}")
+        raise ValueError(f"Unsupported dtype {dtype_obj}")
 
 
 def shared_16x16_to_mma_32x8_layout(i, j):

tilelang/cuda/intrinsics/macro/mma_macro_generator.py

@@ -118,7 +118,8 @@ def __init__(
     def _initialize_k_dim(self, a_dtype=T.float16):
         if isinstance(a_dtype, str):
             a_dtype = DataType(a_dtype)
-        if str(a_dtype) == "float4_e2m1fn":
+        a_dtype_str = str(a_dtype)
+        if a_dtype_str == "float4_e2m1fn":
             if self.chunk < 32:
                 raise ValueError(f"float4_e2m1fn MMA requires chunk >= 32, got chunk={self.chunk}")
             self.k_dim = 32
@@ -300,7 +301,9 @@ def _warp_ld_a_fp64(
         micro_size_k = self.micro_size_k
         local_size_a = self.local_size_a
         a_transposed = self.a_transposed
-        a_dtype_bits = DataType(a_dtype).bits
+        a_dtype_obj = DataType(a_dtype)
+        a_dtype_bits = a_dtype_obj.bits
+        is_fp4_a = str(a_dtype_obj) == "float4_e2m1fn"
         # ldmatrix cannot be used for int8 + trans case.
         ldmatrix_available = not (a_dtype_bits != 16 and a_transposed)
 
@@ -344,8 +347,7 @@ def _warp_ldmatrix_a(
 
                 if ldmatrix_available:
                     num = 4
-                    is_fp4 = str(DataType(a_dtype)) == "float4_e2m1fn"
-                    access_extent = 4 * num if is_fp4 else 2 * num
+                    access_extent = 4 * num if is_fp4_a else 2 * num
                     row_off, col_off = get_ldmatrix_offset("A", tx, 0, stride, a_dtype, a_transposed)
                     src_indices = (
                         tuple(A_other) + (A_base0 + wk + row_off, A_base1 + wi + col_off)
@@ -416,7 +418,9 @@ def _warp_ld_b_fp64(
         micro_size_k = self.micro_size_k
         local_size_b = self.local_size_b
         b_transposed = self.b_transposed
-        b_dtype_bits = DataType(b_dtype).bits
+        b_dtype_obj = DataType(b_dtype)
+        b_dtype_bits = b_dtype_obj.bits
+        is_fp4_b = str(b_dtype_obj) == "float4_e2m1fn"
         thread_binding = self.get_thread_binding()
 
         # legalize shared buffer to region
@@ -464,8 +468,7 @@ def _warp_ldmatrix_b(
 
                 if ldmatrix_available:
                     num = 4 if replicate_b else 2
-                    is_fp4 = str(DataType(b_dtype)) == "float4_e2m1fn"
-                    access_extent = 4 * num if is_fp4 else 2 * num
+                    access_extent = 4 * num if is_fp4_b else 2 * num
                     row_off, col_off = get_ldmatrix_offset("B", tx, 0, stride, b_dtype, b_transposed)
                     src_indices = (
                         tuple(B_other) + (B_base0 + wi + row_off, B_base1 + wk + col_off)
@@ -873,10 +876,11 @@ def __init__(
         self._initialize_transform_kind(transform_kind_a, transform_kind_b)
 
     def _initialize_k_dim(self, a_dtype=T.float16):
-        if str(DataType(a_dtype)) == "float4_e2m1fn":
+        a_dtype_obj = DataType(a_dtype)
+        if str(a_dtype_obj) == "float4_e2m1fn":
             self.k_dim = 32
         else:
-            self.k_dim = 256 // DataType(a_dtype).bits
+            self.k_dim = 256 // a_dtype_obj.bits
 
     def _initialize_local_size(self, m_dim=16, n_dim=16, k_dim=16, warp_size=32):
         self.local_size_a = (m_dim * k_dim) // warp_size