Skip to content

fix(indexer): add validated bounded compatibility policy #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yatesdr wants to merge 2 commits into from
Closed

fix(indexer): add validated bounded compatibility policy #82

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cb4d8ed
feat(indexer): add bounded long-context selection policy
yatesdr Jul 26, 2026
ec8052e
test(indexer): validate compatibility policy selection
yatesdr Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 68 additions & 33 deletions sparkinfer/attention/nsa_indexer/tiled_topk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,18 +41,48 @@
_DEFAULT_TOPK = 2048
_SUPPORTED_TOPK = (512, 1024, 2048)
_RADIX = 256
_SELECTION_POLICY_ENV = "SPARKINFER_NSA_TOPK_SELECTION_POLICY"
_SELECTION_POLICY_EXACT = "exact"
_SELECTION_POLICY_BOUNDED_COMPAT = "bounded_compat"
_SELECTION_POLICIES = frozenset(
{
_SELECTION_POLICY_EXACT,
_SELECTION_POLICY_BOUNDED_COMPAT,
}
)


def _resolve_selection_policy(value: str | None) -> str:
policy = (value or _SELECTION_POLICY_EXACT).strip().lower()
if policy not in _SELECTION_POLICIES:
raise ValueError(
f"{_SELECTION_POLICY_ENV} must be one of "
f"{sorted(_SELECTION_POLICIES)}, got {policy!r}"
)
return policy


_SELECTION_POLICY = _resolve_selection_policy(
os.environ.get(_SELECTION_POLICY_ENV)
)
_BOUNDED_COMPAT = _SELECTION_POLICY == _SELECTION_POLICY_BOUNDED_COMPAT
# Use every SM120 CTA lane for the first histogram. The four-times narrower
# bucket keeps ordinary 32k/64k low-contrast rows on the buffered exact path;
# truly degenerate buckets still use the rescan fallback below.
_COARSE_RADIX_BITS = 10
#
# ``bounded_compat`` intentionally retains the historical 8-bit coarse bucket
# and bounded 4096-candidate refinement. This is an explicit model-compatibility
# policy, not an unsafe overflow: writes remain bounds checked, while candidates
# beyond the fixed buffer budget are omitted. The default remains exact.
_COARSE_RADIX_BITS = 8 if _BOUNDED_COMPAT else 10
_COARSE_RADIX_BINS = 1 << _COARSE_RADIX_BITS
_HIST_SLOTS = _COARSE_RADIX_BINS + 128
# A 1024-thread selector CTA is already limited to one resident block per SM on
# SM120 (1536 threads/SM). Keeping 8192 candidates therefore avoids the common
# long-context overflow without reducing occupancy; the complete shared-memory
# allocation remains below the 99 KiB opt-in block limit. The exact rescan below
# still covers wider or degenerate threshold buckets.
_SMEM_CANDS = 8192
_SMEM_CANDS = 4096 if _BOUNDED_COMPAT else 8192
_SCAN_UNROLL = 4
_SUPERTILE_K_ENV = "SPARKINFER_NSA_TOPK_SUPERTILE_K"
_SUPERTILE_K_DEFAULT = 32768
Expand Down Expand Up @@ -539,6 +569,7 @@ def __init__(
# or the user's final output.
self.is_first = bool(is_first)
self.output_physical_slots = bool(output_physical_slots)
self.bounded_compat = bool(_BOUNDED_COMPAT)

@cute.jit
def __call__(
Expand Down Expand Up @@ -987,8 +1018,9 @@ class SharedStorage:
# buffer first: none of those intermediate outputs survive. This is
# a CTA-uniform branch because every thread reads the same shared
# counter after the barrier above.
if bin_count > Int32(_SMEM_CANDS):
topk = Int32(-1)
if not cutlass.const_expr(self.bounded_compat):
if bin_count > Int32(_SMEM_CANDS):
topk = Int32(-1)

# Stage 2: refine with 8-bit radix passes
for round_idx in cutlass.range_constexpr(4):
Expand Down Expand Up @@ -1140,33 +1172,34 @@ class SharedStorage:
# Exact overflow fallback: the buffered refine above dropped
# winners when more than _SMEM_CANDS candidates shared the coarse
# threshold bucket. Redo the selection exactly by re-scanning.
if bin_count > Int32(_SMEM_CANDS):
_exact_overflow_fallback(
tx,
total_len,
topk_static,
s_hist0,
s_hist1,
s_out,
h0,
ctr,
thr,
ni0,
ni1,
lr,
flat=False,
flat_values=input_tensor,
input_tensor=input_tensor,
carry_values=carry_values,
row_base=row_base,
row_start=row_start,
carry_base=out_base,
chunk_len=length,
block_q=self.block_q,
block_k=self.block_k,
is_tiled=self.is_tiled,
is_first=self.is_first,
)
if not cutlass.const_expr(self.bounded_compat):
if bin_count > Int32(_SMEM_CANDS):
_exact_overflow_fallback(
tx,
total_len,
topk_static,
s_hist0,
s_hist1,
s_out,
h0,
ctr,
thr,
ni0,
ni1,
lr,
flat=False,
flat_values=input_tensor,
input_tensor=input_tensor,
carry_values=carry_values,
row_base=row_base,
row_start=row_start,
carry_base=out_base,
chunk_len=length,
block_q=self.block_q,
block_k=self.block_k,
is_tiled=self.is_tiled,
is_first=self.is_first,
)

cute.arch.sync_threads()
idx0 = Int32(tx)
Expand Down Expand Up @@ -1543,7 +1576,8 @@ def run_tiled_topk(
dynamic_strides=(0,),
),
(
"tiled_topk_v26_wide_coarse",
"tiled_topk_v27_selection_policy",
_SELECTION_POLICY,
topk,
block_q,
block_k,
Expand Down Expand Up @@ -1713,7 +1747,8 @@ def run_row_topk(
"carry_indices", carry_indices_key_tensor, dynamic=True
),
(
"row_topk_v6",
"row_topk_v7_selection_policy",
_SELECTION_POLICY,
topk,
output_gather_table is not None,
),
Expand Down
30 changes: 30 additions & 0 deletions tests/attention/test_nsa_topk_selection_policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
from __future__ import annotations

import pytest

from sparkinfer.attention.nsa_indexer.tiled_topk import (
_resolve_selection_policy,
)


@pytest.mark.parametrize(
("raw", "expected"),
[
(None, "exact"),
("", "exact"),
("exact", "exact"),
(" EXACT ", "exact"),
("bounded_compat", "bounded_compat"),
(" BOUNDED_COMPAT ", "bounded_compat"),
],
)
def test_resolve_selection_policy(raw: str | None, expected: str) -> None:
assert _resolve_selection_policy(raw) == expected


def test_resolve_selection_policy_rejects_unknown_value() -> None:
with pytest.raises(
ValueError,
match="SPARKINFER_NSA_TOPK_SELECTION_POLICY must be one of",
):
_resolve_selection_policy("legacy")
Comment on lines +10 to +30

@coderabbitai coderabbitai Bot Jul 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Test the actual bounded-overflow contract.

These cases only test parsing. Add a CUDA oracle test with a threshold bin exceeding 4,096 candidates: exact must equal reference top-k, while bounded_compat must match the historical bounded selector’s expected outputs. Exercise both tiled and row dispatches.

As per coding guidelines, “Validate correctness gates—including oracles, cosine/top-k equality, nonzero tensors, quantization semantics, and boundary behavior—before interpreting timings.”

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/attention/test_nsa_topk_selection_policy.py` around lines 10 - 30,
Extend the CUDA tests beyond _resolve_selection_policy parsing with an oracle
case containing a threshold bin larger than 4,096 candidates. Exercise both
tiled and row dispatches, asserting exact matches the reference top-k and
bounded_compat matches the historical bounded selector’s expected outputs;
include the required nonzero input and boundary/selection correctness checks
before any timing assertions.

Source: Coding guidelines