From cd0854c4f0fe9100ad9a8a6a99e62c4c387ecb6f Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 24 Sep 2026 16:58:48 -0700 Subject: [PATCH 1/5] [DO NOT MERGE YET] Diagnose Qwen3.5 CUDA export illegal memory access Run the Qwen3.5 MoE export with CUDA_LAUNCH_BLOCKING=1 so the traceback points at the kernel that faults instead of a later CUDA API call. Follow-up commits on this branch will carry the actual fix. --- .ci/scripts/export_model_artifact.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 223dd82951c..e47d8f5d583 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -492,6 +492,8 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then EXPORT_LOG=$(mktemp) TMPDIR="$INDUCTOR_TMPDIR" \ TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ + CUDA_LAUNCH_BLOCKING=1 \ + PYTHONFAULTHANDLER=1 \ python -m executorch.examples.models.qwen3_5_moe.export \ --prequantized "$LOCAL_MODEL_DIR" \ --output-dir "${OUTPUT_DIR}" \ From 2c1d3361649ce78ea9ad63183b59306c15a55bff Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 24 Sep 2026 17:49:13 -0700 Subject: [PATCH 2/5] [DO NOT MERGE YET] Surface the hidden Qwen3.5 export CUDA error Inductor re-raises codegen errors from None, and preserve_rng_state's set_rng_state failure then replaces the original exception, so the log only shows the secondary error. Wrap the export to print the full exception chain and log every CachingAutotuner.run (flagging undersized tensor storage) so the faulting kernel is identifiable. --- .ci/scripts/diag_run_module.py | 66 ++++++++++++++++++++++++++++ .ci/scripts/export_model_artifact.sh | 2 +- 2 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 .ci/scripts/diag_run_module.py diff --git a/.ci/scripts/diag_run_module.py b/.ci/scripts/diag_run_module.py new file mode 100644 index 00000000000..dac955a11e3 --- /dev/null +++ b/.ci/scripts/diag_run_module.py @@ -0,0 +1,66 @@ +# Temporary diagnostic wrapper for the Qwen3.5 CUDA export illegal memory access. +# Usage: python .ci/scripts/diag_run_module.py [args...] +import runpy +import sys +import traceback + +import torch +from torch._inductor.runtime.triton_heuristics import CachingAutotuner + + +def _required_nbytes(t): + if t.numel() == 0: + return 0 + last = t.storage_offset() + for size, stride in zip(t.size(), t.stride()): + last += (size - 1) * stride + return (last + 1) * t.element_size() + + +_orig_run = CachingAutotuner.run + + +def _logged_run(self, *args, **kwargs): + name = self.inductor_meta.get("kernel_name") or getattr(self.fn, "__name__", "?") + bad = [] + for i, a in enumerate((*args, *kwargs.values())): + if isinstance(a, torch.Tensor) and a.is_cuda: + have = a.untyped_storage().nbytes() + need = _required_nbytes(a) + if have < need: + bad.append(f"arg{i} shape={tuple(a.shape)} storage={have} need={need}") + msg = f"[diag] autotuner.run {name} launchers={len(self.launchers)}" + if bad: + msg += " UNDERSIZED: " + "; ".join(bad) + print(msg, file=sys.stderr, flush=True) + return _orig_run(self, *args, **kwargs) + + +CachingAutotuner.run = _logged_run + + +def _print_chain(exc): + seen = set() + depth = 0 + while exc is not None and id(exc) not in seen: + seen.add(id(exc)) + print(f"[diag] ===== exception chain depth {depth}: {type(exc).__name__}", file=sys.stderr) + traceback.print_exception(type(exc), exc, exc.__traceback__, chain=False) + inner = getattr(exc, "inner_exception", None) + if inner is not None and id(inner) not in seen: + exc = inner + else: + exc = exc.__cause__ or exc.__context__ + depth += 1 + + +module = sys.argv[1] +sys.argv = [module] + sys.argv[2:] +try: + runpy.run_module(module, run_name="__main__", alter_sys=True) +except BaseException as e: + if isinstance(e, SystemExit) and not e.code: + raise + _print_chain(e) + sys.stderr.flush() + sys.exit(1) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index e47d8f5d583..a57c0c08c2d 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -494,7 +494,7 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ CUDA_LAUNCH_BLOCKING=1 \ PYTHONFAULTHANDLER=1 \ - python -m executorch.examples.models.qwen3_5_moe.export \ + python .ci/scripts/diag_run_module.py executorch.examples.models.qwen3_5_moe.export \ --prequantized "$LOCAL_MODEL_DIR" \ --output-dir "${OUTPUT_DIR}" \ --dense-prefill dequant \ From 3acb14c72310971125d00800069f88a2a78ea374 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 24 Sep 2026 18:23:00 -0700 Subject: [PATCH 3/5] Rehydrate emptied KV tensors for AOTI's JIT-variant autotune pass With autotune_at_compile_time off, PyTorch 2.14 no longer autotunes through a Python first pass. GraphLowering._run_jit_variant_for_autotune compiles a C++ JIT wrapper and runs it on self.constants and the real inputs directly. The low-memory CUDA export frees KV-cache storage with resize_(0), and #21803 only rehydrated it inside CachingAutotuner.run, so the null KV storage now reaches the kernels and the Qwen3.5 MoE export fails with an illegal memory access. Rehydrate emptied constants and inputs for the duration of the JIT run, the same way CachingAutotuner.run already does. --- backends/cuda/cuda_backend.py | 35 +++++++++++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 35 +++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 4db99b6b040..63c31909864 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -215,6 +215,7 @@ def _force_external_weights_for_fqn_binding(consts_size): return orig_determine_aoti_mmap_flags(consts_size) orig_autotuner_run = _Autotuner.run + orig_run_jit_variant = getattr(_GL, "_run_jit_variant_for_autotune", None) def _is_same_skip_emptied(data, value): # KV buffers freed via resize_(0) all have data_ptr 0, so the stock @@ -259,6 +260,36 @@ def _autotuner_run_with_rehydrated_emptied_args(self, *args, **kwargs): with _rehydrate_emptied_tensors(tensors): return orig_autotuner_run(self, *args, **kwargs) + def _run_jit_variant_with_rehydrated_emptied( + self, wrapper_code, kernel_code, extract_real_inputs, *args, **kwargs + ): + # With autotune_at_compile_time off, PyTorch 2.14 autotunes by running a + # C++ JIT wrapper on self.constants and the real inputs directly, so + # emptied KV tensors never pass through CachingAutotuner.run and their + # null storage reaches the kernels. Keep them valid for the whole run. + with contextlib.ExitStack() as stack: + stack.enter_context( + _rehydrate_emptied_tensors(list(self.constants.values())) + ) + + def _extract_rehydrated_inputs(): + inputs = extract_real_inputs() + stack.enter_context( + _rehydrate_emptied_tensors( + [x for x in inputs if isinstance(x, torch.Tensor)] + ) + ) + return inputs + + return orig_run_jit_variant( + self, + wrapper_code, + kernel_code, + _extract_rehydrated_inputs, + *args, + **kwargs, + ) + def _get_const_synthesize_zeros(self, name): # AOTI serializes each constant via get_original_value_of_constant -> # _to_bytes. For KV buffers we freed with resize_(0) this would otherwise @@ -292,6 +323,8 @@ def _codegen_device_target_aware(self, device): ) _codecache.determine_aoti_mmap_flags = _force_external_weights_for_fqn_binding _Autotuner.run = _autotuner_run_with_rehydrated_emptied_args + if orig_run_jit_variant is not None: + _GL._run_jit_variant_for_autotune = _run_jit_variant_with_rehydrated_emptied prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True try: @@ -305,6 +338,8 @@ def _codegen_device_target_aware(self, device): _codecache.TensorProperties = orig_tensor_properties _codecache.determine_aoti_mmap_flags = orig_determine_aoti_mmap_flags _Autotuner.run = orig_autotuner_run + if orig_run_jit_variant is not None: + _GL._run_jit_variant_for_autotune = orig_run_jit_variant def _is_kv_buffer(name, v) -> bool: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 7b45e677da2..2e1e42d5df4 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -65,6 +65,41 @@ def fake_run(_autotuner, base_arg, view_arg, *, stream): self.assertEqual(result, "run-result") self.assertEqual(base.untyped_storage().nbytes(), 0) + def test_low_memory_jit_autotune_rehydrates_constants_and_inputs(self): + from executorch.backends.cuda.cuda_backend import _compile_time_cpu_clones + from torch._inductor.graph import GraphLowering + + if not hasattr(GraphLowering, "_run_jit_variant_for_autotune"): + self.skipTest("PyTorch has no JIT-variant autotune pass") + + constant = torch.empty(8) + constant.untyped_storage().resize_(0) + graph_input = torch.empty(4) + graph_input.untyped_storage().resize_(0) + original_run = GraphLowering._run_jit_variant_for_autotune + test_case = self + + def fake_run(_graph, wrapper_code, kernel_code, extract_real_inputs, names): + test_case.assertGreater(constant.untyped_storage().nbytes(), 0) + inputs = extract_real_inputs() + test_case.assertGreater(inputs[0].untyped_storage().nbytes(), 0) + test_case.assertEqual(torch.count_nonzero(inputs[0]), 0) + test_case.assertEqual(names, ["kernel"]) + + GraphLowering._run_jit_variant_for_autotune = fake_run + try: + graph = object.__new__(GraphLowering) + graph.constants = {"kv_cache": constant} + with _compile_time_cpu_clones(torch.device("cuda")): + graph._run_jit_variant_for_autotune( + "wrapper", "kernels", lambda: [graph_input], ["kernel"] + ) + finally: + GraphLowering._run_jit_variant_for_autotune = original_run + + self.assertEqual(constant.untyped_storage().nbytes(), 0) + self.assertEqual(graph_input.untyped_storage().nbytes(), 0) + def test_emulate_precision_casts_compile_spec(self): options = CudaBackend.get_aoti_compile_options( [CompileSpec(key="emulate_precision_casts", value=b"OFF")] From eb6e61332b5211658394935e234765aa902d0c95 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 24 Sep 2026 19:17:59 -0700 Subject: [PATCH 4/5] Use compile-time autotune for the Qwen3.5 CUDA export again autotune_at_compile_time=OFF was added with the PyTorch 2.13 pin (#20967) to avoid a 2.13-only compile-time autotune issue. On PyTorch 2.14 it selects the new AOTI lazy-compile path, whose AOT code loads every Triton kernel from a cubin file at runtime and ignores embed_kernel_binary. ExecuTorch only packs the wrapper .so, so the runner fails on the first execute with "CUDA driver error: file not found". Drop the override so AOTI uses its default (compile-time autotune with embedded kernel binaries). --- examples/models/qwen3_5_moe/export.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/examples/models/qwen3_5_moe/export.py b/examples/models/qwen3_5_moe/export.py index 92408a7ea71..4c060ce0137 100644 --- a/examples/models/qwen3_5_moe/export.py +++ b/examples/models/qwen3_5_moe/export.py @@ -1138,9 +1138,8 @@ def _export_cuda(model, config, args): "enable_dynamic_shape": True, "get_mutable_buffer_metadata": _mutable_buffer_metadata_json(model), } - # Avoid PyTorch 2.13's fused cast/reduction compile-time autotune path. Keep - # max autotuning enabled so prefill GEMMs use the CUDA backend's Triton-only - # GEMM lowering instead of emitting libtorch fallback kernels. + # Keep max autotuning enabled so prefill GEMMs use the CUDA backend's + # Triton-only GEMM lowering instead of emitting libtorch fallback kernels. et_prog = to_edge_transform_and_lower( {"decode": decode_ep, "prefill": prefill_ep}, partitioner={ @@ -1151,7 +1150,6 @@ def _export_cuda(model, config, args): CompileSpec("low_memory_mode", b"ON"), CompileSpec("emulate_precision_casts", b"OFF"), CompileSpec("max_autotune", b"ON"), - CompileSpec("autotune_at_compile_time", b"OFF"), ] ) ], @@ -1162,7 +1160,6 @@ def _export_cuda(model, config, args): CompileSpec("low_memory_mode", b"ON"), CompileSpec("emulate_precision_casts", b"OFF"), CompileSpec("max_autotune", b"ON"), - CompileSpec("autotune_at_compile_time", b"OFF"), ] ) ], From 581658eb6e4be67628e28ec16ec295a25cada430 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Thu, 24 Sep 2026 22:27:07 -0700 Subject: [PATCH 5/5] Remove Qwen3.5 export diagnostics and the JIT-variant KV rehydration Compile-time autotune is back on for the Qwen3.5 CUDA export, so AOTI no longer runs the JIT-variant autotune pass and the KV rehydration wrapper is unused. Drop it along with the temporary CUDA_LAUNCH_BLOCKING and exception-chain diagnostics. The only remaining change is the export config. --- .ci/scripts/diag_run_module.py | 66 ------------------------- .ci/scripts/export_model_artifact.sh | 4 +- backends/cuda/cuda_backend.py | 35 ------------- backends/cuda/tests/test_cuda_export.py | 35 ------------- 4 files changed, 1 insertion(+), 139 deletions(-) delete mode 100644 .ci/scripts/diag_run_module.py diff --git a/.ci/scripts/diag_run_module.py b/.ci/scripts/diag_run_module.py deleted file mode 100644 index dac955a11e3..00000000000 --- a/.ci/scripts/diag_run_module.py +++ /dev/null @@ -1,66 +0,0 @@ -# Temporary diagnostic wrapper for the Qwen3.5 CUDA export illegal memory access. -# Usage: python .ci/scripts/diag_run_module.py [args...] -import runpy -import sys -import traceback - -import torch -from torch._inductor.runtime.triton_heuristics import CachingAutotuner - - -def _required_nbytes(t): - if t.numel() == 0: - return 0 - last = t.storage_offset() - for size, stride in zip(t.size(), t.stride()): - last += (size - 1) * stride - return (last + 1) * t.element_size() - - -_orig_run = CachingAutotuner.run - - -def _logged_run(self, *args, **kwargs): - name = self.inductor_meta.get("kernel_name") or getattr(self.fn, "__name__", "?") - bad = [] - for i, a in enumerate((*args, *kwargs.values())): - if isinstance(a, torch.Tensor) and a.is_cuda: - have = a.untyped_storage().nbytes() - need = _required_nbytes(a) - if have < need: - bad.append(f"arg{i} shape={tuple(a.shape)} storage={have} need={need}") - msg = f"[diag] autotuner.run {name} launchers={len(self.launchers)}" - if bad: - msg += " UNDERSIZED: " + "; ".join(bad) - print(msg, file=sys.stderr, flush=True) - return _orig_run(self, *args, **kwargs) - - -CachingAutotuner.run = _logged_run - - -def _print_chain(exc): - seen = set() - depth = 0 - while exc is not None and id(exc) not in seen: - seen.add(id(exc)) - print(f"[diag] ===== exception chain depth {depth}: {type(exc).__name__}", file=sys.stderr) - traceback.print_exception(type(exc), exc, exc.__traceback__, chain=False) - inner = getattr(exc, "inner_exception", None) - if inner is not None and id(inner) not in seen: - exc = inner - else: - exc = exc.__cause__ or exc.__context__ - depth += 1 - - -module = sys.argv[1] -sys.argv = [module] + sys.argv[2:] -try: - runpy.run_module(module, run_name="__main__", alter_sys=True) -except BaseException as e: - if isinstance(e, SystemExit) and not e.code: - raise - _print_chain(e) - sys.stderr.flush() - sys.exit(1) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index a57c0c08c2d..223dd82951c 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -492,9 +492,7 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then EXPORT_LOG=$(mktemp) TMPDIR="$INDUCTOR_TMPDIR" \ TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ - CUDA_LAUNCH_BLOCKING=1 \ - PYTHONFAULTHANDLER=1 \ - python .ci/scripts/diag_run_module.py executorch.examples.models.qwen3_5_moe.export \ + python -m executorch.examples.models.qwen3_5_moe.export \ --prequantized "$LOCAL_MODEL_DIR" \ --output-dir "${OUTPUT_DIR}" \ --dense-prefill dequant \ diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 63c31909864..4db99b6b040 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -215,7 +215,6 @@ def _force_external_weights_for_fqn_binding(consts_size): return orig_determine_aoti_mmap_flags(consts_size) orig_autotuner_run = _Autotuner.run - orig_run_jit_variant = getattr(_GL, "_run_jit_variant_for_autotune", None) def _is_same_skip_emptied(data, value): # KV buffers freed via resize_(0) all have data_ptr 0, so the stock @@ -260,36 +259,6 @@ def _autotuner_run_with_rehydrated_emptied_args(self, *args, **kwargs): with _rehydrate_emptied_tensors(tensors): return orig_autotuner_run(self, *args, **kwargs) - def _run_jit_variant_with_rehydrated_emptied( - self, wrapper_code, kernel_code, extract_real_inputs, *args, **kwargs - ): - # With autotune_at_compile_time off, PyTorch 2.14 autotunes by running a - # C++ JIT wrapper on self.constants and the real inputs directly, so - # emptied KV tensors never pass through CachingAutotuner.run and their - # null storage reaches the kernels. Keep them valid for the whole run. - with contextlib.ExitStack() as stack: - stack.enter_context( - _rehydrate_emptied_tensors(list(self.constants.values())) - ) - - def _extract_rehydrated_inputs(): - inputs = extract_real_inputs() - stack.enter_context( - _rehydrate_emptied_tensors( - [x for x in inputs if isinstance(x, torch.Tensor)] - ) - ) - return inputs - - return orig_run_jit_variant( - self, - wrapper_code, - kernel_code, - _extract_rehydrated_inputs, - *args, - **kwargs, - ) - def _get_const_synthesize_zeros(self, name): # AOTI serializes each constant via get_original_value_of_constant -> # _to_bytes. For KV buffers we freed with resize_(0) this would otherwise @@ -323,8 +292,6 @@ def _codegen_device_target_aware(self, device): ) _codecache.determine_aoti_mmap_flags = _force_external_weights_for_fqn_binding _Autotuner.run = _autotuner_run_with_rehydrated_emptied_args - if orig_run_jit_variant is not None: - _GL._run_jit_variant_for_autotune = _run_jit_variant_with_rehydrated_emptied prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True try: @@ -338,8 +305,6 @@ def _codegen_device_target_aware(self, device): _codecache.TensorProperties = orig_tensor_properties _codecache.determine_aoti_mmap_flags = orig_determine_aoti_mmap_flags _Autotuner.run = orig_autotuner_run - if orig_run_jit_variant is not None: - _GL._run_jit_variant_for_autotune = orig_run_jit_variant def _is_kv_buffer(name, v) -> bool: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 2e1e42d5df4..7b45e677da2 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -65,41 +65,6 @@ def fake_run(_autotuner, base_arg, view_arg, *, stream): self.assertEqual(result, "run-result") self.assertEqual(base.untyped_storage().nbytes(), 0) - def test_low_memory_jit_autotune_rehydrates_constants_and_inputs(self): - from executorch.backends.cuda.cuda_backend import _compile_time_cpu_clones - from torch._inductor.graph import GraphLowering - - if not hasattr(GraphLowering, "_run_jit_variant_for_autotune"): - self.skipTest("PyTorch has no JIT-variant autotune pass") - - constant = torch.empty(8) - constant.untyped_storage().resize_(0) - graph_input = torch.empty(4) - graph_input.untyped_storage().resize_(0) - original_run = GraphLowering._run_jit_variant_for_autotune - test_case = self - - def fake_run(_graph, wrapper_code, kernel_code, extract_real_inputs, names): - test_case.assertGreater(constant.untyped_storage().nbytes(), 0) - inputs = extract_real_inputs() - test_case.assertGreater(inputs[0].untyped_storage().nbytes(), 0) - test_case.assertEqual(torch.count_nonzero(inputs[0]), 0) - test_case.assertEqual(names, ["kernel"]) - - GraphLowering._run_jit_variant_for_autotune = fake_run - try: - graph = object.__new__(GraphLowering) - graph.constants = {"kv_cache": constant} - with _compile_time_cpu_clones(torch.device("cuda")): - graph._run_jit_variant_for_autotune( - "wrapper", "kernels", lambda: [graph_input], ["kernel"] - ) - finally: - GraphLowering._run_jit_variant_for_autotune = original_run - - self.assertEqual(constant.untyped_storage().nbytes(), 0) - self.assertEqual(graph_input.untyped_storage().nbytes(), 0) - def test_emulate_precision_casts_compile_spec(self): options = CudaBackend.get_aoti_compile_options( [CompileSpec(key="emulate_precision_casts", value=b"OFF")]