write-triton-silu-mul-kernel

Skill: Write a Triton SiLU-Mul (SwiGLU) Kernel

Purpose

Guide the agent through implementing a correct, numerically stable Triton kernel that computes y = silu(a) * b, the elementwise activation step inside SwiGLU MLPs used by LLaMA, Mistral, Qwen, Gemma, and similar modern LLMs. The full MLP is down_proj( silu(gate_proj(x)) * up_proj(x) ); this skill covers the fused activation that sits between the two GEMMs. It also generalizes to GeGLU (gelu(a) * b) and ReGLU (relu(a) * b), which share the same kernel structure with a different activation.

Use this when

• You need a fused elementwise kernel that reads a and b once and writes y once, instead of materializing silu(a) as a separate tensor.
• You are writing an inference path where gate_proj and up_proj are computed separately (or as a single fused matmul producing [gate, up]) and the activation is a distinct kernel call between the matmuls.
• The matmul backend (cuBLAS, CUTLASS without a custom epilogue, or a vendor library) does not allow you to fuse the activation into the matmul epilogue.
• The intermediate tensor is wide enough (e.g., intermediate_size of 14336, 28672, or larger) that the bandwidth cost of materializing silu(a) separately is measurable.
• You want a GeGLU or ReGLU variant — same kernel skeleton, different activation function.

Do not use this when

• torch.nn.functional.silu(a) * b under torch.compile already fuses the chain on your PyTorch build. Validate this with TORCH_COMPILE_DEBUG=1 before writing a custom kernel — modern inductor handles this case well.
• A working CUDA implementation already exists in your serving stack. vLLM ships silu_and_mul in csrc/activation_kernels.cu; SGLang and TensorRT-LLM have equivalents. Re-implementing in Triton is only worth it if you need backend portability or kernel-level fusion with an adjacent op.
• You can fuse the activation into the matmul epilogue (CUTLASS epilogue visitor, Triton matmul with custom epilogue). A standalone elementwise kernel always pays an extra round trip to HBM; the epilogue does not.
• The shape is small enough that kernel launch overhead dominates (e.g., B*T*intermediate_size < 1M elements). At that size, any reasonable implementation is fine.
• You need the backward pass for training. The forward kernel is straightforward, but the backward must save a and b (or silu(a) and b) and recompute silu'(a). Plan the autograd function before writing the forward.