[mxfp8 moe training] refactor autograd func

[danielvegamyhre]

Stacked PRs:

• [mxfp8 moe training] register sharding rules for cutedsl 2d quant kernel #4178
• [mxfp8 moe training] remove unused block_size arg #4177
• ->[mxfp8 moe training] refactor autograd func #4176

---

[mxfp8 moe training] refactor autograd func

Summary

The autograd func has become a bit messy with individual "emulated vs non-emulated" branches everywhere, every helper, etc.

This PR refactors the autograd func to have 2 clean code paths: auto vs emulated. The bifurcation happens early, and all code after that point just has one distinct path of kernels/functions, making it easier to reason about and modify.

_compute_fwd, _compute_dgrad, and _compute_wgrad all now have distinct "emulated" vs "auto" paths that are cleanly separated at the beginning of the path.

Tests

• pytest test/prototype/moe_training/test_mxfp8_grouped_mm.py -k dq_fwd_bwd
• pytest test/prototype/moe_training/test_training.py -s

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 001dc09 with merge base 02105d4 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/dtypes/test_nf4.py::TestComm::test_comm

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

[danielvegamyhre]

@drisspg @vkuzo i addressed the comments from #4157 but got in a crazy git state, recreated the PR than wrangle with it, apologies. ready for review.

[vkuzo]

the _auto suffix is weird here, just remove it?

[danielvegamyhre]

how about _compute_{fwd/dgrad/wgrad}_sm100? that is really what i am trying to communicate in the code to the reader:

• _compute_{fwd/dgrad/wgrad} = logical top level function, dispatcher based on kernel preference
• _compute_{fwd/dgrad/wgard}_sm100 = compute with sm100 kernels
• _compute_{fwd/dgrad/wgard}_emulated = compute with emulated logic

[vkuzo]

do the triton kernels match the native pytorch reference exactly? not blocking this PR, i'm just curious

[danielvegamyhre]

Yes, there are no new kernels used in this PR btw, this is still using the existing triton kernel, tested here:

ao/test/prototype/mx_formats/test_kernels.py

Lines 478 to 512 in 4611835

	def test_triton_mxfp8_dim0_randn(M, K, scaling_mode):
	x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
	x_mx_ref, x_s_ref = triton_to_mxfp8_dim0_reference(
	x, block_size=32, scaling_mode=scaling_mode
	)
	x_mx_t, x_s_t = triton_to_mxfp8_dim0(
	x,
	inner_block_size=32,
	scaling_mode=scaling_mode.value.lower(),
	)
	torch.testing.assert_close(x_mx_t, x_mx_ref, rtol=0, atol=0)
	torch.testing.assert_close(x_s_t, x_s_ref, rtol=0, atol=0)
	@pytest.mark.skipif(not has_triton(), reason="unsupported without triton")
	@pytest.mark.skipif(
	not is_sm_at_least_100() and not is_MI350(),
	reason="mxfp8 requires CUDA capability 10.0 or greater or ROCm gfx950 or greater.",
	)
	@pytest.mark.parametrize(
	"scaling_mode", (ScaleCalculationMode.FLOOR, ScaleCalculationMode.RCEIL)
	)
	def test_triton_mxfp8_dim0_zeros(scaling_mode):
	x = torch.zeros(128, 256, dtype=torch.bfloat16, device="cuda")
	x_mx_ref, x_s_ref = triton_to_mxfp8_dim0_reference(
	x, block_size=32, scaling_mode=scaling_mode
	)
	x_mx_t, x_s_t = triton_to_mxfp8_dim0(
	x,
	inner_block_size=32,
	scaling_mode=scaling_mode.value.lower(),
	)
	assert not x_mx_t.isnan().any(), "quantized tensor should not contain NaNs"
	torch.testing.assert_close(x_mx_t, x_mx_ref, rtol=0, atol=0)
	torch.testing.assert_close(x_s_t, x_s_ref, rtol=0, atol=0)

i am planning to integrate the new cutedsl kernel for this in follow ups

[vkuzo]

skimmed it and looks good, lmk if you need a proper review