Lightning-AI / Lightning-AI/lightning-thunder

Replace usage of nvFuser's broadcast_in_dim in executor with broadcast+expand ops

Open
#2,765 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement executors nvfuser
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

cc @mruberry

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.

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.