[Inductor] Remove bf16 fallback for atomic_add

[karthickai]

Stack from ghstack (oldest at bottom):

• -> [Inductor] Remove bf16 fallback for atomic_add #167380

Fixes: #97016

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @mlazos

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit b3e3ede with merge base 1727a71 ():

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

• inductor / unit-test / inductor-pallas-test / test (inductor-pallas, 1, 1, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  Process completed with exit code 1.

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

[shunting314]

This will always pass since result is an alias to output?

[karthickai]

thanks for the catch, I create a new expected tensor instead of reusing output.

[shunting314]

is this config still needed?

[karthickai]

yeah, the config is not needed anymore I removed it as well as this test case (

pytorch/test/inductor/test_cuda_repro.py

Lines 2260 to 2283 in e8d411e

	@skipCUDAIf(
	not SM90OrLater, "uses bfloat16 atomic add instrs which requires SM >= 90"
	)
	@unittest.skipIf(
	config.is_fbcode(),
	"bfloat16 atomic add is supported in fbcode, so we won't fallback",
	)
	def test_index_add_fallback(self):
	def f(x, y):
	return torch.index_select(x, 0, y)
	x = torch.randn(
	2000, 384, dtype=torch.bfloat16, device="cuda", requires_grad=True
	)
	y = torch.ones(713268, dtype=torch.int64, device="cuda")
	x_ref = x.clone().detach().requires_grad_(True)
	y_ref = y.clone().detach()
	out, (_, bw_code) = run_fw_bw_and_get_code(lambda: torch.compile(f)(x, y))
	fc = FileCheck()
	fc.check("aten.index_add")
	fc.run(bw_code)
	self.assertEqual(f(x_ref, y_ref), out)

)

[shunting314]

Do we still need fallback atomic_add for bfloat16 if compute_capacility < 90?

[karthickai]

Yes, we still need fallback. I've added check sm < (9, 0) to fallback.

[mlazos]

Hi why are we removing this?

This is used internally

[karthickai]

Hi why are we removing this?

This is used internally

We are addressing the issue (#97016). Previously, atomic_add did not support bf16, so a fallback was implemented. triton now supports bf16 atomic add, so we are removing this fallback

[mlazos]

Hi why are we removing this?
This is used internally

We are addressing the issue (#97016). Previously, atomic_add did not support bf16, so a fallback was implemented. triton now supports bf16 atomic add, so we are removing this fallback

Okay sounds good, you may want to ensure the config is not used internally before removing it may cause errors when landing this.

[karthickai]

Hi why are we removing this?
This is used internally

We are addressing the issue (#97016). Previously, atomic_add did not support bf16, so a fallback was implemented. triton now supports bf16 atomic add, so we are removing this fallback

Okay sounds good, you may want to ensure the config is not used internally before removing it may cause errors when landing this.

I did an internal code search for the config bfloat16_atomic_adds_enabled, and only the files we touched are affected. I think it’s safe to merge.

[karthickai]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here