NOSSLETTER.techLIVE
back to archive

[Perf][Kernel] Single-launch fused scoring+top-k kernel for sparse-indexer prefill, without materializing logits

Language: Python+4264 -108performance
concepts demonstrated
CUDA kernelsMemory optimizationTop-k selection algorithmsSparse indexing
The Issue & Context

Before this PR, the sparse-indexer prefill step materialized a large fp32 logits matrix for every query row and visible key, causing massive memory usage and slowdowns when handling large key sets (up to 1M keys).

Root Cause Analysis

The dense prefill path computed and stored the full [rows, keys] logits matrix in fp32, requiring repeated kernel launches and large HBM traffic for writing and reading logits, which limited throughput and forced tradeoffs between memory budget and KV cache capacity.

Solution Strategy & Architectural Pattern

This PR implements a single persistent CUDA kernel that fuses scoring and top-k selection for each row without materializing the full logits matrix, ranking compact 24-bit score codes instead and rescanning only rows that cannot be served by the gate exactly in fp32, drastically reducing memory traffic and kernel launches.

Key Takeaway for Contributors

Fusing compute and selection steps into a single kernel and avoiding intermediate large buffers can significantly reduce memory bandwidth and improve throughput in large-scale GPU workloads.