[None][fix] Fix unfused RoPE for yarn models: double rotation and GPT-OSS pairing by Hudayday · Pull Request #16378 · NVIDIA/TensorRT-LLM

Hudayday · 2026-07-14T15:18:44Z

Description

With rope_fusion=False, GPT-OSS (and any yarn-scaled model on the unfused
path) generates position-dependent garbage: short outputs look fine, long
generations degrade into fragments. Two independent bugs stack on this path,
and both are needed for a correct rotation — one fix controls how many
times RoPE is applied, the other how it is applied.

Bug 1: RoPE applied twice (yarn missing from `is_rope()`)

create_attention hands pos_embd_params to the backend whenever the
embedding type is not in PositionEmbeddingType.is_rope(). yarn was missing
from that list, so with rope_fusion=False the Python RotaryEmbedding
rotated q/k once and the C++ QKV preprocess (which maps yarn onto its NeoX
rotation branch) rotated them again before writing the KV cache. The cached
K ends up mscale²-scaled and rotated by twice the angle, so the error grows
with position.

Fix: add yarn to is_rope(). The existing condition then keeps the
rope params away from the backend on the unfused path, and the existing
"rope types must carry rope params" validation now covers yarn too.

Bug 2: wrong pairing convention (GPT-OSS `is_neox=False`)

GPT-OSS applies NeoX-style (rotate-half) RoPE, matching the HF reference,
but declared is_neox=False (GPT-J interleaved). The fused kernel ignores
this flag for yarn and always rotates NeoX-style, which masked the wrong
value since the model landed (#6645); the unfused path honors the flag and
paired the rotary dims incorrectly.

Fix: is_neox=True. No effect on the fused path (the flag is unused
there). modeling_mistral.py also declares is_neox=False for yarn and may
deserve the same audit.

Why both are required

Fixing only Bug 2 still double-rotates; fixing only Bug 1 leaves the single
rotation incorrectly paired. Either alone still produces garbage.

Evidence (GPT-OSS-20B)

Post-fix, the unfused K cache is bitwise-aligned with the fused path at
every position (cos = 1.0, norm ratio = 1.0). Pre-fix it decayed to
cos = −0.18 by position 150 with a norm ratio equal to the yarn mscale
(1.3466) — the fingerprint of a second yarn rotation.
AIME25 avg pass@1 (n=8, official sampling): **70.83 fused vs 70.00
unfused** (parity within noise). Pre-fix unfused scored 0.
Affects any code path that disables rope fusion for a yarn model, e.g.
sparse-attention methods that require unfused RoPE.

Test Coverage

Two regression tests in tests/unittest/_torch/modules/test_rotary_embedding.py
(registered in l0_h100.yml):

unfused yarn ⇒ the backend receives no position-embedding params (RoPE
cannot be applied twice);
the unfused yarn rotation matches the fused kernel's NeoX rotate-half
convention on the same cos/sin table (bit-level comparison).

Summary by CodeRabbit

Bug Fixes
- Corrected YARN rotary positional embedding behavior for GPT-OSS models.
- Prevented duplicate application of YARN embeddings when rotary fusion is disabled.
- Improved numerical correctness for unfused YARN rotary calculations.
Tests
- Added coverage validating YARN embedding configuration and output accuracy.

Two stacked bugs made rope_fusion=False produce position-dependent garbage for GPT-OSS (and would for any yarn-scaled model): 1. Double RoPE: create_attention hands pos_embd_params to the backend whenever the embedding type is not in PositionEmbeddingType.is_rope(). yarn was missing from that list, so with rope_fusion=False the Python RotaryEmbedding rotated q/k once and the C++ QKV preprocess (which maps yarn onto its NeoX rotation branch) rotated them again before writing the KV cache: cached K ended up mscale^2-scaled and rotated by twice the angle, so the error grows with position and long generations degrade into fragments. Adding yarn to is_rope() restores the intended ownership (Python side rotates, backend receives no rope params) and extends the existing params validation to yarn. 2. GPT-OSS declared is_neox=False, but the model applies NeoX-style (rotate-half) RoPE, as in the HF reference. The fused kernel ignores the flag for yarn and always rotates NeoX-style, which masked the wrong value; the unfused path honors it and applied interleaved pairing. (modeling_mistral.py declares is_neox=False for yarn too and may deserve the same audit.) Verified on GPT-OSS-20B: post-fix unfused K cache is bitwise-aligned with the fused path at every position (cos=1.0, norm ratio=1.0); AIME25 avg pass@1 (n=8, official sampling) is 70.83 fused vs 70.00 unfused. Pre-fix unfused scored 0. Regression tests cover backend ownership and pairing-convention parity against the kernel cos/sin table. Signed-off-by: tianruih <tianruih@nvidia.com>

coderabbitai · 2026-07-14T15:26:29Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ef64dfda-d53e-44aa-bf70-d58e818bfbc7

📥 Commits

Reviewing files that changed from the base of the PR and between 671bfec and c1455b4.

📒 Files selected for processing (3)

tensorrt_llm/_torch/models/modeling_gpt_oss.py
tensorrt_llm/functional.py
tests/unittest/_torch/modules/test_rotary_embedding.py

📝 Walkthrough

Walkthrough

YARN is now classified as a RoPE type, GPT-OSS uses NeoX YARN parameterization, and tests cover unfused YARN ownership and numerical rotation behavior.

Changes

YARN RoPE handling

Layer / File(s)	Summary
YARN RoPE configuration `tensorrt_llm/functional.py`, `tensorrt_llm/_torch/models/modeling_gpt_oss.py`	Classifies YARN as RoPE and sets GPT-OSS `RopeParams.is_neox` to `True`, with comments describing fused and unfused behavior.
Unfused YARN validation `tests/unittest/_torch/modules/test_rotary_embedding.py`	Verifies that unfused attention does not pass position-embedding parameters to the backend and that Python YARN rotation matches the NeoX rotate-half reference.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: perkzzheng, yunruis

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main fixes for unfused YaRN RoPE and GPT-OSS rotary pairing.
Description check	✅ Passed	The description includes the issue, fixes, and test coverage and is mostly complete relative to the template.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

Hudayday · 2026-07-14T15:28:46Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-14T15:35:11Z

PR_Github #59223 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-14T19:18:29Z

PR_Github #59223 [ run ] completed with state FAILURE. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47720 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T01:32:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T01:38:14Z

PR_Github #59318 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T06:03:13Z

PR_Github #59318 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47801 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T06:03:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T06:09:58Z

PR_Github #59374 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T07:59:33Z

PR_Github #59374 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47848 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T08:08:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T08:14:40Z

PR_Github #59403 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T10:07:50Z

PR_Github #59403 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47875 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T10:11:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T10:16:26Z

PR_Github #59434 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T12:03:06Z

PR_Github #59434 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47903 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T12:03:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T12:10:01Z

PR_Github #59449 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T13:59:05Z

PR_Github #59449 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47917 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-15T14:30:57Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T14:36:39Z

PR_Github #59460 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-15T16:24:31Z

PR_Github #59460 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #47927 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Hudayday · 2026-07-16T01:34:09Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T01:40:08Z

PR_Github #59570 [ run ] triggered by Bot. Commit: 11e2ba3 Link to invocation

tensorrt-cicd · 2026-07-16T02:46:51Z

PR_Github #59570 [ run ] completed with state SUCCESS. Commit: 11e2ba3
/LLM/main/L0_MergeRequest_PR pipeline #48014 completed with status: 'SUCCESS'

CI Report

Link to invocation

Hudayday requested review from a team as code owners July 14, 2026 15:18

Hudayday requested review from BowenFu, PerkzZheng, brnguyen2, kaiyux and yunruis July 14, 2026 15:18

github-actions Bot assigned Hudayday Jul 14, 2026

Hudayday changed the title ~~[None][fix] Fix unfused RoPE for yarn models: double rotation and GPT-OSS pairing~~ [None][fix] Fix unfused RoPE for yarn models: double rotation and GPT-OSS pairing Jul 14, 2026

Hudayday force-pushed the fix-gptoss-unfused-rope branch from c1455b4 to 11e2ba3 Compare July 14, 2026 15:22

BowenFu approved these changes Jul 15, 2026

View reviewed changes

Hudayday enabled auto-merge (squash) July 15, 2026 06:03

yechank-nvidia approved these changes Jul 15, 2026

View reviewed changes

brnguyen2 approved these changes Jul 16, 2026

View reviewed changes

PerkzZheng approved these changes Jul 17, 2026

View reviewed changes

Hudayday merged commit 29f4552 into NVIDIA:main Jul 17, 2026
8 checks passed

Uh oh!

Conversation

Hudayday commented Jul 14, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Bug 1: RoPE applied twice (yarn missing from is_rope())

Bug 2: wrong pairing convention (GPT-OSS is_neox=False)

Why both are required

Evidence (GPT-OSS-20B)

Test Coverage

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 14, 2026

Walkthrough

Changes

Uh oh!

Hudayday commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Hudayday commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Hudayday commented Jul 14, 2026 •

edited by coderabbitai Bot

Loading

Bug 1: RoPE applied twice (yarn missing from `is_rope()`)

Bug 2: wrong pairing convention (GPT-OSS `is_neox=False`)