Paged attention changes to THD attention #3

Draft

sudhakarsingh27 wants to merge 259 commits into

sudhakarsingh27:te_gemma_generation_supportfrom

cyanguwa:paged_attention

This pull request is big! We're only showing the most recent 250 commits

Commits on Nov 22, 2024

[Core] Add function to convert container to string (#1342)

timmoon10
and
pre-commit-ci[bot]
authored

Commits on Nov 25, 2024

Commits on Nov 27, 2024

Fix cuda graph capture for grouped gemm (#1345)
xrennvidia
authored

Commits on Dec 2, 2024

Update list of CI users (#1340)
timmoon10
authored
Improving communication overlap for the case of multi kernel queue usage (#1308)

youngeunkwon0405
and
pre-commit-ci[bot]
authored

Commits on Dec 4, 2024

Commits on Dec 5, 2024

Commits on Dec 6, 2024

Commits on Dec 10, 2024

[JAX] Use default factory for not sharing mutable default values (#1364)

zlsh80826
and
phu0ngng
authored

Commits on Dec 12, 2024

Commits on Dec 14, 2024

Commits on Dec 16, 2024

Enabling FP8 all-gather for TE Float8Tensor when using Torch FSDP2 (#1358)

youngeunkwon0405
and
pre-commit-ci[bot]
authored

Commits on Dec 17, 2024

Commits on Dec 18, 2024

Commits on Dec 20, 2024

[common/PyTorch] Add cuDNN SWA (left, 0) + padding + bottom right causal (#1378)

cyanguwa
and
pre-commit-ci[bot]
authored

Commits on Jan 2, 2025

Update copyright to include 2025 (#1388)
ksivaman
authored

Commits on Jan 6, 2025

Commits on Jan 7, 2025

fix build issue from previous merge
cyanguwa
committed

Commits on Jan 10, 2025

Take token count quantization of fused attention into consideration for CP results correction (#1396)

xrennvidia
and
pre-commit-ci[bot]
authored

Commits on Jan 11, 2025

[JAX] Test_multiprocessing_encoder with process spawn in bash (#1394)
phu0ngng
authored

Commits on Jan 13, 2025

Commits on Jan 15, 2025

Commits on Jan 16, 2025

[PyTorch] `te.Linear` FP8 DGRAD+RS output bugfix (#1412)

denera
and
pre-commit-ci[bot]
authored

Commits on Jan 17, 2025

[JAX] Consolidate the distributed fused attention test code (#1405)
mgoldfarb-nvidia
authored

Commits on Jan 21, 2025

[PyTorch] Fix AttentionParams comparison logic (#1397)
cyanguwa
authored

Commits on Jan 22, 2025

[PyTorch] Avoid `parameters` function in op backward pass (#1403)

timmoon10
and
pre-commit-ci[bot]
authored

Commits on Jan 24, 2025

[JAX] Support segment_ids/pos as FA inputs (#1406)
zlsh80826
authored

Commits on Jan 27, 2025

[MoE][PyTorch] Add mask-based MoE permutation (#1373)

authored

Commits on Jan 28, 2025

Commits on Jan 30, 2025

Commits on Jan 31, 2025

Support `store_param_remainders` feature from Apex in TE Fused Adam (#1408)

authored

Commits on Feb 7, 2025

Update main branch with TE 2.0 code, update version to 2.1.0.dev0
ptrendx
committed

Commits on Feb 8, 2025

Commits on Feb 10, 2025

WIP: minor fix/preparation for inference/cuda graph
cyanguwa
committed

Commits on Feb 11, 2025

[JAX] Flax module init with a given dtype (#1472)
phu0ngng
authored

Commits on Feb 13, 2025

Fix a bug for D being nullptr in grouped gemm (#1475)

authored

Commits on Feb 14, 2025

Commits on Feb 16, 2025

WIP: flash-attn, non-CG
cyanguwa
committed

Commits on Feb 18, 2025

Commits on Feb 20, 2025

Commits on Feb 21, 2025

WIP: thd_bshd_bshd
cyanguwa
committed

Commits on Feb 22, 2025

Commits on Feb 23, 2025

Commits on Feb 26, 2025

Commits on Mar 1, 2025

Commits on Mar 3, 2025

Commits on Mar 4, 2025

Commits on Mar 6, 2025

Commits on Mar 7, 2025

Commits on Mar 8, 2025

[PyTorch] Enabling Per-Tensor Current Scaling Recipe (#1471)

authored

Commits on Mar 9, 2025

Verified TE2.0 with offloading (#1514)

sanandaraj5597
and
ksivaman
authored

Commits on Mar 12, 2025

Commits on Mar 13, 2025

Commits on Mar 14, 2025

Commits on Mar 15, 2025

Commits on Mar 17, 2025