Skip to content

fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32::MAX (#63) - #64

Merged
isPANN merged 2 commits into
mainfrom
fix/cuda-64bit-indexing-issue-63
Jun 16, 2026
Merged

fix(cuda): 64-bit global indexing so tropical GEMM addresses past i32::MAX (#63)#64
isPANN merged 2 commits into
mainfrom
fix/cuda-64bit-indexing-issue-63

Conversation

@isPANN

@isPANN isPANN commented Jun 16, 2026

Copy link
Copy Markdown
Collaborator

Fixes #63.

Problem

The strided-batched tropical GEMM panicked on large operands that fit in VRAM:

batched GEMM output C stride 2147483648 exceeds i32::MAX; matrix too large for the strided-batched kernel

2147483648 = 2³¹ is one past i32::MAX. This is a kernel addressing ceiling, not OOM — the same family as the gridDim.z cap fixed in #62. It blocked GPU tropical contraction exactly in the large-instance regime where the GPU backend matters (e.g. miso's max-plus contraction at sc_target=32, whose output C stride m·n = 2³¹).

Three 32-bit sites caused it, all on the global-addressing path:

  1. stride kernel params were int,
  2. the batch base used (int)blockIdx.z * stride,
  3. the within-matrix offset OFFSET_COL(row,col,ld) = col*ld+row overflowed int once a single matrix held ≥ 2³¹ elements (so widening only the batch stride is insufficient).

Fix

Widen the global-addressing path to 64-bit; leave the hot inner loop (shared-memory As/Bs + accumulators, all tile-sized) at i32 so there's no cost where it counts.

  • OFFSET_COL64 = (long long)(col)*(ld)+(row) — the cast wraps the multiply (a single IMAD.WIDE on the device), not the already-overflowed sum. Used for every global load/store in the shared tile macros, so batched, single-matrix, and argmax forward kernels are all fixed at once.
  • Batched stride params int → long long; batch base (long long)blockIdx.z * stride.
  • Backward kernels: total/idx → 64-bit; scatter offsets i + (long long)k*M, k + (long long)j*K.
  • Rust: stride_to_i32 → stride_to_i64 with i64 kernel args (ABI stays in sync: i64long long, 8 bytes; cudarc passes each arg by pointer so only size must match). Backward grid_size now computed in u64 via div_ceiltotal as u32 truncated when m·n ≥ 2³². Public signatures unchanged, so downstream crates are unaffected.

Tests

Adds three #[ignore] regression tests crossing 2³¹ (46341² ≈ 2.147e9): batched-forward (the exact repro), single-matrix offset, and backward scatter. They sample device-side / zero bulk inputs on-device to keep host memory small.

cargo test -p tropical-gemm-cuda -- --ignored

Validation (real hardware)

Run on a 6× A800 80GB box, CUDA 12.1 (CUDARC_CUDA_VERSION=12010):

  • All three ignored regression tests pass — the corner cell at offset m·n−1 = 2,147,488,280 (well past i32::MAX) computes the correct value, exactly what i32 indexing would corrupt.
  • Full normal-size suite: 57 passed, 0 failed — no regression from the shared-macro widening.

CI skips CUDA, so the large-allocation tests stay #[ignore]d (they need ≥ 8.6 GB / 16 GB VRAM).

🤖 Generated with Claude Code

…::MAX (#63)

The strided-batched kernel rejected any operand stride >= i32::MAX, panicking on
large matrices that fit in VRAM (e.g. miso's max-plus contraction at sc_target=32,
where the output C stride m*n = 2^31). The ceiling was the kernels' 32-bit
addressing, not memory: stride params were `int`, the batch base used
`(int)blockIdx.z * stride`, and the within-matrix offset `OFFSET_COL(row,col,ld)
= col*ld+row` overflowed `int` once a single matrix held >= 2^31 elements.

Widen the GLOBAL-addressing path to 64-bit while leaving the hot inner loop
(shared-memory As/Bs + accumulators, all tile-sized) at i32:

- Add OFFSET_COL64 (`(long long)(col)*(ld)+(row)` -- cast wraps the multiply, not
  the finished sum) and use it for every global load/store in the shared
  LOAD/STORE tile macros. This fixes batched, single-matrix, and argmax forward
  kernels at once since they share those macros.
- Batched stride params int -> long long; batch base `(long long)blockIdx.z * stride`.
- Backward kernels: `total`/`idx` -> 64-bit, scatter offsets `i + (long long)k*M`
  and `k + (long long)j*K`.
- Rust: stride_to_i32 -> stride_to_i64 with i64 kernel args (ABI stays in sync:
  i64 <-> long long, 8 bytes); backward grid_size computed in u64 via div_ceil
  (`total as u32` truncated when m*n >= 2^32). Public signatures unchanged.

Add three #[ignore] regression tests crossing 2^31 (batched forward, single
matrix, backward scatter); they sample device-side / zero bulk inputs on-device
to keep host memory small. Validated on a 6x A800 80GB box (CUDA 12.1): all three
pass and the full 57-test normal-size suite stays green (no regression from the
shared-macro widening).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@codecov

codecov Bot commented Jun 16, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.29%. Comparing base (ec04a4a) to head (0467db3).

Additional details and impacted files
@@           Coverage Diff           @@
##             main      #64   +/-   ##
=======================================
  Coverage   96.29%   96.29%           
=======================================
  Files          19       19           
  Lines         918      918           
=======================================
  Hits          884      884           
  Misses         34       34           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Extract the duplicated column-major operand build and per-cell sample/verify
shared by the two forward regression tests into `build_big_operands` /
`assert_big_maxplus_cells` (+ `big_a_val`/`big_b_val`), and replace the three
runtime `m*n > i32::MAX` asserts with one compile-time `const _` assert on
BIG_DIM. Behavior-preserving test-only cleanup (the kernel fix is unchanged);
net -8 lines, ~50 lines of copy-paste removed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@isPANN
isPANN merged commit 210b276 into main Jun 16, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CUDA strided-batched tropical GEMM rejects stride ≥ i32::MAX (kernel indexing is 32-bit) — fails on large matrices that fit in VRAM

1 participant