fix(indexer): add validated bounded compatibility policy - #82
Conversation
📝 WalkthroughWalkthroughAdds environment-controlled ChangesSelection policy execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Environment
participant tiled_topk
participant TopkKernel
participant CudaCache
Environment->>tiled_topk: set selection policy
tiled_topk->>TopkKernel: configure exact or bounded_compat mode
TopkKernel->>TopkKernel: handle threshold-bin overflow
tiled_topk->>CudaCache: compile using policy-specific cache key
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with 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.
Inline comments:
In `@tests/attention/test_nsa_topk_selection_policy.py`:
- Around line 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.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 6d01d05c-fabb-4363-8495-eb49b8f3541c
📒 Files selected for processing (2)
sparkinfer/attention/nsa_indexer/tiled_topk.pytests/attention/test_nsa_topk_selection_policy.py
| @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") |
There was a problem hiding this comment.
🎯 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
|
Superseded by SparkInfer #86 and vLLM #189, which provide improved fixes. |
1 participant
Summary
Add an explicit
SPARKINFER_NSA_TOPK_SELECTION_POLICY=bounded_compatpolicy for sparse-MLA deployments whose checkpoints were validated against
the historical bounded radix selector.
The default remains
exact. Unknown values fail closed at import time.Fixes the SparkInfer component of
local-inference-lab/vllm#182.
Root cause
The v20 exact-overflow changes (
1012199e, later optimized by83a58444)changed the selected sparse-attention set at long context. On captured
production tensors:
This model has no separate sliding/local window in sparse layers. The exact
proxy selector therefore reallocates part of the only 2,048-entry sparse
attention budget away from older positions. Exactness for the
E4M3-query/FP8-key proxy did not preserve end-to-end retrieval quality.
Change
bounded_compatexplicitly selects the historical:All shared-memory writes remain bounds-checked. The selection policy is part
of the compile-cache key. No behavior changes unless the new policy is
explicitly enabled.
Reproduction and causal evidence
Frozen production posture:
i8_ring;cached_tokens=0);Results:
bounded_compatAll recovered responses finalized normally with content
738216. Thediscriminator boot remained healthy with a 507,612-token KV pool at 460k,
zero restarts, and no illegal-access, cuBLAS, EngineDead, OOM, traceback,
assertion, or worker-died signatures.
An independent cold generalization ladder on the same process also passed:
The trace-free production candidate then passed the established deepest cold
gate at a 480,000-token admission limit:
stop738216That candidate was derived from the clean
5517197/be0edcarelease image andcontained only PR #80 plus this PR. It exposed a 500,992-token KV pool, remained
healthy with zero restarts, and had zero fatal signatures.
Full reproduction, frozen prompt hashes, boundary traces, learned-tensor
replay, and evidence hashes are in
vLLM issue #182.
Validation
python -m py_compilepasses for the changed module and policy test.speculator graph capture, and passed the frozen causal gate above.
unchanged.
The focused pytest was not run on the authoring Mac because that environment
does not have pytest installed; CI should run
tests/attention/test_nsa_topk_selection_policy.py.Scope
This PR provides a visible compatibility policy, not a claim that bounded
selection should become the universal default. A future default should be
deterministic and justified by end-to-end model quality rather than only
quantized-proxy exactness.
Summary by CodeRabbit
New Features
Bug Fixes
Tests