Skip to content

[PyTorch] Fix EP zero-copy MXFP8 scale offset to be window-relative #3447

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
phu0ngng merged 2 commits into from
Sep 1, 2026
+5 −6
Merged

[PyTorch] Fix EP zero-copy MXFP8 scale offset to be window-relative #3447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3f1dcff
[PyTorch] Fix EP zero-copy MXFP8 scale offset to be window-relative
phu0ngng Aug 31, 2026
4c0b415
Merge branch 'main' into fix/ep_pyt
phu0ngng Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/nccl-extensions
Open in desktop
Submodule nccl-extensions updated 119 files
9 changes: 4 additions & 5 deletions transformer_engine/pytorch/csrc/extensions/ep.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,12 +78,11 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) {
NVTE_CHECK(nccl_sm != nullptr,
"Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to "
"\"NCCL\" before allocating EP payload buffers.");
// rendezvous resolves ``t`` by its storage base, so get_offset() is the allocation's offset in
// the NCCL window. Add ``t``'s own storage offset so a slice/view of a symm-mem allocation
// (e.g. the scale region carved from a shared recv buffer) resolves to its true position in the
// window rather than the allocation base.
// NCCL EP consumes window-relative offsets (the NCCL window starts at the signal pad,
// not at the buffer base). get_window_offset() = buffer_offset + get_offset(); add
// ``t``'s own storage offset for slice/view positioning.
const uint64_t offset =
static_cast<uint64_t>(nccl_sm->get_offset()) +
static_cast<uint64_t>(nccl_sm->get_window_offset()) +
static_cast<uint64_t>(t.storage_offset()) * static_cast<uint64_t>(t.element_size());
return NVTECommWindow{static_cast<ncclWindow_t>(nccl_sm->get_window()), offset};
#else
Expand Down
Loading