perf(indexer): keep exact long-context top-k on the fast path

[voipmonitor]

Summary

• widen the tiled selector's coarse radix from 8 to 10 bits
• increase the shared candidate capacity from 4,096 to 8,192 entries
• skip buffered refinement once overflow has already selected the exact rescan path
• retain the candidate-buffer-free exact fallback introduced by Fix silent candidate-buffer overflow in indexer top-k #55

Why

#55 fixed silent top-k corruption when the threshold bucket exceeded the fixed candidate buffer. On GLM-5.2 long-context prefill, however, the 8-bit coarse bucket commonly exceeded 4,096 entries. That sent ordinary rows through four full exact rescans and reduced 64k prefill throughput.

SM120 already runs this 1,024-thread selector at one CTA per SM. The wider histogram uses all CTA lanes, and the larger candidate arrays keep the complete shared-memory allocation below the 99 KiB opt-in block limit without reducing occupancy. Truly degenerate rows still take the exact fallback.

Validation

• ruff check passed on both changed files
• 7 passed in test_paged_prefill_topk_long_context.py on SM120
• coverage includes 32k/64k low-contrast rows, all-equal overflow, logical two-level fold, and multi-chunk carry fold
• disabling the exact fallback was tested only as a diagnostic and failed at 64k; this PR does not remove or weaken it

Matched TP8/DCP1/MTP0 GLM-5.2 runs:

Variant	64k prefill tok/s
v19, before #55	5,917
v20 with #55	5,782
8,192 candidates + skipped discarded refine	5,877
This PR, default 32k supertile	5,895
This PR + runtime 64k supertile	5,914

The last row uses SPARKINFER_PAGED_INDEX_SUPERTILE_K=65536, which is deliberately not made a library-wide default in this PR. It measured 87.70 tok/s CC1 decode and reduced KV capacity by 2,816 tokens (0.51%) versus the default 32k supertile.

Summary by CodeRabbit

• Bug Fixes

  • Improved top-k selection accuracy for long-context attention workloads, especially in low-contrast cases.
  • Enhanced handling of candidate overflows to ensure exact fallback selection when intermediate buffers are exceeded.
  • Expanded coarse threshold detection to improve reliability and reduce incorrect selections.

• Tests

  • Added regression coverage for larger sequence lengths and overflow scenarios.
  • Clarified validation of selected results against the true top-k cutoff.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: d7bf9c43-95a0-46b8-9257-ee3bacf95c12

📥 Commits

Reviewing files that changed from the base of the PR and between cc9b476 and 51f510d.

📒 Files selected for processing (2)

• sparkinfer/attention/nsa_indexer/tiled_topk.py
• tests/attention/test_paged_prefill_topk_long_context.py

---

📝 Walkthrough

Walkthrough

The tiled top-k kernel now selects coarse threshold buckets with a 10-bit radix histogram, expands shared-memory storage, and sends oversized threshold buckets directly to exact overflow fallback. Long-context regression cases move to larger sequence lengths.

Changes

Wide coarse top-k selection

Layer / File(s)	Summary
Widen coarse histogram selection sparkinfer/attention/nsa_indexer/tiled_topk.py	Coarse-bin conversion, shared histogram storage, prefix scans, and threshold selection now use the wider radix range.
Handle threshold candidates and validate overflow sparkinfer/attention/nsa_indexer/tiled_topk.py, tests/attention/test_paged_prefill_topk_long_context.py	Output scans use coarse bins, oversized candidate buckets trigger exact fallback, the cache label changes, and long-context test boundaries and documentation are updated.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant TiledTopKKernel
  participant SharedHistogram
  participant CandidateBuffer
  participant ExactOverflowFallback
  TiledTopKKernel->>SharedHistogram: build and scan widened coarse histogram
  SharedHistogram-->>TiledTopKKernel: return threshold_bin
  TiledTopKKernel->>CandidateBuffer: collect threshold-bin candidates
  CandidateBuffer-->>TiledTopKKernel: return bin_count
  TiledTopKKernel->>ExactOverflowFallback: rescan when candidate count exceeds _SMEM_CANDS
  ExactOverflowFallback-->>TiledTopKKernel: rebuild exact top-k outputs

Loading

Possibly related PRs

• local-inference-lab/sparkinfer#55: Both changes modify tiled NSA top-k threshold candidate overflow handling and exact fallback behavior.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title matches the main goal: long-context top-k performance while preserving exact results on the fast path.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/tiled-topk-wide-coarse-radix-clean-20260724

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}