diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 9f47d6eb3b..6a81bce7c9 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 9f47d6eb3b60962d8157a579b4caaaa4ae6b19f4 +Subproject commit 6a81bce7c9fe5874e8c852224fa47c077033da72 diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 7fa7f094f0..c74d4ddb6d 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -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(nccl_sm->get_offset()) + + static_cast(nccl_sm->get_window_offset()) + static_cast(t.storage_offset()) * static_cast(t.element_size()); return NVTECommWindow{static_cast(nccl_sm->get_window()), offset}; #else