Study Plan

Triton Row-Wise Softmax

Activation FunctionsGPU Programming

Medium

Implement stable softmax independently across each row of a two-dimensional CUDA tensor using Triton. For input value $x_{r,c}$ in row $r$ and column $c$ , let $C$ be the number of columns and $j$ index columns. The row maximum $m_r$ and normalized output $y_{r,c}$ are:

m_r=\max_{0\le j<C}x_{r,j}.

y_{r,c}=\frac{e^{x_{r,c}-m_r}}{\sum_{j=0}^{C-1}e^{x_{r,j}-m_r}}.

Complete only the empty kernel body; the supplied solve launcher is already implemented. The pointers identify the input and output buffers, row_stride is the input row stride in elements, n_cols is the row width, and BLOCK_SIZE is the compile-time padded width. Return no value; write normalized rows into the contiguous output tensor with the same shape and storage dtype as the input.

Examples

Input: x = [[1, 2, 3], [3, 2, 1]], dtype = float32

Output: [[0.090031, 0.244728, 0.665241], [0.665241, 0.244728, 0.090031]]

Explanation: Reversing a row reverses its probability order.

Input: x = [[-2, 0, 1, 4, -1]], dtype = float16

Output: [[0.0023, 0.016998, 0.046204, 0.928223, 0.006256]]

Input: x = [[8], [-3]], dtype = bfloat16

Output: [[1], [1]]

Hint 1

tl.max and tl.sum reduce values along a block axis.

Hint 2

tl.load supports a masked fallback value for padded columns.

Hint 3

tl.exp evaluates elementwise exponentials.

Requirements

Use Triton with float32 arithmetic for normalization.
Normalize each row independently without modifying the input.
Support padded input rows and write only valid output columns.

Constraints

x and out are separate CUDA tensors with equal shape, device, and float32, float16, or bfloat16 storage dtype.
Input values are finite; input columns and the output tensor are contiguous.
Valid matrices have zero or more rows and 1 through 65,536 columns.
The supplied launcher rejects invalid rank, column stride, or width and handles zero rows.
The supplied launcher converts bfloat16 through temporary float32 buffers.

Try Similar Problems

Masked Triton GELU KernelEasy Fused Tiled Matmul and ReLUHard Tiled Triton Row SumMedium Blockwise Online SoftmaxHard Softmax as a CUDA KernelMedium Online Softmax Kernel in TritonMedium Fused Row-Wise Softmax in TritonMedium Fused Matmul + Bias + ReLU in TritonMedium Implement Swish (SiLU) as a CUDA KernelEasy ReLU as a CUDA KernelEasy Implement Leaky ReLU as a CUDA KernelEasy Implement Sigmoid as a CUDA KernelEasy

Triton Row-Wise Softmax

Activation FunctionsGPU Programming

Medium

m_r=\max_{0\le j<C}x_{r,j}.

y_{r,c}=\frac{e^{x_{r,c}-m_r}}{\sum_{j=0}^{C-1}e^{x_{r,j}-m_r}}.

Examples

Input: x = [[1, 2, 3], [3, 2, 1]], dtype = float32

Output: [[0.090031, 0.244728, 0.665241], [0.665241, 0.244728, 0.090031]]

Explanation: Reversing a row reverses its probability order.

Input: x = [[-2, 0, 1, 4, -1]], dtype = float16

Output: [[0.0023, 0.016998, 0.046204, 0.928223, 0.006256]]

Input: x = [[8], [-3]], dtype = bfloat16

Output: [[1], [1]]

Hint 1

tl.max and tl.sum reduce values along a block axis.

Hint 2

tl.load supports a masked fallback value for padded columns.

Hint 3

tl.exp evaluates elementwise exponentials.

Requirements

Use Triton with float32 arithmetic for normalization.
Normalize each row independently without modifying the input.
Support padded input rows and write only valid output columns.

Constraints

x and out are separate CUDA tensors with equal shape, device, and float32, float16, or bfloat16 storage dtype.
Input values are finite; input columns and the output tensor are contiguous.
Valid matrices have zero or more rows and 1 through 65,536 columns.
The supplied launcher rejects invalid rank, column stride, or width and handles zero rows.
The supplied launcher converts bfloat16 through temporary float32 buffers.

Try Similar Problems

x =

Accepts: array

dtype =

Accepts: string

Examples

Requirements

Constraints

Kernel performance

Triton Row-Wise Softmax

Examples

Requirements

Constraints

Kernel performance