[PyT] Plumbing correct bias dims from TE to cudnn, while adding support for additional bias shapes#2537

Merged

KshitijLakhani merged 25 commits into

NVIDIA:mainfrom

KshitijLakhani:klakhani/fix/bias-shape

Feb 18, 2026

Commits on Feb 6, 2026

Plumbing correct bias dims from TE to cudnn
KshitijLakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
KshitijLakhani
committed
Make changes for cp bias code
KshitijLakhani
committed
Add dbias and dbias_ to run_dpa_with_cp test
KshitijLakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
KshitijLakhani
committed
Fix: Use output_dBias instead of input_dBias to extract the shape
KshitijLakhani
committed
Add guards for bias/bias_/dbias/dbias_ being None
KshitijLakhani
committed
Add support for bias shape 111s in addition to the original 1hss, 11ss, b1ss and bhss
KshitijLakhani
committed
Add support for dbias calculation and variant packing for the dbias shapes b1ss, bhss, 11ss in addition to the already supported 1hss
KshitijLakhani
committed
Add support for 111s bias shape in DPA
KshitijLakhani
committed
Allow fused attn for dbias calculation for 11ss, b1ss, bhss. Disable fused attn if dbias calculation for 111s is required, else enable
KshitijLakhani
committed
Disable requires_grad for bias for shape 111s in tests
KshitijLakhani
committed
Disable bias grad / training flag for 111s bias in the non-CP attn tests. Add bias shape 111s to test_dpa_bias_shapes
KshitijLakhani
committed
Fix to correctly create the bias shape tensor instead of the hard coded shape. Fix the comparison logic shapes for bias/dbias
KshitijLakhani
committed
Add fused attn cp test cases for all supported bias shapes
KshitijLakhani
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[PyT] Plumbing correct bias dims from TE to cudnn, while adding support for additional bias shapes#2537

[PyT] Plumbing correct bias dims from TE to cudnn, while adding support for additional bias shapes#2537
KshitijLakhani merged 25 commits into
NVIDIA:mainfrom
KshitijLakhani:klakhani/fix/bias-shape

Commits on Feb 6, 2026

Plumbing correct bias dims from TE to cudnn

[pre-commit.ci] auto fixes from pre-commit.com hooks

Make changes for cp bias code

Add dbias and dbias_ to run_dpa_with_cp test

[pre-commit.ci] auto fixes from pre-commit.com hooks

Fix: Use output_dBias instead of input_dBias to extract the shape

Add guards for bias/bias_/dbias/dbias_ being None

Add support for bias shape 111s in addition to the original 1hss, 11ss, b1ss and bhss

Add support for dbias calculation and variant packing for the dbias shapes b1ss, bhss, 11ss in addition to the already supported 1hss

Add support for 111s bias shape in DPA

Allow fused attn for dbias calculation for 11ss, b1ss, bhss. Disable fused attn if dbias calculation for 111s is required, else enable

Disable requires_grad for bias for shape 111s in tests

Disable bias grad / training flag for 111s bias in the non-CP attn tests. Add bias shape 111s to test_dpa_bias_shapes

Fix to correctly create the bias shape tensor instead of the hard coded shape. Fix the comparison logic shapes for bias/dbias

Add fused attn cp test cases for all supported bias shapes

[pre-commit.ci] auto fixes from pre-commit.com hooks

Commits on Feb 13, 2026

nit: switch to elif for bias grad conditional

Add CP support for bias/dbias shape 111s

Add support for is_training in CP attn tests

[pre-commit.ci] auto fixes from pre-commit.com hooks

Commits on Feb 17, 2026

nit: Fix incorrect comment

nit: Fix incorrect comment and assert string

Commits on Feb 18, 2026

Create the dbias graph tensor only if it is a cuDNN supported bias shape

Fix the dim that is being compared for the two cp chunks in the test

nit: Reinstate the original test for right side swa