[JAX] [PyT] Tighten SWA checks in DPA, MHA and other APIs before passing onto cuDNN fused attn & unfused attn backends#2970

Open

KshitijLakhani wants to merge 8 commits into

NVIDIA:mainfrom

KshitijLakhani:klakhani/feat/attn-swa-enh-jax

Commits on May 7, 2026

Add check_set_window_size to JAX attn to mimic the behavior of Pyt attn
KshitijLakhani
committed
TransformerLayer APi does not need to check the window size. That contract responsibility can be handled by MHA and lower APIs
KshitijLakhani
committed

Commits on May 8, 2026