Skip to content

[PyTorch] Fix: Resolve EP symm-mem window offset for both old and new torch version - #3466

Open
phu0ngng wants to merge 2 commits into
NVIDIA:mainfrom
phu0ngng:fix_pyt_symm
Open

[PyTorch] Fix: Resolve EP symm-mem window offset for both old and new torch version #3466
phu0ngng wants to merge 2 commits into
NVIDIA:mainfrom
phu0ngng:fix_pyt_symm

Conversation

@phu0ngng

@phu0ngng phu0ngng commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Description

Select get_window_offset() or get_offset() at compile time so EP zero-copy builds against PyTorch versions before and after pytorch/pytorch#189088 moved the symm-mem signal pad to the front of the allocation.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@phu0ngng
phu0ngng requested a review from ksivaman as a code owner September 2, 2026 22:53
@phu0ngng phu0ngng changed the title [PyTorch][Fix] Resolve EP symm-mem window offset for both old and new torch version [PyTorch] Fix: Resolve EP symm-mem window offset for both old and new torch version Sep 2, 2026
…em layouts

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng

phu0ngng commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator Author

/te-ci L1 torch

@phu0ngng
phu0ngng requested a review from ptrendx September 2, 2026 22:57
@greptile-apps

greptile-apps Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR makes the NCCL symmetric-memory window offset calculation compatible across PyTorch API versions.

  • Selects get_window_offset() at compile time when the API is available.
  • Falls back to the older get_offset() API while preserving tensor-view storage offsets.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/pytorch/csrc/extensions/ep.cpp Adds compile-time selection between the old and new symmetric-memory offset APIs without changing the surrounding zero-copy window construction flow.

Reviews (2): Last reviewed commit: "Merge branch 'main' into fix_pyt_symm" | Re-trigger Greptile

@phu0ngng

phu0ngng commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator Author

All EP tests passed. The CI reports the failures FileExistsError: [Errno 17] File exists: '/logs' which is unrelated to this change.

@ptrendx ptrendx added 2.19 2.20 and removed 2.19 labels Sep 3, 2026
@ptrendx

ptrendx commented Sep 4, 2026

Copy link
Copy Markdown
Member

/te-ci pytorch

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants