Repository: pytorch
Issue: compile PixArt-sigma error
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

After torch.compile, when I forward the net it show "TypeError: function takes exactly 12 arguments (17 given)"

### Error logs

Traceback (most recent call last):
  File "/home/sh9/workspace/PixArt-sigma/scripts/inference.py", line 327, in <module>
    
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/sh9/workspace/PixArt-sigma/scripts/inference.py", line 185, in visualize
    z,
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 1204, in sample
    model_prev_list = [self.model_fn(x, t)]
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 451, in model_fn
    return self.data_prediction_fn(x, t)
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 439, in data_prediction_fn
    noise = self.noise_prediction_fn(x, t)
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 433, in noise_prediction_fn
    return self.model(x, t)
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 406, in <lambda>
    self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 331, in model_fn
    noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/dpm_solver.py", line 289, in noise_pred_fn
    output = model(x, t_input, cond, **model_kwargs)
  File "/home/sh9/workspace/PixArt-sigma/scripts/inference.py", line 87, in torch_compile_inference
    model_out = model(x, timestep, y, data_info=data_info, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
    return fn(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/sh9/workspace/PixArt-sigma/diffusion/model/nets/PixArtMS.py", line 165, in forward
    def forward(self, x, timestep, y, mask=None, data_info=None, **kwargs):
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
    return fn(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
    return fn(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 917, in forward
    return compiled_fn(full_args)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 89, in g
    return f(*args)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 106, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 113, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 152, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 906, in __call__
    return self.get_current_callable()(inputs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 784, in run
    return model(new_inputs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 934, in _run_from_cache
    return compiled_graph.compiled_artifact(inputs)
  File "/tmp/torchinductor_sh9/se/csecnsndrjxqycfqndn3wtzltlhczifholvmqn5jffqeh6hin35a.py", line 1797, in call
    triton_poi_fused__to_copy_cat_0.run(arg548_1, buf1, 512, grid=grid(512), stream=stream0)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 635, in run
    self.autotune_to_one_config(*args, grid=grid, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 531, in autotune_to_one_config
    timings = self.benchmark_all_configs(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 507, in benchmark_all_configs
    timings = {
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 508, in <dictcomp>
    launcher: self.bench(launcher, *args, **kwargs)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 479, in bench
    return do_bench(kernel_call, rep=40, fast_flush=True)
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/utils.py", line 170, in do_bench
    return triton_do_bench(*args, **kwargs)[0]
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/triton/testing.py", line 101, in do_bench
    fn()
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/torch/_inductor/triton_heuristics.py", line 472, in kernel_call
    launcher(
  File "<string>", line 7, in launcher
  File "/home/sh9/miniconda3/envs/pixart/lib/python3.10/site-packages/triton/backends/nvidia/driver.py", line 370, in __call__
    self.launch(*args, **kwargs)
TypeError: function takes exactly 12 arguments (17 given)

### Minified repro

- clone PixArt-sigma
- copy code to PixArt-sigma/scripts
- python3 my_code.py

```
import sys
from pathlib import Path
current_file_path = Path(__file__).resolve()
sys.path.insert(0, str(current_file_path.parent.parent))
import torch
from diffusion.model.nets import PixArtMS_XL_2


device = "cuda" if torch.cuda.is_available() else "cpu"
latent_size = 1024 // 8
max_sequence_length = 300
pe_interpolation = {256: 0.5, 512: 1, 1024: 2, 2048: 4}     # trick for positional embedding interpolation
micro_condition = False
weight_dtype = torch.float16
print(f"Inference with {weight_dtype}")

# model setting
model = PixArtMS_XL_2(
    input_size=latent_size,
    pe_interpolation=pe_interpolation[1024],
    micro_condition=micro_condition,
    model_max_length=max_sequence_length,
).to(device).half()

# inference before compile
x = torch.randn(2, 4, 128, 128, dtype=torch.float32, device=device)
timestep = torch.tensor([999., 999.], dtype=torch.float32, device=device)
y = torch.randn(2, 1, 300, 4096, dtype=torch.float32, device=device)
mask = torch.randn(1, 300, dtype=torch.float32, device=device)
data_info={'img_hw': torch.tensor([[1024., 1024.]], dtype=torch.float32, device=device), 
            'aspect_ratio': torch.tensor([[1. ]], dtype=torch.float32, device=device)}
output_1 = model(x, timestep, y, mask, data_info=data_info)


# torch compile
model = torch.compile(model)

# inference
output_2 = model(x, timestep, y, mask, data_info=data_info)
```

### Versions

Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.3) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.19.5
Libc version: glibc-2.31

Python version: 3.10.0 (default, Mar  3 2022, 09:58:08) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-149-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A10
Nvidia driver version: 535.104.12
cuDNN version: Probably one of the following:
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.9.4
/usr/local/cuda-12.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.9.4
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
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 57 bits virtual
CPU(s):                          28
On-line CPU(s) list:             0-27
Thread(s) per core:              2
Core(s) per socket:              14
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           106
Model name:                      Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz
Stepping:                        6
CPU MHz:                         2300.000
BogoMIPS:                        4600.00
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       672 KiB
L1i cache:                       448 KiB
L2 cache:                        17.5 MiB
L3 cache:                        54 MiB
NUMA node0 CPU(s):               0-27
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Not affected
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:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid md_clear arch_capabilities

Versions of relevant libraries:
[pip3] came-pytorch==0.1.3
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] onnx==1.15.0
[pip3] onnx-graphsurgeon==0.5.2
[pip3] onnxruntime-gpu==1.16.3
[pip3] pytorch-triton==3.0.0+45fff310c8
[pip3] torch==2.3.0
[pip3] torchaudio==2.3.0
[pip3] torchvision==0.18.0
[pip3] triton==2.3.0
[conda] came-pytorch              0.1.3                    pypi_0    pypi
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch-triton            3.0.0+45fff310c8          pypi_0    pypi
[conda] torch                     2.3.0                    pypi_0    pypi
[conda] torchaudio                2.3.0                    pypi_0    pypi
[conda] torchvision               0.18.0                   pypi_0    pypi
[conda] triton                    2.3.0                    pypi_0    pypi

cc @ezyang @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @msaroufim