Discussion on "Rebuilding Triton and Helion from Scratch in 4,000 Lines of Python" | Hashnode

Discussion on "Rebuilding Triton and Helion from Scratch in 4,000 Lines of Python" | Hashnode

Responses5

Markdown supported · ⌘↵ to post

Ahmet Özel

The 2,500 and 1,500 element rows are a useful sanity check because they force a ragged tail on every row instead of hiding it in one launch. The next benchmark I would add is a bucketed sweep by tail occupancy, reporting achieved bandwidth and instruction count for 1-3, 4-7, and 8-15 valid lanes. That would separate branch divergence from lost vector width, since the branch can stay warp-uniform while memory transaction efficiency still falls.

Ahmet Özel

The runtime contiguity check is a clever trade: a few integer comparisons are cheap relative to the roughly 500-cycle memory access they unlock. I would be curious to see the cost split for aligned interior blocks versus masked tail blocks, since real workloads often have awkward dimensions rather than clean powers of two. If the check is warp-uniform, it should avoid divergence, but shape-distribution benchmarks would show how often the scalar fallback becomes visible.

Arpit Singh Gautam

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.

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.