moe quant with dedicated kernels [wip]

[HDCharles]

Summary:

current status:
both kernels are working. The padding is a significant issue with compile for the pytorch kernel while the fbgemm kernel doesn't seem compatible with compile. Hopefully this can be handled using the changes mentioned below to avoid the data dependent padding.

todo:
test the no-padding compilable pytorch kernel

change base integration to grouped_gemm (another PR)

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2325

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 1 Unrelated Failure

As of commit 186708f with merge base f0f1f6c ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (disabled by #2352)
  test/integration/test_integration.py::TestSubclass::test_int4_weight_only_quant_subclass_grouped_5_cuda

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

is it possible to call this op without modifying the source model?

is there a gropup_mm for bfloat16 that we can overwrite and dispatch to scaled_grouped_mmm?

[alexsamardzic]

Yes, there is _grouped_mm in PyTorch core that does that.

[HDCharles]

that's a better integration point but i'm not sure i'll be able to complete that before i have to head out on leave.

also i'd probably make that a separate PR instead of combining everything into one since that would be a significant change to the base moe integration.

[alexsamardzic]

PR to hopefully remove need for padding groups is here: pytorch/pytorch#155466.

[ngimel]

heads up, soon we won't need padding once #155466 lands

[ngimel]

we have fbgemm-like kernels available via autotuning in torch.compile, thanks to #155138, do you think we still need separate fbgemm path?

[alexsamardzic]

PR pytorch/pytorch#155466, that makes it possible to avoid padding, is merged. Here is a quick patch to remove padding (note that it also disables FBGEMM altogether, so _scaled_grouped_mm implementation from PyTorch core is used):

096_fuse_moeb-diff.txt

I'm not completely sure that scale tensor shape adjustment I made here are correct, but in any case this patch, together with latest PyTorch used, will make all the tests in test_moe_quant.py pass.

[ngimel]

@alexsamardzic we still would need padding for backward where K could possibly become 0?

[alexsamardzic]

@alexsamardzic we still would need padding for backward where K could possibly become 0?

This PR is not concerned about backward, but I would say @danielvegamyhre is touching on it: #2405. In any case, you have a point, here is a diff for _grouped_mm tests in PyTorch to demonstrate the issue:

diff

diff --git a/test/test_matmul_cuda.py b/test/test_matmul_cuda.py
index 4e64c807425..96667a79440 100644
--- a/test/test_matmul_cuda.py
+++ b/test/test_matmul_cuda.py
@@ -354,15 +354,15 @@ class TestMatmulCuda(TestCase):
     @unittest.skipIf(TEST_WITH_ROCM, "ROCm doesn't support CUTLASS")
     @xfailIfSM100OrLater
     @unittest.skipIf(not SM90OrLater, "Grouped gemm supported on SM90")
-    @parametrize("strided", [False, True])
-    @parametrize("a_row_major", [False, True])
-    @parametrize("b_row_major", [False, True])
-    @parametrize("use_torch_compile", [False, True])
+    @parametrize("strided", [True])
+    @parametrize("a_row_major", [True])
+    @parametrize("b_row_major", [True])
+    @parametrize("use_torch_compile", [True, False])
     def test_grouped_gemm_2d_3d(self, strided, a_row_major, b_row_major, use_torch_compile):
         device = "cuda"
         dtype = torch.bfloat16
         s_int = int(strided)
-        m, n, k, n_groups = 16, 32, 64, 4
+        m, n, k, n_groups = 3, 32, 64, 5
         if a_row_major:
             a = torch.randn(m * n_groups, k * (1 + s_int), device=device, dtype=dtype)[:, :k]
         else:
@@ -388,6 +388,7 @@ class TestMatmulCuda(TestCase):
             a.grad = None
             b.grad = None
             offs = torch.arange(m, n_groups * m + 1, m, device="cuda", dtype=torch.int32)
+            offs = torch.tensor([0, 1, 6, 6, 15], device="cuda", dtype=torch.int32)
             if check_zero_size:
                 offs[0] = offs[1]

If offsets changed say to [1, 3, 5, 6, 15], so that there are no zero sizes, it would work fine. I'm going to see is it possible to further refine these checks, in order to make it work for this case.