Skip to content

Avoid creating tensor in CosmosAttnProcessor2_0 #11761

Open
@chenxiao111222

Description

@chenxiao111222

# 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions