Skip to content

vllm_omni.diffusion.attention.backends.ring.ring_utils ¶

LseLayout `module-attribute` ¶

LseLayout = Literal['bhs', 'bsh']

flatten_varlen_lse ¶

flatten_varlen_lse(lse, cu_seqlens)

ring_kv_block_valid_length ¶

ring_kv_block_valid_length(
    valid_kv_length: int | None,
    block_size: int,
    block_rank: int,
    world_size: int,
) -> int

Return the valid prefix length in one globally ordered ring K/V block.

unflatten_varlen_lse ¶

unflatten_varlen_lse(lse, cu_seqlens, max_seqlen: int)

update_out_and_lse ¶

update_out_and_lse(
    out: Tensor | None,
    lse: Tensor | None,
    block_out: Tensor,
    block_lse: Tensor,
    slice_=None,
    *,
    lse_layout: LseLayout,
) -> tuple[Tensor, Tensor]