CUDA Cast kernel crashes with illegal memory access on tensors with >2^31 elements (int32 overflow) — same family as #28107

[TimPietrusky]

Description

The CUDA Cast kernel crashes with cudaErrorIllegalAddress (error 700) when the input tensor has more than 2^31 (~2.15 billion) elements.

This is the same class of bug that was fixed in Gather by #28108 (response to #28107) — the kernel uses CUDA_LONG = int32_t for its element index, which wraps to negative once the element count crosses INT32_MAX. The bound check if (id < N) then passes for invalid offsets and the kernel reads / writes out of bounds.

Root cause

In onnxruntime/core/providers/cuda/tensor/cast_op.cu, CudaCastStd (line ~244) and CudaCastSat (line ~377) both compute the per-thread element id with the CUDA_LONG-based macro from cu_inc/common.cuh:677:

#define CUDA_LONG int32_t

The launch in Impl_Cast(..., size_t count) passes count through static_cast<int>(num_of_elements) — silent truncation when count >= 2^31.

The exact same pattern is also present in unary_elementwise_ops_impl.cu (used by Cast8, IsInf, IsNan) and in many other elementwise CUDA kernels (Add, Sub, Mul, etc.). They are all candidates for the same fix.

Reproducer

Any causal LM exported via optimum-cli export onnx whose vocab_size × max_seq_length > 2^31.

Concrete: LiquidAI/LFM2.5-1.2B-Instruct-ONNX (vocab 65 536, declared max context 128 000) hits the cliff at exactly seq_length = 32 768:

import numpy as np, onnxruntime as ort
sess = ort.InferenceSession(
    "model_q4f16.onnx",
    providers=[("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider"],
)
P = 32_768  # exactly 2^15
inputs = {
    "input_ids":      np.random.randint(0, 65000, size=(1, P), dtype=np.int64),
    "attention_mask": np.ones((1, P), dtype=np.int64),
    # past_key_values + past_conv inputs zero-initialised…
}
sess.run(["logits"], inputs)        # works
P = 32_769                          # one token over → 2^31 + 65536 elements
inputs["input_ids"] = np.random.randint(0, 65000, size=(1, P), dtype=np.int64)
inputs["attention_mask"] = np.ones((1, P), dtype=np.int64)
sess.run(["logits"], inputs)        # crashes:
# CUDA failure 700: an illegal memory access was encountered
# at .../tensor/cast_op.cu  (logits/Cast node)

Run with CUDA_LAUNCH_BLOCKING=1 to confirm the failing op is /model/graph_outputs/logits/Cast.

The same model also runs fine on CPU EP at seq_length = 32 769, which confirms it's a CUDA EP kernel issue rather than a model export issue.

Why this is hot today

Every causal LM exported via optimum-cli export onnx defaults to emitting full-sequence logits (optimum-onnx/convert.py hardcodes logits_to_keep = 0). Combined with modern vocab sizes (Llama 3.2 = 128 256, Qwen 2.5 = 151 936, LFM2.5 = 65 536) this puts the cliff somewhere between S = 16 K and S = 32 K for nearly every long-context-capable HF ONNX LLM running on ORT CUDA EP.

So this is not a niche corner case — it's the silent reason long-context inference of any standard HF causal-LM ONNX export breaks on ORT CUDA EP today.

Proposed fix

Mirror #28108's pattern for cast_op.cu:

• Switch the per-thread index in CudaCastStd / CudaCastSat from CUDA_LONG to int64_t (or size_t) for both the loop counter and the bound check.
• Pass num_of_elements through as size_t end-to-end without truncating to int.

Same change should be considered for the rest of unary_elementwise_ops_impl.cu and the binary elementwise kernels — they share the same idiom and the same latent bug.

Workaround for users hitting this today

Add a Slice node before the LM head matmul that takes only the last token of the hidden state. Drops the logits tensor from [B, S, V] to [B, 1, V], well below the int32 limit, and is what HF transformers / vLLM / TGI / llama.cpp do anyway under the name logits_to_keep=1.

cc @justinchuby (you fixed the Gather twin in #28108 — this is the same family).

[justinchuby]

Feels like taking only the last token of the hidden state is the right thing to do? (More efficient?) Though we do want to fix the cast limitations.