TL;DR: A fused row softmax loads each valid logit, performs the reduction and normalization on chip, then writes each result. That is one input read and one output write when the working row fits without spills. Use independent input/output strides, mask padding, accumulate in FP32, and define fully masked rows. Reductions and exponentials still cost time; a byte bound is not a benchmark.
How to approach it
State the shape, stride and numerical contract first. Write the masked kernel, explain its launch, then derive the traffic. Close with correctness tests and the resource signal that would make you choose a different row strategy.
A strong answer
A softmax passes a contiguous-tensor test but corrupts the next output row for a sliced input. The input has padding between rows; the newly allocated output does not. Reusing the input stride for the output was the bug. Triton is a Python-embedded language for writing GPU kernels in blocks of elements rather than individual threads. This is why the pointer contract belongs in the answer before performance tuning.
A row you can follow by hand
Logits are the scores being converted into probabilities. Subtracting their maximum preserves those probabilities while avoiding overflow from large positive finite values. Padded lanes load negative infinity and contribute zero. An all-negative-infinity row needs separate treatment: subtracting its maximum would otherwise evaluate negative infinity minus negative infinity.
Kernel and launch contract
The following kernel accepts positive-stride, non-overlapping input and output views. Each valid input is finite or negative infinity; NaN and positive infinity are outside this example's contract. Each all-masked row produces zeros, an explicit application policy rather than a probability distribution. FP32 or FP16 storage was checked below; normalization uses FP32.
import triton
import triton.language as tl
@triton.jit
def row_softmax(X, Y, x_row, x_col, y_row, y_col,
n_cols, BLOCK: tl.constexpr):
row = tl.program_id(0).to(tl.int64)
col = tl.arange(0, BLOCK).to(tl.int64)
valid = col < n_cols
x = tl.load(X + row*x_row + col*x_col,
mask=valid, other=-float('inf')).to(tl.float32)
m = tl.max(x, 0)
# Define an all-masked row as zeros, rather than NaNs.
m = tl.where(m == -float('inf'), 0.0, m)
numerator = tl.exp(x - m)
denominator = tl.sum(numerator, 0)
y = numerator / tl.where(denominator > 0, denominator, 1.0)
tl.store(Y + row*y_row + col*y_col, y, mask=valid)
Launch one program per row with grid (M,), n_cols=N, and BLOCK=triton.next_power_of_2(N). Pass each view's row and column strides in elements, independently. num_warps=4 is a starting configuration to test, not an optimal choice for every width. Skip the launch for an empty tensor; do not call the kernel with an empty row. The host must validate storage bounds, output non-overlap and supported dtype before launch, and catch compilation/resource failures for unsupported widths. The widened offsets avoid doing row-address arithmetic in 32-bit integers.
Triton Programming Model explains how the program's elements map to threads. In this ordinary launch a program is a thread block on one SM; it does not spread the row across multiple SMs. Reductions may use shared memory and synchronization as well as registers. Register pressure includes live temporaries, not just one value per column.
Execution scope: this exact kernel passed 48 CPU-interpreter cases using Triton 3.4.0, with NumPy-backed pointers and only the Torch-specific host-transfer adapter replaced. Cases cover FP32/FP16, different input/output strides, widths 1, 3, 31, 33, 1000 and 4097, padding guards and fully masked rows. This verifies interpreted arithmetic and addresses; it does not establish GPU compilation, race behavior, BF16 support or speed. Those need the target device and a compiled-kernel check.
Count bytes without promising a runtime
M = 32,768 rows; N = 4,096 columns; s = 4 bytes/value
Tensor bytes = M × N × s = 536,870,912 B = 512 MiB
Read + write = 2 × M × N × s = 1,073,741,824 B = 1 GiB
Illustrative bandwidth = 3 × 10^12 B/s
Time floor = bytes / bandwidth = 0.000358 s = 0.358 ms
Check: one read plus one write must be twice the tensor size.
This is a bandwidth floor, not a measured runtime.
Compare with an explicitly separate max, subtract, exp, sum and divide implementation: logical reads are 5MN + 2M elements and writes are 3MN + 2M. Ignoring the small row statistics, its activation traffic is about four times the fused 2MN traffic. Cache reuse can change DRAM transactions. torch.softmax is not this naive five-operation baseline; a tuned library kernel may already do the same fusion.
Memory-Bound vs Compute-Bound Kernels supplies the next question: which resource actually limits this shape? Exponential throughput, reductions, instruction overhead and spills can dominate. Halving storage bytes with FP16 does not guarantee halving elapsed time, and a copy bandwidth comparison does not prove that softmax has no remaining work worth optimizing.
When the row is too wide
Power-of-two padding can nearly double the working set just above a boundary. Compare compiled register/shared-memory use and local-memory traffic across row widths and warp counts. Keep this design while it is correct and fastest for the workload; reverse the choice when spills, reduction overhead or too few independent rows make another implementation faster.
A chunked two-pass alternative first computes the final maximum and denominator through online rescaling, then rereads logits to write normalized outputs. Its basic traffic is two input reads and one output write. An online denominator alone cannot emit final probabilities for earlier chunks before the final denominator is known: it must retain data or revisit it. Kernel Fusion with a following value reduction can avoid storing the probabilities entirely, which is the idea developed in FlashAttention.
What interviewers probe next
- Why separate strides? Allocation and slicing can give input and output different layouts.
- Why FP32? BF16 has seven stored fraction bits; reductions need more precision than that, although this example has not been device-validated for BF16.
- How would you time it? Warm compilation first, use device events with synchronization, record shape/dtype/strides and compare with the library baseline. Keep allocation and transfer timing separate when the application does.
Common mistakes
- Describing one read and one write as only one tensor's worth of traffic.
- Allowing padding or an all-masked row to contaminate the denominator with NaNs.
- Claiming a universal row-width cutoff or a measured speed from peak bandwidth.
Key takeaways
- The 32,768 × 4,096 FP32 example moves at least 1 GiB through its input/output contract.
- Separate stride arguments and output guards catch layout bugs that contiguous tests miss.
- Keep the one-program row strategy only while its measured resource costs beat the alternatives.
The Triton tutorial provides the reference fusion pattern; its hardware results are not measurements of this kernel.