I1. Softmax over real tensor views¶
Prerequisite: beginner lessons 2, 4 and 8. Time: 2-3 sessions. Edit: strided_softmax.py. Goal: handle a view without assuming its logical neighbors are adjacent in memory.
Work out the addresses¶
Consider contiguous storage [0,1,2,3,4,5], viewed as a 2x3 matrix. Transpose it
to obtain [[0,3],[1,4],[2,5]]. The new shape is 3x2 and its strides are (1,3).
| Logical row | Logical column | row*1 + col*3 |
Value |
|---|---|---|---|
| 0 | 0 | 0 | 0 |
| 0 | 1 | 3 | 3 |
| 2 | 0 | 2 | 2 |
| 2 | 1 | 5 | 5 |
The pointer passed to Triton already addresses the first element of the view.
Do not add PyTorch's storage offset a second time. A slice such as
base[1:, 1::2] combines a shifted pointer with a column stride of two.
The old portfolio softmax intentionally rejects noncontiguous columns. In this
workshop the load is X + row*SR + col*SC. The output has a newly allocated
contiguous layout, so its store is Y + row*N + col. Input and output addressing
are different even though the tensors have identical shapes.
Implement in three passes¶
- Add a
@triton.jitkernel with one program per logical row. Pass both input strides and width. Use the next power of two for the reduction block. - Add stable FP32 softmax, negative-infinity tail padding, and masked stores.
- Validate with
learning.workshops.contracts.matrixandwarps, allocate the output, and launch under the input's CUDA device context. Empty row counts return the empty output without launching.
Read-only broadcast inputs with stride zero are valid: several logical entries can read one address. Output entries must still have distinct addresses. This does not authorize overlapping writes to the input.
python -m learning.workshops.check strided_softmax --hint
python -m learning.workshops.check strided_softmax
python -m learning.workshops.benchmark --op strided_softmax --output results/local/i1.json
Experiment, not just correctness¶
The benchmark constructs transposed inputs. Compare direct loads using 4 and 8
warps with a variant that calls x.contiguous() before the same kernel. The copy
belongs inside the timed function. Predict whether the extra traffic can pay for
better access during reduction, then inspect both shapes and both timing modes.
Logical strides explain requested addresses, not actual DRAM transactions. A claim about coalescing efficiency needs generated-code/profiler evidence. Start with the address pattern and a measured runtime; do not invent hardware counters.
Debugging clues: a contiguous case passes but transpose fails -> check SC; only slices fail -> check double-counted offset; values are right but shape/layout is wrong -> check output allocation; probabilities are too small -> inspect padding.
Done when: all checks pass, you draw the transpose addresses without code, and your worksheet reports the cost of the copy rather than excluding it.
Reference: Triton softmax tutorial. Next: derive the backward pass.