Repository navigation
[Feat] Refactor UVA mechanism and introduce metadata table to support remote device-to-device copy - #15
Conversation
- Updated `example_remote_copy.py` to simplify kernel parameters and improve memory allocation using a new `BaseAllocator`. - Introduced new built-in functions `get_rank`, `get_num_ranks`, `get_remote_base_ptr`, and `get_uintptr_t` for better distributed memory management. - Refactored `remote_copy` implementation to utilize the new built-in functions for rank and pointer management. - Added support for dynamic memory allocation in the `BaseAllocator` class, facilitating efficient tensor management in distributed environments. - Enhanced the Python interface for tensor creation to support custom allocators, improving flexibility in memory management. - Updated relevant tests and examples to reflect these changes and ensure compatibility with the new allocator and distributed operations.
WalkthroughThe change set replaces 2D distributed examples with 1D variants, introduces allocator-driven tensor management, adds new distributed intrinsics and device metadata plumbing, removes prior meta-data buffer mapping, extends CUDA codegen/NVSHMEM toggles, adds host-side init_table, and provides a ctypes-based CUDA allocator with a PyTorch extension for tensor-from-pointer. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant U as User Script
participant D as Dist Runtime
participant A as Allocator (ctypes)
participant J as JITKernel
participant W as Wrapper (C)
participant G as GPU Kernel
U->>D: init_dist(...)
U->>A: get_allocator(is_distributed, local_rank, group)
A->>A: cudaMalloc + (optional IPC table build)
U->>J: kernel = compile(...)
U->>J: kernel.initialize(allocator)
J->>W: init_table(table_ptr, table_size)
W->>G: cudaMemcpyToSymbol(meta_data, host_table)
U->>U: dst = tensor(..., allocator)\nsrc = tensor(..., allocator)
U->>J: kernel(dst, src)
rect rgba(200, 230, 255, 0.2)
note over G: Device-side execution
G->>G: rank = tl.get_rank()\nnum = tl.get_num_ranks()
G->>G: base = tl.get_remote_base_ptr(peer)
G->>G: off = tl.get_uintptr_t(ptr) - base_local
G->>G: push/pull via warp ops
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 25
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tilelang/distributed/common/ipc_ext.cpp (4)
51-56: Fix TensorOptions construction and element-size helper.
- .device_index(...) is not a setter; use .device(torch::kCUDA, index).
- Prefer c10::elementSize(dtype) (or at::elementSize) over torch::elementSize.
- auto option_gpu = - at::TensorOptions(at::kCUDA).dtype(dtype).device_index(current_device); - auto size = torch::elementSize(dtype) * + auto option_gpu = + at::TensorOptions().dtype(dtype).device(torch::kCUDA, current_device); + auto size = static_cast<size_t>(c10::elementSize(dtype)) * std::accumulate(shape.begin(), shape.end(), (size_t)1, std::multiplies<>());Also applies to: 53-55
59-61: Avoid throwing from the from_blob deleter.The CUDA_CHECK in the deleter will throw, which is unsafe in destructors. Log instead.
- return at::from_blob( - ptr, shape, [=](void *ptr) { CUDA_CHECK(cudaFree(ptr)); }, option_gpu); + return at::from_blob(ptr, shape, + [=](void *p) { + cudaError_t e = cudaFree(p); + if (e != cudaSuccess) { + std::fprintf(stderr, + "from_blob deleter: cudaFree failed: %s\n", + cudaGetErrorString(e)); + } + }, + option_gpu);
64-69: Memory leak and missing error check in create_ipc_handle.Avoid malloc; check the CUDA call; construct bytearray safely.
-pybind11::bytearray create_ipc_handle(void *ptr) { - cudaIpcMemHandle_t *handle = - (cudaIpcMemHandle_t *)malloc(sizeof(cudaIpcMemHandle_t)); - cudaIpcGetMemHandle(handle, ptr); - pybind11::bytearray result = {handle->reserved, CUDA_IPC_HANDLE_SIZE}; - return result; -} +pybind11::bytearray create_ipc_handle(void *ptr) { + cudaIpcMemHandle_t handle; + CUDA_CHECK(cudaIpcGetMemHandle(&handle, ptr)); + pybind11::bytearray result(reinterpret_cast<char *>(handle.reserved), + sizeof(handle)); + return result; +}
78-81: Replace VLAs with std::vector and initialize local entries.VLAs are non-standard in C++. Also ensure the “self” slot is initialized (e.g., nullptr) to avoid copying uninitialized pointers to device.
- cudaIpcMemHandle_t ipc_handles[device_ids.size()]; - // buffer_ptrs is a pointer on host and points to the address on device - void *buffer_ptrs[device_ids.size()]; + std::vector<cudaIpcMemHandle_t> ipc_handles(num_nvl_ranks); + // buffer_ptrs is a pointer on host and points to the address on device + std::vector<void *> buffer_ptrs(num_nvl_ranks, nullptr); @@ - std::memcpy(ipc_handles[i].reserved, handle_str.c_str(), + std::memcpy(ipc_handles[i].reserved, handle_str.data(), CUDA_IPC_HANDLE_SIZE); CUDA_CHECK(cudaIpcOpenMemHandle(&buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess)); } } @@ - CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs, - sizeof(void *) * device_ids.size(), + CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu, buffer_ptrs.data(), + sizeof(void *) * num_nvl_ranks, cudaMemcpyHostToDevice));Also applies to: 95-99
🧹 Nitpick comments (9)
tilelang/distributed/common/ipc_ext.cpp (1)
1-3: Include reorder is fine; add missing headers now required by this TU.Current code uses std::optional, std::accumulate and std::memcpy but does not include , , or . Add them to avoid ODR/implicit includes.
#include <ATen/ops/from_blob.h> #include <c10/core/ScalarType.h> #include <c10/cuda/CUDAFunctions.h> +#include <optional> +#include <numeric> +#include <cstring>tilelang/utils/cpp/tensor_from_ptr.cpp (2)
92-108: Ownership deleter semantics: document that cudaFree requires cudaMalloc-originated pointers.If callers pass IPC-mapped or driver-allocated pointers here, cudaFree may fail. Recommend clarifying in the docstring and consider a flag enum for free-kind if future needs arise.
69-79: Validate CUDA pointer and remove dead size calculation
- After creating
options, insert acudaPointerGetAttributes(&attr, data_ptr)check and throw ifattr.type != cudaMemoryTypeDeviceorattr.device != device.- Remove the unused size computations (
int64_t itemsize = element_size_bytes(st);andint64_t required_bytes = safe_mul_int64(nelems, itemsize);at lines 89–90).- (Optional) Extend
dtype_from_stringto handle"int16"→at::kShort,"uint16"→at::kUInt16,"uint32"→at::kUInt32, and"uint64"→at::kUInt64.tilelang/__init__.py (1)
90-91: Public re-exports look good. Consider adding to all to make the API explicit.Optional but improves discoverability and static analysis.
from .utils.tensor import tensor # noqa: F401 from .utils.allocator import get_allocator # noqa: F401 +__all__ = [*globals().get("__all__", []), "tensor", "get_allocator"]src/op/distributed.h (1)
214-233: Naming and overlap with existing intrinsics; clarify types.
- New intrinsics use lower_snake_case, while existing ones use PascalCase (e.g., GetPE, GetPENum). This is inconsistent and may confuse users alongside GetPE/GetPENum.
- Please document expected return/value types: get_rank/get_num_ranks -> int64; get_remote_base_ptr -> handle; get_uintptr_t -> uint64.
Would you like me to draft a follow-up that:
- Aliases GetPE/GetPENum to the new intrinsics (or deprecates them), and
- Updates comments to state the exact dtypes for each intrinsic?
tilelang/jit/kernel.py (1)
139-139: Persist the allocator after successful init.
self.allocatoris never set. Keep a reference to the allocator once initialization succeeds.def initialize( self, allocator: BaseAllocator, ): @@ - if result != 0: + if result != 0: error_msg = self.adapter.lib.get_last_error().decode('utf-8') raise RuntimeError(f"Initialization failed: {error_msg}") + # Persist allocator for later use (e.g., consistency checks, re-init). + self.allocator = allocatorsrc/tl_templates/cuda/distributed.h (1)
9-15: Prefer symbolic indices, and (optionally) add a debug bound check forrank.Improves readability and guards against out-of-range
rankduring development.Apply:
- TL_DEVICE uint64_t get_rank() { return meta_data[0]; } - TL_DEVICE uint64_t get_num_ranks() { return meta_data[1]; } - TL_DEVICE uint64_t get_remote_base_ptr(uint64_t rank) { - return meta_data[2 + rank]; - } + static constexpr int kRankIndex = 0; + static constexpr int kNumRanksIndex = 1; + static constexpr int kBasePtrsOffset = 2; + + TL_DEVICE uint64_t get_rank() { return meta_data[kRankIndex]; } + TL_DEVICE uint64_t get_num_ranks() { return meta_data[kNumRanksIndex]; } + TL_DEVICE uint64_t get_remote_base_ptr(uint64_t rank) { +#ifndef NDEBUG + // Bounds check in debug to catch misuses early. + if (rank >= meta_data[kNumRanksIndex]) return 0ull; +#endif + return meta_data[kBasePtrsOffset + rank]; + }tilelang/utils/allocator.py (2)
99-111: Define_table_sizefor non-distributed mode to avoid attribute errors.
table_sizeproperty can raise when not distributed.Apply:
- self._table = None + self._table = None self._buffer_ptrs = None - self._initialized = False + self._table_size = 0 + self._initialized = False
71-73: Remove unused class attributefunc.Dead field; keep the class lean.
Apply:
-class BaseAllocator: - func: Optional[callable] = None +class BaseAllocator:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (21)
examples/distributed/example_pull_warp.py(2 hunks)examples/distributed/example_push_warp.py(2 hunks)src/op/distributed.cc(1 hunks)src/op/distributed.h(1 hunks)src/op/op.h(0 hunks)src/op/remote_copy.cc(3 hunks)src/op/remote_copy.h(0 hunks)src/target/codegen_cuda.cc(10 hunks)src/target/codegen_cuda.h(1 hunks)src/tl_templates/cuda/distributed.h(1 hunks)src/transform/lower_tile_op.cc(1 hunks)tilelang/__init__.py(1 hunks)tilelang/distributed/common/ipc_ext.cpp(1 hunks)tilelang/jit/adapter/wrapper.py(1 hunks)tilelang/jit/kernel.py(3 hunks)tilelang/language/distributed/common.py(1 hunks)tilelang/utils/allocator.py(1 hunks)tilelang/utils/cpp/setup.py(1 hunks)tilelang/utils/cpp/tensor_from_ptr.cpp(1 hunks)tilelang/utils/target.py(2 hunks)tilelang/utils/tensor.py(2 hunks)
💤 Files with no reviewable changes (2)
- src/op/op.h
- src/op/remote_copy.h
🧰 Additional context used
🧬 Code graph analysis (12)
tilelang/language/distributed/common.py (1)
tilelang/language/tir/op.py (1)
call_intrin(119-144)
src/op/distributed.h (2)
tilelang/language/distributed/common.py (2)
get_rank(8-11)get_num_ranks(14-17)src/tl_templates/cuda/distributed.h (1)
get_uintptr_t(17-19)
tilelang/utils/tensor.py (2)
tilelang/utils/allocator.py (4)
BaseAllocator(71-213)device(114-115)initialized(163-164)_allocate_tensor(166-197)tilelang/utils/target.py (1)
parse_device(87-100)
tilelang/utils/target.py (1)
tilelang/utils/allocator.py (1)
device(114-115)
src/op/distributed.cc (2)
tilelang/language/distributed/common.py (2)
get_rank(8-11)get_num_ranks(14-17)src/tl_templates/cuda/distributed.h (1)
get_uintptr_t(17-19)
src/op/remote_copy.cc (2)
tilelang/language/distributed/common.py (1)
get_rank(8-11)src/tl_templates/cuda/distributed.h (1)
get_uintptr_t(17-19)
src/target/codegen_cuda.cc (3)
src/target/codegen_cuda.h (2)
use_nvshmem_(126-156)use_distributed_(124-158)tilelang/language/distributed/common.py (2)
get_rank(8-11)get_num_ranks(14-17)src/tl_templates/cuda/distributed.h (1)
get_uintptr_t(17-19)
src/tl_templates/cuda/distributed.h (2)
src/op/distributed.h (1)
tl(11-233)tilelang/language/distributed/common.py (2)
get_rank(8-11)get_num_ranks(14-17)
tilelang/__init__.py (2)
tilelang/utils/tensor.py (1)
tensor(45-57)tilelang/utils/allocator.py (1)
get_allocator(216-228)
tilelang/utils/cpp/tensor_from_ptr.cpp (1)
tilelang/utils/allocator.py (1)
device(114-115)
tilelang/jit/kernel.py (1)
tilelang/utils/allocator.py (4)
BaseAllocator(71-213)initialized(163-164)table(204-205)table_size(208-209)
tilelang/utils/allocator.py (2)
tilelang/utils/cpp/tensor_from_ptr.cpp (2)
tensor_from_ptr(69-120)tensor_from_ptr(69-71)tilelang/utils/target.py (1)
parse_device(87-100)
🔇 Additional comments (9)
tilelang/distributed/common/ipc_ext.cpp (1)
83-93: Handle-only window vs. all-gather size assumption.The assertion equates device_ids.size() with all_gathered_handles.size(), but rank indexing uses offset + i. If all_gathered_handles is global across ranks, this assert can be wrong. Please confirm expected shape and adjust either the assert or the offset computation.
Would you run/check the gather layout (per-node vs. global) and confirm rank is comparable to offset + i? If global, assert against num_global_ranks and compute offset accordingly.
src/transform/lower_tile_op.cc (1)
452-456: The struct dump will show exactly which fields are present and confirm ifmeta_data_buffer_andbuffer_to_meta_data_index_have indeed been removed. Once we’ve verified the signature, we can check all downstreamPush/Pulllowering sites for matching updates.src/target/codegen_cuda.h (1)
124-127: use_nvshmem_ flag toggled correctly in codegen_cuda.cc
All NVSHMEM operations in src/target/codegen_cuda.cc setthis->use_nvshmem_ = true, ensuring the<nvshmem.h>and<nvshmemx.h>headers are conditionally emitted as intended.src/op/remote_copy.cc (2)
16-16: Include distributed.h: LGTM.Required for new intrinsics.
65-75: Assumption check: symmetric UVA layout across ranks.These translations rely on identical offsets from each rank’s base. If NVSHMEM allocator changes layout, this breaks. Add a brief comment or DCHECK that sizes align across ranks (if metadata is available).
Also applies to: 116-125
src/op/distributed.cc (1)
197-214: LGTM: new intrinsics are registered with correct arity and opaque effect.The registrations for tl.get_rank, tl.get_num_ranks, tl.get_remote_base_ptr, and tl.get_uintptr_t look consistent with the rest of the file.
src/target/codegen_cuda.cc (1)
207-210: Request: enforce host-side bound to match meta_data[1024].Device declares
uint64_t __constant__ meta_data[1024];. Ensure the hostinit_tablerejects sizes >1024 (see kernel.py suggestion) to prevent buffer overflows.src/tl_templates/cuda/distributed.h (1)
7-7: Fix invalid device symbol declaration (duplicateextern).
extern "C" extern __device__is invalid; use a singleexternwith C linkage.Apply:
-extern "C" extern __device__ uint64_t meta_data[1024]; +extern "C" __device__ uint64_t meta_data[1024];Likely an incorrect or invalid review comment.
tilelang/utils/allocator.py (1)
7-9: Extension import path is correct; no changes needed.
|
|
||
|
|
||
| def kernel_(M, N, num_rank, block_M, threads): | ||
| def kernel_(M, num_rank, block_M, threads): |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Drop unused num_rank parameter from kernel_ signature.
Same rationale as in push example.
Apply:
-def kernel_(M, num_rank, block_M, threads):
+def kernel_(M, block_M, threads):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def kernel_(M, num_rank, block_M, threads): | |
| def kernel_(M, block_M, threads): |
🤖 Prompt for AI Agents
In examples/distributed/example_pull_warp.py around line 14, the kernel_
function signature includes an unused parameter num_rank; remove num_rank from
the function parameters and update any internal references or callers
accordingly so the signature becomes def kernel_(M, block_M, threads): and
callers pass only the remaining arguments.
| rank = T.alloc_local([1], "uint64") | ||
| num_rank = T.alloc_local([1], "uint64") | ||
| rank[0] = T.get_rank() | ||
| num_rank[0] = T.get_num_ranks() |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove unused local num_rank allocation.
Apply:
- rank = T.alloc_local([1], "uint64")
- num_rank = T.alloc_local([1], "uint64")
- rank[0] = T.get_rank()
- num_rank[0] = T.get_num_ranks()
+ rank = T.alloc_local([1], "uint64")
+ rank[0] = T.get_rank()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| rank = T.alloc_local([1], "uint64") | |
| num_rank = T.alloc_local([1], "uint64") | |
| rank[0] = T.get_rank() | |
| num_rank[0] = T.get_num_ranks() | |
| rank = T.alloc_local([1], "uint64") | |
| rank[0] = T.get_rank() |
🤖 Prompt for AI Agents
In examples/distributed/example_pull_warp.py around lines 22 to 25, remove the
unused local variable allocation for num_rank: delete the T.alloc_local([1],
"uint64") for num_rank and the subsequent assignment num_rank[0] =
T.get_num_ranks(); if any later code needs the number of ranks, call
T.get_num_ranks() directly where used or store it in a used variable name,
otherwise simply remove the allocation and assignment to eliminate the unused
local.
| warp_idx = T.get_thread_binding(0) // 32 | ||
| warp_copy_size = T.ceildiv(block_M, threads // 32) | ||
| warp_start = bx * block_M + warp_copy_size * warp_idx | ||
| T.pull_warp( | ||
| src=T.address_of(src[bid * block_M, 0]), | ||
| dst=T.address_of(dst[bid * block_M, 0]), | ||
| size=block_M * N, | ||
| src=T.address_of(src[warp_start]), | ||
| dst=T.address_of(dst[warp_start]), | ||
| size=warp_copy_size, | ||
| src_pe=rank[0] ^ 1, |
There was a problem hiding this comment.
Tail-safe copy to avoid OOB on the last block/warp.
Mirror the guarded sizing from push.
Apply:
- warp_idx = T.get_thread_binding(0) // 32
- warp_copy_size = T.ceildiv(block_M, threads // 32)
- warp_start = bx * block_M + warp_copy_size * warp_idx
+ warp_idx = T.get_thread_binding(0) // 32
+ warp_cnt = threads // 32
+ tile_start = bx * block_M
+ tile_size = T.min(block_M, M - tile_start)
+ warp_copy_size = T.ceildiv(tile_size, warp_cnt)
+ warp_start = tile_start + warp_copy_size * warp_idx
+ copy_size = T.max(0, T.min(warp_copy_size, tile_size - warp_copy_size * warp_idx))
T.pull_warp(
- src=T.address_of(src[warp_start]),
- dst=T.address_of(dst[warp_start]),
- size=warp_copy_size,
+ src=T.address_of(src[warp_start]),
+ dst=T.address_of(dst[warp_start]),
+ size=copy_size,
src_pe=rank[0] ^ 1,
unroll_factor=4)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| warp_idx = T.get_thread_binding(0) // 32 | |
| warp_copy_size = T.ceildiv(block_M, threads // 32) | |
| warp_start = bx * block_M + warp_copy_size * warp_idx | |
| T.pull_warp( | |
| src=T.address_of(src[bid * block_M, 0]), | |
| dst=T.address_of(dst[bid * block_M, 0]), | |
| size=block_M * N, | |
| src=T.address_of(src[warp_start]), | |
| dst=T.address_of(dst[warp_start]), | |
| size=warp_copy_size, | |
| src_pe=rank[0] ^ 1, | |
| warp_idx = T.get_thread_binding(0) // 32 | |
| warp_cnt = threads // 32 | |
| tile_start = bx * block_M | |
| tile_size = T.min(block_M, M - tile_start) | |
| warp_copy_size = T.ceildiv(tile_size, warp_cnt) | |
| warp_start = tile_start + warp_copy_size * warp_idx | |
| copy_size = T.max(0, T.min(warp_copy_size, tile_size - warp_copy_size * warp_idx)) | |
| T.pull_warp( | |
| src=T.address_of(src[warp_start]), | |
| dst=T.address_of(dst[warp_start]), | |
| size=copy_size, | |
| src_pe=rank[0] ^ 1, | |
| unroll_factor=4) |
🤖 Prompt for AI Agents
In examples/distributed/example_pull_warp.py around lines 26 to 33, the
pull_warp call uses warp_copy_size unguarded and can read past the final
block/warp; mirror the push-side guarded sizing by computing the remaining
elements for this warp and using the min(remaining, warp_copy_size) (or zero if
none) as the size argument to T.pull_warp, so the pull never requests more
elements than exist in the block/end of the tensor.
| kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads)) | ||
| kernel.initialize(allocator=allocator) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Update compile-call to new kernel_ signature.
Apply:
- kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads))
+ kernel = tilelang.compile(kernel_(M, BLOCK_M, threads))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads)) | |
| kernel.initialize(allocator=allocator) | |
| kernel = tilelang.compile(kernel_(M, BLOCK_M, threads)) | |
| kernel.initialize(allocator=allocator) |
🤖 Prompt for AI Agents
In examples/distributed/example_pull_warp.py around lines 53 to 54, the call
kernel = tilelang.compile(kernel_(M, num_ranks, BLOCK_M, threads)) needs to be
updated to match the new kernel_ function signature: open the kernel_
definition, reorder or add/remove arguments to match its current parameters
(e.g., move or rename BLOCK_M/threads or include new defaults), then call
tilelang.compile with the corrected kernel_ invocation and keep
kernel.initialize(allocator=allocator) unchanged.
|
|
||
|
|
||
| def kernel_(M, N, num_rank, block_M, threads): | ||
| def kernel_(M, num_rank, block_M, threads): |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Drop unused num_rank parameter from kernel_ signature.
It’s unused and shadowed; rely on intrinsics inside the kernel.
Apply:
-def kernel_(M, num_rank, block_M, threads):
+def kernel_(M, block_M, threads):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def kernel_(M, num_rank, block_M, threads): | |
| def kernel_(M, block_M, threads): |
🤖 Prompt for AI Agents
In examples/distributed/example_push_warp.py around line 14, the kernel_
function signature currently includes an unused and shadowed parameter
`num_rank`; remove `num_rank` from the signature (i.e., change def kernel_(M,
block_M, threads):) and update any local references that expected that
positional argument to rely on intrinsics inside the kernel instead. Also adjust
any calls or tests that pass `num_rank` to kernel_ so they omit that argument,
and update related docstrings/comments if they mention the parameter.
| def _allocate_tensor(self, | ||
| shape: Tuple[int, ...], | ||
| dtype: torch.dtype, | ||
| take_ownership: bool = False) -> torch.Tensor: | ||
|
|
||
| numel = _prod_shape(shape) | ||
| itemsize = _element_size_bytes(dtype) | ||
| bytes_needed = numel * itemsize | ||
|
|
||
| bytes_alloc = _align_up(bytes_needed, self._align) | ||
|
|
||
| if not isinstance(self._ptr, ctypes.c_void_p): | ||
| raise TypeError("self._ptr must be ctypes.c_void_p") | ||
| cur_ptr_val = int(self._ptr.value) | ||
| if cur_ptr_val == 0: | ||
| raise RuntimeError("null device pointer") | ||
|
|
||
| dtype_str = _dtype_to_str.get(dtype) | ||
| if dtype_str is None: | ||
| dtype_str = str(dtype).split(".")[-1] | ||
|
|
||
| t = tensor_from_ptr(cur_ptr_val, | ||
| list(shape) if isinstance(shape, tuple) else [shape], dtype_str, | ||
| self._device, take_ownership) | ||
|
|
||
| if take_ownership: | ||
| self._ptr = ctypes.c_void_p(0) | ||
| else: | ||
| new_ptr_val = cur_ptr_val + bytes_alloc | ||
| self._ptr.value = new_ptr_val | ||
|
|
||
| return t |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add capacity checks to prevent device OOM/overrun and keep pointer alignment.
Currently, allocations can silently walk past the reserved size.
Apply:
- numel = _prod_shape(shape)
+ numel = _prod_shape(shape)
itemsize = _element_size_bytes(dtype)
bytes_needed = numel * itemsize
-
- bytes_alloc = _align_up(bytes_needed, self._align)
+ bytes_alloc = _align_up(bytes_needed, self._align)
if not isinstance(self._ptr, ctypes.c_void_p):
raise TypeError("self._ptr must be ctypes.c_void_p")
cur_ptr_val = int(self._ptr.value)
if cur_ptr_val == 0:
raise RuntimeError("null device pointer")
+ base_ptr_val = int(self._base_ptr.value)
+ used_bytes = cur_ptr_val - base_ptr_val
+ if used_bytes < 0 or used_bytes + bytes_alloc > self.size:
+ raise MemoryError(
+ f"Allocator out of memory: requested {bytes_alloc} bytes, "
+ f"used {used_bytes} / capacity {self.size}")
dtype_str = _dtype_to_str.get(dtype)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _allocate_tensor(self, | |
| shape: Tuple[int, ...], | |
| dtype: torch.dtype, | |
| take_ownership: bool = False) -> torch.Tensor: | |
| numel = _prod_shape(shape) | |
| itemsize = _element_size_bytes(dtype) | |
| bytes_needed = numel * itemsize | |
| bytes_alloc = _align_up(bytes_needed, self._align) | |
| if not isinstance(self._ptr, ctypes.c_void_p): | |
| raise TypeError("self._ptr must be ctypes.c_void_p") | |
| cur_ptr_val = int(self._ptr.value) | |
| if cur_ptr_val == 0: | |
| raise RuntimeError("null device pointer") | |
| dtype_str = _dtype_to_str.get(dtype) | |
| if dtype_str is None: | |
| dtype_str = str(dtype).split(".")[-1] | |
| t = tensor_from_ptr(cur_ptr_val, | |
| list(shape) if isinstance(shape, tuple) else [shape], dtype_str, | |
| self._device, take_ownership) | |
| if take_ownership: | |
| self._ptr = ctypes.c_void_p(0) | |
| else: | |
| new_ptr_val = cur_ptr_val + bytes_alloc | |
| self._ptr.value = new_ptr_val | |
| return t | |
| bytes_alloc = _align_up(bytes_needed, self._align) | |
| if not isinstance(self._ptr, ctypes.c_void_p): | |
| raise TypeError("self._ptr must be ctypes.c_void_p") | |
| cur_ptr_val = int(self._ptr.value) | |
| if cur_ptr_val == 0: | |
| raise RuntimeError("null device pointer") | |
| # --- begin capacity check to prevent overruns --- | |
| base_ptr_val = int(self._base_ptr.value) | |
| used_bytes = cur_ptr_val - base_ptr_val | |
| if used_bytes < 0 or used_bytes + bytes_alloc > self.size: | |
| raise MemoryError( | |
| f"Allocator out of memory: requested {bytes_alloc} bytes, " | |
| f"used {used_bytes} / capacity {self.size}") | |
| # --- end capacity check --- | |
| dtype_str = _dtype_to_str.get(dtype) |
🤖 Prompt for AI Agents
In tilelang/utils/allocator.py around lines 166 to 197, the allocator currently
advances self._ptr without checking remaining capacity or maintaining alignment;
add checks to compute remaining_bytes = (base_ptr + self._size) - cur_ptr_val
(use the original base pointer stored on the allocator) and ensure bytes_alloc
<= remaining_bytes, raising a RuntimeError (device OOM) if not, and do not
modify self._ptr on failure; also ensure the pointer advance preserves alignment
by aligning the new_ptr_val (or align cur_ptr_val before allocation) and verify
the aligned allocation still fits inside capacity before assigning self._ptr
(and when take_ownership is True, still validate bytes_alloc <= remaining_bytes
before zeroing the pointer).
| def get_allocator(size: int = 2**30, | ||
| device: str = "cuda", | ||
| is_distributed: bool = True, | ||
| local_rank: int = 0, | ||
| num_local_ranks: int = 1, | ||
| group: Optional[dist.ProcessGroup] = None) -> BaseAllocator: |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Default is_distributed=False to avoid assertion traps on casual use.
Current defaults assert because group=None.
Apply:
-def get_allocator(size: int = 2**30,
- device: str = "cuda",
- is_distributed: bool = True,
+def get_allocator(size: int = 2**30,
+ device: str = "cuda",
+ is_distributed: bool = False,
local_rank: int = 0,
num_local_ranks: int = 1,
group: Optional[dist.ProcessGroup] = None) -> BaseAllocator:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def get_allocator(size: int = 2**30, | |
| device: str = "cuda", | |
| is_distributed: bool = True, | |
| local_rank: int = 0, | |
| num_local_ranks: int = 1, | |
| group: Optional[dist.ProcessGroup] = None) -> BaseAllocator: | |
| def get_allocator(size: int = 2**30, | |
| device: str = "cuda", | |
| is_distributed: bool = False, | |
| local_rank: int = 0, | |
| num_local_ranks: int = 1, | |
| group: Optional[dist.ProcessGroup] = None) -> BaseAllocator: |
🤖 Prompt for AI Agents
In tilelang/utils/allocator.py around lines 216 to 221, change the function
signature to default is_distributed to False and guard any assertions or code
paths that assume a ProcessGroup: update the signature parameter is_distributed:
bool = False, and ensure any assertion like assert group is not None (or code
that uses group) is only executed when is_distributed is True (or provide a
clear fallback for non-distributed use). This prevents accidental assertion
failures when callers omit group for casual/local usage.
| from setuptools import setup | ||
| from torch.utils.cpp_extension import BuildExtension, CUDAExtension | ||
|
|
||
| ABI = 0 # from torch._C._GLIBCXX_USE_CXX11_ABI == False |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Derive ABI from the current PyTorch build to avoid ODR/ABI mismatches.
Hard-coding ABI = 0 can break on toolchains where PyTorch is built with the C++11 ABI enabled.
-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from setuptools import setup
+import torch
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
@@
-ABI = 0 # from torch._C._GLIBCXX_USE_CXX11_ABI == False
+ABI = int(getattr(torch._C, "_GLIBCXX_USE_CXX11_ABI", 0))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| ABI = 0 # from torch._C._GLIBCXX_USE_CXX11_ABI == False | |
| from setuptools import setup | |
| import torch | |
| from torch.utils.cpp_extension import BuildExtension, CUDAExtension | |
| ABI = int(getattr(torch._C, "_GLIBCXX_USE_CXX11_ABI", 0)) |
🤖 Prompt for AI Agents
In tilelang/utils/cpp/setup.py around line 4, ABI is hard-coded to 0 which can
cause ODR/ABI mismatches; change it to derive the value from the current PyTorch
build by importing torch and setting ABI = 1 if torch._C has attribute
_GLIBCXX_USE_CXX11_ABI and it's True, otherwise ABI = 0, and include a safe
fallback when torch or the attribute is unavailable (e.g., default to 0 but
log/warn).
|
|
||
| def parse_device(device: Union[str, torch.device, int]) -> int: | ||
| if isinstance(device, str): | ||
| if device.startswith("cuda"): | ||
| return torch.cuda.current_device() | ||
| elif device == "cpu": | ||
| return -1 | ||
| else: | ||
| raise ValueError(f"unknown device string: {device}") | ||
| elif isinstance(device, torch.device): | ||
| return device.index if device.type == "cuda" else -1 | ||
| elif isinstance(device, int): | ||
| return device | ||
| else: | ||
| raise TypeError("device must be str|torch.device|int") No newline at end of file |
There was a problem hiding this comment.
Make parse_device robust: support 'cuda:N', handle None index, and validate range.
Current impl returns current_device for any "cuda..." string and can return None for torch.device with no index. This can misroute allocations.
-def parse_device(device: Union[str, torch.device, int]) -> int:
- if isinstance(device, str):
- if device.startswith("cuda"):
- return torch.cuda.current_device()
- elif device == "cpu":
- return -1
- else:
- raise ValueError(f"unknown device string: {device}")
- elif isinstance(device, torch.device):
- return device.index if device.type == "cuda" else -1
- elif isinstance(device, int):
- return device
- else:
- raise TypeError("device must be str|torch.device|int")
+def parse_device(device: Union[str, torch.device, int]) -> int:
+ def _validate_cuda_idx(idx: int) -> int:
+ if idx < 0:
+ raise ValueError(f"cuda device index must be >= 0, got {idx}")
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA is not available")
+ if idx >= torch.cuda.device_count():
+ raise ValueError(f"cuda device index {idx} out of range [0,{torch.cuda.device_count()-1}]")
+ return idx
+
+ if isinstance(device, str):
+ if device == "cpu":
+ return -1
+ if device.startswith("cuda"):
+ if ":" in device:
+ _, idx_s = device.split(":", 1)
+ return _validate_cuda_idx(int(idx_s))
+ return _validate_cuda_idx(torch.cuda.current_device())
+ raise ValueError(f"unknown device string: {device}")
+ if isinstance(device, torch.device):
+ if device.type == "cuda":
+ idx = torch.cuda.current_device() if device.index is None else device.index
+ return _validate_cuda_idx(int(idx))
+ return -1
+ if isinstance(device, int):
+ return -1 if device == -1 else _validate_cuda_idx(device)
+ raise TypeError("device must be str|torch.device|int")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def parse_device(device: Union[str, torch.device, int]) -> int: | |
| if isinstance(device, str): | |
| if device.startswith("cuda"): | |
| return torch.cuda.current_device() | |
| elif device == "cpu": | |
| return -1 | |
| else: | |
| raise ValueError(f"unknown device string: {device}") | |
| elif isinstance(device, torch.device): | |
| return device.index if device.type == "cuda" else -1 | |
| elif isinstance(device, int): | |
| return device | |
| else: | |
| raise TypeError("device must be str|torch.device|int") | |
| def parse_device(device: Union[str, torch.device, int]) -> int: | |
| def _validate_cuda_idx(idx: int) -> int: | |
| if idx < 0: | |
| raise ValueError(f"cuda device index must be >= 0, got {idx}") | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA is not available") | |
| if idx >= torch.cuda.device_count(): | |
| raise ValueError(f"cuda device index {idx} out of range [0,{torch.cuda.device_count()-1}]") | |
| return idx | |
| if isinstance(device, str): | |
| if device == "cpu": | |
| return -1 | |
| if device.startswith("cuda"): | |
| if ":" in device: | |
| _, idx_s = device.split(":", 1) | |
| return _validate_cuda_idx(int(idx_s)) | |
| return _validate_cuda_idx(torch.cuda.current_device()) | |
| raise ValueError(f"unknown device string: {device}") | |
| if isinstance(device, torch.device): | |
| if device.type == "cuda": | |
| idx = torch.cuda.current_device() if device.index is None else device.index | |
| return _validate_cuda_idx(int(idx)) | |
| return -1 | |
| if isinstance(device, int): | |
| return -1 if device == -1 else _validate_cuda_idx(device) | |
| raise TypeError("device must be str|torch.device|int") |
| if allocator is not None: | ||
| assert allocator.initialized(), "Allocator is not initialized" | ||
| if device is not None: | ||
| device = parse_device(device) | ||
| assert allocator.device == device, f"Allocator device must be the " \ | ||
| f"same as the device of the tensor, but got {allocator.device} != {device}" | ||
| return allocator._allocate_tensor(shape, dtype) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Replace asserts with explicit exceptions and clarify device check.
Asserts can be stripped with -O; raise explicit exceptions and avoid mutating the device parameter.
- if allocator is not None:
- assert allocator.initialized(), "Allocator is not initialized"
- if device is not None:
- device = parse_device(device)
- assert allocator.device == device, f"Allocator device must be the " \
- f"same as the device of the tensor, but got {allocator.device} != {device}"
- return allocator._allocate_tensor(shape, dtype)
+ if allocator is not None:
+ if not allocator.initialized():
+ raise RuntimeError("Allocator is not initialized")
+ if device is not None:
+ dev_id = parse_device(device)
+ if allocator.device != dev_id:
+ raise ValueError(
+ f"Allocator device must match tensor device: {allocator.device} != {dev_id}"
+ )
+ return allocator._allocate_tensor(shape, dtype)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if allocator is not None: | |
| assert allocator.initialized(), "Allocator is not initialized" | |
| if device is not None: | |
| device = parse_device(device) | |
| assert allocator.device == device, f"Allocator device must be the " \ | |
| f"same as the device of the tensor, but got {allocator.device} != {device}" | |
| return allocator._allocate_tensor(shape, dtype) | |
| if allocator is not None: | |
| if not allocator.initialized(): | |
| raise RuntimeError("Allocator is not initialized") | |
| if device is not None: | |
| dev_id = parse_device(device) | |
| if allocator.device != dev_id: | |
| raise ValueError( | |
| f"Allocator device must match tensor device: {allocator.device} != {dev_id}" | |
| ) | |
| return allocator._allocate_tensor(shape, dtype) |
🤖 Prompt for AI Agents
In tilelang/utils/tensor.py around lines 49 to 55, replace the assert checks
with explicit exceptions and avoid reassigning the incoming device parameter:
check that allocator.initialized() and if not raise a RuntimeError (or
appropriate custom exception) with a clear message; if device is provided, parse
it into a new local variable (e.g., parsed_device = parse_device(device)) and
compare allocator.device to parsed_device, raising a ValueError with a
descriptive message on mismatch; finally return
allocator._allocate_tensor(shape, dtype) unchanged.
This pull request refactors the distributed remote copy (push/pull) kernel examples and simplifies the underlying implementation for symmetric communication. It introduces new TileLang intrinsics for querying rank and remote pointers, removes the need for explicit metadata buffers, and updates the CUDA code generation to use these intrinsics. The changes also update the example scripts to use the new APIs and simplify their tensor handling.
The most important changes are:
Distributed Remote Copy Kernel Refactor
get_rank,get_num_ranks,get_remote_base_ptr, andget_uintptr_tfor querying process rank and remote buffer addresses, replacing the previous metadata-based approach. [1] [2]PushWarpOpandPullWarpOplowering to use the new intrinsics for symmetric communication, eliminating the need formeta_data_bufferandbuffer_to_meta_data_index. [1] [2] [3] [4] [5]Example Script Simplification
example_push_warp.pyandexample_pull_warp.pyto remove metadata handling, use 1D tensors, and rely on the new intrinsics for distributed communication. These scripts now require exactly two ranks and use a simplified kernel launch. [1] [2] [3] [4]Code Generation and Build
These changes streamline distributed kernel implementation, make the codebase cleaner and easier to maintain, and simplify the user-facing examples.
Summary by CodeRabbit
New Features
Refactor
Chores
Changes