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",
    valid_kv_length: int | None = None,
)

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",
    valid_kv_length: int | None = None,
)