On the cost split: the check is a handful of integer compares per 16-byte group, so on an aligned interior block it is pure overhead in front of the vector load, a few ALU ops hidden under the memory latency; on a masked tail the affected groups drop to per-lane predicated scalar. I have not broken that into a clean fast-path-versus-fallback microbenchmark, and you are right that a shape-distribution sweep is the honest way to show it. The one data point I can offer in that direction: the softmax and layernorm examples deliberately use non-power-of-two row lengths (2500 and 1500), so every row carries a masked tail rather than a single tail block for the whole launch, and they still sit at the memory-bandwidth ceiling next to Triton. So at least for bandwidth-bound kernels the fallback stays hidden behind the DRAM access even when it fires on every row. Isolating the two paths across a spread of awkward shapes is a good addition and something I would like to run properly.
@arpitsinghgautam
Arpit Singh Gautam
I work on efficiency × reliability for foundation models, mostly quantization and what it quietly breaks
About
I'm Arpit Singh Gautam, a data scientist in the CSG CTO Lab. I work on efficient LLM inference and the reliability of quantized models, things like RL-based mixed-precision quantization, KV-cache optimization, and what post-training quantization quietly does to a model's calibration, factual recall, and robustness to prompt injection.
Available for
Collaborations on LLM efficiency and reliability, open-source GPU kernel work, paper discussions, and 1:1 mentoring. Always happy to talk quantization, inference systems, or compilers.
Arpit Singh Gautam's blogs
Comments
AS
AS
Thanks Ahmet, good questions, and you have the divergence point right. On an interior block the branch is effectively warp-uniform: the offset stride comes from the same arange expression for every lane, and the alignment term reduces to "is the base pointer 16-byte aligned," which is the same across the warp, so the only per-lane term is "are all lanes in this group unmasked." In the interior that is trivially true, every thread takes the vector path, and there is no divergence. Even a misaligned base fails uniformly, so all lanes take the scalar path together. The only place lanes can actually disagree is the boundary block where the mask goes ragged, so divergence is confined to that one block per launch (or the tail of each row, for a row kernel), and the fallback is predicated rather than branchy, so it does not cascade.