Implement persistent matmul scheduling

[jacobhinkle]

Stacked on #3642

This is a followup to #3792 that implements persistent scheduling.

There is a current limitation that affects both persistent scheduling and "grid swizzling": if MatmulOp or LinearOp are present in the fusion, we will hit inlining errors. This is because in that case we have a non-trivial AxisMapping on the MmaOp. The missing input dimensions are not tracked through the scheduling transforms (merges and splits) required for either grid swizzling or persistent scheduling. Because of this, I introduced three new parametrized tests matching the original MLPBenchmarkTests but with _BroadcastInputs suffix. These tests use fusedMultiplySum instead of linear. The persistent variant of the non BroadcastInputs tests are skipped until we fix the inlining issue.

I currently observe a correctness issue in the MLPBenchmarkTest.FwdEpilogueFusion_BroadcastInputs test regardless of parametrization. This means that we are getting incorrect results even for data parallel scheduling. I confirmed this test also fails on main. I currently skip this test with a warning mesage.

[jacobhinkle]

!test

[github-actions]

Review updated until commit 0cbb3e6

Description

• Implement persistent scheduling for Hopper matmul kernels

• Add support for warp specialization on Hopper by default

• Parametrize MLP Benchmark tests to include persistent and warp specialization configurations

• Add new tests for broadcast inputs in MLP benchmarks

---

Changes walkthrough 📝

Relevant files

Enhancement

hopper_multi_matmul.cpp Add persistent scheduling support                                                csrc/scheduler/hopper_multi_matmul.cpp Include matmul_heuristic.h Remove temporary check for persistent scheduling Add persistent kernel scheduling logic Update block parallelization based on tiling strategy +41/-16  matmul_utils.cpp Set warp specialization default                                                    csrc/scheduler/matmul_utils.cpp Set warp specialization as default on Hopper +4/-0      test_matmul.cpp Update MLPBenchmarkTest for persistent kernels                      tests/cpp/test_matmul.cpp Parametrize MLPBenchmarkTest to include persistent and warp specialization Add tests for broadcast inputs Skip persistent kernel tests for unsupported operations Adjust parameters for smem and register constraints +235/-80

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The PR introduces persistent kernel scheduling, but there are known issues with inlining errors when MatmulOp or LinearOp are present in the fusion. This should be addressed to ensure correctness. if (params_->tiling_strategy != MatmulParams::TilingStrategy::OneTilePerCTA) { NVF_CHECK( params_->splitk_factor == 1, "Hopper matmul scheduler does not support scheduling persistent split-K kernels"); } NVF_CHECK( params_->tiling_strategy != MatmulParams::TilingStrategy::DistributeStagesAcrossSMs, "Hopper matmul scheduler does not support distributing stages across SMs a la stream-K"); Failing Tests The FwdEpilogueFusion_BroadcastInputs test is currently failing. This should be investigated and resolved to ensure the correctness of the persistent kernel implementation. TEST_P(MLPBenchmarkTest, FwdEpilogueFusion_BroadcastInputs) { GTEST_SKIP() << "THIS TEST IS CURRENTLY FAILING" << std::endl; Test Skips Multiple tests are skipped due to unsupported features or known issues. These should be addressed or documented to ensure comprehensive testing. auto options = at::TensorOptions().dtype(at::kBFloat16).device(at::kCUDA); auto a_ref = at::randn({M, K}, options);

[rdspring1]

LGTM.

Looks like some overlap with #3642. Do you plan to merge it first?

[rdspring1]

Is this rename from a bad merge?

[jacobhinkle]

I just hadn't merged #3642 in a while I think.

[rdspring1]

I thought using warp specialization by default was causing some test failures.

[jacobhinkle]

Not anymore. I think that was before integrating the warp tile split.

[jacobhinkle]

!test

[jacobhinkle]

Grrr. Failures on ampere. Will fix before merging.

[jacobhinkle]

!test