flashinfer.comm.pack_strided_memory

flashinfer.comm.pack_strided_memory(ptr: int, segment_size: int, segment_stride: int, num_segments: int, dtype: dtype, dev_id: int) → Tensor

Pack a strided device allocation as a PyTorch tensor view.

A fresh DLPack capsule is created on every call, so each tensor returned by this function consumes its own capsule.

Parameters:

  • ptr (int) – Device pointer (e.g. one returned by cudaMalloc).

  • segment_size (int) – Size of each segment in bytes.

  • segment_stride (int) – Stride between consecutive segments in bytes.

  • num_segments (int) – Number of segments to expose.

  • dtype (torch.dtype) – Element dtype of the resulting tensor.

  • dev_id (int) – CUDA device ID hosting ptr.

Returns:

A tensor that views the provided device memory.

Return type:

torch.Tensor