Use nvFuser executor decisions to pass on op execution to a different backend and retire hybrid `torch_compile_cat_ex` executor. about lightning-thunder HOT 2 OPEN

Another alternative option which doesn't involve nvFuser magically having the context on what is RoPE and what is the best group of ops to send to another executor -

User defined backend as a decorator.

In litGPT for example, if we have a pytorch function that does RoPE, we change the function to something like

@thunder.jit(force_backend=torch_compile_executor)
def apply_rope(inputs):
........

so when thunder is constructing the trace, any operations traced inside this function are pushed to the forced executor backend.

Pros:

• User gets control over execution
• Solves pattern matching for at least semi-ninja users. users can do this for anything and use custom kernels registered as backends in Thunder.
• Alleviates our RoPE perf issues

Cons:

• Adds complexity to user behavior and usage.
• Might force changes into how traces are constructed and transforms are done.

from lightning-thunder.

triage review —

• it would be good to talk through the issue and possible solutions during a technical discussion
• @kiya00 is gathering some additional benchmarking data that may inform our thinking here
• https://github.com/Lightning-AI/lightning-thunder/blob/main/notebooks/zero_to_thunder.ipynb has an extensibility section that may inform the rope mapping to execution (and how to divert it)

from lightning-thunder.