feat: support and optimize deepseek v4 MTP for mlu.#1952
Open
phantomlei3 wants to merge 14 commits into
Open
Conversation
- unify packed and split compression state access\n- preserve fallback and block swap behavior with regression tests
- support torch_mlu_ops 1.12.2 causal-conv null-block arguments\n- pass new TF32 and state-layout flags through gated-delta kernels
b0bdb15 to
632a2c5
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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
<type>: <subject>.Pre-commit Checks
pre-commitby runningpip install pre-commitor an equivalent command.pre-commit install.pre-commit run --all-filesand fixed any reported issues.Self Review
.agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.mainbranch.Build and Test Coverage
python setup.py build testhas passed on a CUDA machine.python setup.py build testhas passed on an NPU machine.python setup.py build testhas passed on an MLU machine.Reviewer Notes