Code

Add support for Flex Attention#1675

Code

Draft

ShashankMosaicML wants to merge 92 commits intomosaicml:mainmosaicml/llm-foundry:mainfrom ShashankMosaicML:shashank/flexattentionShashankMosaicML/llm-foundry:shashank/flexattentionCopy head branch name to clipboard

+1,242-126

Commits

Commits on Nov 18, 2024

• adding flex attention

  ShashankMosaicML

  committed
  f416539

  Copy full SHA for f416539

• registrifying score mods

  ShashankMosaicML

  committed
  ac3a884

  Copy full SHA for ac3a884

• registrifying attention mask mods

  ShashankMosaicML

  committed
  31b27e2

  Copy full SHA for 31b27e2

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  c8fffa5

  Copy full SHA for c8fffa5

Commits on Nov 19, 2024

• bug_fix

  ShashankMosaicML

  committed
  86dce3b

  Copy full SHA for 86dce3b

• bug_fix

  ShashankMosaicML

  committed
  cb8f4a6

  Copy full SHA for cb8f4a6

• lint

  ShashankMosaicML

  committed
  902850a

  Copy full SHA for 902850a

• configuring test

  ShashankMosaicML

  committed
  9c9708d

  Copy full SHA for 9c9708d

• configuring tests

  ShashankMosaicML

  committed
  f1ff430

  Copy full SHA for f1ff430

• bug fix

  ShashankMosaicML

  committed
  e537f5a

  Copy full SHA for e537f5a

• fixing alibi

  ShashankMosaicML

  committed
  c527dd7

  Copy full SHA for c527dd7

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  15e303e

  Copy full SHA for 15e303e

• configuring further tests

  ShashankMosaicML

  committed
  c4ef5d9

  Copy full SHA for c4ef5d9

• refactoring

  ShashankMosaicML

  committed
  6b37427

  Copy full SHA for 6b37427

• adding warnings and errors

  ShashankMosaicML

  committed
  e30fe7a

  Copy full SHA for e30fe7a

• gating tests on torch version

  ShashankMosaicML

  committed
  924a53c

  Copy full SHA for 924a53c

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  57048e3

  Copy full SHA for 57048e3

• reorganizing function defs

  ShashankMosaicML

  committed
  67a2aea

  Copy full SHA for 67a2aea

• refactoring

  ShashankMosaicML

  committed
  04f3a62

  Copy full SHA for 04f3a62

• passing in dicts of mask and score mods

  ShashankMosaicML

  committed
  ab6c58c

  Copy full SHA for ab6c58c

• making mask and score mods configurable via yaml

  ShashankMosaicML

  committed
  3b3827d

  Copy full SHA for 3b3827d

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  be43e8d

  Copy full SHA for be43e8d

Commits on Nov 20, 2024

• adding torch.compile

  ShashankMosaicML

  committed
  2264f91

  Copy full SHA for 2264f91

• ..

  ShashankMosaicML

  committed
  e274d9f

  Copy full SHA for e274d9f

• ..

  ShashankMosaicML

  committed
  a26bb4f

  Copy full SHA for a26bb4f

• undoing comment out

  ShashankMosaicML

  committed
  d5ab7d3

  Copy full SHA for d5ab7d3

Commits on Nov 26, 2024

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  d40e978

  Copy full SHA for d40e978

• adding torch comile

  ShashankMosaicML

  committed
  5f13e7b

  Copy full SHA for 5f13e7b

• temporary commit commenting out block mask and score mod

  ShashankMosaicML

  committed
  ca8e173

  Copy full SHA for ca8e173

• undoing prev temp commit

  ShashankMosaicML

  committed
  f5486ff

  Copy full SHA for f5486ff

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  fdced3a

  Copy full SHA for fdced3a

Commits on Nov 27, 2024

• speeding up block mask generation

  ShashankMosaicML

  committed
  c53db63

  Copy full SHA for c53db63

• precompilining create block mask

  ShashankMosaicML

  committed
  ec5900d

  Copy full SHA for ec5900d

• minor

  ShashankMosaicML

  committed
  02ad3b6

  Copy full SHA for 02ad3b6

• compiling mask and flex attn once for the entire model

  ShashankMosaicML

  committed
  13a5fc8

  Copy full SHA for 13a5fc8

• ..

  ShashankMosaicML

  committed
  2ae6027

  Copy full SHA for 2ae6027

• ..

  ShashankMosaicML

  committed
  0c5150a

  Copy full SHA for 0c5150a

• making sequence id transforms configurable

  ShashankMosaicML

  committed
  ff28304

  Copy full SHA for ff28304

• ..

  ShashankMosaicML

  committed
  23ba20f

  Copy full SHA for 23ba20f

• ..

  ShashankMosaicML

  committed
  72c45ae

  Copy full SHA for 72c45ae

• ..

  ShashankMosaicML

  committed
  73066a4

  Copy full SHA for 73066a4

Commits on Dec 1, 2024

• ..

  ShashankMosaicML

  committed
  9f616f7

  Copy full SHA for 9f616f7

• converting mods from dict to list

  ShashankMosaicML

  committed
  94ecade

  Copy full SHA for 94ecade

Commits on Dec 2, 2024

• switching off seq id masking if configured so

  ShashankMosaicML

  committed
  4b30130

  Copy full SHA for 4b30130

• fix bug

  ShashankMosaicML

  committed
  9daf068

  Copy full SHA for 9daf068

• fix bug

  ShashankMosaicML

  committed
  67aa900

  Copy full SHA for 67aa900

• adding global and local window mask

  ShashankMosaicML

  committed
  65a0425

  Copy full SHA for 65a0425

• ..

  ShashankMosaicML

  committed
  3443b69

  Copy full SHA for 3443b69

• fixing test

  ShashankMosaicML

  committed
  f6b3705

  Copy full SHA for f6b3705

• ..

  ShashankMosaicML

  committed
  d5ff138

  Copy full SHA for d5ff138

• flex attn softcap only int

  ShashankMosaicML

  committed
  43cb0d1

  Copy full SHA for 43cb0d1

• ..

  ShashankMosaicML

  committed
  f623a1f

  Copy full SHA for f623a1f

• ..

  ShashankMosaicML

  committed
  0fea56a

  Copy full SHA for 0fea56a

• ..

  ShashankMosaicML

  committed
  eb6e792

  Copy full SHA for eb6e792

• ..

  ShashankMosaicML

  committed
  5852da0

  Copy full SHA for 5852da0

Commits on Dec 3, 2024

• Merge branch 'mosaicml:main' into shashank/flexattention

  ShashankMosaicML

  authored
  615a904

  Copy full SHA for 615a904

Commits on Dec 4, 2024

• simplifying design

  ShashankMosaicML

  committed
  0474027

  Copy full SHA for 0474027

• removing check_seq_id_attn_mask

  ShashankMosaicML

  committed
  70aa0c7

  Copy full SHA for 70aa0c7

• ..

  ShashankMosaicML

  committed
  fc8a120

  Copy full SHA for fc8a120

• ..

  ShashankMosaicML

  committed
  5f88093

  Copy full SHA for 5f88093

• fixing tests

  ShashankMosaicML

  committed
  661f7f6

  Copy full SHA for 661f7f6

• ..

  ShashankMosaicML

  committed
  fef3a5d

  Copy full SHA for fef3a5d

• ..

  ShashankMosaicML

  committed
  f6c66e8

  Copy full SHA for f6c66e8

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  eacca42

  Copy full SHA for eacca42

• allowing block overrides for flex attention

  ShashankMosaicML

  committed
  4385f18

  Copy full SHA for 4385f18

• ..

  ShashankMosaicML

  committed
  e17d1ff

  Copy full SHA for e17d1ff

• configuring tests, fixing bugs

  ShashankMosaicML

  committed
  58760fc

  Copy full SHA for 58760fc

• fixing bug when using past kv caches

  ShashankMosaicML

  committed
  f4ad493

  Copy full SHA for f4ad493

• bug fix

  ShashankMosaicML

  committed
  67f9aae

  Copy full SHA for 67f9aae

• ..

  ShashankMosaicML

  committed
  5fcbc18

  Copy full SHA for 5fcbc18

• ..

  ShashankMosaicML

  committed
  8dfdedb

  Copy full SHA for 8dfdedb

Commits on Dec 5, 2024

• fixing score mod bug

  ShashankMosaicML

  committed
  8912cb2

  Copy full SHA for 8912cb2

• ..

  ShashankMosaicML

  committed
  18c4bb9

  Copy full SHA for 18c4bb9

• ..

  ShashankMosaicML

  committed
  bf1cb6c

  Copy full SHA for bf1cb6c

• ..

  ShashankMosaicML

  committed
  5093efd

  Copy full SHA for 5093efd

• ..

  ShashankMosaicML

  committed
  96b8f82

  Copy full SHA for 96b8f82

• ..

  ShashankMosaicML

  committed
  f1ad991

  Copy full SHA for f1ad991

• ..

  ShashankMosaicML

  committed
  18afcc5

  Copy full SHA for 18afcc5

• configuring with torch 2.5.1 and 2.6.0.dev

  ShashankMosaicML

  committed
  434aa83

  Copy full SHA for 434aa83

• configuring more tests with torch 2.5.1 and 2.6.0.dev

  ShashankMosaicML

  committed
  216fcb9

  Copy full SHA for 216fcb9

• ..

  ShashankMosaicML

  committed
  438e0f3

  Copy full SHA for 438e0f3

Commits on Dec 6, 2024

• ..

  ShashankMosaicML

  committed
  2bb25ee

  Copy full SHA for 2bb25ee

• ..

  ShashankMosaicML

  committed
  9831b5e

  Copy full SHA for 9831b5e

• ..

  ShashankMosaicML

  committed
  ad601e4

  Copy full SHA for ad601e4

• ..

  ShashankMosaicML

  committed
  77115c5

  Copy full SHA for 77115c5

Commits on Dec 8, 2024

• figuring out d_model and seq lengths for which flex attention works

  ShashankMosaicML

  committed
  dfde51b

  Copy full SHA for dfde51b

• adding todos

  ShashankMosaicML

  committed
  d1d04ce

  Copy full SHA for d1d04ce

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  5eca05f

  Copy full SHA for 5eca05f

Commits on Dec 9, 2024

• adding test for local global attention

  ShashankMosaicML

  committed
  718d89d

  Copy full SHA for 718d89d

Commits on Dec 17, 2024

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  135abd7

  Copy full SHA for 135abd7

Commits on Jan 7, 2025

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  369e818

  Copy full SHA for 369e818

Commits on Jan 23, 2025

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  8a62ca4

  Copy full SHA for 8a62ca4