[Probe] Peer and multicast addresses are silently wrong when ranks allocate in different orders

[Rachmanino]

Summary

BaseAllocator derives peer addresses from the local rank's own bump-pointer offset, so every peer address and every multicast offset is only correct while all ranks have allocated exactly the same bytes in the same order. Nothing checks that, and a divergence silently produces cross-GPU accesses at the wrong address.

tilelang/distributed/allocator.py:652 (peer tensors):

peer_ptr_val = int(self._buffer_ptrs[i]) + current_offset

tilelang/distributed/allocator.py:443 (multicast tensors) has the same shape:

mcast_t = tensor_from_ptr(self._mcast_ptr, shape, dtype_str, self._device, False)
local_t = tensor_from_ptr(self._mcast_phys_ptr + current_offset, shape, dtype_str, self._device, False)

Reproduction (2 GPUs)

Rank 1 allocates a private scratch buffer before the shared one, which is enough to desynchronize the bump pointers. Rank 0 then writes through its peer view of rank 1's shared tensor.

import torch
import torch.distributed as dist
import tilelang
from testing.python.distributed._utils import distributed_test


@distributed_test(nprocs=2)
def test_divergent_allocation_order(local_rank: int, num_ranks: int):
    from tilelang.distributed.host import init_dist

    rank, _, group = init_dist(local_rank, num_ranks)
    allocator = tilelang.get_allocator(
        size=1 << 22, device=f"cuda:{local_rank}", is_distributed=True,
        local_rank=local_rank, num_local_ranks=num_ranks, group=group)
    try:
        # Only rank 1 allocates this, so its bump pointer moves ahead by 4 KB.
        if rank == 1:
            allocator._allocate_tensor((1024,), torch.int32)

        shared = allocator._allocate_tensor((16,), torch.int32, return_peers=True)

        MAGIC = 0x5A5A
        shared[rank].fill_(0)
        torch.cuda.synchronize(); dist.barrier(group)
        if rank == 0:
            shared[1].fill_(MAGIC)          # write through the peer view of rank 1
        torch.cuda.synchronize(); dist.barrier(group)
        if rank == 1:
            got = int(shared[1][0])
            print(f"rank1_sees={got:#x} expected={MAGIC:#x} ok={got == MAGIC}")
    finally:
        allocator.close()
        dist.destroy_process_group()

Observed:

rank1_sees=0x0 expected=0x5a5a ok=False

Rank 0's current_offset is 0, so its peer view of rank 1 resolves to rank1_base + 0 — which is rank 1's scratch buffer, not the shared tensor that actually lives at rank1_base + 4096. The write lands 4 KB early, rank 1 reads its own tensor and sees zero. No error is raised on either rank.

Why this matters more for multicast

For the multicast path the same divergence is worse than a wrong address. A multimem access at MC offset X touches every rank's physical memory at offset X, so all ranks must agree that offset X holds the same logical tensor. If the per-rank bump pointers diverge, a multimem.ld_reduce silently reduces across different tensors on different ranks, and the result looks plausible rather than obviously broken.

Note on the asymmetry

This is a real gap rather than plain user error, because the allocator already validates its other cross-rank invariants collectively and carefully. _validate_distributed_configuration (allocator.py:226) all-gathers and checks size, align, use_vmm, mcast_size, num_local_ranks, local_rank and device, and reports every offending rank. Against that, per-allocation offsets — the one invariant that produces silent cross-GPU corruption when broken — are not checked at all, and the SPMD requirement is not documented anywhere.

Requiring rank-identical allocation order is entirely reasonable for an SPMD allocator. The problem is that violating it is undetectable.

Options

1. Document only. State in the allocator docstring and in docs/distributed_api_reference.md that every rank must perform an identical sequence of _allocate_tensor / _allocate_mcast_tensor calls, and that peer and multicast views are undefined otherwise. Zero runtime cost.
2. Validate under a debug flag. All-gather (current_offset, bytes_alloc) on each allocation and raise if the ranks disagree, gated behind an env var (e.g. TILESCALE_CHECK_ALLOC_SYMMETRY=1) so the collective is off by default. Turns silent corruption into an immediate, precise error during development.
3. Validate always. Correct but adds a collective to every allocation, which is likely unacceptable for allocation-heavy code.

My suggestion is 1 + 2: document the requirement, and make it cheaply checkable when someone is debugging a wrong-result problem. I did not implement either, since which of these is right depends on how much per-allocation overhead you are willing to accept.

Environment

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