Skip to content

Commit cb94e69

Browse files
authored
Use compile-time autotune for the Qwen3.5 CUDA export on PyTorch 2.14 (#23140)
### Summary `test-model-cuda-e2e-Qwen3.5-35B-A3B-HQQ-INT4-quantized-int4-tile-packed` (and the Qwen3.5 `benchmark-cuda` job) have failed on every main run since the PyTorch 2.14 pin bump (#22501). The Qwen3.5 export set `autotune_at_compile_time=OFF`. That was added with the 2.13 pin (#20967) to avoid a 2.13-only compile-time autotune issue. On PyTorch 2.14, OFF selects AOTInductor's new lazy-compile path, which breaks this model in two ways: 1. **Export crash.** Autotuning now runs a C++ JIT variant of the graph (`GraphLowering._run_jit_variant_for_autotune`) directly on `self.constants` and the real inputs. The low-memory CUDA export frees KV-cache storage (`resize_(0)`) before compile. The #21803 rehydration only hooks `CachingAutotuner.run`, so the null KV storage reaches the kernels: ``` _run_jit_variant_for_autotune -> compiled_fn(input_handles) RuntimeError: CUDA driver error: an illegal memory access was encountered ``` The logged error was only the secondary one, from `preserve_rng_state` -> `set_rng_state`. 2. **Unrunnable artifact.** Even with the KV storage rehydrated, the lazy 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`. ### Fix Drop the `autotune_at_compile_time=OFF` override. On 2.14, AOTI defaults it to on for accelerators, so kernels are autotuned at compile time on generated example inputs and embedded in the `.so`, the same as the other CUDA models. The obsolete 2.13 comment is removed too. ### Test plan CI on this PR (earlier commits carried temporary diagnostics, which have since been removed): - `test-model-cuda-e2e-Qwen3.5-35B-A3B-HQQ-INT4-quantized-int4-tile-packed` passes: - export GPU peak 18682 MB (limit 21504 MB) - runner produces correct output, decode 113.8 tok/s - `Qwen3.5-MoE serving smoke passed`
1 parent f626041 commit cb94e69

1 file changed

Lines changed: 2 additions & 5 deletions

File tree

‎examples/models/qwen3_5_moe/export.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1138,9 +1138,8 @@ def _export_cuda(model, config, args):
11381138
"enable_dynamic_shape": True,
11391139
"get_mutable_buffer_metadata": _mutable_buffer_metadata_json(model),
11401140
}
1141-
# Avoid PyTorch 2.13's fused cast/reduction compile-time autotune path. Keep
1142-
# max autotuning enabled so prefill GEMMs use the CUDA backend's Triton-only
1143-
# GEMM lowering instead of emitting libtorch fallback kernels.
1141+
# Keep max autotuning enabled so prefill GEMMs use the CUDA backend's
1142+
# Triton-only GEMM lowering instead of emitting libtorch fallback kernels.
11441143
et_prog = to_edge_transform_and_lower(
11451144
{"decode": decode_ep, "prefill": prefill_ep},
11461145
partitioner={
@@ -1151,7 +1150,6 @@ def _export_cuda(model, config, args):
11511150
CompileSpec("low_memory_mode", b"ON"),
11521151
CompileSpec("emulate_precision_casts", b"OFF"),
11531152
CompileSpec("max_autotune", b"ON"),
1154-
CompileSpec("autotune_at_compile_time", b"OFF"),
11551153
]
11561154
)
11571155
],
@@ -1162,7 +1160,6 @@ def _export_cuda(model, config, args):
11621160
CompileSpec("low_memory_mode", b"ON"),
11631161
CompileSpec("emulate_precision_casts", b"OFF"),
11641162
CompileSpec("max_autotune", b"ON"),
1165-
CompileSpec("autotune_at_compile_time", b"OFF"),
11661163
]
11671164
)
11681165
],

0 commit comments

Comments
 (0)