Lightning-AI / Lightning-AI/lightning-thunder

Make prims.shape fusible with nvFuser executor

Open
#2,730 0 comments 0 reactions 1 assignee View on GitHub

@beverlylytle is already working on this.

Since Nov 12, 2025.

enhancement fusion logic nvfuser primitives
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature

Currently, `prims.shape` is not supported by the nvFuser executor leading to simple torch.add operation to not be fused with nvFuser if it involved broadcasting and it's not a simple translation to `prims.add`.

### Motivation

We need `prims.shape` to occur in the nvFuser fusible regions as a lot of operations have broadcasting behavior and for dynamic shapes efficient execution shape retrieval must be in the fused region and not outside.

Test program:
```python
import torch, thunder
from functools import partial
from thunder import prims

a = torch.randn(1, 1024, device="cuda", dtype=torch.float32)
b = torch.randn(2, 1024, device="cuda", dtype=torch.float32)

@partial(thunder.jit, cache="symbolic values")
def f(a, b):
(i0, i1) = prims.shape(a)
(i2, i3) = prims.shape(b)
a = prims.broadcast_in_dim(a, (i2, i1), (0, 1))
return prims.add(a, b)

f(a, b)

print(f._lc_cs.last_traces[-1])
```
Prints:
```py
def computation(a, b):
# a: "cuda:0 f32[[IntegerProxy name=i0, value=1, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# b: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i3, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"

# /tmp/ipython-input-1588267087.py:10: (i0, i1) = prims.shape(a)
(i0, i1) = shape(a)
del i0

# /tmp/ipython-input-1588267087.py:11: (i2, i3) = prims.shape(b)
(i2, i3) = shape(b)
del i3
[t9] = nvFusion0(a, i2, i1, b)
# t7 = prims.broadcast_in_dim(a, (i2, i1), (0, 1)) # t7: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
# t9 = prims.add(t7, b) # t9: "cuda:0 f32[[IntegerProxy name=i2, value=2, static=CONSTRAINT.CONSTRAINABLE], [IntegerProxy name=i1, value=1024, static=CONSTRAINT.CONSTRAINABLE]]"
del i2, i1
return (t9,)
```

cc @tfogal

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.