Repository: pytorch
Issue: Support Mutation on Epilogue Fusion on template
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

It seems that `torch.compile` does not fuse matmul followed by `index_add`, even though scatter (`index_put_`) is considered pointwise and should theoretically be fusable.

When attempting to fuse `op1 (matmul)` -> `op2 (index_put_)`, the inductor throws a fusion error due to `index_put_` involving mutation https://github.com/pytorch/pytorch/blob/0b75b7ff2b8ab8f40e433a52b06a671d6377997f/torch/_inductor/scheduler.py#L3102

Here's the code snippet:

```python
import torch

def matmul(a, b, n, d):
    C = torch.matmul(a, b)
    e = d.index_add_(dim=0, index=n, source=C)
    return e

a = torch.randn(100, 100, device="cuda")
b = torch.randn(100, 100, device="cuda")
d = torch.randn(100, 100, device="cuda")
n = torch.randint(100, (100,), device="cuda")

torch._inductor.epilogue_fusion_first = True
torch._inductor.epilogue_fusion = True
torch._inductor.benchmark_epilogue_fusion = True
compiled = torch.compile(matmul, mode="max-autotune-no-cudagraphs")
result = compiled(a, b, n, d)
```

When I manually inserted scatter into the matmul template, it produced the correct result, which suggests these two operations should be fusable. However, if I remove the mutation check `node2.has_aliasing_or_mutation()`, another error occurs during fusion (see attached).

Is there a strict correctness requirement preventing the fusion of mutation nodes with templates, or is this something that could be addressed?

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @jansel 


### Error logs

## Error Logs
```
$ TORCH_LOGS="output_code,fusion" python test.py

  warnings.warn(
V1212 13:20:34.264000 791207 site-packages/torch/_inductor/scheduler.py:2272] [0/0] [__fusion] ===== attempting fusion (1/10): 2 nodes =====
V1212 13:20:34.265000 791207 site-packages/torch/_inductor/scheduler.py:2547] [0/0] [__fusion] fuse_nodes_once, candidates:
V1212 13:20:34.265000 791207 site-packages/torch/_inductor/scheduler.py:2549] [0/0] [__fusion]   SchedulerNode(name='op0')
V1212 13:20:34.266000 791207 site-packages/torch/_inductor/scheduler.py:2549] [0/0] [__fusion]   SchedulerNode(name='op1'), Scatter(['[100, 100]', 'origins=OrderedSet([index_put])'])
V1212 13:20:34.266000 791207 site-packages/torch/_inductor/scheduler.py:2673] [0/0] [__fusion] found 1 possible fusions
V1212 13:20:34.267000 791207 site-packages/torch/_inductor/scheduler.py:2558] [0/0] [__fusion] fusing op0 with op1
V1212 13:20:34.267000 791207 site-packages/torch/_inductor/scheduler.py:2279] [0/0] [__fusion] completed fusion round (1/10): fused 2 nodes into 1 nodes
V1212 13:20:34.267000 791207 site-packages/torch/_inductor/scheduler.py:2279] [0/0] [__fusion]
V1212 13:20:34.267000 791207 site-packages/torch/_inductor/scheduler.py:2286] [0/0] [__fusion] ===== fusion complete (1 iterations) =====
Traceback (most recent call last):
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1446, in _call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/repro/after_dynamo.py", line 129, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/__init__.py", line 2234, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1253, in compile_fx
    return compile_fx(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1521, in compile_fx
    return aot_autograd(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/backends/common.py", line 72, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1071, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1056, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 522, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 759, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 179, in aot_dispatch_base
    compiled_fw = compiler(fw_module, updated_flat_args)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1350, in fw_compiler_base
    return _fw_compiler_base(model, example_inputs, is_inference)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1421, in _fw_compiler_base
    return inner_compile(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 475, in compile_fx_inner
    return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/repro/after_aot.py", line 85, in debug_wrapper
    inner_compiled_fn = compiler_fn(gm, example_inputs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 661, in _compile_fx_inner
    compiled_graph = FxGraphCache.load(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 1334, in load
    compiled_graph = compile_fx_fn(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 570, in codegen_and_compile
    compiled_graph = fx_codegen_and_compile(gm, example_inputs, **fx_kwargs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 878, in fx_codegen_and_compile
    compiled_fn = graph.compile_to_fn()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1913, in compile_to_fn
    return self.compile_to_module().call
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1839, in compile_to_module
    return self._compile_to_module()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1845, in _compile_to_module
    self.codegen_with_cpp_wrapper() if self.cpp_wrapper else self.codegen()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1784, in codegen
    self.scheduler.codegen()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/scheduler.py", line 3383, in codegen
    return self._codegen()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/scheduler.py", line 3445, in _codegen
    self.get_backend(device).codegen_template(node, epilogue)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/codegen/cuda_combined_scheduling.py", line 75, in codegen_template
    return self._triton_scheduling.codegen_template(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_inductor/codegen/simd.py", line 1461, in codegen_template
    kernel, render = template_node.node.make_kernel_render(template_node.node)
TypeError: 'NoneType' object is not callable

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/jaeyeon/Playground/halide_pointcloud/template_fusion_test/test.py", line 22, in <module>
    result = compiled(a, b, n,d)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 465, in _fn
    return fn(*args, **kwargs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1269, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1064, in __call__
  1
    result = self._inner_convert(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 526, in __call__
    return _compile(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 924, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 666, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_utils_internal.py", line 87, in wrapper_function
    return function(*args, **kwargs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 699, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1322, in transform_code_object
    transformations(instructions, code_options)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 219, in _fn
    return fn(*args, **kwargs)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 634, in transform
    tracer.run()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2796, in run
    super().run()
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 983, in run
    while self.step():
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 895, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2987, in RETURN_VALUE
    self._return(inst)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2972, in _return
    self.output.compile_subgraph(
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1117, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1369, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1416, in call_user_compiler
    return self._call_user_compiler(gm)
  File "/home/jaeyeon/miniconda3/envs/halide/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1465, in _call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e) from e
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
TypeError: 'NoneType' object is not callable

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

## ir_post_fusion.txt

```
op0_op1: FusedSchedulerNode(SchedulerNode,SchedulerNode)
op0_op1.writes =
    [   MemoryDep('buf0', c0, {c0: 10000}, None),
        MemoryDep('buf1', c1 + 100*tmp0, {c0: 100, c1: 100}, atomic_add)]
op0_op1.unmet_dependencies = []
op0_op1.met_dependencies =
    [   MemoryDep('arg3_1', c0, {c0: 100}, None),
        StarDep(name='arg0_1', mode=None),
        StarDep(name='arg1_1', mode=None),
        StarDep(name='arg2_1', mode='atomic_add')]
op0_op1.outputs = [
    buf0: MultiTemplateBuffer
    buf0.layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
    buf0.users = [NodeUser(node=SchedulerNode(name='op1'), can_inplace=False, is_weak=False)]
    buf1: ComputedBuffer
    buf1.layout = MutationLayoutSHOULDREMOVE('cuda', torch.float32, size=[100, 100], stride=[100, 1])
    buf1.mutations = ['arg2_1']
    buf1.users = [NodeUser(node=OUTPUT, can_inplace=False, is_weak=False)]
]
op0_op1.snodes[0] =
op0: SchedulerNode(MultiTemplateBuffer)
op0.writes = [MemoryDep('buf0', c0, {c0: 10000}, None)]
op0.unmet_dependencies = []
op0.met_dependencies = [StarDep(name='arg0_1', mode=None), StarDep(name='arg1_1', mode=None)]
op0.outputs = [
    buf0: MultiTemplateBuffer
    buf0.layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
    buf0.users = [NodeUser(node=SchedulerNode(name='op1'), can_inplace=False, is_weak=False)]
]
op0.group.device = cuda:0
op0.group.iteration = (10000, 1)
op0.sizes = ([100, 100], ())
arg1_1_layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
arg0_1_layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
buf0_layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
op0 Unfinalized multi template buffer
op0_op1.snodes[1] =
op1: SchedulerNode(ComputedBuffer)
op1.writes = [MemoryDep('buf1', c1 + 100*tmp0, {c0: 100, c1: 100}, atomic_add)]
op1.unmet_dependencies = [MemoryDep('buf0', c0, {c0: 10000}, None)]
op1.met_dependencies =
    [   MemoryDep('arg3_1', c0, {c0: 100}, None),
        StarDep(name='arg2_1', mode='atomic_add')]
op1.outputs = [
    buf1: ComputedBuffer
    buf1.layout = MutationLayoutSHOULDREMOVE('cuda', torch.float32, size=[100, 100], stride=[100, 1])
    buf1.mutations = ['arg2_1']
    buf1.users = [NodeUser(node=OUTPUT, can_inplace=False, is_weak=False)]
]
op1.group.device = cuda:0
op1.group.iteration = (10000, 1)
op1.sizes = ([100, 100], [])
arg3_1_layout = FixedLayout('cuda', torch.int64, size=[100], stride=[1])
buf0_layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
arg2_1_layout = FixedLayout('cuda', torch.float32, size=[100, 100], stride=[100, 1])
buf1_layout = MutationLayoutSHOULDREMOVE('cuda', torch.float32, size=[100, 100], stride=[100, 1])
class op1_loop_body:
    var_ranges = {z0: 100, z1: 100}
    index0 = z0
    index1 = 100*z0 + z1
    index2 = 100*indirect0 + z1
    def body(self, ops):
        get_index = self.get_index('index0')
        load = ops.load('arg3_1', get_index)
        set_indirect0 = self.set_indirect0(load)
        get_index_1 = self.get_index('index1')
        load_1 = ops.load('buf0', get_index_1)
        get_index_2 = self.get_index('index2')
        store = ops.store('buf1', get_index_2, load_1, 'atomic_add')
        return store
op1 Triton code:
    import triton
    import triton.language as tl
    from triton.compiler.compiler import AttrsDescriptor

    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, instance_descriptor, DeviceProperties

    @triton_heuristics.pointwise(
        size_hints=[16384],
        filename=__file__,
        triton_meta={'signature': {0: '*i64', 1: '*fp32', 2: '*fp32', 3: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=86, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=1536, multi_processor_count=82), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3), equal_to_1=())]},
        inductor_meta={'autotune_hints': set(), 'kernel_name': 'Placeholder.DESCRIPTIVE_NAME', 'mutated_arg_names': ['out_ptr0'], 'no_x_dim': False, 'num_load': 2, 'num_reduction': 0, 'backend_hash': 'EB0670E067E4DCC123BDA265AA7F4C62765F41E5CCF5A66A7FA518D5515E450E', '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': True, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
        min_elem_per_thread=0
    )
    @triton.jit
    def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
        xnumel = 10000
        xoffset = tl.program_id(0) * XBLOCK
        xindex = xoffset + tl.arange(0, XBLOCK)[:]
        xmask = xindex < xnumel
        x1 = (xindex // 100)
        x2 = xindex
        x0 = xindex % 100
        tmp0 = tl.load(in_ptr0 + (x1), xmask, eviction_policy='evict_last')
        tmp6 = tl.load(in_ptr1 + (x2), xmask)
        tmp1 = tl.full([XBLOCK], 100, tl.int32)
        tmp2 = tmp0 + tmp1
        tmp3 = tmp0 < 0
        tmp4 = tl.where(tmp3, tmp2, tmp0)
        tl.device_assert(((0 <= tmp4) & (tmp4 < 100)) | ~(xmask), "index out of bounds: 0 <= tmp4 < 100")
        tl.atomic_add(out_ptr0 + (x0 + (100*tmp4)), tmp6, xmask, sem='relaxed')
op0_op1 Unfinalized multi template buffer
```

### Versions

```
Collecting environment information...
PyTorch version: 2.5.1
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: 10.0.0-4ubuntu1
CMake version: version 3.31.0-rc1
Libc version: glibc-2.31

Python version: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-89-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.66
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 535.183.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      39 bits physical, 48 bits virtual
CPU(s):                             20
On-line CPU(s) list:                0-19
Thread(s) per core:                 2
Core(s) per socket:                 10
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              165
Model name:                         Intel(R) Core(TM) i9-10900K CPU @ 3.70GHz
Stepping:                           5
CPU MHz:                            3700.000
CPU max MHz:                        5300.0000
CPU min MHz:                        800.0000
BogoMIPS:                           7399.70
Virtualization:                     VT-x
L1d cache:                          320 KiB
L1i cache:                          320 KiB
L2 cache:                           2.5 MiB
L3 cache:                           20 MiB
NUMA node0 CPU(s):                  0-19
Vulnerability Gather data sampling: Vulnerable: No microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Vulnerable: Clear CPU buffers attempted, no microcode; SMT vulnerable
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: 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:                Vulnerable: No microcode
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 dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp pku ospke md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] jax-triton==0.2.0
[pip3] numpy==2.0.1
[pip3] nvidia-cublas-cu12==12.6.3.3
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.5.1.17
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-nccl-cu12==2.23.4
[pip3] nvidia-nvjitlink-cu12==12.6.77
[pip3] nvtx==0.2.10
[pip3] pynvjitlink-cu12==0.4.0
[pip3] torch==2.5.1
[pip3] torch_cluster==1.6.3+pt24cu121
[pip3] torch-geometric==2.6.1
[pip3] torch_scatter==2.1.2+pt24cu121
[pip3] torch_sparse==0.6.18+pt24cu121
[pip3] torch_spline_conv==1.2.2+pt24cu121
[pip3] torchaudio==2.5.1
[pip3] torchsparse==2.1.0
[pip3] torchvision==0.20.1
[pip3] triton==3.1.0
[conda] blas                      1.0                         mkl
[conda] cuda-cudart               12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-cudart-dev           12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-cudart-static        12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-cupti                12.1.62                       0    nvidia/label/cuda-12.1.0
[conda] cuda-cupti-static         12.1.62                       0    nvidia/label/cuda-12.1.0
[conda] cuda-libraries            12.1.0                        0    nvidia/label/cuda-12.1.0
[conda] cuda-libraries-dev        12.1.0                        0    nvidia/label/cuda-12.1.0
[conda] cuda-libraries-static     12.1.0                        0    nvidia/label/cuda-12.1.0
[conda] cuda-nvrtc                12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-nvrtc-dev            12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-nvrtc-static         12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] cuda-nvtx                 12.1.66                       0    nvidia/label/cuda-12.1.0
[conda] cuda-opencl               12.1.56                       0    nvidia/label/cuda-12.1.0
[conda] cuda-opencl-dev           12.1.56                       0    nvidia/label/cuda-12.1.0
[conda] cuda-runtime              12.1.0                        0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] jax-triton                0.2.0                    pypi_0    pypi
[conda] libcublas                 12.1.0.26                     0    nvidia/label/cuda-12.1.0
[conda] libcublas-dev             12.1.0.26                     0    nvidia/label/cuda-12.1.0
[conda] libcublas-static          12.1.0.26                     0    nvidia/label/cuda-12.1.0
[conda] libcufft                  11.0.2.4                      0    nvidia/label/cuda-12.1.0
[conda] libcufft-dev              11.0.2.4                      0    nvidia/label/cuda-12.1.0
[conda] libcufft-static           11.0.2.4                      0    nvidia/label/cuda-12.1.0
[conda] libcurand                 10.3.2.56                     0    nvidia/label/cuda-12.1.0
[conda] libcurand-dev             10.3.2.56                     0    nvidia/label/cuda-12.1.0
[conda] libcurand-static          10.3.2.56                     0    nvidia/label/cuda-12.1.0
[conda] libcusolver               11.4.4.55                     0    nvidia/label/cuda-12.1.0
[conda] libcusolver-dev           11.4.4.55                     0    nvidia/label/cuda-12.1.0
[conda] libcusolver-static        11.4.4.55                     0    nvidia/label/cuda-12.1.0
[conda] libcusparse               12.0.2.55                     0    nvidia/label/cuda-12.1.0
[conda] libcusparse-dev           12.0.2.55                     0    nvidia/label/cuda-12.1.0
[conda] libcusparse-static        12.0.2.55                     0    nvidia/label/cuda-12.1.0
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch
[conda] libnvjitlink              12.1.105                      0    nvidia
[conda] libnvjitlink-dev          12.1.55                       0    nvidia/label/cuda-12.1.0
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-fft                   1.3.11                   pypi_0    pypi
[conda] mkl-random                1.2.8                    pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.11          py310h5eee18b_0
[conda] mkl_random                1.2.8           py310h1128e8f_0
[conda] numpy                     2.0.1                    pypi_0    pypi
[conda] numpy-base                2.0.1           py310hb5e798b_1
[conda] nvidia-cublas-cu12        12.6.3.3                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.6.80                  pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.6.77                  pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.5.1.17                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.3.0.4                 pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.7.1.2                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.5.4.2                 pypi_0    pypi
[conda] nvidia-nccl-cu12          2.23.4                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.6.77                  pypi_0    pypi
[conda] nvtx                      0.2.10                   pypi_0    pypi
[conda] pynvjitlink-cu12          0.4.0                    pypi_0    pypi
[conda] pytorch                   2.5.1           py3.10_cuda12.1_cudnn9.1.0_0    pytorch
[conda] pytorch-cuda              12.1                 ha16c6d3_6    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     2.5.1                    pypi_0    pypi
[conda] torch-cluster             1.6.3+pt24cu121          pypi_0    pypi
[conda] torch-geometric           2.6.1                    pypi_0    pypi
[conda] torch-scatter             2.1.2+pt24cu121          pypi_0    pypi
[conda] torch-sparse              0.6.18+pt24cu121          pypi_0    pypi
[conda] torch-spline-conv         1.2.2+pt24cu121          pypi_0    pypi
[conda] torchaudio                2.5.1                    pypi_0    pypi
[conda] torchsparse               2.1.0                    pypi_0    pypi
[conda] torchtriton               3.1.0                     py310    pytorch
[conda] torchvision               0.20.1                   pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
```