[Common] Fix TMA synchronization in quantization kernels

[Oleg-Goncharov]

Description

This PR fixes shared-memory reuse races and incomplete completion handling in kernels using TMA shared-to-global transfers.

The issue was not isolated to a single kernel. The same TMA pipeline and synchronization pattern had been copied and adapted across multiple NVFP4, FP8, and MXFP8 kernels. As a result, the same underlying synchronization bug appeared in several related implementations.

Motivation and root cause

The affected kernels use shared-memory ring buffers as the source of asynchronous TMA shared-to-global transfers. A typical pipeline iteration performs the following operations:

1. All threads cooperatively write an output tile into shared memory.
2. A designated leader thread issues the TMA shared-to-global transfer.
3. The leader commits the transfer into a bulk async-group.
4. A later iteration reuses the same shared-memory buffer.

Bulk async-groups are maintained per issuing thread. Only the leader thread issues the TMA operation and calls cp_async_bulk_commit_group(). Therefore, only that thread has the corresponding outstanding async-groups.

Some kernels called cp_async_bulk_wait_group_read<N>() without a subsequent CTA-wide synchronization point before all threads started overwriting the reused shared-memory buffer. Calling the wait from non-issuing threads does not solve the problem because those threads do not own the leader's async-groups.

Consequently, the following race was possible:

• The leader observed that an older TMA operation had finished reading a shared-memory buffer.
• Other threads had no synchronization dependency on that observation.
• Those threads could begin overwriting the buffer while the TMA engine was still reading it.

This could result in rare, timing-dependent output corruption, particularly under global-memory backpressure.

The existing fence.proxy.async.shared::cta instructions were correct and remain necessary. They make cooperative generic-proxy shared-memory writes visible to the asynchronous TMA proxy before a transfer is issued. However, a proxy fence does not communicate the leader's later async-group completion to the rest of the CTA and therefore cannot by itself protect shared-memory buffer reuse.

In several kernels, the output-buffer wait was also coupled to the conditional prefetch of the next input stage. This meant that the wait could be skipped on the final pipeline stage even though output-buffer lifetime and input-prefetch state are independent concerns.

A related latent issue existed in the tuned NVFP4 kernel: the transposed output ring was hard-coded to two buffers while the allowed number of outstanding TMA groups was controlled by PREFETCH_STAGES. Increasing the prefetch depth could therefore make the ring smaller than required by the TMA wait depth.

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

This PR introduces the following changes:

• Added an explicit issuer-to-CTA completion handoff before shared-memory output buffers are reused:

  • The thread that issued and committed the TMA operation performs cp_async_bulk_wait_group_read<N>().
  • A subsequent CTA-wide __syncthreads() ensures that no cooperative writer can overwrite the buffer before the issuer has observed completion of the relevant TMA read.

• Moved output-ring synchronization out of input-prefetch conditionals.

  • Output-buffer lifetime is now handled independently from the availability of a next input stage.
  • This also covers the final pipeline iteration.

• Corrected wait depths to derive from the actual number of shared-memory output buffers rather than unrelated input-prefetch constants.

• Added full issuer-only cp_async_bulk_wait_group() drains at kernel tails.

  • This ensures that all committed TMA shared-to-global operations have completed before the CTA exits or releases associated shared-memory state.
  • The completion is followed by a CTA rendezvous where other threads may still access or release related shared-memory state.

• Fixed the tuned NVFP4 transposed-output ring size.

  • BUFFS_NUM_OUT_TR now follows PREFETCH_STAGES + 1.
  • Static assertions document and enforce the relationship between TMA wait depth and output-ring capacity.

• Added compile-time pipeline invariants for MXFP8 kernels.

  • Persistent grouped kernels now verify that resetting the shared-memory ring between jobs is compatible with the number of stages.
  • Kernels that currently rely on having no output-ring wrap explicitly enforce that constraint.

• Added compiler "memory" clobbers to the inline PTX wrappers for:

  • cp.async.bulk.commit_group
  • cp.async.bulk.wait_group
  • cp.async.bulk.wait_group.read
  • async proxy fences

  These clobbers prevent the compiler from moving memory operations across synchronization instructions whose memory effects are otherwise not visible to the C++ compiler.

• Replaced final source-read-only waits with full completion waits in one-shot TMA store paths, including fused grouped requantization and square-blockwise quantize-transpose.

• Audited all current TMA shared-to-global call sites.

  • The affected NVFP4, FP8, and MXFP8 paths were updated.
  • Existing warp-specialized paths that already use an explicit mbarrier-based producer/consumer handoff were preserved.
  • Two remaining wait_group_read<0>() calls are intentional intermediate waits protecting reuse of dynamic shared memory for dbias reduction; they are not incomplete kernel-tail drains.

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

[Oleg-Goncharov]

/te-ci

[greptile-apps]

Greptile Summary

The PR strengthens completion handling for TMA shared-to-global transfers so shared-memory output rings are not reused before the issuing thread observes completion.

• Adds issuer-only completion waits followed by CTA-wide handoffs before cooperative buffer reuse.
• Replaces kernel-tail read-completion waits with full completion drains.
• Derives wait depth and ring capacity from output-buffer counts and adds compile-time pipeline invariants.
• Adds compiler memory clobbers to bulk-async and async-proxy synchronization wrappers.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/util/ptx.cuh	Adds compiler memory clobbers to TMA bulk-group and async-proxy synchronization wrappers.
transformer_engine/common/cast/fp8/quantize_fp8.cuh	Aligns waits with output-ring capacity, hands issuer completion to cooperative writers, and fully drains tail transfers.
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh	Protects output-ring reuse with issuer-only waits and CTA-wide synchronization.
transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh	Adds convergent completion handoffs around specialized output-ring reuse and a full tail drain.
transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh	Moves output-buffer lifetime synchronization ahead of ring reuse and drains outstanding stores before teardown.
transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh	Sizes the transposed-output ring from prefetch depth and enforces the required compile-time relationship.
transformer_engine/common/cast/fused_group_requantize.cu	Upgrades the one-shot TMA store tail from source-read completion to full transfer completion.

Sequence Diagram

sequenceDiagram
    participant CTA as Cooperative CTA threads
    participant Leader as TMA issuer thread
    participant SMEM as Shared-memory output ring
    participant TMA as TMA engine
    CTA->>SMEM: Cooperatively write output tile
    CTA->>CTA: Async-proxy fence and rendezvous
    Leader->>TMA: Issue shared-to-global transfer
    Leader->>Leader: Commit bulk async-group
    Leader->>Leader: Wait until reusable group completes
    Leader->>CTA: CTA-wide completion handoff
    CTA->>SMEM: Reuse completed ring buffer
    Leader->>Leader: Fully drain remaining groups at tail
    Leader->>CTA: Final rendezvous

Loading

_{Reviews (3): Last reviewed commit: "Fix TMA shared-memory reuse synchronizat..." | Re-trigger Greptile}

[Oleg-Goncharov]

/te-ci