Skip to content

feat: add layerwise split kv cache sharding. - #2260

Open
shifengmin wants to merge 5 commits into
xLLM-AI:mainfrom
shifengmin:bugfix/layerwise-kv-broadcast-tp16-validated_final
Open

feat: add layerwise split kv cache sharding.#2260
shifengmin wants to merge 5 commits into
xLLM-AI:mainfrom
shifengmin:bugfix/layerwise-kv-broadcast-tp16-validated_final

Conversation

@shifengmin

@shifengmin shifengmin commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add --layerwise_split_size to shard persistent KV by a static layer owner inside each attention TP group, while non-owners share one scratch layer so PD-Disaggregation stay compatible.
  • Depends on the companion ATB PR that implements owner attention.

Dependency

Depends on xLLM-AI/xllm-atb-layers#14

Related Issues

Change Type

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

Pull Request Checklist

PR Title and Commit Messages

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

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.

Changed-file validation passed:

pre-commit run --files tests/core/runtime/acl_graph_executor_test.cpp \
  xllm/core/runtime/acl_graph_executor_impl.cpp \
  xllm/core/runtime/acl_graph_executor_impl.h
clang-format.............................................................Passed

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: a full python setup.py build --device npu completed successfully on an Ascend development machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

This intentionally falls back to eager execution instead of caching one graph per DP distribution. The latter can create an unbounded graph cache for normal scheduling variance.

Shard persistent KV by static layer owner and keep one shared scratch layer so PD cache ids stay aligned. Prefill, decode, and spec-verify gather queries to the owner, run attention on the owner's paged cache, and broadcast the output.

Reject incompatible layerwise_split_size configs at startup. Disable layerwise split only for draft workers.

Co-authored-by: Cursor <cursoragent@cursor.com>
shifengmin and others added 4 commits August 18, 2026 15:56
Pin the submodule to xLLM-AI/xllm-atb-layers#14 on main (523a71d), including the default-path graph gate.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant