[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 repairs synchronization around asynchronous TMA shared-to-global stores in FP8, MXFP8, and NVFP4 quantization kernels.

• Makes the issuing thread wait for its thread-local bulk async groups and hands completion to cooperative writers through CTA-wide barriers.
• Separates output-ring reuse synchronization from input prefetching and adds complete tail drains.
• Aligns wait depths and ring capacities, adds compile-time pipeline invariants, and adds compiler memory clobbers to PTX synchronization wrappers.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness, build, or security regression identified in the changed synchronization paths.

The changed kernels consistently make the TMA issuer observe completion before a CTA rendezvous permits shared-buffer reuse, and the examined ring sizes, wait depths, control flow, tail drains, and shared-memory calculations remain internally consistent.

Important Files Changed

Filename	Overview
transformer_engine/common/util/ptx.cuh	Adds compiler memory clobbers to bulk async-group and async-proxy synchronization wrappers without changing their current emitted PTX immediates.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Restricts async-group waits to the issuing thread, publishes completion through CTA barriers, drains stores at the persistent-kernel tail, and enforces the current ring-reset invariant.
transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh	Moves output-ring completion handling ahead of cooperative reuse and replaces incomplete tail read waits with full issuer-side drains.
transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh	Protects output-ring reuse independently of prefetch availability and fully drains TMA stores before barrier destruction.
transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh	Derives both output-ring capacities from prefetch depth while keeping kernel layout and launch-time shared-memory sizing consistent.
transformer_engine/common/cast/fp8/quantize_fp8.cuh	Corrects ring-derived wait depths, adds issuer-to-CTA completion handoffs, and fully drains outstanding stores at kernel completion.

Sequence Diagram

sequenceDiagram
  participant CTA as Cooperative CTA threads
  participant Leader as TMA-issuing leader
  participant TMA as TMA engine
  CTA->>CTA: Populate shared output tile
  CTA->>Leader: CTA synchronization
  Leader->>TMA: Issue shared-to-global transfer
  Leader->>Leader: Commit bulk async group
  Leader->>Leader: Wait before ring-buffer reuse
  TMA-->>Leader: Shared-memory read completed
  Leader->>CTA: CTA-wide completion handoff
  CTA->>CTA: Safely overwrite reused buffer
  Leader->>TMA: Full tail drain
  TMA-->>Leader: All transfers completed
  Leader->>CTA: Final CTA rendezvous

Loading

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