bug(cuda): int overflow in weight-upload/dequant byte-size products for >2 GiB tensors

[jamesburton]

Problem

While auditing the CPU backend for the #429 bug class across the other backends, I found the same shape on the CUDA weight-upload/dequant path. Two sites, plus one consistency case.

The native CUDA kernels are clean — every weight-row base in native/kernels/ is already (size_t)row * blocks_per_row * BYTES, and I confirmed the widening survives into the checked-in native/ptx/*.ptx (mul.wide.s32 present at those sites). These are all managed-side.

1. CudaKernels.cs:607 — F16 passthrough copy

public void LaunchDequantToF16(nint src, QuantizationType srcDtype,
                                 nint dst, int totalElements, nint stream)
{
    case QuantizationType.F16:
        CudaDriverApi.cuMemcpyDtoD_v2(dst, src, (nuint)(totalElements * 2)).ThrowOnError();

totalElements * 2 is int * int. It wraps negative above 1,073,741,823 elements (a 2 GiB F16 tensor), and (nuint) then sign-extends that to a ~16 EiB byte count for cuMemcpyDtoD_v2.

Reachability — Llama 3.1 405B LM head (vocab 128256 × hidden 16384) = 2,101,346,304 elements, which wraps. Worth noting Llama 3.1 70B's LM head (128256 × 8192 = 1,050,673,152) sits only 23M elements under the limit — a 2.2% margin. Anything slightly wider than 70B trips it.

2. CudaWeights.cs:245 — element count in UploadAndDequant

int totalElements = outputDim * inputDim;

Wraps at 2^31 elements. The 405B LM head fits by ~2%, so this needs something wider still — but it is worth fixing because every byte size derived from it is already widened ((long)totalElements * sizeof(ushort) etc.), leaving the element count as the sole narrow product. That asymmetry is the kind of thing that reads as intentional and gets copied.

Note this one cannot simply be widened: totalElements is passed as int to LaunchConvertF32ToF16 and LaunchDequantToF16, so a long count could not be honoured without changing those signatures. Failing loudly is the proportionate fix; genuinely supporting >2^31-element tensors is a larger piece of work.

3. CudaTransformerModel.cs:290 — consistency only

nint lastHidden = _state.HiddenState + (nint)((seqLen - 1) * hiddenSize * h);

An activation offset, so self-limiting — wrapping needs a 2 GB single-shot activation buffer. Flagging it only because sibling lines in the same file already cast (long), so the inconsistency invites the wrong conclusion about which form is correct.

What I checked and ruled out

• All 63 weight-row bases across 57 .cu files — already (size_t).
• MoE expert offsets, managed and native. The managed ones are already long, which matters: DeepSeek-V3's 256 experts × 8.26 MB = 2.1 GB would have wrapped an int. The native MoE kernels take a per-expert pointer array, so no expert * stride arithmetic exists in device code at all.
• KV-cache offsets (CudaKvCache, CudaPagedKvCache, CudaQuantizedKvCache, CudaKvBlockPool) — safe, but only because rowBytes/blockBytes are declared long several lines above the use. A grep for (long) on the offset line itself gives a false positive here.
• CudaDriverApi P/Invoke signatures — all sizes nuint, all pointers nint. No int-vs-size_t truncation at the boundary.

Separate finding worth recording

long is 32-bit in the device code this repo ships. native/build.ps1 drives nvcc on Windows, which follows the MSVC LLP64 model, so sizeof(long) == 4 in .cu files. There are ~44 sites in native/kernels/ using bare long as though it widened (dequant_i2_s.cu, dequant_pq2_0.cu, pq2_0_repack.cu, pq2_0_gemv.cu, elementwise_f32.cu, gemma4_f32.cu).

I bounded every one and none is reachable — the I2_S/PQ2_0 cases need n*k >= 2.7e11 elements — so I am not proposing any change. But the same source compiled on Linux (LP64) would widen, so a kernel can be correct upstream and silently 32-bit in the shipped PTX. size_t / int64_t / long long are the safe choices here.

Acceptance criteria

• Widen site 1; make site 2 fail loudly rather than wrap; widen site 3 for consistency with its siblings.
• No change to the native kernels — they are already correct.

Caveat on verification

I have no NVIDIA hardware to hand (my dev box is AMD), so these are unexercised at the overflowing scale. Site 1 needs a >2 GiB F16 tensor, which is out of reach on the 12 GB and 16 GB cards I can borrow. The fixes are type-level and the arithmetic is shown above, but I want to be explicit that no CUDA kernel was executed.

Priority

Low-to-moderate. Site 1 is the one that matters: it is reachable on any model above ~70B, and the failure mode is a wild cuMemcpyDtoD_v2 size rather than a clean error.