Skip to content

Add logit scaling sdpa using FlexAttention for Gemma2 #32877

Description

@ArthurZucker

Feature request

With the recent publish of FlexAttention, think it makes senses to support it with sdpa path of Gemma2.
Leaving it open to the community:

  • implement the path
  • make sure integration tests work
  • add some doc

Related to #32865

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions