Skip to content

Add Kimi-k2-thinking and k2.5 and k2.6 checkpoint conversion support.#3768

Merged
copybara-service[bot] merged 1 commit intomainfrom
gagik-kimi
Apr 30, 2026
Merged

Add Kimi-k2-thinking and k2.5 and k2.6 checkpoint conversion support.#3768
copybara-service[bot] merged 1 commit intomainfrom
gagik-kimi

Conversation

@gagika
Copy link
Copy Markdown
Collaborator

@gagika gagika commented Apr 28, 2026

Description

  • Extends convert_deepseek_family_ckpt.py to convert HuggingFace Kimi checkpoints (kimi-k2-thinking, kimi-k2.5-text,
    kimi-k2.6-text) to the MaxText DeepSeek family layout.
  • Adds an int4 dequantization helper (dequantize_pack_quantized_int4) for the routed-expert projections, which ship as
    compressed-tensors pack-quantized weights with per-group symmetric scales (group_size=32).
  • K2.5 and K2.6 are extracted from their multimodal KimiK25ForConditionalGeneration wrapper — the text branch shape and quantization layout are identical, so K2.6 reuses the K2.5 key map.

Tests

  • pytest tests/unit/dequantize_pack_quantized_int4_test.py
  • Convert each variant end-to-end and verify forward-pass logits match HF on a single prompt
  • Kimi k2.5 logit matching KL divergence = [0.00043962 0.00289483]: https://paste.googleplex.com/6351128208998400

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@github-actions
Copy link
Copy Markdown

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

@gagika gagika requested a review from igorts-git as a code owner April 29, 2026 14:31
@github-actions
Copy link
Copy Markdown

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

Copy link
Copy Markdown
Collaborator

@richjames0 richjames0 left a comment

Choose a reason for hiding this comment

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

Great testing Gagik. One nth comment

@github-actions
Copy link
Copy Markdown

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions
Copy link
Copy Markdown

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

Copy link
Copy Markdown
Collaborator

@shuningjin shuningjin left a comment

Choose a reason for hiding this comment

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

Thank you for expanding support to multiple variants (K2-Thinking, K2.5, K2.6) on top of Kimi-K2! The int4 dequantization with unit tests, checkpoint conversion, and user guide updates all look great.

@copybara-service copybara-service Bot merged commit ff916b8 into main Apr 30, 2026
29 checks passed
@copybara-service copybara-service Bot deleted the gagik-kimi branch April 30, 2026 17:11
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