flashinfer.green_ctx.empty_per_locality_domain

flashinfer.green_ctx.empty_per_locality_domain(part_shape: Tuple[int, ...], dtype: dtype, dev: device) → Tensor

Allocate an uninitialized tensor of shape [num_domains, *part_shape] whose part i lives in locality domain i’s memory (one virtual address range over per-domain physical memory). Each part is contiguous; the stride of dim 0 is the part size rounded up to the allocation granularity.

The memory comes from the CUDA driver (cuMemCreate / cuMemMap), not from PyTorch’s caching allocator: it does not count in torch.cuda.memory_allocated() and cannot be reused by other tensors. It is released when the tensor and all its views are freed (after a device synchronize; if that happens during CUDA graph capture, at the next allocation or release outside capture).

Example: flashinfer.gdn_decode.ugpu_state_pool() builds a GDN state pool from two parts, so that each die reads and writes only its own half of every state.

Parameters:

  • part_shape – Shape of one part.

  • dtype – Element type.

  • dev – CUDA device with at least two locality domains.

Returns:

The tensor.