[Bug]: torch.compile fails for Gemma3n on pytorch 2.8

[shinoss]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

run

vllm serve google/gemma-3n-E2B-it -tp 1

on torch==2.8.0:

(EngineCore_DP0 pid=2805721)   File "/data/users/yhshin/gitrepos/vllm/vllm/compilation/backends.py", line 330, in call_module
(EngineCore_DP0 pid=2805721)     compiler_manager.compile(
(EngineCore_DP0 pid=2805721)                      ^^^^^^^^
(EngineCore_DP0 pid=2805721)   File "/data/users/yhshin/gitrepos/vllm/vllm/compilation/backends.py", line 179, in compile
(EngineCore_DP0 pid=2805721)     compiled_graph, handle = self.compiler.compile(
(EngineCore_DP0 pid=2805721)                              ^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=2805721)   File "/data/users/yhshin/gitrepos/vllm/vllm/compilation/compiler_interface.py", line 217, in compile
(EngineCore_DP0 pid=2805721)     compiled_graph.save(path=path, format="unpacked")
(EngineCore_DP0 pid=2805721)   File "/home/yhshin/uv_env/vllm/lib/python3.12/site-packages/torch/_inductor/standalone_compile.py", line 73, in save
(EngineCore_DP0 pid=2805721)     assert len(cache_info.aot_autograd_artifacts) == 1, cache_info
(EngineCore_DP0 pid=2805721)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=2805721) torch._dynamo.exc.BackendCompilerFailed: backend='<vllm.compilation.backends.VllmBackend object at 0x7f5ad84e7d10>' raised:
(EngineCore_DP0 pid=2805721) AssertionError: CacheInfo(artifacts=defaultdict(<class 'list'>, {'inductor': ['frtb6xof4xglqbvww362mucq4jpbcm7hdkky7duecxlfqy7turzg'], 'autotune':...

fails in standalone_compile added in pytorch/pytorch#150670 due to no artifacts returned after compiling the second submodule (submod_2) of Gemma3n.

@zou3519 mentioned this happens when any output of a subgraph returns a view of the input

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[shinoss]

@heheda12345 gemma3n issue wasn't due to fast prefill changes, looks like a general issue with pytorch 2.8

[shinoss]

pytorch/pytorch#161994 may fix this

[zou3519]

I want to point out that the problem (a cache issue) existed in PyTorch 2.7 too. It's not just a PyTorch 2.8 problem. PyTorch 2.8 makes it error loudly so that the user knows the problem is there.

[jetnet]

It's iteresting: the release notes state "Gemma 3n supports this and that", but when you try to run it - it just crashes.
I'm wondering, how you guys run your tests then? Could you pls share your Dockerfile?

[gstrat88]

you can try with this option
export TORCHDYNAMO_DISABLE=1

[jetnet]

you can try with this option export TORCHDYNAMO_DISABLE=1

Thanks, the magic param helped! But, and this is a big but:

• v.10.2 + TORCHDYNAMO_DISABLE=1
  • Avg prompt throughput: 7.5 tokens/s, Avg generation throughput: 3.7 tokens/s
• v0.10.1.1-timm-fix
  • Avg prompt throughput: 36.6 tokens/s, Avg generation throughput: 19.6 tokens/s

So, the question remains - how do you guys test? :)

[gstrat88]

a huge but i would say, but for my current case it is acceptable
this is propably the best solution v0.10.1.1-timm-fix thanks for that!

[ProExpertProg]

@ElizaWszola can you post your repro instructions as well

[Lucaskabela]

I patched pytorch/pytorch#161994, onto 2.9.0 and then ran

vllm serve google/gemma-3n-E2B-it -tp 1

And got a successful build; then ran a minor benchmark to make sure everything was working:

vllm bench serve   --backend openai-chat   --endpoint-type openai-chat   --model google/gemma-3n-E2B-it   --endpoint /v1/chat/completions   --dataset-name random-mm   --num-prompts 100   --max-concurrency 10   --random-prefix-len 25   --random-input-len 300   --random-output-len 40   --random-range-ratio 0.2   --random-mm-base-items-per-request 2   --random-mm-num-mm-items-range-ratio 0.5   --random-mm-limit-mm-per-prompt '{"image": 4, "video": 0}'   --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'   --request-rate inf   --ignore-eos   --seed 42