more general inplace support (index_copy_ in litgpt fails to trace)

[ali-alshaar7]

Note: If you have a model or program that is not supported yet but should be, please use the program coverage template.

🐛 Bug

To Reproduce

Steps to reproduce the behavior:

1. Go to '...'
2. Run '....'
3. Scroll down to '....'
4. See error

Code sample

import thunder, torch, litgpt

with torch.device('cuda'):
    m = litgpt.GPT.from_name('Llama-3.2-1B').bfloat16().requires_grad_(False)
    m.max_seq_length = 1024
    m.set_kv_cache(16)
inp1 = torch.ones(2, 16, device="cuda", dtype=torch.int32)
inp_pos1 = torch.stack([torch.arange(16, device="cuda") for _ in range(2)])
jm = thunder.jit(m)
jm(inp1, inp_pos1)

fail with

File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 547, in get_computation_and_inputs
    orig_to_view_swap_map = check_inplace_to_views(computation_trc)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/functionalization.py", line 100, in check_inplace_to_views
    check(
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/baseutils.py", line 146, in check
    raise exception_type(s())
NotImplementedError: in-place op of `torch.Tensor.index_copy_` to `<TensorProxy(name="t130", dtype=thunder.dtypes.bfloat16, shape=(8, 1024, 64))>`, a view tensor of `<TensorProxy(name="t128", dtype=thunder.dtypes.bfloat16, shape=(2, 8, 1024, 64))>` is not supported because 524288 != 1048576

Expected behavior

it compiles

Environment

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[t-vi]

This is expected as you will only use the first of the 16 batch entries in the kv cache and then write to it.
Note that the two values 16 are for two different quantities: batch size for KV cache and sequence length.

(It would be nice to have general inplace resolution, but we currently don't.)