[PIR][Pass] Add fused_flash_attn_pass

[yuanlehome]

PR Category

Inference

PR Types

New features

Description

Pcard-71500

支持如下五种pattern的融合。

         Q          K           V
         |          |           |
     transpose  transpose   transpose
         |          |           |
       scale    transpose       |
         |          |           |
         -- matmul--            |
              |                 |
    mask --- add                |
              |                 |
            cast                |
              |                 |
           softmax              |
              |                 |
             cast               |
              |                 |
              ------matmul------
                      |
                     out
         Q   K   V   None   mask
         |   |   |     |      |
         ------flash_attn------
                   |
                  out

         Q          K           V
         |          |           |
     transpose  transpose   transpose
         |          |           |
       scale    transpose       |
         |          |           |
         -- matmul--            |
              |                 |
    mask --- add                |
              |                 |
              |                 |
           softmax              |
              |                 |
              |                 |
              ------matmul------
                      |
                     out
         Q   K   V   None   mask
         |   |   |     |      |
         ------flash_attn------
                   |
                  out

         Q          K           V
         |          |           |
     transpose  transpose   transpose
         |          |           |
         |      transpose       |
         |          |           |
         -- matmul--            |
              |                 |
            scale               |
              |                 |
    mask --- add                |
              |                 |
            cast                |
              |                 |
           softmax              |
              |                 |
             cast               |
              |                 |
              ------matmul------
                      |
                     out
         Q   K   V   None   mask
         |   |   |     |      |
         ------flash_attn------
                   |
                  out

         Q          K           V
         |          |           |
     transpose  transpose   transpose
         |          |           |
         |    transpose         |
         |          |           |
         -- matmul--            |
              |                 |
            scale               |
              |                 |
    mask --- add                |
              |                 |
              |                 |
           softmax              |
              |                 |
              |                 |
              ------matmul------
                      |
                     out
         Q   K   V   None   mask
         |   |   |     |      |
         ------flash_attn------
                   |
                  out

              transpose
                     |
          -----------+----------
          |          |           |
        slice       slice      slice
          |          |           |
          Q          K           V
          |          |           |
          |       transpose      |
          |          |           |
          -- matmul--            |
               |                 |
             scale               |
               |                 |
     mask --- add                |
               |                 |
            softmax              |
               |                 |
               ------matmul------
                       |
                   transpose
                       |
                      out
            transpose
                |
          ------+------
          |     |     |
        slice slice slice
          |     |     |
          Q     K     V              mask
          |     |     |               |
    tranpose tranpose tranpose        |
          |     |     |               |
          -------flash_attn------------
                    |
                   out

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.