fixing an error from rebasing

NVIDIA · Dec 24, 2024 · a3c68a0 · a3c68a0
1 parent 494aea7
commit a3c68a0
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/csrc/scheduler/hopper_multi_matmul.cpp b/csrc/scheduler/hopper_multi_matmul.cpp
@@ -523,13 +523,10 @@ void HopperMultipleMatmulScheduler::scheduleEpilogue() {
       dc->setMemoryType(MemoryType::Local);
       d_smem->setMemoryType(MemoryType::Shared);
 
-      // Set LoadStoreOp
-      // TODO: extend support when mma is not cast to half
-      NVF_CHECK(
-          dataTypeSize(dc->dtype()) == 2,
-          "We support use_smem_epilogue on Hopper only when the output is 16-bit");
+      auto store_with_stmatrix = dataTypeSize(dc->dtype()) == 2;
 
       if (store_with_stmatrix) {
+        // Set LoadStoreOp
         d_smem->definition()->as<LoadStoreOp>()->setOpType(
             LoadStoreOpType::StMatrix);
       }