torch.compile: ConstDictVariable should support runtime dict lookup for dynamic keys to avoid recompilation
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
## 🚀 Feature / Improvement
### Problem
When a `torch.compile`'d `nn.Module` uses a float attribute as a dictionary key (e.g. `cache[self.key]`), Dynamo is forced to specialize on each concrete float value at compile time. This is because `ConstDictVariable.__getitem__` resolves dict lookups at trace time using `_HashableTracker`, which requires compile-time hash/equality matching. Each distinct `self.key` value produces a separate cache entry with a guard like `___as_tensor(self.key).item() == N.0`, leading to O(N) recompilations for N module instances.
With `inline_inbuilt_nn_modules=True` (the default), the module itself gets a `TYPE_MATCH` guard (not `ID_MATCH`), so there are no `id_matched_objs` to bucket cache entries by instance. All compilations count against the same `recompile_limit` (default 8), hitting the limit at the 9th instance.
### Repro
```python
import torch
import torch.nn as nn
cache = {}
class Module(nn.Module):
def __init__(self, key: float):
super().__init__()
self.key = key
cache[key] = torch.randn(16)
def forward(self, x):
return x + cache[self.key]
x = torch.randn(16)
for key in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]:
model = torch.compile(Module(key))
model(x)
```
**Output:**
```
torch._dynamo hit config.recompile_limit (8)
function: 'forward' (repro.py:14)
last reason: 0/7: ___as_tensor(self.key).item() == 8.0
```
### Proposed improvement
Teach `ConstDictVariable` to emit a **runtime** `dict.__getitem__` call when the lookup key is dynamic (e.g. an `UnspecializedPythonVariable` backed by a float), instead of requiring compile-time resolution. This would allow a single compiled graph to handle all instances regardless of the concrete value of `self.key`, eliminating recompilation entirely for this pattern.
Concretely, this would involve changes in `torch/_dynamo/variables/dicts.py`:
- In `getitem_const` / `getitem_const_raise_exception_if_absent`, detect when the key argument is dynamic (no compile-time hash match possible)
- Instead of raising `unimplemented` / `ObservedKeyError`, emit a graph node that performs `dict.__getitem__(cache, self.key)` at runtime
- Guard that the dict contains the expected keys (via `DICT_KEYS_MATCH` or similar)
### Why not fix the cache accounting instead?
This was explored in #171598 but rejected by reviewers. The two-tier cache system (`recompile_limit` / `accumulated_recompile_limit`) is being phased out with `inline_inbuilt_nn_modules=True` as the default. Modifying cache accounting to hide real recompilations was considered worse than surfacing them. The right fix is to eliminate the recompilation at its source.
### Related
- Original user report: #171593
- Attempted cache accounting fix (closed): #171598
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @Lucaskabela @jataylo @oulgen @jamesjwu @aorenste @anijain2305 @laithsakka @masnesral @coconutruben @aditvenk @williamwen42
Contributor guide
Assessment
This issue has not been assessed yet.