Ifu dev 20260706 v2.17

[AllenFarcas]

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

[github-actions]

Claude Walkthrough

Intent. IFU (integrate-from-upstream) of NVIDIA TransformerEngine release_v2.17 (156 commits, v2.15 → v2.17) into the ROCm dev branch. The bulk of the diff (~453 files, +70k/-18k) is the upstream merge itself; the rest is a build-fix commit plus ~25 targeted ROCm-specific patches to make v2.17 compile, run, and pass CI on gfx942/gfx950, plus a CI infrastructure refactor.

Key changes.

• Upstream v2.17 merge (fbcaca7a): 78 conflicts resolved. Notable reconciliations — MXFP4 enum values shifted to 10/11 so upstream NVFP4 can take 8/9; _Linear refactored onto upstream's LinearFwdArgs/LinearBwdArgs dataclasses with ROCm fields folded in; the test Tensor helper adopted upstream's Buffer/CudaPtr model; new upstream NCCL Expert-Parallelism subsystem force-disabled on ROCm.
• Build-fix pass (0cf6b69c): guards for CUDA-only kernels newly introduced upstream — row-scaled NVFP4 amax (transformer_engine/common/cast/dispatch/quantize.cuh), variable-shape grouped swizzle (transformer_engine/common/swizzle/swizzle.cu), TMA multi-tensor transpose in transformer_engine/common/fused_attn/flash_attn.cu, <cuda_pipeline.h> in transformer_engine/common/fused_router/async_loader.h. Restores sum()/max() helpers dropped by upstream's warp-reduce rewrite and gates warp-shuffle widths for wave64. Reconciles nvte_fused_attn_fwd/bwd with new format-descriptor params. MXFP4Quantizer updated to the new Quantizer virtuals.
• CI refactor (71eb59d6): the sGPU test runner previously inlined in .github/workflows/rocm-ci.yml is extracted into .github/scripts/run_parallel_sgpu.sh + .github/scripts/ci_sgpu_jobs.conf (see .github/workflows/rocm-ci.yml:188); adds progress logging, persistent JAX compilation cache, and per-suite .rc files instead of FAIL_* sentinels.
• NVFP4 correctness fixes: portable nvfp4::compute_rowwise_amax extracted so ROCm can enable row-scaled NVFP4 quantization (transformer_engine/common/cast/nvfp4/rowwise_amax_nvfp4.cuh, 93f88883); arch-aware E4M3 max restored in the fnuz dequant scale (transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh, 4e94fd7a); fnuz-aware reference global scale (transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py, be3feea9); inplace_swizzle_scale_for_gemm ROCm early-return to fix NoneType scale_inv crash (transformer_engine/pytorch/csrc/extensions/swizzle.cpp, bebd2475); CPU-offload dequant recursion fix (9c2bbdad); accept NVFP4TensorStorage in update_quantized (cdd4e9ca).
• Fused-attention fixes: neutralize non-finite per-step LSE/grads under is_hip_extension() guard for gfx942 striped ring-attention with sliding window (transformer_engine/jax/cpp_extensions/attention.py, 19d3daec); bwd layout fix for CPU-offload reload (a6bbb7c9); don't force FlashAttention THD pad-between-seqs leg on ROCm (990bf8bf).
• Other correctness: FSDP2 all-gather thread scale_inv through metadata for Float8 current-scaling (transformer_engine/pytorch/tensor/fsdp2_allgather_tensor.py, 2f2b5fc6); route grouped-tensor GEMM through te_general_grouped_gemm on ROCm (306a371c); add missing roles= param to MXFP4BlockScalingRecipeState (c99b4d74); Triton current-scaling quantize fix (transformer_engine/pytorch/triton_kernels/cast.py, fdc9510a); circular import fix in transformer_engine/pytorch/quantized_tensor.py (f88abf74).
• Test guards / xfails for ROCm-unsupported v2.17 features: 4over6 NVFP4, grouped NVFP4, Blackwell-only RHT swizzled-SF, standalone JAX TopK, NVFP4 async DCP parity, primary-FP8 + FSDP2 forward-memory accumulation, large-M grouped-GEMM allocations, cpp operator tests missing scale accessors.

Walkthrough.

The overwhelming majority of the diff is the mechanical upstream merge; the interesting parts are what the ROCm fork had to change on top.

C++ common library. The v2.17 merge brings row-scaled NVFP4, a new grouped-swizzle kernel with variable shapes, a TMA-based multi-tensor transpose path, and a new fused-router API. The build-fix commit guards these behind USE_ROCM/__HIP_DEVICE_COMPILE__ where the underlying primitives (Blackwell CUtensorMap, wave32-only shuffles, <cuda_pipeline.h>) have no HIP equivalent. Row-scaled NVFP4 is then unlocked in a follow-up (93f88883) by extracting the amax kernel out of the Blackwell-only quantize-transpose header into rowwise_amax_nvfp4.cuh, since the ROCm blockwise quantize/dequantize kernels already handle row-scaled inputs — this was the piece answering a review comment about "did we not support rowwise NVFP4?" (see the tip commit 6369177f). Separately, the fnuz dequant fix (4e94fd7a) is subtle: upstream hardcoded E4M3_MAX=448 in the dequant factor, but fnuz archs have maxNorm=240, so the un-fixed code was scaling outputs by ~1.87×; the fix restores detail::TypeExtrema<fp8e4m3>::max on HIP device compilation.

Fused attention. transformer_engine/common/fused_attn_rocm/fused_attn.cpp was reconciled with upstream's new format-descriptor params on nvte_fused_attn_fwd/bwd (strict on layout, permissive on FP8 scale-inv). transformer_engine/common/include/transformer_engine/fused_attn.h retains the NVTE_F16_max512_seqlen enumerator with a restored explanatory comment (6369177f) — upstream v2.17 dropped the max512 impl/dispatch but the enumerator is kept so the ROCm fork's CUDA-path bindings still compile. The JAX ring-attention NaN fix (19d3daec) neutralizes non-finite per-step output/LSE (fwd) and grads (bwd), gated on is_hip_extension() and a no-op on finite values, fixing 5 test_context_parallel_ring_attn failures with window_size=(20,0) + THD_SEPARATE + PADDING_CAUSAL.

PyTorch bindings. MXFP4Quantizer::create_tensor / create_grouped_tensor were rewritten to match the new base Quantizer virtual signatures (device/pin_memory, precomputed tensor_offsets). Multiple entry points that only exist upstream (grouped-tensor GEMM, fused grouped-linear path, primary-FP8 + FSDP2 fwd memory accumulation) are guarded so ROCm falls through to the general grouped-gemm path.

FSDP2 all-gather. transformer_engine/pytorch/tensor/fsdp2_allgather_tensor.py previously reconstructed scale_inv via 1 / quantizer.scale, which crashed on Float8CurrentScalingQuantizer because current scaling computes its scale in-kernel and exposes no .scale attribute. Since fsdp_pre_all_gather already runs an amax reduction (every rank's per-tensor scale_inv is identical), the fix threads the scale_inv through the all-gather metadata and uses it in post_all_gather for the current-scaling path; delayed scaling is unchanged.

CI infrastructure. The old inline heredoc in rocm-ci.yml that fanned out Torch/JAX/examples/core on separate GPUs is replaced by a general driver run_parallel_sgpu.sh reading a job table from ci_sgpu_jobs.conf, plus run_examples.sh for the examples slot. Per-job .rc files replace the FAIL_* sentinel scheme. Timeouts adjusted (sGPU 360→270 min but inner step 180→240 min; mGPU 360→210 min). JAX picks up a persistent compilation cache.

Testing. No new test functionality is added; the delta is mostly guards — skips and xfails for ROCm-unsupported v2.17 features (grouped NVFP4, 4over6, Blackwell-only RHT, standalone JAX TopK, primary-FP8+FSDP2 fwd mem accumulation, NVFP4 async DCP parity under torch 2.8), plus stale -k filter fixes after upstream's recipe-id rename (bd8c4042). Row-scaled NVFP4 is stated verified by OperatorTestRowScaled/FusedCastTransposeNVFP4TestSuite (32/32). The upstream merge commit explicitly notes it was static-validated only (no GPU in the resolution env); the follow-up ~25 commits are the observed-failure fixes from actually running CI.

Notes for reviewers.

• Base is dev, not main. This large IFU will land on dev; expect a subsequent dev → main promotion after downstream soak.
• ROCm guard style is inconsistent by design: some guards are USE_ROCM (C++ preprocessor), some are IS_HIP_EXTENSION (build-time Python), some are is_hip_extension() (runtime). The build-fix commit normalizes guard comments to "Disabled on ROCm" but does not homogenize the guard mechanism — follow each file's existing convention per CLAUDE.md.
• MXFP4 enum renumbering (8/9 → 10/11) is a wire-format concern: any external consumer serializing NVTEDType values will need to be aware. NVFP4 now occupies the old MXFP4 slots.
• New CUDA-only subsystems (NCCL Expert Parallelism, upstream's grouped-tensor GEMM path, TMA transpose in flash_attn.cu) are stubbed/disabled — porting them to ROCm is deferred.
• The upstream merge commit's own note ("full ROCm HIP build + GPU runtime tests not yet run") is now outdated; subsequent CI runs on this branch drove the ~25 fix commits listed above. Treat CI-green on this PR as the actual signal.

---

Generated by Claude. To request a code review, comment /claude review.

[github-actions]

Automated review summary

Scope reviewed: ROCm-side commits on top of the NVIDIA v2.17 upstream merge (commits fbcaca7a..6369177f, ~26 commits touching ~69 files). Upstream NVIDIA content brought in by the merge itself was not re-reviewed.

Code review: no additional correctness / safety findings beyond what @wenchenvincent has already flagged. The 21 existing inline questions on backend selection, ROCm-specific skips, fnuz handling, upstream API deltas, and dqkv layout still need author responses — this review does not duplicate them.

Copyright headers: 11 files need updates — see inline comments. All are ROCm-modified/added files whose headers still carry only the NVIDIA line. Per repo policy, an AMD line at the current year (2026) should sit above the NVIDIA line, and new files added by ROCm need both headers.

[ipanfilo]

Should be only False for ROCm

[AllenFarcas]

Gated use_cublasmp to (False,) if IS_HIP_EXTENSION else (False, True) on all six comm/gemm-overlap test parametrizations. On ROCm only the False cases are generated; CUDA still gets False+True.

[ipanfilo]

TODO probably is not applicable - we know why the code is different. So it should rather be other comment, similar to test_cast_master_weights_to_fp8() which later solved the same problem

[ipanfilo]

The message should rather say it is not supported on ROCm

[AllenFarcas]

Corrected. Applied the same correction to the sibling module-level skip in test_nvfp4_rht_quantize_swizzle_fusion.py for consistency.

[wangye805]

nit: usually for those nv side feature we didn't enabled during IFU, we put TODO so later grep can be easier. But if you keep a list of NV side feature then this is also fine

[AllenFarcas]

I enriched the jax.py comment to explain the gate, matching the fuller comment already in build_tools/pytorch.py, and kept # Disabled on ROCm as the grep marker.

[wangye805]

Why the model 126m cannot run mxfp8 recipe?

[AllenFarcas]

126m does run mxfp8

The parametrize now uses ids=recipe_id which produces a readable name based id from the recipe class instead of recipe0, so the first recipe now shows up as MXFP8BlockScaling. I updated the filter from recipe0 to MXFP8BlockScaling so it keeps selecting the same MXFP8 cases.

[wangye805]

Why previously we can use const Tensor& but now need Tensor&? do we edit some contents in test

[AllenFarcas]

Upstream dropped the const qualifier from to_cpu() and the *_cpu_dptr() accessors in tests/cpp/test_common.h

compareResults_mxfp4 calls those on test, so once they became non-const the const Tensor& parameter stopped compiling and had to become Tensor&.

[wangye805]

Is this a new test from upstream? If so, can you file a ticket and put a TODO comment with this link here

[AllenFarcas]

Working on it right now, will push an updated commit soon.

[wangye805]

Need similar todo to github ticket for this fsdp2 issue

[AllenFarcas]

Working on it right now, will push an updated commit soon.

[wangye805]

Why this file looks like a new file?

[AllenFarcas]

Row-scaled NVFP4 is new in upstream v2.17, so the IFU pulled it in and ROCm needed the row-wise amax to support it but upstream keeps that kernel inside the Blackwell-only quantize_transpose_nvfp4.cuh, which doesn't compile on ROCm, so instead of taking it as-is we extracted the portable amax into this new header.

[wangye805]

I think this issue is worth a new git ticket for fully understanding.

If it's a CK issue in gfx942, we may need to report to CK team.

[wangye805]

Try hipcub

[AllenFarcas]

Will implement as a feature after the IFU is merged.

[wangye805]

Is this a bug fix over NV side?

[AllenFarcas]

Yes, bug fix over NV. Upstream's tensor.to("cuda") on the NVFP4 subclass re-enters __torch_dispatch__ -> dequantize -> .to(), recursing infinitely when running the test_dcp_output_parity[NVFP4BlockScaling-sync] test.

We move the raw buffers via the metadata dict and rebuild NVFP4TensorStorage directly to avoid the recursion. Worth upstreaming.

[wangye805]

Why the original .shape need to be changed to .size()?

[AllenFarcas]

tensor is a quantized storage subclass (NVFP4/MXFP8) which overrides size() to infer the true logical shape from its raw packed buffers, while .shape isn't overridden, so it returns the raw packed storage dims and can be wrong or warn when only partial data is present.

_setup_conditional_transpose_storage needs the logical shape to size the transpose buffer, so we switched to the storage-safe .size() (and len(shape) instead of .ndim).

When running test_quantize in tests/pytorch/triton_kernels/test_cast.py I caught this.

[AllenFarcas]

@ipanfilo @wangye805 related to your comment on tests/pytorch/distributed/test_fusible_ops.py line 1049 "TODO probably is not applicable - we know why the code is different. So it should rather be other comment, similar to test_cast_master_weights_to_fp8() which later solved the same problem"

I cannot comment there but ill add here my reply: I replaced the TODO with an explanatory comment in the # ROCm: style used in test_cast_master_weights_to_fp8.py: # ROCm: fnuz FP8 (gfx942) caps E4M3 at 240; OCP (gfx950+/CUDA) caps at 448., so the reference max must branch on is_fp8_fnuz(). Also removed an unused transformer_engine_torch import left behind.