[Kernel] Add submanifold sparse 3D convolution (bf16 implicit GEMM, gfx950) - #942
Draft
jiacao-amd wants to merge 10 commits into
Draft
[Kernel] Add submanifold sparse 3D convolution (bf16 implicit GEMM, gfx950)#942jiacao-amd wants to merge 10 commits into
jiacao-amd wants to merge 10 commits into
Conversation
added 4 commits
July 31, 2026 01:35
…fx950)
Self-contained spconv-style operator: raw voxel coordinates in, dense output
features out, kernel map built internally. Matches the input convention of
spconv / MinkowskiEngine / TorchSparse so it can drop into their call sites.
Three stages, four GPU kernels:
coords --[map]--> LUT --[compact]--> pair list --[gemm]--> output
MAPPING. Two interchangeable builders, routed by bounding-box volume, because
neither dominates:
dense presence grid O(1) lookup, but O(volume) memory AND fill, so its
cost tracks the grid rather than the point count.
2-4x faster on compact grids; cannot address past
2^31 cells.
Z-Delta Spira, arXiv 2511.20834 5.2. One sort, then K^2
binary searches instead of K^3, using
packed(q+d) == packed(q) + packed(d) so queries never
unpack. O(N) memory, cost flat in grid volume.
MEASURED: 195^3 dense 41 us / zdelta 160 us; 595x595x245 dense 93 / zdelta 180;
2095^2x695 dense CANNOT RUN (3.05e9 cells, 12.2 GB of presence, and the linear
index overflows int32) while zdelta does it in 203 us. That last grid is PTv3
on SemanticKITTI, so the second mapper is a functional requirement outdoors,
not a tuning option. Keys are packed into 32 bits when the padded extent allows
(1.15-1.27x) and int64 otherwise.
COMPACTION. An indexed [tile, kv, row] layout runs all 16 rows of a tile even
when only 2-3 have a neighbour for that tap: MEASURED 5.4-6.2x wasted MACs,
i.e. 16-21% of the multiplies useful, while executed throughput is already
45-53% of peak -- the hardware saturates computing zeros. Relisting only the
valid pairs, grouped by tap so a tile shares one weight slab, makes every MFMA
tile full. Padding each tap's run to a multiple of 16 costs 1.002x.
GEMM. mfma_f32_16x16x32_bf16 with fp32 accumulate, one wave per block, gather
and scatter both fused in: A is read through in_rows[], results are atomically
added through out_rows[], and no per-pair intermediate is materialised.
Four things that shaped the design, each measured rather than assumed:
* No LDS. With one 64-lane wave and a 16-wide tile every operand element is
read by exactly one lane -- reuse is 1.00x in fp32 and again in bf16 after
vectorisation. Staging through LDS is a pure detour that also sets
LDS/block = C_IN*64 B, which collapsed occupancy to 5 blocks/CU at C_IN=512
(isolated by compiling for a padded C_IN: 17x slower at identical work).
* Weights are packed [kv, t, k_outer, u, k_inner]. The obvious K-major
[kv, t, u, ci] gives each lane its 8 consecutive k but puts adjacent lanes
C_IN*2 bytes apart, so a wave reads 16 half-used cache lines: 50%
coalescing. Interleaving k_inner innermost makes it 8 fully-used lines.
* Loads for K_UNROLL k-steps are hoisted above their MFMAs. Issuing each MFMA
right after its load left the wave on L2 latency (24-27 cycles against an
8-cycle issue rate) even though DRAM traffic is 7-16% of peak and L2 hits
91-96%. This is the no-LDS way to get prefetch.
* The centre tap is compacted along with the rest. Leaving it on a separate
indexed kernel looks right -- SubM makes it 100% dense already -- but that
ignores the extra launch and the extra read-modify-write: folding it in
measured 1.40-2.46x faster and removed a whole kernel.
dtype is bf16 because production sparse conv is 16-bit: spconv ships
fp32/fp16/bf16, its own benchmark tables are F16/TF32, and NVIDIA's shipped
CenterPoint logs are float16 and int8 with no fp32 variant. An fp32 GEMM was
built first and measured 1.6-2.9x slower at every shape. Rounding the inputs
costs ~2.5e-3 relative error per layer; across a 22-layer stack with
normalisation and residuals that compounds to ~1.3e-2 (cosine similarity
0.99994, argmax agreement 98.6%). NOT yet validated against a pretrained model
-- that gate should be cleared before relying on it for accuracy-sensitive work.
Sized against MEASURED production shapes, not synthetic ones: nuScenes
CenterPoint over all 6019 val frames is median 91,090 active voxels (p95
114,827, max 135,306), and PTv3 captures give ScanNet ~148k and SemanticKITTI
~120k. Earlier tuning had used a 10k-point toy input, which understates both N
and neighbour density by enough to invert several conclusions.
End to end on those shapes (warm cache, includes mapping):
nuScenes C=128 N=32371 47 us
nuScenes C=256 N=26509 101 us
nuScenes C=512 N=15052 172 us
SemKITTI C=128 N=94517 178 us
Host-side torch is used for allocation, torch.sort (Z-Delta) and the compaction
scan; the GPU kernels are pure FlyDSL. Replacing those host ops with
hand-written kernels was tried and measured 2-8x SLOWER -- torch's
nonzero/cumsum are parallel scans and the replacement serialised on atomics.
Mechanical cleanup per the kernel-code-cleanup skill, in the two mapping kernels
and the akv-compaction tail they share:
arith.cmpi/andi/addi/muli -> Python operators on typed fx values
arith.constant(n, T.i32) -> fx.Int32(n)
fx.Index(...) -> fx.Int32(...)
vector.extract/from_elements-> fx.Vector(v)[i] / fx.Vector.from_elements([...])
arith.constant_vector -> fx.Vector.filled(...)
gpu.thread_idx.x -> gpu.thread_id("x")
const_expr(<python int>) -> the int itself (builders take plain ints)
fx.Index was not just verbose but wrong-ish here: buffer_load documents its
offset as i32, so every fx.Index(...) forced an implicit index->i32 cast. The
offsets are now i32 throughout.
scf.IfOp is kept explicit, and the reason is structural rather than stylistic.
_collect_assigned_vars (ast_rewriter.py) treats the RECEIVER OF A METHOD CALL
inside a branch as loop/branch state, not just assigned names:
def visit_Call(self, node):
base_name = RegionAnalyzer._get_call_base(node.func)
if base_name is not None and base_name != "self":
add_unique(invoked_args, base_name)
Every branch body here reads or writes a GTensor (lut_.store, pr_.load, ...),
so the GTensor lands in result_names and scf_if_dispatch rejects it with
"state variable 'lut_' is GTensor, not an MLIR Value". Three formulations were
tried under FLYDSL_RUNTIME_ENABLE_CACHE=0 -- assignments inside the branch,
loads fully inlined so the branch assigns nothing, and the innermost branch
whose values are dead afterwards -- and all three fail the same way. This is a
GTensor/rewriter incompatibility, not something callers can write around; fixing
it means changing kernels/common/tensor_shim.py, which four other kernels share.
The conditions are now typed compares with .ir_value() at the boundary.
GTensor / buffer_ops also stay for that same shared-shim reason.
Verified against a golden capture taken before the migration: the LUT, the
compacted in_rows and out_rows are BIT-IDENTICAL on four shapes including K=5.
Float output differs by ~1e-7, which is below the kernel's own run-to-run noise
(~2e-6): buffer_atomic_add reassociates, so re-running unmodified code differs
by the same amount. 39/39 tests pass with FLYDSL_RUNTIME_ENABLE_CACHE=0.
Attempted the same fx-operator migration there and reverted it. Its keys are KEY_T (i32 or i64 depending on grid extent) and GTensor.load returns a raw value typed by that dtype rather than an fx wrapper, so operator compares do not apply and fx.Int32 offsets take a different buffer_load path -- the result fails to compile. Recorded in-code so the next reader does not retry it.
…face
Same mechanical cleanup as the map kernels, applied to kernel_bf16_cmp:
fx.Index(...) -> fx.Int32(...) (buffer_load wants i32; every
fx.Index was forcing an index->i32 cast)
arith.constant / cmpi -> fx.Int32(n) / typed compares
arith.constant_vector -> fx.Vector.filled(4, 0.0, fx.Float32)
const_expr(<python int>) -> the int (make_layout/vec_load take plain ints)
dropped `lane = tid` alias and the duplicate mfma_row/mfma_col computation
The EPILOGUE is left on fx.Index / vector.extract deliberately, with the reason
in-code. Migrating it is op-count-identical but measured 185 us vs 170 us at
C_OUT=512 (median-of-7, range 184.8-186.4 so well outside noise) -- an
instruction-scheduling effect around the predicated atomics, the same class of
exception PR ROCm#913 documents for mxfp_moe gemm1 and fp8_4wave. Bisected by
reverting prologue/hot-loop/epilogue independently: the hot loop is free, the
epilogue is not.
Verified: LUT, in_rows and out_rows BIT-IDENTICAL to the pre-migration golden on
four shapes incl. K=5; float output within the kernel's own atomic-order noise.
39/39 tests pass cold. Perf median-of-5 vs before: C=128 1.00x, C=256 1.02x,
C=512 1.02x.
jiacao-amd
force-pushed
the
spconv-bf16-implicit-gemm
branch
from
July 31, 2026 01:36
c2e72f3 to
1ad9e42
Compare
added 6 commits
July 31, 2026 18:06
…piler _compile_fused_tiled_kernel took kv_total (= K**3) and recovered K with round(kv_total ** (1/3)), which needed an assert to catch a bad inverse. The caller already has kernel_size and was squaring-then-cubing it away, so pass it directly -- K**3 is exact and the assert has nothing left to guard. This also matches _compile_zdelta_kernel, which already took kernel_size. Behaviour-preserving: LUT/in_rows/out_rows bit-identical on four shapes incl. K=5, 39/39 tests pass cold.
An earlier commit (b0d838b) claimed scf.IfOp had to stay because _collect_assigned_vars treats a method-call receiver inside a branch as branch state, so any GTensor used in the body lands in result_names and scf_if_dispatch rejects it with "state variable 'lut_' is GTensor, not an MLIR Value". The diagnosis was right; the conclusion was not. The receiver only becomes state if it is a symbol from an ENCLOSING scope -- _collect_assigned_vars filters through in_active_symbols(). Constructing the GTensor inside the branch that uses it makes the name local to that branch, so nothing is collected and the plain `if` lowers normally. Deeper branches each need their own local view; that is why the first attempt at this still failed one level down. All nine scf.IfOp sites convert -- the scatter kernel, both mapping kernels (4-deep nesting in zdelta: o_ok / binary-search / p_ok / hit), and the GEMM epilogue. The scf import is now unused and removed. The epilogue comment about fx.Index / vector.extract is unchanged and still accurate: this commit only replaces the branch construct there, not the index arithmetic that carried the 9% cost. Verified with FLYDSL_RUNTIME_ENABLE_CACHE=0 throughout, since a stale compile cache is what produced the wrong conclusion the first time: - 39/39 tests pass - LUT, mask, active_kv_ids and active_count BIT-IDENTICAL to the pre-change baseline on four shapes including K=5 - a real runtime branch survives lowering (one llvm.cond_br in the final IR), so the `if` is not being constant-folded away - no perf change at the shape that motivated the original epilogue note: N=100k C=128 174.5 -> 174.7 us, C=512 1143.4 -> 1145.0 us (median-of-7; 0.1%, within run-to-run noise)
The four lru_cache sizes here (4/8/16/32) were the smallest in the whole
kernel library -- 4, 8 and 16 were each the only occurrence of that value
across all 40 compile caches under kernels/. The convention in this very
directory is 64 for the helper kernels and 256 for the main one:
conv3d_implicit.py compile_transpose_ncdhw_ndhwc 64
compile_conv3d_implicit 256
conv3d_implicit_fp8.py all three 64
This now follows it: the three mapping kernels get 64, the GEMM gets 256.
The old numbers came from counting the key space and adding a small margin,
which is the wrong thing to optimize. A cache entry is one compiled function
object; an eviction is a full recompile, measured here at 709 ms against a
0.05 ms cached call -- roughly 14,000x. The trade is that asymmetric, so
every other kernel in the repo just picks a comfortably large number.
Counting also got one of them wrong. The zdelta key includes n_bits, which
is derived from N at run time:
n_bits = max(1, int(N - 1).bit_length() + 1)
A PTv3-style stage sweep (N = 338..100k) yields 7 distinct n_bits, times
kernel_size in (3, 5) = 14 keys against maxsize=16. Two spare slots, and
the failure mode is not an error but an intermittent 709 ms stall.
39/39 tests pass with FLYDSL_RUNTIME_ENABLE_CACHE=0.
build_compacted_pairs was the largest single stage of the map path -- 297 us of the 443 us total at nuScenes scale (N=91k), against 38 us for the GEMM it feeds. Profiling it statement by statement showed the time was not in data movement: two host-device syncs cost 79 us (34%) and ~17 torch ops carry ~75 us of pure dispatch, while a single pass that only reads the lut takes 25 us. An earlier attempt to fuse this into one kernel failed because it replaced the prefix sum with atomic counters and serialised 380k pairs onto 27 of them. The fix is not a better atomic but a smaller scan: with one thread per (kv, tile), the scan runs on the [KV, num_tiles] counts instead of the KV*num_tiles*block_m lut slots -- 16x fewer elements -- so an ordinary cumsum suffices and no device-wide scan inside a kernel is needed. _compile_pair_count count valid rows per (kv, tile) cumsum exclusive scan over the small counts _compile_pair_write each thread writes its rows from its reserved offset _compile_tile_kv emit tile_kv (replaces repeat_interleave, 54 us on its own) Pairs stay ordered (kv major, tile, row), which is exactly what nonzero() produced, so this is a bit-identical replacement rather than an equivalent one. Verified against the old implementation inlined from HEAD, on 11 shapes covering N=1, N=7, N=16, N=17 (block_m boundaries), k=3 and k=5, batched [N,4] coords, and nuScenes at N=91090: in_rows, out_rows, tile_kv and n_tiles_c all match exactly. 39/39 tests pass with FLYDSL_RUNTIME_ENABLE_CACHE=0. build_compacted_pairs 297 -> 120 us (2.5x) end-to-end, no key 602 -> 297 us (2.0x) end-to-end, with key 51 -> 51 us (unchanged -- the stage is cached) The two kernels now take 22 us combined, close to the 25 us single-read floor; the remaining ~98 us is the cumsum and allocation glue. Note this only helps the first layer of an indice_key-sharing stack, or callers that pass no key at all.
Drops the user-facing indice_key argument. The cache is now keyed on (id(coords), coords._version, block_m, kernel_size), so reuse is automatic and a different or mutated coordinate set is a miss rather than a hit. indice_key was copied from spconv, which is the outlier among the libraries that ship this. torchsparse's Conv3d takes no key at all and keys its kmaps on (input.stride, kernel_size, stride, dilation); MinkowskiEngine keys on the coordinate map identity. NVIDIA's deployed libspconv does carry a `rulebook` string, but only as an ONNX graph attribute copied verbatim from the training module's indice_key -- the inference API is `forward(stream)` and no caller ever names a rulebook. The old key had no coordinate identity in it, and the hit-time guard compared only N, so two different coordinate sets with equal N silently shared one rulebook. Measured on box=32 N=400 C=64: max|diff| 0.388 against |gold|max 0.525, with no exception raised. Under the new key the same sequence gives 2.98e-08. Note that spatial_shape is deliberately still absent from the key: it selects a storage stride and a bounds check, both invariant over occupied cells, and 60 randomized configs (k=3 and k=5, tight vs loose vs oversized shapes, including a fully dense 3x3x3 cube where every point is on a face) produce bit-identical LUTs -- adding it would only force needless rebuilds. id() is recycled once a tensor is freed, so a weakref.finalize drops both cache entries first; this is the scheme _WPACK_CACHE already uses for weights. Verified: after three short-lived coord tensors die, both caches are empty, so a recycled id cannot hit a stale entry. In-place mutation is caught by _version. Two views of the same rows get different ids, which costs a rebuild but is never a wrong hit. A 4-layer block with the caller passing nothing now costs 464.7 us (116.2 us/layer) against 51.6 us steady-state -- the rulebook is built once. Before this, the same code without indice_key rebuilt it every layer. 39/39 tests pass with FLYDSL_RUNTIME_ENABLE_CACHE=0.
build_compacted_pairs ran a torch cumsum plus a blocking .item() between its
count and write kernels. That block measured 88-183 us, 45-73% of the whole
compaction, against 17-129 us for the three fx kernels around it.
Replace it with a decoupled scan over the [KV, num_tiles] counts: a per-chunk
LDS scan on the full grid, a KV-way parallel scan of the chunk totals, a serial
pass over the KV row totals (KV is 27 or 125), then a grid-parallel fold into
the slot offsets. n_tiles_c is published to totals[0], so the host reads one
scalar instead of syncing on a reduction.
compaction ptv3 N=10k 112.5 -> 67.5 us 1.67x
nusc N=90k 121.7 -> 68.2 us 1.78x
nusc N=720k 317.9 -> 172.0 us 1.85x
N=300k 233.8 -> 140.9 us 1.66x
N=500k 329.4 -> 207.9 us 1.58x
end to end ptv3 N=10k 175 -> 111 us
nusc N=91k 310 -> 220 us
nusc N=135k 367 -> 276 us
Compaction is not free, so this moves where it pays for itself: measured
against an uncompacted pair list feeding the same gemm, the break-even shifts
from C~192-256 down to C~128-192. At N=150k, C=128 goes from -47 us to +16 us.
Also drop mask/active_kv_ids/active_count from both map kernels. Nothing
consumed them -- the gemm never took them and the coords entry point unpacked
them into throwaway locals -- but producing them cost a block barrier and a
serialized tid==0 loop over all KV taps. Dense map 33.1 -> 21.8 us at N=10k.
test_compacted_path_matches_indexed toggled FLYDSL_SPCONV_COMPACT to A/B the
compacted and uncompacted paths, but nothing in kernels/ has ever read that
variable, so both halves ran the same path and the test could not fail. Rebuild
it against a pair list derived straight from the lut. Verified it now catches a
compaction bug: dropping one row in the write kernel fails all four cases.
jiacao-amd
force-pushed
the
spconv-bf16-implicit-gemm
branch
from
August 1, 2026 00:04
2d1d4cd to
f9cb4f1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a self-contained submanifold sparse 3D convolution (spconv-style): raw voxel
coordinates in, dense output features out, kernel map built internally. Matches the
input convention of spconv / MinkowskiEngine / TorchSparse so it drops into their call
sites. Four GPU kernels, one public entry point.
Draft — see Open questions at the bottom before merging.
Design, and the measurements behind it
Sized against measured production shapes rather than synthetic ones: nuScenes
CenterPoint over all 6019 val frames is median 91,090 active voxels (p95 114,827, max
135,306), and PTv3 captures give ScanNet ~148k and SemanticKITTI ~120k. Earlier tuning
had used a 10k-point toy input, which understates both N and neighbour density by enough
to invert several conclusions.
Two mappers, both required.
build_lut_autoroutes by bounding-box volume:The dense presence grid is 2-4x faster on compact grids but is O(volume) in memory and
fill; Z-Delta (Spira, arXiv 2511.20834 §5.2) sorts
packed keys once and does K² binary searches instead of K³, flat in grid volume. The
last row is PTv3 on SemanticKITTI, so the second mapper is a functional requirement
outdoors, not a tuning option.
Compaction is the largest single win (2.4-3.3x). An indexed
[tile, kv, row]layoutruns all 16 rows of a tile when only 2-3 have a neighbour for that tap — measured
5.4-6.2x wasted MACs, i.e. 16-21% of multiplies useful, while executed throughput was
already 45-53% of peak. The hardware was saturated computing zeros.
No LDS, deliberately. With one 64-lane wave and a 16-wide tile every operand element
is read by exactly one lane — reuse is 1.00x in fp32 and again in bf16 after
vectorisation. Isolated by compiling for a padded C_IN so useful work is unchanged: LDS
alone made it 17x slower by collapsing occupancy to 5 blocks/CU at C_IN=512. An
LDS-transposed epilogue was also built and measured 8x slower, with padding making no
difference — bounded by the round-trip, not by bank conflicts.
Weight layout is
[kv, t, k_outer, u, k_inner]. The obvious K-major[kv, t, u, ci]gives each lane its 8 consecutive k but puts adjacent lanes C_IN*2 bytesapart, so a wave reads 16 half-used cache lines: 50% coalescing, and that was 84% of the
kernel's requested traffic. Interleaving k_inner innermost makes it 8 fully-used lines —
1.8x, a pure packing change.
dtype is bf16 because production sparse conv is 16-bit: spconv ships fp32/fp16/bf16,
its own benchmark tables are F16/TF32, and NVIDIA's shipped CenterPoint logs are
perf-float16-*.logandperf-int8.logwith no fp32 variant. An fp32 GEMM was builtfirst and measured 1.6-2.9x slower at every shape.
Performance
End to end through the public API on production shapes (warm cache, includes mapping):
A 6-layer stage sharing one
indice_keyruns in 0.596 ms.API style
Commits 3-5 port the kernels to the current
fxsurface per the kernel-code-cleanupskill (
fx.Index→fx.Int32,arith.*→ operators,vector.*→fx.Vector), −34lines net. Verified bit-identical against a golden capture and perf-neutral
(median-of-5: 1.00x / 1.02x / 1.02x).
Three sites are left legacy on purpose, each documented in-code:
scf.IfOpeverywhere._collect_assigned_varstreats the receiver of a methodcall inside a branch as branch state, so any
lut_.store(...)puts aGTensorintoresult_namesandscf_if_dispatchrejects it. Three formulations were tried cold;all fail identically. Fixing this means changing
kernels/common/tensor_shim.py,shared by four other kernels.
KEY_T(i32 or i64 by grid extent) andGTensor.loadreturns a raw value typed by that dtype, so operator compares don'tapply and
fx.Int32offsets take a differentbuffer_loadpath.at C_OUT=512 (median-of-7, range 184.8-186.4) — the same instruction-scheduling class
of exception Port GEMM/MoE/conv kernels to the layout API #913 documents for
mxfp_moegemm1 andfp8_4wave.Also note this kernel does not use
buffer_ops/create_buffer_resourcedirectly, so thelayout-API port in #913 does not apply to it; it goes through the shared
GTensorshim.Test plan
39 tests: generic shapes (K=3 and K=5, non-divisible channel counts, degenerate C=1),
[N,4]batched coords, explicit vs derivedspatial_shape, dense-vs-Z-Delta bitequality, the >2^31-cell grid, 32- vs 64-bit key selection, weight-cache invalidation and
eviction, and full-tap coverage of the compacted pair list.
Open questions before this leaves draft
~2.5e-3; across a 22-layer stack with normalisation and residuals it compounds to
~1.3e-2 (cosine similarity 0.99994, argmax agreement 98.6%) — but that is synthetic
weights, not a real mIoU comparison. This should gate any accuracy-sensitive use.
generated coordinates, not a real SemanticKITTI frame.
torch.sortand the compaction scan; the GPUkernels are pure FlyDSL. Replacing those host ops with hand-written kernels was tried
and measured 2-8x slower (torch's nonzero/cumsum are parallel scans; the
replacement serialised on atomics). Flagging in case the repo prefers a different
dependency posture.
🤖 Generated with Claude Code