Code

Add support for Flex Attention#1675

Code

Draft

ShashankMosaicML wants to merge 91 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

• registrifying score mods

  ShashankMosaicML

  committed
  ac3a884

• registrifying attention mask mods

  ShashankMosaicML

  committed
  31b27e2

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

  ShashankMosaicML

  authored
  c8fffa5

Commits on Nov 19, 2024

• bug_fix

  ShashankMosaicML

  committed
  86dce3b

• bug_fix

  ShashankMosaicML

  committed
  cb8f4a6

• lint

  ShashankMosaicML

  committed
  902850a

• configuring test

  ShashankMosaicML

  committed
  9c9708d

• configuring tests

  ShashankMosaicML

  committed
  f1ff430

• bug fix

  ShashankMosaicML

  committed
  e537f5a

• fixing alibi

  ShashankMosaicML

  committed
  c527dd7

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

  ShashankMosaicML

  authored
  15e303e

• configuring further tests

  ShashankMosaicML

  committed
  c4ef5d9

• refactoring

  ShashankMosaicML

  committed
  6b37427

• adding warnings and errors

  ShashankMosaicML

  committed
  e30fe7a

• gating tests on torch version

  ShashankMosaicML

  committed
  924a53c

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

  ShashankMosaicML

  authored
  57048e3

• reorganizing function defs

  ShashankMosaicML

  committed
  67a2aea

• refactoring

  ShashankMosaicML

  committed
  04f3a62

• passing in dicts of mask and score mods

  ShashankMosaicML

  committed
  ab6c58c

• making mask and score mods configurable via yaml

  ShashankMosaicML

  committed
  3b3827d

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

  ShashankMosaicML

  authored
  be43e8d

Commits on Nov 20, 2024

• adding torch.compile

  ShashankMosaicML

  committed
  2264f91

• ..

  ShashankMosaicML

  committed
  e274d9f

• ..

  ShashankMosaicML

  committed
  a26bb4f

• undoing comment out

  ShashankMosaicML

  committed
  d5ab7d3

Commits on Nov 26, 2024

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

  ShashankMosaicML

  authored
  d40e978

• adding torch comile

  ShashankMosaicML

  committed
  5f13e7b

• temporary commit commenting out block mask and score mod

  ShashankMosaicML

  committed
  ca8e173

• undoing prev temp commit

  ShashankMosaicML

  committed
  f5486ff

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

  ShashankMosaicML

  authored
  fdced3a

Commits on Nov 27, 2024

• speeding up block mask generation

  ShashankMosaicML

  committed
  c53db63

• precompilining create block mask

  ShashankMosaicML

  committed
  ec5900d

• minor

  ShashankMosaicML

  committed
  02ad3b6

• compiling mask and flex attn once for the entire model

  ShashankMosaicML

  committed
  13a5fc8

• ..

  ShashankMosaicML

  committed
  2ae6027

• ..

  ShashankMosaicML

  committed
  0c5150a

• making sequence id transforms configurable

  ShashankMosaicML

  committed
  ff28304

• ..

  ShashankMosaicML

  committed
  23ba20f

• ..

  ShashankMosaicML

  committed
  72c45ae

• ..

  ShashankMosaicML

  committed
  73066a4

Commits on Dec 1, 2024

• ..

  ShashankMosaicML

  committed
  9f616f7

• converting mods from dict to list

  ShashankMosaicML

  committed
  94ecade

Commits on Dec 2, 2024

• switching off seq id masking if configured so

  ShashankMosaicML

  committed
  4b30130

• fix bug

  ShashankMosaicML

  committed
  9daf068

• fix bug

  ShashankMosaicML

  committed
  67aa900

• adding global and local window mask

  ShashankMosaicML

  committed
  65a0425

• ..

  ShashankMosaicML

  committed
  3443b69

• fixing test

  ShashankMosaicML

  committed
  f6b3705

• ..

  ShashankMosaicML

  committed
  d5ff138

• flex attn softcap only int

  ShashankMosaicML

  committed
  43cb0d1

• ..

  ShashankMosaicML

  committed
  f623a1f

• ..

  ShashankMosaicML

  committed
  0fea56a

• ..

  ShashankMosaicML

  committed
  eb6e792

• ..

  ShashankMosaicML

  committed
  5852da0

Commits on Dec 3, 2024

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

  ShashankMosaicML

  authored
  615a904

Commits on Dec 4, 2024

• simplifying design

  ShashankMosaicML

  committed
  0474027

• removing check_seq_id_attn_mask

  ShashankMosaicML

  committed
  70aa0c7

• ..

  ShashankMosaicML

  committed
  fc8a120

• ..

  ShashankMosaicML

  committed
  5f88093

• fixing tests

  ShashankMosaicML

  committed
  661f7f6

• ..

  ShashankMosaicML

  committed
  fef3a5d

• ..

  ShashankMosaicML

  committed
  f6c66e8

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  eacca42

• allowing block overrides for flex attention

  ShashankMosaicML

  committed
  4385f18

• ..

  ShashankMosaicML

  committed
  e17d1ff

• configuring tests, fixing bugs

  ShashankMosaicML

  committed
  58760fc

• fixing bug when using past kv caches

  ShashankMosaicML

  committed
  f4ad493

• bug fix

  ShashankMosaicML

  committed
  67f9aae

• ..

  ShashankMosaicML

  committed
  5fcbc18

• ..

  ShashankMosaicML

  committed
  8dfdedb

Commits on Dec 5, 2024

• fixing score mod bug

  ShashankMosaicML

  committed
  8912cb2

• ..

  ShashankMosaicML

  committed
  18c4bb9

• ..

  ShashankMosaicML

  committed
  bf1cb6c

• ..

  ShashankMosaicML

  committed
  5093efd

• ..

  ShashankMosaicML

  committed
  96b8f82

• ..

  ShashankMosaicML

  committed
  f1ad991

• ..

  ShashankMosaicML

  committed
  18afcc5

• configuring with torch 2.5.1 and 2.6.0.dev

  ShashankMosaicML

  committed
  434aa83

• configuring more tests with torch 2.5.1 and 2.6.0.dev

  ShashankMosaicML

  committed
  216fcb9

• ..

  ShashankMosaicML

  committed
  438e0f3

Commits on Dec 6, 2024

• ..

  ShashankMosaicML

  committed
  2bb25ee

• ..

  ShashankMosaicML

  committed
  9831b5e

• ..

  ShashankMosaicML

  committed
  ad601e4

• ..

  ShashankMosaicML

  committed
  77115c5

Commits on Dec 8, 2024

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

  ShashankMosaicML

  committed
  dfde51b

• adding todos

  ShashankMosaicML

  committed
  d1d04ce

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  5eca05f

Commits on Dec 9, 2024

• adding test for local global attention

  ShashankMosaicML

  committed
  718d89d

Commits on Dec 17, 2024

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  135abd7

Commits on Jan 7, 2025

• Merge branch 'main' into shashank/flexattention

  ShashankMosaicML

  authored
  369e818