Lightning-AI / Lightning-AI/lightning-thunder
Replace usage of nvFuser's broadcast_in_dim in executor with broadcast+expand ops
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
🚀 Feature
Replace usage of nvFuser's broadcast_in_dim in the executor with broadcast + expand operations.
Motivation
Currently, the Thunder executor uses nvFuser's broadcast_in_dim for broadcasting, as seen here.
Jacob suggested (in this PR review) that using explicit broadcast and expand ops would be preferable. This change would improve the developer experience by creating a closer match between Python-defined fusions and nvFuser's "math/kernel" representation.
Pitch
Update the executor logic to replace the direct usage of broadcast_in_dim with a combination of broadcast and expand operations where applicable.
Alternatives
Continue using broadcast_in_dim as it is currently implemented.
Additional context
- Current implementation: https://github.com/Lightning-AI/lightning-thunder/blob/9ca2b0d65e491f564dc905bfe80b1777f50daafb/thunder/executors/nvfuserex_impl.py#L1180-L1189
- Suggestion source: https://github.com/NVIDIA/Fuser/pull/5507#pullrequestreview-3478756789
cc @mruberry
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read thunder/executors/nvfuserex_impl.py around lines 1180-1189 and the referenced nvFuser review for context. Trace the executor’s current broadcast_in_dim path and replace it with broadcast plus expand where applicable; done means the executor uses the explicit operations and better matches the Python-defined fusion representation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100