Sparse convolution in FP32

Sparse convolution is most of a VoxelNeXt frame. On the densest qualification frame, after the detection head had been optimized, sparse convolution kernels took 86% of the device time and building neighbour maps another 4.6%. This page describes how those kernels were built, in exact FP32, and what was measured along the way.

All numbers are from one RTX 4060 Laptop GPU on nuScenes 10-sweep frames. Each entry is a measurement recorded in the project notebook at the time it was made, so later numbers include the earlier changes.

• 0.54×first correct kernel against spconvone residual block, 20 frames
• 1.283×backbone prefix against spconv20 held-out frames, all won
• 1.0e-5maximum absolute error in FP32against the spconv reference

The starting point

A submanifold convolution computes, for every active voxel, a sum over its 27 neighbours of the neighbour’s features multiplied by the weights for that kernel offset. The arithmetic is simple. The difficulty is finding which neighbours exist and arranging the work so the GPU can process many voxels together.

The first DriveRT kernel was a correctness reference: one GPU thread per output row and channel, scanning all 27 offsets and all input channels. On one sparse frame it beat spconv, but over 20 representative frames it took 1.85 times as long. spconv’s default on this model is mask implicit GEMM, so that was the target to beat, not its older native algorithm.

Two early alternatives were measured and rejected. A gather, GEMM and scatter path launched 54 small GEMMs per block and spent as much time moving features as computing them. A pair-parallel kernel removed the host loop but added contention on shared output rows and was slower on all 20 frames.

Neighbour masks and sorted rows

For a 3×3×3 kernel, the neighbours of a voxel can be written as a 27-bit mask. The coordinate map builder produces this mask while it builds the neighbour table, and a GPU radix sort orders the rows by mask. A kernel then processes a tile of rows together and only loads the offsets present in the union of their masks.

Real frames have many distinct masks, from 2,303 at 17k voxels to 29,915 at 100k voxels, but most rows still fall into tiles of rows with identical masks. Sorting changes only the order of work. Outputs are written back to their original rows, so the result is the same.

On the first residual block this took DriveRT from behind spconv to 1.02×, then to 1.14× by giving each thread several output rows, and to 1.17× with a 32-row tile and four rows per thread.

Some variants were correct and slower:

• keeping the first convolution’s output in sorted order, which added an indirection on every neighbour read of the second convolution;
• 64-row tiles, whose larger mask unions cost more than the extra weight reuse saved;
• scheduling exact runs of identical masks separately from a cleanup pass for the rest, which cost more in extra launches and small tails than it saved in wasted offsets.

One kernel family, configured per stage

The same pipelined kernel was then instantiated for 16, 32, 64 and 128 channels. Each width has its own tile shape: rows per tile, output channels per thread, and the number of input channels staged at once. Tile shapes were chosen by offline search on five calibration frames, frozen, and then evaluated on 20 held-out frames.

change                                 DriveRT     spconv    speedup   frames won
strided family + C128 tile search      18.161 ms   20.093 ms  1.106×     20/20
C32 pipelined, 12 rows × 2 channels    16.990 ms   20.317 ms  1.196×     20/20
C16 pipelined, 8 rows × 1 channel      15.903 ms   20.404 ms  1.283×     20/20

These are conv1 to conv4 of the backbone, with dynamic coordinate maps and exact FP32. The maximum absolute error against spconv stayed at 1.0e-5. After the C16 change every sparse stage was faster than its spconv counterpart on the profiled frame.

The best tile differs between stages of the same width. For the four C128 stages the search selected a 96-row tile on the larger domains and a 64-row tile on the smallest one. This is why selection is done per stage and recorded in the Plan, not written as a rule in the kernel.

Calibrated bit order

A mask tile executes the union of its rows’ offsets, so rows with similar masks should be adjacent. Which bit of the mask is most significant changes which rows end up together. The submanifold kernels already used a calibrated permutation of the 27 bits. The strided kernels, which change the coordinate set, still sorted by the plain mask value.

With the calibrated order the four larger strided stages ran 1.10× to 1.16× faster, with the same coordinate sets and bit-identical features. The smallest stage became slower, so it keeps the plain order. On the complete detector this saved 0.14 to 0.17 ms per frame.

Weight layout and memory access

The largest C128 stages were still about 1.15× slower than a dense GEMM library on identical shapes. Three measurements located the cost before any code changed:

• replacing the neighbour table with an identity map changed latency by only 3.5%, so the indirection was not the problem;
• only 3.7% of rows in the largest C128 domain have all 27 neighbours, so a separate fast path for interior voxels would not apply to enough rows;
• the weights were stored output-major, so each thread read its eight weights as eight separate memory loads.

Storing the weights as [offset][input][output] and giving each thread consecutive output channels turned those loads into two 16-byte reads. The output is bit-identical to the previous kernel, and the complete detector improved by 0.05 to 0.11 ms.

A later profile of the same kernel found a second problem in the same place. With eight consecutive channels per thread, the reads of neighbouring threads landed in different groups of shared memory and conflicted. Two separated runs of four channels made the reads contiguous:

shared-memory conflicts   8,695,533  ->  20,950
shared-load wavefronts      21.89 M  ->  13.20 M

Generalizing this kernel from C128 to all widths helped at 32 channels and not at 64. On the complete FP32 path the change saved 0.50 ms on the calibration frames and 0.64 ms on held-out frames.

Offset-group partitioning

Even with sorted rows, a 64-row C128 tile executed 1.21 times the useful work, because rows in one tile rarely have identical masks. The opposite extreme, one pass per kernel offset, has no waste but writes the output 27 times. Partitioning the offsets into a few groups, with the rows ordered separately for each group, sits in between.

Three groups reduced the executed work on the largest C128 domain by 12.2%, and its residual block became 9.6% faster. The difference is the extra accumulation passes over the output. On the complete detector the gain was 0.48 to 0.54 ms per frame, measured in three interleaved comparisons, and every frame improved.

The first version applied the ReLU only to rows present in the last group. Rows without an offset in that group kept unrectified values, exactly 18,253 minus 12,404 rows on the test frame. The last launch now spans every row.

The grouped map was first more expensive to prepare than the time it saved on small domains. Building every group’s order with one radix sort instead of one sort per group reduced preparation on the largest domain from 0.330 ms to 0.102 ms, against 0.098 ms without grouping.

Grouping is selected where it pays: on the C128 and C64 domains. On C32 it was 0.01 to 0.03 ms slower, because at that width the extra launches and ordering cost more than the reduced arithmetic.

Ideas that were measured and rejected

• Low-rank weights along the offset axis. The C128 weights reshaped to 27 × 16,384 are close to full rank: rank 24 of 27 still leaves 14% relative error.
• Dense blocks, and Winograd on top of them. Filling a block with zeros is exact for submanifold convolution, but even the densest domain inflates the work 1.82× at 2×2×2 blocks, and Winograd needs 4×4×4 input tiles, where inflation is 4.47×.
• Reusing neighbour maps between frames. The voxel index sets of consecutive keyframes have a Jaccard overlap of only 0.14 to 0.38, because ego motion shifts every index.
• Hybrid sparse and block-dense execution. A dense block evaluates all 27 offsets per cell while a sparse row evaluates about 16 plus tile waste, so dense execution starts 1.36× behind even for a full block.
• Wider microtiles for C32 and C64. They issue fewer memory loads per multiply-add and were still 0.09 to 0.49 ms slower in the complete detector.
• Larger input staging for C128. A 32-offset staging tile won the stage-level search and lost the complete-detector comparison, 27.98 ms against 27.85 ms.
• Skipping absent neighbours with a predicate. It raised the register count per thread from 124 to 170 and was slower.
• Spatially sorted output coordinates. Map preparation grew from 1.87 to 2.04 ms and the convolutions did not get faster.

Measuring kernels correctly

Several results in this project were first wrong because of the measurement, not the kernel:

• a tactic probe passed the production weight layout to every candidate, so a kernel expecting another layout was timed on data it could not read and reported plausible speedups with outputs wrong by whole units;
• a kernel selected its tile from the tensor width, so a candidate registered for one width ran at another and the search recorded it as that domain’s winner;
• a search keyed by operation contract gave every C128 domain the same key, so the last domain measured overwrote the others;
• without a minimum gain, the search selected differences of 0.008 ms.

Each of these is now checked: weights are prepared per candidate, width is enforced on every dispatch path, the search key includes the domain, and a candidate has to beat the incumbent by a fixed margin.

Where it stands

At the end of the exact-FP32 work, the C128 kernels ran within about 1.10× of a dense GEMM library on the same shapes, with 3.5% of that being the neighbour gather. Every kernel width was limited by how fast the GPU could issue work rather than by memory, near a local optimum among about forty registered tile variants. Further FP32 gains need a different algorithm rather than a different tiling.

The next step was reduced precision. TF32 keeps FP32 storage and changes how multiplications round. It is qualified on the reference model and still in progress.

Back to DriveRT · Detection head and row-selective evaluation