[Relax] Correct YaRN RoPE frequency scaling formula to align with the original paper

[Liberatedwinner]

Summary

Fixed frequency calculations for RoPE (YaRN) scaling and correct range finding.

Description

Greetings:

This PR corrects the mathematical formulation of the YaRN RoPE scaling. I have verified that this change eliminates the discrepancy observed when comparing against PyTorch baseline (an implementation of gpt-oss).

in yarn_find_correction_range()

low, high

Removed tir.floor and tir.ceil operations in yarn_find_correction_dim(). In YaRN paper, there is no floor or ceil function within calculations of those values. In gpt-oss, the implementation uses floating-point values for these thresholds to ensure smooth interpolation in the ramp function. Rounding them caused quantization errors in the ramp mask.

in rope_freq_yarn()

freq_inter

Currently, the implementation calculates the inverse frequency as:

freq_inter = tir.const(1, "float32") / tir.power(
    scaling_factor * theta, d * 2 % d_range / tir.const(d_range, "float32")
)

This implies scale is also affected by the exponent, leading to non-uniform scaling across dimensions.

According to the YaRN method (and an implementation of gpt-oss), the scaling factor should be applied linearly:

exponent = d * 2 % d_range / tir.const(d_range, "float32")
freq_power = tir.power(theta, exponent)
freq_inter = tir.const(1, "float32") / (scaling_factor * freq_power)

d_range

The yarn_find_correction_range() function was incorrectly using the current dimension index d to calculate thresholds. This caused the ramp boundaries to shift dynamically per dimension. It has been corrected to use the total dimension size (d_range) to ensure consistent frequency thresholds.

Before:

yarn_find_correction_range(..., d, ...)

After:

yarn_find_correction_range(..., d_range, ...)

Thank you very much for reading.

[gemini-code-assist[bot]]

Summary of Changes

Hello @Liberatedwinner, I'm Gemini Code Assist[^1]! 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 critical mathematical inaccuracies in the YaRN RoPE frequency scaling implementation. By aligning the formulas for frequency calculation and correction range determination with the original YaRN paper and a verified gpt-oss baseline, it eliminates discrepancies and ensures accurate, smoothly interpolated positional embeddings for large language models.

Highlights

• YaRN RoPE Formula Correction: The frequency scaling formula for YaRN RoPE has been corrected to align with the original paper and gpt-oss implementation, resolving observed discrepancies.
• Removed Floor/Ceil in Range Calculation: The tir.floor and tir.ceil operations were removed from the low and high calculations within yarn_find_correction_range, as these were not present in the original YaRN paper and caused quantization errors.
• Linear Scaling Factor Application: The scaling_factor in rope_freq_yarn is now applied linearly, ensuring uniform scaling across dimensions as per the YaRN method, instead of being affected by the exponent.
• Consistent Frequency Thresholds: The yarn_find_correction_range function now correctly uses the total dimension size (d_range) instead of the current dimension index (d) to calculate thresholds, ensuring consistent ramp boundaries.

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 :thumbsup: and :thumbsdown: 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.

[^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.

[tlopex]

cc @MasterJH5574