From 47a5ff603f918d1454ef4caa30e5ddf2f6bccfc1 Mon Sep 17 00:00:00 2001 From: ceo Date: Sat, 17 Oct 2026 22:35:08 +0900 Subject: [PATCH 1/2] [Common] Use .shared::cta for 2D TMA loads on sm_12x On sm_120/sm_121, ptxas lowers cp.async.bulk.tensor.2d with a .shared::cluster destination into a run-time branch: a native UTMALDG when the destination is the executing CTA, and otherwise a call to the driver syscall __cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast. The first launch of any kernel that contains that call raises the context's per-thread stack limit from 1,024 B to 14,624 B, and the driver keeps the matching local memory reserved for the lifetime of the context, outside the framework's caching allocator: 1,196 MiB on a 60-SM RTX 5080 Laptop GPU, 400 MiB on a 20-SM RTX 5050. On sm_90a and sm_100a both destination forms compile to the same native code. cp_async_bulk_tensor_2d_global_to_shared always targets the executing CTA's own shared memory (its addresses come from __cvta_generic_to_shared), so on sm_12x it can issue the .shared::cta form, which cp_async_bulk_tensor_2d_global_to_shared_cta already uses. This covers every TMA cast kernel that loads through the helper or through copy_2d_to_shared{,x2,x3}: NVFP4, MXFP8, FP8 and the Hadamard amax kernels. Other architectures compile unchanged. Add an sm_12x regression test that checks, in a fresh process, that the first NVFP4 cast (1D and 2D block scaling) leaves the stack limit unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TadBEfH5zvffbgRV7hesgt Signed-off-by: ceo --- .../nvfp4/test_sm12x_tma_stack_limit.py | 72 +++++++++++++++++++ transformer_engine/common/util/ptx.cuh | 15 ++++ 2 files changed, 87 insertions(+) create mode 100644 tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py diff --git a/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py b/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py new file mode 100644 index 00000000000..28e8bafec30 --- /dev/null +++ b/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py @@ -0,0 +1,72 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Regression test: TMA cast kernels must not raise the CUDA stack limit on SM120/SM121. + +On sm_12x, ptxas lowers a cp.async.bulk.tensor load with a .shared::cluster destination to a +run-time branch whose remote-CTA side is a driver syscall. The first launch of a kernel that +contains it raises the context's per-thread stack limit (1 KB -> ~14.6 KB), and the driver keeps +the matching local memory reserved for the lifetime of the context. The limit is per context and +never shrinks, so each check runs in a fresh process. +""" + +import subprocess +import sys +import textwrap + +import pytest +import torch + +import transformer_engine.pytorch as te + +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +_CHILD = textwrap.dedent( + """ + import ctypes + import sys + + import torch + import transformer_engine.pytorch as te + + libcuda = ctypes.CDLL("libcuda.so.1") + + + def stack_limit(): + torch.cuda.synchronize() + value = ctypes.c_size_t() + assert libcuda.cuCtxGetLimit(ctypes.byref(value), 0) == 0 # CU_LIMIT_STACK_SIZE + return value.value + + + x = torch.randn(1024, 2048, device="cuda", dtype=torch.bfloat16) + quantizer = te.NVFP4Quantizer( + rowwise=True, columnwise=True, with_2d_quantization=sys.argv[1] == "2d" + ) + before = stack_limit() + quantizer(x) + print(before, stack_limit(), flush=True) + """ +) + + +def _is_sm12x() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12 + + +@pytest.mark.skipif(not _is_sm12x(), reason="Only SM120/SM121 lower cluster TMA loads to a syscall") +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +@pytest.mark.parametrize("scaling", ["1d", "2d"]) +def test_nvfp4_cast_keeps_stack_limit(scaling: str) -> None: + """The first NVFP4 cast in a process must not raise the per-thread stack limit.""" + result = subprocess.run( + [sys.executable, "-c", _CHILD, scaling], capture_output=True, text=True, check=False + ) + fields = result.stdout.split() + assert len(fields) >= 2, f"child failed (exit {result.returncode}):\n{result.stderr[-2000:]}" + before, after = (int(v) for v in fields[-2:]) + assert after == before, ( + f"The first {scaling} NVFP4 cast raised the per-thread stack limit from {before} to" + f" {after} bytes" + ) diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 74fd2261f58..51f7e8ee400 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -237,11 +237,26 @@ __device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( // barrier condition: // - leader must arrive (i.e. 1 thread as set above) // - TMA hardware substracts bytes from expect_tx counter, must reach zero +#if __CUDA_ARCH__ >= 1200 + // sm_12x: use the .shared::cta destination. dst_shmem and mbar are always in the executing + // CTA, so both forms copy the same bytes. For .shared::cluster, ptxas emits a run-time branch + // whose remote-CTA side calls a driver syscall + // (__cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast). The first launch of any kernel that + // contains it raises the per-thread stack limit from 1 KB to ~14.6 KB, and the driver keeps + // the matching local memory (~1.2 GiB on a 60-SM RTX 5080 Laptop GPU) reserved for the life + // of the context. + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global.tile" + ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_shmem_ptr), + "l"(tensor_map_ptr), "r"(offset_x), "r"(offset_y), "r"(mbar_ptr) + : "memory"); +#else asm volatile( "cp.async.bulk.tensor.2d.shared::cluster.global.tile" ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_shmem_ptr), "l"(tensor_map_ptr), "r"(offset_x), "r"(offset_y), "r"(mbar_ptr) : "memory"); +#endif // #if __CUDA_ARCH__ >= 1200 #else NVTE_DEVICE_ERROR("cp_async_bulk_tensor_2d_global_to_shared is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) From 5a2365240245bb649dd2fb737a7be22d2082e07a Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Wed, 4 Nov 2026 13:41:41 -0800 Subject: [PATCH 2/2] Remove SM12x TMA stack-limit regression test Signed-off-by: Przemek Tredak --- .../nvfp4/test_sm12x_tma_stack_limit.py | 72 ------------------- 1 file changed, 72 deletions(-) delete mode 100644 tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py diff --git a/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py b/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py deleted file mode 100644 index 28e8bafec30..00000000000 --- a/tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py +++ /dev/null @@ -1,72 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Regression test: TMA cast kernels must not raise the CUDA stack limit on SM120/SM121. - -On sm_12x, ptxas lowers a cp.async.bulk.tensor load with a .shared::cluster destination to a -run-time branch whose remote-CTA side is a driver syscall. The first launch of a kernel that -contains it raises the context's per-thread stack limit (1 KB -> ~14.6 KB), and the driver keeps -the matching local memory reserved for the lifetime of the context. The limit is per context and -never shrinks, so each check runs in a fresh process. -""" - -import subprocess -import sys -import textwrap - -import pytest -import torch - -import transformer_engine.pytorch as te - -nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) - -_CHILD = textwrap.dedent( - """ - import ctypes - import sys - - import torch - import transformer_engine.pytorch as te - - libcuda = ctypes.CDLL("libcuda.so.1") - - - def stack_limit(): - torch.cuda.synchronize() - value = ctypes.c_size_t() - assert libcuda.cuCtxGetLimit(ctypes.byref(value), 0) == 0 # CU_LIMIT_STACK_SIZE - return value.value - - - x = torch.randn(1024, 2048, device="cuda", dtype=torch.bfloat16) - quantizer = te.NVFP4Quantizer( - rowwise=True, columnwise=True, with_2d_quantization=sys.argv[1] == "2d" - ) - before = stack_limit() - quantizer(x) - print(before, stack_limit(), flush=True) - """ -) - - -def _is_sm12x() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12 - - -@pytest.mark.skipif(not _is_sm12x(), reason="Only SM120/SM121 lower cluster TMA loads to a syscall") -@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) -@pytest.mark.parametrize("scaling", ["1d", "2d"]) -def test_nvfp4_cast_keeps_stack_limit(scaling: str) -> None: - """The first NVFP4 cast in a process must not raise the per-thread stack limit.""" - result = subprocess.run( - [sys.executable, "-c", _CHILD, scaling], capture_output=True, text=True, check=False - ) - fields = result.stdout.split() - assert len(fields) >= 2, f"child failed (exit {result.returncode}):\n{result.stderr[-2000:]}" - before, after = (int(v) for v in fields[-2:]) - assert after == before, ( - f"The first {scaling} NVFP4 cast raised the per-thread stack limit from {before} to" - f" {after} bytes" - )