Revert "perf(ROCm): add is_rdna() detection and optimize CE loss for RDNA GPUs"

[GoldenGrapeGentleman]

Hi @danielhanchen, sorry for the noise — wanted to flag a post-merge finding.

This reverts the num_warps change from #4123 for RDNA GPUs in the chunked cross-entropy forward path.

Reason

After stricter benchmarking (10 trials × 500 iters on W7900 gfx1100), the change turns out to have no meaningful effect:

warps=16 (PR #4123)	warps=32 (original)	Δ
0.1935ms	0.1942ms	~0.4%

The kernel is memory-bandwidth bound (~80–96% HBM utilization). The performance plateau starts at warps=16, so both values are equivalent at the bandwidth ceiling — the original warps=32 is already sufficient.

More importantly, the 16 if is_cdna() or is_rdna() else 32 expression risks misleading future contributors: RDNA has warp_size=32 (same as NVIDIA), not 64 like CDNA, so the two architectures should not be grouped together in warp logic.

[gemini-code-assist]

Summary of Changes

Hello, 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 reverts a prior performance optimization aimed at RDNA GPUs within the cross-entropy loss kernel. The decision to revert was based on post-merge benchmarks demonstrating that the num_warps = 16 setting for RDNA GPUs yielded no significant performance gain compared to the original num_warps = 32. This indicates the kernel is primarily bottlenecked by memory bandwidth, rendering the specific warp count less impactful beyond a certain threshold. The change restores the original num_warps configuration and removes the unnecessary is_rdna() detection.

Highlights

• Reverted RDNA GPU Optimization: The num_warps change introduced in a previous PR for RDNA GPUs in the chunked cross-entropy forward path has been reverted.
• Performance Justification: Benchmarks showed no meaningful performance improvement from the previous change, indicating the kernel is memory-bandwidth bound and the original num_warps=32 was sufficient.
• Code Cleanup: The is_rdna() detection import and its usage in the num_warps calculation have been removed, simplifying the logic.

🧠 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

• unsloth/kernels/cross_entropy_loss.py
  • Removed is_rdna import.
  • Reverted num_warps assignment logic to its prior state.

Activity

• No specific activity (comments, reviews, progress) has been recorded for this pull request.

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 reverts a previous performance optimization for RDNA GPUs in the cross-entropy loss kernel. The original change, which adjusted num_warps, was found to have no significant performance impact. The code changes correctly revert the num_warps calculation and remove the associated is_rdna utility function import. The revert is well-justified by the provided benchmarks and the implementation is correct. I have no suggestions for improvement.

[danielhanchen]

The benchmarks clearly show warps=16 vs warps=32 is within noise for RDNA (memory-bandwidth-bound at ~80-96% HBM utilization). The semantic point is also correct -- RDNA warp_size=32 matches NVIDIA, not CDNA's warp_size=64, so they should not be grouped in warp-count logic.

After this merges, the only remaining caller of is_rdna() is the Gemma3 NaN correctness workaround in loader.py -- which needs the broad detection to stay safe.

Clean 3-line revert, zero regression risk. LGTM.