Issue Station
Sign In
Tools hereToolbox

pytorch - 💡(How to fix) Fix CuTe/Quack `_fused_rms_norm` override hangs while compiling very large normalized shape

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…

Error Message

So _fused_rms_norm_cond should probably return False for unsupported huge N so the call falls back to the ATen CUDA implementation, or the cutlass DSL call should complete (potentially with an error) in a reasonable amount of time.

Fix Action

Fix / Workaround

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 48 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 48 On-line CPU(s) list: 0-47 Vendor ID: AuthenticAMD Model name: AMD EPYC 7413 24-Core Processor CPU family: 25 Model: 1 Thread(s) per core: 2 Core(s) per socket: 24 Socket(s): 1 Stepping: 1 Frequency boost: enabled CPU(s) scaling MHz: 79% CPU max MHz: 3630.8101 CPU min MHz: 1500.0000 BogoMIPS: 5299.85 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm cons tant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_le gacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs i bpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local use r_shstk clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_ spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca debug_swap Virtualization: AMD-V L1d cache: 768 KiB (24 instances) L1i cache: 768 KiB (24 instances) L2 cache: 12 MiB (24 instances) L3 cache: 128 MiB (4 instances) NUMA node(s): 1 NUMA node0 CPU(s): 0-47 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Not affected Vulnerability Mds: Not affected Vulnerability Meltdown: Not affected Vulnerability Mmio stale data: Not affected Vulnerability Reg file data sampling: Not affected Vulnerability Retbleed: Not affected Vulnerability Spec rstack overflow: Mitigation; Safe RET Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Not affected

Code Example

import torch

DISABLE_NATIVE_OVERRIDE = False  # Set True: the same input completes.

if DISABLE_NATIVE_OVERRIDE:
    from torch._native.registry import deregister_op_overrides

    deregister_op_overrides(disable_op_symbols="_fused_rms_norm")

x = torch.empty((1, 268_435_456), device="cuda", dtype=torch.bfloat16)
print(torch.version.git_version, torch.cuda.get_device_name(), flush=True)

y, rstd = torch.ops.aten._fused_rms_norm(x, [x.shape[-1]], None, 1e-5)
torch.cuda.synchronize()
print("ok", y.shape, rstd.shape)

---

Thread 1 (Thread 0x756d68200300 (LWP 2153) "python"):
#0  0x00007566c9508b73 in ?? ()
   from /usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/_mlir/_mlir_libs/_cutlass_ir.cpython-312-x86_64-linux-gnu.so
#1  0x0000000000000000 in ?? ()

---

PyTorch version: 2.13.0a0+git7168b60
Is debug build: False
CUDA used to build PyTorch: 13.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.11.0-17-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.3.33
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA H100 80GB HBM3
Nvidia driver version: 595.58.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_tensor_ir.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_ext.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.23.0
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               48
On-line CPU(s) list:                  0-47
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7413 24-Core Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   24
Socket(s):                            1
Stepping:                             1
Frequency boost:                      enabled
CPU(s) scaling MHz:                   79%
CPU max MHz:                          3630.8101
CPU min MHz:                          1500.0000
BogoMIPS:                             5299.85
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm cons
tant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_le
gacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs i
bpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local use
r_shstk clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_
spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca debug_swap
Virtualization:                       AMD-V
L1d cache:                            768 KiB (24 instances)
L1i cache:                            768 KiB (24 instances)
L2 cache:                             12 MiB (24 instances)
L3 cache:                             128 MiB (4 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-47
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; Safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] intel-openmp==2021.4.0
[pip3] mkl==2021.1.1
[pip3] mkl-devel==2021.1.1
[pip3] mkl-include==2021.1.1
[pip3] mypy==1.16.0
[pip3] mypy_extensions==1.1.0
[pip3] numpy==1.26.4
[pip3] nvidia-cuda-cupti==13.3.35
[pip3] nvidia-cudnn-frontend==1.24.1
[pip3] onnx==1.21.0
[pip3] onnx-ir==0.1.16
[pip3] onnxscript==0.6.2
[pip3] optree==0.13.0
[pip3] tbb==2021.13.1
[pip3] torch==2.13.0a0+git7168b60
[pip3] torch_c_dlpack_ext==0.1.5
[pip3] triton==3.7.0+git88b227e2
[conda] Could not collect
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

I noticed during a sanity check that torch.ops.aten._fused_rms_norm can hang indefinitely on H100 when the Python-native CuTe/Quack override is enabled and the normalized dimension is very large.

The reproducing case is:

import torch

DISABLE_NATIVE_OVERRIDE = False  # Set True: the same input completes.

if DISABLE_NATIVE_OVERRIDE:
    from torch._native.registry import deregister_op_overrides

    deregister_op_overrides(disable_op_symbols="_fused_rms_norm")

x = torch.empty((1, 268_435_456), device="cuda", dtype=torch.bfloat16)
print(torch.version.git_version, torch.cuda.get_device_name(), flush=True)

y, rstd = torch.ops.aten._fused_rms_norm(x, [x.shape[-1]], None, 1e-5)
torch.cuda.synchronize()
print("ok", y.shape, rstd.shape)

The call does not return for >2.5h on a server system with an AMD EPYC 7413 CPU. Attaching gdb while it is stuck showed the main Python thread inside the cutlass DSL MLIR extension:

Thread 1 (Thread 0x756d68200300 (LWP 2153) "python"):
#0  0x00007566c9508b73 in ?? ()
   from /usr/local/lib/python3.12/dist-packages/nvidia_cutlass_dsl/python_packages/cutlass/_mlir/_mlir_libs/_cutlass_ir.cpython-312-x86_64-linux-gnu.so
#1  0x0000000000000000 in ?? ()

Disabling only the Python-native _fused_rms_norm override by setting DISABLE_NATIVE_OVERRIDE = True in the repro makes the op complete as expected.

So _fused_rms_norm_cond should probably return False for unsupported huge N so the call falls back to the ATen CUDA implementation, or the cutlass DSL call should complete (potentially with an error) in a reasonable amount of time. I'll reach out separately to the cutlass DSL team to see whether this might be a bug on their side. I'll update this issue once I have a link to that.

CC: @ngimel , @drisspg I believe you were working on this integration, so adding you FYI / if this is a known issue

Versions

PyTorch version: 2.13.0a0+git7168b60
Is debug build: False
CUDA used to build PyTorch: 13.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 23 2026, 19:04:32) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.11.0-17-generic-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 13.3.33
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA H100 80GB HBM3
Nvidia driver version: 595.58.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_tensor_ir.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_ext.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.23.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.23.0
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               48
On-line CPU(s) list:                  0-47
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7413 24-Core Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   24
Socket(s):                            1
Stepping:                             1
Frequency boost:                      enabled
CPU(s) scaling MHz:                   79%
CPU max MHz:                          3630.8101
CPU min MHz:                          1500.0000
BogoMIPS:                             5299.85
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm cons
tant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_le
gacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs i
bpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local use
r_shstk clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_
spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca debug_swap
Virtualization:                       AMD-V
L1d cache:                            768 KiB (24 instances)
L1i cache:                            768 KiB (24 instances)
L2 cache:                             12 MiB (24 instances)
L3 cache:                             128 MiB (4 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-47
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; Safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] intel-openmp==2021.4.0
[pip3] mkl==2021.1.1
[pip3] mkl-devel==2021.1.1
[pip3] mkl-include==2021.1.1
[pip3] mypy==1.16.0
[pip3] mypy_extensions==1.1.0
[pip3] numpy==1.26.4
[pip3] nvidia-cuda-cupti==13.3.35
[pip3] nvidia-cudnn-frontend==1.24.1
[pip3] onnx==1.21.0
[pip3] onnx-ir==0.1.16
[pip3] onnxscript==0.6.2
[pip3] optree==0.13.0
[pip3] tbb==2021.13.1
[pip3] torch==2.13.0a0+git7168b60
[pip3] torch_c_dlpack_ext==0.1.5
[pip3] triton==3.7.0+git88b227e2
[conda] Could not collect

cc @ezyang @gchanan @kadeng @msaroufim @slayton58

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING