Study Plan

Blockwise Online Softmax

Activation FunctionsGPU Programming

Hard

Implement stable row-wise softmax in CUDA using online summaries. A summary contains the maximum $m$ and the shifted exponential sum $\ell$ of its values:

\ell=\sum_j e^{x_j-m}.

Here, $j$ indexes values and $x_j$ is one value. For another block with summary $m_b,\ell_b$ , the combined summary is:

m_{\mathrm{new}}=\max(m,m_b).

\ell_{\mathrm{new}}=e^{m-m_{\mathrm{new}}}\ell+e^{m_b-m_{\mathrm{new}}}\ell_b.

Here, subscript $b$ denotes the other block and subscript new denotes their union. Once $m$ and $\ell$ summarize the complete row, its normalized output at column $j$ is:

y_j=\frac{e^{x_j-m}}{\ell}.

Complete the empty CUDA kernel; the supplied solve launcher is already implemented. Input and output are preallocated matrices with the supplied dimensions and element strides. The kernel returns no value and writes the probabilities into output, preserving the input shape and storage dtype.

Examples

Input: input = [[1, 2, 3], [-1, 0, 1]], dtype = float32, block_size = 64

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

Explanation: The rows differ by a constant shift, so they have the same probabilities.

Input: input = [[1000, 1001, 1002, 999]], dtype = float32, block_size = 128

Output: [[0.087144, 0.236883, 0.643914, 0.032059]]

Input: input = [[-3, 2, 0.5, 4, -1]], dtype = float16, block_size = 32

Output: [[0.000777, 0.115356, 0.025742, 0.852539, 0.005745]]

Hint 1

expf evaluates the exponential in float32.

Hint 2

__syncthreads coordinates access to shared-memory summaries.

Hint 3

__half2float and __bfloat162float convert low-precision input values.

Requirements

Use CUDA with float32 arithmetic for the online summaries and normalization.
Preserve float32, float16, or bfloat16 storage in the output.
Support the supplied input and output strides without modifying the input.

Constraints

Input values are finite and input and output pointers reference separate CUDA buffers.
rows and cols are nonnegative integers; an empty dimension requires no writes.
Row and column strides are measured in elements.
dtype_code is 0 for float32, 1 for float16, and 2 for bfloat16.
block_size is a power of two from 32 through 1024; the supplied launcher handles invalid launch metadata.

Try Similar Problems

Fused Tiled Matmul and ReLUHard Triton Row-Wise SoftmaxMedium Masked Triton GELU KernelEasy 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 Implement Exact GELU as a CUDA KernelEasy

Blockwise Online Softmax

Activation FunctionsGPU Programming

Hard

Implement stable row-wise softmax in CUDA using online summaries. A summary contains the maximum $m$ and the shifted exponential sum $\ell$ of its values:

\ell=\sum_j e^{x_j-m}.

Here, $j$ indexes values and $x_j$ is one value. For another block with summary $m_b,\ell_b$ , the combined summary is:

m_{\mathrm{new}}=\max(m,m_b).

\ell_{\mathrm{new}}=e^{m-m_{\mathrm{new}}}\ell+e^{m_b-m_{\mathrm{new}}}\ell_b.

Here, subscript $b$ denotes the other block and subscript new denotes their union. Once $m$ and $\ell$ summarize the complete row, its normalized output at column $j$ is:

y_j=\frac{e^{x_j-m}}{\ell}.

Examples

Input: input = [[1, 2, 3], [-1, 0, 1]], dtype = float32, block_size = 64

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

Explanation: The rows differ by a constant shift, so they have the same probabilities.

Input: input = [[1000, 1001, 1002, 999]], dtype = float32, block_size = 128

Output: [[0.087144, 0.236883, 0.643914, 0.032059]]

Input: input = [[-3, 2, 0.5, 4, -1]], dtype = float16, block_size = 32

Output: [[0.000777, 0.115356, 0.025742, 0.852539, 0.005745]]

Hint 1

expf evaluates the exponential in float32.

Hint 2

__syncthreads coordinates access to shared-memory summaries.

Hint 3

__half2float and __bfloat162float convert low-precision input values.

Requirements

Use CUDA with float32 arithmetic for the online summaries and normalization.
Preserve float32, float16, or bfloat16 storage in the output.
Support the supplied input and output strides without modifying the input.

Constraints

Input values are finite and input and output pointers reference separate CUDA buffers.
rows and cols are nonnegative integers; an empty dimension requires no writes.
Row and column strides are measured in elements.
dtype_code is 0 for float32, 1 for float16, and 2 for bfloat16.
block_size is a power of two from 32 through 1024; the supplied launcher handles invalid launch metadata.

Try Similar Problems

input =

Accepts: array

dtype =

Accepts: string

block_size =

Accepts: number

Examples

Requirements

Constraints

Kernel performance

Blockwise Online Softmax

Examples

Requirements

Constraints

Kernel performance