[JAX] Make all jax attention calls use non-packed common calls (#2358)

* fix
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* add notes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

for more information, see https://pre-commit.ci

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

for more information, see https://pre-commit.ci



* small fixes
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

transformer_engine/common/fused_attn/fused_attn.cpp

@@ -29,7 +29,7 @@ transformer_engine::Tensor make_tensor_view(const transformer_engine::Tensor *so
   return view;
 }
 
-// Helper function to calculate stride for packed QKV tensor unpacking
+// Helper function to calculate stride in bytes for packed QKV tensor unpacking
 size_t calculate_qkv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype,
                             size_t h, size_t d) {
   size_t stride = 0;