Skip to content

Avoid creating tensor in CosmosAttnProcessor2_0 #11761

@chenxiao111222

Description

@chenxiao111222
Contributor

# 4. Prepare for GQA
query_idx = torch.tensor(query.size(3), device=query.device)
key_idx = torch.tensor(key.size(3), device=key.device)
value_idx = torch.tensor(value.size(3), device=value.device)
key = key.repeat_interleave(query_idx // key_idx, dim=3)
value = value.repeat_interleave(query_idx // value_idx, dim=3)

        # 4. Prepare for GQA
        query_idx = query.size(3)
        key_idx = key.size(3)
        value_idx = value.size(3)
        key = key.repeat_interleave(query_idx // key_idx, dim=3)
        value = value.repeat_interleave(query_idx // value_idx, dim=3)

Speedup ~10% here in Cosmos2TextToImagePipeline and Cosmos2VideoToWorldPipeline.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

      Development

      Participants

      @yiyixuxu@chenxiao111222

      Issue actions

        Avoid creating tensor in CosmosAttnProcessor2_0 · Issue #11761 · huggingface/diffusers