[LeoCristallo]

What does this PR do?

Fuse the staged silu(gate) * up activation in the CUDA merged-FC1 Triton MoE forward into one pointwise kernel. It reads the merged FC1 views directly and writes only the final activation, avoiding the intermediate SiLU tensor and one kernel launch.

The SiLU result is rounded to the activation dtype before multiplication. Routing weights remain applied after FC2; saved tensors and backward formulas are unchanged. ROCm and MLU retain the ATen activation path.

Checklist Before Starting

• Related issues/PRs: None.
• PR title follows [{modules}] {type}: {description}.

Test

Baseline/candidate checks use 38898e9 without/with this patch. The fused-model sync selection also passes at the final d02a95e base.

Extended the existing CI-selected tests/ops/moe_experts/test_moe_experts.py with exact staged-rounding, input-preservation, empty/strided/clamped-input, and merged-consumer forward/backward comparisons.

• MoE experts, MoE-LoRA experts and ops-base suites: baseline 182 passed / 27 availability skips; candidate 198 passed / 27 availability skips.
• Three fused-model implicit-sync tests pass on both baseline and candidate, including the final d02a95e base.
• The existing DeepSeek-V4 Inductor expert VJP test passes for all three seeds on both baseline and candidate.
• Case-only compatibility probes pass on both sides: non-reentrant checkpoint replay, torch.compile(..., backend="eager", fullgraph=True) forward/backward, and CUDA graph forward/backward with replay.
• Separate two-update deterministic Qwen3 MoE training comparisons match recorded loss and gradient norm exactly for both datasets.
• make quality and configured pre-commit checks on the proposed files pass.

Training comparison uses the Qwen3 MoE configuration (48 layers, hidden size 2048, expert intermediate size 768, 128 experts, top-k 8), random initialization and BF16. Each dataset uses a nominal sequence budget of 81,920. Baseline, candidate and baseline repeat each run ten updates; the first six are warmup and the last four are measured. Input files, initialization and settings are identical, and effective input lengths match at every compared update. Training measurements use official 38898e9 with and without this patch.

All ratios below are relative to the first baseline within the same dataset. Peak allocated memory is the candidate peak over the four measured updates, normalized to the corresponding baseline peak.

Dataset	Baseline throughput	Candidate throughput	Candidate peak allocated
Synthetic tokens	1.0000×	1.0039×	1.0000×
FineWeb-Edu documents	1.0000×	1.0098×	1.0001×

Candidate throughput improves by about 0.39% on synthetic inputs and 0.98% on documents versus the reported baseline. Every matched steady update is faster, with zero allocation retries in all measured runs. Peak allocated memory is effectively unchanged.

API and Usage Example

No public API, registry selection or configuration changes. The existing fused_triton merged implementation uses the internal helper on its supported path.

Design & Code Changes

• Add a forward-only, stride-aware fused SiLU/product helper with explicit intermediate dtype rounding and no input mutation.
• Invoke it from the existing merged custom-autograd forward.
• Preserve GEMMs, clamp handling, routing-weight placement, saved tensors and backward.
• Document the activation invariant and extend existing ops tests.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied pre-commit checks to the proposed files.
• Added/updated documentation.
• Extended an existing CI-selected test under the ops suite manifests.

Summary by CodeRabbit

• New Features

  • CUDA merged MoE paths now use a fused SwiGLU activation while preserving existing output rounding behavior. Float16 and bfloat16 inputs retain their expected output shape and layout.
  • ROCm and MLU continue to use the existing activation behavior.

• Tests

  • Added coverage for output correctness, input preservation, and forward and gradient equivalence across merged paths.
  • Added checks for compiled SwiGLU initialization across supported entry points.

• Documentation

  • Documented activation behavior, supported input layouts, and mutation guarantees.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: ByteDance-Seed/VeOmni
• Branch: main
• File: .coderabbit.yaml
• Commit: 5bbd362
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Advanced
• Run ID: 2cdf78c4-e7fc-4e10-9fc5-f5205d38bdb9

📥 Commits

Reviewing files that changed from the base of the PR and between 3f27a29 and 1c571ee.

📒 Files selected for processing: 1

• tests/distributed/test_torch_compile.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 3 remain after this review.

📝 Walkthrough

The standard and EP merged-FC1 paths now use a shared SwiGLU helper. NVIDIA CUDA uses a Triton kernel; other devices use the staged ATen expression. Tests cover activation results, gradients, compilation, and lazy imports.

Changes: Merged-FC1 SwiGLU

Layer / File(s)	Summary
CUDA Triton SwiGLU helper veomni/ops/kernels/moe_experts/shared/swiglu.py, veomni/ops/kernels/moe_experts/shared/_triton_swiglu.py	Adds merged_swiglu_forward. The helper validates the inputs, uses ATen on non-NVIDIA devices, and lazily loads the Triton kernel on NVIDIA CUDA. The kernel rounds SiLU to the input dtype before multiplication.
Merged-FC1 integration and validation veomni/ops/kernels/moe_experts/standard/triton.py, veomni/ops/kernels/moe_experts/shared/ep.py, tests/ops/moe_experts/test_moe_experts.py, tests/distributed/test_torch_compile.py, tests/ops/base/test_op_import.py, docs/design/fused_moe_ops.md	The standard and EP merged-FC1 paths call the helper. Tests check output properties, rounding, forward results, gradients, compiled entry points, and lazy imports. The documentation describes the activation path and unchanged backward formulas.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant MergedFc1Path
  participant SwiGLUHelper
  participant TritonKernel
  MergedFc1Path->>SwiGLUHelper: Pass gate and up tensors
  SwiGLUHelper->>TritonKernel: Launch for nonempty NVIDIA CUDA inputs
  TritonKernel->>SwiGLUHelper: Return merged activation
  SwiGLUHelper->>MergedFc1Path: Return activation

Loading

Possibly related PRs

• ByteDance-Seed/VeOmni#1315: Adds swiglu_oai activation support in the same shared module and merged-FC1 call paths.

Suggested reviewers: zx55

Merge Risk: ⚪ Minimal · up to 1c571

No merge-blocking issue was identified in the changed compiled SwiGLU test.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 52.63% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 19 functions across 8 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: fusing the merged MoE forward activation. It is specific, relevant, and suitable for a pull request history.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Coach257]

Updated in af0731669 (rebased onto latest main):

1. Consolidate merged_swiglu_forward into veomni/ops/kernels/moe_experts/shared/swiglu.py instead of introducing a separate single-caller shared/triton_swiglu.py module, and encapsulate the NVIDIA CUDA vs. non-NVIDIA (ROCm / MLU) fallback inside merged_swiglu_forward (with lazy Triton/libdevice JIT initialization so importing shared/swiglu.py stays Triton-free).
2. Share the fused forward activation across both non-EP and EP merged-FC1 paths: EPMergedFc1GroupGemm.forward in veomni/ops/kernels/moe_experts/shared/ep.py now also routes through swiglu.merged_swiglu_forward(fc1_1_output, fc1_2_output) alongside MergedFc1TritonFusedMoeExpertFunction.forward in standard/triton.py, and the regression test in tests/ops/moe_experts/test_moe_experts.py verifies both non-EP and EP paths.

[FoolPlayer]

[P2] On the first compiled training call, this initializer writes module globals and constructs a JITFunction inside the traced custom autograd forward, so fullgraph=True fails unless an eager call has already warmed the cache. Initialize the kernel outside tracing while preserving optional imports, and add a fresh-process compiled-first regression.

[LeoCristallo]

Fixed the compiled-first initialization while keeping the shared helper and fallback for both EP and non-EP consumers. Only the Triton JIT definition is now behind a lazy private-module import. Added fresh-process Inductor regressions without eager warmup; these and the existing MoE and Inductor VJP tests pass.

[FoolPlayer]

[P2] Verified by static inspection: the lazy JIT factory and module-global writes are removed, and _triton_swiglu.py defines the kernel with @​triton.jit during import. This addresses the reported cold-start tracing failure; the new subprocess regressions exercise compiled-first calls without eager warmup.