skills.rest
KrxGu
KrxGuCOMMUNITY
Explore all skills by this author →
Website𝕏@krishgupta72GitHub

Write a Triton Softmax Kernel

Implement a numerically stable Triton softmax kernel with fp32 accumulation and multi-block online rescaling.

54|7|Updated Apr 10, 2026
Search Tags:#softmax#triton#numerical-stability#gpu-kernel#kernel-optimization#masked-softmax
One-click install
npx skills add https://github.com/KrxGu/kernel-skills --skill write-a-triton-softmax-kernel
Or copy as Structured Prompt for Agent▼
Please help me install this Agent Skill.
Skill: Write a Triton Softmax Kernel
Source: https://github.com/KrxGu/kernel-skills/tree/main/skills/triton/write-triton-softmax-kernel
Command: npx skills add https://github.com/KrxGu/kernel-skills --skill write-a-triton-softmax-kernel

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

Implements a numerically stable, single-program-per-row Triton softmax kernel that avoids separate reduction passes, handles masked attention, and prevents overflow/precision loss during exponentiation and accumulation.

Core Features & Use Cases

  • Numerical stability: casts inputs to fp32 for reductions, subtracts row max before exponentiation, and uses out-of-bounds loads with -inf to avoid phantom contributions.
  • Masked and fused softmax: applies additive masks before reduction and supports fusion with downstream elementwise operations (e.g., scale or multiply-by-V).
  • Multi-block online softmax: handles rows larger than BLOCK_SIZE with running-max/rescaled-sum updates or multi-pass strategies and documents correctness/performance tradeoffs.
  • Use Case: Replace a failing or unfused softmax in attention kernels with a Triton implementation that preserves correctness for large rows and masked sequences while remaining bandwidth-efficient.

Quick Start

Implement a Triton row-wise softmax kernel that uses fp32 accumulation, applies additive masks before max reduction, handles rows larger than BLOCK_SIZE with online max/sum updates, and validates outputs against torch.nn.functional.softmax.

Frequently Asked Questions about Write a Triton Softmax Kernel

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I implement a numerically stable softmax kernel in Triton?▼

A numerically stable softmax kernel in Triton requires casting inputs to fp32 for accumulation, subtracting the row max before exponentiation, and masking out-of-bounds elements with -inf to prevent phantom contributions.

How does online softmax handle rows larger than BLOCK_SIZE in Triton?▼

Online softmax handles rows larger than BLOCK_SIZE by computing running-max and rescaled-sum updates across multiple passes, ensuring correct accumulation across blocks while maintaining numerical stability.

Why does my Triton masked softmax produce incorrect attention weights?▼

Incorrect masked softmax outputs often occur when additive masks are applied after reduction. Applying the additive mask before the max reduction ensures proper exclusion of padded positions.

Can I fuse elementwise operations with a softmax kernel in Triton?▼

Yes, this single-program-per-row Triton softmax supports fusion with downstream elementwise operations like scaling or multiply-by-V to improve memory bandwidth efficiency.

What is the best way to validate a custom Triton softmax kernel?▼

The best way to validate a custom Triton softmax kernel is to compare its outputs against torch.nn.functional.softmax to ensure numerical correctness.

Related Skills

Explore More Software Engineering Skills→
JKhyroJKhyro

optimizing-attention-flash

Optimize transformer attention to reduce sequence-length bottlenecks and GPU memory pressure.

Community
Advanced
NVIDIANVIDIA
2.8k

kernel-triton-writing

Design, implement, verify and benchmark Triton kernels with fixed contracts and built-in harnesses.

Official
Advanced
mindspore-aimindspore-ai
258

cuda-c-examples-torch

Compile and call CUDA C kernels from PyTorch using load_inline.

Official
Advanced
xchang1121xchang1121
6

triton-cuda-examples-torch

Implement Triton CUDA kernels for PyTorch neural network layers.

Community
Advanced