From 91b65380752fbb77d7437568203f1dce3ce5a705 Mon Sep 17 00:00:00 2001 From: David Fan Date: Sat, 6 Jun 2026 02:16:43 +0000 Subject: [PATCH 1/4] Fix CUDA illegal memory access crash when TopK inputs contain NaN When vision models produce NaN logits (e.g., due to corrupt/edge-case images), the TopK kernels crash with illegal memory access because: - TopK_Pair.p is initialized to INT_MAX for tie-breaking - NaN comparisons always return false, so no Insert succeeds - p stays at INT_MAX, causing scores_in[INT_MAX] = OOB write - This corrupts GPU state and kills all subsequent inferences Fix: Guard all scores_in writes with p < vocab_size check, and clamp output indices to 0 when no valid element is found. This ensures: 1. No OOB memory access (prevents crash) 2. No GPU state corruption (subsequent inferences still work) 3. Downstream embedding lookups use token 0 (valid, produces garbage text for that sample but doesn't crash) Repro: Qwen2.5-VL-3B with a 1500x1204 image where ort-extensions bicubic resize produces incorrect pixel values, causing NaN in the vision encoder that propagates to all-NaN logits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/cuda/cuda_topk_distributed_select_sort.cuh | 18 ++++++++++++++---- src/cuda/cuda_topk_select_sort.cuh | 13 ++++++++++--- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/src/cuda/cuda_topk_distributed_select_sort.cuh b/src/cuda/cuda_topk_distributed_select_sort.cuh index 8a1a01c164..61b5cd078f 100644 --- a/src/cuda/cuda_topk_distributed_select_sort.cuh +++ b/src/cuda/cuda_topk_distributed_select_sort.cuh @@ -112,8 +112,13 @@ __global__ void GetTopKKernelDistributedSelectSort(float* scores_in, float* scor if (tid == 0) { distributed_scores_out_curr[ite] = top_k_sequence.u; - distributed_indices_out_curr[ite] = top_k_sequence.p; - scores_in[top_k_sequence.p] = MIN_FLOAT; + // Clamp invalid index to 0 for safety (NaN inputs leave p at INT_MAX) + distributed_indices_out_curr[ite] = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0; + // Guard against NaN inputs: if all scores are NaN, no valid index is found + // and p remains INT_MAX (its init value). Writing to that index would be OOB. + if (top_k_sequence.p < vocab_size) { + scores_in[top_k_sequence.p] = MIN_FLOAT; + } __threadfence_block(); } __syncthreads(); @@ -181,8 +186,13 @@ __global__ void GetTopKKernelDistributedSelectSort(float* scores_in, float* scor int vocab_index = top_k_sequence_reduced.p_indirection; scores_out[ite] = top_k_sequence_reduced.u; - indices_out[ite] = vocab_index; - shared_distributed_scores_out[index] = MIN_FLOAT; + // Guard against NaN: if no valid element was found, p_indirection stays INT_MAX. + // Clamp to 0 to avoid downstream OOB in embedding lookups. + indices_out[ite] = (vocab_index >= 0 && vocab_index < vocab_size) ? vocab_index : 0; + // Guard against NaN inputs: if all scores are NaN, p remains -1 (init value). + if (index >= 0 && index < num_top_k_shards * k) { + shared_distributed_scores_out[index] = MIN_FLOAT; + } __threadfence_block(); } diff --git a/src/cuda/cuda_topk_select_sort.cuh b/src/cuda/cuda_topk_select_sort.cuh index 17553540fa..345ea4bd5a 100644 --- a/src/cuda/cuda_topk_select_sort.cuh +++ b/src/cuda/cuda_topk_select_sort.cuh @@ -79,8 +79,10 @@ __global__ void GetTop1Kernel(const float* scores_in, float* scores_out, int* in // Thread 0 writes the final result. No fence or write-back to scores_in is needed. if (tid == 0) { + // Guard against NaN inputs: clamp invalid index to 0 + int safe_p = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0; scores_out[batch] = top_k_sequence.u; - indices_out[batch] = top_k_sequence.p; + indices_out[batch] = safe_p; } } @@ -104,9 +106,14 @@ __global__ void GetTopKKernel(volatile float* scores_in, float* scores_out, int* TopK_Pair top_k_sequence = BlockReduce(temp_storage).Reduce(partial, reduce_topk_op); if (tid == 0) { + // Guard against NaN inputs: if all scores are NaN, no valid index is found + // and p remains INT_MAX (its init value). Clamp to index 0 to avoid OOB. + int safe_p = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0; scores_out[ite + batch * k] = top_k_sequence.u; - indices_out[ite + batch * k] = top_k_sequence.p; - scores_in[batch * vocab_size + top_k_sequence.p] = -FLT_MAX; + indices_out[ite + batch * k] = safe_p; + if (top_k_sequence.p < vocab_size) { + scores_in[batch * vocab_size + top_k_sequence.p] = -FLT_MAX; + } __threadfence_block(); } From e96b7ea14ede616de0da606df1acff4d364010d6 Mon Sep 17 00:00:00 2001 From: David Fan Date: Sat, 6 Jun 2026 02:51:58 +0000 Subject: [PATCH 2/4] Fix misleading comment: p_indirection init value is -1, not INT_MAX Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/cuda/cuda_topk_distributed_select_sort.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cuda/cuda_topk_distributed_select_sort.cuh b/src/cuda/cuda_topk_distributed_select_sort.cuh index 61b5cd078f..6d91c5a7a5 100644 --- a/src/cuda/cuda_topk_distributed_select_sort.cuh +++ b/src/cuda/cuda_topk_distributed_select_sort.cuh @@ -186,10 +186,10 @@ __global__ void GetTopKKernelDistributedSelectSort(float* scores_in, float* scor int vocab_index = top_k_sequence_reduced.p_indirection; scores_out[ite] = top_k_sequence_reduced.u; - // Guard against NaN: if no valid element was found, p_indirection stays INT_MAX. + // Guard against NaN: if no valid element was found, p_indirection stays at -1 (init value). // Clamp to 0 to avoid downstream OOB in embedding lookups. indices_out[ite] = (vocab_index >= 0 && vocab_index < vocab_size) ? vocab_index : 0; - // Guard against NaN inputs: if all scores are NaN, p remains -1 (init value). + // Guard against NaN inputs: if all scores are NaN, p stays at -1 (init value). if (index >= 0 && index < num_top_k_shards * k) { shared_distributed_scores_out[index] = MIN_FLOAT; } From 8beb105352b8cc0f8b2a08294c4907aeab24183c Mon Sep 17 00:00:00 2001 From: David Fan Date: Mon, 8 Jun 2026 23:51:33 +0000 Subject: [PATCH 3/4] Add NaN safety unit test for CUDA TopK kernels Adds a test that feeds all-NaN input scores to each TopK algorithm and verifies: 1. No CUDA illegal memory access (the original crash scenario) 2. All output indices are valid (within [0, vocab_size)) This guards the defensive bounds checks added in the TopK kernels against accidental removal in the future. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- test/cuda/cuda_topk_tests.cpp | 121 ++++++++++++++++++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/test/cuda/cuda_topk_tests.cpp b/test/cuda/cuda_topk_tests.cpp index 3b62caa90c..96d17494ff 100644 --- a/test/cuda/cuda_topk_tests.cpp +++ b/test/cuda/cuda_topk_tests.cpp @@ -213,4 +213,125 @@ TEST(TopKTests, ParityTests) { RunParityTests(params); } } + +// Test that TopK kernels handle all-NaN inputs without crashing or producing OOB indices. +// This guards against the bug where NaN comparisons always return false, leaving the +// index at its sentinel value (INT_MAX or -1), which would cause an illegal memory access. +void RunNaNSafetyTest(const TopKTestParams& params) { + std::cout << "\n--- Running NaN Safety Test with batch_size=" << params.batch_size + << ", vocab_size=" << params.vocab_size << ", k=" << params.k << " ---\n"; + + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + size_t total_vocab_size = static_cast(params.batch_size) * params.vocab_size; + size_t topk_size = static_cast(params.batch_size) * params.k; + + auto scores_in_d = Generators::CudaMallocArray(total_vocab_size); + + // Fill input with all NaN values + std::vector scores_in_h(total_vocab_size, std::numeric_limits::quiet_NaN()); + CUDA_CHECK(cudaMemcpy(scores_in_d.get(), scores_in_h.data(), scores_in_h.size() * sizeof(float), cudaMemcpyHostToDevice)); + + auto topk_data = std::make_unique(params.batch_size, params.vocab_size, stream); + + auto test_algo = [&](const std::string& name, auto func) { + // Re-upload NaN data since some algorithms modify the input in-place + CUDA_CHECK(cudaMemcpy(scores_in_d.get(), scores_in_h.data(), scores_in_h.size() * sizeof(float), cudaMemcpyHostToDevice)); + + func(); + + topk_data->CompactOutput(params.batch_size, params.k, stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + // Verify no CUDA errors (the main goal: no illegal memory access) + cudaError_t err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << "CUDA error after " << name << ": " << cudaGetErrorString(err); + + // Verify all output indices are valid (within [0, vocab_size)) + std::vector actual_indices_h(topk_size); + CUDA_CHECK(cudaMemcpy(actual_indices_h.data(), topk_data->topk_indices_compact.get(), + actual_indices_h.size() * sizeof(int), cudaMemcpyDeviceToHost)); + + for (size_t i = 0; i < topk_size; ++i) { + ASSERT_GE(actual_indices_h[i], 0) + << name << ": output index[" << i << "] = " << actual_indices_h[i] << " is negative"; + ASSERT_LT(actual_indices_h[i], params.vocab_size) + << name << ": output index[" << i << "] = " << actual_indices_h[i] + << " is out of bounds (vocab_size=" << params.vocab_size << ")"; + } + + std::cout << " [PASS] " << name << " (NaN Safety)" << std::endl; + }; + + test_algo(Generators::cuda::select_sort::kAlgorithmName, [&]() { + Generators::cuda::select_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + + test_algo(Generators::cuda::per_batch_radix_sort::kAlgorithmName, [&]() { + Generators::cuda::per_batch_radix_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + + if (Generators::cuda::hybrid_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) { + test_algo(Generators::cuda::hybrid_sort::kAlgorithmName, [&]() { + Generators::cuda::hybrid_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + } + + if (Generators::cuda::iterative_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) { + test_algo(Generators::cuda::iterative_sort::kAlgorithmName, [&]() { + Generators::cuda::iterative_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + } + + if (Generators::cuda::cascaded_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) { + test_algo(Generators::cuda::cascaded_sort::kAlgorithmName, [&]() { + Generators::cuda::cascaded_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + } + + if (Generators::cuda::flash_convergent::IsSupported(params.batch_size, params.vocab_size, params.k)) { + test_algo(Generators::cuda::flash_convergent::kAlgorithmName, [&]() { + Generators::cuda::flash_convergent::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + } + + if (Generators::cuda::distributed_select_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) { + test_algo(Generators::cuda::distributed_select_sort::kAlgorithmName, [&]() { + Generators::cuda::distributed_select_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + } + + // Test the default RunTopK dispatcher + test_algo("DEFAULT", [&]() { + Generators::cuda::RunTopK(topk_data.get(), stream, scores_in_d.get(), + params.vocab_size, params.batch_size, params.k); + }); + + CUDA_CHECK(cudaStreamDestroy(stream)); +} + +TEST(TopKTests, NaNSafetyTests) { + // Test representative configurations that exercise different code paths + std::vector test_cases = { + {1, 200, 1}, // Top-1 path (GetTop1Kernel) + {1, 200, 4}, // Small vocab, k > 1 (GetTopKKernel) + {4, 2000, 1}, // Multi-batch Top-1 + {4, 2000, 16}, // Multi-batch general TopK + {1, 200000, 1}, // Large vocab Top-1 (triggers distributed_select_sort) + {1, 200000, 32}, // Large vocab general TopK (triggers distributed_select_sort) + {32, 20000, 64}, // Large batch, large k + }; + + for (const auto& params : test_cases) { + RunNaNSafetyTest(params); + } +} #endif From b5adc69b7203e1360d28141a80a8e0e89443bf91 Mon Sep 17 00:00:00 2001 From: David Fan Date: Tue, 9 Jun 2026 02:06:57 +0000 Subject: [PATCH 4/4] Fix clang-format violations in NaN safety test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- test/cuda/cuda_topk_tests.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/test/cuda/cuda_topk_tests.cpp b/test/cuda/cuda_topk_tests.cpp index 96d17494ff..60e4fe778d 100644 --- a/test/cuda/cuda_topk_tests.cpp +++ b/test/cuda/cuda_topk_tests.cpp @@ -321,13 +321,13 @@ void RunNaNSafetyTest(const TopKTestParams& params) { TEST(TopKTests, NaNSafetyTests) { // Test representative configurations that exercise different code paths std::vector test_cases = { - {1, 200, 1}, // Top-1 path (GetTop1Kernel) - {1, 200, 4}, // Small vocab, k > 1 (GetTopKKernel) - {4, 2000, 1}, // Multi-batch Top-1 - {4, 2000, 16}, // Multi-batch general TopK - {1, 200000, 1}, // Large vocab Top-1 (triggers distributed_select_sort) - {1, 200000, 32}, // Large vocab general TopK (triggers distributed_select_sort) - {32, 20000, 64}, // Large batch, large k + {1, 200, 1}, + {1, 200, 4}, + {4, 2000, 1}, + {4, 2000, 16}, + {1, 200000, 1}, + {1, 200000, 32}, + {32, 20000, 64}, }; for (const auto& params : test_cases) {