Skip to content

[NPU] Support Kimi-K2-Thinking model - #399

Open
yuxinshan wants to merge 1 commit into
vllm-project:ascendfrom
yuxinshan:kimi_k2_thinking
Open

[NPU] Support Kimi-K2-Thinking model#399
yuxinshan wants to merge 1 commit into
vllm-project:ascendfrom
yuxinshan:kimi_k2_thinking

Conversation

@yuxinshan

@yuxinshan yuxinshan commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Summary

Support Kimi-K2-Thinking model.
Related issue: #395

Architecture

To solve the problem of the parallel mode:

  • docker/npu_patch/megatron-bridge.patch

To solve the problem of the megatron to hf:

  • vime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py
  • vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py

To support Kimi-K2-Thinking model:

  • vime/backends/megatron_utils/model_provider.py
  • vime/utils/megatron_bridge_utils.py

Validation

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces NPU support and compatibility enhancements, including a pure PyTorch fallback for int4 quantization, safe NPU device routing, and a patch for handling Kimi K2 architectures. The review feedback highlights critical issues that need to be addressed: correcting the invalid deepdim parameter to keepdim in the quantization fallback function to prevent runtime errors, safely checking for NPU availability to avoid AttributeError on non-NPU environments, and adding a defensive check for None bridge instances.

Comment thread vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py Outdated
Comment thread vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py Outdated
Comment thread vime/utils/megatron_bridge_utils.py
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