flashinfer.quantization.segment_packbits¶

flashinfer.quantization.segment_packbits(x: Tensor, indptr: Tensor, bitorder: str = 'big') → Tuple[Tensor, Tensor]¶

将一批二值段数据打包到 uint8 数组中的比特位。

对于每个段，此函数与 numpy.packbits 的语义相同。

参数:

x (torch.Tensor) – 要打包的 1D 二值数组，形状 (indptr[-1],)。
indptr (torch.Tensor) – x 中每个段的索引指针，形状 (batch_size + 1,)。 x 中的第 i 个段是 x[indptr[i]:indptr[i+1]]。
bitorder (str) – 输出的比特顺序（“bit”/”little”）。默认值为“big”。

返回值:

y (torch.Tensor) – 一个 uint8 打包数组，形状：(new_indptr[-1],)。 y[new_indptr[i]:new_indptr[i+1]] 包含打包的比特位 x[indptr[i]:indptr[i+1]]。
new_indptr (torch.Tensor) – y 中每个打包段的新索引指针，形状 (batch_size + 1,)。保证 new_indptr[i+1] - new_indptr[i] == (indptr[i+1] - indptr[i] + 7) // 8。

示例

>>> import torch
>>> from flashinfer import segment_packbits
>>> x = torch.tensor([1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1], dtype=torch.bool, device="cuda")
>>> x_packed, new_indptr = segment_packbits(x, torch.tensor([0, 4, 7, 11], device="cuda"), bitorder="big")
>>> list(map(bin, x_packed.tolist()))
['0b10110000', '0b100000', '0b11010000']
>>> new_indptr
tensor([0, 1, 2, 3], device='cuda:0')

注意

由于该函数依赖于数据，因此不支持 torch.compile。

参见

packbits