Expert Kernels

Select the expert compute backend via moe_implementation. This controls which kernel handles the batch grouped GEMM over all active experts.

---

Backend Comparison

Backend	Kernel	EP compatible	torch.compile	LoRA	Notes
triton	Triton group_gemm_same_nk	Yes	Partial†	Yes	Triton grouped GEMM; works with recompute paths
quack	CuTe DSL kernels	Yes	No	Yes	SM90/SM100 kernels; optional package availability is checked at runtime
native	torch._grouped_mm	Yes	Yes (compute path)	Yes	PyTorch grouped GEMM, without the Quack extension
eager	Per-expert PyTorch matmuls	Yes	Yes	Yes	Reference/debugging path; EP uses the registered eager EP compute function

model:
  moe_implementation: triton   # example; benchmark against native/quack for the target shape

†Triton: The group_gemm_same_nk kernel is compatible with torch.compile. However, EP dispatch/combine (AllToAll, DeepEP) is excluded from compile graphs because it involves dynamic shapes (variable token counts per rank per step).

---

Quack Backend

Quack uses CuTe DSL-based GEMM kernels written by Tri Dao (the author of Flash Attention). The kernels are built on PyTorch’s CuTe DSL and compiled for SM90+ (H100) and SM100+ (B200). They are not available on older GPUs.

What makes Quack different

Unlike the Triton backend which compiles kernels via Python-defined Triton JIT, Quack kernels are written directly in CuTe DSL — the same low-level CUDA tensor abstraction used for Flash Attention. This gives:

• Persistent kernels with pingpong double-buffering (H100) — the kernel occupies SMs continuously across multiple expert GEMMs, overlapping memory loads of the next tile with computation of the current tile
• cu_seqlens format — uses Flash Attention’s cumulative sequence length format to group tokens by expert, avoiding unnecessary tensor copies
• SwiGLU fusion — gate_proj and up_proj GEMMs plus the SiLU activation can be fused into a single kernel dispatch via gemm_gated / gemm_act
• SM100 (B200) support — different tile config: tile_m=256, tile_n=256, cluster_m=2, cluster_n=1, pingpong=False

Default tile config per architecture:

Architecture	tile_m	tile_n	cluster_m	cluster_n	pingpong
SM90 (H100)	128	192	2	1	True
SM100 (B200)	256	256	2	1	False

Autotuning

By default, autotuning is disabled (XORL_QUACK_TUNED=0). MoE routing produces variable token counts per expert per step — each unique total_M would trigger a full ~60-config benchmark, making training orders of magnitude slower. The fixed default config is chosen to be performant across typical token budgets.

Enable autotuning only when shapes are stable (e.g. evaluation with fixed batch size) or when the tuning cache is already warm:

XORL_QUACK_TUNED=1 torchrun ... -m xorl.cli.train config.yaml

Registered training variants

The current registries expose these Quack variants. This table does not imply support for every dtype, architecture, quantization, or topology combination:

Variant	Class / Function	Use case
Standard	quack_expert_forward	Default forward, no EP
EP	QuackEPGroupGemm.apply	EP dispatch + compute
LoRA	quack_moe_lora_forward	LoRA adapters, no EP
EP + LoRA	QuackEPGroupGemmWithLoRA.apply	EP + LoRA
moe_act	quack_expert_forward_moe_act	Gradient checkpointing, no EP
EP + moe_act	QuackEPGroupGemmMoeAct.apply	EP + gradient checkpointing

Pairing with DeepEP

Quack handles expert FFN compute while DeepEP handles dispatch/combine, so they can be evaluated together on supported H100-class deployments. Their SM allocations compete, however; compare this pairing with Quack + AllToAll and with other expert backends on the same model and token shape.

model:
  moe_implementation: quack
  ep_dispatch: deepep
  deepep_num_sms: 48    # more SMs for DeepEP; fewer left for Quack — tune based on profiling

Limitations

• Requires SM90+ (H100) or SM100+ (B200) — not available on A100 or older
• Not compatible with torch.compile — custom CUDA/CuTe extensions cannot be traced
• Capability checks still apply — LoRA and hybrid shared-LoRA paths are registered, but dtype, quantization, architecture, and dispatch combinations can impose narrower constraints
• Autotuning disabled by default — variable token counts make per-step autotuning prohibitively slow

---

Triton Grouped GEMM Details

The core of the triton backend is src/xorl/ops/group_gemm/kernel/group_gemm.py. The primary kernel group_gemm_same_nk_kernel handles the case where all experts share the same K (input) and N (output) dimensions.

Inputs and outputs:

• A: [total_tokens, K] — all tokens routed to this rank, concatenated
• B: [G, K, N] — expert weight tensor in GKN layout
• cumsum_M: [G] — cumulative expert token counts (prefix sum)
• C: [total_tokens, N] — output

Tile strategy: Each thread block processes a BLOCK_M × BLOCK_N output tile for one expert. An inner loop steps through BLOCK_K slices of A and B, accumulating the partial sum in registers.

Autotuned configs: BLOCK_M ∈ {64, 128}, BLOCK_N ∈ {128, 256}, BLOCK_K ∈ {32, 64} — autotuned per (N, K) shape at first use and cached.

Optional activation fusion: When moe_act gradient checkpointing is active, the kernel stores the post-gate×up intermediate activations to a separate buffer during the forward pass (controlled by the STORE_ACTIVATIONS flag). The custom backward function recomputes only the down projection from these stored intermediates.

---

torch.compile with MoE

torch.compile is applied per decoder block when enable_compile: true.

What is compiled:

• Attention layers (Q/K/V projections, flash attention wrapper, output projection)
• Router gate linear and softmax
• Expert FFN compute (Triton group_gemm_same_nk is torch.compile-friendly)
• Dense shared-expert FFN layers (if present)

What is excluded from compile graphs:

• EP dispatch/combine collectives (alltoall_pre_dispatch, alltoall_post_combine, DeepEP dispatch/combine) — dynamic shapes from variable per-rank token counts break static graph tracing
• Routing replay record/replay logic — mutable external state, decorated with @torch.compiler.disable

Example compile config:

model:
  moe_implementation: triton   # or native; quack is incompatible with compile
train:
  enable_compile: true

The native backend (torch._grouped_mm) compiles cleanly end-to-end. The quack backend uses custom CUDA extensions and must not be used with enable_compile: true.

---

MoE Gradient Checkpointing (gradient_checkpointing_method)

gradient_checkpointing_method controls what is recomputed in the backward pass. For MoE models with EP, this has a large impact on throughput because recompute_full_layer recomputes AllToAll dispatch + combine on every backward layer.

Methods

Method	Backward recomputes	EP comm recomputed	Retained activation state
recompute_full_layer (default)	Entire decoder layer	Yes — dispatch/combine are replayed	lowest
recompute_before_dispatch	Attention + router only; retains dispatch + expert + combine	No	higher
no_recompute	Nothing	No	highest

# No-recompute candidate (use only after a shape-matched fit check):
train:
  enable_gradient_checkpointing: true
  gradient_checkpointing_method: no_recompute

# Balanced memory/speed — recompute attn+router, keep dispatch+expert+combine:
train:
  enable_gradient_checkpointing: true
  gradient_checkpointing_method: recompute_before_dispatch

# Lowest retained-activation state (start here for memory-constrained runs):
train:
  enable_gradient_checkpointing: true
  gradient_checkpointing_method: recompute_full_layer

How recompute_before_dispatch works (Megatron-style)

In the expert MLP forward: gate_output = x @ gate_proj, up_output = x @ up_proj, postact = silu(gate_output) * up_output, output = postact @ down_proj.

With recompute_before_dispatch, the attention and router outputs are recomputed in backward, but MoE dispatch and combine are not recomputed. This avoids re-executing AllToAll communication, which is the dominant cost at multi-node.

The retained tensors and avoided work differ by model and backend. Measure peak memory and end-to-end step time at the intended sequence length before changing the default.

Why topology matters

recompute_full_layer replays dispatch and combine during backward; recompute_before_dispatch retains their outputs instead. Avoiding that replay can matter more when EP communication crosses nodes, but the result depends on fabric, token balance, overlap, retained-memory pressure, and expert-kernel cost. Treat the choice as a measured memory/throughput tradeoff rather than a fixed speedup.

---

Source

File	Description
src/xorl/models/layers/moe/experts.py	MoEExperts — GKN weight tensors, backend dispatch, EP forward, moe_act
src/xorl/ops/group_gemm/kernel/group_gemm.py	group_gemm_same_nk_kernel — Triton autotuned grouped GEMM with optional STORE_ACTIVATIONS
src/xorl/ops/__init__.py	triton_moe_forward, triton_moe_lora_forward, backend registry
src/xorl/ops/group_gemm/kernel/quack.py	quack_group_gemm_same_nk — CuTe GEMM wrapper; XORL_QUACK_TUNED env var
src/xorl/ops/quack/gemm_interface.py	Quack GEMM interface — SM90/SM100 configs, autotuner, gated activation kernels
src/xorl/models/layers/moe/backend/quack.py	quack_expert_forward — non-EP quack MoE forward
src/xorl/ops/moe/quack.py	QuackEPGroupGemm, QuackEPGroupGemmMoeAct — EP + moe_act variants