[Common] Use .shared::cta for 2D TMA loads on sm_12x

[kimchi-developer]

Description

On sm_120 and sm_121 (RTX 50 series, RTX PRO Blackwell, GB10), the first launch of a TE cast kernel that loads
through ptx::cp_async_bulk_tensor_2d_global_to_shared raises the CUDA context's per-thread stack limit from
1,024 B to 14,624 B. That covers every NVFP4, MXFP8 and FP8 TMA cast kernel and the Hadamard amax kernels. The
driver then keeps local memory for that stack resident for the life of the context. On an RTX 5080 Laptop GPU
(60 SMs), the first NVFP4 quantization takes 1,196 MiB of device memory (400 MiB on a 20-SM RTX 5050), none of
it visible to PyTorch's allocator (torch.cuda.memory_reserved() does not grow). The model does not need that
memory: it is reserved for a driver syscall that TE never calls.

This PR makes the helper use the .shared::cta destination on sm_12x. The stack limit then stays at 1,024 B, that
memory stays free, and outputs are bitwise identical.

Related: #2255 (NVFP4 on RTX 50 series), #3300 (NVFP4 RHT for SM120/SM121, whose sm_12x path loads through this
helper and is also covered).

Root cause

cp_async_bulk_tensor_2d_global_to_shared (transformer_engine/common/util/ptx.cuh) issues

cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [dst], [tensorMap, {x, y}], [mbar];

On sm_12x, ptxas lowers the .shared::cluster destination into a run-time branch. A destination in the executing CTA
gets a native UTMALDG; any other destination calls (CALL.ABS.NOINC) the driver-provided
__cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast, which the kernel records as an EIATTR_EXTERNS entry. TE
never takes the remote branch, because its destinations are always CTA-local. Still, any kernel that contains the
call is launched with the syscall's stack requirement. At that first launch the driver raises the context-wide stack
limit and reserves local memory for it: roughly SMs × max resident threads per SM × stack size, here
60 × 1,536 × 14,624 B. Both last for the lifetime of the context.

The lowering reproduces without TE (tma_forms.cu, attached: one kernel per destination form; nvcc 13.4.92, -O3):

target	.shared::cta	.shared::cluster
sm_90a	56 SASS instructions, UTMALDG, no call, REG 10, STACK 0	identical SASS
sm_100a	56 SASS instructions, UTMALDG, no call, REG 10, STACK 0	identical SASS
sm_120, sm_120a, sm_120f, sm_121a	56 SASS instructions, UTMALDG, no call, REG 10, STACK 0	80 SASS instructions: CTA-rank check, UTMALDG and CALL.ABS.NOINC → __cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast; REG 24, STACK 8

Runtime of the same pair (tma_stack_repro.cu, NVIDIA GeForce RTX 5050, 20 SMs, CUDA 13.4): .shared::cta keeps the stack limit at 1024 → 1024 B with 0.0 MiB taken; .shared::cluster raises it 1024 → 14624 B and the device loses 400.0 MiB. Both copies are verified correct, and cudaFuncGetAttributes().localSizeBytes reports only 8 B for the cluster kernel.

Fix

dst_shmem and mbar always come from __cvta_generic_to_shared, which yields .shared::cta window addresses. A
.shared::cta address used in the .shared::cluster form already refers to the executing CTA's own shared memory.
So on sm_12x the helper can issue the .shared::cta form of the same copy: same bytes, same mbarrier transaction
count, no remote path. TE already uses that form in cp_async_bulk_tensor_2d_global_to_shared_cta (FP8
block-scaling grouped quantize), and the 1D bulk-copy helper cp_async_bulk_tensor_1d_global_to_shared uses the
.shared::cta destination on every arch ≥ sm_100. The .shared::cta destination needs PTX ISA 8.6 (CUDA 12.8),
which every sm_12x build already requires.

The change is confined by the preprocessor to __CUDA_ARCH__ >= 1200, so sm_90/100/103/107/110 code compiles exactly
as before. In the minimal test, sm_90a and sm_100a produce identical SASS for both forms. Switching the helper to
.shared::cta unconditionally, like the 1D helper, would therefore be equally valid if maintainers prefer a single
form. I kept the narrow version so that datacenter targets need no re-validation.

Scope

The helper is used directly, or through copy_2d_to_shared{,x2,x3}, by every TMA cast kernel (main @ eb3c860):

• direct: cast/core/grouped_tma.cuh (192, 197), cast/mxfp8/dequantize_mxfp8.cuh (113, 136),
  cast/mxfp8/group_dequantize_mxfp8.cuh (282, 301), cast/mxfp8/group_scaled_swiglu_mxfp8.cuh (309, 313, 333,
  337), cast/mxfp8/specialized/quantize_mxfp8.cuh (955, 1503, 1534), cast/nvfp4/group_quantize_transpose_nvfp4.cuh
  (169, 193), cast/nvfp4/quantize_transpose_nvfp4.cuh (345, 369),
  cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh (505, 580), and the wrappers in util/ptx.cuh
  (1981, 2005, 2009, 2032, 2036, 2040);
• via copy_2d_to_shared{,x2,x3}: cast/fp8/gated_fp8.cuh, cast/fp8/quantize_fp8.cuh, cast/mxfp8/gated_mxfp8.cuh,
  cast/mxfp8/quantize_mxfp8.cuh, cast/nvfp4/group_quantize_transpose_nvfp4.cuh,
  cast/nvfp4/quantize_transpose_nvfp4.cuh, hadamard_transform/{hadamard_transform,group_hadamard_transform,graph_safe_group_hadamard_transform}.cu.

Static scan of the sm_120 slice of libtransformer_engine.so (cuobjdump -elf; kernels whose EIATTR_EXTERNS
names a __cuda_syscall_* symbol):

kernels referencing __cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast	before	after
TE 2.20.2, NVTE_CUDA_ARCHS=120 (sm_120f) → with this change (sm_120a)	8,010	0
main @ eb3c860, NVTE_CUDA_ARCHS=120 (sm_120f), without → with this change	8,195	0

By family on main before the change: group_quantize_mxfp8_kernel 4,896, quantize_mxfp8_kernel 1,224, quantize_gated_mxfp8_kernel 576, quantize_transpose_nvfp4_tuned_1D_kernel 288, quantize_transpose_nvfp4_2D_kernel 256, cast_fp8_gated_kernel 240, quantize_transpose_nvfp4_kernel 176, group_compute_fused_amax_kernel 112, cast_fp8_2D_kernel 102, quantize_mxfp8_kernel_cast_only 96, dequantize_mxfp8_kernel 96, group_scaled_swiglu_mxfp8_kernel 48, compute_fused_amax_kernel 29, group_dequantize_mxfp8_kernel 24, group_quantize_transpose_nvfp4_kernel 8, GroupHadamardAmaxTmaKernel 8, GraphSafeGroupHadamardAmaxTmaKernel 8, HadamardAmaxTmaKernel 8. With the change, the 28 kernels that still reference a syscall are the out-of-scope sites below: 22 × 4D unicast and 6 × 2D multicast.

Evidence (driver 615.71.09, CUDA 13.4, PyTorch 2.16.0.dev20261016+cu134)

TE 2.20.2 without vs with this change, on an RTX 5080 Laptop GPU (sm_120, 60 SMs), on fixed synthetic inputs with
the stochastic-rounding/RHT RNG pinned:

	without	with this change
per-thread stack limit after the first NVFP4 cast	14,624 B	1,024 B
stack limit after every NVFP4 quantizer variant¹ and te.Linear fwd+bwd under FP8 current scaling, NVFP4, NVFP4+RHT+SR+2D	14,640 B	1,024 B
device memory taken by the first NVFP4 cast	1,196 MiB, all outside the caching allocator	0 MiB
outputs: FP4 data, E4M3 scales and amax of the 5 quantizer variants; y, dgrad, wgrad of the 3 te.Linear runs	—	39/39 tensors byte-identical
quantize time² plain / 2D / SR / SR+RHT / RHT (µs)	24.7 / 24.8 / 167.4 / 350.4 / 27.1	23.1 / 22.1 / 163.0 / 350.5 / 26.8

¹ plain 1D, 2D (16×16) weight, stochastic rounding, SR+RHT, RHT input; x = 4096×1280 BF16, grads 4096×5120.
² CUDA events, mean of 50 calls after 5 warm-ups, single run: an incidental observation, not a claim of this PR.

An Nsight Compute launch trace of one NVFP4 te.Linear forward+backward shows the mechanism directly.
launch__stack_size is 1,024 for every kernel before the first quantize_transpose_nvfp4_tuned_1D_kernel<0,0,1,0>
launch. From that launch onward it is 14,624 for every kernel, unrelated PyTorch elementwise and reduction kernels
included.

main @ eb3c860 without vs with this change, built for NVTE_CUDA_ARCHS=120, on an NVIDIA GeForce RTX 5050 (sm_120, 20 SMs), same driver and toolkit:

	without	with this change
per-thread stack limit after the first NVFP4 cast (repro_stack_limit.py)	14,624 B	1,024 B
device memory outside the caching allocator, first NVFP4 cast	400.0 MiB	0.0 MiB
stack limit after 10 ops³ in one process	14,624 B	1,024 B
outputs of those ops	—	39/39 tensors byte-identical
test_sm12x_tma_stack_limit.py	2 failed in 5.86s	2 passed in 5.97s
quantize time⁴ NVFP4 1D / 2D / SR, MXFP8 (µs)	21.16 / 20.99 / 37.38 / 10.2	20.73 / 20.39 / 35.84 / 9.49

³ nvfp4_1d, nvfp4_rowwise, nvfp4_2d, nvfp4_sr, mxfp8_q, mxfp8_dq, mxfp8_2d, fp8_cs, fp8_ds, linear_fp8cs; 1024×2048 BF16 inputs, te.Linear 2048→1024 with batch 512.
⁴ CUDA events; median of 5 × 100 calls after 20 warm-ups; an incidental observation, not a claim of this PR.

How to reproduce

# sm_120/sm_121 GPU; TE built for sm_120 (the default arch list includes it)
python repro_stack_limit.py
#   without: stack limit (bytes/thread): 1024 -> 14624
#            device memory outside the caching allocator: +1196.0 MiB   (RTX 5080 Laptop GPU)
#   with:    stack limit (bytes/thread): 1024 -> 1024
#            device memory outside the caching allocator: +0.0 MiB
#   (NVIDIA GeForce RTX 5050, 20 SMs, main: without → +400.0 MiB; with → +0.0 MiB)

repro_stack_limit.py (attached) quantizes one 1024×2048 BF16 tensor with te.NVFP4Quantizer. Before and after the
cast, it prints cuCtxGetLimit(CU_LIMIT_STACK_SIZE) and total − free − memory_reserved(). tma_stack_repro.cu
(attached) reproduces the effect without TE: nvcc -std=c++17 -O3 -arch=sm_120a tma_stack_repro.cu -lcuda, then run
it with cta and with cluster.

Tests

tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py (new; L0 runs tests/pytorch/nvfp4 as a directory). On sm_12x,
it runs the first NVFP4 cast (1D and 2D block scaling) in a fresh process, because the limit is per context and
never shrinks. It asserts that CU_LIMIT_STACK_SIZE is unchanged, and skips on other GPUs. On main @ eb3c860
(NVIDIA GeForce RTX 5050): 2 failed in 5.86s without the fix (The first 1d NVFP4 cast raised the per-thread stack limit from 1024 to 14624 bytes), 2 passed in 5.97s with it.

Out of scope (same lowering, separate concerns)

• fused_attn/flash_attn.cu has a file-local cp_async_bulk_tensor_4d_global_to_shared that uses the
  .shared::cluster form too. On sm_12x, its 22 multi_tensor_transpose_to_bhsd::transpose_to_bhsd_kernel<…>
  instantiations still reference __cuda_syscall_cp_async_bulk_tensor_4d_tile_unicast. The TMA path of
  multi_tensor_transpose_to_bhsd is taken for sm_arch >= 100, which includes sm_12x, so the same change most
  likely applies there. It deserves its own PR with attention-path validation.
• gemm/cutlass_grouped_gemm.cu: 6 CUTLASS SM90 grouped-GEMM kernels (cluster shape 1×2×1), built for the sm_120
  arch-specific target, reference __cuda_syscall_cp_async_bulk_tensor_2d_tile_multicast. That is genuine multicast,
  which .shared::cta cannot express.

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

• transformer_engine/common/util/ptx.cuh: cp_async_bulk_tensor_2d_global_to_shared issues the .shared::cta
  destination form when __CUDA_ARCH__ >= 1200; every other architecture keeps the .shared::cluster form.
• tests/pytorch/nvfp4/test_sm12x_tma_stack_limit.py: an sm_12x stack-limit regression test (first NVFP4 cast,
  1D and 2D block scaling).

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 (no user-visible API or option changes)
• My changes generate no new warnings (identical nvcc/ptxas warning counts per category with and without the
  change; none mention ptx.cuh)
• 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: the new test passes. The existing
  tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py fails on sm_120 at main @ eb3c860 with and without this change
  (16 failed, 1 skipped in 2.06s / 16 failed, 1 skipped in 2.17s; quantize_transpose_nvfp4.cuh:1854 "Columnwise RHT requires an RHT matrix."). That is a
  pre-existing issue, unrelated to this change.

---

References (outside TE)

• fix(memory): bulk copies into shared::cta, chunk-capture in the T2 plan kekzl/imp#2640 (merged 2026-10-08) reports the same lowering on an RTX 5090: cudaLimitStackSize 1024 → 14448 B
  and ~3.5 GB of local memory, both cleared by moving bulk copies to .shared::cta.
• Typed CTA-shared pointers so TMA G2S can use .shared::cta (SM12x setmaxnreg loss, #1238) cuda-rust#1337 tabulates the TMA global→shared forms on sm_120a vs sm_100a: only .shared::cta avoids the
  syscall path on SM12x. SM12x does support peer destinations through the emulated path, so a by-target switch is
  sound only where every destination is CTA-local, as it is for this helper.

Reproduction files

The three reproduction files referenced above (repro_stack_limit.py, tma_stack_repro.cu, and tma_forms.cu) are available in this public gist.