From cb4d8edb239c208752a13f9875940027cb3fa984 Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 26 Jul 2026 12:58:10 -0400 Subject: [PATCH 1/2] feat(indexer): add bounded long-context selection policy --- .../attention/nsa_indexer/tiled_topk.py | 95 ++++++++++++------- 1 file changed, 62 insertions(+), 33 deletions(-) diff --git a/sparkinfer/attention/nsa_indexer/tiled_topk.py b/sparkinfer/attention/nsa_indexer/tiled_topk.py index 4b82a1f5b..83b903653 100644 --- a/sparkinfer/attention/nsa_indexer/tiled_topk.py +++ b/sparkinfer/attention/nsa_indexer/tiled_topk.py @@ -41,10 +41,34 @@ _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, + } +) +_SELECTION_POLICY = os.environ.get( + _SELECTION_POLICY_ENV, _SELECTION_POLICY_EXACT +).strip().lower() +if _SELECTION_POLICY not in _SELECTION_POLICIES: + raise ValueError( + f"{_SELECTION_POLICY_ENV} must be one of " + f"{sorted(_SELECTION_POLICIES)}, got {_SELECTION_POLICY!r}" + ) +_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 accidental overflow: writes remain bounds checked, while +# candidates beyond the deterministic 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 @@ -52,7 +76,7 @@ # 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 @@ -539,6 +563,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__( @@ -987,8 +1012,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): @@ -1140,33 +1166,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) @@ -1543,7 +1570,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, @@ -1713,7 +1741,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, ), From ec8052e41c2ef8740698b53f003a986f14449ffc Mon Sep 17 00:00:00 2001 From: derek Date: Sun, 26 Jul 2026 13:52:05 -0400 Subject: [PATCH 2/2] test(indexer): validate compatibility policy selection --- .../attention/nsa_indexer/tiled_topk.py | 28 ++++++++++------- .../test_nsa_topk_selection_policy.py | 30 +++++++++++++++++++ 2 files changed, 47 insertions(+), 11 deletions(-) create mode 100644 tests/attention/test_nsa_topk_selection_policy.py diff --git a/sparkinfer/attention/nsa_indexer/tiled_topk.py b/sparkinfer/attention/nsa_indexer/tiled_topk.py index 83b903653..b97ec0d1d 100644 --- a/sparkinfer/attention/nsa_indexer/tiled_topk.py +++ b/sparkinfer/attention/nsa_indexer/tiled_topk.py @@ -50,14 +50,21 @@ _SELECTION_POLICY_BOUNDED_COMPAT, } ) -_SELECTION_POLICY = os.environ.get( - _SELECTION_POLICY_ENV, _SELECTION_POLICY_EXACT -).strip().lower() -if _SELECTION_POLICY not in _SELECTION_POLICIES: - raise ValueError( - f"{_SELECTION_POLICY_ENV} must be one of " - f"{sorted(_SELECTION_POLICIES)}, got {_SELECTION_POLICY!r}" - ) + + +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; @@ -65,9 +72,8 @@ # # ``bounded_compat`` intentionally retains the historical 8-bit coarse bucket # and bounded 4096-candidate refinement. This is an explicit model-compatibility -# policy, not an accidental overflow: writes remain bounds checked, while -# candidates beyond the deterministic buffer budget are omitted. The default -# remains exact. +# 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 diff --git a/tests/attention/test_nsa_topk_selection_policy.py b/tests/attention/test_nsa_topk_selection_policy.py new file mode 100644 index 000000000..31377b990 --- /dev/null +++ b/tests/attention/test_nsa_topk_selection_policy.py @@ -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")