Skip to content

Added concatenated layout support - #46

Closed
IlyasMoutawwakil wants to merge 2 commits into
Dao-AILab:mainfrom
IlyasMoutawwakil:concatenated-layout
Closed

IlyasMoutawwakil wants to merge 2 commits into
Dao-AILab:mainfrom
IlyasMoutawwakil:concatenated-layout

Conversation

@IlyasMoutawwakil

Copy link
Copy Markdown

We wanna add sonic-moe support in Transformers, but the experts layout we are using in our modeling across MoEs is concatenated [gate;up] and not the interleaved one (see #45) changing the weight layout to interleaved breaks a couple things for us and seems to be an adoption pain point in other projects as well so we are suggesting this PR 🤗

bf16_experts_tflops

@IlyasMoutawwakil
IlyasMoutawwakil marked this pull request as ready for review April 14, 2026 13:43
@vasqu

vasqu commented Apr 14, 2026

Copy link
Copy Markdown

cc @tridao

@tridao

tridao commented Apr 14, 2026

Copy link
Copy Markdown
Member

this fix is pretty interesting, it's simpler than I thought (by just changing the layout of the weight tensor).
@GarlGuo we should check if this creates extra TMA instructions. I guess it doubles the TMA instructions, which may not matter? Maybe the perf impact will be minimal if the N dimension is not contiguous (not the one with stride 1)?

@tridao

tridao commented Apr 14, 2026

Copy link
Copy Markdown
Member

we'll need to think through how this affects the rest of the kernels (esp bwd). But if it's just about changing the tensor layout (and eating the extra TMA instruction) I think we should support this concat layout!

@tridao

tridao commented Apr 14, 2026

Copy link
Copy Markdown
Member

@GarlGuo is about to land a big refactor in the next couple of days so if we're to merge it it'll be after the refactor.

@tridao

tridao commented Apr 16, 2026

Copy link
Copy Markdown
Member

confirming that concat layout doesnt' even generate extra TMA instruction, it just changes TMA 3D to TMA 4D. That's why perf stays the same. This is great!

@GarlGuo

GarlGuo commented Apr 16, 2026

Copy link
Copy Markdown
Member

I will block this PR until #39 is merged (which might happen tmrw). After that I will work on the concatenated layout support.

@tridao

tridao commented Apr 16, 2026

Copy link
Copy Markdown
Member

i have sth in quack to support concat layout, will open PR soon

@tridao

tridao commented Apr 18, 2026

Copy link
Copy Markdown
Member

@IlyasMoutawwakil i adapted this to quack (Dao-AILab/quack#107), adding you as co-author if that's ok.

@GarlGuo

GarlGuo commented Apr 19, 2026

Copy link
Copy Markdown
Member

@IlyasMoutawwakil I have created a new PR as #47. I will also add you as a co-author if that's ok.

@IlyasMoutawwakil

Copy link
Copy Markdown
Author

awesome ! works for me and thank you both 🙏

@GarlGuo

GarlGuo commented Apr 20, 2026

Copy link
Copy Markdown
Member

@IlyasMoutawwakil #47 has been merged and is available on the latest pypi package: https://pypi.org/project/sonic-moe/

@vasqu

vasqu commented Apr 20, 2026

Copy link
Copy Markdown

Awesome!! We will start integrating SonicMoE with huggingface/transformers#45433 🔥

@GarlGuo

GarlGuo commented Apr 20, 2026

Copy link
Copy Markdown
Member

@IlyasMoutawwakil @vasqu I will close this PR and issue #45 for now. Please let me know if more help is needed.

@vasqu

vasqu commented Apr 20, 2026

Copy link
Copy Markdown

Sounds good, can you also ref the PR in the issue for others that might stumble on it?

@GarlGuo

GarlGuo commented Apr 20, 2026

Copy link
Copy Markdown
Member

@vasqu Sounds good. I have left a note.

@GarlGuo GarlGuo closed this Apr 20, 2026
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.

4 participants