I need you to write a numerically tight softmax for the RTX PRO 6000 (SM120 Blackwell, GDDR7, 1.8 TB/s). The reference implementation is in reference.py and your kernel goes in solution.py. You can make whatever mess you want in this directory — scratch files, intermediate kernels, build artifacts, profiling traces — but the final answer has to be in solution.py with the same Model, get_inputs, and get_init_inputs interface as reference.py.

The op is softmax along the last dim of a 2D fp32 tensor. The reference computes ground truth in fp64 and casts back to fp32, and you have to match it within atol=rtol=1e-5 — that's a tighter bar than default fp32 (1e-4) on purpose. With long reductions naive fp16 accumulation drifts past this; fp32 accumulation with subtract-max stability is enough on most shapes; on the largest vocabs you may need compensated (Kahan-style) summation to stay under the bar. The shapes you have to handle are batch=32 vocab=4096 (sanity), batch=16 vocab=32768 (GPT-2 class), batch=8 vocab=131072 (Llama-3 vocab), batch=4 vocab=262144 (256K, DeepSeek-V3 / Gemma-3 class — naive fp16 sum drifts past 1e-5 here), and batch=8 vocab=131072 with extreme logits (a few very large positives per row to stress max-subtract — exping before subtracting overflows). The check and benchmark scripts handle the extreme-flag input generation; you just need to read x and return y.

This needs to be a real custom kernel. Don't import or call torch.nn.functional.softmax, torch.softmax, F.softmax, liger_kernel.softmax, liger_kernel.transformers.softmax, or any .softmax( method on a tensor. Try Triton, CUDA C++ via load_inline, or inline PTX — whatever fits. Anything you're uncertain about, look up PTX docs, clone Liger-Kernel or other reference repos, read library headers, and investigate.

Your flywheel is implement, profile (ncu, nsys, torch.profiler — whatever's useful) and time it with benchmark.py, verify correctness by running `python check.py` and reading the output, then iterate. Don't substitute your own one-off correctness snippets for check.py — it iterates over every shape, your spot-check almost certainly won't. If `python check.py` hasn't printed PASS, you're not done. Take as long as you need to actually push the number up.