Skip to content

vllm_omni.diffusion.attention.backends.ring_pytorch_attn

logger module-attribute

logger = init_logger(__name__)

RingAttentionFunc

Bases: Function

Ring Attention autograd function using PyTorch SDPA (inference only, no backward).

forward staticmethod

forward(
    ctx,
    group,
    q,
    k,
    v,
    sm_scale,
    is_causal,
    op_type,
    joint_tensor_key=None,
    joint_tensor_value=None,
    joint_strategy="front",
)

ring_pytorch_attn_func

ring_pytorch_attn_func(
    q,
    k,
    v,
    dropout_p=0.0,
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),
    softcap=0.0,
    alibi_slopes=None,
    deterministic=False,
    return_attn_probs=False,
    group=None,
    op_type="efficient",
    joint_tensor_key=None,
    joint_tensor_value=None,
    joint_strategy="front",
)