Skip to content

feat: support and optimize deepseek v4 MTP for mlu.#1952

Open
phantomlei3 wants to merge 14 commits into
xLLM-AI:mainfrom
phantomlei3:feat/deepseek_v4_mtp
Open

feat: support and optimize deepseek v4 MTP for mlu.#1952
phantomlei3 wants to merge 14 commits into
xLLM-AI:mainfrom
phantomlei3:feat/deepseek_v4_mtp

Conversation

@phantomlei3

@phantomlei3 phantomlei3 commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Description

  • Merge compress_state / compress_index_state: The new MLU compressor kernel now receives a single continuous state_cache and a single block table. The legacy split KV/score tensors can no longer directly satisfy this interface. NPU will continue to use the split branch, keeping the shape, dtype, ND format, and getter completely unchanged.

  • SWA (Sliding Window Attention) window release requires reserving an additional num_speculative_tokens slots. This is to ensure compatibility with the MTP strategy during batch expansion. In batch expansion scenarios, we must prevent the blocks at the left boundary of the window from being released prematurely. This additional reserved space for the batch will not affect NPU functionality.

  • MTP hidden state alignment: We aligned the MTP hidden state logic with the DeepSeek-V4 MTP implementation from the vLLM/SGLang community. The correct approach should utilize the pre-hc-head flattened hidden state. Using the previous logic resulted in a noticeably lower acceptance rate. Currently, this fix has only been applied to USE_MLU. The NPU side should also verify whether this modification needs to be synced.

  • 合并 compress_state / compress_index_state: 新版 MLU compressor kernel 接收单个连续 state_cache 和单个 block table,旧的 KV/score 分离 tensor 无法直接满足接口。 NPU 仍走 split 分支,shape dtype、ND format 和 getter 都保持不变

  • SWA窗口释放需要额外预留num_speculative_tokens个位置,这里是为了兼容batch expansion的MTP策略,因为在batch expansion场景下要避免窗口左边界 block 过早释放,后续bat额外增加的预留空间不会影响NPU功能。

  • MTP hidden state对vllm/sglang社区的deepseek v4 MTP逻辑进行对齐,正确的方法应该使用pre-hc-head flattened hidden,用之前的逻辑会发现接受率有明显过低的问题,但目前只给USE_MLU做了,需要NPU这边也检查一下是否需要一起修改

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

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.

2 participants