Skip to content

[BugFix][Attention] Build SFA indexer DCP metadata independently - #16325

Merged
weijinqian0 merged 1 commit into
vllm-project:mainfrom
Biuapha:fix/sfa-indexer-dcp-metadata
Sep 17, 2026
Merged

weijinqian0 merged 1 commit into
vllm-project:mainfrom
Biuapha:fix/sfa-indexer-dcp-metadata

Conversation

@Biuapha

@Biuapha Biuapha commented Sep 11, 2026 •

Copy link
Copy Markdown
Collaborator

What this PR does / why we need it?

After #15669 separated the SFA indexer cache and metadata builder, replicated DCP can expose an address mismatch: SFA uses a rank-local KV-cache layout, while the indexer cache is replicated across DCP ranks. Building indexer metadata from the SFA-local block table or slot mapping can therefore read or write the wrong indexer cache locations.

This PR makes the indexer builder derive its complete cache view from common attention metadata and the explicitly supplied CP context:

CommonAttentionMetadata + CP context -> indexer builder -> AscendSFAIndexerMetadata

The builder owns the final block table and write-ready slots, RoPE tables, query/key sequence lengths, PCP decode boundary, and LI-C8 cache-write grouping metadata. Mutable derived buffers are keyed by the persistent common slot-mapping buffer so graph capture, replay, and MTP draft steps use the appropriate storage.

Existing parallel calculations are shared through stateless helpers: get_cp_local_query_key_lens and build_pcp_ordered_slot_mapping in context_parallel/common_cp.py, and replicated-DCP address / PCP global-view helpers in context_parallel/sfa_dcp_utils.py. SFA and indexer builders call those helpers and retain their own metadata buffers.

Why sfa_v1.py changes

Change Reason
Move LI-C8 group metadata construction to the indexer builder; remove the old SFA-side construction and field assignments. The grouping must describe the indexer's final slots and block size. Removing the unused SFA-side call also avoids duplicate auxiliary-metadata work; SFA's own C8 KV-write path remains intact.
Look up the sharing target's indexer metadata, then the draft indexer's own metadata; remove fallback to the SFA layer's metadata. KV-sharing draft layers can have their own metadata even when physical cache storage is shared. SFA metadata is not a valid substitute for the indexer's replicated cache view.
Remove forward-time assignments of actual_seq_lengths_query, actual_seq_lengths_key, and num_decode_tokens. These fields are already constructed by the indexer builder. Copying them from SFA would preserve a dependency on SFA's construction and buffer ownership.
Read cos / sin inside indexer forward from indexer metadata instead of passing SFA's tensors. This completes the indexer's metadata ownership, including CP slicing and graph/MTP storage. It reuses the existing RoPE calculation.

The direct DCP bug fix is the replicated cache-address construction. The RoPE interface change and removal of unused SFA auxiliary fields complete the independent-metadata design; they do not introduce a different SFA attention algorithm.

The proposer retains metadata builders for cache-only layers. It identifies these groups through the existing backend contract (get_impl_cls() is None), not a concrete KV-cache-spec type, so an independent tail-cache spec is not mistaken for the primary attention group. The backend query runs in the proposer's vLLM configuration context. During MTP, common sequence/position/slot state advances once through the primary attention group, then each cache-only builder constructs its own metadata from that updated common view. Cache-only groups are excluded from primary graph/backend selection, not from metadata construction. This remains the existing one-main-attention-group path; it does not add generic multi-KV-cache scheduling or heterogeneous full-graph support. The Step3.5 and Gemma4 proposers are unchanged.

Does this PR introduce any user-facing change?

It fixes indexer cache addressing and removes reliance on SFA metadata being used as an indexer fallback or patched into indexer metadata during forward. There is no public API or configuration-format change.

The separate compiled MoE reduction issue is addressed by #16495. The hardware comparison below includes that patch only in an isolated test overlay; it is not included in this PR.

How was this patch tested?

Revision and source checks

  • Latest-main conflict resolution (ec6e301d, 2026-09-15): rebased the single PR commit onto main 0e0910b0. The only textual conflict was the existing _make_runner() helper in tests/ut/worker/test_model_runner_v2.py: main [BugFix][Test] Fix MRV2 CPU UT regression #16574 already supplies adaptive_verification and use_fia, so this PR now retains only the still-missing attn_groups = [] initialization. git range-diff shows no production-logic change from the rebase. Ruff check passed for all 19 changed files, the conflict file passes Ruff format check, git diff --check passes, and the paired vLLM pin remains 84030bbe3d74d99bad477a3d2e37a973ccd8865c.

  • DSpark adaptive graph CI fix (ec6e301d, 2026-09-15): the failed A3 eight-card case was test_glm5_2_dspark_eager[adaptive]; MTP and fixed DSpark passed in the same job. Adaptive FULL graph records its graph-shaped token width in positions.shape[0], while common_attn_metadata.num_input_tokens can retain the compact count. Before the split, indexer RoPE came from SFA, whose builder already handles this contract. The independent indexer builder now applies the same width when constructing its own positions/RoPE/slot view. One focused unit regression was added; Ruff check, targeted format check, and git diff --check pass. Local full pytest was not claimed because the available Windows helper environment lacks the paired vLLM dependency set; CI is the authoritative rerun. No DSpark/model-runner/CI test behavior or assertion was changed. The first CI revision exposed four existing test doubles that omit the optional method field; the final predicate uses getattr(..., None) for that field as well, without changing those fixtures. That CPU-UT run reported 4624 passed / 4 failed before this one-line compatibility correction.

  • Current CI follow-up (f0945801, 2026-09-15): compared with 242e9751, this adds only three initialization lines to the existing _make_runner() test helper in tests/ut/worker/test_model_runner_v2.py: attn_groups = [], adaptive_verification = None, and use_fia = False. No production code or PR base changed; the PR remains one signed-off commit. The previous CPU-UT job reported 4611 passed / 3 failed after CI rebased onto main a643db116. Both the tested runner code and newly added tests came from main, not this PR: their lightweight __new__ fixture omitted fields normally initialized by the real runner. The fix preserves all assertions and changes neither production defaults nor execution paths. Ruff check, format check and git diff --check pass; the new pre-commit/mypy job and CPU-UT job both completed successfully. This does not claim that every NPU matrix job has finished. An independent CPU reproduction used main a643db116 and its paired vLLM 84030bbe3d74d99bad477a3d2e37a973ccd8865c: the complete test_model_runner_v2.py reproduced 33 passed / 3 failed before the fix, then 36 passed / 0 failed after the same three-line patch. The complete test_attn_utils_v2.py also passed 23/23. No assertions were removed, dependencies installed, or NPU probes performed. Logs: /mnt/share/g00672350/refactor/pr16325-ci-242e9751-20260915/112_cpu_model_runner_v2_unfixed.log, 112_cpu_model_runner_v2_fixed.log, and 112_cpu_attn_utils_v2_fixed.log. The archive identities were verified: before cd13ae104, after cbcd83bc3; these reproduce the CI main + PR combination without changing this PR's base. Targeted Ruff checks pass. The repository-wide auto-formatting command was not run because it could modify unrelated files; full pre-commit/mypy remains covered by CI.

  • Earlier mypy fix (242e9751): declared consumes_pcp_context = False on the existing _PrefillStateBuilder test double; no production changes. The pre-commit job including mypy passed. This is previous-run evidence, not a claim that all checks for the new head have passed.

  • Current CPU regression (112, 2026-09-15): the complete 165-item focused Eagle/SFA/indexer suite passed under the repository's standard CPU conftest profile: 155 passed, 10 skipped, 0 failed. It ran on exact source 242e9751 with paired vLLM a97dacb; current production code is identical. Log: /mnt/share/g00672350/refactor/pr16325-validation-242e9751-20260915/112_cpu_eagle_sfa_indexer_standard_profile.log. An earlier A5-profile CPU attempt had two custom-op-dispatch expectation failures; the authoritative result is the independent full rerun under the standard CPU profile, not combined partial results. The priority cache-only group/context set also passed 10/10.

  • Current NPU precision/acceptance (111, 2026-09-15): MTP1, TP8/DP1, DSA-CP off, target FULL_DECODE_ONLY with compilation enabled and breakable disabled: 24/24 passed, including 9/9 original >=2K exact retrieval answers and 4/4 additional exact answers at 12,761 / 25,801 / 38,051 / 39,441 input tokens. The full suite also contains short sanity cases. Raw Prometheus counter deltas were 64 drafts, 64 drafted tokens, 58 accepted tokens: 90.625% draft-token acceptance, mean acceptance length 1.90625 (including the bonus token). This is a small long-input/short-output regression workload, not a matched pre-refactor acceptance-rate baseline or a throughput benchmark. Both SFA-C8 and LI-C8 were enabled; the speculative configuration retained enforce_eager=true for the drafter. No GSM8K.

  • Current MTP5/DSA-CP-off result (111, 2026-09-15): the same TP8/DP1, W4A4C8, FULL_DECODE_ONLY + [BugFix][MoE] Keep runtime reduction inside complete custom op #16495 configuration passed 24/24, including all 9 original >=2K exact retrieval cases and all 4 extra 12,761 / 25,801 / 38,051 / 39,441-token cases. Counter deltas: 41 drafts, 205 drafted tokens, 103 accepted tokens; aggregate draft acceptance 50.2439%, mean acceptance length 3.5122. Acceptance by draft position 1-5 was 92.6829% / 58.5366% / 51.2195% / 39.0244% / 9.7561% (denominator 41 drafts for each position). These small, mostly short-answer requests do not establish equality to the historical interval metrics or to a matched pre-refactor/eager baseline. Evidence directory: /mnt/share/g00672350/refactor/pr16325-validation-242e9751-20260915/111_mtp5_graph_dsa_off_retry1. The initial off attempt failed before readiness because memory from the preceding stopped service had not yet been released; retry began only after all eight cards were healthy and free. No memory threshold was lowered, no device/container was reset, and no unrelated process was stopped.

  • NPU source identity and remaining cases: the tested archive is exact PR revision 242e9751 plus authorized [BugFix][MoE] Keep runtime reduction inside complete custom op #16495 patch d9bae5c4614df6959caf93d9b6f6ad05778fd52f (isolated overlay da1e17c53d7362766162ab0583ea267ba0880f02), paired with vLLM a97dacb7106ee49f39f3d1fc6ae1800ff724e01d. Current PR ec6e301d additionally contains the focused DSpark-adaptive indexer graph-width fix above; the prior NPU results do not validate that new line-level change. Evidence: /mnt/share/g00672350/refactor/pr16325-validation-242e9751-20260915/111_mtp1_graph_dsa_off_retry2/{status.json,results/summary.json,metrics_before.txt,metrics_after.txt,service.log}. MTP5 + DSA-CP-on at TP4/DP2 was rejected before readiness by the paired vLLM's spec-decode/SP graph-shape validation (query length 6 versus TP4); no precision pass or failure is claimed for that configuration. Its failed service was stopped. The MTP5/DSA-CP-off case subsequently passed as reported above. The author approved TP2/DP4 for the MTP5 + DSA-CP-on graph case. Its validation supervisor was externally terminated with exit 137 before readiness; no requests or precision/acceptance result were obtained. Container/host checks did not show OOM. The author confirmed external process termination and requested that validation wait, so further NPU launches/retries are now paused. Only task-owned test processes are being cleaned up; no device/container reset or unrelated process cleanup is performed. The TP2/DP4 result remains unvalidated, and the topology difference must be retained in any later result. 112 cards 1/2 report UB/RAS 81AF8000 alarms, so 112 is used only for CPU tests. NPU validation is now paused again at the author's request; the old recurring task remains paused and will not automatically restart tests.

  • Cache-only group follow-up (061cf2ab4, 2026-09-15): relative to 249c63d3, only llm_base_proposer.py (+20/-14) and its existing Eagle proposer test file (+38/-5) changed. Removed the concrete SFA-indexer-spec check, used the existing backend capability for primary selection and the MTP cache-only group list, and retained the existing single-main-group restriction. No broad proposer/CP/Gemma4 refactor was restored.

  • Extended the existing focused test to cover all six permutations of main attention + SFA indexer + an independent SlidingWindowSpec-shaped cache-only tail group, configuration scope, per-layer metadata ownership, graph backend selection and the no-executable-group error. The test's six cases passed locally against the actual production method ASTs with runtime imports/spec classes/config context isolated. The old HEAD was also checked and does select an independent tail spec as primary when it comes first. This is isolated method-level testing, not a full runtime pytest result.

  • Ruff 0.14.0 check, format --check and git diff --check passed on both files. format.sh ci could not proceed because local pre-commit is not installed. NPU/model validation and the recurring task remain paused; no new precision or acceptance-rate result is claimed.

  • The author's 111 /mnt/share/g00672350/refactor/prci registers KpoolTailSpec from models/glm5next/kv_cache.py, where it derives from SlidingWindowSpec; its layer refers to vllm.v1.attention.backends.mla.indexer.KpoolTailBackend. That backend class was not present in the vLLM source mapped by the specified xhg_refactor_sfa container at inspection time. The new selection path supports such a tail backend when it follows the cache-only get_impl_cls() is None contract; complete KPool runtime/cache-addressing compatibility has not been validated.

  • Previous test-only cleanup (249c63d3, 2026-09-15): production code is unchanged from restored revision 199a459c. Test changes relative to the PR base were reduced from 11 files (+793/-47) to 9 files (+502/-49).

  • Removed the extra standalone common-CP helper tests, KPool wrapper argument-binding test and large proposer dummy-capture scaffold. Merged duplicate PCP decode-boundary, DSA padding and capture-buffer cases through parameterization; reused the existing Model Runner V2 context-propagation test. Direct indexer cache-address/C8/buffer and split-MTP metadata regressions, plus required existing-test signature/mock adaptations, remain.

  • Ruff 0.14.0 check, format --check and git diff --check passed for that test-only cleanup. Runtime pytest and NPU validation were not rerun; validation and the recurring task remain paused at the author's request. CI/UT and hardware results below are historical evidence, not new results for this test-only revision.

  • 2026-09-15 restoration, before the test-only cleanup: restored the exact pre-generalization commit 199a459c7e3b2ac749e05e6163b5dc7bcdc401bc. The broader proposer/Gemma4/MLA-DCP generalization is no longer included. Validation and the recurring task are paused at the author's request; the results below are previously recorded evidence, not newly rerun tests.

  • Single commit: ec6e301d183c5da1023652744b4e9d6395172b86, based on current main 0e0910b060f58a8400397331462c306040dc34e4.

  • At 199a459c, the CPU-UT follow-up changed only two test files relative to e889fc608bf1b60854816d7d9f6bbdd09429a39a; its production files were identical to the hardware-tested pre-squash head c30f4f421e0eea29c043f3ee8131c0333eb5100e. git diff --check passed.

  • The previous CPU-UT run had six failures caused by stale test doubles: four SFA forward cases used the old indexer call signature, and two Eagle cases omitted the indexer cache spec's block_size. The focused recheck passed: SFA forward 4/4; Eagle proposer 8 passed, 10 skipped in the existing 111 container with batch-invariant mode to avoid its missing custom-op registration. Ruff check and format --check passed on both touched tests. The historical PR CI run passed pre-commit and the full CPU-UT job; the six previous test-double failures are gone. Device-test matrix jobs were still running when that evidence was recorded; this is not a current CI status report.

  • Pre-generalization CP-helper checks: six dependency-free common-CP tests and 1,500 exact before/after DSA-CP comparisons passed. They cover all ranks at TP1/2/4/8, main/MTP steps, padding, stable buffer reuse, independent buffers, and SFA/indexer calculations. The PCP orchestration regression passed against both revisions. These checks execute repository method bodies with distributed/NPU imports excluded.

  • Ruff 0.14.0 check, format --check, and AST parsing passed on the eight Python files touched by the final helper cleanup. The full format.sh ci command was unavailable locally because pre-commit was not installed.

  • An earlier 11-file focused pytest run at c9916556c, using vLLM a97dacb7106ee49f39f3d1fc6ae1800ff724e01d, reported 204 passed, 13 skipped, 0 failed. It covered attention/indexer, proposer, worker metadata plumbing and GLM/DeepSeek indexers. This earlier run is not a full pytest result for the final tree.

A5 long-input precision checks (2026-09-14)

Setup: A5 node 112, Ascend 950DT, GLM-5.2 W4A4C8, vLLM main a97dacb7106ee49f39f3d1fc6ae1800ff724e01d, EP enabled, max_model_len=40960, max_num_batched_tokens=4096, max_num_seqs=8, and breakable graph disabled.

The following settings apply to every row in this NPU matrix, both with and without the #16495 overlay, and to the four additional 12K-39K retrieval checks below:

  • MTP: disabled. No speculative-decoding configuration was supplied. These are non-MTP precision results; they do not validate MTP-enabled execution or the MTP + C8 combination.
  • C8: enabled for both SFA and the indexer. The model is W4A4C8, with additional_config.enable_sparse_sfa_c8=True and additional_config.enable_sparse_li_c8=True.
  • The tested Ascend/vLLM source commits are paired correctly. Both Ascend test commits (c30f4f421e0eea29c043f3ee8131c0333eb5100e and overlay 3bc93d4dae7a8fe397d3262880ac9da78ca83375) pin vLLM a97dacb7106ee49f39f3d1fc6ae1800ff724e01d in .github/vllm-main-verified.commit. The NPU checkout has that exact vLLM HEAD, and the service logs confirm imports from /mnt/share/g00672350/refactor/vllm-a97dacb/vllm/. The bot-managed vLLM footer reflects the base branch and is not the version used for these hardware results.

The reference script was used unchanged: eight serial requests, four concurrent requests, then eight concurrent requests. Prompt lengths were 17/14/32/752/1832/2192/2677/3700 tokens. The nine requests at 2,192 / 2,677 / 3,700 tokens were also checked for exact expected-code output.

Code Topology DSA-CP requested / effective Execution Script passes >=2K exact
This PR's tree TP8/DP1 off / off eager 20/20 9/9
This PR's tree TP8/DP1 off / off FULL_DECODE_ONLY 2/20 0/9
This PR's tree TP8/DP1 on / off eager 20/20 9/9
This PR's tree TP8/DP1 on / off FULL_DECODE_ONLY 3/20 0/9
This PR's tree TP4/DP2 on / on eager 20/20 9/9
This PR's tree TP4/DP2 on / on FULL_DECODE_ONLY 20/20 9/9
This PR + #16495 TP8/DP1 off / off FULL_DECODE_ONLY 20/20 9/9
This PR + #16495 TP4/DP2 on / on FULL_DECODE_ONLY 20/20 9/9

All requests returned HTTP 200. The two failing graph rows produced corrupted text; the script's substring check can count a corrupted short answer as a pass, so the exact long-input result is the stronger signal.

With this vLLM configuration, TP8/DP1 makes use_sequence_parallel_moe false and Ascend explicitly disables requested DSA-CP. TP4/DP2 enables it. The table therefore records different topologies rather than implying a controlled DSA-CP toggle at one TP/DP shape.

The overlay was #16495 head d9bae5c4614df6959caf93d9b6f6ad05778fd52f cherry-picked onto the tested pre-squash head, producing test commit 3bc93d4dae7a8fe397d3262880ac9da78ca83375. In its TP8/DP1 graph mode, four additional original-style sequential retrieval prompts at 12,761 / 25,801 / 35,271 / 39,441 tokens all returned the exact expected code with finish_reason=stop.

Evidence on the shared A5 mount: /mnt/share/g00672350/refactor/pr16325-longinput-testkit-112/results_*/summary.json. Overlay service log: /mnt/share/g00672350/refactor/longinput_112_tp8dp1_graph_dsa_off_plus16495_3bc93d.log.

A5 MTP1/MTP5 long-input precision checks (2026-09-14)

A5 node 111, Ascend 950DT, GLM-5.2 W4A4C8, TP8/DP1, EP enabled, DSA-CP off, FULL_DECODE_ONLY graph mode, breakable graph disabled. Both enable_sparse_sfa_c8 and enable_sparse_li_c8 were True. The server used Ascend test overlay 3bc93d4dae7a8fe397d3262880ac9da78ca83375 (this PR + #16495) with its matching vLLM a97dacb7106ee49f39f3d1fc6ae1800ff724e01d; max_model_len=40960, max_num_batched_tokens=4096, and max_num_seqs=8. The engine logs confirmed effective MTP speculative settings of num_spec_tokens=1 and num_spec_tokens=5, respectively, and successful decode-graph capture.

MTP draft tokens Original serial/concurrent cases Original >=2K exact answers Additional exact long inputs
1 20/20 9/9 4/4
5 20/20 9/9 4/4

The additional prompt lengths were 12,761 / 25,801 / 38,051 / 39,441 tokens. Each returned HTTP 200, finish_reason=stop, and content exactly equal to the expected retrieval code; both full 24-case suites passed. The original 20 cases used the same eight serial, four concurrent, and eight concurrent requests as the non-MTP matrix. MTP-enabled hardware evidence covers this TP8/DP1, DSA-CP-off configuration with #16495 applied; it does not establish MTP precision for the PR alone or for DSA-CP-on configurations.

MTP draft acceptance observed during the long-input runs

The server's SpecDecoding metrics log lines show that speculative decoding was active. During the sustained request windows, MTP1 had roughly 90–95% average draft acceptance and 1.9–2.0 mean acceptance length. MTP5 had roughly 57–66% average draft acceptance and 3.9–4.3 mean acceptance length; its per-position acceptance declined from about 90–94% at draft position 1 to about 31–44% at position 5. These are interval metrics from the service logs, not an aggregate over the 24 requests. Alongside the 24/24 exact-output checks for each setting, they show no obvious acceptance collapse. No matched pre-refactor acceptance-rate baseline was collected, so these numbers do not establish acceptance-rate equivalence with the old implementation.

Evidence on the shared A5 mount: /mnt/share/g00672350/refactor/mtp1_graph_111_results_v2/summary.json and /mnt/share/g00672350/refactor/mtp5_graph_111_results/summary.json. Service logs: /mnt/share/g00672350/refactor/mtp1_graph_111.log and /mnt/share/g00672350/refactor/mtp5_graph_111.log.

A5 Model Runner V2 long-input precision checks (2026-09-14)

A5 node 111, the existing xhg_refactor_sfa container, GLM-5.2 W4A4C8, both sparse SFA-C8 and LI-C8 enabled, VLLM_USE_V2_MODEL_RUNNER=1, FULL_DECODE_ONLY graph mode, breakable graph disabled, EP enabled, and the same Ascend #16325 + #16495 test overlay 3bc93d4dae7a8fe397d3262880ac9da78ca83375 paired with vLLM a97dacb7106ee49f39f3d1fc6ae1800ff724e01d. Worker logs explicitly confirmed the V2 model-runner path and successful decode-graph capture. The overlay contains the same PR production code as restored revision 199a459c; #16495 is not part of this PR.

Topology DSA-CP MTP draft tokens Long-input exact-output cases
TP8/DP1 off 0 24/24
TP8/DP1 off 1 24/24
TP8/DP1 off 5 24/24
TP4/DP2 on 0 24/24

Each suite comprised the prior 20 serial/concurrent requests plus four exact-code retrieval prompts of 12,761 / 25,801 / 38,051 / 39,441 tokens. All reported cases returned HTTP 200, finish_reason=stop, and the expected output. These results validate the tested Model Runner V2 graph configurations, not eager mode or this PR without #16495. The TP4/DP2 + DSA-CP-on + MTP1/MTP5 combinations were not run because the A5 nodes became occupied by other services; no result is claimed for them.

Evidence on the shared A5 mount: /mnt/share/g00672350/refactor/mrv2_mtp0_graph_111_results/summary.json, mrv2_mtp1_graph_111_results/summary.json, mrv2_mtp5_graph_111_results/summary.json, and mrv2_tp4dp2_dsa_mtp0_graph_111_results/summary.json. Corresponding service logs use the same names without _results/summary.json and with .log suffix.

Validation limits: the A5 profile rejects replicated SFA DCP before startup, so replicated-DCP address transforms are covered by unit tests rather than this hardware matrix. MTP was disabled in the preceding non-MTP matrix; the separate MTP1/MTP5 checks are reported above. GSM8K and full repository pytest have not been rerun on the final tree; the long-input checks are not a throughput benchmark.

Earlier combined-build GSM8K evidence

These results belong to combined validation commit 5f0326bb3368df11a3067b9fe8c8bc88f9e5463a, base fb00c6be102f31e356bea4ab0f5a64e5091608c6, and vLLM b2f685834a6456197e7033966fdef52a23f1abcd, before the indexer and MoE fixes were separated. They are not final-tree GSM8K results.

Single-node W4A4C8, TP4/DP2, sparse SFA/LI-C8 enabled, MTP and MLAPO disabled, max_num_batched_tokens=4096, max_num_seqs=8, breakable graph disabled:

DSA-CP Execution Long-input script GSM8K
off eager 20/20 178/200 (89.0%)
off FULL_DECODE_ONLY 20/20 178/200 (89.0%)
on eager 20/20 183/200 (91.5%)
on FULL_DECODE_ONLY 20/20 180/200 (90.0%)

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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 addresses a bug in the SFA DCP indexer cache addressing by decoupling the indexer's metadata construction from the primary SFA attention layer. By allowing the indexer to independently build its own replicated block table and slot mapping, the system ensures that the indexer receives the correct addressing even when physical caches are replicated across DCP ranks. The changes include new stateless helper utilities for address transformation and an opt-in mechanism for propagating PCP context, ensuring architectural consistency without impacting existing APIs.

Highlights

  • Independent SFA Indexer Metadata: Decoupled the SFA indexer metadata construction from the main SFA attention layer, allowing it to build its own replicated block table and slot mapping independently.
  • Stateless Address Helpers: Extracted replicated-view address transforms into stateless helpers in a new utility module to support independent metadata generation.
  • PCP Context Propagation: Enabled explicit PCP context opt-in for metadata builders, ensuring correct PCP+DCP ordering for independent indexer builds.
  • Regression Fixes: Restored correct SFA DCP indexer cache addressing and maintained LI-C8 reshape optimization support.
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 Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@github-actions

Copy link
Copy Markdown
Contributor

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.


Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [BugFix] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

  • Fixes #<issue_number>
  • Closes #<issue_number>
  • Resolves #<issue_number>
  • Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

@gemini-code-assist gemini-code-assist Bot left a comment

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.

Code Review

Suggested PR Title:

[Attention][Feature] Support SFA replicated indexer under DCP and PCP

Suggested PR Summary:

### What this PR does / why we need it?
This PR introduces support for the SFA replicated indexer under Decode Context Parallel (DCP) and Prefill Context Parallel (PCP) configurations. It refactors SFA DCP address calculations into a new stateless utility module (`sfa_dcp_utils.py`) and updates `AscendSFAIndexerMetadataBuilder` to independently construct replicated block tables and slot mappings under DCP. Additionally, it enables metadata builders to opt-in to consuming PCP context via a `consumes_pcp_context` attribute, which is now integrated into the attention utility builder.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
New unit tests have been added in `tests/ut/attention/test_indexer.py` and `tests/ut/worker/test_attn_utils_v2.py` to verify SFA indexer metadata building under DCP/PCP and ensure proper propagation of PCP context.

I have no further feedback to provide as there are no review comments.

@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from f56f80c to 148918f Compare September 11, 2026 03:35
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@Biuapha Biuapha added the ready-precise run selected e2e test for pr label Sep 11, 2026
@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from 148918f to 2275e63 Compare September 11, 2026 12:56
@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from 2275e63 to 5f0326b Compare September 11, 2026 13:10
@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from 5f0326b to b43f342 Compare September 11, 2026 14:52
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

1 similar comment
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@drslark drslark left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM.

@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from 920cab0 to 53c01c2 Compare September 15, 2026 16:21
Comment thread vllm_ascend/attention/context_parallel/sfa_dcp_utils.py
Comment thread vllm_ascend/attention/indexer.py Outdated
Comment thread vllm_ascend/attention/indexer.py
Comment thread vllm_ascend/spec_decode/llm_base_proposer.py Outdated
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

@Biuapha
Biuapha force-pushed the fix/sfa-indexer-dcp-metadata branch from edb4d3d to 7659003 Compare September 16, 2026 07:43
Build replicated DCP cache addresses and complete RoPE, parallel sequence,
PCP decode-boundary and LI-C8 metadata in the indexer builder. Reuse
stateless CP helpers while preserving independent graph and draft buffers.

Remove SFA metadata fallback and forward-time backfilling. Preserve split
indexer metadata groups for KV-sharing draft layers and advance common
MTP state once before building each group's metadata.

Add regression coverage for cache layouts, CP calculations, metadata
ownership, draft group handling and indexer call sites.

Signed-off-by: Biuapha <1731372716@qq.com>
if should_update_next_steps:
cache_only_groups = (
[group for group in self.draft_attn_groups if self._is_cache_only_draft_attn_group(group)]
if self.method == "mtp"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

What about here? It is still related to mtp judgement.

@lijiahang226 lijiahang226 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM.

@weijinqian0
weijinqian0 merged commit 8a36237 into vllm-project:main Sep 17, 2026
30 checks passed
tanjiangshan added a commit to tanjiangshan/vllm-ascend that referenced this pull request Sep 17, 2026
- routing matrix for the default gate: non-PD W8A8Dynamic (with and
  without C8) and MXFP8 resolve to PROLOG_V3, unquantized stays NATIVE,
  KV producers take the fused path too (quantized) or stay NATIVE
  (unquantized); the KV-consumer cases and the weight-disposal guard
  are unchanged;
- full-forward UT in test_sfa_prolog_v3_forward.py (mirrors
  test_sfa_nope_forward.py: the real AscendSFAImpl.forward runs on tiny
  dimensions with the NPU ops replaced by pure-PyTorch reference
  implementations, and the output is checked against an independently
  computed expectation):
  - decode-only step: the npu_mla_prolog_v3 reference is invoked once
    with the expected wiring (int64 cache indices, fused weight
    products, MXFP8 quant branch, kv_cache_quant_mode), raw hidden
    states are handed to the indexer k path, and end-to-end numerics
    are verified against the reference chain;
  - prefill step: the fused path is taken as well (P-nodes included);
  - the NATIVE fallback chain is still covered end to end (unquantized
    layers keep the per-layer chain for prefill);
  - two forwards on one metadata share a single int64 slot conversion
    (storage-aliased, no re-cast);
- adapt the indexer reuse test to the vllm-project#16325 forward signature (cos/sin
  now come from the indexer metadata) and update the forward_k mocks in
  test_mla.py to the 3-tuple return.

Signed-off-by: huamus <1943805462@qq.com>
tanjiangshan added a commit to tanjiangshan/vllm-ascend that referenced this pull request Sep 17, 2026
- routing matrix for the default gate: non-PD W8A8Dynamic (with and
  without C8) and MXFP8 resolve to PROLOG_V3, unquantized stays NATIVE,
  KV producers take the fused path too (quantized) or stay NATIVE
  (unquantized); the KV-consumer cases and the weight-disposal guard
  are unchanged;
- full-forward UT in test_sfa_prolog_v3_forward.py (mirrors
  test_sfa_nope_forward.py: the real AscendSFAImpl.forward runs on tiny
  dimensions with the NPU ops replaced by pure-PyTorch reference
  implementations, and the output is checked against an independently
  computed expectation):
  - decode-only step: the npu_mla_prolog_v3 reference is invoked once
    with the expected wiring (int64 cache indices, fused weight
    products, MXFP8 quant branch, kv_cache_quant_mode), raw hidden
    states are handed to the indexer k path, and end-to-end numerics
    are verified against the reference chain;
  - prefill step: the fused path is taken as well (P-nodes included);
  - the NATIVE fallback chain is still covered end to end (unquantized
    layers keep the per-layer chain for prefill);
  - two forwards on one metadata share a single int64 slot conversion
    (storage-aliased, no re-cast);
- adapt the indexer reuse test to the vllm-project#16325 forward signature (cos/sin
  now come from the indexer metadata) and update the forward_k mocks in
  test_mla.py to the 3-tuple return.

Signed-off-by: huamus <1943805462@qq.com>
tanjiangshan added a commit to tanjiangshan/vllm-ascend that referenced this pull request Sep 18, 2026
- routing matrix for the default gate: non-PD W8A8Dynamic (with and
  without C8) and MXFP8 resolve to PROLOG_V3, unquantized stays NATIVE,
  KV producers take the fused path too (quantized) or stay NATIVE
  (unquantized); the KV-consumer cases and the weight-disposal guard
  are unchanged;
- full-forward UT in test_sfa_prolog_v3_forward.py (mirrors
  test_sfa_nope_forward.py: the real AscendSFAImpl.forward runs on tiny
  dimensions with the NPU ops replaced by pure-PyTorch reference
  implementations, and the output is checked against an independently
  computed expectation):
  - decode-only step: the npu_mla_prolog_v3 reference is invoked once
    with the expected wiring (int64 cache indices, fused weight
    products, MXFP8 quant branch, kv_cache_quant_mode), raw hidden
    states are handed to the indexer k path, and end-to-end numerics
    are verified against the reference chain;
  - prefill step: the fused path is taken as well (P-nodes included);
  - the NATIVE fallback chain is still covered end to end (unquantized
    layers keep the per-layer chain for prefill);
  - two forwards on one metadata share a single int64 slot conversion
    (storage-aliased, no re-cast);
- adapt the indexer reuse test to the vllm-project#16325 forward signature (cos/sin
  now come from the indexer metadata) and update the forward_k mocks in
  test_mla.py to the 3-tuple return.

Signed-off-by: huamus <1943805462@qq.com>
weiguihua2 pushed a commit that referenced this pull request Sep 18, 2026
### What this PR does / why we need it?

A5 SFA DCP decode can select fewer keys on a rank than its local KV
length. The sparse index tensor then contains a valid prefix followed by
`-1` padding. Using the KV length as the softmax extent includes invalid
entries and produces incorrect per-rank normalization for the DCP merge.

This patch ports the focused fixes from our standalone SFA operator
experiment into the existing custom operator:

- Bound sparse-mode-0 LSE computation by the actual valid index prefix.
- Initialize skipped queries to zero output and zero softmax max/sum
(reconstructed LSE is negative infinity), including queries with a
nonempty cache but no selected local keys.
- Initialize `n2Size` before output initialization.
- Make prefill KV and RoPE views contiguous after the packed DCP gather.
- Advertise `SFA_DCP_REPLICATED_INDEXER` on A5 while retaining the
platform capability check.

The current custom operator already accepts PA_BSND with
`return_softmax_lse`, so no host-tiling change is needed on this base.
No Mooncake changes or bulk operator-source refresh are included.

### Does this PR introduce _any_ user-facing change?

Enables the A5 SFA DCP capability and fixes its BF16 KV-cache path. The
validated target is full GLM5.2 with W4A8 weights and **both SFA C8 and
indexer C8 disabled**, TP8/DCP8/EP, eager execution, no MTP or PD
disaggregation. This PR does not claim C8 support.

The full-model experiment also applied #16325 (tested revision
`ec6e301d183c5da1023652744b4e9d6395172b86`) for indexer metadata. Those
changes are not included here; the newer revision of that PR has not
been validated as part of this work.

### How was this patch tested?

- Historical experiment: the equivalent numerical fixes in standalone
ops-transformer, together with contiguous prefill inputs, the capability
bypass and the above #16325 revision, passed all six full 78-layer
GLM5.2 functional cases on A5 (arithmetic, factual answer, translation,
Python generation, long-context retrieval and repeated arithmetic).
- This is a minimal port to the custom operator on main `4c5ee3320`, not
the exact operator build used in that experiment. **The port has not yet
been recompiled or rerun on NPU.** Historical results are not evidence
of a completed regression for this exact commit.
- Added an A5 operator regression comparing output and LSE against a
PyTorch reference for 0, 1, 127, 128, 129 and 257 selected keys,
including a nonempty cache with no selected keys.
- Added a prefill regression using non-contiguous packed-cache split
views; updated the hardware capability expectation.
- Ruff 0.14.0 check/format, Python syntax parsing and `git diff --check`
passed. The repository-wide `bash format.sh ci` could not run because
pre-commit is not installed in this local environment. New unit/NPU
tests have not been executed in this PR preparation.

Performance, graph mode, MTP, C8 and PD combinations remain outside the
completed validation scope.

- vLLM main:
vllm-project/vllm@84030bb

---------

Signed-off-by: chengruiqi (C) <c00913489@china.huawei.com>
Co-authored-by: chengruiqi (C) <c00913489@china.huawei.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants