rnnt_loss: per-sequence logits offset is computed in int and overflows for large-vocabulary batches

[happyarts]

Summary

rnnt_loss slices its (B, T, U, D) logits tensor with an int offset. Once B * T * U * D exceeds 2^31 the offset wraps for the later sequences in the batch, so the kernels read and write outside the tensor. Nothing catches it — the result is silent corruption rather than an error.

This is reachable with a normal large-vocabulary training step: B=32, T=400, U=80, D=4096 is 4.19e9 elements, and the offset for b = 16 already lands at 2.10e9, just past INT_MAX. In fp16 (which compute() accepts, cpu/compute.cpp:37-38) that is 8.4 GB of logits plus another 8.4 GB for the empty_like gradients — 16.8 GB, comfortable on a 40 GB accelerator.

Found by code audit while working on #4208 / #4209 (the CPU forced_align 32-bit overflow); same bug class. Analysis only — I have not reproduced it at runtime, as I do not have a machine that size.

The line

src/libtorchaudio/rnnt/cpu/cpu_kernels.h:139-142, where B, maxT, maxU and D are all const int& bound to Options members:

for (int b = 0; b < B; ++b) {
  seqLogits.push_back(
      TensorView<const DTYPE>({maxT, maxU, D}, logits + b * maxT * maxU * D));

b * maxT * maxU * D is evaluated entirely in int. The dimensions come straight from the tensor (cpu/compute.cpp:91-93: maxSrcLen_ = logits.size(1), maxTgtLen_ = logits.size(2), numTargets_ = logits.size(3)), each of which is fine on its own — it is the product that overflows. Signed overflow is UB; in practice the pointer lands far before the buffer and ComputeLogProbsOneSequence then reads and writes there.

compute() validates dtype, device, contiguity, dims and length agreement, but nothing bounds the product.

Related, lower priority

• Workspace size accounting is int throughout, and the allocation and its own guard derive from the same wrapped value, so the check cannot catch the overflow: options.h:66-68 (BTU()), workspace.h:30-35, then cpu/compute.cpp:104-110 allocates from ComputeSizeFromOptions(options) while workspace.h:39-40 asserts needed_size <= size against it. Both sides wrap identically. This needs B * T * U (no D) past 2^31, which is far less reachable than the case above — but if it ever happens it produces an undersized buffer that passes validation.
• Indexer4D in rnnt/cpu/kernel_utils.h:45-63 and rnnt/gpu/kernel_utils.h is dead code — grep finds no instantiation anywhere in src/. Only Indexer2D/Indexer3D are used, and those index (B, T, U) without D, so they need the same unreachable magnitude as the workspace. Worth deleting rather than widening.

Suggested fix

Compute the offset in 64-bit — the minimal change is at the slice site:

const int64_t seqStride = static_cast<int64_t>(maxT) * maxU * D;
... TensorView<const DTYPE>({maxT, maxU, D}, logits + b * seqStride));

TensorView's own indexing would need the same treatment to be fully safe for a single sequence over 2^31 elements, though that is a much larger tensor. Alternatively, or in the meantime, a guard in compute() computing B * T * U * D in int64_t and failing with a clear message would turn silent corruption into a loud error. Happy to send a PR for whichever you prefer.

Versions

Current main (4e3e282). The arithmetic predates the stable-ABI migration.