Skip to content

vllm_omni.diffusion.distributed.utils

build_local_sp_padding_mask 

build_local_sp_padding_mask(
    batch_size: int,
    local_seq_len: int,
    sp_original_seq_len: int | None,
    sp_padding_size: int,
    sequence_parallel_rank: int,
    device,
)

Build a per-rank SP padding mask that matches the local shard shape.

get_local_device 

get_local_device() -> device

Return the torch device for the current rank based on detected device type.