Repository: pytorch/pytorch

Issue:
 Strange “IndexError” when compiling
### 🐛 Describe the bug

I ran into a strange error where the compiler failed to identify/match the signature of the function.

```python
import torch
import torch.nn as nn

class Model(nn.Module):
   def __init__(self, n_embd=768, bias=False):
        super().__init__()
        self.n_head = 6
        self.embd = nn.Embedding(50257, 768)
        self.c_attn = nn.Linear(in_features=n_embd, out_features=3 * n_embd, bias=bias)
        self.dropout = 0

    def forward(self, x):
        x = self.embd(x)
        (B, T, C) = x.size()
        q, k, v = self.c_attn(x).chunk(chunks=3, dim=-1)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        y = torch.nn.functional.scaled_dot_product_attention(
            query=q, key=k, value=v, attn_mask=None, dropout_p=self.dropout, is_causal=True
        )
        return y
```

The forward pass works correctly:

```python
model = Model().cuda()
batch_size = 2
x = torch.randint(0, 50257, (batch_size, 1024)).cuda()
y = model(x)
```

However, if I compile the model:

```python
model = torch.compile(model)
_ = model(x)
```

The compiler throws an indexing error (see below for error message).

Strangely enough, the compiler works correctly when I remove the parameter names from the function call (i.e. “query/key/value”)

```python
import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, n_embd=768, bias=False):
        super().__init__()
        self.n_head = 6
        self.embd = nn.Embedding(50257, 768)
        self.c_attn = nn.Linear(in_features=n_embd, out_features=3 * n_embd, bias=bias)
        self.dropout = 0

    def forward(self, x):
        x = self.embd(x)
        (B, T, C) = x.size()
        q, k, v = self.c_attn(x).chunk(chunks=3, dim=-1)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        y = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, attn_mask=None, dropout_p=self.dropout, is_causal=True
        )
        return y
```

```python
model = Model().cuda()
batch_size = 2
x = torch.randint(0, 50257, (batch_size, 1024)).cuda()
model = torch.compile(model)
_ = model(x)
y = model(x)
```

Posted this on the PyTorch Forum too:

* [Strange “IndexError” when compiling](https://discuss.pytorch.org/t/strange-indexerror-when-compiling/173191)

## Full Error Message

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:324, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, hooks, frame)
    323 try:
--> 324     out_code = transform_code_object(code, transform)
    325     orig_code_map[out_code] = code

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py:445, in transform_code_object(code, transformations, safe)
    443 propagate_line_nums(instructions)
--> 445 transformations(instructions, code_options)
    446 return clean_and_assemble_instructions(instructions, keys, code_options)[1]

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:311, in _compile.<locals>.transform(instructions, code_options)
    299 tracer = InstructionTranslator(
    300     instructions,
    301     code,
   (...)
    309     mutated_closure_cell_contents,
    310 )
--> 311 tracer.run()
    312 output = tracer.output

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1738, in InstructionTranslator.run(self)
   1737 _step_logger()(logging.INFO, f"torchdynamo start tracing {self.f_code.co_name}")
-> 1738 super().run()

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:588, in InstructionTranslatorBase.run(self)
    584 self.output.push_tx(self)
    585 while (
    586     self.instruction_pointer is not None
    587     and not self.output.should_exit
--> 588     and self.step()
    589 ):
    590     pass

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:552, in InstructionTranslatorBase.step(self)
    551     unimplemented(f"missing: {inst.opname}")
--> 552 getattr(self, inst.opname)(inst)
    554 return inst.opname != "RETURN_VALUE"

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:342, in break_graph_if_unsupported.<locals>.decorator.<locals>.wrapper(self, inst)
    341 try:
--> 342     return inner_fn(self, inst)
    343 except Unsupported as excp:

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:1026, in InstructionTranslatorBase.CALL_FUNCTION_KW(self, inst)
   1025 assert len(kwargs) == len(argnames)
-> 1026 self.call_function(fn, args, kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py:486, in InstructionTranslatorBase.call_function(self, fn, args, kwargs)
    485     raise AssertionError(f"Attempt to trace forbidden callable {inner_fn}")
--> 486 self.push(fn.call_function(self, args, kwargs))

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/variables/torch.py:484, in TorchVariable.call_function(self, tx, args, kwargs)
    481 if self.value == torch._C._nn.scaled_dot_product_attention:
    482     # See:[Note] SDPA_flash's meta function returns incorrect Philox seed and offset
    483     # in pytorch/torch/_meta_registrations.py
--> 484     fake_query = args[0].as_proxy().node.meta["example_value"]
    485     fake_key = args[1].as_proxy().node.meta["example_value"]

IndexError: list index out of range

from user code:
   File "/tmp/ipykernel_24294/2653966783.py", line 19, in forward
    y = torch.nn.functional.scaled_dot_product_attention(

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True


The above exception was the direct cause of the following exception:

InternalTorchDynamoError                  Traceback (most recent call last)
Cell In[5], line 2
      1 model = torch.compile(model)
----> 2 _ = model(x)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:82, in OptimizedModule.forward(self, *args, **kwargs)
     81 def forward(self, *args, **kwargs):
---> 82     return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:209, in _TorchDynamoContext.__call__.<locals>._fn(*args, **kwargs)
    207 dynamic_ctx.__enter__()
    208 try:
--> 209     return fn(*args, **kwargs)
    210 finally:
    211     set_eval_frame(prior)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:337, in catch_errors_wrapper.<locals>.catch_errors(frame, cache_size)
    334             return hijacked_callback(frame, cache_size, hooks)
    336 with compile_lock:
--> 337     return callback(frame, cache_size, hooks)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:404, in convert_frame.<locals>._convert_frame(frame, cache_size, hooks)
    402 counters["frames"]["total"] += 1
    403 try:
--> 404     result = inner_convert(frame, cache_size, hooks)
    405     counters["frames"]["ok"] += 1
    406     return result

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:104, in wrap_convert_context.<locals>._fn(*args, **kwargs)
    102 torch.fx.graph_module._forward_from_src = fx_forward_from_src_skip_result
    103 try:
--> 104     return fn(*args, **kwargs)
    105 finally:
    106     torch._C._set_grad_enabled(prior_grad_mode)

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:262, in convert_frame_assert.<locals>._convert_frame_assert(frame, cache_size, hooks)
    259 global initial_grad_state
    260 initial_grad_state = torch.is_grad_enabled()
--> 262 return _compile(
    263     frame.f_code,
    264     frame.f_globals,
    265     frame.f_locals,
    266     frame.f_builtins,
    267     compiler_fn,
    268     one_graph,
    269     export,
    270     hooks,
    271     frame,
    272 )

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/utils.py:163, in dynamo_timed.<locals>.dynamo_timed_inner.<locals>.time_wrapper(*args, **kwargs)
    161     compilation_metrics[key] = []
    162 t0 = time.time()
--> 163 r = func(*args, **kwargs)
    164 time_spent = time.time() - t0
    165 # print(f"Dynamo timer: key={key}, latency={latency:.2f} sec")

File /opt/conda/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py:394, in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, hooks, frame)
    392 except Exception as e:
    393     exception_handler(e, code, frame)
--> 394     raise InternalTorchDynamoError() from e

InternalTorchDynamoError:
```










### Versions

Collecting environment information...
PyTorch version: 2.0.0.dev20230220+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.35

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-60-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A6000
Nvidia driver version: 525.60.13
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   46 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          36
On-line CPU(s) list:             0-35
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
CPU family:                      6
Model:                           85
Thread(s) per core:              2
Core(s) per socket:              18
Socket(s):                       1
Stepping:                        7
CPU max MHz:                     4800.0000
CPU min MHz:                     1200.0000
BogoMIPS:                        6000.00
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512_vnni md_clear flush_l1d arch_capabilities
Virtualization:                  VT-x
L1d cache:                       576 KiB (18 instances)
L1i cache:                       576 KiB (18 instances)
L2 cache:                        18 MiB (18 instances)
L3 cache:                        24.8 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-35
Vulnerability Itlb multihit:     KVM: Mitigation: VMX disabled
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] intel-extension-for-pytorch==1.13.100
[pip3] numpy==1.24.1
[pip3] torch==2.0.0.dev20230220+cu118
[pip3] torchaudio==2.0.0.dev20230221+cu118
[pip3] torchvision==0.15.0.dev20230221+cu118
[conda] intel-extension-for-pytorch 1.13.100                 pypi_0    pypi
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] torch                     2.0.0.dev20230220+cu118          pypi_0    pypi
[conda] torchaudio                2.0.0.dev20230221+cu118          pypi_0    pypi
[conda] torchvision               0.15.0.dev20230221+cu118          pypi_0    pypi

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh


Discussion:
cc @drisspg 
Will take a look
I will push a fix 
@drisspg  thank you!

Resolving patch:
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
--- a/torch/_dynamo/variables/torch.py
+++ b/torch/_dynamo/variables/torch.py
@@ -481,9 +481,34 @@ def get_state_from_generator():
             if self.value == torch._C._nn.scaled_dot_product_attention:
                 # See:[Note] SDPA_flash's meta function returns incorrect Philox seed and offset
                 # in pytorch/torch/_meta_registrations.py
-                fake_query = args[0].as_proxy().node.meta["example_value"]
-                fake_key = args[1].as_proxy().node.meta["example_value"]
-                fake_value = args[2].as_proxy().node.meta["example_value"]
+                all_kwargs = kwargs.copy()
+                all_kwargs.update(
+                    dict(
+                        zip(
+                            (
+                                "query",
+                                "key",
+                                "value",
+                                "attn_mask",
+                                "dropout_p",
+                                "is_causal",
+                            ),
+                            args,
+                        )
+                    )
+                )
+                fake_query = all_kwargs["query"].as_proxy().node.meta["example_value"]
+                fake_key = all_kwargs["key"].as_proxy().node.meta["example_value"]
+                fake_value = all_kwargs["value"].as_proxy().node.meta["example_value"]
+                fake_mask = all_kwargs.get("attn_mask")
+                if isinstance(fake_mask, TensorVariable):
+                    fake_mask = fake_mask.as_proxy().node.meta["example_value"]
+                else:
+                    fake_mask = None
+                dropout_p = kwargs.get("dropout_p")
+                dropout_p = dropout_p.value if dropout_p is not None else 0.0
+                is_causal = kwargs.get("is_causal")
+                is_causal = is_causal.value if is_causal is not None else False
                 # We look through the stack to find a cuda autocast context
                 # If we do we will convert the fake tensors to torch.float16
                 is_cuda_autocast_context = False
@@ -502,15 +527,10 @@ def get_state_from_generator():
                     fake_value = fake_value.clone().to(amp_dtype)
 
                 backend_choice = torch._fused_sdp_choice(
-                    fake_query, fake_key, fake_value
+                    fake_query, fake_key, fake_value, fake_mask, dropout_p, is_causal
                 )
                 if backend_choice == torch.backends.cuda.SDPBackend.FLASH_ATTENTION:
-                    dropout_p = kwargs.get("dropout_p")
-                    # Lets see if they passed it in as not an arg
-                    if len(args) >= 5:
-                        dropout_p = args[4]
-
-                    if dropout_p is not None and dropout_p.value != 0.0:
+                    if dropout_p is not None and dropout_p != 0.0:
                         unimplemented(
                             "FlashAttention with dropout is not supported in cuda graphs"
                         )


Test patch:
diff --git a/test/dynamo/test_dynamic_shapes.py b/test/dynamo/test_dynamic_shapes.py
--- a/test/dynamo/test_dynamic_shapes.py
+++ b/test/dynamo/test_dynamic_shapes.py
@@ -60,6 +60,11 @@ def make_dynamic_cls(cls):
     # Cannot call sizes() on tensor with symbolic sizes/strides
 )
 
+unittest.expectedFailure(
+    DynamicShapesMiscTests.test_parsing_sdpa_dynamic_shapes
+    # Cannot call sizes() on tensor with symbolic sizes/strides
+)
+
 
 # DynamicShapesSubGraphTests
 unittest.expectedFailure(
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
--- a/test/dynamo/test_misc.py
+++ b/test/dynamo/test_misc.py
@@ -3145,6 +3145,53 @@ def forward(self, query, key, value):
         self.assertEqual(compiled.device.index, 0)
         self.assertEqual(compiled.dtype, torch.float16)
 
+    @unittest.skipIf(
+        not PLATFORM_SUPPORTS_FUSED_SDPA or not SM80OrLater,
+        "Can't run fused SDPA on this platform",
+    )
+    def test_parsing_sdpa(self):
+        class MyModule(torch.nn.Module):
+            def forward(self, query, key, value):
+                out = F.scaled_dot_product_attention(query, key, value, None, 0, True)
+                out = F.scaled_dot_product_attention(
+                    query=query,
+                    key=key,
+                    value=value,
+                    attn_mask=None,
+                    dropout_p=0,
+                    is_causal=True,
+                )
+                out = F.scaled_dot_product_attention(
+                    query,
+                    key=key,
+                    value=value,
+                    attn_mask=None,
+                    dropout_p=0,
+                    is_causal=True,
+                )
+                out = F.scaled_dot_product_attention(
+                    query, key, value, None, dropout_p=0, is_causal=True
+                )
+                return out
+
+        device = "cuda"
+        dtype = torch.float16
+        seq_len_q = 1
+        seq_len_k = 1
+        head_dim = 8
+        query = torch.ones(
+            1, 8, seq_len_q, head_dim, device=device, dtype=dtype, requires_grad=True
+        )
+        key = torch.ones(
+            1, 8, seq_len_k, head_dim, device=device, dtype=dtype, requires_grad=True
+        )
+        value = torch.ones(
+            1, 8, seq_len_k, head_dim, device=device, dtype=dtype, requires_grad=True
+        )
+        module = MyModule()
+        opt_mod = torch._dynamo.optimize("inductor")(module)
+        opt_mod(query, key, value)
+
     def test_autocast_cpu(self):
         class MyModule(torch.nn.Module):
             def forward(self, x):