[compile] Allow strings in custom ops without regressing compilation times#38123
Open
zou3519 wants to merge 1 commit intovllm-project:mainfrom
Open
[compile] Allow strings in custom ops without regressing compilation times#38123zou3519 wants to merge 1 commit intovllm-project:mainfrom
zou3519 wants to merge 1 commit intovllm-project:mainfrom
Conversation
zou3519
requested review from
BoyuanFeng,
LucasWilkinson,
MatthewBonanni,
ProExpertProg,
sighingnow,
tdoublep and
youkaichao
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a ModuleName opaque type to improve torch.compile behavior by hoisting layer names as graph inputs, preventing per-layer recompilation for custom operations. This change involves updating various torch.ops.vllm calls across attention, KV cache, and Mamba mixer modules to use _encode_layer_name and _resolve_layer_name for consistent handling of layer names. A critical issue was identified where the kv_cache_dummy_dep variable could be undefined, leading to an UnboundLocalError.
zou3519
force-pushed
the
opaque_everything
branch
from
35626a5 to
99e2154
Compare
ProExpertProg
approved these changes
Mar 25, 2026
| qkvz_output_size: int, | ||
| ba_output_size: int, | ||
| layer_name: str, | ||
| layer_name: _layer_name_type, |
Collaborator
There was a problem hiding this comment.
Nit: can we call this LayerName or LayerNameType?
zou3519
force-pushed
the
opaque_everything
branch
from
99e2154 to
14a1342
Compare
5 tasks
5 tasks
|
This pull request has merge conflicts that must be resolved before it can be |
zou3519
force-pushed
the
opaque_everything
branch
2 times, most recently
from
807e2e7 to
594eb86
Compare
…times This is a follow-up to vllm-project#35475 to extend the fix to all custom operators, not just the MOE custom ops. Previously, string inputs to custom ops would regress compilation times. The problem goes: - a transformer model (e.g. llama3-70b) has 80 identical layers - we capture a full graph and the split the graph on the attention operations. - this produces 81 subgraphs: the middle 79 graphs are all identical (aside from graph inputs - parameters and buffers) - vLLM-compile produces 1 compiled artifact for all of the middle 79 subgraphs. - If a custom operator with a layer_name string appears in the graph, then this causes the middle 79 subgraphs to now be unique, so vLLM-compile ends up producing 79 compiled artifacts for them. In PyTorch 2.11, we have added a special class (the OpaqueObject type). The idea is that instead of passing strings to custom operators, we can pass a special LayerName object to the custom operator. This signifies to the compiler that it should "lift" the LayerName object to being a graph input and not bake the value directly into the graph. More notes: - the LayerName object used to be called ModuleName. I renamed it here, LayerName seemed more appropriate. - VLLM_USE_LAYERNAME=0 turns this feature off. This option is here just in case something breaks. I'll probably remove it in the next month. Signed-off-by: Richard Zou <zou3519@gmail.com>
zou3519
force-pushed
the
opaque_everything
branch
from
594eb86 to
04fc04b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a follow-up to #35475 to extend the fix to all custom operators, not just the MOE custom ops.
Previously, string inputs to custom ops would regress compilation times. The problem goes:
In PyTorch 2.11, we have added a special class (the OpaqueObject type). The idea is that instead of passing strings to custom operators, we can pass a special LayerName object to the custom operator. This signifies to the compiler that it should "lift" the LayerName object to being a graph input and not bake the value directly into the graph.
More things: