Repository navigation
[Performance][SFA] Bypass redundant LightningIndexer scoring for eligible full-visible short-prefill cases - #15813
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a performance optimization for Ascend SFA by implementing a fast path for specific full-visible short-prefill scenarios. By precomputing and reusing round-robin indices, the system avoids unnecessary LightningIndexer scoring. The feature is strictly gated behind a new configuration flag and includes robust fallback mechanisms for unsupported cases, such as multi-sequence batches or specific MTP layer configurations, ensuring system stability and correctness. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Attention][Feature] Implement SFA full visible index bypass optimizationSuggested PR Summary:
### What this PR does / why we need it?
This pull request introduces an optimization to bypass the SFA indexer scoring for short prefills when the SFA full visible index bypass is enabled. It adds `_get_full_visible_topk_indices` to `AscendSFAImpl` to retrieve pre-computed top-k indices, avoiding redundant indexer calculations. Feedback suggests lazy-initializing the `_full_visible_index_table` on the correct device during the forward pass to prevent device mismatch issues when modules are moved across devices, along with corresponding test updates.
### Does this PR introduce _any_ user-facing change?
No, this is an internal performance optimization.
### How was this patch tested?
Tested with new unit tests in `tests/ut/attention/test_sfa_v1.py`, `tests/ut/ops/test_mla.py`, `tests/ut/patch/worker/test_patch_deepseek_v2.py`, and `tests/ut/test_ascend_config.py`.| if ( | ||
| type(self) is not AscendSFAImpl | ||
| or not get_ascend_config().enable_sfa_full_visible_index_bypass | ||
| or not self.allow_short_prefill_indexer_scoring_skip | ||
| or self.skip_topk | ||
| or not self.has_indexer | ||
| or self.enable_sparse_sfa_c8 | ||
| or self.enable_sparse_li_c8 | ||
| or attn_metadata.attn_state | ||
| not in (AscendAttentionState.PrefillNoCache, AscendAttentionState.PrefillCacheHit) | ||
| or attn_metadata.block_size != SFA_FULL_VISIBLE_TEMPLATE_BLOCK_SIZE | ||
| or attn_metadata.seq_lens_cpu is None | ||
| or attn_metadata.seq_lens_cpu.numel() != 1 | ||
| or attn_metadata.num_actual_tokens != topk_num_tokens | ||
| or self._full_visible_index_table is None | ||
| or self._full_visible_index_table.device != device | ||
| ): | ||
| return None | ||
| kv_length = int(attn_metadata.seq_lens_cpu[0]) | ||
| context_length = kv_length - attn_metadata.num_actual_tokens | ||
| if context_length < 0 or kv_length > SFA_INDEXER_SPARSE_COUNT: | ||
| return None | ||
| return self._full_visible_index_table[context_length + 1 : kv_length + 1].unsqueeze(1) |
There was a problem hiding this comment.
In PyTorch, plain tensor attributes (like self._full_visible_index_table) are not automatically moved to the target device when .to(device) or .npu() is called on the parent nn.Module. Since models are often initialized on CPU first and then moved to NPU, or initialized on npu:0 and run on another rank in multi-GPU/TP setups, self._full_visible_index_table will remain on the wrong device. This causes the device mismatch check (self._full_visible_index_table.device != device) to silently fail and fall back to the slow path, completely defeating the performance optimization.
To make this robust, we should lazy-initialize or update self._full_visible_index_table on the fly during the forward pass if it is None or on a different device.
if (
type(self) is not AscendSFAImpl
or not get_ascend_config().enable_sfa_full_visible_index_bypass
or not self.allow_short_prefill_indexer_scoring_skip
or self.skip_topk
or not self.has_indexer
or self.enable_sparse_sfa_c8
or self.enable_sparse_li_c8
or attn_metadata.attn_state
not in (AscendAttentionState.PrefillNoCache, AscendAttentionState.PrefillCacheHit)
or attn_metadata.block_size != SFA_FULL_VISIBLE_TEMPLATE_BLOCK_SIZE
or attn_metadata.seq_lens_cpu is None
or attn_metadata.seq_lens_cpu.numel() != 1
or attn_metadata.num_actual_tokens != topk_num_tokens
):
return None
table = self._full_visible_index_table
if table is None or table.device != device:
table = self._get_or_create_full_visible_index_table(device)
self._full_visible_index_table = table
kv_length = int(attn_metadata.seq_lens_cpu[0])
context_length = kv_length - attn_metadata.num_actual_tokens
if context_length < 0 or kv_length > SFA_INDEXER_SPARSE_COUNT:
return None
return table[context_length + 1 : kv_length + 1].unsqueeze(1)There was a problem hiding this comment.
Thanks for checking device placement. This table is not a CPU tensor that relies on a later Module.to() call: the backend constructor explicitly selects torch.device("npu", torch.npu.current_device()) and copies the table to that device. The cache is keyed by device. The worker initialization also explicitly selects its NPU before creating the model runner.
For this bounded fast path, a missing table or device mismatch deliberately falls back to the existing LightningIndexer path. Creating a table in this helper would introduce table construction and a device transfer into the request path, which this PR intentionally keeps as a slice/view lookup.
We therefore retain the initialization-time allocation and fallback checks. Arbitrary device migration after backend construction is not accelerated by this path; a mismatch disables the optimization rather than returning indices on the wrong device. This is a source-based explanation, not a new cross-device runtime test claim.
Sources: sfa_v1.py at 213451e lines 538–594; worker/worker.py lines 421–497.
| metadata.seq_lens_cpu = torch.tensor([64]) | ||
| impl._full_visible_index_table = None | ||
| self.assertIsNone(impl._get_full_visible_topk_indices(metadata, 64, torch.device("cpu"))) |
There was a problem hiding this comment.
Update the test to assert that _get_full_visible_topk_indices successfully lazy-initializes the index table when it is initially None, rather than returning None.
| metadata.seq_lens_cpu = torch.tensor([64]) | |
| impl._full_visible_index_table = None | |
| self.assertIsNone(impl._get_full_visible_topk_indices(metadata, 64, torch.device("cpu"))) | |
| metadata.seq_lens_cpu = torch.tensor([64]) | |
| impl._full_visible_index_table = None | |
| result = impl._get_full_visible_topk_indices(metadata, 64, torch.device("cpu")) | |
| self.assertIsNotNone(result) | |
| self.assertIsNotNone(impl._full_visible_index_table) |
There was a problem hiding this comment.
This assertion intentionally protects the fallback contract described in the implementation thread: if the persistent table is missing, the helper returns None and the caller uses the existing LightningIndexer path. Changing the assertion to require lazy initialization would change that contract and add allocation to the request path, so we are keeping the fallback test.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1eaa5d2ff7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # ---- A-family (envs fallback): default = envs module value, before-validator injects ---- | ||
| enable_fused_mc2: int = 0 | ||
| enable_mlapo: bool = True | ||
| enable_sfa_full_visible_index_bypass: bool = False |
There was a problem hiding this comment.
When users want to opt into this user-facing optimization, a repo-wide search finds this switch only in implementation and unit-test code, with no entry in docs/source/user_guide/configuration/additional_config.md or the applicable model guides. Add it to the configuration reference and document its single-sequence, 2,048-token, non-C8 limitations and additional device-memory allocation so users can enable it safely.
AGENTS.md reference: AGENTS.md:L367-L371
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed — this is a user-facing opt-in and the current configuration reference is missing an entry.
The entry needs to cover the default False, an --additional-config example, the single-sequence prefill restriction, the 2,048-token limit on cached context plus new tokens, normal-layer permission, decode/MTP/C8 and other unsupported-path fallbacks, and the persistent per-process/device table allocation (16,785,408 bytes, approximately 16.008 MiB), shared across layers rather than allocated per request.
This remains a documentation follow-up at the current PR HEAD; it does not require changing the runtime algorithm.
There was a problem hiding this comment.
Addressed in the current rebased head 180baa0795cba9801a9a9953017f5fd3e1ac1e70. docs/source/user_guide/configuration/additional_config.md now documents enable_sfa_full_visible_index_bypass, including the default-off opt-in, single-sequence prefill eligibility, total visible KV <= 2048 (including cached context), MTP/C8/CP/unsupported-mode fallbacks, and the shared per-process/device table memory cost. The Markdown link and Read the Docs checks pass on this head.
| context_length = kv_length - attn_metadata.num_actual_tokens | ||
| if context_length < 0 or kv_length > SFA_INDEXER_SPARSE_COUNT: | ||
| return None | ||
| return self._full_visible_index_table[context_length + 1 : kv_length + 1].unsqueeze(1) |
There was a problem hiding this comment.
This returns a row-slice view of a class-level persistent table shared by every layer (and never reset per request). If anything downstream writes into topk_indices in place - e.g. _update_indexcache_topk_indices or a kernel - it would corrupt other layers'/requests' index rows. Worth confirming all consumers are read-only, or returning a clone.
There was a problem hiding this comment.
Thanks for pointing out this potential aliasing issue.
We traced the consumers of the full-visible index tensor. In the inspected bypass path, _update_indexcache_topk_indices uses copy_ into the separately allocated topk_indices_buffer; squeeze only changes the source view. The persistent full-visible table is not the cache-update destination.
We verified whole-table invariance with separate host and native canaries at PR commit 1eaa5d2ff765fa80551536f1990560ae93caed63. The native canary covered 8 cases and 28 synchronized SFA consumer calls, including overlapping short/longer/short slices, two extracted backend objects, and cache update before consumption. The entire 2049 × 2048 table remained exactly equal to an independently allocated pre-consumption snapshot.
These results cover the extracted helper/cache/adapter path with the installed native SFA extension on Ascend910B2 in eager BF16 mode; they do not establish full-model, graph, or future-consumer safety. Based on this bounded evidence, we keep the zero-copy view instead of adding an unconditional clone.
Thanks for the careful review.
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
213451e to
180baa0
Compare
github-actions
Bot
added
documentation
|
The earlier merge conflict is now resolved. This PR is rebased directly on current The remaining |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
…dexer backend Signed-off-by: Ruiqiu Zheng <191817791+Ruiqiu-Zheng@users.noreply.github.com>
180baa0 to
84c10e8
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
2 participants
What this PR does / why we need it?
This PR rebases the opt-in full-visible short-prefill LightningIndexer bypass onto current main after the indexer/SFA ownership split in #15669.
When the entire visible KV window fits within
topk=2048, LightningIndexer scoring cannot remove any valid KV entry. For this bounded case,AscendSFAIndexerBackend.forward()returns a slice of a precomputed round-robin full-visible index table after the normal indexer cache write, while every unsupported or unproven case falls back to the incumbent LightningIndexer path. The rebased patch does not put bypass ownership back intoattention/sfa_v1.py.The feature is disabled by default through
enable_sfa_full_visible_index_bypass=False. Current admission is deliberately fail-closed: single sequence, supported prefill mode, total visible KV<=2048, block size 128, supported model/layer/device/table geometry. MTP/speculative paths, LI C8, PCP/DSA-CP, missing or mismatched metadata/device/table state, visible KV>2048, and other unproven cells keep the native path.The shared RR table is device-keyed and read-only by contract. The worker establishes the NPU device before model/model-runner construction; request-time device mismatch declines the fast path rather than allocating a replacement table. Downstream index-cache publication copies into a separate buffer, so the shared table is not the mutation destination.
Does this PR introduce any user-facing change?
Yes. It documents the new opt-in
enable_sfa_full_visible_index_bypassadditional configuration option. The default isFalse, so existing behavior is unchanged unless explicitly enabled.Validation
The semantic-rebased source is frozen against the R52 candidate. Final-source host hygiene passed (
ruff format --check, Ruff E/F,py_compile, andgit diff --check).Single-NPU native validation on the rebased source covers full-visible cases
(context, query) = (0,64), (128,64), (0,2048). The bypass and native LightningIndexer select the same valid KV set for every row; downstream native SFA outputs are equivalent within the recorded tolerance; repeated calls are exact; the shared table remains unchanged; and>2048falls back. Native LI is allowed to return a different ordering for the same full-visible selected set, so the correctness oracle is selected-set equality plus downstream SFA equivalence rather than elementwise index-order equality.Bounded single-NPU LI+SFA path measurements on the validated path previously showed approximately 10.75% (
K=64), 8.44% (K=1024), and 6.02% (K=2048) latency reduction, with all measured blocks positive. These are component/path measurements only, not full-model TTFT or production-performance claims.Scope / non-claims
This PR does not claim full-model or production improvement, and it does not change MTP/speculative, LI C8, PCP, DSA-CP, or long-context behavior. Unsupported cases remain on the incumbent native path.