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

Write CUDA Softmax Kernel

Implement numerically stable CUDA softmax kernels with fp32 accumulation and masking.

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

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

Implements a correct, numerically stable, and high-performance CUDA softmax along the last dimension for 2D/3D tensors, preventing overflow/underflow, handling masked rows, and supporting fp16/bf16 inputs with fp32 accumulation to avoid common softmax correctness bugs in attention and kernel-fusion scenarios.

Core Features & Use Cases

  • Numerical stability: always subtracts the row maximum before exponentiation and accumulates in fp32 for low-precision inputs.
  • Masked softmax: supports additive and boolean masks with explicit handling for fully-masked rows and documented fallback behavior.
  • Decomposition & performance: guidance for warp/block assignment, warp-level reductions, optional online (single-pass) or two-pass algorithms, and host dispatch parameters for common attention shapes.
  • Precision & backward pass: fp16/bf16 input handling with fp32 accumulation and a reusable reduction pattern for the backward gradient computation.
  • Use Case: fuse softmax into attention kernels for low memory footprint flash-attention patterns or implement a standalone masked softmax when library routines are insufficient.

Quick Start

Ask the agent to generate a CUDA kernel that computes numerically stable row-wise softmax with optional additive or boolean masking, using fp32 accumulation for fp16/bf16 inputs and warp+block reductions for performance.

Frequently Asked Questions about Write CUDA Softmax Kernel

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

FAQPage Schema
How do I implement a numerically stable CUDA softmax kernel for fp16 inputs?▼

To implement a numerically stable CUDA softmax for fp16, subtract the row maximum before exponentiation and use fp32 accumulation. This prevents overflow and underflow common in low-precision attention computations.

What is the best way to handle fully masked rows in a custom CUDA softmax kernel?▼

Handling fully masked rows in a CUDA softmax kernel requires explicit logic to avoid NaN outputs. The implementation provides documented fallback behavior for additive and boolean masks when an entire row is masked.

Can I fuse masked softmax directly into my CUDA attention kernel?▼

Yes, you can fuse masked softmax into CUDA attention kernels. The implementation supports standalone 2D/3D masked softmax and fusion patterns for low memory footprint flash-attention scenarios.

How does warp reduction improve performance in CUDA softmax computation?▼

Warp reduction improves CUDA softmax performance by optimizing row-wise summation. The kernel provides warp and block reduction helpers, supporting both online single-pass and two-pass algorithms for common attention shapes.

Why does my CUDA softmax kernel produce incorrect results with bf16 inputs?▼

CUDA softmax kernels produce incorrect bf16 results due to insufficient accumulation precision. Accumulating in fp32 for bf16 inputs ensures correct exponentiation and summation, avoiding common softmax correctness bugs.

Do I need to configure shared memory and block parameters for a CUDA softmax kernel?▼

Yes, configuring shared memory and block/grid parameters is needed. The implementation includes host dispatch parameters, allowing you to tune block assignment and shared memory for specific attention tensor layouts.

Related Skills

Explore More Software Engineering Skills→
OpenMOSSOpenMOSS
410

add-jit-kernel

Add a lightweight JIT CUDA kernel with host wrapper, tests, and benchmarks to SGLang's jit_kernel module.

Official
Advanced
KrxGuKrxGu
54

Write a Triton Softmax Kernel

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

Community
Advanced
zli5460zli5460
2

optimizing-attention-flash

Integrate Flash Attention kernels into PyTorch transformer models.

Community
Intermediate
mindspore-aimindspore-ai
258

triton-cuda-examples-torch

Integrate Triton CUDA kernels with PyTorch models for custom operators.

Official
Intermediate