Repository navigation
NVFP4 quantize on sm_120 reserves ~3.4 GiB of device memory via a TMA driver syscall #3765
Description
Environment: Transformer Engine 2.20.2 (PyPI wheel transformer_engine_cu13), PyTorch 2.12.0+cu130, CUDA 13.2 toolkit,
driver 595.71.05, NVIDIA GeForce RTX 5090 (sm_120, 32 GB).
Symptom
The first NVFP4 quantization in a process permanently takes ~3.3 GiB of device memory outside the PyTorch caching
allocator (about 10% of a 32 GB card):
import torch
from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
x = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16)
free0, total = torch.cuda.mem_get_info()
NVFP4Quantizer(rowwise=True, columnwise=False)(x); torch.cuda.synchronize()
free1, _ = torch.cuda.mem_get_info()
print((free0 - free1) / 2**30) # ~3.3 GiB; cudaLimitStackSize goes from 1024 to 14608 bytesFP8 current scaling and MXFP8 quantization do not show it; cuBLAS GEMMs (BF16/FP8/NVFP4) and the PyTorch SDPA
backends do not either. CUDA_MODULE_LOADING=EAGER does not change it (it is not module code).
Cause
- The launch of
quantize_transpose_nvfp4_tuned_1D_kernelraisescudaLimitStackSizefrom 1024 B to 14608 B, and the
driver reserves that much local memory for every resident thread (170 SMs x 1536 threads), although the kernel
itself reports 8 bytes of local memory. - The kernel's ELF lists an extern
__cuda_syscall_cp_async_bulk_tensor_2d_tile_unicast: the 2D TMA load in
ptx::cp_async_bulk_tensor_2d_global_to_shared(transformer_engine/common/util/ptx.cuh) uses ashared::cluster
destination (cp.async.bulk.tensor.2d.shared::cluster.global.tile...). On sm_120 this compiles to
UTMALDGplusCALL.ABS.NOINCinto a driver syscall, and the syscall's stack requirement causes the reservation. - A standalone 20-line TMA kernel reproduces it (
tma_syscall_repro.cu, below): with ashared::clusterdestination
the launch costs +3.41 GiB; withshared::ctathe SASS is a bareUTMALDG, local memory is 0 and nothing is reserved.
Same result for-arch=sm_120,sm_120aandsm_120f.
Proposed fix
TE already has cp_async_bulk_tensor_2d_global_to_shared_cta (used only by the FP8 block-scaling grouped kernels).
The callers of the 2D shared::cluster helper launch with cluster size 1 and pass CTA-local shared memory and mbarrier
addresses, so switching the helper (or the callers) to the shared::cta destination is semantically identical:
- "cp.async.bulk.tensor.2d.shared::cluster.global.tile"
+ "cp.async.bulk.tensor.2d.shared::cta.global.tile"Patch validated
Rebuilt libtransformer_engine.so from v2.20.2 with only this change (sm_120, CUDA 13.2), RTX 5090:
| stock wheel | patched | |
|---|---|---|
| device memory outside the PyTorch allocator after NVFP4 quantize | 3.92 GiB | 0.61 GiB (CUDA context only) |
| quantized outputs (rowwise; rowwise+columnwise+RHT; stochastic rounding) | bit-identical (data, scales, amax) | |
| Qwen3-0.6B-shaped full training step, 4K tokens, NVFP4 recipe | 155.4 ms | 155.3 ms |
| same, 16K tokens, no activation recomputation | OOM (3/3 runs) | 964.3 ms |
Impact
On 32 GB consumer Blackwell cards this is the difference between fitting and not fitting full-parameter training of
~1.7B models with fp32 master weights and AdamW (usable memory 27.4 GiB vs ~30.7 GiB).
Standalone reproducer (tma_syscall_repro.cu, no TE dependency)
// REPORT 26.2: on the RTX 5090 (sm_120, driver 595.71, CUDA 13.2) a 2D TMA load written with a shared::cluster destination
// compiles to CALL.ABS.NOINC (a driver syscall) next to UTMALDG; its first launch raises cudaLimitStackSize from 1 KiB to
// 14608 B and the driver reserves local memory for every resident thread of the GPU: +3.41 GiB, kept for the process lifetime.
// The shared::cta form compiles to a bare UTMALDG and reserves nothing. TE 2.20.2 uses the shared::cluster form in
// ptx::cp_async_bulk_tensor_2d_global_to_shared (common/util/ptx.cuh), which every NVFP4 quantize kernel calls.
// nvcc -O2 -std=c++17 -arch=sm_120a tma_syscall_repro.cu -o r_cluster -lcuda && ./r_cluster
// nvcc -O2 -std=c++17 -arch=sm_120a -DCTA tma_syscall_repro.cu -o r_cta -lcuda && ./r_cta
// cuobjdump -sass r_cluster | grep -E "CALL|UTMALDG"
#ifdef CTA
#define TMA_DST "shared::cta"
#else
#define TMA_DST "shared::cluster"
#endif
#include <cuda.h>
#include <cuda_runtime.h>
#include <cudaTypedefs.h>
#include <cstdio>
#include <cstdint>
__global__ void k(const __grid_constant__ CUtensorMap map, float* out) {
__shared__ alignas(128) float tile[32 * 32];
__shared__ alignas(8) uint64_t bar;
uint32_t sb = (uint32_t)__cvta_generic_to_shared(&bar), st = (uint32_t)__cvta_generic_to_shared(tile);
if (threadIdx.x == 0) {
asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" :: "r"(sb));
asm volatile("fence.proxy.async.shared::cta;");
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" :: "r"(sb), "r"(32 * 32 * 4));
asm volatile("cp.async.bulk.tensor.2d." TMA_DST ".global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];"
:: "r"(st), "l"(&map), "r"(0), "r"(0), "r"(sb) : "memory");
asm volatile("{ .reg .pred p; W: mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0; @!p bra W; }" :: "r"(sb));
}
__syncthreads();
out[threadIdx.x] = tile[threadIdx.x];
}
int main() {
float *src, *out; cudaMalloc(&src, 64 * 64 * 4); cudaMalloc(&out, 1024 * 4); cudaMemset(src, 0, 64 * 64 * 4);
size_t lim0, lim1, f0, f1, t; cudaDeviceGetLimit(&lim0, cudaLimitStackSize); cudaMemGetInfo(&f0, &t);
PFN_cuTensorMapEncodeTiled_v12000 enc; cudaDriverEntryPointQueryResult q;
cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", (void**)&enc, 12000, cudaEnableDefault, &q);
CUtensorMap m; cuuint64_t dims[2] = {64, 64}, strides[1] = {64 * 4}; cuuint32_t box[2] = {32, 32}, es[2] = {1, 1};
CUresult r = enc(&m, CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 2, src, dims, strides, box, es, CU_TENSOR_MAP_INTERLEAVE_NONE,
CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
cudaFuncAttributes a; cudaFuncGetAttributes(&a, k);
k<<<1, 1024>>>(m, out); cudaError_t e = cudaDeviceSynchronize();
cudaDeviceGetLimit(&lim1, cudaLimitStackSize); cudaMemGetInfo(&f1, &t);
printf("%-15s encode %d launch %s localSizeBytes %zu stack limit %zu -> %zu device memory used +%.2f GiB\n", TMA_DST, r, cudaGetErrorString(e),
a.localSizeBytes, lim0, lim1, (double)(f0 - f1) / (1 << 30));
}Output on the RTX 5090:
shared::cluster encode 0 launch no error localSizeBytes 8 stack limit 1024 -> 14608 device memory used +3.41 GiB
shared::cta encode 0 launch no error localSizeBytes 0 stack limit 1024 -> 1024 device memory used +0.00 GiB
Activity
0z5a commentedon Oct 31, 2026
I checked which architectures this affects. Compiling your standalone reproducer with CUDA 13.3 and grepping
cuobjdump -sass:-archshared::clustershared::ctasm_90a,sm_100a,sm_103a,sm_110aUTMALDG, noCALL.ABS, no__cuda_syscallsymbolsm_120a,sm_121aUTMALDG+CALL.ABS, 3__cuda_syscallsymbolsUTMALDGonlyFor this reproducer, the two forms compile to the same SASS on
sm_100aandsm_103a, instruction for instruction. Onsm_90aandsm_110a, I only checked that neither form has the syscall. On a B300 (sm_103), both binaries leavecudaLimitStackSizeat 1024 and reserve no device memory (+0.00 GiB). I didn't compare SASS of TE's own kernels. By the reproducer, though, the syscall and the reservation are specific to SM12x.The helper has more callers than the NVFP4 kernels. Through
copy_2d_to_shared/copy_2d_to_sharedx2/copy_2d_to_sharedx3, it is also used by:_ctahelper)hadamard_transform.cu,group_hadamard_transform.cu,graph_safe_group_hadamard_transform.cu)None of these files launches with a cluster dimension. The code under
common/that does use clusters doesn't call this helper:cast/mxfp8/specialized/cast_bidim.cuhcustomized_pipeline.cuhSo the
shared::ctadestination is equivalent for every caller, not just the NVFP4 ones. I have no SM12x GPU, so I couldn't check which of those other paths actually run on SM12x and hit the reservation there.The change needs no newer toolkit either:
cp_async_bulk_tensor_2d_global_to_shared_ctaalready emits theshared::ctaform, guarded only by__CUDA_ARCH__ >= 900.KafKafrnZ commentedon Nov 3, 2026
One more SM12x data point, since 0z5a had no SM12x GPU: a different card, OS, driver branch and toolkit, using your standalone repro unchanged.
RTX 5070 Ti Laptop GPU (sm_120, 46 SMs, 12 GB), Windows 11, driver 617.14, nvcc 13.4.59 + MSVC 14.51. Same output for
-arch=sm_120,sm_120aandsm_120f, 3 runs each:localSizeBytesshared::clustershared::ctaCALL.ABS.NOINCsyscall, so a newer driver or toolkit is no workaround.Two things from reading
main(b021980) that bear on where the fix goes. I haven't run TE itself on this card.copy_2d_to_shared*wrappers,ptx::cp_async_bulk_tensor_2d_global_to_sharedis called directly incast/core/grouped_tma.cuh,mxfp8/dequantize_mxfp8.cuh,mxfp8/group_dequantize_mxfp8.cuh,mxfp8/group_scaled_swiglu_mxfp8.cuh,mxfp8/specialized/quantize_mxfp8.cuhand the NVFP4 kernels. Changing the helper itself covers all of them; changing only the wrapper callers would miss these.cp_async_bulk_tensor_1d_global_to_sharedalready emitscp.async.bulk.shared::cta.global(its comment still saysshared::cluster). So FP8 rowwise casts throughquantize_1Ddon't hit the syscall, which fits your FP8 result. The FP8 dispatch gateis_supported_by_CC_100()issm_arch >= 100, so SM12x gets the same TMA kernels there as sm_100.(I used an AI assistant, Claude Code, to help run and write this up; the numbers are from my machine.)