fix(cuda): guard exhausted candidates in TopK stage 2

[mzggr0914]

Skip the invalidation store when a reduction returns NOT_FOUND. Add regression coverage for exhausted rows, masked values, and multiple k values.

Validation: 12 TopK tests passed on Windows CUDA.

AI-assisted analysis, implementation, and test execution.

[Copilot]

Copilot review overview

🟢 Approval recommended

No unresolved blocking issues were identified.

Review effort: Lite
Findings: None

What changed in this PR

Fixes CUDA TopK stage 2 handling for exhausted candidates and adds regression coverage.

Changes:

• Skip invalidation stores when reduction returns NOT_FOUND.
• Add tests for masked values, repeated calls, row positions, data types, and multiple k values.

File	Description
tests/​ops_test.cc	Adds comprehensive TopK regression coverage.
src/​ops/​topk_gpu.cu	Prevents invalid memory writes for exhausted reductions.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jordimas]

Could you also run the new regression test under CUDA Compute Sanitizer, ideally before and after the fix? The invalid write may not reliably cause an output mismatch, so this would help confirm the regression is caught.

[mzggr0914]

Could you also run the new regression test under CUDA Compute Sanitizer, ideally before and after the fix? The invalid write may not reliably cause an output mismatch, so this would help confirm the regression is caught.

Thanks, agreed — the output assertions alone may miss the invalid write. I also ran the regression under Compute Sanitizer memcheck in the local validation; the before/after logs are attached.

The before-fix build contains the new regression test on d44d2d0, but not the kernel guard. The fixed kernel and test sources match fbfd438.

Environment: Windows x64, RTX 4070 SUPER, CUDA 12.8.93, Compute Sanitizer 2025.1.0.0; Release build with -lineinfo.

Build	Test filter	Allocator	Result
Before fix	CUDA/OpDeviceFPTest.TopKWithExhaustedCandidates/float16	cuda_malloc_async	Invalid 2-byte global write in topk_stage_2, topk_gpu.cu:262; exit 99
After fix	*TopK*	cuda_malloc_async	12 tests passed; 0 memcheck errors; exit 0
After fix	*TopK*	cub_caching	12 tests passed; 0 memcheck errors; exit 0

All runs used --tool memcheck --error-exitcode 99. The after-fix suite includes the same CUDA FP16 regression. The attachment contains the full logs and reproduction commands.

This confirms that memcheck catches the regression without relying on an output mismatch.

after-memcheck-async.txt
after-memcheck-cub.txt
before-memcheck.txt

[mzggr0914]

please let me know if anything else is needed from my side.
@jordimas

[jordimas]

Thanks, great job!