[CUDA] Add NVFP4 support in QMoE

docs/ContribOperators.md

@@ -4935,7 +4935,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>activation_type</tt> : string</dt>
 <dd>Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu</dd>
 <dt><tt>block_size</tt> : int</dt>
-<dd>Size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). If provided, both hidden_size and inter_size must be divisible by the block size. Otherwise, there is no blocking and a whole column shares one scaling factor. </dd>
+<dd>Size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). Both hidden_size and inter_size must be divisible by the block size. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor. </dd>
 <dt><tt>expert_weight_bits</tt> : int</dt>
 <dd>Number of bits used in quantized weights. Supported values are 2, 4, and 8. Default is 4 bits</dd>
 <dt><tt>k</tt> : int</dt>

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp

@@ -379,7 +379,7 @@ struct MixedGroupedGemmInputUtils {
           constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
           asm volatile(
               "{\n"
-              "  lop3 .b32 %0, %2, %4, %5, %6;\n"
+              "  lop3.b32 %0, %2, %4, %5, %6;\n"
               "  xor  .b32 %1, %3, %5;        \n"
               "}\n"
               : "=r"(scale_pos_[0]), "=r"(scale_pos_[1])

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h

@@ -107,6 +107,23 @@ struct LayoutDetailsB<TypeA, uint4b_t, Arch, typename platform::enable_if<Arch::
   using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
 };
 
+// MXFP4 (e2m1) weights with FP16/BF16 activations on Turing+. Mirrors the uint4b_t 4-bit
+// interleaved layout: weights are stored in the SM80 column-major tile-interleaved layout and
+// dequantized after the smem load via OpMultiplyAddDequantizeInterleavedBToA.
+template <typename TypeA, typename Arch>
+struct LayoutDetailsB<TypeA, cutlass::float_e2m1_t, Arch, typename platform::enable_if<Arch::kMinComputeCapability >= 75>::type> {
+  static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits<TypeA>::value;
+
+ private:
+  static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits<cutlass::float_e2m1_t>::value;
+  static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK;
+
+ public:
+  using Layout = layout::ColumnMajorTileInterleave<ThreadblockK, ColumnsInterleaved>;
+  static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits<cutlass::float_e2m1_t>::value;
+  using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA;
+};
+
 }  // namespace kernel
 }  // namespace gemm
 }  // namespace cutlass

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h

@@ -180,7 +180,8 @@ struct MoeFCGemm {
               int64_t const* total_tokens_including_expert, int64_t gemm_n, int64_t gemm_k,
               GemmCoord* host_problem_sizes = nullptr)
         : problem_count(problem_count), threadblock_count(threadblock_count), group_size(group_size), output_op(output_op), ptr_A(const_cast<ElementA*>(ptr_A)), ptr_B(const_cast<ElementB*>(ptr_B)), weight_scales(const_cast<ElementScale*>(weight_scales)), weight_zeros(const_cast<ElementScale*>(weight_zeros)), ptr_C(const_cast<ElementC*>(ptr_C)), C_is_broadcast{C_is_broadcast}, ptr_D(ptr_D), total_tokens_including_expert(total_tokens_including_expert), gemm_n(gemm_n), gemm_k(gemm_k), host_problem_sizes(nullptr) {
-      if (platform::is_same<uint8_t, ElementB>::value || platform::is_same<uint4b_t, ElementB>::value) {
+      if (platform::is_same<uint8_t, ElementB>::value || platform::is_same<uint4b_t, ElementB>::value ||
+          platform::is_same<cutlass::float_e2m1_t, ElementB>::value) {
         assert(weight_scales);
       }
       this->gather_A_indices = nullptr;
@@ -283,7 +284,8 @@ struct MoeFCGemm {
   }
 
   static Status can_implement(Arguments const& args) {
-    if constexpr (platform::is_same<uint8_t, ElementB>::value || platform::is_same<uint4b_t, ElementB>::value) {
+    if constexpr (platform::is_same<uint8_t, ElementB>::value || platform::is_same<uint4b_t, ElementB>::value ||
+                  platform::is_same<cutlass::float_e2m1_t, ElementB>::value) {
       if (args.weight_scales == nullptr) {
         CUTLASS_TRACE_HOST("MoeFCGemm::can_implement() - weight scales are required for uint8_t and uint4b_t");
         return Status::kInvalid;

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h

@@ -118,8 +118,9 @@ struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, Ele
   static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
                 "Mma multistage must dequantize after ldsm");
 
-  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-                "Element B must be uint8 or uint4");
+  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value ||
+                    platform::is_same<ElementB, cutlass::float_e2m1_t>::value,
+                "Element B must be uint8, uint4 or float_e2m1");
 
   static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
                                                                   ? cutlass::arch::CacheOperation::Global
@@ -218,8 +219,9 @@ struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, Ele
   static_assert(platform::is_same<Operator, arch::OpMultiplyAddDequantizeInterleavedBToA>::value,
                 "Mma multistage must dequantize after ldsm");
 
-  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-                "Element B must be uint8 or uint4");
+  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value ||
+                    platform::is_same<ElementB, cutlass::float_e2m1_t>::value,
+                "Element B must be uint8, uint4 or float_e2m1");
 
   static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
                                                                   ? cutlass::arch::CacheOperation::Global

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h

@@ -116,8 +116,9 @@ struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, Ele
   static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value,
                 "Element A must be fp16 or bf16");
 
-  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-                "Element B must be uint8 or uint4");
+  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value ||
+                    platform::is_same<ElementB, cutlass::float_e2m1_t>::value,
+                "Element B must be uint8, uint4 or float_e2m1");
 
   using OperatorInfo = arch::DetagOperator<Operator_>;
   using Operator = typename OperatorInfo::Operator;
@@ -201,8 +202,9 @@ struct DqMma<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, Ele
   static_assert(platform::is_same<ElementA, half_t>::value || platform::is_same<ElementA, bfloat16_t>::value,
                 "Element A must be fp16 or bf16");
 
-  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value,
-                "Element B must be uint8 or uint4");
+  static_assert(platform::is_same<ElementB, uint8_t>::value || platform::is_same<ElementB, uint4b_t>::value ||
+                    platform::is_same<ElementB, cutlass::float_e2m1_t>::value,
+                "Element B must be uint8, uint4 or float_e2m1");
 
   using OperatorInfo = arch::DetagOperator<Operator_>;
   using Operator = typename OperatorInfo::Operator;

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h

@@ -221,6 +221,91 @@ struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, uint4b_t, LayoutB, kAli
   using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
+////////////////////////////////////////////////////////////////////////////////
+/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & MXFP4 (e2m1) weight,
+/// mma pipelined (stage=2)
+template <
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, ElementAccumulator,
+                  layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+
+  using Mma = DqMma<half_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
+                    kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
+                    WarpShape, InstructionShape, 2, Operator>;
+
+ public:
+  using MmaCore = typename Mma::MmaCore;
+  using IteratorA = typename Mma::IteratorA;
+  using IteratorB = typename Mma::IteratorB;
+  using ThreadblockMma = typename Mma::ThreadblockMma;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & MXFP4 (e2m1) weight,
+/// mma multistage (stage>=3)
+template <
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    int kStages,
+    /// Shared memory clear option
+    SharedMemoryClearOption SharedMemoryClear>
+struct DefaultMma<cutlass::half_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, ElementAccumulator,
+                  layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
+                  false, SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<half_t>::value;
+
+  using Mma = DqMma<half_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, half_t, layout::RowMajor,
+                    kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
+                    WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+
+ public:
+  using MmaCore = typename Mma::MmaCore;
+  using IteratorA = typename Mma::IteratorA;
+  using IteratorB = typename Mma::IteratorB;
+  using ThreadblockMma = typename Mma::ThreadblockMma;
+};
+
 #ifdef ENABLE_FP8
 ////////////////////////////////////////////////////////////////////////////////
 /// Specialization for row-major output (OperatorClass TensorOp), fp8 activation & int4 weight, mma multistage

onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h

@@ -331,6 +331,52 @@ struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, uint4b_t, LayoutB,
   using ThreadblockMma = typename Mma::ThreadblockMma;
 };
 
+////////////////////////////////////////////////////////////////////////////////
+/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & MXFP4 (e2m1) weight,
+/// mma pipelined (stage=2)
+template <
+    typename LayoutA, int kAlignmentA, typename LayoutB, int kAlignmentB, typename ElementAccumulator,
+    typename ArchTag, typename ThreadblockShape, typename WarpShape, typename InstructionShape, typename Operator>
+struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, ElementAccumulator,
+                  layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, Operator> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+
+  using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
+                    kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
+                    WarpShape, InstructionShape, 2, Operator>;
+
+ public:
+  using MmaCore = typename Mma::MmaCore;
+  using IteratorA = typename Mma::IteratorA;
+  using IteratorB = typename Mma::IteratorB;
+  using ThreadblockMma = typename Mma::ThreadblockMma;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & MXFP4 (e2m1) weight,
+/// mma multistage (stage>=3)
+template <
+    typename LayoutA, int kAlignmentA, typename LayoutB, int kAlignmentB, typename ElementAccumulator,
+    typename ArchTag, typename ThreadblockShape, typename WarpShape, typename InstructionShape, typename Operator,
+    int kStages, SharedMemoryClearOption SharedMemoryClear>
+struct DefaultMma<cutlass::bfloat16_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, ElementAccumulator,
+                  layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape, WarpShape, InstructionShape, kStages, Operator,
+                  false, SharedMemoryClear> {
+ private:
+  static constexpr int kAlignmentScale = 128 / sizeof_bits<bfloat16_t>::value;
+
+  using Mma = DqMma<bfloat16_t, LayoutA, kAlignmentA, cutlass::float_e2m1_t, LayoutB, kAlignmentB, bfloat16_t, layout::RowMajor,
+                    kAlignmentScale, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, ArchTag, ThreadblockShape,
+                    WarpShape, InstructionShape, kStages, Operator, SharedMemoryClear>;
+
+ public:
+  using MmaCore = typename Mma::MmaCore;
+  using IteratorA = typename Mma::IteratorA;
+  using IteratorB = typename Mma::IteratorB;
+  using ThreadblockMma = typename Mma::ThreadblockMma;
+};
+
 }  // namespace threadblock
 }  // namespace gemm
 }  // namespace cutlass