pytorch / pytorch/pytorch

reuse_hash_fn leaves a subscripted capture pinned to the traced index

Open
#192,906 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged high priority module: correctness (silent) module: dynamo module: higher order operators oncall: pt2 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.