Fix register_local_buffer API break in Triton alltoallv kernel (#1996)#1996
Open
snarayankh wants to merge 1 commit intometa-pytorch:mainfrom
Open
Fix register_local_buffer API break in Triton alltoallv kernel (#1996)#1996snarayankh wants to merge 1 commit intometa-pytorch:mainfrom
snarayankh wants to merge 1 commit intometa-pytorch:mainfrom
Conversation
meta-cla
bot
added
the
CLA Signed
Contributor
|
@snarayankh has exported this pull request. If you are a Meta employee, you can view the originating Diff in D99914167. |
meta-codesync
bot
changed the title
snarayankh
force-pushed
the
export-D99914167
branch
from
91fe099 to
c9a80c2
Compare
snarayankh
force-pushed
the
export-D99914167
branch
from
c9a80c2 to
d21d7b1
Compare
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 8, 2026
…pytorch#1996) Summary: Pull Request resolved: meta-pytorch#1996 The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
force-pushed
the
export-D99914167
branch
2 times, most recently
from
9a3497a to
83371f8
Compare
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 8, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
force-pushed
the
export-D99914167
branch
from
83371f8 to
af5368e
Compare
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: Pull Request resolved: meta-pytorch#1996 The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
force-pushed
the
export-D99914167
branch
from
af5368e to
d23be1f
Compare
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
force-pushed
the
export-D99914167
branch
from
d23be1f to
9745eee
Compare
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
snarayankh
pushed a commit
to snarayankh/torchcomms
that referenced
this pull request
Apr 9, 2026
…pytorch#1996) Summary: The `register_local_buffer()` API was changed to return a single `int64` (a device pointer to a `RegisteredBuffer` struct) instead of the previous 3-tuple `(base_ptr, size, nccl_win)`. This broke the Triton alltoallv kernel and all its consumers. The fix: - Remove dead kernel parameters `src_base_ptr` and `src_size` which were never used in the kernel body - Extract the `backend_window` (ncclWindow_t) field from the device-side `RegisteredBuffer` struct at byte offset 16 via cudaMemcpy D2H. This is the value that `put_block_direct` and `put_warp_chunked_direct` externs need for NVLink inline-PTX memcpy and GIN RDMA fallback - Fix all `deregister_local_buffer(*src_info)` splat-unpack calls to pass the single int handle directly This is a pure API adaptation with zero kernel behavioral change. The externs receive the identical `ncclWindow_t` value as before. The NVLink inline-PTX optimization in `put_block_direct` is preserved (no switch to the slower `put_block` path). RegisteredBuffer struct layout (all backends): offset 0: base_ptr (void*, 8 bytes) offset 8: size (size_t, 8 bytes) offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct offset 24: lkey (uint32_t, 4 bytes) Files changed: - device_alltoallv_dynamic.py: kernel params, host-side field extraction, kernel launch - alltoallv_op.py: type annotation, deregister call - test_device_alltoallv_dynamic_e2e.py: deregister calls - benchmark_device_alltoallv_dynamic.py: deregister calls Differential Revision: D99914167
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
The
register_local_buffer()API was changed to return a singleint64(a device pointer to a
RegisteredBufferstruct) instead of the previous3-tuple
(base_ptr, size, nccl_win). This broke the Triton alltoallvkernel and all its consumers.
The fix:
src_base_ptrandsrc_sizewhich werenever used in the kernel body
backend_window(ncclWindow_t) field from the device-sideRegisteredBufferstruct at byte offset 16 via cudaMemcpy D2H. This isthe value that
put_block_directandput_warp_chunked_directexternsneed for NVLink inline-PTX memcpy and GIN RDMA fallback
deregister_local_buffer(*src_info)splat-unpack calls to passthe single int handle directly
This is a pure API adaptation with zero kernel behavioral change. The
externs receive the identical
ncclWindow_tvalue as before. The NVLinkinline-PTX optimization in
put_block_directis preserved (no switch tothe slower
put_blockpath).RegisteredBuffer struct layout (all backends):
offset 0: base_ptr (void*, 8 bytes)
offset 8: size (size_t, 8 bytes)
offset 16: backend_window (void*, 8 bytes) ← extracted for put_block_direct
offset 24: lkey (uint32_t, 4 bytes)
Files changed:
kernel launch
Differential Revision: D99914167