[RFC]: Kernel Optimization for Diffusion DiT and MoE LLM

Motivation.

Many forward_cuda paths in the diffusion layers/ folder simply delegate to forward_native, so the CUDA path runs under torch.compile-on-native; whether that is the fastest option per op — vs a Triton kernel or an existing upstream CUDA kernel — has not been measured. Several common patterns at call sites (adaLN modulation, GroupNorm + SiLU pairs, QK-RMSNorm + RoPE, AllReduce + RMSNorm) also run as separate kernel launches that have known fused replacements. Both questions apply to the diffusion DiTs and to the MoE LLM decoder shared by the Qwen3-Omni Thinker and Talker. This RFC unifies single-op and fused-op kernel work under one decision framework.

Proposed Change.

Goals

• Single-op kernels (already in the layers/ folder): RMSNorm, LayerNorm, RMSNormVAE, RotaryEmbeddingWAN, AdaLayerNorm.
• Fused-op kernels (exposed as new layer classes in the same folder, called from existing DiT / VAE / MoE LLM call sites): adaLN scale/shift modulation, GroupNorm + SiLU, QK-RMSNorm + RoPE, AllReduce + RMSNorm.
• For each candidate op, benchmark the implementations that already exist:
  1. native PyTorch + torch.compile
  2. Triton — if a kernel is available (e.g., from SGLang)
  3. existing CUDA — if an existing directly-callable kernel is available (e.g., from upstream vllm._custom_ops or flashinfer)
• Two benchmark stages per row:
  1. µbench — synthetic tensors at the activation shapes used by representative models. Cheap; runs on any modern GPU.
  2. E2E — wire the candidate into the actual model and measure end-to-end latency / quality. A "win" is ≥ 1% e2e improvement over native + torch.compile. Microbench gains that don't translate to ≥ 1% e2e are not worth wiring.
• Decide per op:
  • No candidate clears the 1% e2e bar → no PR. File a sub-issue with bench numbers; the row closes as ❌.
  • Triton wins ≥ 1% e2e → open a PR that ports the Triton kernel into the layers/ folder.
  • Existing CUDA wins ≥ 1% e2e → open a PR that calls the upstream op directly from forward_cuda (or, for cross-layer fused ops, from the relevant attention/decoder call site). vllm-omni does not use torch.compile custom fusion passes, so candidates that exist only as compile-pass rewrites do not qualify — only directly-callable Python/C++ entry points count.
• Writing new CUDA kernels is out of scope. We don't benchmark hypothetical new CUDA; if all three above lose, the row stays open until someone upstreams a kernel.
• Keep forward_native numerically authoritative; non-CUDA paths (NPU/XPU/MUSA) untouched. Equivalence vs forward_native within bf16 tolerance (atol=rtol=1e-2) gates µbench/E2E acceptance.
• Register every new op via torch.library.custom_op + torch.library.register_fake so torch.compile doesn't graph-break.

Architecture

For each candidate op (single-op and fused-op) in the diffusion layers/ folder:

   Benchmark on representative shapes
       │
       ├─►  native + torch.compile
       ├─►  Triton            (if SGLang has one)
       └─►  existing CUDA     (if vllm._custom_ops / flashinfer has a directly-callable one)
                            │
                            ▼
   ┌──────────────────┬──────────────────────┬──────────────────────────────┐
   │  Native wins     │  Triton wins         │  Existing CUDA wins          │
   ├──────────────────┼──────────────────────┼──────────────────────────────┤
   │  No PR.          │  PR: port Triton     │  PR: call the upstream op    │
   │  File sub-issue  │  kernel into         │  directly from forward_cuda  │
   │  with bench      │  layers/.            │  / call site.                │
   │  results.        │                      │                              │
   └──────────────────┴──────────────────────┴──────────────────────────────┘

   Out of scope: new CUDA kernels (not benchmarked) and torch.compile custom fusion passes (vllm-omni does not use them).

Tracking Matrix

Workflow for a contributor claiming a row:

1. Comment on the RFC issue with the GPU/accelerator you have access to (microbench can run on any modern GPU; e2e needs the target model's full VRAM).
2. µbench: time each listed candidate (native + torch.compile, Triton ref if listed, existing CUDA ref if listed) on the activation shapes that representative models use. Post results.
3. E2E: wire the leading candidate into the actual model and measure end-to-end latency / quality. The win bar is ≥ 1 % e2e improvement over native + torch.compile.
4. PR cell: if no candidate clears 1 % e2e → file a sub-issue with bench numbers (row closes as ❌). Otherwise → open a PR that ports the Triton kernel or wires the existing CUDA path (row closes as ✅ with PR link).

Each row carries three independent trackers: µbench, E2E, PR. Symbols per cell: 🙋 unclaimed · ⏳ in progress (link if available) · ✅ done (link to results / PR) · ❌ negative result · — N/A (no candidate to evaluate).

Type	Op	Triton ref	Existing CUDA ref	µbench	E2E	PR
Single	RMSNorm	sglang rms_norm_fn, triton_one_pass_rms_norm	vllm rms_norm, fused_add_rms_norm; flashinfer rmsnorm, fused_add_rmsnorm	❌@GuoningHuang	❌@GuoningHuang	❌@GuoningHuang
Single	LayerNorm	sglang _layer_norm_fwd_1pass_kernel [1]	—	⏳@Sendoh-code	⏳@Sendoh-code	⏳@Sendoh-code
Single	RMSNormVAE	sglang rms_norm_fn [2]	—	⏳@Sendoh-code	⏳@Sendoh-code	⏳@Sendoh-code
Single	RotaryEmbeddingWAN	sglang apply_rotary_embedding [3]	vllm-flash-attn apply_rotary_emb; flashinfer apply_rope_with_cos_sin_cache [4]	⏳@divyanshsinghvi	⏳@divyanshsinghvi	⏳@divyanshsinghvi
Single	AdaLayerNorm	see LayerNorm + adaLN scale/shift rows	—	⏳@zxdreamer	⏳@zxdreamer	⏳@zxdreamer
Fused	adaLN scale/shift modulation	sglang scale_shift.py (sgl-project/sglang#18257)	—	❌@sphinxkkkbc	❌@sphinxkkkbc	❌@sphinxkkkbc
Fused	GroupNorm + SiLU	sglang group_norm_silu.py (sgl-project/sglang#22814)	—	⏳@GuoningHuang	⏳@GuoningHuang	⏳@GuoningHuang
Fused	QK-RMSNorm + RoPE	—	vllm fused_qk_norm_rope — direct call [5]	⏳@divyanshsinghvi	⏳@divyanshsinghvi	⏳@divyanshsinghvi
Fused	AllReduce + RMSNorm	—	flashinfer trtllm_allreduce_fusion, wrapped by vllm flashinfer_all_reduce.py [6]	⏳@fywc	⏳@fywc	⏳@fywc

Notes:

1. The kernel is @triton.jit-private; norm.py only exposes rms_norm_fn publicly. A thin Python wrapper around _layer_norm_fwd_1pass_kernel needs to be added before benchmarking LayerNorm.
2. rms_norm_fn normalizes along the last dim. RMSNormVAE is channel-first; the contributor needs a transpose adapter.
3. sglang's apply_rotary_embedding flattens x to (B*S, H, D) internally and grids over B*S. Per-batch cos/sin works after the caller reshapes cos to (B*S, D/2) (one .view()).
4. vllm-flash-attn and flashinfer both require 2D cos/sin (seqlen, D/2). Wan ships 3D batched cos/sin (per-batch positions differ). Base RotaryEmbedding.forward_cuda's cos = cos[0] workaround silently corrupts Wan's per-batch RoPE — that's why RotaryEmbeddingWan.forward_cuda deliberately routes to forward_native.
5. fused_qk_norm_rope operates on a packed qkv tensor — wiring requires rewriting the affected attention modules to feed qkv and accept the in-place output. Upstream's torch.compile pass (qk_norm_rope_fusion) does this auto-rewrite; we don't use compile passes, so the rewrite must be manual.
6. The directly-callable FlashInfer fused all-reduce + RMSNorm exists; upstream's torch.compile pass (allreduce_rms_fusion) is one consumer of it, but the underlying op is invokable without the pass. Wiring spans the all-reduce boundary in MoE/MLP decoder blocks — call-site rewrite, not a forward_cuda swap.

Feedback Period.

No response

CC List.

@ZJY0516 @hsliuustc0106 @princepride

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.