Skip to content

feat(magi2): add fused SwiGLU7 activation - #7252

Merged
hsliuustc0106 merged 1 commit into
vllm-project:mainfrom
yeahdongcn:xd/magi2-swiglu7
Sep 26, 2026
Merged

hsliuustc0106 merged 1 commit into
vllm-project:mainfrom
yeahdongcn:xd/magi2-swiglu7

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Sep 8, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Add MAGI-2's packed SwiGLU7 activation as a diffusion CustomOp and use the
fused Triton implementation by default for eager GPU execution. CPU execution
and torch.compile keep the native PyTorch expression.

The implementation preserves the released formula: FP32 arithmetic, upper-only
gate clamp, symmetric linear clamp, linear + 1, and one final output cast.
It uses no MUSA-only API and does not require an environment variable.

Validation

python -m pytest -q -o addopts='' \
  tests/diffusion/layers/test_swiglu7.py -m 'core_model and cpu'

python -m pytest -q -o addopts='' \
  tests/diffusion/layers/test_swiglu7.py -m cuda

python -c 'import torchada, pytest; raise SystemExit(pytest.main([
  "-o", "addopts=", "-q", "tests/diffusion/layers/test_swiglu7.py",
  "-m", "musa"]))'

Results on commit 2a2eceb93:

  • CPU: 19 passed
  • NVIDIA H20 / CUDA / Triton 3.7.1: 14 passed
  • MTT S5000 host 10.20.36.18, GPU 7 (48 SM) / MUSA / Triton 3.2.0: 14 passed

Microbenchmark:

python -m benchmarks.kernels.benchmark_magi2_swiglu7 \
  --tokens 256 --num-tests 30
Device Native p50 Fused p50 Speedup
MTT S5000 host 10.20.36.18 GPU 7 (48 SM), Triton 3.2.0 167.720 us 93.580 us 1.792x
NVIDIA H20, Triton 3.7.1 82.352 us 46.800 us 1.760x

Both runs passed numerical parity. The benchmark uses alternating native/fused
calls with a 512 MiB cache flush outside timed regions; it is an eager operator
diagnostic rather than a full-video throughput claim.

Scope

No MUSA-specific namespace patch, MATE dependency, mHC, MoE, EP, sampler, or
platform-wide change is included. Independent of #7206, #7249, and #7261;
#7156 remains the integration reference.

@vllm-omni-review-bot

Copy link
Copy Markdown

This PR appears to belong to: docs/design/module/benchmarking.md, docs/design/module/diffusion/diffusion_model_integration.md, docs/design/module/diffusion/index.md.

Module owners: @wtomin @RuixiangMa @david6666666

Routing: @wtomin via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @RuixiangMa via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description

@yeahdongcn, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Sep 14, 2026
Comment thread vllm_omni/diffusion/models/magi2/layers.py Outdated
Comment thread tests/diffusion/layers/test_swiglu7.py Outdated
@hsliuustc0106 hsliuustc0106 added the enhancement New feature or request label Sep 14, 2026

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The kernel, fallback matrix, and tests all verify clean on my side (19 CPU + 14 CUDA tests pass; parity exact on CPU and within ULP-scale tolerances on GPU). The only blocker is the unregistered MAGI2_USE_FUSED_SWIGLU7 env var — the environment-inventory drift test fails at this head, so the cpu lane would go red on merge.

Independently reproduced the perf claim on L20X (triton 3.7.1): I=1280 gives 1.94× (256 tok) to 3.31× (4096 tok), and the dense-MLP width I=12288 reaches 8.08× at 4096 tokens with parity passing — the reported 1.7× understates the gain at production widths. Approve-able once the registry entry lands.

@yeahdongcn
yeahdongcn force-pushed the xd/magi2-swiglu7 branch 2 times, most recently from 477b240 to 2a2eceb Compare September 16, 2026 04:48
@yeahdongcn yeahdongcn changed the title feat(magi2): add opt-in SwiGLU7 CustomOp feat(magi2): add fused SwiGLU7 activation Sep 16, 2026
@vllm-omni-review-bot

Copy link
Copy Markdown

Omni ReviewBot: no human activity for 7 days

@yeahdongcn this pull request has had no human commit, comment or review since 2026-09-16. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

linear = tl.load(x_ptr + 2 * offsets + 1, mask=mask, other=0).to(tl.float32)
# Comparisons preserve NaNs; gate is only upper-clamped, unlike linear.
gate = tl.where(gate > 7.0, 7.0, gate)
linear = tl.where(linear < -7.0, -7.0, linear)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not use a general swiglu with 7 as an input

@hsliuustc0106 hsliuustc0106 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Env-var blocker resolved by removing the flag entirely, and the tests now run in the core_model tier. Approving; the general-SwiGLU-shape question above stands as a non-blocking follow-up.

@hsliuustc0106 hsliuustc0106 added ready label to trigger buildkite CI cuda-test Used to trigger vllm-omni cuda CI separately. labels Sep 25, 2026
@vllm-omni-review-bot

vllm-omni-review-bot commented Sep 25, 2026 •

Copy link
Copy Markdown

Omni ReviewBot: superseded

The CI failure noted on 2a2eceb93cfd refers to an earlier head; the pull request now points at 3aadf82b2fc7.

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
@hsliuustc0106 hsliuustc0106 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Sep 25, 2026
@hsliuustc0106
hsliuustc0106 merged commit e75278d into vllm-project:main Sep 26, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cuda-test Used to trigger vllm-omni cuda CI separately. diffusion codes related to diffusion models enhancement New feature or request high priority high priority issue, needs to be done asap Kernel optimization Codes related to optimize kernel execution to improve hardware utilization ready label to trigger buildkite CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants