indexer: bound two-level fold workspace - #87
Merged
Conversation
Select the existing exact streaming carry when the logical-output two-level candidate slab exceeds a configurable runtime budget. Preserve the parallel fold for smaller workloads and add planner and GPU correctness coverage.\n\nInspired by Josh Cartu's adaptive-fold study.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe paged indexer now plans two-level folding from environment-configured policy and memory limits, integrates the plan into ChangesTwo-level fold policy
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Environment
participant index_topk_fp8
participant _plan_two_level_fold
participant FoldBuffers
Environment->>_plan_two_level_fold: Read fold mode and memory budget
index_topk_fp8->>_plan_two_level_fold: Provide fold dimensions and output mode
_plan_two_level_fold-->>index_topk_fp8: Return slices, total_slices, and reason
index_topk_fp8->>FoldBuffers: Allocate fold state when slices are available
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
sparkinfer/attention/nsa_indexer/paged.py (1)
909-916: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
fold_plan.reasonis computed but never surfaced.The plan's
reasonfield (e.g. "candidate slab exceeds the memory budget", "forced") captures exactly why the runtime chose two-level fold vs. streaming carry, butindex_topk_fp8never logs or exposes it — only the unit tests read.reasondirectly. Given this is a new auto-switching behavior that changes memory/perf characteristics silently based on env config, surfacing the reason (e.g., via a debug log or counter) would help diagnose unexpected route flips in production.🤖 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 `@sparkinfer/attention/nsa_indexer/paged.py` around lines 909 - 916, Surface fold_plan.reason in index_topk_fp8 after _plan_two_level_fold returns, using the module’s existing debug logging or metrics mechanism so operators can identify why two-level fold was selected. Preserve the current routing behavior and include the reason for both budget-driven and forced decisions.
🤖 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 `@sparkinfer/attention/nsa_indexer/paged.py`:
- Around line 102-178: Centralize supertile geometry by adding a shared
_chunk_geometry helper and use it in both _plan_two_level_fold and
index_topk_fp8’s execution loop, preserving the existing chunk boundaries and
token offsets. Add total_slices to _TwoLevelFoldPlan, populate it in
_plan_two_level_fold, and replace the downstream sum over two_level_slices with
that field.
- Around line 67-99: Update the invalid-mode ValueError in
_read_two_level_fold_policy to document all accepted values, including the
off/false/no and on/true/yes aliases, while preserving the existing validation
behavior and reported raw value.
---
Nitpick comments:
In `@sparkinfer/attention/nsa_indexer/paged.py`:
- Around line 909-916: Surface fold_plan.reason in index_topk_fp8 after
_plan_two_level_fold returns, using the module’s existing debug logging or
metrics mechanism so operators can identify why two-level fold was selected.
Preserve the current routing behavior and include the reason for both
budget-driven and forced decisions.
🪄 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: 05ec6c78-192e-4e9e-a595-33baeb7f57c4
📒 Files selected for processing (3)
sparkinfer/attention/nsa_indexer/paged.pytests/attention/test_paged_indexer_integration.pytests/attention/test_paged_prefill_topk_long_context.py
26 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
auto, forced two-level, and forced carry controlsThis does not add a new selector or change top-k precision. Both routes are exact. It only chooses which existing fold implementation handles the intermediate candidates.
Motivation
The parallel two-level fold allocates FP32 values plus int32 indices for every query row, fold slice, and top-k entry. On large logical-output DCP prefill profiles, that transient slab can reach several GiB per rank and is included in vLLM's peak non-KV memory profile. The resulting peak reduces the KV cache even though the allocation is not persistent.
Physical-slot output, including the current DCP1 route, already requires streaming carry and is unchanged.
The policy was inspired by Josh Cartu's adaptive-fold study. This PR ports the idea onto current SparkInfer
master, derives the complete fold geometry internally, accounts for both candidate tensors and row lengths, validates configuration and dimensions, and adds current integration coverage.Controls
auto: use two-level fold while the candidate allocation fits; otherwise use carry0: force streaming carry1: force historical two-level behaviorE2E results
Paired runs used TP8/DCP4/MTP0, FP8 KV, A16, a 1,048,576-token model limit, GPUs 0-7, and otherwise identical source and configuration.
The throughput deltas are within run noise. The KV capacity gain was stable across boots.
Validation
Summary by CodeRabbit