From 3f1dcff028d5be01d643240c62f5e32aa0661592 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 31 Aug 2026 11:17:23 -0700 Subject: [PATCH] [PyTorch] Fix EP zero-copy MXFP8 scale offset to be window-relative Signed-off-by: Phuong Nguyen --- 3rdparty/nccl-extensions | 2 +- transformer_engine/pytorch/csrc/extensions/ep.cpp | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) 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