Skip to content

[Triton] Fix triton launch grids under multi gpu scenarios - #1003

Open
vasqu wants to merge 1 commit into
state-spaces:mainfrom
vasqu:fix-device-autocast
Open

[Triton] Fix triton launch grids under multi gpu scenarios#1003
vasqu wants to merge 1 commit into
state-spaces:mainfrom
vasqu:fix-device-autocast

Conversation

@vasqu

@vasqu vasqu commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

It's essentially the same underlying issue as Dao-AILab/flash-attention#523

Triton can have issues to select the correct underlying device

@kingtweak69

Copy link
Copy Markdown

!

@Chessing234

Copy link
Copy Markdown
Contributor

nice catch on the launch grid — can you add a tiny multi-gpu smoke test or at least a comment with the failing shape so this doesn't regress quietly?

@vasqu

vasqu commented Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

cc @tridao whenever you have time

Seems like nemotron h tried to fix it in a wrong way so I'd rather have this upstream (they pinned the complete cuda stream)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants