Code

[JAX] THD ring attention#1454

Code

Merged

zlsh80826 merged 12 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom zlsh80826:rewang/thd-ring-attnzlsh80826/TransformerEngine:rewang/thd-ring-attnCopy head branch name to clipboardMar 3, 2025

Commits

Commits on Feb 24, 2025

• Support THD + ring attention for self attn

  Show description for b9cb819
  zlsh80826

  committed
  b9cb819

  Copy full SHA for b9cb819

• Consolidate reorder strategy

  Show description for 5eb9c37
  zlsh80826

  committed
  5eb9c37

  Copy full SHA for 5eb9c37

• Fix dataclass frozen issue

  Show description for 748b59c
  zlsh80826

  committed
  748b59c

  Copy full SHA for 748b59c

• Remove redundant code

  Show description for bab7237
  zlsh80826

  committed
  bab7237

  Copy full SHA for bab7237

• Use AttnBiasType, AttnMaskType, QKVLayout in cpp_extension

  Show description for 71293bf
  zlsh80826

  committed
  71293bf

  Copy full SHA for 71293bf

• Fix lint

  Show description for b822315
  zlsh80826

  committed
  b822315

  Copy full SHA for b822315

• Refine P2P helper check_supported

  Show description for f6721f4
  zlsh80826

  committed
  f6721f4

  Copy full SHA for f6721f4

• Add segment_ids/pos check

  Show description for 889500a
  zlsh80826

  committed
  889500a

  Copy full SHA for 889500a

• Fixup

  Show description for 312b95f
  zlsh80826

  committed
  312b95f

  Copy full SHA for 312b95f

• Add dual chunk swap example

  Show description for fc2ebcb
  zlsh80826

  committed
  fc2ebcb

  Copy full SHA for fc2ebcb

• Align different reorder code structure

  Show description for 2da1075
  zlsh80826

  committed
  2da1075

  Copy full SHA for 2da1075

Commits on Feb 27, 2025

• Merge branch 'main' into rewang/thd-ring-attn

  phu0ngng

  authored
  c248f53

  Copy full SHA for c248f53