wip

Author: liqiangxl

tests/cpp/test_cluster.cpp

@@ -103,7 +103,7 @@ INSTANTIATE_TEST_SUITE_P(
     });
 
 using ClusterReductionTestAutoScheduler = ClusterReductionTest;
-TEST_P(ClusterReductionTestAutoScheduler, SoftmaxDtypeHiddenSize) {
+TEST_P(ClusterReductionTestAutoScheduler, Softmax) {
   auto [hidden_size, dtype] = GetParam();
   int batch_size = scheduler_utils::safeDiv(deviceSMCount(), 8);
   auto fusion_ptr = std::make_unique<Fusion>();
@@ -112,9 +112,9 @@ TEST_P(ClusterReductionTestAutoScheduler, SoftmaxDtypeHiddenSize) {
   auto tv0 = makeContigTensor(2, dtype);
   fusion.addInput(tv0);
   auto tv1 = maybeCastOp(DataType::Float, tv0);
-    auto tv2 = softmax(tv1, 1);
-    auto tv3 = maybeCastOp(DataType::BFloat16, tv2);
-    fusion.addOutput(tv3);
+  auto tv2 = softmax(tv1, 1);
+  auto tv3 = maybeCastOp(DataType::BFloat16, tv2);
+  fusion.addOutput(tv3);
   auto unscheduled_fusion_copy = fusion;
 
   auto options =
@@ -129,7 +129,6 @@ TEST_P(ClusterReductionTestAutoScheduler, SoftmaxDtypeHiddenSize) {
   }
   testValidate(&unscheduled_fusion_copy, outputs, {t0});
 }
-
 INSTANTIATE_TEST_SUITE_P(
     ,
     ClusterReductionTestAutoScheduler,

tests/cpp/utils.h

@@ -745,6 +745,9 @@ bool isSchedulerInUse(
     const nvfuser::FusionKernelRuntime* kernel_rt,
     const SchedulerType& scheduler_type);
 
+// Disable magic zero
+const CompileParams matmul_cparams{DataType::Int32, 255, false};
+
 // Utility to generate tensor with bias applied on the input tensor
 TensorView* biasEpilogue(TensorView* tensor, TensorView* bias);