Release 2.15

[VeeraRajasekhar]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[Copilot]

Pull request overview

Release cut for TransformerEngine 2.15.0, including updates across PyTorch/JAX attention paths, grouped MLP/linear FP8+MXFP8 behavior, and scale swizzle utilities, plus expanded CI/QA coverage.

Changes:

• Add/extend MXFP8-related attention support (new QKV formats/layouts, softmax offset grad plumbing, CP updates) across Python and C++/CUDA backends.
• Refine grouped linear/MLP behavior (env-gated single-param mode, cuDNN frontend gating, updated quantizer usage, wgrad path updates).
• Add multi-tensor scale swizzle helpers + tests; adjust build/QA scripts for new coverage and release versioning.

Reviewed changes

Copilot reviewed 61 out of 63 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
transformer_engine/pytorch/utils.py	Add env-gated resolver for grouped-linear single-param flags (with warnings).
transformer_engine/pytorch/triton/init.py	Export mhc from the Triton wrapper package.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Adjust scale shape/offset computations for quantized grouped tensors (incl. block scaling).
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Update fused grouped-MLP forward gating and cuDNN kwargs handling.
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py	Update fused grouped-MLP backward gating, wgrad kernel selection, and tensor packing edge cases.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Gate single grouped weight/bias via env var; refine weight quantizer usage + recipe-state reset behavior.
transformer_engine/pytorch/ops/basic/basic_linear.py	Refine weight quantizer usage + recipe-state reset behavior for BasicLinear.
transformer_engine/pytorch/ops/_common.py	Consolidate cuDNN frontend version gate for fused grouped-MLP support.
transformer_engine/pytorch/newton_schulz.py	Change coefficient API to structured tuples and flatten only at C++ call boundary.
transformer_engine/pytorch/module/grouped_linear.py	Gate single grouped weight/bias via env var in module wrapper.
transformer_engine/pytorch/csrc/util.h	Add unchecked multi-tensor swizzle API declaration.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Implement unchecked swizzle path and add pybind-facing multi-tensor swizzle helpers; improve grouped swizzle padding handling.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Bind multi-tensor swizzle + fused-attn tensor utilities into the PyTorch extension.
transformer_engine/pytorch/csrc/extensions.h	Extend fused attention signatures and declare multi-tensor transpose/pad APIs.
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Extend fused-attn Python wrapper API with output/scale formats and richer backward outputs.
transformer_engine/pytorch/attention/multi_head_attention.py	Adjust FP8 output decisions to account for MXFP8 scaling constraints.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Extend DPA recipe routing for MXFP8; update layout detection and CP call wiring.
transformer_engine/jax/csrc/extensions/attention.cpp	Fix ragged vs non-ragged fused-attn tensor shapes and pass new format args through.
transformer_engine/jax/attention.py	Replace O(T²) mask path with O(T·S) per-segment aggregation path for THD seqlen/offset metadata.
transformer_engine/common/util/pybind_helper.h	Add new QKV format/layout enum values to Python bindings.
transformer_engine/common/transformer_engine.cpp	Make scale shape checking optionally skippable in input validation.
transformer_engine/common/include/transformer_engine/swizzle.h	Add unchecked multi-tensor swizzle API for cases with intentionally mismatched scale shapes.
transformer_engine/common/include/transformer_engine/fused_attn.h	Extend fused-attn API with new formats/layouts and add a C++ AttentionShape helper.
transformer_engine/common/fused_attn/utils.h	Add MXFP8 padding helpers, stride generation by format, and extend fused-attn descriptor keys.
transformer_engine/common/fused_attn/utils.cu	Add stride generation support for BHSD layout.
transformer_engine/common/fused_attn/fused_attn.cpp	Plumb new fused-attn format/layout args; extend backend eligibility for BHSD and newer cuDNN versions.
transformer_engine/common/fused_attn/fused_attn_fp8.h	Update FP8 fused-attn forward/backward signatures to match new format/layout plumbing.
transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h	Remove cuDNN version compile-guard wrappers from header.
transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu	Remove cuDNN version compile-guard wrappers from compilation unit.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h	Extend f16 arbitrary-seqlen signatures with output/grad formats and dqkv layout.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Plumb new format/layout args through arbitrary-seqlen kernels and caching keys.
transformer_engine/common/fused_attn_rocm/fused_attn.cpp	Accept (but ignore) new fused-attn format/layout args on ROCm for API parity.
transformer_engine/common/common.h	Add integer vectorized-load macro and update CheckInputTensor signature default arg.
transformer_engine/common/CMakeLists.txt	Move flash_attn.cu into arch-specific CUDA source list.
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh	Add synchronization before dbias reuse path.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Fix scale indexing for swizzled scales + zero OOB padding; add sync before dbias reuse path.
tests/pytorch/utils.py	Handle None outputs in compare helper.
tests/pytorch/distributed/run_newton_schulz.py	Update reference to new coefficient tuple format.
tests/pytorch/attention/test_attention.py	Expand attention test matrix (new dims, new scaling mode, softmax offset grad, new formats).
tests/pytorch/attention/test_attention_with_cp.py	Expand CP attention test coverage (new configs, determinism plumbing, MXFP8 scaling mode).
tests/pytorch/attention/run_attention_with_cp.py	Add MXFP8 quantizers + determinism env plumbing; improve NaN/Inf asserts.
tests/jax/test_custom_call_compute.py	Make grouped-dense bias grad reference accumulate in fp32 to match primitive.
tests/cpp/operator/test_swizzle.cu	Add compact grouped-scale swizzle test and broaden narrow-path coverage.
tests/cpp/operator/test_multi_swizzle.cu	Add new C++ tests for multi-tensor swizzle/unswizzle/roundtrip.
tests/cpp/operator/CMakeLists.txt	Include new multi-swizzle test; disable it on ROCm for now.
qa/L0_pytorch_unittest/test.sh	Enable grouped-linear single-param mode for relevant unit tests; add MHC unittest invocation.
qa/L0_pytorch_lint/test.sh	Add a lint QA job for cpplint + pylint.
qa/L0_pytorch_debug_unittest/test.sh	Enable grouped-linear single-param mode for debug unittests.
docs/api/pytorch.rst	Add API docs entries for Triton MHC fused functions.
ci/pytorch.sh	Enable grouped-linear single-param mode for fusible-ops and sanity tests in CI.
build_tools/VERSION.txt	Bump version from 2.15.0.dev0 to 2.15.0.

Comments suppressed due to low confidence (1)

transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py:317

• grouped_gemm_wgrad_kernel() is annotated/used as an optional kernel selector, but it can currently raise ImportError when grouped_gemm_wgrad_wrapper_sm100 is unavailable. Since callers pass cudnn_wgrad_kernel_fn=self.grouped_gemm_wgrad_kernel() unconditionally, this can crash at runtime even though a safe fallback path exists (i.e., returning None so the non-cuDNN wgrad path is used). Catch the import failure and return None instead.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

Claude review — release_v2.15_rocm sync (PR #656)

Scope reviewed: 3 ROCm-specific commits on top of ~19 upstream v2.15 commits:

• 8d83ffb5 — resolve v2.15 merge conflicts + port fused_attn_rocm/fused_attn.cpp nvte_fused_attn_{fwd,bwd} to the new format-aware upstream C API.
• b20f0145 — guard the new upstream TMA transpose path (flash_attn.cu) behind #ifndef __HIP_PLATFORM_AMD__ and move get_max_dynamic_smem() out of the CUDA-only guard in swizzle.cu so ROCm-compiled multi-tensor swizzle can link.
• 08537961 — guard new upstream FP8/MXFP8 arch heuristics (utils.py) with IS_HIP_EXTENSION, skip test_multi_swizzle.cu on ROCm, and gate test_fusible_ops.py/test_sanity.py with NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 to match upstream L0.

Verdict: The ROCm-specific commits look well-scoped and correctly documented. The nvte_fused_attn_{fwd,bwd} port accepts the new o_format/do_format/dqkv_layout/*_scale_inv_format parameters with clear (void) casts and a rationale comment (ROCm CK/AOTriton reject FP8/MXFP8 in is_{ck,aotriton}_backend_supported, so the extra format hints are moot). Signatures in common/include/transformer_engine/fused_attn.h and common/fused_attn_rocm/fused_attn.cpp match exactly. TMA guards in flash_attn.cu cover the struct, the entire TMA path (constants, kernel, create_strided_tensor_map), and the dispatch site in multi_tensor_transpose_to_bhsd. The NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 placements match qa/L0_pytorch_unittest/test.sh.

Duplicating prior Copilot findings not repeated here: (1) eager mhc import in transformer_engine/pytorch/triton/__init__.py (breaks environments without triton), (2) backward_grouped_mlp.py:317 ImportError fallthrough.

Copyright headers: 1 file needs an update — see inline comment.