[Kernel] fp8 conv3d: 8-wave GEMM pipeline + BIG_IN fix#860
Draft
jiacao-amd wants to merge 23 commits into
Draft
Conversation
Generalize the BF16 and FP8 implicit-GEMM conv3d kernels from a single hardcoded 128x128/2x4-wave shape to a compile-time-parameterized tile (TILE_M, TILE_N, WAVE_M, WAVE_N) and add a manual per-shape autotuner. - Replace the hand-scheduled 2x2-half BF16 pipeline (phase_*_prefetch + peeled prologue/main/epilogue, welded to QM_STEPS==2/QN_STEPS==1) with a generic double-buffered 2-stage loop over a flat acc[MI_M*MI_N] MFMA grid. - Generalize gather/commit/read/store to range_constexpr over MI_M/MI_N and per-thread LDG counts; add tile legality asserts (MFMA divisibility, LDG vectorization, LDS budget). - FP8: same flat-acc generalization, g2s halves -> row-block loop, Mfma16x16x128(MI_M, MI_N), flat_work_group_size from BLOCK_THREADS. - New kernels/conv/conv3d_autotune.py: benchmark a small candidate list with do_bench, cache the best tile per shape (in-memory + JSON disk), reusing flydsl.autotune. Enabled via autotune=True or FLYDSL_CONV3D_AUTOTUNE=1; default path stays (128,128,2,4) for zero behavior change. - TILE_K and filter size (kt/kh/kw) stay compile-time constants. - Add tile-sweep + autotune-cache tests and scripts/bench_conv3d_tiles.py.
Buffer load offsets are 32-bit, so a single buffer resource can address
at most ~2 GB. The existing BIG_IN / BIG path rebased per image (nbase =
m_offset // dhw), which is a no-op for n=1 and fails entirely when a
single image exceeds 2^31 elements.
Fix: per-block time-plane rebase for bf16 and fp8:
* conv3d_implicit_8wave (bf16):
- compile_transpose_ncdhw_ndhwc: BIG = n*c*s > 2^31 path adds per-block
(nb, c0, s0) rebase so read/write residuals are tile-local (<900M <2^31).
- compile_conv3d_implicit_8wave: BIG_IN path computes base_t (block's
first time-plane, shifted down by pt) and rebases x_rsrc so gather_a
residual (di*d + in_t - base_t)*h*w*c stays within int32 (<200M).
* conv3d_implicit_8wave_fp8:
- compile_pack_activation_ncdhw_bf16_to_ndhwc_fp8: same BIG per-block
(nb, c0, s0) rebase for read (BF16 input) and write (FP8 byte output).
- compile_conv3d_implicit_8wave_fp8: BIG_IN path adds
_make_fp8_buffer_tensor_from_addr helper and computes a rebased x_div
so g2s_a_block's g_elem residual stays within int32.
Verified on gfx950 (MI355X):
- 1x1024x240x160x90 conv (1,3,3) and (3,1,1): PASS (was crash/wrong)
- 1x1024x120x160x90, 1x2048x60x80x45: no regression
- All 14 bf16 unit tests pass; 6/7 fp8 tests pass (1 is pre-existing flaky)
For n==1 (no split-K), pack a lane's 4 contiguous, 8-byte-aligned output values into one vec4 bf16 buffer_store instead of 4 scalar stores. The epilogue VMEM-store was ~12% of kernel stalls on the short-reduction 3x1x1 path; this brings the 3x1x1 core (256x256x4x4 tile) from 3.18ms to 2.85ms (854->951 TFLOPS), beating MIOpen (2.92ms). Gated on n==1, not split-K, and dhw % 4 == 0; scalar fallback otherwise. Add conv2d_implicit / conv1d_implicit as thin wrappers that reshape to the degenerate 3D form (D=T=1, and H=R=1 for 1D) and reuse the same kernel and fast paths -- zero kernel duplication. Guard temporal_only_fast off when BIG_IN (>2^31 elems): the large-tensor rebase uses the generic n/t/h/w decode, which the temporal fast decode does not carry, so >2^31 3x1x1 inputs fall through to the correct generic path. Tests: add conv2d/conv1d vs torch coverage (26 total, all pass).
Previously the temporal-only fast decode was disabled under BIG_IN (>2^31 elements), falling back to the generic path. Instead, rebase the temporal fast-path address against the same x_base_elem origin the BIG_IN x_rsrc is built from, so the residual element offset stays within int32. This keeps the faster temporal decode for large 3x1x1 inputs while staying correct. Verified vs torch.nn.functional.conv3d on a >2^31-element input (1x640x240x160x90, 3x1x1): allclose, maxdiff 1.0 (matches the generic path). Adds a large_shape regression test.
…upport Port the BF16 conv3d optimizations to the FP8 kernel: vectorized (vec4) epilogue store for n==1, the temporal-only 3x1x1 fast decode path, 64-bit output addressing for >2^31-element outputs (BIG_OUT), and a split-K overflow guard. Add conv2d_implicit_fp8 / conv1d_implicit_fp8 thin wrappers mirroring the BF16 API, plus 2D/1D FP8 correctness tests. Tidy the BF16 kernel comments.
Make conv3d_implicit_8wave(_fp8) the single public entry that dispatches
1D/2D/3D by filter rank; move the rank-specific bodies to private
_conv{1,2,3}d_impl(_fp8) helpers so the 2D/1D reshape wrappers no longer
recurse through the public name. Merge _choose_splitk into _resolve_splitk
(behavior verified bit-identical over the full input space). Also drop
stale explanatory comments and trailing whitespace.
Co-Authored-By: Claude <noreply@anthropic.com>
Port hgemm_splitk's async-copy deep-pipeline recipe into the bf16 conv3d kernel. Replace the sync global->VGPR->LDS hop (gather_a/commit_a) with raw_ptr_buffer_load_lds direct global->LDS DMA, freeing the A-tile VGPRs and letting the software pipeline go from 2 to 4 stages. Padding is masked by OOB-routing: the buffer resources are rebuilt with the real num_records and invalid im2col taps are routed past the bounds so the hardware bounds check writes 0 to LDS (the DMA path has no VGPR step for a select mask). Gated by USE_ASYNC = not BIG_IN and X/W byte sizes <= 2^31; BIG_IN and >2^31 tensors keep the proven sync path. Measured on MI355X (gfx950), tile 256x256x4x4, controlled back-to-back A/B: c2048 1x3x3 M216k 842 -> 924 TFLOPS (+9.7%) c2048 3x3x3 M216k 848 -> 967 TFLOPS (+13.9%) 26/26 conv tests pass.
The buffer_load_lds voffset is unsigned 32-bit (verified), so no 64-bit soffset split is needed to reach past the old 2^31-byte limit. Extend the async path to two more regimes: - not-BIG_IN tensors up to ~4.29GB: raw resource with real num_records, padding taps routed to an OOB sentinel just under 2^32. - BIG_IN with n==1: reuse the existing per-block rebase (relative offsets <~0.3GB) and give the rebased resource a fixed 2GB num_records, between the max legal tap and the OOB sentinel, so padding still zeroes. n>1 BIG_IN and >~4.29GB weights keep the sync fallback (di can jump a whole batch past 2^32). All real video-VAE shapes now take the async path. Measured on MI355X (gfx950), tile 256x256x4x4: 512 [240,320,180] 1x3x3 728 -> 784 TFLOPS (+7.7%, 14GB) 1024 [240,160,90] 1x3x3 799 -> 887 TFLOPS (+11.0%, 7GB) 1024 [120,160,90] 1x3x3 804 -> 902 TFLOPS (+12.2%, 3.5GB) 2048 [60,80,45] 3x3x3 847 -> 977 TFLOPS (+15.4%) 26/26 conv tests pass; large shapes verified correct (allclose).
3x1x1 with stride=1 padding=(1,0,0) reduces to a GEMM after unfolding the temporal dimension: treat x as (N*C, 1, D, H*W), apply F.unfold(kernel=(3,1), padding=(1,0)) to get A=(N*D*H*W, C*3), then compute A @ W.T via hgemm_splitk. Weight ordering: weight[k,c,kt].reshape(k, c*3) matches A's column ordering where C varies fast (kt varies slow) at each output position. Routes to hgemm_splitk when A fits in the i32 byte range (<~4.29GB) and K is a multiple of 128/256 (TILE_N constraint). Falls through to the async im2col path for large shapes (A>4.29GB: 512/1024 240-frame channels). The only shape currently using hgemm: 2048[60,80,45] 3x1x1 (A=2.65GB). Larger shapes still use async im2col (same performance as before this patch). 26/26 conv tests pass; correctness verified numerically (max diff < 1e-6 fp32).
…76% regression) The unfold approach allocates a 2.65GB contiguous A matrix (M*3C) on the GPU, whose write cost far exceeds the im2col saving for M=216k. Measured: 834 TF (async im2col) -> 198 TF (unfold+hgemm), a 76% regression. The correct fix requires a FlyDSL kernel that scatters directly from NCDHW into GEMM without materializing the full A — deferred to future work. The 3x1x1 path now falls through to the existing async im2col kernel.
Removing the torch.matmul fast path made 1x1x1 8.7ms (78 TF) vs 3.4ms (201 TF) -- 2.6x slower -- because the async im2col path pays an NCDHW->NDHWC transpose (~5-6ms) that 1x1x1 never needs (x is already channel-major). Comment updated to explain the measurement and why FlyDSL GEMM is not used (i32 overflow at 14GB A, K=48 not multiple of 128).
- Rename conv3d_implicit_8wave -> conv3d_implicit (and _fp8, tests) since the 8-wave detail is no longer the defining trait. - The async global->LDS DMA pipeline is now the only path: replace the USE_ASYNC conditional with asserts (W_BYTES/X_BYTES < 2^32 sentinel; BIG_IN requires n==1) and delete the sync gather/commit helpers + their prologue/ main-loop/epilogue. - Drop the "async" qualifier from identifiers/comments (BIG_ASYNC_NR->BIG_IN_NR, _async_a/_async_b -> _load_a/_load_b) now that there is one pipeline. - Keep PIPE_STAGES=4 fixed: an A/B of 2/3/4 stages showed deeper is faster even for short-K 3x1x1 (memory-bound -> more dependent on deep prefetch), so the adaptive-shallow idea was dropped. 26/26 conv tests pass; VAE-shape benchmarks unchanged vs before the refactor.
Content of the rename from the previous commit: - The async global->LDS DMA pipeline is now the only path: replace the USE_ASYNC conditional with asserts (W_BYTES/X_BYTES < 2^32 sentinel; BIG_IN requires n==1) and delete the sync gather/commit helpers + prologue/loop/epilogue. - Drop the "async" qualifier from identifiers/comments now that there is one pipeline (BIG_ASYNC_NR->BIG_IN_NR, _async_a/_async_b -> _load_a/_load_b). - Keep PIPE_STAGES=4 fixed: A/B of 2/3/4 stages showed deeper is faster even for short-K 3x1x1 (memory-bound -> more dependent on deep prefetch). - Same cleanup applied to the fp8 kernel + both tests. 26/26 conv tests pass; VAE-shape benchmarks unchanged vs before the refactor.
Adds a new standalone fp8 conv3d kernel that ports the fp8_gemm_8wave
hand-scheduled 8-wave pipeline into implicit-GEMM conv3d:
- `kernels/conv/conv3d_implicit_fp8_gemm8w.py`: new kernel
`conv3d_implicit_fp8_gemm8w` (512 threads / 8 waves, 256x256x128 tile,
8-buffer LDS ping-pong). Accepts fp8 (E4M3FN) input natively; includes
a tiled NCDHW->NDHWC fp8 transpose replacing the slow torch.permute
pre-pass. Direct-store epilogue (no CShuffle). Supports BIG_IN shapes
(n*c*d*h*w > 2^31). Rigid constraints: crs%128==0, k%256==0.
- `tests/kernels/test_conv3d_implicit_fp8_gemm8w.py`: correctness tests
(6 cases: 3D, 2D, bias, partial-M) vs fp8-cast conv reference.
- `kernels/conv/conv3d_implicit_fp8.py`: fix BIG_IN rebase for the
general fp8 conv kernel. `_make_fp8_buffer_tensor_from_addr` used
`BitcastOp(f8_ptr_ty.ir_type, ...)` but `fx.PointerType` is an
`ir.Type` subclass with no `.ir_type` attribute; reconstructed using
`make_ptr` to match the `make_buffer_tensor` pattern.
Measured on MI355X (gfx950), kernel-only (pre-packed fp8 inputs):
1024->1024 [120,160,90] 1x3x3:
main fp8 conv kernel: 1287 TF (25.3 ms)
fp8 gemm8w kernel: 2125 TF (15.3 ms) (+1.65x vs main fp8 conv)
End-to-end vs MIOpen bf16 (1024->1024 [120,160,90]):
1x3x3: MIOpen 674 TF, gemm8w 1991 TF (2.95x)
3x1x1: MIOpen 574 TF, gemm8w 1379 TF (2.40x)
Co-Authored-By: Claude <noreply@anthropic.com>
Port fp8_gemm_4wave's 256-thread / 4-wave interleaved-cluster pipeline into conv3d, sibling of conv3d_implicit_fp8_gemm8w: same conv im2col A-loader, transpose, and direct-store epilogue, but the 4-wave GEMM core (2x2 wave grid, 256x256x128 tile, hand-scheduled interleaved cluster, XCD block-swizzle). Uses the SSA MFMA atom, NOT the GEMM's AGPR-pinned Mfma16x16x128AGPR: the 4-wave layout gives each of 256 threads the full 256-accumulator (256x256) fragment set, and the conv im2col address arithmetic adds enough VGPR pressure that the AGPR path's tied `=a,v,v,0` accumulators spill to scratch and corrupt output (verified via ISA: scratch_store a[...]). The 8-wave kernel avoids this by splitting the tile across 512 threads. Benchmark (MI355X/gfx950, kernel-only, pre-packed fp8): 4-wave is consistently 0.42-0.77x the 8-wave across shapes -- same root cause (heavy register spilling, 310 VGPR spills even on the SSA path). Both beat the general fp8 conv on most shapes. The 8-wave pipeline remains the right choice for conv. - kernels/conv/conv3d_implicit_fp8_gemm4w.py: new kernel + 1D/2D/3D dispatch - tests/kernels/test_conv3d_implicit_fp8_gemm4w.py: 6 correctness cases, all pass - scripts/bench_conv3d_fp8_4w_vs_8w.py: kernel-only 4w vs 8w vs general fp8 conv Co-Authored-By: Claude <noreply@anthropic.com>
- _resolve_splitk: remove the MAX_TILES_PER_SPLIT=54 cap that forced split-K on large-K shapes even when the (M,N) grid already filled the GPU (the cap's premise -- a long k-loop is bad -- is false; the fp8_gemm pipeline runs a full 72-iter k-loop as its fast path). The grid-underfill test is now authoritative. Also drop the npq<4096 gate so small underfilled grids can still split-K. Net: C1024 D120 1x3x3 auto sk 2->1, 827 -> 1284 TF (+55%). - gemm8w: add WGM grouped-M workgroup L2-swizzle (ported from the bf16 conv) + hoist the im2col output-spatial decomposition out of the K loop + skip compile-time-degenerate padding checks. C1024 D120 1x3x3 2228 -> 2310 TF.
The public conv3d_implicit_fp8 entry now dispatches to the hand-scheduled fp8_gemm 8-wave conv (conv3d_implicit_fp8_gemm8w, ~1.7x faster) whenever its rigid 256x256x128 constraints hold (k%256==0, crs%128==0, c%16==0, crs//128>=2). Unaligned shapes (and explicit tile=...) fall through to the general kernel, which keeps arbitrary-shape support + split-K + per-shape tile autotune. - gemm8w has a fixed tile, so its only autotune knob is WGM (workgroup L2-swizzle). The autotune path sweeps WGM_VALUES via the shared autotune_conv3d, reusing the cache framework. - Sync conv3d_autotune.py to the PR820 autotuner (schema v3 + (candidate, wgm) pair support + dry-run-before-bench), so both bf16 and fp8 share one autotuner. Named conv3d_autotune.py here; PR820 renames it to conv3d_implicit_autotune.py and the rename resolves on rebase after PR820 merges. - tests: add gemm8w dispatch + WGM-autotune coverage to the fp8 test only. Co-Authored-By: Claude <noreply@anthropic.com>
gemm8w previously required aligned shapes (k%256==0, crs%128==0, crs//128>=2). Generalize it to cover every shape the old general fp8 conv did, so it can fully replace that kernel: - N-partial (k % 256 != 0): grid_n now ceils; the tail n-tile's OOB k columns read 0 from the weight buffer and the epilogue already masks col >= k. - K-partial (crs % 128 != 0): the k-loop runs over crs_pad (crs rounded up to 128); the tail tile's k in [crs, crs_pad) is zeroed on A (k_col < crs im2col mask) and on B (host zero-pads the weight's crs dim to crs_pad). - tiny-K (crs <= 128): padded up to 2 full tiles so the 3-stage pipeline's K_ITERS>=2 precondition holds (the 2nd tile reads zeros). Verified rel ~1e-5 across aligned / N-partial / K-partial / tiny-K / tiny-N (k=32) shapes -- strictly better than the old kernel (which was 0.03-0.31 on the c96/c384 partial cases). Dispatch in conv3d_implicit_fp8 now routes all c%16==0 shapes to gemm8w.
The fp8_gemm 8-wave conv (previously conv3d_implicit_fp8_gemm8w) is now the sole conv3d_implicit_fp8 implementation. It was generalized to cover every shape the old general kernel did (N-partial, K-partial, tiny-K), is ~1.7x faster on aligned shapes, and is more accurate on partial shapes (rel ~1e-5 vs the old kernel's 0.03-0.31 on c96/c384), so the old double-buffered kernel and the dispatch/ fallback layer are removed. - conv3d_implicit_fp8_gemm8w.py -> conv3d_implicit_fp8.py (old file deleted); the three shared host helpers (_normalize_3, _prep_weight_fp8, _make_fp8_buffer_tensor_from_addr) are inlined; all "gemm8w" names dropped (compile_conv3d_implicit_fp8 / conv3d_implicit_fp8_kernel / conv3d_implicit_fp8). - Public entry keeps the same name and signature; wgm=<int> forces the WGM L2-swizzle grouping, autotune=True sweeps WGM_VALUES via the shared autotuner. - tests: rewrite the tile-config sweep as a WGM-config sweep, add WGM-autotune coverage, drop the obsolete test_conv3d_implicit_fp8_gemm8w.py (its coverage is folded into the fp8 test). Only the fp8 test changed. - conv3d_implicit_fp8_gemm4w.py imports the helpers from the merged module now. Co-Authored-By: Claude <noreply@anthropic.com>
The 4-wave port was only ever a comparison experiment -- it is 0.42-0.77x the 8-wave kernel across shapes (the 256-accumulator-per-thread layout spills under the conv im2col address pressure), so it has no production value. Remove the kernel and its test; the 8-wave pipeline (now conv3d_implicit_fp8) is the kernel. Co-Authored-By: Claude <noreply@anthropic.com>
It only compared the removed 4-wave kernel against the 8-wave one; no longer relevant now that the 4-wave port is gone. Co-Authored-By: Claude <noreply@anthropic.com>
Two independent >2^31 addressing bugs made large activations produce NaNs and GPU memory faults (e.g. 512->512 [240,320,180], 7 GB in / 14 GB out): 1. OOB im2col sentinel too small. Padding / out-of-bounds taps used a sentinel element offset of 0x7FFFFF80, but the rebased BIG_IN buffer has num_records = 2 GB (0x80000000 elements at 1 B/fp8). Since the offset is unsigned, 0x7FFFFF80 is IN-BOUNDS, so padding taps read live data -> NaNs on >2 GB activations. Use 0xF0000000 (~4 G), which exceeds the 2 GB rebased resource and any non-BIG_IN buffer. 2. fp8 transpose i32 byte-offset overflow. The NCDHW->NDHWC transpose computed ss*c (up to ~7e9) and cc*s in the i32 buffer offset, overflowing for >2 GB tensors and scrambling / faulting the store. Added an i64 raw-pointer load/store path (TR_BIG) for >2 GB tensors. Verified finite + rel ~1e-5 on BIG_IN (3.5 GB), BIG_IN+BIG_OUT (7/14 GB), and normal shapes. Adds a `large_shape` BIG_IN regression test. Co-Authored-By: Claude <noreply@anthropic.com>
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
Replaces the general FP8 (E4M3FN) implicit-GEMM conv3d kernel with a port of the
fp8_gemm_8wavehand-scheduled 8-wave matmul pipeline.conv3d_implicit_fp8keeps the same public entry point and signature; the old double-buffered kernel
is removed.
Kernel:
kernels/conv/conv3d_implicit_fp8.pyDesign: Ports
fp8_gemm_8wave's pipeline — 512 threads / 8 waves, 256×256×128 tile, 8-buffer LDS half-block ping-pong (S2RLoader/G2SLoader),Mfma16x16x128, direct-store epilogue. The only conv-specific component is the A-operand loader (conv_a_g2s): instead of a contiguous GEMM read, it computes the im2col activation address per LDS chunk and deposits it into the GEMM's swizzled half-block LDS layout so the reusedS2RLoaderreads correct fragments. The B (weight) operand is a plain KTRSC(k, crs)matrix, soG2SLoaderis reused verbatim.Shape coverage: only requires
c % 16 == 0. Arbitrary shapes are handled — N-partial (knot a multiple of 256, grid ceils + epilogue mask), K-partial (crsnot a multiple of 128, weight zero-padded tocrs_pad+ im2colk_col < crsmask), and tiny-K (crs <= 128, padded up to 2 tiles).Extras:
torch.permutepre-pass).wgm=<int>forces it;autotune=True/FLYDSL_CONV3D_AUTOTUNE=1sweepsWGM_VALUESand caches the winner per shape).Performance (MI355X / gfx950, kernel-only, pre-packed fp8):
End-to-end vs MIOpen bf16 (incl. fp8 transpose):
Tests
tests/kernels/test_conv3d_implicit_fp8.py— correctness vs fp8-cast conv reference across aligned / N-partial / K-partial / tiny-K / tiny-N shapes, plus WGM-config and WGM-autotune coverage. All pass.Test plan
python3 -m pytest tests/kernels/test_conv3d_implicit_fp8.py -vpython3 -m pytest tests/kernels/test_conv3d_implicit.py -v(bf16 regression)bash scripts/check_python_style.sh🤖 Generated with Claude Code