fix(cuda): chunk batched tropical GEMM over the gridDim.z 65535 cap - #61
Merged
isPANN merged 3 commits intoJun 14, 2026
Merged
Conversation
The strided-batched kernels map the batch dimension to blockIdx.z, whose grid extent CUDA caps at 65535 on all compute capabilities. A contraction whose batch exceeds that launched with grid.z > 65535 and failed at launch with CUDA_ERROR_INVALID_VALUE. Large tensor-network contractions (space complexity >= ~30) routinely exceed it, so the failure surfaced as an intermittent panic (the batch/m/n split depends on the contraction order). Split batched launches into chunks of at most 65535 batch elements, advancing each operand's base by the chunk start so blockIdx.z stays in [0, chunk). Fixes both the value path (launch_kernel_batched_impl, used by all four scalar-type macro impls) and the argmax path (launch_gemm_external_batched_with_argmax_f32). No kernel (.cu) change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## perf/forward-batched-gemm #61 +/- ##
============================================================
Coverage ? 96.29%
============================================================
Files ? 19
Lines ? 918
Branches ? 0
============================================================
Hits ? 884
Misses ? 34
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
This PR fixes CUDA launch failures for very large strided-batched tropical GEMM workloads by chunking batched kernel launches so the batch dimension mapped to blockIdx.z never exceeds CUDA’s gridDim.z limit (65535). This prevents intermittent CUDA_ERROR_INVALID_VALUE failures when batch > 65535, particularly in large tensor-network contractions.
Changes:
- Introduces a
MAX_GRID_DIM_Zcap and chunkslaunch_kernel_batched_impllaunches, using cudarc slice views to offset A/B/C bases per chunk. - Chunks the external-pointer batched argmax path (
launch_gemm_external_batched_with_argmax_f32) by advancing raw base pointers per chunk. - Minor refactoring around grid calculation / lifetime management to support per-chunk launches.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+159
to
+163
| let stride_a = (m * k) as i32; | ||
| let stride_b = (k * n) as i32; | ||
| let stride_c = (m * n) as i32; | ||
| // Per-batch element extents, used to offset each chunk's operand base. | ||
| let (sa, sb, sc) = (m * k, k * n, m * n); |
Comment on lines
+753
to
+755
| let stride_a_i32 = stride_a as i32; | ||
| let stride_b_i32 = stride_b as i32; | ||
| let stride_c_i32 = stride_c as i32; |
| // launch builder. | ||
| let b_ptr: u64 = b.device_ptr(); // B becomes "A" in kernel | ||
| let a_ptr: u64 = a.device_ptr(); // A becomes "B" in kernel | ||
| // Per-batch element strides (A/B/C/argmax are all rows*cols per batch). |
Address Copilot review on PR #61. The chunked batched launches derive operand base offsets from the usize per-batch stride, while the kernels read the stride as i32. An `as i32` cast that wrapped would desynchronise the two and silently corrupt addressing. Add a shared `stride_to_i32` helper that converts fallibly and returns DimensionMismatch on overflow, and apply it in both the value path (launch_kernel_batched_impl) and the external argmax path (launch_gemm_external_batched_with_argmax_f32). The kernel signature is i32, so a stride past i32::MAX is unrepresentable on-device regardless; fail loudly at the boundary. Also correct the argmax-path stride comment: external A/B carry the (possibly padded) DLPack stride, not rows*cols. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Quality cleanup of the strided-batched GEMM work (no behavior change):
kernels/tropical_gemm.cu
- Factor the tiled max-/min-plus / max-mul inner kernel into one shared
TROPICAL_GEMM_BODY macro. It was previously copy-pasted 8x (4 scalar
types x {single-matrix, batched}); the 8 kernel families are now thin
wrappers that supply only the signature and operand base expressions
(`A/B/C` vs `A + blockIdx.z*strideA`). Net -505 lines.
- Add add_f32/mul_f32/add_f64/mul_f64 helpers so the body takes MUL_FN in
function form for every type (symmetric with the int helpers) instead of
special-casing the float `+`/`*` operator.
src/kernels.rs
- launch_gemm_external_batched_with_argmax_f32: reuse CudaContext::
grid_dims_f32 instead of re-deriving the tile count with a hardcoded 64
(which would silently desync if the block size changed).
- launch_kernel_batched_impl: collapse the `grid: (u32,u32,u32)` parameter
to a single `grid_xy: u32`. The caller-passed z (`batch as u32`) was dead
(recomputed per chunk as the chunk size) and y was always 1.
Verified on A800 (sm_80, CUDA 12.1): 57/57 lib tests pass. PTX/SASS diff of
all 43 kernels vs the pre-refactor source — 31 non-batched kernels are
byte-identical; the 12 batched kernels have identical instruction histograms
and identical register/smem usage (only benign prologue scheduling differs).
No correctness or performance regression.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
isPANN
added a commit
that referenced
this pull request
Jun 14, 2026
…e batch) (#62) * feat(cuda): forward strided-batched tropical GEMM (one launch, blockIdx.z=batch) Adds TROPICAL_GEMM_BATCHED_{F32,F64,I32,I64} forward kernels (no argmax), mirroring the single-matrix forward macros plus blockIdx.z batch indexing and per-batch strides, and wires them through CudaKernel::launch_gemm_batched + launch_kernel_batched_impl (non-swapped column-major convention, raw CudaSlice operands so callers feed already-contiguous device buffers with no copy). Lets the omeinsum device path run a whole node's batch in ONE launch instead of a host-side per-slice clone_dtod loop, and the output may be allocated uninitialized since the kernel fully writes every element. Correctness test test_tropical_gemm_batched_matches_single checks the batched kernel against the trusted single-matrix path + a CPU max-plus reference, incl. a non-block-aligned shape (edge tiles). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(cuda): chunk batched tropical GEMM over the gridDim.z 65535 cap (#61) * fix(cuda): chunk batched tropical GEMM over gridDim.z 65535 cap The strided-batched kernels map the batch dimension to blockIdx.z, whose grid extent CUDA caps at 65535 on all compute capabilities. A contraction whose batch exceeds that launched with grid.z > 65535 and failed at launch with CUDA_ERROR_INVALID_VALUE. Large tensor-network contractions (space complexity >= ~30) routinely exceed it, so the failure surfaced as an intermittent panic (the batch/m/n split depends on the contraction order). Split batched launches into chunks of at most 65535 batch elements, advancing each operand's base by the chunk start so blockIdx.z stays in [0, chunk). Fixes both the value path (launch_kernel_batched_impl, used by all four scalar-type macro impls) and the argmax path (launch_gemm_external_batched_with_argmax_f32). No kernel (.cu) change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * cuda: reject batched strides past i32::MAX instead of truncating Address Copilot review on PR #61. The chunked batched launches derive operand base offsets from the usize per-batch stride, while the kernels read the stride as i32. An `as i32` cast that wrapped would desynchronise the two and silently corrupt addressing. Add a shared `stride_to_i32` helper that converts fallibly and returns DimensionMismatch on overflow, and apply it in both the value path (launch_kernel_batched_impl) and the external argmax path (launch_gemm_external_batched_with_argmax_f32). The kernel signature is i32, so a stride past i32::MAX is unrepresentable on-device regardless; fail loudly at the boundary. Also correct the argmax-path stride comment: external A/B carry the (possibly padded) DLPack stride, not rows*cols. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor(cuda): dedupe GEMM kernel macros and batched launch wiring Quality cleanup of the strided-batched GEMM work (no behavior change): kernels/tropical_gemm.cu - Factor the tiled max-/min-plus / max-mul inner kernel into one shared TROPICAL_GEMM_BODY macro. It was previously copy-pasted 8x (4 scalar types x {single-matrix, batched}); the 8 kernel families are now thin wrappers that supply only the signature and operand base expressions (`A/B/C` vs `A + blockIdx.z*strideA`). Net -505 lines. - Add add_f32/mul_f32/add_f64/mul_f64 helpers so the body takes MUL_FN in function form for every type (symmetric with the int helpers) instead of special-casing the float `+`/`*` operator. src/kernels.rs - launch_gemm_external_batched_with_argmax_f32: reuse CudaContext:: grid_dims_f32 instead of re-deriving the tile count with a hardcoded 64 (which would silently desync if the block size changed). - launch_kernel_batched_impl: collapse the `grid: (u32,u32,u32)` parameter to a single `grid_xy: u32`. The caller-passed z (`batch as u32`) was dead (recomputed per chunk as the chunk size) and y was always 1. Verified on A800 (sm_80, CUDA 12.1): 57/57 lib tests pass. PTX/SASS diff of all 43 kernels vs the pre-refactor source — 31 non-batched kernels are byte-identical; the 12 batched kernels have identical instruction histograms and identical register/smem usage (only benign prologue scheduling differs). No correctness or performance regression. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <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
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.
Problem
The strided-batched tropical GEMM kernels map the batch dimension to
blockIdx.z(tropical_gemm.cu: "Uses blockIdx.z for batch index"). CUDA capsgridDim.zat 65535 on every compute capability. A batched launch withbatch > 65535therefore fails at launch time withCUDA_ERROR_INVALID_VALUE.Large tensor-network contractions (space complexity ≳ 30, e.g. directly contracting a high-treewidth MIS component on an 80 GB GPU instead of slicing) routinely produce a batch dimension above 65535. Because the exact
batch/m/nsplit depends on the contraction order, the failure showed up as an intermittent panic — the same instance would sometimes succeed and sometimes crash:Fix
Split each batched launch into chunks of at most 65535 batch elements, advancing every operand's base by the chunk start so
blockIdx.z ∈ [0, chunk)continues to index the correct batch slice. No.cukernel change — only the host-side launch loop.Covers both batched entry points:
launch_kernel_batched_impl— the value path, shared by all four scalar-type (f32/f64/i32/i64)CudaKernelmacro impls. Chunks via cudarc sub-slice views.launch_gemm_external_batched_with_argmax_f32— the argmax (config-recovery) path. Chunks via base-pointer offsets.Verification
Reproduced the crash with a high-
sc_targetGPU contraction (omeinsum → miso MIS solve, reg-5 source, A800 80 GB), where it was intermittent at sc≥30. With this patch:fallback-latest) and builds with nvcc 12.1.🤖 Generated with Claude Code