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.

repeat_kv ¶

repeat_kv(hidden_states: Tensor, n_rep: int) -> Tensor

Expand GQA KV from num_kv_heads to num_heads.

Input: (batch, seqlen, num_key_value_heads, head_dim) Output: (batch, seqlen, num_key_value_heads * n_rep, head_dim)