Skip to content

Force blksize=128 for linear attention on ascend#4705

Merged
lvhan028 merged 3 commits into
InternLM:mainfrom
jinminxi104:force_128_blksize
Jun 25, 2026
Merged

Force blksize=128 for linear attention on ascend#4705
lvhan028 merged 3 commits into
InternLM:mainfrom
jinminxi104:force_128_blksize

Conversation

@jinminxi104

Copy link
Copy Markdown
Collaborator

Force blksize=128 for linear attention on ascend

Motivation

Please describe the motivation of this PR and the goal you want to achieve through this PR.

Modification

Please briefly describe what modification is made in this PR.

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

  1. Pre-commit or other linting tools are used to fix the potential lint issues.
  2. The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
  3. If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
  4. The documentation has been modified accordingly, like docstring or example tutorials.

Copilot AI review requested due to automatic review settings June 25, 2026 04:49

Copilot AI 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.

Pull request overview

This PR updates the PyTorch executor’s KV-cache block-size adjustment logic to special-case Ascend devices when running models that use linear attention (detected via non-empty states_shapes), forcing the KV cache block size to 128.

Changes:

  • Add an Ascend + linear-attention guard in ExecutorBase._adjust_block_size() that warns and overrides cache_config.block_size (and kernel_block_size) to 128.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread lmdeploy/pytorch/engine/executor/base.py
Comment thread lmdeploy/pytorch/engine/executor/base.py Outdated
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

@grimoire grimoire 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

@lvhan028 lvhan028 merged commit d4fc277 into InternLM:main Jun 25, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants