pub const ARGMAX_NB: usize = 256;
Number of pass-1 blocks for the parallel argmax (fan-out across SMs to saturate HBM). 256 blocks x 256 threads = 65536 threads covering the 248K-vocab scan in ~4 strided loads/thread.