reuse_hash_fn leaves a subscripted capture pinned to the traced index
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`stamp_out_subgraph` remaps argument sources on a `reuse_hash_fn` cache hit, but a capture selected by a subscript stays pointed at the index the region traced with, so every reuse reads the first call's element.
```python
import torch
ncr = torch.compiler.nested_compile_region
class Pool:
def __init__(self, buffers):
self.buffers = buffers
class Layer(torch.nn.Module):
def __init__(self, layer_id, pool):
super().__init__()
self.layer_id = layer_id
self.pool = pool
def forward(self, x):
return x.sin() + self.pool.buffers[self.layer_id]
# The key ignores layer_id on purpose: the caller is asserting "same region,
# different slot", which is the reason to reach for reuse_hash_fn on a KV cache.
@ncr(reuse_hash_fn=lambda layer, x: 0)
def gn(layer, x):
return layer(x)
n = 4
pool = Pool([torch.ones(4) * (i + 1) * 10 for i in range(n)])
layers = [Layer(i, pool) for i in range(n)]
def fn(x):
return sum(gn(layer, x) for layer in layers)
x = torch.zeros(4)
print(fn(x)) # 100.0
print(torch.compile(fn, backend="aot_eager", fullgraph=True)(x)) # 40.0
```
### Actual vs expected
`buffers = [10, 20, 30, 40]`, eager gives `100.0`, compiled gives `40.0`, which is slot 0 read four times.
One could argue a constant key is the caller's mistake. The defect is narrower than that: argument sources *are* parameterized on a hash key hit while subscripted captures silently are not, and the mismatch produces a wrong number rather than a refusal. Parameterizing the selected element is also what #191781 item 3 asks for.
#192631 fixes this path. Filing separately so the behaviour is findable, since the fix is not otherwise described as a bug.
### Versions
main at 98ebd10b379, CPU, `backend="aot_eager"`.
cc @ezyang @gchanan @kadeng @msaroufim @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @amjames @jataylo @azahed98 @ydwu4 @bdhirsh @bobrenjc93 @aorenste
Contributor guide
Research direction
Start with the issue's reproduction and trace the reuse_hash_fn cache-hit path through stamp_out_subgraph. Check how argument sources and subscripted captures are remapped, then rerun the aot_eager fullgraph example; done means compiled reuse selects each layer's buffer and matches eager output. Issue #192631 is noted as fixing this path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100