Training in progress, step 680

.gitattributes

@@ -52,3 +52,6 @@ torchinductor_ch-epfl-345354-j/fxgraph/xb/fxbygtwj5f3jfikmhfoz7vd3mtceabjyw5ggsc
 torchinductor_ch-epfl-345354-j/triton/0/35ZDQUBPR56EFY64BDG6UB7OKJODSPJPYCTXXIKSOWZ2CA3EPWGA/triton_poi_fused__to_copy_cos_mul_sin_1.cubin filter=lfs diff=lfs merge=lfs -text
 torchinductor_ch-epfl-345354-j/triton/0/SS6ERMBZYPQUFPYGWYCICDOLY35WYNXDIHS3C45P5EZJOFLNB5SQ/triton_poi_fused__to_copy_cos_mul_sin_1.cubin filter=lfs diff=lfs merge=lfs -text
 torchinductor_ch-epfl-345354-j/triton/0/YGRXAKJ5T6UCGDY6RY3GMSSD4JVJFIRM6WYZVP7IT63SMM7NWHSA/triton_poi_fused__to_copy_cos_mul_sin_1.cubin filter=lfs diff=lfs merge=lfs -text
+torchinductor_ch-epfl-345354-j/fxgraph/3o/f3odjl3g6dddsqtuut46oufxq75twuknijzu35ozpse6s2ndf7rh/ssax6qk4l2jsc4ovqf64hhenn3rpvkiqhluzs7qfxcc6dzbigdx filter=lfs diff=lfs merge=lfs -text
+torchinductor_ch-epfl-345354-j/fxgraph/7y/f7yf7bpzpliuud5ylt3deq7gx7odarlyvhr2zdhcorkmqjxqlpv7/inrn63eoovhcqnqsled5nyfdmzwunnsaud5rcn37go5egvyuze3 filter=lfs diff=lfs merge=lfs -text
+torchinductor_ch-epfl-345354-j/fxgraph/p3/fp32tea4g6rsybgfroshvlbxseks7iiqpb4cmbb2d3lgralck774/vbrtzv3sqiuroqkjniy4wufikxjgiptmxb4qwrykjuzdrp2mvvt filter=lfs diff=lfs merge=lfs -text

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:26045c44c8f7b94239544f69255de30ea5e05c7f8e23f0c01a67e755eaa0beba
+oid sha256:17e7ee8610ecc91fe729f86ae30fbe910c6bf626a6ee7655d8d6496c22403797
 size 1192135096

torchinductor_ch-epfl-345354-j/bw/cbw2ylnzyqqvpuutlhdcnot22eonl4ebucgx7hpqh66fjy7zgkhb.py

@@ -0,0 +1,361 @@
+"""
+Compile-time auto-tuning block: 
+
+import torch
+from torch._dynamo.testing import rand_strided
+from torch._dynamo.utils import preserve_rng_state
+from torch._inductor.select_algorithm import AlgorithmSelectorCache
+from torch._inductor.async_compile import AsyncCompile
+
+async_compile = AsyncCompile()
+generate_example_value = AlgorithmSelectorCache.generate_example_value
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+
+
+triton_poi_fused_add_mul_neg_slice_backward_0 = async_compile.triton('triton_poi_fused_add_mul_neg_slice_backward_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)
+''', device_str='cuda')
+
+
+triton_poi_fused_add_mul_neg_slice_backward_1 = async_compile.triton('triton_poi_fused_add_mul_neg_slice_backward_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)
+''', device_str='cuda')
+
+async_compile.wait(globals())
+del async_compile
+
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+with torch.cuda._DeviceGuard(0):
+    torch.cuda.set_device(0)
+    stream0 = get_raw_stream(0)
+    from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+    stream0 = get_raw_stream(0)
+    tangents_2 = generate_example_value((4, 8, 177, 128), (181248, 22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    primals_5 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    primals_3 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf0 = generate_example_value((4, 8, 177, 128), (181248, 22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    triton_poi_fused_add_mul_neg_slice_backward_0.run(tangents_2, primals_5, primals_3, buf0, 128, 22656, 724992, stream=stream0)
+    del tangents_2, primals_5, primals_3, buf0
+
+    stream0 = get_raw_stream(0)
+    tangents_1 = generate_example_value((4, 16, 177, 128), (362496, 22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    primals_5 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    primals_3 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf1 = generate_example_value((4, 16, 177, 128), (362496, 22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    triton_poi_fused_add_mul_neg_slice_backward_1.run(tangents_1, primals_5, primals_3, buf1, 128, 22656, 1449984, stream=stream0)
+    del tangents_1, primals_5, primals_3, buf1
+
+"""
+# AOT ID: ['13_backward']
+from ctypes import c_void_p, c_long, c_int
+import torch
+import math
+import random
+import os
+import tempfile
+from math import inf, nan
+from cmath import nanj
+from torch._inductor.hooks import run_intermediate_hooks
+from torch._inductor.utils import maybe_profile
+from torch._inductor.codegen.memory_planning import _align as align
+from torch import device, empty_strided
+from torch._inductor.async_compile import AsyncCompile
+from torch._inductor.select_algorithm import extern_kernels
+from torch._inductor.codegen.multi_kernel import MultiKernelCall
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+
+aten = torch.ops.aten
+inductor_ops = torch.ops.inductor
+_quantized = torch.ops._quantized
+assert_size_stride = torch._C._dynamo.guards.assert_size_stride
+empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
+alloc_from_pool = torch.ops.inductor._alloc_from_pool
+async_compile = AsyncCompile()
+empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/j6/cj6yl63qj32gvorczxkszjnojwnybujlaozjrnzwzik4ukxq4rss.py
+# Topologically Sorted Source Nodes: [], Original ATen: [aten.neg, aten.slice_backward, aten.add, aten.mul]
+# Source node to ATen node mapping:
+# Graph fragment:
+#   %neg_2 : [num_users=1] = call_function[target=torch.ops.aten.neg.default](args = (%slice_5,), kwargs = {})
+#   %full_default : [num_users=2] = call_function[target=torch.ops.aten.full.default](args = ([%primals_9, %primals_10, %primals_1, %primals_2], 0), kwargs = {dtype: torch.bfloat16, layout: torch.strided, device: cuda:0, pin_memory: False})
+#   %slice_scatter_default : [num_users=1] = call_function[target=torch.ops.aten.slice_scatter.default](args = (%full_default, %neg_2, 3, %floordiv, 9223372036854775807), kwargs = {})
+#   %slice_scatter_default_1 : [num_users=1] = call_function[target=torch.ops.aten.slice_scatter.default](args = (%full_default, %slice_6, 3, 0, %floordiv), kwargs = {})
+#   %add_82 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%slice_scatter_default, %slice_scatter_default_1), kwargs = {})
+#   %mul_69 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%tangents_2, %unsqueeze), kwargs = {})
+#   %add_83 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add_82, %mul_69), kwargs = {})
+triton_poi_fused_add_mul_neg_slice_backward_0 = async_compile.triton('triton_poi_fused_add_mul_neg_slice_backward_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)
+''', device_str='cuda')
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/fj/cfjpgg4ctgzyci7wzylrvpq3i7o2o5r6brzgt5u6bzfhltzrjid3.py
+# Topologically Sorted Source Nodes: [], Original ATen: [aten.neg, aten.slice_backward, aten.add, aten.mul]
+# Source node to ATen node mapping:
+# Graph fragment:
+#   %neg_3 : [num_users=1] = call_function[target=torch.ops.aten.neg.default](args = (%slice_7,), kwargs = {})
+#   %full_default_2 : [num_users=2] = call_function[target=torch.ops.aten.full.default](args = ([%primals_6, %primals_7, %primals_1, %primals_2], 0), kwargs = {dtype: torch.bfloat16, layout: torch.strided, device: cuda:0, pin_memory: False})
+#   %slice_scatter_default_2 : [num_users=1] = call_function[target=torch.ops.aten.slice_scatter.default](args = (%full_default_2, %neg_3, 3, %floordiv, 9223372036854775807), kwargs = {})
+#   %slice_scatter_default_3 : [num_users=1] = call_function[target=torch.ops.aten.slice_scatter.default](args = (%full_default_2, %slice_8, 3, 0, %floordiv), kwargs = {})
+#   %add_88 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%slice_scatter_default_2, %slice_scatter_default_3), kwargs = {})
+#   %mul_71 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%tangents_1, %unsqueeze), kwargs = {})
+#   %add_89 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%add_88, %mul_71), kwargs = {})
+triton_poi_fused_add_mul_neg_slice_backward_1 = async_compile.triton('triton_poi_fused_add_mul_neg_slice_backward_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)
+''', device_str='cuda')
+
+
+async_compile.wait(globals())
+del async_compile
+
+def call(args):
+    primals_1, primals_2, primals_6, primals_7, primals_9, primals_10, floordiv, add_78, primals_3, primals_5, tangents_1, tangents_2 = args
+    args.clear()
+    s0 = primals_1
+    s1 = primals_2
+    s7 = primals_6
+    s8 = primals_7
+    s13 = primals_9
+    s14 = primals_10
+    assert_size_stride(primals_3, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(primals_5, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(tangents_1, (s7, s8, s0, s1), (s0*s1*s8, s0*s1, s1, 1))
+    assert_size_stride(tangents_2, (s13, s14, s0, s1), (s0*s1*s14, s0*s1, s1, 1))
+    with torch.cuda._DeviceGuard(0):
+        torch.cuda.set_device(0)
+        ps0 = s0*s1
+        buf0 = empty_strided_cuda((s13, s14, s0, s1), (s0*s1*s14, s0*s1, s1, 1), torch.bfloat16)
+        # Topologically Sorted Source Nodes: [], Original ATen: [aten.neg, aten.slice_backward, aten.add, aten.mul]
+        triton_poi_fused_add_mul_neg_slice_backward_0_xnumel = s0*s1*s13*s14
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_mul_neg_slice_backward_0.run(tangents_2, primals_5, primals_3, buf0, s1, ps0, triton_poi_fused_add_mul_neg_slice_backward_0_xnumel, stream=stream0)
+        del tangents_2
+        buf1 = empty_strided_cuda((s7, s8, s0, s1), (s0*s1*s8, s0*s1, s1, 1), torch.bfloat16)
+        # Topologically Sorted Source Nodes: [], Original ATen: [aten.neg, aten.slice_backward, aten.add, aten.mul]
+        triton_poi_fused_add_mul_neg_slice_backward_1_xnumel = s0*s1*s7*s8
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_mul_neg_slice_backward_1.run(tangents_1, primals_5, primals_3, buf1, s1, ps0, triton_poi_fused_add_mul_neg_slice_backward_1_xnumel, stream=stream0)
+        del primals_3
+        del primals_5
+        del tangents_1
+    return (None, None, None, None, None, None, None, buf1, None, None, buf0, )
+
+
+def benchmark_compiled_module(times=10, repeat=10):
+    from torch._dynamo.testing import rand_strided
+    from torch._inductor.utils import print_performance
+    primals_1 = 177
+    primals_2 = 128
+    primals_6 = 4
+    primals_7 = 16
+    primals_9 = 4
+    primals_10 = 8
+    floordiv = 64
+    add_78 = 64
+    primals_3 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    primals_5 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    tangents_1 = rand_strided((4, 16, 177, 128), (362496, 22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    tangents_2 = rand_strided((4, 8, 177, 128), (181248, 22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    fn = lambda: call([primals_1, primals_2, primals_6, primals_7, primals_9, primals_10, floordiv, add_78, primals_3, primals_5, tangents_1, tangents_2])
+    return print_performance(fn, times=times, repeat=repeat)
+
+
+if __name__ == "__main__":
+    from torch._inductor.wrapper_benchmark import compiled_module_main
+    compiled_module_main('None', benchmark_compiled_module)

torchinductor_ch-epfl-345354-j/da/0df06b8658127e5ce915248b1384d7c58762d116d843ac58b2d4be91c249301b.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 1024, "num_warps": 4, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 25}

torchinductor_ch-epfl-345354-j/da/cdapkyn6r3i3fmlikyctcbjr336rmrrljtxcaqcwrc7auok6odb4.py

@@ -0,0 +1,46 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)

torchinductor_ch-epfl-345354-j/fj/cfjpgg4ctgzyci7wzylrvpq3i7o2o5r6brzgt5u6bzfhltzrjid3.py

@@ -0,0 +1,48 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)

torchinductor_ch-epfl-345354-j/fj/fa90c6e1e9df4537c3b4e7a14a704c2df4c98a82068ce1f3abb21af8a35247a0.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 512, "num_warps": 8, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 51}

torchinductor_ch-epfl-345354-j/fxgraph/3o/f3odjl3g6dddsqtuut46oufxq75twuknijzu35ozpse6s2ndf7rh/ssax6qk4l2jsc4ovqf64hhenn3rpvkiqhluzs7qfxcc6dzbigdx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94817f415c5e932171d5817dc39c8d6ee2faa9103ae9997e05b885e1e42830ee
+size 599455

torchinductor_ch-epfl-345354-j/fxgraph/7y/f7yf7bpzpliuud5ylt3deq7gx7odarlyvhr2zdhcorkmqjxqlpv7/inrn63eoovhcqnqsled5nyfdmzwunnsaud5rcn37go5egvyuze3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4362df6c8e754e2836125907d6e0a3666d46b640a0fb11377f33ba435714c937
+size 600150

torchinductor_ch-epfl-345354-j/fxgraph/p3/fp32tea4g6rsybgfroshvlbxseks7iiqpb4cmbb2d3lgralck774/vbrtzv3sqiuroqkjniy4wufikxjgiptmxb4qwrykjuzdrp2mvvt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8633cd77282291741496a31cb50a855d2643e6cb8790b470a4d3238bf4cad66
+size 517209

torchinductor_ch-epfl-345354-j/j6/42b81ae014f6f1c38797163131fc53a645adf2cc0424a9255831be278e34af78.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 512, "num_warps": 8, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 51}

torchinductor_ch-epfl-345354-j/j6/cj6yl63qj32gvorczxkszjnojwnybujlaozjrnzwzik4ukxq4rss.py

@@ -0,0 +1,48 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_mul_neg_slice_backward_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 6, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_mul_neg_slice_backward_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x0 = (xindex % ks0)
+    x3 = xindex
+    x4 = (xindex % ks1)
+    tmp19 = tl.load(in_ptr0 + (x3), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp20 = tl.load(in_ptr2 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp0 = x0
+    tmp1 = ks0 // 2
+    tmp2 = tmp0 >= tmp1
+    tmp3 = tl.load(in_ptr0 + (x3 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp4 = tl.load(in_ptr1 + (x4 + (-1)*(ks0 // 2)), xmask & tmp2, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp5 = tmp3 * tmp4
+    tmp6 = -tmp5
+    tmp7 = tl.full(tmp6.shape, 0.0, tmp6.dtype)
+    tmp8 = tl.where(tmp2, tmp6, tmp7)
+    tmp9 = 0.0
+    tmp10 = tl.where(tmp2, tmp8, tmp9)
+    tmp11 = tmp0 < tmp1
+    tmp12 = tl.load(in_ptr0 + (ks0 + x3 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp13 = tl.load(in_ptr1 + (ks0 + x4 + (-1)*(ks0 // 2)), xmask & tmp11, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp14 = tmp12 * tmp13
+    tmp15 = tl.full(tmp14.shape, 0.0, tmp14.dtype)
+    tmp16 = tl.where(tmp11, tmp14, tmp15)
+    tmp17 = tl.where(tmp11, tmp16, tmp9)
+    tmp18 = tmp10 + tmp17
+    tmp21 = tmp19 * tmp20
+    tmp22 = tmp18 + tmp21
+    tl.store(out_ptr0 + (x3), tmp22, xmask)

torchinductor_ch-epfl-345354-j/lm/65cc8fd7bae3abb18e0f1b23c8e1c60711fe732716b5724936da88bfeb6409b4.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 1024, "num_warps": 4, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 38}

torchinductor_ch-epfl-345354-j/lm/clmeyn6qnatdy2hjtvp2smmgjoj5zqvy23kyx4cn7bbcafehg7wr.py

@@ -0,0 +1,46 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)

torchinductor_ch-epfl-345354-j/mc/cmcqiflrggybh55qgleeqzkbj3rlpmylxaqgl47omwv5z3y4ifv5.py

@@ -0,0 +1,353 @@
+"""
+Compile-time auto-tuning block: 
+
+import torch
+from torch._dynamo.testing import rand_strided
+from torch._dynamo.utils import preserve_rng_state
+from torch._inductor.select_algorithm import AlgorithmSelectorCache
+from torch._inductor.async_compile import AsyncCompile
+
+async_compile = AsyncCompile()
+generate_example_value = AlgorithmSelectorCache.generate_example_value
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+
+
+triton_poi_fused_add_cat_mul_0 = async_compile.triton('triton_poi_fused_add_cat_mul_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+triton_poi_fused_add_cat_mul_1 = async_compile.triton('triton_poi_fused_add_cat_mul_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+async_compile.wait(globals())
+del async_compile
+
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+with torch.cuda._DeviceGuard(0):
+    torch.cuda.set_device(0)
+    stream0 = get_raw_stream(0)
+    from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+    stream0 = get_raw_stream(0)
+    primals_8 = generate_example_value((4, 16, 177, 128), (362496, 128, 2048, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    primals_3 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    primals_5 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf0 = generate_example_value((4, 16, 177, 128), (362496, 128, 2048, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    triton_poi_fused_add_cat_mul_0.run(primals_8, primals_3, primals_5, buf0, 128, 2048, 177, 1449984, stream=stream0)
+    del primals_8, primals_3, primals_5, buf0
+
+    stream0 = get_raw_stream(0)
+    primals_11 = generate_example_value((4, 8, 177, 128), (181248, 128, 1024, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    primals_3 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    primals_5 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf1 = generate_example_value((4, 8, 177, 128), (181248, 128, 1024, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    triton_poi_fused_add_cat_mul_1.run(primals_11, primals_3, primals_5, buf1, 128, 1024, 177, 724992, stream=stream0)
+    del primals_11, primals_3, primals_5, buf1
+
+"""
+# AOT ID: ['13_forward']
+from ctypes import c_void_p, c_long, c_int
+import torch
+import math
+import random
+import os
+import tempfile
+from math import inf, nan
+from cmath import nanj
+from torch._inductor.hooks import run_intermediate_hooks
+from torch._inductor.utils import maybe_profile
+from torch._inductor.codegen.memory_planning import _align as align
+from torch import device, empty_strided
+from torch._inductor.async_compile import AsyncCompile
+from torch._inductor.select_algorithm import extern_kernels
+from torch._inductor.codegen.multi_kernel import MultiKernelCall
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+
+aten = torch.ops.aten
+inductor_ops = torch.ops.inductor
+_quantized = torch.ops._quantized
+assert_size_stride = torch._C._dynamo.guards.assert_size_stride
+empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
+alloc_from_pool = torch.ops.inductor._alloc_from_pool
+async_compile = AsyncCompile()
+empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/ua/cuaufinvbemnicqkymrm4ujjq5b6ltwjbdc4jcgwj2zxj4usa5kx.py
+# Topologically Sorted Source Nodes: [mul, cat, mul_1, q_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+# Source node to ATen node mapping:
+#   cat => cat
+#   mul => mul_8
+#   mul_1 => mul_29
+#   q_embed => add_36
+# Graph fragment:
+#   %mul_8 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_8, %unsqueeze), kwargs = {})
+#   %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%neg, %slice_1], -1), kwargs = {})
+#   %mul_29 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%cat, %unsqueeze_1), kwargs = {})
+#   %add_36 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_8, %mul_29), kwargs = {})
+triton_poi_fused_add_cat_mul_0 = async_compile.triton('triton_poi_fused_add_cat_mul_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/lm/clmeyn6qnatdy2hjtvp2smmgjoj5zqvy23kyx4cn7bbcafehg7wr.py
+# Topologically Sorted Source Nodes: [mul_2, cat_1, mul_3, k_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+# Source node to ATen node mapping:
+#   cat_1 => cat_1
+#   k_embed => add_72
+#   mul_2 => mul_38
+#   mul_3 => mul_59
+# Graph fragment:
+#   %mul_38 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_11, %unsqueeze), kwargs = {})
+#   %cat_1 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%neg_1, %slice_3], -1), kwargs = {})
+#   %mul_59 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%cat_1, %unsqueeze_1), kwargs = {})
+#   %add_72 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_38, %mul_59), kwargs = {})
+triton_poi_fused_add_cat_mul_1 = async_compile.triton('triton_poi_fused_add_cat_mul_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+async_compile.wait(globals())
+del async_compile
+
+def call(args):
+    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9, primals_10, primals_11 = args
+    args.clear()
+    s0 = primals_1
+    s1 = primals_2
+    s7 = primals_6
+    s8 = primals_7
+    s13 = primals_9
+    s14 = primals_10
+    assert_size_stride(primals_3, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(primals_5, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(primals_8, (s7, s8, s0, s1), (s0*s1*s8, s1, s1*s8, 1))
+    assert_size_stride(primals_11, (s13, s14, s0, s1), (s0*s1*s14, s1, s1*s14, 1))
+    with torch.cuda._DeviceGuard(0):
+        torch.cuda.set_device(0)
+        ps0 = s1*s8
+        buf0 = empty_strided_cuda((s7, s8, s0, s1), (s0*s1*s8, s1, s1*s8, 1), torch.bfloat16)
+        # Topologically Sorted Source Nodes: [mul, cat, mul_1, q_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+        triton_poi_fused_add_cat_mul_0_xnumel = s0*s1*s7*s8
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_cat_mul_0.run(primals_8, primals_3, primals_5, buf0, s1, ps0, s0, triton_poi_fused_add_cat_mul_0_xnumel, stream=stream0)
+        del primals_8
+        ps1 = s1*s14
+        buf1 = empty_strided_cuda((s13, s14, s0, s1), (s0*s1*s14, s1, s1*s14, 1), torch.bfloat16)
+        # Topologically Sorted Source Nodes: [mul_2, cat_1, mul_3, k_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+        triton_poi_fused_add_cat_mul_1_xnumel = s0*s1*s13*s14
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_cat_mul_1.run(primals_11, primals_3, primals_5, buf1, s1, ps1, s0, triton_poi_fused_add_cat_mul_1_xnumel, stream=stream0)
+        del primals_11
+    return (buf0, buf1, primals_3, primals_5, s0, s1, s7, s8, s13, s14, s1 // 2, s1 + (-1)*(s1 // 2), )
+
+
+def benchmark_compiled_module(times=10, repeat=10):
+    from torch._dynamo.testing import rand_strided
+    from torch._inductor.utils import print_performance
+    primals_1 = 177
+    primals_2 = 128
+    primals_3 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    primals_4 = 1
+    primals_5 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    primals_6 = 4
+    primals_7 = 16
+    primals_8 = rand_strided((4, 16, 177, 128), (362496, 128, 2048, 1), device='cuda:0', dtype=torch.bfloat16)
+    primals_9 = 4
+    primals_10 = 8
+    primals_11 = rand_strided((4, 8, 177, 128), (181248, 128, 1024, 1), device='cuda:0', dtype=torch.bfloat16)
+    fn = lambda: call([primals_1, primals_2, primals_3, primals_4, primals_5, primals_6, primals_7, primals_8, primals_9, primals_10, primals_11])
+    return print_performance(fn, times=times, repeat=repeat)
+
+
+if __name__ == "__main__":
+    from torch._inductor.wrapper_benchmark import compiled_module_main
+    compiled_module_main('None', benchmark_compiled_module)

torchinductor_ch-epfl-345354-j/sw/cswtunku7iwygay5azmeyjzubvchodclgrov3ytwoqfke3cbjeek.py

@@ -0,0 +1,357 @@
+"""
+Compile-time auto-tuning block: 
+
+import torch
+from torch._dynamo.testing import rand_strided
+from torch._dynamo.utils import preserve_rng_state
+from torch._inductor.select_algorithm import AlgorithmSelectorCache
+from torch._inductor.async_compile import AsyncCompile
+
+async_compile = AsyncCompile()
+generate_example_value = AlgorithmSelectorCache.generate_example_value
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+
+
+triton_poi_fused_add_cat_mul_0 = async_compile.triton('triton_poi_fused_add_cat_mul_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+triton_poi_fused_add_cat_mul_1 = async_compile.triton('triton_poi_fused_add_cat_mul_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+async_compile.wait(globals())
+del async_compile
+
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+with torch.cuda._DeviceGuard(0):
+    torch.cuda.set_device(0)
+    stream0 = get_raw_stream(0)
+    from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+    stream0 = get_raw_stream(0)
+    arg7_1 = generate_example_value((4, 16, 177, 128), (362496, 128, 2048, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    arg2_1 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    arg4_1 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf0 = generate_example_value((4, 16, 177, 128), (362496, 128, 2048, 1), 'cuda:0', torch.bfloat16, 0, (4, 16, 177, 128))
+    triton_poi_fused_add_cat_mul_0.run(arg7_1, arg2_1, arg4_1, buf0, 128, 2048, 177, 1449984, stream=stream0)
+    del arg7_1, arg2_1, arg4_1, buf0
+
+    stream0 = get_raw_stream(0)
+    arg10_1 = generate_example_value((4, 8, 177, 128), (181248, 128, 1024, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    arg2_1 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    arg4_1 = generate_example_value((1, 177, 128), (22656, 128, 1), 'cuda:0', torch.bfloat16, 0, (1, 177, 128))
+    buf1 = generate_example_value((4, 8, 177, 128), (181248, 128, 1024, 1), 'cuda:0', torch.bfloat16, 0, (4, 8, 177, 128))
+    triton_poi_fused_add_cat_mul_1.run(arg10_1, arg2_1, arg4_1, buf1, 128, 1024, 177, 724992, stream=stream0)
+    del arg10_1, arg2_1, arg4_1, buf1
+
+"""
+# AOT ID: ['12_inference']
+from ctypes import c_void_p, c_long, c_int
+import torch
+import math
+import random
+import os
+import tempfile
+from math import inf, nan
+from cmath import nanj
+from torch._inductor.hooks import run_intermediate_hooks
+from torch._inductor.utils import maybe_profile
+from torch._inductor.codegen.memory_planning import _align as align
+from torch import device, empty_strided
+from torch._inductor.async_compile import AsyncCompile
+from torch._inductor.select_algorithm import extern_kernels
+from torch._inductor.codegen.multi_kernel import MultiKernelCall
+import triton
+import triton.language as tl
+from torch._inductor.runtime.triton_heuristics import start_graph, end_graph
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+from torch._C import _cuda_getCurrentRawStream as get_raw_stream
+
+aten = torch.ops.aten
+inductor_ops = torch.ops.inductor
+_quantized = torch.ops._quantized
+assert_size_stride = torch._C._dynamo.guards.assert_size_stride
+empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
+empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
+empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
+reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
+alloc_from_pool = torch.ops.inductor._alloc_from_pool
+async_compile = AsyncCompile()
+empty_strided_p2p = torch._C._distributed_c10d._SymmetricMemory.empty_strided_p2p
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/sy/csyofljuv4nybcwnbyfhmdfufvm6xfyypzxkd745lbsexq35kdpk.py
+# Topologically Sorted Source Nodes: [mul, cat, mul_1, q_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+# Source node to ATen node mapping:
+#   cat => cat
+#   mul => mul_8
+#   mul_1 => mul_29
+#   q_embed => add_36
+# Graph fragment:
+#   %mul_8 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%arg7_1, %unsqueeze), kwargs = {})
+#   %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%neg, %slice_1], -1), kwargs = {})
+#   %mul_29 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%cat, %unsqueeze_1), kwargs = {})
+#   %add_36 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_8, %mul_29), kwargs = {})
+triton_poi_fused_add_cat_mul_0 = async_compile.triton('triton_poi_fused_add_cat_mul_0', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+# kernel path: /tmp/torchinductor_ch-epfl-345354-j/da/cdapkyn6r3i3fmlikyctcbjr336rmrrljtxcaqcwrc7auok6odb4.py
+# Topologically Sorted Source Nodes: [mul_2, cat_1, mul_3, k_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+# Source node to ATen node mapping:
+#   cat_1 => cat_1
+#   k_embed => add_72
+#   mul_2 => mul_38
+#   mul_3 => mul_59
+# Graph fragment:
+#   %mul_38 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%arg10_1, %unsqueeze), kwargs = {})
+#   %cat_1 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%neg_1, %slice_3], -1), kwargs = {})
+#   %mul_59 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%cat_1, %unsqueeze_1), kwargs = {})
+#   %add_72 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul_38, %mul_59), kwargs = {})
+triton_poi_fused_add_cat_mul_1 = async_compile.triton('triton_poi_fused_add_cat_mul_1', '''
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 1048576}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_1', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)
+''', device_str='cuda')
+
+
+async_compile.wait(globals())
+del async_compile
+
+def call(args):
+    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1 = args
+    args.clear()
+    s0 = arg0_1
+    s1 = arg1_1
+    s7 = arg5_1
+    s8 = arg6_1
+    s13 = arg8_1
+    s14 = arg9_1
+    assert_size_stride(arg2_1, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(arg4_1, (1, s0, s1), (s0*s1, s1, 1))
+    assert_size_stride(arg7_1, (s7, s8, s0, s1), (s0*s1*s8, s1, s1*s8, 1))
+    assert_size_stride(arg10_1, (s13, s14, s0, s1), (s0*s1*s14, s1, s1*s14, 1))
+    with torch.cuda._DeviceGuard(0):
+        torch.cuda.set_device(0)
+        ps0 = s1*s8
+        pool1 = empty_strided_cuda((s7, s8, s0, s1), (s0*s1*s8, s1, s1*s8, 1), torch.bfloat16)
+        buf0 = pool1  # alloc
+        # Topologically Sorted Source Nodes: [mul, cat, mul_1, q_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+        triton_poi_fused_add_cat_mul_0_xnumel = s0*s1*s7*s8
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_cat_mul_0.run(arg7_1, arg2_1, arg4_1, buf0, s1, ps0, s0, triton_poi_fused_add_cat_mul_0_xnumel, stream=stream0)
+        del arg7_1
+        ps1 = s1*s14
+        pool0 = empty_strided_cuda((s13, s14, s0, s1), (s0*s1*s14, s1, s1*s14, 1), torch.bfloat16)
+        buf1 = pool0  # alloc
+        # Topologically Sorted Source Nodes: [mul_2, cat_1, mul_3, k_embed], Original ATen: [aten.mul, aten.cat, aten.add]
+        triton_poi_fused_add_cat_mul_1_xnumel = s0*s1*s13*s14
+        stream0 = get_raw_stream(0)
+        triton_poi_fused_add_cat_mul_1.run(arg10_1, arg2_1, arg4_1, buf1, s1, ps1, s0, triton_poi_fused_add_cat_mul_1_xnumel, stream=stream0)
+        del arg10_1
+        del arg2_1
+        del arg4_1
+    return (buf0, buf1, )
+
+
+def benchmark_compiled_module(times=10, repeat=10):
+    from torch._dynamo.testing import rand_strided
+    from torch._inductor.utils import print_performance
+    arg0_1 = 177
+    arg1_1 = 128
+    arg2_1 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    arg3_1 = 1
+    arg4_1 = rand_strided((1, 177, 128), (22656, 128, 1), device='cuda:0', dtype=torch.bfloat16)
+    arg5_1 = 4
+    arg6_1 = 16
+    arg7_1 = rand_strided((4, 16, 177, 128), (362496, 128, 2048, 1), device='cuda:0', dtype=torch.bfloat16)
+    arg8_1 = 4
+    arg9_1 = 8
+    arg10_1 = rand_strided((4, 8, 177, 128), (181248, 128, 1024, 1), device='cuda:0', dtype=torch.bfloat16)
+    fn = lambda: call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1])
+    return print_performance(fn, times=times, repeat=repeat)
+
+
+if __name__ == "__main__":
+    from torch._inductor.wrapper_benchmark import compiled_module_main
+    compiled_module_main('None', benchmark_compiled_module)

torchinductor_ch-epfl-345354-j/sy/4d7e479de298eba85b0b0b838eb227f9e0129c56097eee3ae7375c0bd11c1c7d.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 1024, "num_warps": 4, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 25}

torchinductor_ch-epfl-345354-j/sy/csyofljuv4nybcwnbyfhmdfufvm6xfyypzxkd745lbsexq35kdpk.py

@@ -0,0 +1,46 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': True, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)

torchinductor_ch-epfl-345354-j/ua/5c2b35108d0f6d838f30c1b016c317aa650f5f275110f17a9aff57ce54e364e6.best_config

@@ -0,0 +1 @@
+{"XBLOCK": 1024, "num_warps": 4, "num_stages": 1, "configs_hash": "3ca5c3e34d35093f3c9ab2829a9faeebad5e61c4ca13d5ed6053d7b71ce60d5a", "found_by_coordesc": false, "time_taken_ms": 38}

torchinductor_ch-epfl-345354-j/ua/cuaufinvbemnicqkymrm4ujjq5b6ltwjbdc4jcgwj2zxj4usa5kx.py

@@ -0,0 +1,46 @@
+
+import triton
+import triton.language as tl
+
+from torch._inductor.runtime import triton_helpers, triton_heuristics
+from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
+from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
+triton_helpers.set_driver_to_gpu()
+
+@triton_heuristics.pointwise(
+    size_hints={'x': 2097152}, 
+    filename=__file__,
+    triton_meta={'signature': {'in_ptr0': '*bf16', 'in_ptr1': '*bf16', 'in_ptr2': '*bf16', 'out_ptr0': '*bf16', 'ks0': 'i32', 'ks1': 'i32', 'ks2': 'i32', 'xnumel': 'i32', 'XBLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=42, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]]}]},
+    inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_cat_mul_0', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': '8D9A40F96256AE993B0CB3DAC1136935BA540F7848683690590C84AF795CC5ED', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
+    min_elem_per_thread=0
+)
+@triton.jit
+def triton_poi_fused_add_cat_mul_0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr):
+    xoffset = tl.program_id(0) * XBLOCK
+    xindex = xoffset + tl.arange(0, XBLOCK)[:]
+    xmask = xindex < xnumel
+    x4 = xindex
+    x0 = (xindex % ks0)
+    x2 = ((xindex // ks1) % ks2)
+    x5 = xindex // ks0
+    tmp0 = tl.load(in_ptr0 + (x4), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp1 = tl.load(in_ptr1 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp17 = tl.load(in_ptr2 + (x0 + ks0*x2), xmask, eviction_policy='evict_last').to(tl.float32)
+    tmp2 = tmp0 * tmp1
+    tmp3 = x0
+    tmp4 = tl.full([1], 0, tl.int64)
+    tmp5 = tmp3 >= tmp4
+    tmp6 = ks0 + (-1)*(ks0 // 2)
+    tmp7 = tmp3 < tmp6
+    tmp8 = tl.load(in_ptr0 + (ks0*x5 + (ks0 // 2) + (x0)), xmask & tmp7, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp9 = -tmp8
+    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
+    tmp11 = tl.where(tmp7, tmp9, tmp10)
+    tmp12 = tmp3 >= tmp6
+    tmp13 = ks0
+    tmp14 = tmp3 < tmp13
+    tmp15 = tl.load(in_ptr0 + (ks0*x5 + (x0 + ((-1)*ks0) + (ks0 // 2))), xmask & tmp12, eviction_policy='evict_last', other=0.0).to(tl.float32)
+    tmp16 = tl.where(tmp7, tmp11, tmp15)
+    tmp18 = tmp16 * tmp17
+    tmp19 = tmp2 + tmp18
+    tl.store(out_ptr0 + (x4), tmp19, xmask)