fix the path change in torch v2.10 for spec dec by yeyu-nvidia · Pull Request #863 · NVIDIA/Model-Optimizer

yeyu-nvidia · 2026-02-06T18:24:04Z

What does this PR do?

Type of change:
bug fix

Overview:
torch v2.10 changes the path for _SDPAMerger. will need to use the new path for import

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed.
Is this change backward compatible?: Yes/No
Did you write any new necessary tests?: Yes/No
Did you add or update any necessary documentation?: Yes/No
Did you update Changelog?: Yes/No

Additional Information

Summary by CodeRabbit

Chores
- Updated internal import references to reflect organizational changes in dependencies.

Signed-off-by: Ye Yu <yeyu@nvidia.com>

coderabbitai · 2026-02-06T18:24:23Z

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

📝 Walkthrough

Walkthrough

A single import path update for the _SDPAMerger class in the eagle utilities module reflects a reorganization of torch's distributed tensor experimental modules, moving it under a new _context_parallel submodule.

Changes

Cohort / File(s)	Summary
Import Path Update `examples/speculative_decoding/eagle_utils.py`	Updated `_SDPAMerger` import from `torch.distributed.tensor.experimental._attention` to `torch.distributed.tensor.experimental._context_parallel._attention` due to module reorganization.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~2 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title 'fix the path change in torch v2.10 for spec dec' accurately summarizes the main change: updating an import path for speculative decoding due to torch v2.10 changes.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch yeyu/fix_bug_5875785

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

examples/speculative_decoding/eagle_utils.py (1)
681-689: ⚠️ Potential issue | 🔴 Critical

Version gate rejects torch v2.10, making the import fix unreachable.

The version check restricts patch_ring_attention_for_ttt() to torch 2.8.x only (> 2.7.1 and < 2.9.0). When running on torch v2.10 (the version this PR targets), this raises a RuntimeError before reaching line 710 where _SDPAMerger is used. If the intent is to also support torch v2.10, this version range needs to be expanded.
     if not (
         Version(torch.__version__) > Version("2.7.1")
-        and Version(torch.__version__) < Version("2.9.0")
+        and Version(torch.__version__) < Version("2.11.0")
     ):
The error message on lines 686–688 should also be updated to reflect the newly supported version range.

🤖 Fix all issues with AI agents

In `@examples/speculative_decoding/eagle_utils.py`:
- Line 34: The module currently does an unconditional top-level import of
_SDPAMerger from
torch.distributed.tensor.experimental._context_parallel._attention which raises
ImportError on torch < 2.10; instead gate the import by checking
torch.__version__ (using packaging.version.Version already imported) or wrap in
try/except ImportError and fall back to the older path
torch.distributed.tensor.experimental._attention._SDPAMerger, assigning the
available class to the same symbol name _SDPAMerger so downstream code using
_SDPAMerger continues to work.

coderabbitai · 2026-02-06T18:28:26Z

 from PIL import Image
 from scripts.ar_validate import validate_ar
-from torch.distributed.tensor.experimental._attention import _SDPAMerger
+from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger


⚠️ Potential issue | 🔴 Critical

Unconditional import breaks backward compatibility with torch < 2.10.

This top-level import executes at module load time. For torch versions before 2.10, torch.distributed.tensor.experimental._context_parallel._attention likely doesn't exist, so this will raise an ImportError and prevent the entire module from loading. Since the old path was torch.distributed.tensor.experimental._attention._SDPAMerger, you need a version-gated or fallback import.

Proposed fix: conditional import

-from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger +from packaging.version import Version as _Version + +if _Version(torch.__version__) >= _Version("2.10.0"): + from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger +else: + from torch.distributed.tensor.experimental._attention import _SDPAMerger

Note: torch is already imported on line 28, and packaging.version.Version is imported on line 31, but the torch import must appear before this conditional. Alternatively, use a try/except ImportError fallback:

-from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger +try: + from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger +except ImportError: + from torch.distributed.tensor.experimental._attention import _SDPAMerger

🤖 Prompt for AI Agents

In `@examples/speculative_decoding/eagle_utils.py` at line 34, The module currently does an unconditional top-level import of _SDPAMerger from torch.distributed.tensor.experimental._context_parallel._attention which raises ImportError on torch < 2.10; instead gate the import by checking torch.__version__ (using packaging.version.Version already imported) or wrap in try/except ImportError and fall back to the older path torch.distributed.tensor.experimental._attention._SDPAMerger, assigning the available class to the same symbol name _SDPAMerger so downstream code using _SDPAMerger continues to work.

codecov · 2026-02-06T18:35:11Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.44%. Comparing base (24e3587) to head (e9ec9ee).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #863   +/-   ##
=======================================
  Coverage   73.44%   73.44%           
=======================================
  Files         197      197           
  Lines       20657    20657           
=======================================
  Hits        15172    15172           
  Misses       5485     5485

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

h-guo18 · 2026-02-06T19:04:04Z

Current CP patch depend on exactly torch2.8.0 and will raise error if torch version doesn't match:https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/eagle_utils.py#L685

If there's an import error, I think it's better to add a try-except logic, such that: cp=1 work on all torch version, and cp>1 raise error with torch!=2.8

yeyu-nvidia · 2026-02-06T19:08:03Z

Current CP patch depend on exactly torch2.8.0 and will raise error if torch version doesn't match:https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/eagle_utils.py#L685

If there's an import error, I think it's better to add a try-except logic, such that: cp=1 work on all torch version, and cp>2 raise error with torch!=2.8

Does CP patch not work for torch.2.10.0? Forcing user to torch2.8.0 is not a good idea and may not be maintainable.

yeyu-nvidia · 2026-02-06T19:15:25Z

I see this function needs to adapt to torch.2.10 https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/speculative_decoding/eagle_utils.py#L677. Is there anything else blocking?

Signed-off-by: Ye Yu <yeyu@nvidia.com>

h-guo18 · 2026-02-06T21:41:32Z

 from PIL import Image
 from scripts.ar_validate import validate_ar
-from torch.distributed.tensor.experimental._attention import _SDPAMerger
+from torch.distributed.tensor.experimental._context_parallel._attention import _SDPAMerger


We can remove this import for torch<2.10 compatibility

h-guo18 · 2026-02-06T21:45:46Z

Instead of _SDPAMerger, we can use _attention._SDPAMerger here after version check for torch<2.10 compatibility

Signed-off-by: Ye Yu <yeyu@nvidia.com>

h-guo18 · 2026-02-06T22:29:02Z

Could you also update the test container version to 2.10 for test coverage on cp=2

Model-Optimizer/.github/workflows/example_tests.yml

Lines 89 to 109 in ac30686

    
           ##### Speculative Decoding Example Tests (requires 25.08 image) ##### 
        
           speculative-decoding-pr: 
        
             needs: [check-file-changes, wait-checks] 
        
             if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' 
        
             uses: ./.github/workflows/_example_tests_runner.yml 
        
             secrets: inherit 
        
             with: 
        
               docker_image: "nvcr.io/nvidia/pytorch:25.08-py3" 
        
               example: speculative_decoding 
        
               pip_install_extras: "[hf,dev-test]" 
        
               runner: linux-amd64-gpu-l4-latest-1 
        
           speculative-decoding-non-pr: 
        
             if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }} 
        
             uses: ./.github/workflows/_example_tests_runner.yml 
        
             secrets: inherit 
        
             with: 
        
               docker_image: "nvcr.io/nvidia/pytorch:25.08-py3" 
        
               example: speculative_decoding 
        
               pip_install_extras: "[hf,dev-test]" 
        
               runner: linux-amd64-gpu-h100-latest-2

Signed-off-by: Ye Yu <yeyu@nvidia.com>

**Type of change:** bug fix **Overview:** torch v2.10 changes the path for _SDPAMerger. will need to use the new path for import  ```python ```   - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No  - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No    * **Chores** * Updated internal import references to reflect organizational changes in dependencies.  --------- Signed-off-by: Ye Yu <yeyu@nvidia.com>

## What does this PR do? **Type of change:** bug fix **Overview:** torch v2.10 changes the path for _SDPAMerger. will need to use the new path for import ## Usage  ```python # Add a code snippet demonstrating how to use this ``` ## Testing  ## Before your PR is "*Ready for review*"  - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No  - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No  ## Additional Information   ## Summary by CodeRabbit * **Chores** * Updated internal import references to reflect organizational changes in dependencies.  --------- Signed-off-by: Ye Yu <yeyu@nvidia.com>

## What does this PR do? **Type of change:** bug fix **Overview:** torch v2.10 changes the path for _SDPAMerger. will need to use the new path for import ## Usage  ```python # Add a code snippet demonstrating how to use this ``` ## Testing  ## Before your PR is "*Ready for review*"  - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No  - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No  ## Additional Information   ## Summary by CodeRabbit * **Chores** * Updated internal import references to reflect organizational changes in dependencies.  --------- Signed-off-by: Ye Yu <yeyu@nvidia.com> Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

fix the path change in torch v2.10

4b5e662

Signed-off-by: Ye Yu <yeyu@nvidia.com>

yeyu-nvidia requested a review from a team as a code owner February 6, 2026 18:24

yeyu-nvidia requested a review from h-guo18 February 6, 2026 18:24

coderabbitai bot reviewed Feb 6, 2026

View reviewed changes

yeyu-nvidia added 3 commits February 6, 2026 11:36

update import path to torch v2.10.0

8b646ef

Signed-off-by: Ye Yu <yeyu@nvidia.com>

formatting

bc4eeb4

Signed-off-by: Ye Yu <yeyu@nvidia.com>

update test torch version

8eca187

Signed-off-by: Ye Yu <yeyu@nvidia.com>

h-guo18 reviewed Feb 6, 2026

View reviewed changes

yeyu-nvidia added 3 commits February 6, 2026 13:56

import from different paths for different pytorch versions

c6079be

Signed-off-by: Ye Yu <yeyu@nvidia.com>

update test

ccb9e43

Signed-off-by: Ye Yu <yeyu@nvidia.com>

pin to torch 2.10.0 for cp

759bb87

Signed-off-by: Ye Yu <yeyu@nvidia.com>

h-guo18 approved these changes Feb 6, 2026

View reviewed changes

update docker image pytorch version for tests

a14419f

Signed-off-by: Ye Yu <yeyu@nvidia.com>

yeyu-nvidia requested a review from a team as a code owner February 6, 2026 22:36

yeyu-nvidia requested a review from kevalmorabia97 February 6, 2026 22:36

kevalmorabia97 approved these changes Feb 6, 2026

View reviewed changes

kevalmorabia97 changed the title ~~fix the path change in torch v2.10~~ fix the path change in torch v2.10 for spec dec Feb 6, 2026

Merge branch 'main' into yeyu/fix_bug_5875785

e9ec9ee

yeyu-nvidia enabled auto-merge (squash) February 9, 2026 17:10

yeyu-nvidia merged commit a8f5314 into main Feb 9, 2026
46 of 52 checks passed

yeyu-nvidia deleted the yeyu/fix_bug_5875785 branch February 9, 2026 17:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix the path change in torch v2.10 for spec dec#863

fix the path change in torch v2.10 for spec dec#863
yeyu-nvidia merged 9 commits intomainfrom
yeyu/fix_bug_5875785

yeyu-nvidia commented Feb 6, 2026 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Feb 6, 2026 •

edited

Loading

Review skipped

Walkthrough

Changes

Estimated code review effort

Uh oh!

coderabbitai bot left a comment

Uh oh!

coderabbitai bot Feb 6, 2026

Uh oh!

codecov bot commented Feb 6, 2026 •

edited

Loading

Uh oh!

h-guo18 commented Feb 6, 2026 •

edited

Loading

Uh oh!

yeyu-nvidia commented Feb 6, 2026

Uh oh!

yeyu-nvidia commented Feb 6, 2026

Uh oh!

h-guo18 Feb 6, 2026

Uh oh!

h-guo18 Feb 6, 2026

Uh oh!

h-guo18 commented Feb 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

yeyu-nvidia commented Feb 6, 2026 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

Estimated code review effort

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai bot Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

h-guo18 commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yeyu-nvidia commented Feb 6, 2026

Uh oh!

yeyu-nvidia commented Feb 6, 2026

Uh oh!

h-guo18 Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

h-guo18 Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

h-guo18 commented Feb 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yeyu-nvidia commented Feb 6, 2026 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Feb 6, 2026 •

edited

Loading

codecov bot commented Feb 6, 2026 •

edited

Loading

h-guo18 commented Feb 6, 2026 •

edited

Loading