Skip to content

Use compile-time autotune for the Qwen3.5 CUDA export on PyTorch 2.14 #23140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gasoonjia merged 6 commits into from
Sep 25, 2026
Merged

Use compile-time autotune for the Qwen3.5 CUDA export on PyTorch 2.14 #23140

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cd0854c
[DO NOT MERGE YET] Diagnose Qwen3.5 CUDA export illegal memory access
Gasoonjia Sep 24, 2026
2c1d336
[DO NOT MERGE YET] Surface the hidden Qwen3.5 export CUDA error
Gasoonjia Sep 25, 2026
3acb14c
Rehydrate emptied KV tensors for AOTI's JIT-variant autotune pass
Gasoonjia Sep 25, 2026
eb6e613
Use compile-time autotune for the Qwen3.5 CUDA export again
Gasoonjia Sep 25, 2026
581658e
Remove Qwen3.5 export diagnostics and the JIT-variant KV rehydration
Gasoonjia Sep 25, 2026
5b24ba1
Merge branch 'main' into gasoonjia/fix-qwen35-cuda-export
Gasoonjia Sep 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions examples/models/qwen3_5_moe/export.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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={
Expand All @@ -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"),
]
)
],
Expand All @@ -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"),
]
)
],
Expand Down
Loading