vllm_omni.diffusion.utils.kv_utils ¶

Utilities for batching variable-length tensors in diffusion attention.

left_pad_stack ¶

left_pad_stack(
    tensors: list[Tensor],
) -> tuple[Tensor, Tensor | None]

Left-pad and stack variable-length tensors. Only dim 0 may vary, and it's assumed that all tensors are the same dtype & use the same device.

Returns (stacked, mask) where mask is a 2D boolean mask, and both tensors are on the device of the provided tensors. If all tensors are the same length, None is returned for mask.