pytorch - 💡(How to fix) Fix [vllm] [torch nightly] AOTAutogradCache.save: _pickle.PicklingError on triton launcher (DFlashDraftModel speculative-decode path) [1 participants]

Summary

On a recent torch nightly snapshot, vLLM's test_can_initialize_large_subset[DFlashDraftModel] fails during engine initialization because AOTAutogradCache.save cannot pickle a triton-generated launcher function:

_pickle.PicklingError: Can't pickle <function launcher at 0x7fc56cd2aca0>: attribute lookup launcher on __main__ failed

The failure is on the speculative-decode DFlashDraftModel path (Qwen3-DFlash). The remaining 182/183 model-initialization tests in the same job pass.

Environment

• torch: nightly (vLLM build v0.20.1rc1.dev48+g8a8c9b564.d20260429)
• CUDA: 13.0 / Driver: 570.133.20
• Python: 3.12.13
• Platform: Ubuntu 22.04 x86_64
• vLLM main, commit 8a8c9b564e

Reproduction

Failing test:

tests/models/test_initialization.py::test_can_initialize_large_subset[DFlashDraftModel]

Pytest summary: 1 failed, 182 passed, 5 deselected, 200 warnings in 3522.34s

The test creates a vLLM LLM instance with DFlashDraftModel as the spec-decode draft. During the dummy-run profiling, vLLM compiles qwen3_dflash via torch.compile / AOTAutograd, and the cache-save step on the resulting AOTAutograd entry trips the pickling error.

Traceback (key frames, abridged)

File "vllm/v1/engine/core.py", line 250, in _initialize_kv_caches
  available_gpu_memory = self.model_executor.determine_available_memory()
File "vllm/v1/worker/gpu_worker.py", line 370, in determine_available_memory
  self.model_runner.profile_run()
File "vllm/v1/worker/gpu_model_runner.py", line 5587, in _dummy_run
  self.drafter.dummy_run(...)
File "vllm/v1/spec_decode/dflash.py", line 246, in dummy_run
  self.model(...)
File "vllm/model_executor/models/qwen3_dflash.py", line 549, in forward
  return self.model(input_ids, positions, inputs_embeds)
File "vllm/compilation/decorators.py", line 623, in __call__
  self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
File "vllm/compilation/wrapper.py", line 183, in aot_compile
  return self._compiled_callable.aot_compile((args, kwargs))
File "torch/_dynamo/eval_frame.py", line 868, in aot_compile
  return aot_compile_fullgraph(...)
File "torch/_dynamo/aot_compile.py", line 387, in aot_compile_fullgraph
  compiled_fn = backend(...)
...
File "torch/_inductor/standalone_compile.py", line 462, in standalone_compile
  compiled_fn = compile_fx(...)
File "torch/_inductor/compile_fx.py", line 3022, in _compile_fx_main
  return dynamo_common.aot_autograd(...)
File "torch/_functorch/aot_autograd.py", line 1234, in aot_module_simplified
  compiled_fn, _ = aot_stage2_compile(...)
File "torch/_functorch/_aot_autograd/graph_compile.py", line 482, in aot_stage2_inference
  entry = _cache_inference_info(...)
File "torch/_functorch/_aot_autograd/graph_compile.py", line 536, in _cache_inference_info
  AOTAutogradCache.save(...)
File "torch/_functorch/_aot_autograd/autograd_cache.py", line 1299, in save
  AOTAutogradCache._handle_save_error(e, remote, is_bypass=False)
File "torch/_functorch/_aot_autograd/autograd_cache.py", line 1271, in _handle_save_error
  raise e
File "torch/_functorch/_aot_autograd/autograd_cache.py", line 1279, in save
  content = AOTAutogradCache._pickle_entry(entry, remote)
File "torch/_functorch/_aot_autograd/autograd_cache.py", line 1237, in _pickle_entry
  return pickle.dumps(entry)
_pickle.PicklingError: Can't pickle <function launcher at 0x7fc56cd2aca0>:
                       attribute lookup launcher on __main__ failed

Reproducibility

• Buildkite: https://buildkite.com/vllm/ci/builds/63497#019dd7d4-1952-41f4-826c-0b8dfe51be22
• Job: Torch Nightly Basic Models Tests (Extra Initialization) 1 (soft-failed; build itself is passed because the job is non-blocking)
• vLLM main commit: 8a8c9b564e (2026-04-29 nightly)

Related

A similar _pickle.PicklingError on triton launcher was filed against the torch 2.12 test wheel in pytorch/pytorch#180911 (now closed). That issue never reproduced in the 2.12 test-channel CI runs. This nightly failure has the same exception family but is exercised through the DFlashDraftModel speculative-decode path specifically — possibly a code path the earlier fix didn't cover, or a regression in nightly post-fix.

Diagnosis request

The cached AOTAutograd entry contains a triton-generated launcher function defined in a child __main__-scoped JIT module that pickle can't reach. Could a maintainer:

1. Confirm whether the launcher codegen now produces a non-pickleable function by default on nightly (regression vs. release branch).
2. Check whether the AOTAutogradCache save path should be defensively bypassing entries that contain triton functions, or whether the launcher should be made pickleable.

Links

• vLLM build: https://buildkite.com/vllm/ci/builds/63497
• Job: https://buildkite.com/vllm/ci/builds/63497#019dd7d4-1952-41f4-826c-0b8dfe51be22
• Related closed issue: pytorch/pytorch#180911
• First Failure April 16: https://buildkite.com/vllm/ci/builds/61575
• Success on April 14: https://buildkite.com/vllm/ci/builds/61210

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @bdhirsh @bobrenjc93 @aorenste

Update: bisect + bundled-torch-nightly evidence

Bisect across vLLM main nightly Full CI runs

The PicklingError reliably reproduces from 2026-04-17 onwards. The regression itself was introduced in the torch nightly built between 2026-04-14 and 2026-04-17.

Confirmed bundled torch nightly version in 63497

By inspecting the postmerge docker image config (ECR public manifest → image config blob → history), and cross-referencing the :docker: build image torch nightly job log:

torch==2.13.0.dev20260428+cu130
torchvision==0.27.0.dev20260428+cu130
torchaudio==2.11.0.dev20260428+cu130

Image build timeline:

The wheel dev20260428 was cut ~15 hours before fcea8749c22f landed. At image-build time (April 29 06:18 UTC) the April 29 nightly hadn't published yet, so dev20260428 was the most recent available — and it predates the fix.

Relationship to pytorch/pytorch#180911 (closed)

• The closing commit fcea8749c22f (PR pytorch/pytorch#181463) is a bypass workaround: it catches PicklingError from triton kernels in AOTAutogradCache.save and converts the failure to a warning + continue (cache becomes a no-op for that entry).
• Effect when the fix reaches the wheel: this hard error should become:

  W AOTAutograd cache unable to serialize compiled graph:
    Can't pickle <function launcher at 0x...>: attribute lookup launcher on __main__ failed

  and the test should pass.

Expected fix arrival in vLLM CI

The first PyTorch nightly that should include fcea8749c22f is 2.13.0.dev20260429+cu130 (cut from April 29 main HEAD). Once that wheel publishes and the next vLLM postmerge :docker: build image torch nightly job picks it up, this failure should convert from fatal to a soft warning.

Real (non-bypass) fix

The bypass ships, but the underlying issue is that triton's codegen now produces a launcher function defined in a __main__-scoped JIT module that pickle can't reach. A proper fix would either move the launcher definition into a stable importable module so attribute lookup launcher on __main__ succeeds, or have AOTAutogradCache._pickle_entry exclude triton launcher functions explicitly.

Where the issue lives

• Torch-side. The exception originates entirely inside torch/_functorch/_aot_autograd/autograd_cache.py. vLLM's only role is invoking torch.compile(..., fullgraph=True) on qwen3_dflash, which produces a graph that contains triton-launched kernels — a normal supported usage. No vLLM code change is needed.