Fix SM120 FP8 output ordering before DelayedScaling amax update

[AlbertYang514]

Summary

Add device-scope ordering between the FP8 output stores and the subsequent DelayedScaling amax update in the SM120 NVRTC optimized cast_transpose kernel.

Problem

On SM120, a checkpointed DelayedScaling workload can intermittently corrupt the first FP8 dgrad GEMM after activation recomputation. The loss and gradients can remain finite, and the affected layer or microbatch can vary with execution timing, but the parameter gradient norm can be transiently far outside its normal range.

This change is scoped to that observed optimized RTC path. It does not claim to fix every FP8 gradient anomaly or to affect every GPU architecture.

Root cause

The optimized cast_transpose kernel writes both rowwise and columnwise FP8 outputs and then performs a relaxed atomic update of the DelayedScaling amax value. The atomic update does not by itself order the preceding ordinary global-memory output stores. Without a device-scope fence, a later dgrad consumer can intermittently observe incorrect FP8 output visibility even though a later execution with the same operands is normal.

Fix

Insert a single device-scope fence after the output stores and before the amax reduction/atomic:

__threadfence();

There is no host synchronization, stream synchronization, public API change, or change to the DelayedScaling math.

Reproducer

The new standalone Transformer Engine pytest uses official PyTorch APIs on SM120 with:

• te.Linear modules and DelayedScaling;
• TE activation checkpointing;
• a large shape and 64-microbatch accumulation;
• three fixed seeds;
• finite-loss, input-gradient, parameter-gradient, and gradient-norm checks.

The same test fails for all three seeds without the fence and passes for all three seeds with the fence.

Validation

• Regression before fix: 3/3 seeds fail with finite but corrupted gradients.
• Regression after fix: 3/3 seeds pass.
• 14-layer / 98-Linear soak: 10/10 seeds pass; final gradient norm is 0.00529164-0.00529794.
• 28-layer / 196-Linear soak: 3/3 seeds pass; final gradient norm is 0.00757305-0.00757593.
• Checkpoint-off, TE reentrant, TE non-reentrant, TE single-segment, and native PyTorch reentrant controls pass.
• BF16 and CurrentScaling controls pass.
• Related Float8Tensor/FP8Recipe tests: 77 passed, 2 skipped.
• Activation-recompute FP8 update-ownership tests: 8 passed.
• DelayedScaling, CurrentScaling, cold-transition, optimizer, and _extra_state round trips are exactly consistent where exact equality is required.
• The two independent fixes were validated together on an external 1.7B decoder workload for independent 10- and 50-optimizer-step replays: all 3,840 microbatches had one forward and one backward global FP8 update, gradients remained finite, and no checkpoint was written.

Scope and risk

The production change is limited to the optimized RTC cast_transpose path when an amax pointer is present. The fence has explicit device-scope ordering semantics and does not alter the values computed by the kernel or any API. The main risk is the localized cost of a device fence; control paths and the larger external workload remained stable in validation.

Relationship to #3213

This issue is independent of #3213.

#3213 fixes duplicate global FP8 state-update ownership during activation checkpoint recomputation. This PR fixes missing device-scope ordering between FP8 output stores and the subsequent DelayedScaling amax atomic update in the SM120 NVRTC optimized cast-transpose kernel.

This PR is based directly on the current upstream main branch and does not depend on #3213. The two fixes were also validated together on the original external workload.

Documentation reference

CUDA documents that legacy atomic functions use relaxed memory ordering and do not act as memory fences, while __threadfence() provides device-scope ordering:

• Legacy atomic functions
• Memory fence functions

[greptile-apps]

Greptile Summary

This PR fixes a device-scope memory ordering race in the SM120 NVRTC-optimized cast_transpose kernel by inserting a single __threadfence() between the FP8 output stores and the relaxed atomicMaxFloat amax update, and adds a targeted regression test to validate the fix.

• Kernel fix: __threadfence() is correctly placed after the last __syncthreads() in the transposed-output loop and before reduce_max + atomicMaxFloat, ensuring all global writes to output_c and output_t are device-visible before the amax atomic is published. All threads call the fence (since amax_ptr is a uniform parameter), and the subsequent warp-level reduction in reduce_max provides the necessary inter-thread ordering before thread 0 reaches the atomic.
• Test: The new pytest creates a 14-layer SM120 stress network with activation checkpointing and DelayedScaling, runs three seeds × 3 steps × 64 microbatches, and checks per-microbatch corruption bounds plus a final gradient-norm range (0.001 < final_norm < 0.05). The statically compiled non-RTC fallback in cast_transpose.cu has the same structural pattern (global stores → __syncthreads() → atomicMaxFloat) without a __threadfence(), so it carries a latent ordering risk on any path where the RTC kernel is not selected.

Confidence Score: 5/5

Safe to merge; the kernel change is a single targeted fence with no effect on computed values or public APIs, and the test validates the specific failure mode.

The production change is minimal — one __threadfence() call gated on an existing null-pointer check — and its placement (after all block-scope syncs, before the amax atomic) is correct. The validation and new regression test cover the specific failure scenario.

The statically compiled fallback in transformer_engine/common/transpose/cast_transpose.cu (lines 203-209) has the same pre-fence pattern and would benefit from the same treatment if the RTC path is ever unavailable on SM120.

Important Files Changed

Filename	Overview
transformer_engine/common/transpose/rtc/cast_transpose.cu	Inserts __threadfence() before the amax reduction/atomic to enforce device-scope ordering between all FP8 output stores and the DelayedScaling amax update; placement after the last __syncthreads() and before reduce_max is correct.
tests/pytorch/test_delayed_scaling_dgrad_ordering.py	New SM120-only regression test using 14-layer stress network with activation checkpointing and DelayedScaling; checks per-microbatch corruption bounds and a final gradient-norm range; the final norm assertion bounds (0.001, 0.05) are appropriately widened from an earlier narrow version.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant K as cast_transpose kernel (SM120)
    participant GM as GPU Global Memory
    participant Atom as atomicMaxFloat (amax_ptr)
    participant DGEMM as dgrad GEMM kernel

    K->>GM: store output_c (rowwise FP8)
    K->>GM: store output_t (columnwise FP8, via shmem)
    Note over K: __syncthreads() — block-scope barrier only
    Note over K: __threadfence() — device-scope fence (NEW)
    K->>Atom: reduce_max + atomicMaxFloat(amax_ptr, amax)
    DGEMM->>GM: read FP8 outputs (guaranteed visible after fence)
    DGEMM->>Atom: read amax_ptr (for next-step scaling)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant K as cast_transpose kernel (SM120)
    participant GM as GPU Global Memory
    participant Atom as atomicMaxFloat (amax_ptr)
    participant DGEMM as dgrad GEMM kernel

    K->>GM: store output_c (rowwise FP8)
    K->>GM: store output_t (columnwise FP8, via shmem)
    Note over K: __syncthreads() — block-scope barrier only
    Note over K: __threadfence() — device-scope fence (NEW)
    K->>Atom: reduce_max + atomicMaxFloat(amax_ptr, amax)
    DGEMM->>GM: read FP8 outputs (guaranteed visible after fence)
    DGEMM->>Atom: read amax_ptr (for next-step scaling)

Loading

_{Reviews (2): Last reviewed commit: "test: relax SM120 dgrad norm bounds" | Re-trigger Greptile}