[Probe] Silent wrong-address remote accesses: unbounded PE index, and sync_*_handles leaving the local base NULL

[Rachmanino]

Two independent problems in the distributed addressing path, grouped because they are the same shape: exported API surface that accepts bad input or produces a bad table, and the result is a silent access to the wrong device address rather than an error.

Every remote address the compiler emits has the form

peer_base + (local_addr - local_base)

built by RemapRemoteAddress / GetOffsetFromLocalBase in src/op/distributed_utils.h:43-57, which every non-TMA remote primitive goes through (T.ld, T.st, T.put_*, T.get_*, remote atomics, T.wait_*). Problem 1 corrupts peer_base; problem 2 corrupts local_base. Neither is detected.

Both are reproduced below.

1. An out-of-range PE is accepted at every layer and reads uninitialised constant memory

src/tl_templates/cuda/distributed/distributed.h:7,15-17:

extern "C" { __constant__ uint64_t meta_data[1024]; }

TL_DEVICE uint64_t get_remote_base_ptr(uint64_t rank) { return meta_data[2 + rank]; }

The metadata layout is [0] = rank, [1] = num_ranks, [2 + i] = peer i's base pointer, so only the first 2 + num_ranks entries are ever written. get_remote_base_ptr indexes with no bound check.

Nothing upstream checks the PE either. This compiles on a 4-rank job:

@T.prim_func
def main(A: T.Tensor((8,), "float32"), B: T.Tensor((8,), "float32")):
    with T.Kernel(1, threads=1):
        v = T.alloc_local((1,), "float32")
        T.ld(A[0], v[0], scope="sys", sem="weak", src_pe=99)
        T.st(B[0], v[0], scope="sys", sem="weak", dst_pe=99)

Emitted:

tl::ld<Semantic::WEAK, Scope::SYS, false, false>((tl::get_remote_base_ptr(99) + (tl::get_uintpt...

So the kernel reads meta_data[101] — inside the array but never initialised — treats the garbage as a base pointer, and reads/writes a wild device address. For a PE above 1021 it indexes past the array entirely.

The Python API does not validate it (it cannot know the world size), the lowering does not, and the device does not.

For contrast, the TMA path is bounded: it expands a descriptor dispatch capped at kMaxRemoteTMADescriptors = 8. That chain used to fall through to peer 7 for any out-of-range PE, which was fixed in 62ba8f5 by guarding the last case. The non-TMA paths have no equivalent bound.

Suggested fixes

• Reject a constant PE outside [0, 1022) during lowering. The metadata array bound is a compile-time fact even though the world size is not, and this catches the easy typo case (dst_pe=99) for free.
• Optionally, a device-side rank < meta_data[1] compare under a debug flag. It is one comparison against a constant-cache value already being read.
• At minimum, document in docs/distributed_api_reference.md that a non-negative PE must be < num_ranks and that violating it is undefined, alongside the existing note about the 8-peer remote-TMA limit.

2. sync_vmm_handles / sync_ipc_handles build a peer table whose local slot is NULL

src/shared_memory/shared_memory.cc:909-918 (VMM; the IPC version at :940-950 is identical in shape):

std::vector<void *> buffer_ptrs(rank_count, nullptr);
ScopedPeerMappings mappings_guard(buffer_ptrs, true);

for (size_t i = 0; i < rank_count; ++i) {
  if (i != rank_index) {                       // local slot is skipped ...
    buffer_ptrs[i] = ...open_vmm_handle_impl(h);
  }
}

void **gpu_ptr = reinterpret_cast<void **>(table_address);
SM_CUDA_CHECK(cudaMemcpy(gpu_ptr, buffer_ptrs.data(), table_bytes, cudaMemcpyHostToDevice));

The vector is zero-initialised and the loop skips i == rank_index, so the local rank's entry stays null and is then copied to the device along with the rest.

Reproduced with num_ranks=1, which isolates the question because the import loop body never runs:

import torch, tilelang
from tilelang.distributed.shared_memory import _sync_ipc_handles_raw, _sync_vmm_handles_raw

table = torch.full((1,), 0xDEADBEEF, dtype=torch.uint64, device="cuda:0")
_sync_ipc_handles_raw(0, 1, table.data_ptr(), bytes(64))
torch.cuda.synchronize()
print(hex(int(table[0])))          # -> 0x0, sentinel overwritten with null

table2 = torch.full((1,), 0xDEADBEEF, dtype=torch.uint64, device="cuda:0")
_sync_vmm_handles_raw(0, 1, table2.data_ptr(), bytes(8 + 64))
torch.cuda.synchronize()
print(hex(int(table2[0])))         # -> 0x0

before    : table[0]=0xdeadbeef
after ipc : table[0]=0x0
after vmm : table[0]=0x0

A table built this way makes local_base zero, so every remote address becomes peer_base + absolute_local_VA — far outside any mapping.

This is not internal dead code:

• both are exported in tilelang/distributed/shared_memory/__init__.py __all__ as _sync_vmm_handles / _sync_ipc_handles
• maint/scripts/smoke_test_distribution.py:17,21 lists both as required global functions, so every published wheel is verified to expose them
• the only test coverage is argument validation (testing/python/distributed/shared_memory/test_vmm_ops.py:164-165 checks the error strings for bad rank / num_ranks)

Meanwhile the path actually used by the allocator is correct: BaseAllocator._init_table → import_handles explicitly sets self._peer_ptr_values[peer_rank] = self._base_ptr.value for its own rank (tilelang/distributed/allocator.py:574). So nothing in the project consumes the broken functions.

Suggested fixes

Note this is not a one-line fix: the signature is (rank, num_ranks, buffer_ptrs_gpu_addr, packed_handles) and has no way to learn the local base pointer, so filling the slot requires an API change. Two options:

1. Delete both functions, since _init_table supersedes them. Also drop them from __all__ and from the required-symbol list in maint/scripts/smoke_test_distribution.py.
2. Add a local base pointer parameter and write buffer_ptrs[rank_index] from it, if these are meant to stay as a supported C++-side path.

I did not pick one, since it depends on whether these are intended to remain public API.

Environment

• branch release/v0.0.2-tilelang-550e25d (PR [Release] TileScale 0.0.2 on upstream TileLang 550e25d #61), reproduced at dee401b0
• 8×B200, CUDA 13.2