fix deepgemm set_num_sms import

[shihaobai]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @shihaobai, 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 enhances the robustness of the _set_num_sms_for_deep_gemm method by addressing potential issues with the import path of the set_num_sms function from the deep_gemm library. It ensures that the system can correctly configure GPU Streaming Multiprocessors (SMs) by providing a resilient import strategy and improving error visibility through detailed logging, making the code more stable against library updates.

Highlights

• Improved deep_gemm import reliability: The pull request introduces a fallback mechanism for importing the set_num_sms function from the deep_gemm library. This ensures the function can be successfully imported even if its internal module path changes in future deep_gemm versions, by first trying deep_gemm.jit_kernels.utils and then deep_gemm directly.
• Enhanced error handling: The generic except block within the _set_num_sms_for_deep_gemm method has been replaced with a more specific except BaseException as e. This change ensures that any exceptions encountered during the set_num_sms operation are caught, and a warning message including the exception details is logged, preventing silent failures and aiding in debugging.

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 improves the import logic for set_num_sms from deepgemm by trying multiple import paths and enhances error handling by logging exceptions instead of silently passing. However, the exception handling can be further improved by catching more specific exceptions (ImportError and Exception) instead of using a bare except and BaseException, which are overly broad.

[gemini-code-assist]

While the error handling is improved, it can be made more robust by being more specific with the exceptions being caught.

1. The bare except on line 162 is too broad and can hide unexpected errors. It's better to catch the specific ImportError.
2. Catching BaseException on line 169 is generally discouraged as it can trap system-exiting exceptions like SystemExit and KeyboardInterrupt. It's safer to catch Exception for application-level errors.

Here's a suggested refactoring that addresses both points:

Suggested change

	try:
	# set num sms for deep_gemm
	from deep_gemm.jit_kernels.utils import set_num_sms
	try:
	from deep_gemm.jit_kernels.utils import set_num_sms
	except:
	from deep_gemm import set_num_sms
	deepep_sms = int(os.getenv("DEEPEP_SMS", deep_ep.Buffer.num_sms))
	device_sms = get_device_sm_count()
	deep_ep.Buffer.set_num_sms(deepep_sms)
	set_num_sms(device_sms - deepep_sms)
	except:
	pass
	except BaseException as e:
	logger.warning(f"set num sms for deep_gemm failed: {e}")
	try:
	try:
	from deep_gemm.jit_kernels.utils import set_num_sms
	except ImportError:
	from deep_gemm import set_num_sms
	deepep_sms = int(os.getenv("DEEPEP_SMS", deep_ep.Buffer.num_sms))
	device_sms = get_device_sm_count()
	deep_ep.Buffer.set_num_sms(deepep_sms)
	set_num_sms(device_sms - deepep_sms)
	except Exception as e:
	logger.warning(f"set num sms for deep_gemm failed: {e}")