[ROCm] Remove incorrect ROCm lowering for scaled_matmul to prevent segfault

[Ruturaj4]

The scaled_matmul primitive was incorrectly registering a cuDNN-specific
GPU lowering for the ROCm platform. This lowering creates a custom call
to __op$block_scaled_dot, which is a cuDNN operation that only exists
on NVIDIA GPUs. When running on ROCm, XLA would crash with a segfault
when attempting to execute this operation.

Changes:

• Remove ROCm platform lowering registration from scaled_matmul_stablehlo.py
• Add ROCm skip check in ScaledDotGeneralTest.setUp() to skip tests that
  require cuDNN block scaled dot operations

This fix resolves the segfault seen during parallel pytest execution on
8-GPU ROCm systems.

[gemini-code-assist]

Summary of Changes

Hello @Ruturaj4, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue where the scaled_matmul primitive was causing segfaults on ROCm systems due to an incorrect GPU lowering registration. The fix involves removing the problematic ROCm-specific lowering that relied on a cuDNN-exclusive operation and adding a mechanism to skip relevant tests on ROCm, thereby enhancing the stability and compatibility of JAX with ROCm environments.

Highlights

• ROCm Lowering Removal: Removed the incorrect ROCm platform lowering registration for the scaled_matmul primitive, which was erroneously attempting to use a cuDNN-specific operation (__op$block_scaled_dot) not available on ROCm, leading to segfaults.
• Test Skipping for ROCm: Implemented a skip check in ScaledDotGeneralTest.setUp() to prevent tests requiring cuDNN block scaled dot operations from running on ROCm devices, ensuring test stability and correctness.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• jax/_src/cudnn/scaled_matmul_stablehlo.py
  • Removed ROCm platform lowering registration for _scaled_matmul_p.
• tests/scaled_matmul_stablehlo_test.py
  • Added a conditional test skip for ROCm devices in ScaledDotGeneralTest when scaled_matmul requires cuDNN.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly addresses a segfault on ROCm platforms by removing an incorrect cuDNN-specific lowering for scaled_matmul. The accompanying change to skip the relevant tests in ScaledDotGeneralTest on ROCm is also appropriate. The fix is clear and directly solves the issue. I have one minor suggestion to improve the clarity of the platform-specific checks in the test file.

[gemini-code-assist]

Since the rocm and cuda platforms are mutually exclusive, using elif for the cuda check would make the logic slightly clearer and more efficient by avoiding an unnecessary check on ROCm platforms.

Suggested change

	if jtu.test_device_matches(["rocm"]):
	self.skipTest("scaled_matmul requires cuDNN which is not available on ROCm")
	if jtu.test_device_matches(["cuda"]):
	if jtu.test_device_matches(["rocm"]):
	self.skipTest("scaled_matmul requires cuDNN which is not available on ROCm")
	elif jtu.test_device_matches(["cuda"]):