Code

[JAX][Common] Support GQA#578

Code

Merged

cyanguwa merged 14 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom zlsh80826:rewang/gqa-cleanzlsh80826/TransformerEngine:rewang/gqa-cleanCopy head branch name to clipboardJan 16, 2024

+613-718

Commits

Commits on Jan 14, 2024

• Support num_gqa_groups arguments

  Show description for 2cb5a08
  zlsh80826

  committed
  2cb5a08

  Copy full SHA for 2cb5a08

• Add GQA support on the JAX bridge code

  Show description for 9adee47
  zlsh80826

  committed
  9adee47

  Copy full SHA for 9adee47

• Fix the kv stride of the arbitrary backend

  Show description for 2e0bc84
  zlsh80826

  committed
  2e0bc84

  Copy full SHA for 2e0bc84

• Complete rewrite fused attention tests and add GQA coverage

  Show description for 53e7b65
  zlsh80826

  committed
  53e7b65

  Copy full SHA for 53e7b65

• Support unfused GQA

  Show description for 5b76943
  zlsh80826

  committed
  5b76943

  Copy full SHA for 5b76943

• Calculate seqlen before the primitive for the better perf

  Show description for ae6d209
  zlsh80826

  committed
  ae6d209

  Copy full SHA for ae6d209

• Add GQA layer tests

  Show description for 8d5d407
  zlsh80826

  committed
  8d5d407

  Copy full SHA for 8d5d407

• Apply code style checks for te_jax

  Show description for 3d0c38e
  zlsh80826

  committed
  3d0c38e

  Copy full SHA for 3d0c38e

• Apply code style checks for tests

  Show description for db2cc1f
  zlsh80826

  committed
  db2cc1f

  Copy full SHA for db2cc1f

• Add num_gqa_groups doc

  Show description for 21ed876
  zlsh80826

  committed
  21ed876

  Copy full SHA for 21ed876

• Refine the qkv_type

  Show description for 844eb42
  zlsh80826

  committed
  844eb42

  Copy full SHA for 844eb42

• Correct the variable naming

  Show description for 65df768
  zlsh80826

  committed
  65df768

  Copy full SHA for 65df768

• Handle Max512 CAUSAL

  Show description for bbe7066
  zlsh80826

  committed
  bbe7066

  Copy full SHA for bbe7066

• Add WAR for the latest jax image

  Show description for edb474f
  zlsh80826

  committed
  edb474f

  Copy full SHA for edb474f