Ifu dev 20260828 v2.19 by matthiasdiener · Pull Request #721 · ROCm/TransformerEngine

matthiasdiener · 2026-08-28T21:01:06Z

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

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

…opk_weight tensor (#3187) * expose user-provided weights * adding pool based symm allocation; remove the persistent buffer in EpBuffer * add zero copy tests Signed-off-by: YangFei1990 <feiw@nvidia.com> --------- Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Migrate NCCL EP submodule to NVIDIA/nccl-extensions * Drop PYTHONPATH override from EP test, example, and bench launchers * Drop cross-mode recv comparison in EP zero-copy IdentityAllSymm test Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * [Common] Rename 3rdparty/nccl submodule directory to nccl-extensions Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

… (#3171) * [Common/PyTorch] Support power-of-2 scales in grouped FP8 block-scaling quantize The default Float8BlockScaling recipe constrains scales to powers of 2, so the fused grouped path must honor the flag to stay numerically consistent with the unfused path. Thread a runtime pow_2_scales argument through the grouped quantize kernels (the shared scale helper already implements the rounding) and drop the force_pow_2_scales rejections. Also add a quantization-config parameter to nvte_group_quantize_dbias, which previously had no way to receive force_pow_2_scales or amax_epsilon on the bgrad path. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Enable fused grouped FP8 block-scaling path in GroupedLinear module Admit Float8BlockQuantizer in the fused GroupedTensor path on Hopper. The existing usage flags already match the Hopper TN-only mapping and the grouped GEMM selects transposed columnwise storage for NN/NT layouts, so only the path predicate changes. The fused path is an explicit opt-in via NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM, so raise on Blackwell (SM100/SM110) instead of silently falling back; the fused path has no MXFP8-broadcast emulation. Extend the fused dbias path (tex.bgrad_group_quantize) to FP8 block scaling when dgrad is required (dbias is computed in the rowwise pass). Add fp8_block_scaling to the fused-path tests with a Hopper-only gate, assert the fused path engages via a group_quantize spy, and add a Blackwell error-path test. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Enable FP8 block-scaling in GroupedLinear fusible op Replace the blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state with a per-op supports_float8_block_scaling flag and opt in the GroupedLinear op. Mirror the module-path predicate and fused-bgrad changes; since the graph-safe flow is default-on here (no env-var opt-in), other architectures fall back to the split-quantize flow instead of raising. Force use_split_accumulator=True for FP8 block-scaling operands in general_grouped_gemm_for_grouped_tensor, matching non-grouped general_gemm: cuBLAS has no fast-accum FP8 block-scaling algorithm, so the ops-layer forward failed algo selection without it. Add fp8_block_scaling coverage to the ops GroupedLinear tests. The CUDA-graph-safe test skips it for now: the replayed wgrad for the last expert diverges between replays depending on process allocation history; under investigation. Graph capture remains covered by the module-path test. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Use persistent workspaces in grouped-tensor GEMM general_grouped_gemm_for_grouped_tensor allocated its setup workspace (the cuBLAS per-group pointer/dimension arrays) and its cuBLAS workspace with per-call torch.empty. Under make_graphed_callables the forward and backward graphs share one capture memory pool, and a per-call allocation's block returns to that pool as soon as the Python reference dies, so blocks alias across the two graphs and captured kernels from one graph overwrite the GEMM metadata the other graph reads at replay. Observed as allocation-history-dependent failures in the ops-layer GroupedLinear cuda-graph test: capture-time cublasLtMatmulAlgoGetHeuristic NOT_SUPPORTED errors and corrupted wgrad outputs. This is also the likely mechanism behind the FP8 block-scaling wgrad corruption under CUDA graphs previously observed on Hopper and attributed to cuBLAS. Cache the setup workspace per (device, group size) and reuse the cached per-device cuBLAS workspace from the non-grouped path; consecutive GEMMs reusing one workspace are ordered by the stream. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Fix grouped FP8 block-scaling CUDA-graph deadlock via per-role cuBLAS workspaces The grouped-tensor GEMM path shared one persistent cuBLAS workspace across all grouped matmuls. cuBLAS's grouped GEMM keeps a grid-synchronization flag in the first bytes of that workspace and zeros it (via a captured memset) before each matmul. When the dgrad and wgrad grouped matmuls of a GroupedLinear backward share one workspace inside a replayed CUDA graph, that flag is aliased between the two matmuls; on the second graph replay the second matmul's cooperative kernel deadlocks with cuBLAS 13.6 (and corrupts the last expert's wgrad on cuBLAS < 13.6). The two matmuls are strictly stream-ordered (single stream, all-DEFAULT graph edges, no programmatic dependent launch), so this is shared-workspace reuse, not concurrent co-scheduling. Give dgrad/forward (slot 0) and wgrad (slot 1) distinct persistent cuBLAS workspaces, dedicated to the grouped path. Each slot remains a single persistent allocation, so CUDA-graph capture safety is preserved. Also drop the cuBLAS-version gate that skipped the FP8 block-scaling GroupedLinear CUDA-graph test, so it now exercises the fix on all supported cuBLAS versions. Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: document split-accumulator override, fix stale dbias comment - general_grouped_gemm_for_grouped_tensor: expand the comment to state that the fused grouped FP8 block-scaling GEMM forces use_split_accumulator=True and intentionally overrides the caller-supplied value, consistent with the Float8BlockScaling recipe (which fixes it True for fprop/dgrad/wgrad). - Float8BlockScaling recipe docstring: document that FP8 block scaling always uses split accumulation and that the fused grouped GEMM path ignores any caller- or recipe-supplied use_split_accumulator value. - GroupedLinear ops backward: correct the stale "BF16/FP16 path" comment; that branch also handles quantized paths where bgrad fusion did not apply (e.g. FP8 block scaling without a dgrad pass). Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Revert fusible-ops FP8 block-scaling; scope PR to GroupedLinear module Restrict this PR to the GroupedLinear module fused-quantize path. Revert the fusible-ops FP8 block-scaling enablement -- the BasicOperation opt-in gate, the GroupedLinear op support, and the fusible-ops test coverage -- back to main. Enabling fusible-ops FP8 block-scaling for both grouped and non-grouped paths is deferred to a separate PR. The blanket FP8 block-scaling rejection in BasicOperation.reset_recipe_state is restored. The split-accumulator guard in general_grouped_gemm_for_grouped_tensor is retained: it is correct for the module's FP8 block-scaling grouped GEMM. Signed-off-by: Alp Dener <adener@nvidia.com> * [PyTorch] Isolate grouped wgrad cuBLAS workspace by NT layout, not out-discreteness _get_grouped_cublas_workspace slots were keyed on is_discrete_out as a proxy for "this is the wgrad GEMM", which only holds when wgrad writes a list of per-expert grads. With single_grouped_weight=True, wgrad writes a single grouped weight-grad (GroupedTensor out, not a list), so is_discrete_out is False and it collided with dgrad on slot 0 -- reintroducing the FP8 block-scaling grid-sync-flag aliasing deadlock/corruption under CUDA-graph replay. Key the slot on the wgrad layout (NT / transb) instead: fprop (TN) and dgrad (NN) share slot 0, wgrad (NT) is always isolated on slot 1. Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Address review: isolate grouped cuBLAS workspace per layout; drop redundant test spy - _get_grouped_cublas_workspace now keys the persistent workspace on the grouped GEMM layout, so fprop (TN), dgrad (NN), and wgrad (NT) each get a distinct workspace. The previous NT-vs-rest scheme left fprop and dgrad sharing one workspace; those have also been reported to conflict under CUDA-graph replay. Documents that the deadlock is deterministic and present through cuBLAS 13.7. - Drop the group_quantize call-counting spy in test_grouped_linear_grouped_tensor_path_matches_legacy; fused-path engagement is covered by the graph-safe test. Signed-off-by: Alp Dener <adener@nvidia.com> * updated grouped GEMM workspace comment on stale TMA descriptor related deadlocks Signed-off-by: Alp Dener <adener@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

update nccl-ext submodule name Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* [JAX] Schedule EP dispatch/combine on XLA collective stream * [JAX] Gate EP collective-stream annotation on JAX/XLA version Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* Generalized Tensor Parallelism (GTP) init commit Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * GTP + gmm fusion Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * [fix] Respect per-op activation-offload markers in fused grouped MLP Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean: rename GTP weight-sharding axis to gtp_remat Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "[fix] Respect per-op activation-offload markers in fused grouped MLP" This reverts commit 8bb26f047f690bb8a7e5884a1f9bf419a34fe002. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Make TE GTP-agnostic at construction Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * GTP+nvfp4: fix GTP backward GEMM scaling-mode mismatch for bf16-gathered weights Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Make TE runtime GTP-agnostic via a DistributedWeight protocol Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean - Take a single leader weight in the DistributedWeight dispatchers - Gather the FC2 grouped weight late in the fused grouped MLP Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Simplify the NVFP4 gather post-process; Materialize the EGTP FC1 weight before the NVFP4 dgrad dispatch Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Code clean - Rename gather coalescing flag grouped -> external_coalescing; - Clean up DistributedWeight wiring in TE modules - Restructure _all_gather_nvfp4 Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * fix comments Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Support DistributedWeight in the fusible grouped-linear ops path - Add a self-contained dispatch test with a fake DistributedWeight implementer Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> * Unify distributed-weight wgrad finalize to return a graph-safe dummy - `finalize_weight_grads` now accepts a weight list or a bare leader, mirroring materialize_weight_for_backward; - Centralize the in-place / dummy / async-None finalize contract in DistributedWeight.finalize_group_grads and delegate the dispatcher docstring to it. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> --------- Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Signed-off-by: Tim Moon <tmoon@nvidia.com>

fix nproc Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Avoid invalid make -j 0 in NCCL EP submodule build Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…domain grouping (#3226) * Support axes orthogonal to EP in ep_bootstrap via mesh-derived domain grouping * Fix import-time XLA backend init in EP domain grouping test Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…s. (#3057) * Add token-linear THD fused RoPE path Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * Add THD RoPE full-layer benchmark Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * Cover CP ranks in THD RoPE token-linear tests Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * Address THD RoPE dispatch review feedback Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address THD RoPE linear-grid review feedback Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> * Refine THD RoPE linear-grid dispatch Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> --------- Signed-off-by: plugyawn <progyan.das@iitgn.ac.in> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>

… graph breaks (#3189) * Make get_attention_backend traceable by torch.compile without graph breaks - Read NVTE_* env vars via os.environ.get instead of os.getenv so dynamo installs guards on the values (os.getenv reads are not guarded and would bake stale backend selections into compiled graphs). - Wrap tex.get_fused_attn_backend in a torch.compiler.assume_constant_result helper so the pybind call does not graph-break. - Mark get_device_compute_capability/get_cudnn_version with assume_constant_result for the same reason. - Use a no-op logger when compiling (logging.Logger methods graph-break) and skip debug-log blocks that call int()/str() on pybind enums. - Add test in tests/pytorch/test_torch_compile.py checking fullgraph=True tracing and recompilation on NVTE_* env var changes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Return int from the fused-attn backend probe and test symbolic-int path Comparing the pybind enum returned through assume_constant_result against module-level enum values generates guards dynamo cannot evaluate (crash when the comparison is true, i.e. when cuDNN rejects the config). The wrapper now returns a plain int, comparisons use precomputed int values, and the enum for callers is reconstructed by a second assume_constant_result helper that is never compared during tracing. Also document that os.environ.get (vs os.getenv) is intentional, and drop the guard_scalar specialization of numeric args: symbolic scalars (automatic dynamic) now graph break at the probe instead of forcing a full recompile per seqlen value; the test covers that path without fullgraph and checks the selection stays correct. A second test monkeypatches tex.get_fused_attn_backend to verify the baked result is trace-time-only and actually drives selection. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Replace the FusedAttnBackend dict with a python-side IntEnum The pybind NVTE_Fused_Attn_Backend enum is not traceable by torch.compile: its C-implemented __eq__ cannot be traced, and a pybind enum instance baked through assume_constant_result produces guards dynamo cannot evaluate when compared against module-level enum values. FusedAttnBackend is now a plain python IntEnum generated at import time from tex.NVTE_Fused_Attn_Backend.__members__ (values always in sync with the C enum), and all remaining direct uses of the pybind enum on the python side are replaced with it. Name lookup (FusedAttnBackend["FP8"]) behaves the same as with the previous dict, and the backend value never crosses into a pybind call, so no boundary conversion is needed. This removes the previous int-based workaround in get_attention_backend (_fused_attn_backend_from_int and the precomputed int table): the assume_constant_result wrapper now simply returns the IntEnum and comparisons are traceable directly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Soften the FusedAttnBackend enum transition following the DType pattern Mirror how constants.DType migrated off the pybind enum: explicit IntEnum members pinned to the C values with an import-time sync assert, an __eq__ override comparing by integer value against NVTE_Fused_Attn_Backend (with matching __ne__/__hash__) so mixed comparisons stay equivalent regardless of the pybind11 version, and a cast() classmethod. fused_attn_fwd/bwd normalize their fused_attention_backend argument through cast(), so external callers still passing the pybind enum keep working. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Add FP8 backend-selection coverage to the torch.compile tests test_get_attention_backend_traceable_fp8 compiles the selection with fullgraph=True for AttentionParams(fp8=True) with a DelayedScaling(fp8_dpa) recipe, covering the FP8-only branch (run_config env reads, recipe filters, get_fp8_te_dtype) and checks that flipping NVTE_UnfusedDPA_Emulate_FP8 recompiles and keeps matching eager. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Simplify the fused-attn probe wrapper and drop obsolete logging guards The is_compiling() guards around the available/selected-backend debug logs protected int() on the pybind enum, which crashed dynamo during tracing. With FusedAttnBackend now a python IntEnum, int() on it and str() on the flash-attn PkgVersion both trace cleanly (verified under fullgraph=True), so the logging blocks return to their upstream shape. The probe wrapper also reuses FusedAttnBackend.cast() and a shorter docstring. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Consolidate backend-selection compile tests; keep probe args literal Merge the three get_attention_backend tests into one covering: fullgraph tracing with the probe consulted at trace time only, env var flips (F16 and FP8), attention-param changes, and a forced No_Backend result driving the selection. The bitmask output now also encodes the fused sub-backend, which previously went unchecked. The probe wrapper takes layout/bias/mask/softmax as string keys and resolves the pybind enums internally, so every argument is a literal or a python enum - required for assume_constant_result(specialize_args=True) (pytorch#189042) to derive value guards once available. Scalars must stay concrete until then: the test pins specialize_int/float=True, because a symbolic scalar currently graph breaks at the probe and dynamo's resume then corrupts the returned fused backend (binds the wrapper function object instead of its result; surfaced by the sub-backend bits, minimal repro exists). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drop references to unreleased dynamo features from comments The probe-argument and static-scalar comments referenced assume_constant_result(specialize_args=True), which is not part of any released PyTorch; describe the current behavior only. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Assert selection results only in the compile test Drop the call-counting monkeypatch and its assertions; compiled-vs-eager output equality is what matters and already fails on stale selections. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Address review: add info()/error() no-ops to _NoOpLogger shino16: _NoOpLogger does not subclass logging.Logger, so any traced logger.info()/logger.error() call under torch.compile would raise AttributeError. Add the two no-op methods for completeness. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

[PyTorch] Enable FP8 block scaling across fusible ops Enables FP8 block scaling (FP8BS) in composable fusible-ops for both non-grouped and grouped linears. - Drop the blanket FP8BS NotImplementedError in ops/op.py; FP8BS now uses the generic recipe-state path like MXFP8. - Add `Float8BlockwiseQTensorStorage.view()` so quantized norm outputs reshape without dequantizing (mirrors `MXFP8TensorStorage.view()`). - Route grouped FP8BS through the Hopper graph-safe path (cuBLAS 13.4+). Fuse dbias in the rowwise pass when dgrad is required. - Exercise FP8BS across the fusible-ops test suite with 128-divisible sizes. Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>

…KV inputs (#3200) * [PyTorch] DotProductAttention: declarative packed qkv/kv inputs Fused QKV projections naturally produce one packed buffer, but DotProductAttention forces callers to slice it into q/k/v views that TE then reverse-engineers with pointer-based layout detection (get_qkv_layout inspects data_ptr/storage_offset on every forward, which graph-breaks under torch.compile and adds CPU overhead). Let callers declare the packing instead (JAX-style): * DotProductAttention.forward gains optional qkv_layer (fully packed QKV: [b,s,3,h,d]/[s,b,3,h,d]/[b,s,h,3,d]/[s,b,h,3,d] dense, [t,3,h,d]/ [t,h,3,d] thd), kv_layer (packed KV used with query_layer), and qkv_interleave_dim (-3 or -2; explicit knob rather than shape inference since h==3 or hg==2 would be ambiguous). * Q/K/V are derived as zero-copy select() views and the exact layout enum (bs3hd, bsh3d, sb3hd, bshd_bs2hd, t3hd, ...) is constructed declaratively -- it is truthful by construction, so get_qkv_layout is never called on this path, including for thd and FP8 DPA. * combine_and_quantize no longer re-combines what is already combined: a new optional combined= argument carries the caller's original packed buffer, which is quantized directly instead of rebuilding the packed buffer from q/k/v views via combine_tensors (a raw set_ with a silent adjacency/interleave assumption). The packed original is threaded from DPA.forward through FusedAttention to FusedAttnFunc.forward; all legacy call sites are untouched (combined=None preserves exact behavior), and backward combine calls are unchanged (gradients have no pre-packed original). Tests: dense fwd+grad bit-exactness vs separate contiguous q/k/v for bs3hd/bsh3d/sb3hd/kv-packed/GQA (fused + flash), validation errors, torch.compile (no data_ptr/UntypedStorage graph breaks), FP8 combined-vs-views bit equivalence, and detection-free declared t3hd. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] MultiheadAttention: pass packed projection output to DPA declaratively Adopt the new DotProductAttention packed API inside MultiheadAttention: * self-attention (np == ng): the fused QKV projection output, already viewed as [.., h, 3, d] (qkv_weight_interleaved) or [.., 3, h, d], is handed to DPA directly as qkv_layer with the matching qkv_interleave_dim (-2 / -3) -- no SplitAlongDim slicing in MHA and no pointer-based layout detection in DPA. * cross-attention: the packed KV projection output is exposed as [.., hg, 2, d] / [.., 2, hg, d] and passed as kv_layer. * The pass-through only engages when no per-tensor operation needs the individual q/k/v slices: it is skipped for RoPE, QK normalization, KV caching (inference_params), CPU offloading, GQA (np != ng, not a uniform 3-interleave) and quantized (FP8) projection outputs; those keep the legacy sliced-views path unchanged. Tests: MHA self (interleaved + non-interleaved) and cross (both interleaves) are bit-exact vs the same MHA with packed inputs converted back to separate contiguous q/k/v (output, input grad, weight grads); spy asserts the packed argument and interleave dim actually reach DPA; GQA and RoPE fall back to the views path. TransformerLayer regression suite unchanged; test_kv_cache failures on this device are pre-existing on origin/main (verified). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] DotProductAttention: factor packed-input handling into _unpack_packed_qkv Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] DotProductAttention: validate packed-input last-dim stride Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] DotProductAttention: pass packed qkv/kv buffers as separate arguments Rename the single layout-dependent packed_qkv plumbing argument (which held the full QKV buffer for *3* layouts but the KV buffer for *_2* layouts) into explicit packed_qkv/packed_kv, mirroring the public qkv_layer/kv_layer API. combine_and_quantize's combined= is split into combined_qkv=/combined_kv= accordingly, and _unpack_packed_qkv no longer returns the packed tensor since the callers already hold qkv_layer/kv_layer. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Move packed qkv/kv input tests into test_attention.py Fold the tests from the new test_dpa_packed_inputs.py file into the existing attention test suite as a dedicated section, reusing its imports. No new test file; test logic unchanged apart from renaming the module-level constants (_B/_S/_H/_D/_DTYPE -> _PACKED_*) to avoid collisions. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Deprecate pointer-based detection of packed qkv layouts get_qkv_layout now emits a DeprecationWarning when it recognizes q/k/v as views of a packed buffer purely from data pointers/strides/offsets (detected *3*/*_2* layouts), pointing callers at the declarative qkv_layer/kv_layer API. Separate q/k/v tensors (hd_hd_hd layouts) never warn since there is nothing to declare. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Fix implicit string concatenation lint warning Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Trim packed-input test section to equivalence and MHA tests Drop the API-validation, torch.compile graph-break, combine_and_quantize equivalence and thd no-detection spy tests; keep the dense/flash bit-exact equivalence tests and the MHA packed pass-through/fallback tests. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Drop MHA packed pass-through tests Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Fold packed-input tests into test_dpa_qkv_layout via a declarative param Parametrize test_dpa_qkv_layout and test_dpa_qkv_layout_thd with declarative={views,declarative}: the declarative mode passes the packed buffer to DotProductAttention via qkv_layer/kv_layer (declared layout, gradients read off the packed buffer) instead of slicing it into q/k/v views for pointer-based detection. This reuses the whole existing config matrix (masks, bias, SWA, cross-attention, thd, all backends) for the declarative API, replacing the dedicated packed-input test section. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Shrink declarative packed-input tests to a dedicated small matrix Revert the declarative parametrization of test_dpa_qkv_layout(_thd) (which doubled their whole config x layout product) and instead add test_dpa_qkv_layout(_thd)_declarative covering all packed layouts on a trimmed config dimension: one self-attention and one cross-attention config (kv_layer path) for dense, one config for thd. Past the input handling the backend code is identical to the views mode, so the full config matrix added no coverage. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Use explicit q/k/v grad variables in the DPA test harness Replace the .grad-holder objects substituted for q/k/v in declarative packed mode with q_grad/k_grad/v_grad variables computed right after backward, used uniformly by all return paths. 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 * [PyTorch] Address review: warning stacklevel, offload suppression, explicit k_norm - get_qkv_layout deprecation warning: add stacklevel=2 and skip it while CPU offloading is enabled (offloading forces MultiheadAttention onto the sliced-views fallback, so the caller has no migration option there). - MultiheadAttention: gate the packed pass-through on k_norm explicitly instead of relying on q_norm/k_norm being created together. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Address review nits: clarify DPA packed-input docstrings - DotProductAttention.forward: mark query/key/value_layer as Optional and note they are required only when no packed input (qkv_layer/kv_layer) is given (Charlene). - combine_and_quantize: describe combined_qkv/combined_kv in terms of qkv_group=1 / qkv_group=2 layouts instead of '3'/'2' layouts (Charlene). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* [PyTorch] Make UnfusedDotProductAttention compatible with torch.compile + CUDA graphs Refactor TE custom kernels used by the unfused attention path so that `torch.compile(fullgraph=True, mode="reduce-overhead")` can trace the forward and backward and capture them into CUDA graphs without graph breaks. - softmax.py / softmax.cpp: register all `scaled_*_softmax_{forward,backward}` kernels as `torch.library.custom_op`s with fake impls and an autograd binding that mirrors the previous `torch.autograd.Function`s. The C++ backward kernels now allocate a fresh output buffer instead of writing in-place into `output_grad`, so the ops no longer alias their inputs (required by `torch.library.custom_op` and inductor cudagraph trees). - utils.py: convert `ConvertTHDtoBSHD` / `ConvertBSHDtoTHD` to `torch.library.custom_op`s, with thin wrapper classes that keep the existing `.apply(...)` callsite syntax. Drop the `int(cu_seqlens[-1].item())` from the hot path of `ConvertBSHDtoTHD.apply` -- under `torch.compile` it created an unbacked SymInt, which made the Inductor partitioner emit `None` placeholders for output buffers and caused `cudagraph_trees` to assert. `num_tokens` is now passed in by the caller as a regular (Sym)Int. - backends.py: in the THD branch of unfused DPA, capture `total_tokens_q = query_layer.shape[0]` before overwriting `query_layer` with the BSHD form, and thread it back into `ConvertBSHDtoTHD.apply` at the end of the forward. - test_torch_compile.py: add `test_unfused_dpa_torch_compile`, parametrized over qkv layouts (`bshd_bshd_bshd`, `sbhd_sbhd_sbhd`, `thd_thd_thd`, `bs3hd`, `sbh3d`), that compiles `UnfusedDotProductAttention.forward` directly with `fullgraph=True, mode="reduce-overhead"` and runs forward+backward several times so the CUDA graphs are recorded and replayed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Made-with: Cursor * [PyTorch] torch.compile support for the FP8-emulation path of UnfusedDotProductAttention Make the FP8-emulation path (NVTE_UnfusedDPA_Emulate_FP8=1) of UnfusedDotProductAttention traceable by torch.compile(fullgraph=True). - backends.py: register the quantize+dequantize roundtrips used by FP8EmulationFunc as torch.library custom ops (te_fp8_emu::roundtrip_<QuantizerClass> and te_fp8_emu::roundtrip_qkv_<QuantizerClass>) taking the quantizer as a value-opaque argument, with fake impls for tracing. Ops are registered only for the value-opaque quantizer classes (Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer (delayed scaling) carries scale/amax tensor state, is not value-opaque, and deliberately keeps the plain eager path -- FP8 emulation with delayed scaling is not supported under torch.compile. - backends.py: dispatch helpers `_fp8_emu_roundtrip{,_qkv}` key on `type(quantizer).__qualname__` so they stay traceable for opaque quantizer arguments; FP8EmulationFunc forward/backward now call them (onnx_forward unchanged). - backends.py: the joint q/k/v roundtrip clones any output whose storage is shared with an input or another output, checking storage identity directly -- the dequantized q/k/v can be views into one combined buffer, and view metadata (`_base`) is not populated under the torch-dispatch mode AOTAutograd runs custom ops with, so a `_base`-guarded clone triggered the custom-op aliasing deprecation warning under torch.compile. - UnfusedDotProductAttention.forward: only query FP8GlobalStateManager.get_fp8_recipe() when fp8_meta["local_recipes"] is absent. - test_torch_compile.py: add test_unfused_dpa_fp8_emulation_torch_compile (current scaling + mxfp8, sbhd/bshd layouts; compiled fullgraph forward+backward must match eager) and test_unfused_dpa_fp8_emulation_delayed_scaling_eager guarding the eager delayed-scaling path after the refactor. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Run UnfusedDotProductAttention as an eager island when fp8_output=True With fp8_output=True the backend returns a Float8Tensor -- a tensor subclass that cannot cross a torch.compile graph boundary -- so the forward dispatches to a torch._dynamo.disable'd wrapper, the same mechanism DotProductAttention and FusedAttention use module-wide. With fp8_output=False the dispatcher is resolved at trace time and adds no graph break. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Test FP8-emulation compile path also with reduce-overhead (cudagraphs) Parametrize test_unfused_dpa_fp8_emulation_torch_compile over compile mode (default, reduce-overhead), run 3 iterations so the CUDA graphs are recorded and replayed. The te_fp8_emu roundtrip ops for current scaling are pure (no mutated args), so inductor cudagraphs capture them; verified no cudagraph skips with TORCH_LOGS=cudagraphs. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Drop FP8 torch.compile support in UnfusedDotProductAttention; run FP8 as an eager island FP8 in the unfused backend (emulation and Float8Tensor output) is not supported under torch.compile: the forward dispatcher routes fp8=True and/or fp8_output=True to a torch._dynamo.disable'd wrapper, same as DotProductAttention does module-wide. Remove the FP8-emulation compile tests. The te_fp8_emu::* custom ops taking value-opaque quantizers stay as the eager implementation of FP8EmulationFunc. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Remove te_fp8_emu custom ops; restore upstream FP8EmulationFunc The ops existed solely to make the FP8-emulation path traceable by torch.compile; since FP8 in the unfused backend now always runs as an eager island, they are dead machinery (plus import-time registration and output clones the plain eager path never needed). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Shorten the total_tokens_q comment Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Drop _scale_to_tensor; pass the softmax scale as a plain float The tex softmax kernels take 'float scale_factor' directly. The 0-D tensor wrapping was a leftover of the old autograd.Function idiom, where the float had to be a tensor only to fit save_for_backward; the custom ops keep the scale on ctx as a plain attribute. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Drop redundant num_tokens comment (rationale documented at the callsite) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Review cleanups: black formatting, drop unused import/duplicate dict, silence W0613 - run black over the four changed files (earlier commits skipped pre-commit) - drop unused 'import os' in test_torch_compile.py - drop duplicated module-level _default_causal_mask dict in softmax.py - del unused 'output' arg in the conversion setup_context helpers 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 * [PyTorch] Address review: document Convert{THD,BSHD} apply() arity in docstrings Greptile P2: the ConvertTHDtoBSHD/ConvertBSHDtoTHD class docstrings said callsites keep the .apply(...) syntax without reflecting the actual argument list. Spell out the apply() signatures so the required args (incl. num_tokens / max_seqlen) are explicit. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Fix UnboundLocalError on thd inference in UnfusedDotProductAttention The thd output conversion (q_format=='thd') passes total_tokens_q to the new ConvertBSHDtoTHD custom op, but total_tokens_q was only assigned on the training 'thd' input branch, not the inference 'thd_2bshd' branch, so thd KV-cache inference raised UnboundLocalError. Capture total_tokens_q once right after q_format is known, before any layout conversion: for both 'thd' and 'thd_2bshd' the query enters in thd layout so query_layer.shape[0] is the total query token count (a backed SymInt, unlike cu_seqlens_q[-1].item() which would sync the GPU and break torch.compile + cudagraphs). 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>

* optimize for gemm's conditional enablement Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * avoid code repeatition Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove redundant test Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address review comment Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

… backward (#3206) * [common][PyTorch] NVFP4: enable row-scaled transpose quantization for backward The #2931 row-scaled NVFP4 path only produced the rowwise forward activation; its columnwise/transpose output was rejected. That made the per-token activation unusable in the backward wgrad GEMM, so row-scaled training had to fall back to a dequantized/high-precision backward. This change extends the existing row-scaled path to also emit the columnwise (transpose) direction, so a training Linear with row_scaled_activation=True now quantizes the forward activation row-scaled in both directions and the wgrad GEMM consumes the row-scaled transpose directly. It is a minimal extension of #2931 (no new CUTLASS kernels, no grouped path, no RHT/4over6 transpose fusion). Signed-off-by: Cael Ling <caell@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Cael Ling <caell@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* fix(pytorch): add missing f-string prefixes to error messages Signed-off-by: Andrew White <andrewh@cdw.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Andrew White <andrewh@cdw.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Simplify NCCL header/lib CMake logic NCCL is now always required for all native build paths (due to a change sometime between v2.16 and main). This updates CMake to only search for the headers in a single place. This does not impact the NCCL EP dependency resolution. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Add additional NCCL header search paths This will now also search Python's site-packages, making it much easier to pull in headers from `nvidia-nccl-cu{12,13}`. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Make TE linker script be provided when linking TE Previously this was provided every time the C++ or CUDA compiler was invoked for everything built by CMake, including compiler detection and compilation, instead of just when linking the TE library. Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* [PyTorch] Preserve grouped weight initialization metadata Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt> * [PyTorch] Mark late grouped weights for delayed wgrad Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt> * [PyTorch] Preserve grouped MXFP8 columnwise usage Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt> * Fix bug where grouped MLP used same quantizer in grouped linear op and weight param Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Fix inconsistent line endings GitHub inserted CRLFs. Microsoft... Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: CarlosGomes98 <carlosmiguel.gomes@live.com.pt> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <tmoon@nvidia.com>

…s with pickles (#3245) Allow Flash Attention tests to use checkpoints with pickles Signed-off-by: Tim Moon <tmoon@nvidia.com>

…#3240) * test: cover lazy NCCL setup for Newton-Schulz Construct the cuSOLVERMp context before any other distributed collective so the distributed test exercises lazy ProcessGroupNCCL communicator initialization. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * fix: initialize borrowed NCCL comm for Newton-Schulz Materialize the ProcessGroupNCCL communicator before borrowing its raw handle, retain the process group for the context lifetime, and validate the communicator size and rank during native context creation. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* Add NCCL EP eager mode and drop-on-overflow group policy * [PyTorch] Rename EP prepare outputs to tokens_per_expert and total_recv_tokens * [PyTorch] Assert ep_bootstrap ran before EpBuffer construction Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…(#3266) * [PyTorch] Fix Float8BlockwiseQTensor.shape for columnwise-only tensors Float8BlockwiseQTensor.shape is a fast path added to avoid PyObject lookups. For a columnwise-only tensor it returned the raw columnwise buffer shape, but blockwise stores columnwise data transposed, so the property disagreed with both Float8BlockwiseQTensorStorage.size() and the wrapper's own metadata: make_empty((128, 256)) yielded .shape == (256, 128) while .size() == (128, 256). Apply the same reorder size() does. Float8Tensor and NVFP4Tensor already de-transpose in their equivalent fast paths; MXFP8Tensor needs no change because it stores columnwise data untransposed. Cover the invariant for every quantization scheme and usage combination with test_shape_matches_size, rather than only for blockwise. Introduced in 9dac78e ("CPU Overhead Optimizations", #2559). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip row-scaled NVFP4 columnwise-only case and speed up the 2D shape path Row-scaled NVFP4 accepts set_usage(rowwise=False) but its allocator asserts on rowwise usage, so test_shape_matches_size hit an NVTE_CHECK failure instead of skipping. Filter that combination out before set_usage. Also give Float8BlockwiseQTensor.shape a 2D fast path: indexing torch.Size directly measures ~186 ns/call against ~311 ns for building an intermediate list, on the columnwise-only branch. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Fix size() on the columnwise-only path for FP8 and blockwise Both storages forwarded *args straight to the underlying buffer, then reordered the result. That works for the rowwise branch, where the buffer matches the logical shape, but not for the transposed one: size(dim) returned an int from the buffer, which the reorder then tried to index, so every size(dim) call raised TypeError. Forwarding dim is also wrong in principle there, since buffer dim i is not logical dim i. Rebuild the logical shape in full first, then index into it. Float8TensorStorage also flattened the transpose-only shape to 2D, which disagreed with both dim() and Float8Tensor.shape for rank >= 3; it now applies the same rotation the shape property does. Reachable from quantize() followed by update_usage(rowwise_usage=False). NVFP4 is left alone: it reports columnwise-only tensors flattened on purpose and warns about it, and shape and size() agree there. Extend test_shape_matches_size over 2D and 3D shapes, asserting that shape and size() describe the same tensor and that size(dim) agrees with it for positive and negative dims. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>

The FA-version job reloads checkpoints generated by test_attention.py, so the legacy delayed-scaling FP8 metadata is trusted inside this test. Mirror the L0 opt-in across every L3 attention execution path to preserve the secure runtime default while preventing false version-matrix failures. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Make NVTE tensor handle pool size configurable Signed-off-by: hongbinl <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Validate tensor handle pool env vars Signed-off-by: hongbinl <hongbinl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify tensor handle pool env parsing Signed-off-by: hongbinl <hongbinl@nvidia.com> * Rerun CI after BIA artifact mount outage Signed-off-by: hongbinl <hongbinl@nvidia.com> --------- Signed-off-by: hongbinl <hongbinl@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Gate FA4 and stabilize attention test imports FA4 can be installed on SM8x even though its current implementation rejects those GPUs. Disable selection and skip dedicated FA4 tests there so A100 and L40 use supported attention backends. FA4 and CUTLASS can also expose a generic utils package on sys.path. Prepend the Transformer Engine test helper directory in the context-parallel test so collection resolves the intended utilities. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Isolate FlashAttention CI backends Moving images can install FA4 alongside older FlashAttention generations, which mixes a shared Python namespace and can make context-parallel reference runs compile an unsupported backend. Isolate the L3 version matrix, keep current CP comparisons on FA2/FA3, and temporarily reject symmetric D512 FA4 on Blackwell until upstream kernel support is complete. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Prepend PyTorch test utility imports FA4 and its CUTLASS dependency expose a top-level utils module after Transformer Engine imports. Appending the test root can therefore bind these late imports to the installed module and fail collection. Give the repository helper precedence in the four test files that exhibited this ordering. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Recognize FA3 sliding-window CP support The all-gather and a2a guards use the FA2 package version check to recognize FlashAttention support, so an isolated FA3 run is rejected even though FA3 implements sliding-window attention. Accept the explicit FA3 backend in both guards. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Keep CP backend selection at CI boundaries The CP runner must honor an explicit generation selected by its caller, particularly the existing B200 L3 FA4 lane. Remove its internal V4 override, restore the L3 SM100 selection changed in 0f6c71eb, and disable V4 only for the L1 suite that still targets FA2/FA3. This keeps per-generation L3 isolation intact without making the shared runner silently override directed coverage. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Align FlashAttention CI coverage by architecture Keep L0 on the mature FA2 path while L3 owns newer-generation coverage. Restrict H100 L3 to FA3 and B200 L3 to non-CP FA4 so unsupported H100 FA4 kernels and mislabeled Blackwell CP results do not obscure the intended signal. Make FA4-specific tests honor backend enablement to prevent silent fallback under an FA4 label. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Guard FA4 against incompatible CUTLASS installs Package metadata can report FA4 present even when a later dependency install leaves its transitive CUTLASS stack unusable. Reject the known b24/CUTLASS combination below the stable 4.6.2 release and treat a nested interface ImportError as an unavailable optional backend so unrelated Transformer Engine imports can continue. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Squashed to single commit for review. Original PR: andrewwhitecdw/TransformerEngine#8 Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>

… Op in TE Sequential (#3320) * produce/consume extra output Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * allow for fusions with producer/consumer being part of same fuser with error handling tests Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * minor cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * dispatch combine impl Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fusible ops test Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * keep just ops infra changes Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * cleanup with residual tests Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comment Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * update to cleaner documentation Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * update docs Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * pin channels through channel version Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * unecessary handling removal Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * simplify Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * doc update + extra_grad = None case Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * no need to check staleness in every forward call Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * remove redundant tests Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert from bad names Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * keep simple Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * unecessary checks Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * minor doc Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * Update docs/examples/op_fuser/op_fuser.rst Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * Update transformer_engine/pytorch/ops/fuser.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * address review comments + extra output being configurable to be outputted Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * address review comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * not picklable Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * simplify.. lock it permanently Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * a bit of doc Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

* Support MXFP8 2D quantization Signed-off-by: kunlunl <kunlunl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix comments Signed-off-by: kunlunl <kunlunl@nvidia.com> * Fix comments Signed-off-by: kunlunl <kunlunl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support explicit 2D MXFP8 grouped quantization Signed-off-by: kunlunl <kunlunl@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: kunlunl <kunlunl@nvidia.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com>

…ear (#3324) * [PyTorch] Enable NVFP4 row-scaled (per-token) backward for GroupedLinear Extend the row-scaled NVFP4 support added for dense Linear to the MoE GroupedLinear module, so the wgrad is computed in NVFP4 instead of falling back to high precision. Signed-off-by: Cael Ling <caell@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unhelpful comments Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Remove unnecessary comment Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Cael Ling <caell@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

Move ffi type definition before ffi target definitions Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

use pinned memory; use explicit cuda sync Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

* [Common] Ensure quantization kernels handle noop properly Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * nit Signed-off-by: Kaining Zhong <kainingz@nvidia.com> --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

…tests (#3398) Support synchronous collectives in HLO collective bytes assert tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> (cherry picked from commit 27b12df7674e7f95beb4d4adfbcfb29d2f880e52)

… (#3373) fix: SBHD reorder skip uses original shape instead of swapped tensor Use tensor.shape[seq_dim] instead of shape[seq_dim] when deciding whether a Striped SBHD case is large enough. Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> (cherry picked from commit a8ff8244165a9766adadbe6ea070c3bb3a3ff5b5)

* Bump min compute_on version Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Bump to 0.11.1 Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit aa490dca9119a1d1841e55d657b7b575fd4e7a2f)

* Fix build issues with 26.08 JAX container images Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Update test to better cover change Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Remove backwards compat fix for jax 0.10.1 and 0.11.0 (see NVIDIA/TransformerEngine#3406) Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> (cherry picked from commit 250ef1bccf02e9fda4cb47078fce4e3c953bd150)

…cs (#3409) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 17c351c34ccfb32506c7ef3cdefb4e51b4b4f686)

* Surface distributed comm-overlap rank errors Comm-overlap launchers captured child stdout and stderr but raised only stderr, which could leave the originating rank context out of pytest and JUnit failures. Preserve the existing result predicates while attaching bounded tails of both streams, and record layer-worker exceptions so torchrun can report the rank-local traceback. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Isolate comm-overlap launcher environment The launcher mutated os.environ and then called os.unsetenv, which leaves Python's environment mapping unchanged. Conditional backend flags could therefore leak into later parameterized children. Build a child-only environment instead so each launch gets its intended overrides while preserving the parent process environment. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep FA4 out of L1 distributed tests The moving PyTorch image installs FA4 by default, implicitly expanding an L1 suite that historically covered earlier attention backends. Export the FA4 selector at the suite boundary so every pytest and torchrun child retains the intended backend scope while dedicated attention suites own FA4 coverage. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit 54f2f36a088485f38ffaeaab473ec38d8d42212d)

…uncher pass count (#3415) [PyTorch] Scale test_multi_process_ep outer timeout with launcher pass count Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Co-authored-by: fheinecke <23390735+fheinecke@users.noreply.github.com> (cherry picked from commit d0067d0da6cb01811b8072a7d304f0e00344ce2b)

…ith 26.08 JAX image (2.19 only) Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* Add support in lower level JAX API for returning max logit and softmax aux to the user from TE JAX fused attn output Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add support for returning reduced per head max logit. Plumb max logit and softmax through the JAX fused attn primitives Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add max logit to JAX fused attn FFI and set it in the workspace Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add first pass tests for max logit and softmax aux tensor outputs in JAX fused attn tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Reject aux returns with score_mod Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Handle SM120 max-logit layout Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Drop softmax aux return Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Modify static args in fused attn tests for jax Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Inline the choice of what is to be returned and remove redundant function for it Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Expose JAX max logit Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Support CP max logit Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Expand JAX max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Remove JAX max logit integration tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Broaden JAX CP max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Reduce JAX max logit across DP Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Simplify JAX max logit return Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Document JAX max logit reductions Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine JAX max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Rename JAX max logit buffer Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Check JAX attention tensor pack capacity Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit 0459a095509ee0bc68a5d5abbaea4aa6af88e4f6)

* Add distributed Muon optimizer Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Muon closure and reference test Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer distributed API handling Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer docs and params typing Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tensor-parallel Newton-Schulz wrapper Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Move Newton-Schulz wrapper into optimizers Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Use tensor-parallel Newton-Schulz in Muon Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Limit Muon branch to Newton-Schulz TP Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Run Newton-Schulz distributed cases in one launch Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Add single-GPU Newton-Schulz coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Newton-Schulz compatibility and replicated coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit 5df5a0d8f13cec533afb9073967701529ff401a9)

…in HLO (#3412) * [JAX] Fix counting of synced and wrapped collectives in HLO Signed-off-by: Alex Y. Chan <alechan@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alex Y. Chan <alechan@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit ff22bea00da8dc3f9e93faf52f80bcb9a97d2d6f)

… (#3425) [JAX] Fix sync-tagged collective start classification Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> (cherry picked from commit d7340a43c9f928f96884c3d842d2738406d6703f)

* update_filter_fp8_thd_attention Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * add for fp8+thd debug Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * uncomment the configs Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Re-plumb FP8+THD ragged-offset support on top of merged main Mirrors the F16 arbitrary_seqlen ragged-offset pattern in the FP8 path: - Backend selector: enable FP8+THD for cuDNN >= 9.23 on sm >= 100 - fwd/bwd _impl: ragged detection, batch/seqlen bucketing, set_ragged_offset() on Q/K/V/O/dO/dQ/dK/dV/Stats, workspace allocation for ragged offsets, cu_seqlens_padded_to_offsets kernel - fwd/bwd dispatchers: accept num_tokens_q/kv, cu_seqlens_padded, compute max_batch/max_tokens, THD Stats shape Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address FP8 THD review feedback Use cuDNN 9.23 as the FP8 THD ragged-offset gate and prefer int64 offsets, matching cuDNN guidance for the new path. Restrict FP8 THD backend selection to padding masks, align ragged offset tuple order with the F16 convention, and enable zero-fill for FP8 THD comparison tests. Suppress the forward FP8 graph-builder fn_size lint using the same local pattern already used by the backward builder, because refactoring the full graph construction is outside this review cleanup. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 THD ragged version scope Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 ragged offset kernel calls after main merge The shared conversion kernel now accepts RaggedOffsetMultipliers, so construct it in FP8 forward and backward instead of passing the removed scalar argument list. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 direct-seqlen merge integration Keep the actual batch when cuDNN consumes user cu-seqlens because a bucketed batch would read past the buffers. Reuse the aligned fallback workspace and keep SM120 stats dense so allocation matches the graph layout. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Disable fast zero fill in FP8 THD comparison The optimized mha_fill path reads CUDA cu_seqlens from host C++ and segfaults for THD. A controlled A/B passed with False while the enabled path exited 139. Keep the comparison test on the safe path until a graph-safe zero-fill implementation lands. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Align FP8 THD MHA test token counts FP8 Linear flattens THD inputs to [t, h*d], so adjust generated sequence lengths before building cu_seqlens to make the total token count divisible by eight for both forward and backward. Keep fast zero fill disabled in the MHA helper to avoid the known host dereference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Enable FP8 THD context parallel attention Remove obsolete FP8+THD+CP gates now that the fused-attention path supports this combination. Replace the unsafe host-side suffix calculation with stream-ordered zeroing, and preserve FP8 metadata and token-major THD layout in the all-gather path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix FP8 THD CP gradient padding P2P FP8 backward densely combines rank-local partial gradients, which can repopulate inter-sequence dK/dV padding after native zero initialization. Reconstruct every local padding interval from actual and padded cu-seqlens and clear dQ/dK/dV after reduction. Initially limit cleanup to FP8 because high-precision control runs kept padding zero; later commits can broaden that policy independently. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable FP8 THD inference on Hopper with cuDNN 9.25 cuDNN 9.25 provides a working Hopper FP8 THD forward kernel, while 9.23 selects a plan that traps with an illegal instruction. Keep Hopper backward gated and preserve the existing cuDNN 9.23 requirement on Blackwell. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Cover FP8 THD CP inference on Hopper Reuse four existing Hopper cases for forward-only validation instead of adding a new Cartesian test axis. Blackwell retains forward-and-backward coverage, while Hopper runs the cuDNN 9.25-supported inference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Support bottom-right causal FP8 THD attention Admit the bottom-right causal mask only on SM100+, where the cuDNN FP8 backend supports it. Treat the mask as both causal and padding when constructing the frontend graph. Place delayed-scaling P2P half-gradients into their per-sequence THD halves so causal CP backward does not copy half-sized tensors into full buffers. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fuse delayed FP8 THD half-gradient placement Extend the native THD gradient correction kernel with copy-and-zero operations for raw byte gradients. This replaces per-step Python index construction with one vectorized CUDA launch while preserving the inactive-half zeroing required by causal P2P backward. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Limit FP8 THD attention to Blackwell Remove the Hopper forward-only selector and test specialization because this PR targets complete FP8 THD support, including backward. Apply THD gradient-padding cleanup regardless of FP8 state so enabling the cleanup does not change non-FP8 padding semantics. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Exercise default zero fill in FP8 THD tests Reserve seven positions in the final random THD sequence so aligning total tokens for cuBLAS only increases the sequence length without exceeding the configured maximum. Remove the temporary fast_zero_fill overrides now that the default path passes MHA and DPA backward validation. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Vectorize THD CP gradient padding cleanup Build padding masks with an on-device search instead of launching an operation per sequence, and reuse the KV mask for dK and dV. This removes batch-size-dependent overhead while preserving valid gradients. Add reference, mutation, and CUDA graph coverage for the helper. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Clean THD gradients before reverse A2A Per-step THD sequence metadata describes gradients in sequence order. Zero inter-sequence padding before reverse A2A converts those gradients to CP-rank order, preventing valid rank-local gradient rows from being cleared. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Address FP8 THD review feedback Remove redundant zero-fill and backend conditions while preserving strided THD gradient handling. Reset the shared delayed-scaling amax explicitly because ATen view zeroing does not update quantizer metadata. Co-locate the flattened THD padding mask with the dense padding-mask utilities and clarify raw-byte gradient correction semantics. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retry transient FP8 THD pool NaNs Long-lived context-parallel workers can retain state across heterogeneous cases, while the same FP8 THD case succeeds in a fresh worker. Reuse the existing one-retry mechanism only for the specific NaN assertion; persistent NaNs still fail on the fresh retry. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Normalize THD masks in FP8 attention tests THD inputs carry variable sequence lengths, so backend selection requires the padding-aware form of each logical mask. Copy each shared model config before translating the mask to avoid leaking mutations across parametrized cases. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Enable fused FP8 THD selection on SM90 cuDNN 9.23 adds FP8 THD forward support on Hopper. Relax the architecture gate to exercise SM90 support; the following commit narrows this to forward-only after validation showed Hopper backward must remain unavailable. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Limit SM90 fused FP8 THD to forward cuDNN supports fused FP8 THD forward on Hopper but rejects the backward path. Keep SM100+ eligible for training while allowing SM90 only for inference so unsupported training configurations skip during backend selection. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit ffcb3c0013dbdead74a144c5aef890f62be86d2a)

* TE/JAX MoEBlock optimizations Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Keep arch guard consistent in EP tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com> (cherry picked from commit 2fd4604bd6b9c232baa6c352385eeca8405a39d9)

Merge upstream NVIDIA TransformerEngine release_v2.19 into the ROCm dev branch and resolve all conflicts, preserving ROCm adaptations while adopting upstream improvements. Key resolutions: - grouped_linear (module + ops): drop removed ROCm path-support probes; gate the ROCm inline split-quantize path under IS_HIP_EXTENSION and use upstream's _split_quantize helpers on CUDA; fall back to the split path when use_grouped_tensor=True on ROCm (no cuBLASLt grouped-tensor GEMM). - graph.py: keep the ROCM-25129 warmup-stream reuse while adopting upstream's pre_warmup_hook / capture-time hook plumbing. - FusedAttnBackend (cpp_extensions/fused_attn.py): adopt upstream's IntEnum but define members per-platform so the import-time C++/python sync assertion passes on both the ROCm and CUDA enums. - dot_product_attention (backends/context_parallel/utils): keep ROCm CK/ AOTriton backend selection and IS_HIP_EXTENSION guards; adopt upstream's is_version_supported() FA gating. - pytorch csrc attention.cpp: keep the TECUDAGuard masquerade; drop dead mha_fill now that upstream zero-fills via TensorWrapper::zero_/fill_. - jax csrc attention.cpp: adopt upstream's dynamic aux-tensor packing and return_max_logit handling under the ROCm backend guard; keep (void) memset casts. - jax base.py: single platform-aware FFI registration loop after EpInstanceState registration. VERSION.txt -> 2.19.0.dev0. Replace 3rdparty/nccl submodule with 3rdparty/nccl-extensions per upstream. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Copilot

Pull request overview

Integrates Transformer Engine v2.19 development changes across quantization, distributed execution, kernels, build infrastructure, and documentation.

Changes:

Adds hybrid/identity quantization, MXFP8 2D/FSDP support, and distributed-weight integration.
Expands expert parallelism, fused attention, activation, and context-parallel capabilities.
Updates tests, CI, dependencies, examples, and documentation.

Reviewed changes

Copilot reviewed 109 out of 269 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`transformer_engine/pytorch/utils.py`	Adds compile-safe hardware queries.
`transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py`	Annotates NVFP4 inner tensors.
`transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py`	Adds MXFP8 FSDP buffer handling.
`transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py`	Rejects unsupported grouped quantizers.
`transformer_engine/pytorch/tensor/nvfp4_tensor.py`	Adds metadata and allocation specifications.
`transformer_engine/pytorch/tensor/float8_blockwise_tensor.py`	Adds allocation metadata and shape handling.
`transformer_engine/pytorch/tensor/_quantization_helpers.py`	Adds efficient view-shape resolution.
`transformer_engine/pytorch/tensor/__init__.py`	Exports hybrid and identity tensors.
`transformer_engine/pytorch/quantization.py`	Adds alignment and MXFP8 2D logic.
`transformer_engine/pytorch/optimizers/fused_adam.py`	Advances empty parameter-group steps.
`transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py`	Refines return annotations.
`transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py`	Refines return annotations.
`transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py`	Refines return annotations.
`transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py`	Refines return annotations.
`transformer_engine/pytorch/ops/fused/backward_linear_add.py`	Handles absent extra gradients.
`transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py`	Handles unused extra outputs.
`transformer_engine/pytorch/ops/basic/rmsnorm.py`	Adds quantizer fallback handling.
`transformer_engine/pytorch/ops/basic/make_extra_output.py`	Supports absent extra gradients.
`transformer_engine/pytorch/ops/basic/layer_norm.py`	Adds quantizer fallback handling.
`transformer_engine/pytorch/ops/basic/add_extra_input.py`	Refines return annotations.
`transformer_engine/pytorch/ops/basic/activation.py`	Introduces fused scaled unary activations.
`transformer_engine/pytorch/ops/basic/__init__.py`	Exports grouped-path capability query.
`transformer_engine/pytorch/ops/_common.py`	Filters normalization quantizers.
`transformer_engine/pytorch/module/fp8_unpadding.py`	Documents recipe-specific alignment.
`transformer_engine/pytorch/module/fp8_padding.py`	Documents recipe-specific alignment.
`transformer_engine/pytorch/module/_common.py`	Adds wgrad quantizer helpers.
`transformer_engine/pytorch/module/__init__.py`	Exports grouped-path capability query.
`transformer_engine/pytorch/dynamo/__init__.py`	Exports tensor specifications.
`transformer_engine/pytorch/distributed_weight.py`	Defines distributed-weight extension protocol.
`transformer_engine/pytorch/custom_recipes/reference_utils.py`	Updates reference utility description.
`transformer_engine/pytorch/custom_recipes/reference_current_scaling.py`	Renames and reorganizes reference APIs.
`transformer_engine/pytorch/custom_recipes/quantization.py`	Removes superseded GEMM definitions.
`transformer_engine/pytorch/custom_recipes/gemm.py`	Hosts custom GEMM types.
`transformer_engine/pytorch/csrc/type_converters.cpp`	Distinguishes initialized empty storage.
`transformer_engine/pytorch/csrc/extensions/swizzle.cpp`	Supports variable-shape grouped swizzling.
`transformer_engine/pytorch/csrc/extensions/softmax.cpp`	Avoids backward input aliasing.
`transformer_engine/pytorch/csrc/extensions/pybind.cpp`	Exposes grouped and scaled operations.
`transformer_engine/pytorch/csrc/extensions.h`	Updates native extension declarations.
`transformer_engine/pytorch/csrc/common.h`	Adds MXFP8 2D state.
`transformer_engine/pytorch/attention/__init__.py`	Exports fused MLA projection.
`transformer_engine/pytorch/__init__.py`	Exports new public PyTorch APIs.
`transformer_engine/jax/version_utils.py`	Detects collective-stream support.
`transformer_engine/jax/router.py`	Corrects auxiliary-loss gradient forwarding.
`transformer_engine/jax/quantize/tensor.py`	Supports grouped transport views.
`transformer_engine/jax/flax/transformer.py`	Adds maximum-logit attention output.
`transformer_engine/jax/csrc/extensions/pybind.cpp`	Extends EP bootstrap binding.
`transformer_engine/jax/csrc/extensions/ep.cpp`	Adds EP overflow and receive totals.
`transformer_engine/jax/csrc/extensions.h`	Updates JAX native declarations.
`transformer_engine/jax/cpp_extensions/quantization.py`	Supports shard-local MXFP8 groups.
`transformer_engine/jax/cpp_extensions/gemm.py`	Handles compound mesh axes.
`transformer_engine/jax/cpp_extensions/base.py`	Corrects FFI registration ordering.
`transformer_engine/debug/pytorch/debug_quantization.py`	Corrects directional usage updates.
`transformer_engine/common/util/cuda_runtime.cpp`	Discovers wheel-provided CUDA headers.
`transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu`	Adds portable stochastic FP4 fallback.
`transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu`	Supports transpose-only quantization.
`transformer_engine/common/recipe/__init__.py`	Adds MXFP8 2D and custom alignment options.
`transformer_engine/common/normalization/common.h`	Declares MXFP8 normalization toggle.
`transformer_engine/common/normalization/common.cpp`	Implements cuDNN MXFP8 dtype toggle.
`transformer_engine/common/newton_schulz/newton_schulz.cpp`	Validates NCCL communicator metadata.
`transformer_engine/common/include/transformer_engine/transformer_engine.h`	Extends quantization configuration API.
`transformer_engine/common/include/transformer_engine/fused_attn.h`	Clarifies gradient-correction behavior.
`transformer_engine/common/include/transformer_engine/ep.h`	Extends expert-parallel configuration.
`transformer_engine/common/include/transformer_engine/comm_window.h`	Adds scale-window metadata.
`transformer_engine/common/include/transformer_engine/cast.h`	Extends grouped quantize-dbias API.
`transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu`	Cleans includes and comments.
`transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu`	Removes unused includes.
`transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu`	Cleans includes and comments.
`transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu`	Removes unused includes.
`transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu`	Cleans includes and comments.
`transformer_engine/common/fused_attn/fused_attn_fp8.h`	Extends FP8 attention metadata.
`transformer_engine/common/fused_attn/context_parallel.cu`	Adds FP8 copy/zero correction.
`transformer_engine/common/common.h`	Adds MXFP8 config and pointer alignment.
`transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh`	Supports row-scaled columnwise output.
`transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh`	Preserves shared-memory pointer provenance.
`transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh`	Uses shared-memory alignment helper.
`transformer_engine/common/cast/fp8/quantize_fp8.cuh`	Honors no-op flags in optimized kernel.
`transformer_engine/common/cast/fp8/gated_fp8.cuh`	Uses shared-memory alignment helper.
`transformer_engine/common/cast/cast_grouped_dbias.cu`	Forwards quantization configuration.
`transformer_engine/common/__init__.py`	Improves CUDA wheel header discovery.
`tests/pytorch/utils.py`	Adds block-scaling tolerances.
`tests/pytorch/test_onnx_export.py`	Adjusts FP16 export tolerances.
`tests/pytorch/test_fused_rope.py`	Expands long-table RoPE coverage.
`tests/pytorch/test_fused_optimizer.py`	Tests empty-group step tracking.
`tests/pytorch/test_float8blockwisetensor.py`	Tests invalid FSDP scale geometry.
`tests/pytorch/test_float8_current_scaling_exact.py`	Updates custom-recipe imports.
`tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py`	Tests stochastic FP4 adjacency.
`tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py`	Updates reference imports.
`tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py`	Covers row-scaled transpose support.
`tests/pytorch/nvfp4/test_nvfp4_group_quantize.py`	Adds grouped RHT coverage.
`tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py`	Updates reference imports.
`tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py`	Tests bilateral row-scaled GEMMs.
`tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py`	Tests padded swizzled scales.
`tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py`	Tests extreme E8M0 codes.
`tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py`	Strengthens checkpoint memory validation.
`tests/pytorch/distributed/test_torch_fsdp2.py`	Adds hybrid FSDP2 suites.
`tests/pytorch/distributed/test_sanity.py`	Prevents helper import shadowing.
`tests/pytorch/distributed/test_newton_schulz.py`	Consolidates distributed test launching.
`tests/pytorch/distributed/test_fusible_ops.py`	Prevents helper import shadowing.
`tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py`	Prevents helper import shadowing.
`tests/pytorch/distributed/test_ep.py`	Scales EP suite timeout.
`tests/pytorch/distributed/test_comm_gemm_overlap.py`	Improves subprocess isolation and diagnostics.
`tests/pytorch/distributed/run_test_ep.sh`	Adds eager, overflow, and MXFP8 passes.
`tests/pytorch/distributed/run_numerics_exact.py`	Updates custom-recipe imports.
`tests/pytorch/distributed/run_layer_with_overlap.py`	Records distributed worker failures.
`tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py`	Adds hybrid recipe helpers.
`tests/pytorch/distributed/fsdp2_tests/conftest.py`	Parameterizes hybrid recipes.
`tests/pytorch/debug/test_sanity.py`	Tests debug quantizer usage routing.
`tests/pytorch/attention/test_kv_cache.py`	Prevents helper import shadowing.
`tests/pytorch/attention/run_attention_with_cp.py`	Propagates explicit padding state.
`tests/jax/test_fused_attn_score_mod.py`	Updates fake attention signature.
`tests/jax/test_distributed_dense.py`	Adds compound DP/FSDP sharding.
`tests/jax/run_te_ep_moe.sh`	Allows test-file override.
`tests/jax/multi_process_launch_ep.sh`	Corrects editable-install guidance.
`tests/jax/distributed_test_base.py`	Adds eight-device mesh coverage.
`tests/cpp/test_common.cu`	Supports columnwise row-scaled NVFP4 tests.
`tests/cpp/operator/test_cast_float8blockwise_grouped.cu`	Tests power-of-two grouped scales.
`tests/cpp_distributed/test_ep_common.h`	Updates EP test hidden size.
`setup.py`	Builds the NCCL extensions submodule.
`README.rst`	Documents cuDNN frontend requirement.
`qa/L2_jax_unittest/test.sh`	Pins dependencies and timeout plugin.
`qa/L2_jax_distributed_unittest/test.sh`	Isolates softmax collective workaround.
`qa/L1_pytorch_hybrid_distributed_unittest/test.sh`	Adds hybrid distributed QA.
`qa/L1_pytorch_distributed_unittest/test.sh`	Separates hybrid and attention coverage.
`qa/L1_jax_distributed_unittest/test.sh`	Updates collective workaround.
`qa/L0_pytorch_debug_unittest/test.sh`	Selects FlashAttention generation.
`qa/L0_jax_unittest/test.sh`	Pins dependencies and removes duplicate flag.
`examples/pytorch/ep/run_test_ep.sh`	Removes source-tree path override.
`examples/pytorch/ep/ep_moe.py`	Updates EP buffer APIs.
`examples/pytorch/ep/bench/run_nccl_ep_bench.sh`	Uses NCCL extensions paths.
`examples/pytorch/ep/bench/run_ep_bench.sh`	Removes source-tree path override.
`examples/pytorch/ep/bench/ep_bench.py`	Updates EP buffer APIs.
`examples/jax/ep/run_test_ep.sh`	Corrects editable-install guidance.
`examples/jax/ep/ep_moe.py`	Handles receive-total output.
`examples/jax/ep/bench/ep_bench.py`	Handles receive-total output.
`examples/jax/encoder/requirements.txt`	Constrains NLTK version.
`docs/examples/te_jax_integration.rst`	Publishes attention tutorials.
`docs/examples/jax/attention.rst`	Adds JAX attention guide.
`docs/examples/jax/attention.out`	Adds tutorial output.
`docs/examples/jax/attention_context_parallel.out`	Adds CP tutorial output.
`docs/envvars.rst`	Documents new runtime controls.
`CONTRIBUTING.rst`	Corrects spelling.
`build_tools/VERSION.txt`	Advances version to 2.19 development.
`build_tools/build_ext.py`	Isolates non-incremental CMake builds.
`benchmarks/linear/benchmark_graph_safe_grouped_mlp.py`	Corrects benchmark command examples.
`.gitmodules`	Switches to NCCL extensions.
`.github/workflows/trigger-ci.yml`	Adds authorized CI users.
`.github/workflows/build.yml`	Updates CUDA, JAX, and NCCL environments.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+  /*! Whether to use 2D block scaling for MXFP8 */
+  kNVTEQuantizationConfigMXFP82DQuantization = 10,
 #ifdef USE_ROCM
  /*! Whether to apply Hadamard transform before MXFP4 quantization */
-  kNVTEQuantizationConfigMXFP4UseHadamard = 10,
+  kNVTEQuantizationConfigMXFP4UseHadamard = 11,


+  struct ncclWindow_vidmem*
+      scale_window;      /*!< Window for a block-scaled tensor's scale-inverse, or NULL for raw. */
+  uint64_t scale_offset; /*!< Byte offset of the scale-inverse within scale_window. */


 void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output,
-                               NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream);
+                               NVTEGroupedTensor dbias, NVTETensor workspace,
+                               const NVTEQuantizationConfig quant_config, cudaStream_t stream);


@@ -0,0 +1,117 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.


@@ -0,0 +1,28 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.


@@ -0,0 +1,33 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.


KshitijLakhani and others added 30 commits July 21, 2026 14:19

Changed VERSION to 2.19.0.dev0 (#3228)

870a68e

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Fix nccl-extensions submodule name (#3239)

8280f02

update nccl-ext submodule name Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Update list of authorized CI users (#3241)

8abfa2e

Signed-off-by: Tim Moon <tmoon@nvidia.com>

[Common] Fix Build: NCCL EP build to respect MAX_JOBS (#3138)

4adad4c

fix nproc Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Build Fix: Avoid invalid make -j 0 in NCCL EP submodule build (#3243)

c39f36d

Avoid invalid make -j 0 in NCCL EP submodule build Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[JAX] Keep MoE aux-loss cotangent scalar (#3237)

5d51790

[Common] Make MXFP8 specialized kernel respect noop flag as well (#3246)

ea41e08

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

[PyTorch] Allow multi-version Flash Attention tests to use checkpoint…

3c4f9af

…s with pickles (#3245) Allow Flash Attention tests to use checkpoints with pickles Signed-off-by: Tim Moon <tmoon@nvidia.com>

[CI] Fix old NVCC dependency in GH CI (#3230)

c7766dd

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

sudhakarsingh27 and others added 25 commits August 13, 2026 10:13

[JAX] Move ffi type definition before ffi target definitions (#3382)

86e3487

Move ffi type definition before ffi target definitions Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>

Fix FlashAttention local version validation (#3356)

a0a2d08

Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

[Pytorch] Optimize CPU overhead for NCCL EP eager mode (#3378)

e61a2fb

use pinned memory; use explicit cuda sync Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

Change VERSION to 2.19.0

55b2cba

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

[JAX] Support synchronous collectives in HLO collective bytes assert …

bcc72f3

…tests (#3398) Support synchronous collectives in HLO collective bytes assert tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> (cherry picked from commit 27b12df7674e7f95beb4d4adfbcfb29d2f880e52)

Loosen fp16 atol in ONNX export tests for onnxruntime 1.29 CPU numeri…

08d6906

…cs (#3409) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 17c351c34ccfb32506c7ef3cdefb4e51b4b4f686)

Disable async on distributed softmax test on 8x GPU jobs due to bug w…

ca49c3a

…ith 26.08 JAX image (2.19 only) Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

[JAX] Fix classification of sync-tagged collective start instructions…

42e953e

… (#3425) [JAX] Fix sync-tagged collective start classification Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com> (cherry picked from commit d7340a43c9f928f96884c3d842d2738406d6703f)

matthiasdiener self-assigned this Aug 28, 2026

matthiasdiener added the ci-level 3 CI test level 3 label Aug 28, 2026

matthiasdiener requested a balanced review from Copilot August 28, 2026 21:03

Copilot started reviewing on behalf of matthiasdiener August 28, 2026 21:04 View session

Copilot AI reviewed Aug 28, 2026

View reviewed changes

Navigation Menu

Uh oh!

Ifu dev 20260828 v2.19 - #721

Ifu dev 20260828 v2.19#721
matthiasdiener wants to merge 106 commits into
devfrom
IFU-dev-20260828-v2.19

matthiasdiener commented Aug 28, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants

		@@ -0,0 +1,117 @@
		# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

		@@ -0,0 +1,28 @@
		# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

		@@ -0,0 +1,33 @@
		# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Uh oh!

Conversation

matthiasdiener commented Aug 28, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants