A custom TPU kernel must give the same answer as the compiler path it replaces. For softmax with a mask, that question has more edges than it first seems.
The TPU Kernel Parity Lab compares a Pallas masked softmax kernel with JAX and XLA under stated test conditions. Its benchmark dataset is built for TPU v5e. The source and the dataset are both linked from that page.
This post explains what parity means for this kernel. It also shows how the checks from other Computelab.co studies apply.
Pallas and XLA in one sentence each
XLA is the compiler that JAX uses by default. It takes a softmax written in jax.numpy and fuses it into TPU code on its own.
Pallas is the JAX kernel language. It lets an engineer write the tiling, the memory moves and the loop order by hand.
A Pallas kernel can choose a different tile shape or reduction order than XLA. Each choice can change the floating point result by a small amount. Parity work measures that amount and decides if it is acceptable.
Why masked softmax is a hard parity target
Softmax subtracts the row maximum, takes the exponent, then divides by the row sum. The mask removes some positions from the row before that happens.
Each step is a place where two correct kernels can differ. The row maximum must ignore masked positions. The row sum is a reduction, so its order changes the low bits. The divide spreads any error in the sum across every output in the row.
Edge rows need their own checks. A row where every position is masked has no valid maximum. A row with one valid position should return exactly one there. A mask that cuts across a tile boundary tests whether the kernel tracks the mask through its tiling.
These are the cases where a hand written kernel and a compiler path are most likely to disagree.
The reference comes first
Comparing Pallas against XLA alone is not enough. If both drift the same way, they agree with each other and still miss the true answer.
The studies on this site compare against a float64 reference instead. The Tesla T4 copy overlap study checked its result against a float64 CPU product. The Kernelwitness study measured three summation orders against a float64 reference.
The same approach fits softmax. Compute the masked softmax in float64 on the host. Then measure both the Pallas kernel and the XLA path against it.
Set the bound by measurement
A fixed percent tolerance is risky for a reduction. On the Kernelwitness histogram, a one percent relative tolerance accepted 1,091 of 4,096 wrong bins. A measured rule accepted 4.
That rule measured correct summation orders in ULP against float64. The bound was set at four times the worst observed distance. In that study the worst correct order sat at 108 ULP.
For a Pallas softmax, sample the legal tile shapes and reduction orders the kernel may use. The XLA path is one more correct order. The bound comes from the spread of all of them.
The tolerance post walks through this calibration step by step.
Guard the output buffer
Masked kernels often skip work for masked tiles. That is a valid speedup, and it is also where missing writes hide.
Fill the output with a NaN pattern before each run, and use fresh inputs each round. Any element the kernel fails to write then fails the check. The poison buffer post covers this gate.
Where this fits in porting
The engine ports across CUDA, HIP, NKI, Pallas, Triton and SYCL, as listed on the FAQ. A softmax that moves from Triton or CUDA to Pallas on TPU needs the same parity proof.
A related softmax result sits on the case studies page. There, a Triton softmax ran about 1.35x faster than PyTorch in float16 GPU graph timing.
The steps stay the same on every chip. Build a float64 reference and sample the correct orders. Set a measured ULP bound, poison the output, and record the conditions.
Run your workload
Bring the kernel or solver run that costs you the most time. Run a workload with Computelab.co.