Repository: pytorch
Issue: capture_triton gets skipped by dynamo
Labels: triaged,actionable,oncall: pt2,module: dynamo,module: user triton

### 🐛 Describe the bug

Instead of just capturing the triton kernel, dynamo skips capture_triton.

Repro below:

```python
import torch

import triton
import triton.language as tl

@triton.jit
def triton_(out, inp, numel, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < numel
    vals = tl.load(inp + offsets, mask)
    result = vals * vals
    tl.store(out + offsets, vals, mask)


def call_triton(inp):
    out = torch.empty_like(inp)
    BLOCK_SIZE = 256
    numel = inp.numel()

    def grid(META):
        return (triton.cdiv(numel, META["BLOCK_SIZE"]),)

    torch._library.capture_triton(triton_)[grid](out, inp, numel, BLOCK_SIZE)
    return inp

x = torch.randn(15, device="cuda")
call_triton(x)
torch.compile(call_triton, fullgraph=True)(x)
```

Error message:
```
  File "/home/dberard/local/scripts/capture_triton_skipped.py", line 29, in <module>                                                                                                  15:08:41 [40/10111]
    torch.compile(call_triton, fullgraph=True)(x)
  File "/home/dberard/local/pytorch/torch/_dynamo/eval_frame.py", line 465, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 1278, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 526, in __call__
    return _compile(
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 924, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 666, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/home/dberard/local/pytorch/torch/_utils_internal.py", line 87, in wrapper_function
    return function(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 699, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/dberard/local/pytorch/torch/_dynamo/bytecode_transformation.py", line 1322, in transform_code_object
    transformations(instructions, code_options)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 219, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 634, in transform
    tracer.run()
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2796, in run
    super().run()
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 983, in run
    while self.step():
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 895, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 582, in wrapper
    return inner_fn(self, inst)
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1602, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 830, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/home/dberard/local/pytorch/torch/_dynamo/variables/functions.py", line 729, in call_function
    unimplemented(msg)
  File "/home/dberard/local/pytorch/torch/_dynamo/exc.py", line 297, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: 'skip function capture_triton in file /home/dberard/local/pytorch/torch/_library/triton.py'

from user code:
   File "/home/dberard/local/scripts/capture_triton_skipped.py", line 24, in call_triton
    torch._library.capture_triton(triton_)[grid](out, inp, numel, BLOCK_SIZE)
```

### Versions

main branch ~sep 11, H100

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec @oulgen @aakhundov @zou3519