Fix CUDA illegal memory access crash in TopK when inputs contain NaN

[jiafatom]

Problem

When vision models produce all-NaN logits (e.g., due to corrupt or edge-case images triggering incorrect pixel values in the image preprocessing pipeline), the CUDA TopK kernels crash with an illegal memory access error that corrupts GPU state and kills all subsequent inferences.

Root Cause

• TopK_Pair.p is initialized to INT_MAX for tie-breaking (smaller index wins ties)
• NaN comparisons (NaN > x) always return false, so no Insert succeeds
• p stays at INT_MAX, causing scores_in[INT_MAX] → out-of-bounds write
• This corrupts GPU state, crashing the current and all subsequent inferences

Repro

• Model: Qwen2.5-VL-3B-Instruct (ONNX, CUDA EP)
• Image: 1500×1204 PNG from AI2D dataset (sample 926)
• The ort-extensions bicubic resize (1500→1512 width) produces incorrect pixel values for this image
• Incorrect pixels → NaN in vision encoder (feature index 2008) → propagates to 100% NaN logits → TopK crash

Fix

Guard all scores_in writes with p < vocab_size check, and clamp output indices to 0 when no valid element is found:

1. GetTop1Kernel: Clamp indices_out to 0 if p >= vocab_size
2. GetTopKKernel: Clamp output index to 0, guard scores_in write with bounds check
3. GetTopKKernelDistributedSelectSort Stage 1: Same pattern for distributed indices
4. GetTopKKernelDistributedSelectSort Stage 2: Guard shared memory writes and output indices

Behavior After Fix

• ✅ No crash or OOB memory access
• ✅ No GPU state corruption (subsequent inferences work normally)
• ✅ Downstream embedding lookups use token 0 (valid token, produces garbage text for that sample but doesn't crash)

This is an ORT CUDA EP issue, not a genai bug per se. The genai TopK kernel is just the first thing that detects the corruption. The fix would need to be in ORT's attention kernel for large sequence lengths, or genai should cap/resize images that would produce too many patches.

Confirmed! The root cause is vision model fp16 numerical instability, NOT an ort-extensions bug.

Even with perfect, HF-identical pixel values, the vision.onnx model in fp16 produces NaN at output position [1316, 0] (patch 1316 onward). This is a model precision issue — certain image patterns trigger fp16 overflow in the vision transformer's attention/normalization layers.

Testing

Verified locally on A100-SXM4-80GB with the repro image:

• Before fix: CUDA error: an illegal memory access was encountered → all subsequent GPU ops fail
• After fix: Generation completes (garbage output due to NaN), second inference succeeds normally

Files Changed

• src/cuda/cuda_topk_select_sort.cuh — GetTop1Kernel + GetTopKKernel
• src/cuda/cuda_topk_distributed_select_sort.cuh — Distributed TopK Stage 1 + Stage 2

[Copilot]

Pull request overview

This PR hardens the CUDA TopK selection-sort kernels against all-NaN logits to prevent out-of-bounds writes that can trigger illegal memory access and corrupt subsequent GPU work, improving runtime resilience for edge-case vision inputs.

Changes:

• Clamp invalid Top-1 / Top-K output indices to 0 when no valid max element is found (e.g., all inputs are NaN).
• Add bounds guards around in-place scores_in[...] “blanking” writes to prevent OOB when the selected index remains at its sentinel value.
• Apply the same NaN-safety pattern to both the single-block selection-sort path and the distributed selection-sort (Stage 1 + Stage 2) path.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
src/cuda/cuda_topk_select_sort.cuh	Adds NaN-safe index clamping and guards the in-place blanking write to prevent OOB in Top-1/Top-K kernels.
src/cuda/cuda_topk_distributed_select_sort.cuh	Adds NaN-safe clamping/guards in distributed TopK Stage 1 and Stage 2 to avoid OOB writes and downstream invalid indices.