Matmul default scheduling [1]

[Priya2698]

Initial PR for Issue #1669.

1. Adds a EnableOption::MatmulExprEval to turn on expression evaluation for matmul while the API in under progress.
2. Currently, we only evaluate the MmaOp. The next PRs will amend this to look ahead and evaluate Mma + Cast, which is what we should see in the fusion definitions. See discussion here. In the absence of this we may have casts such as bfloat->float->bfloat.

[naoyam]

Does this mean any matmul pattern including mma + epilogue will be handled by the expression evaluator? Shouldn't it only take care of the mma part?

[wujingyue]

There are two problems with taking care of just the mma:

1. at::matmul doesn't do HH->S. We could plug in another backend that supports HH->S for EE though.
2. MMA is never alone in these GPT models (e.g. LLaMA). It's always part of a linear layer or an SDPA. nvFuser doesn't do SDPA well and we will have to offload it to another executor for quite some time, so scratch that. A linear layer however comes with this MMA->BiasAdd pattern. In order for its performance to be on par with framework-not-giving-nvFuser-the-matmul, we have to execute MMA+epilogue in one call.

Wdyt, @naoyam? Extending EE to do MMA+epilogue is the most obvious way to me to solve the above problems. But I could definitely be wrong.

[naoyam]

Extending EE to do MMA+epilogue

How would it do that? Does Aten support matmul with some epilogue op?

[wujingyue]

Oh yeah, it's called at::addmm. I realized I was wrong about Relu. torch.nn.Linear doesn't do Relu, so the pattern would be matmul+biasadd which is what at::addmm does.

[naoyam]

So, how would you handle matmul+epilogue patterns that are accepted by the nvFuser matmul scheduler but there's no corresponding aten version? Would that end up doing mma and the epilogue op separately?

[naoyam]

That makes sense. Or, maybe we might want to have a separate scheduler for EE? We try the native matmul scheduler first, and then the EE matmul scheduler, and then the other schedulers. No particular preference but just my two cents.

[wujingyue]

Yes, that's a certainly great option to consider. It makes things more composable at a risk of being harder to share logics with MatmulScheduler. I hope the preference will be more obvious when we know what the heuristics look like!

[Priya2698]

Adding on to @wujingyue's comments, the plan for the next PRs is:

1. Support Mma + Cast -> avoid roundtrip casting (half-> float -> half) by checking mma_out->uses() and if it is either a castOp or pointwise ops with inputs of the same type (half), then skip casting back the output of at::matmul. This will not execute matmul + bias in a single call.
2. Handle common epilogue fusions: We will need to pattern match and evaluate within the MmaOp accordingly. test_matmul_scheduler.cpp currently has a few cases that I will start with: mma + bias, mma + bias + relu/gelu, mma + relu.
3. Epilogue fusions that are not be supported: They can still be computed through EE but should ideally not be plumbed down to the matmul scheduler.

[wujingyue]

Do 2 only when it's needed. You should double check this but I think Llama for example uses linear with bias off and none of the linear layers in our benchmarks do relu or gelu

[Priya2698]

Llama2 has bias=False in the linear layers but some GPT configs have bias=True.

[wujingyue]

Nice work! The PR currently shows as a draft. I'll hold off my reviews when it's ready.

[wujingyue]

LGTM with some comments to be resolved.

[jacobhinkle]

LGTM. Just some comments on the tests.

[drzejan2]

Thanks for preparing this change, they look good to me after caught up to the latest discussions around matmul scheduler.

[protonu]

Nit: Will it make sense to make this C++ test file a part of the test_matmul target?
https://github.com/NVIDIA/Fuser/blob/9fac12bdd98a63fdc88dde265b8add6a0e3f41cf/CMakeLists.txt#L480C1-L489C41

[Priya2698]

Nit: Will it make sense to make this C++ test file a part of the test_matmul target? https://github.com/NVIDIA/Fuser/blob/9fac12bdd98a63fdc88dde265b8add6a0e3f41cf/CMakeLists.txt#L480C1-L489C41

Thanks for the suggestion! Moved the test file.

[jjsjann123]

FYI, this pr seems to break CI on V100.

[Priya2698]

Thanks for pointing this out. Let's revert this until I identify the patch.

[jjsjann123]

FYI, these are the failing tests from CI:

00:00:43 [  FAILED  ] 3 tests, listed below:
00:00:43 [  FAILED  ] MatmulATenEvaluationTest.SingleMmaOp
00:00:43 [  FAILED  ] MatmulATenEvaluationTest.MmaOpAndCast
00:00:43 [  FAILED  ] MatmulATenEvaluationTest.MatmulWithBias

[Priya2698]

Yes, these are the tests I added to check functionality.
Looking into it.

[wujingyue]

Thanks for pointing this out. Let's revert this until I identify the patch.

Thank you! I can't say enough great things about revert-and-debug-later!

[Priya2698]

Are we looking to support V100 through the default path? CC: @kevinstephano
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD can be used in the tests to fix the error. The error surfaced because matmulScheduler does not support V100 (See:

Fuser/csrc/scheduler/matmul_utils.cpp

Line 44 in 77caa57

inline std::optional<MmaMacro> getMmaOp(

).

If we wish to support V100, we can have appropriate checks in the heuristic verification to allow other architectures when we are using expression evaluator.

[wujingyue]

Are we looking to support V100 through the default path?

I'd do it. This is actually something that can be supported with much less effort in fallback mode than in codegen. Sounds like a low-hanging fruit to me.

[wujingyue]

Are we looking to support V100 through the default path? CC: @kevinstephano NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD can be used in the tests to fix the error. The error surfaced because matmulScheduler does not support V100 (See:

Fuser/csrc/scheduler/matmul_utils.cpp

Line 44 in 77caa57

inline std::optional<MmaMacro> getMmaOp(

).
If we wish to support V100, we can have appropriate checks in the heuristic verification to allow other architectures when we are using expression evaluator.

FYI, I suspect it's not just V100. https://github.com/NVIDIA/Fuser/actions/runs/7910657123/job/21593586846 seems to be the same error but for H100.

[protonu]

On option could be use to use something like this (https://github.com/NVIDIA/Fuser/blob/88727dc828684f5a62d7f1837a610b7589f629d1/test/test_combine_mul_sum.cpp#L40C1-L57C3) to reduce the machines these tests run on, in case you guys don't plan on adding support for V100/H100.

[Priya2698]

On option could be use to use something like this (https://github.com/NVIDIA/Fuser/blob/88727dc828684f5a62d7f1837a610b7589f629d1/test/test_combine_mul_sum.cpp#L40C1-L57C3) to reduce the machines these tests run on, in case you guys don't plan on adding support for V100/H100.

Thanks for the suggestion, I am moving forward with supporting any architecture since it's simple enough.
For this PR, the fix will be to skip the heuristic computation when the MatmulExprEval is set. In the next PRs, I'll refactor some code to always use the default scheduling for any unsupported architecture.