Skip to content

refactor: move npu context parallel partition into model side. - #2003

Merged
shifengmin merged 14 commits into
xLLM-AI:mainfrom
shifengmin:feat/npu-cp-model-closure
Jul 27, 2026
Merged

refactor: move npu context parallel partition into model side.#2003
shifengmin merged 14 commits into
xLLM-AI:mainfrom
shifengmin:feat/npu-cp-model-closure

Conversation

@shifengmin

@shifengmin shifengmin commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator
  • Move input padding and output reorder to after embedding and after the last decoder layer via a model-side CP closure (npu_cp_plan.h), so the scheduler, worker output processing, and MTP no longer special-case CP.
  • Delete the legacy worker-side cp_input_partition and its test

Description

Related Issues

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Allowed types: feat, bugfix, docs, test, refactor, chore, style, revert, perf, model, build, release.
The subject should use clear English, start with a verb, include at least 4 words, and end with ..

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

If you are unsure how to set up pre-commit, see the pre-commit documentation.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

@gemini-code-assist

gemini-code-assist Bot commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

shifengmin added a commit to shifengmin/xllm that referenced this pull request Jul 23, 2026
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
d661e59 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
shifengmin added a commit to shifengmin/xllm that referenced this pull request Jul 23, 2026
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
2685a4b 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from 6fd3e60 to 7f34a96 Compare July 23, 2026 05:50
@shifengmin
shifengmin marked this pull request as ready for review July 24, 2026 02:43
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from 7f28e51 to 77c9d4d Compare July 24, 2026 02:44
shifengmin added a commit to shifengmin/xllm that referenced this pull request Jul 24, 2026
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
2e1cbc3 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from 77c9d4d to f6cf582 Compare July 24, 2026 08:40
shifengmin added a commit to shifengmin/xllm that referenced this pull request Jul 27, 2026
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
7f6690b 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from 948c9ba to 3383cfe Compare July 27, 2026 06:35
shifengmin added a commit to shifengmin/xllm that referenced this pull request Jul 27, 2026
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
2dbc776 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from b1c0c93 to 5b91ef9 Compare July 27, 2026 09:10
shifengmin and others added 4 commits July 27, 2026 17:34
@shifengmin @cursoragent
refactor: move npu context parallel partition into model side.
74ada3a 
- Move input padding and output reorder to after embedding and after the
  last decoder layer via a model-side CP closure (npu_cp_closure.h), so the
  scheduler, worker output processing, and MTP no longer special-case CP.
- Delete the legacy worker-side cp_input_partition and its test (Phase E).
- Fix SFA prev/next split to use local_padded seq lens so short or unaligned
  prompts no longer hit c10::IndexError; clamp gather indices defensively.
- Loosen NPU validate_model_cp to allow tp_size>1 with dp_size==1 (CP x TP
  orthogonal), requiring npu_kernel_backend=ATB and enable_graph=false.
- Isolate TP1 draft CP plan/tensors/slots from the target CP group, and route
  DeepSeek V3.2 CP draft to deepseek_v32_mtp.
- Add unit tests for cp_group_ranks, npu_cp_capability, npu_cp_closure,
  ep_padding_plan, and prefill_plan.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
fix: disable LM head CP allgather after model-side gather_restore
35addbb 
Model-side CP already restores global-real hidden before the LM head
(npu_cp::gather_restore after the last decoder layer). PR xLLM-AI#2003 cleared
outputHidden but left contextParallelInfo = ATTN_CP enabled, so the ATB
LmHead graph still built a nested prefill CP AllGather on the already-
global hidden (gated by gatherAhead && contextParallelInfo.IsEnabled(),
independent of outputHidden). That second AllGather treated the full
sequence as a CP shard and scaled hidden by cp_size again, peaking at
O(cp_size * T * H) and OOM'ing on long prefill (e.g. cp=32, T~30000,
H=6144, bf16 -> ~10.99 GiB).

Drop the ATTN_CP assignment so contextParallelInfo stays default-
constructed (rankIds empty -> IsEnabled()==false) and the ATB LmHead
graph skips the CP AllGather node. gatherAhead (prefill indices Gather)
and ATTN_TP (vocab column parallel) are unchanged.

See docs/pr2003_lmhead_nested_cp_allgather_bug.md.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
bugfix: build ATB attention CP metadata from local_padded seq lens.
6f00676 
The ATB-facing attention metadata (q_seq_lens / kv_seq_lens / q_cu_seq_lens /
q_max_seq_len / kv_max_seq_len) must be built from local_padded_seq_lens so the
lengths stay consistent with the hidden row count (local_padded_token_num),
CpEpPadding, and SFA actual_seq_lengths. The legacy code used local_real, which
mismatched for non-aligned prompts where local_real < local_padded: the ATB
graph sized AllGather / attention buffers by kv_seq_lens (real) while the
actual KV data held local_padded rows, causing SDMA DDR out-of-range crashes
and garbled output. local_real is retained in the plan for restore / sampling /
logging but must not be written into the decoder variant pack.

Add CPU comparison unit tests (npu_cp_prefill_plan_cmp_test,
npu_cp_ep_padding_cmp_test, npu_cp_slots_cmp_test) covering the prefill plan,
EP padding, and slots / gather-index computation under chunked prefill + prefix
cache + kvsplit.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
refactor: shard npu embedding and lm head by dp-local tp to close cp …
0873340 
…boundary.

Model-side CP localizes hidden after embedding and restores global-real
hidden before the lm head, so both layers always see the full sequence and
must be CP-unaware but DP-aware. Previously they sharded by the CP-local TP
(size tp), replicating the weights across all CP ranks. Now shard across the
dp-local-TP group (world / dp_size = cp_size * tp_size) -- the full TP width
within one DP -- so weights are partitioned across cp*tp ranks instead of
replicated.

- mapping_npu: word_embed_tp_/lm_head_tp_ group_size = world/dp = cp*tp;
  to_json emits lm_head_tp_ (was attn_tp_).
- word_embedding/lm_head/column_parallel_linear impls: use WORD_EMBED_TP /
  LM_HEAD_TP (ATB path) and manual dp-local-TP (TORCH path).
- loaders + get_padded_vocab_size + NpuLmHeadImpl padded_vocab_size_: shard
  and pad by world/dp = cp*tp.
- cp_size == 1 collapses to tp_size (backward compatible).
- LM_HEAD_TP is only consumed by NpuLmHeadImpl; decoder last-layer AllGather
  uses MLP_TP (enableDpOut=false), so no double gather with gather_restore.

Validated on prefill-1 (cp=32, ep=32, epd=2, MTP=3, kv_split=4): short chat
and ~32k needle-in-haystack pass with no garbled output.

Co-authored-by: Cursor <cursoragent@cursor.com>
shifengmin and others added 8 commits July 27, 2026 17:34
@shifengmin @cursoragent
refactor: consolidate npu cp plan into plan helper.
84f6077 
Fold scattered CP prefill/ep-padding fields into npu::cp::Plan and prepare_cp_input so model/runtime share one plan accessor path.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
refactor: unify npu cp metadata planning.
dbfad25 
Consolidate CP input sharding, attention metadata, CP-EP metadata, cache-slot preparation, and output merging into NpuCpPlan with one worker-owned build entry.

Replace the scattered compatibility helpers and legacy names, preserve graph tensor byte parity with focused tests, and ensure MTP target and draft leaves each build and shard exactly once.

Validation: clang-format --dry-run --Werror; git diff --check; online prefill test passed (confirmed by user).
Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
refactor: fold npu cp plan prepare into plan and guard model insertion.
5ddae1d 
Move WorkerImpl::prepare_npu_cp_plan into NpuCpPlan::prepare(), fed by a
cached per-worker CpPlanRuntimeConfig, so the worker calls the plan in one
line after the global attention-metadata consumers. shard_model_input now
rewrites hidden/positions in place, merge_model_output uses the cp_group
bound at prepare() time, and the worker applies attention metadata in
place -- models drop modified_input_params and guard the CP path with
if (cp_plan.enabled()), leaving the non-CP forward untouched.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
fix: link parallel_state in spec_input_builder_test.
94ef840 
NpuCpPlan::to is now pulled in by ModelInputParams::to after the CP
refactor, so the test needs :parallel_state to link.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin @cursoragent
style: trim verbose comments on the npu cp path.
313e485 
Remove historical Legacy/refactor essays and compress remaining CP
comments to short one-liners so the model-side CP path stays readable.

Co-authored-by: Cursor <cursoragent@cursor.com>
@shifengmin
refactor: isolate legacy expert array preparation for glm4 mtp.
92ee07e 
- reuse the original model input parameters across the MTP forward path
- prepare the legacy expert array only for GLM4 MoE MTP
@shifengmin
bugfix: preserve mlu cp validation behavior.
18a2934
@shifengmin
bugfix: fix npu backend validation in mlu builds.
e253c01
@shifengmin
shifengmin force-pushed the feat/npu-cp-model-closure branch from 5b91ef9 to e253c01 Compare July 27, 2026 09:35
DongheJin
DongheJin previously approved these changes Jul 27, 2026
View reviewed changes
@shifengmin
bugfix: fix kv split cache slot ownership.
6700682 
- derive shard ownership from logical cache slots

- cover multi-sequence and chunked-prefill layouts
@shifengmin
shifengmin merged commit a865489 into xLLM-AI:main Jul 27, 2026
7 of 26 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Kang-Meng Kang-Meng Kang-Meng approved these changes
@DongheJin DongheJin DongheJin approved these changes
@yq33victor yq33victor Awaiting requested review from yq33victor yq33victor is a code owner
@liutongxuan liutongxuan Awaiting requested review from liutongxuan liutongxuan is a code owner
@walsonyang walsonyang Awaiting requested review from walsonyang walsonyang is a code owner
@RobbieLeung RobbieLeung Awaiting requested review from RobbieLeung RobbieLeung is a code owner
@JimHsiung JimHsiung Awaiting requested review from JimHsiung JimHsiung is a code owner
@XuZhang99 XuZhang99 Awaiting requested review from XuZhang99 XuZhang99 is a code owner
@yingxudeng yingxudeng Awaiting requested review from yingxudeng yingxudeng is a code owner
@DragonFive DragonFive Awaiting requested review from DragonFive DragonFive is a code owner
@Clement-Wang26 Clement-Wang26 Awaiting requested review from Clement-Wang26 Clement-Wang26 is a code owner
@zhang-minchao zhang-minchao Awaiting requested review from zhang-minchao zhang-minchao is a code owner
@xiao-yu-chen xiao-yu-chen Awaiting requested review from xiao-yu-chen xiao-yu-chen is a code owner
@liujinguang0125 liujinguang0125 Awaiting requested review from liujinguang0125 liujinguang0125 is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@shifengmin @Kang-Meng @DongheJin