Skip to content

feat(pytorch): support Kimi K2.6 - #4846

Open
qescccczmr wants to merge 8 commits into
InternLM:mainfrom
qescccczmr:feat/kimi-k2.6-pytorch
Open

feat(pytorch): support Kimi K2.6 #4846
qescccczmr wants to merge 8 commits into
InternLM:mainfrom
qescccczmr:feat/kimi-k2.6-pytorch

Conversation

@qescccczmr

@qescccczmr qescccczmr commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Add Kimi-K2.6 support to the PyTorch backend, including text and multimodal inference, compressed-tensors W4A16 MoE, EAGLE3 speculative decoding, and serving integration.

The implementation supports Kimi-K2.6's MLA attention, MoE routing, vision encoder, multimodal projector, and external EAGLE3 draft model.

Changes

  • Add Kimi-K2.6 configuration, model loading, native language model, vision tower, and multimodal projector.
  • Add compressed-tensors W4A16 MoE model construction, weight loading, and Triton kernels.
  • Add Kimi-K2.6 EAGLE3 speculative decoding with shared target embeddings and auxiliary hidden states.
  • Add coverage for compressed-tensors W4A16 kernels, MoE routing, FA3 attention, model loading, and serving.

Acc Benchmarks

Academic

Dataset Metric LMDeploy HF reference
GPQA Diamond Accuracy, TP8, 4-run average, CompassAcademic 91.67% 90.50%
AIME 2026 Accuracy, TP8 97.60% 96.40%

LMDeploy and Hugging Face were evaluated with TP8 under aligned evaluation settings.

LMDeploy reaches or slightly exceeds the Hugging Face reference on both academic benchmarks, with no observed text accuracy regression.

Multimodal

Dataset Metric LMDeploy HF reference
MathVision Accuracy 84.54% 87.40%

Kimi-K2.6 successfully completes the MathVision multimodal evaluation.

Acceptance of EAGLE3

Draft model:

lightseekorg/kimi-k2.6-eagle3.1-mla

Aggregate results

Engine Mean accept length Strict draft acceptance
LMDeploy 2.93204 64.4015%
SGLang 2.93131 64.3769%
LMDeploy − SGLang +0.00074 +0.0246 pp

Per-position acceptance rates:

Engine Position 0 Position 1 Position 2
LMDeploy 81.43% 63.41% 48.36%
SGLang 81.52% 63.21% 48.40%

Runtime Validation

  • Kimi-K2.6 TP8 text inference passed on 8 NVIDIA H200 GPUs.
  • Kimi-K2.6 TP8+EP8 text inference completed model loading, warmup, and generation successfully.
  • EAGLE3 completed all 30 AIME24 requests without engine errors.
  • MathVision multimodal evaluation completed successfully.

Assistance

Assisted with Codex, reviewed manually.

@qescccczmr
qescccczmr force-pushed the feat/kimi-k2.6-pytorch branch 4 times, most recently from dd00426 to 707e933 Compare August 17, 2026 11:13
@qescccczmr
qescccczmr force-pushed the feat/kimi-k2.6-pytorch branch from 707e933 to fb64e19 Compare August 17, 2026 13:15
@qescccczmr
qescccczmr force-pushed the feat/kimi-k2.6-pytorch branch 3 times, most recently from fb64e19 to 33315ea Compare August 19, 2026 11:34
@qescccczmr qescccczmr changed the title Feat/kimi k2.6 pytorch feat(pytorch): support Kimi K2.6 inference and EAGLE3 Aug 19, 2026
@qescccczmr qescccczmr changed the title feat(pytorch): support Kimi K2.6 inference and EAGLE3 feat(pytorch): support Kimi K2.6 Aug 19, 2026
@qescccczmr
qescccczmr marked this pull request as ready for review August 19, 2026 14:11
Comment thread lmdeploy/hf_configs/__init__.py
Comment thread lmdeploy/pytorch/kernels/cuda/fused_noaux_tc.py Outdated
Comment thread lmdeploy/pytorch/kernels/cuda/__init__.py Outdated
Comment thread lmdeploy/pytorch/quantization/__init__.py Outdated
Comment thread lmdeploy/pytorch/spec_decode/spec_agent.py Outdated
Comment thread tests/pytorch/kernel/test_fa3_attention.py Outdated
# DeepseekV2ModelConfigBuilder also recognizes ``kimi_k2``. Register the
# narrower Kimi builder immediately before it without changing the shared
# DeepSeek builder.
_builders = AutoModelConfigBuilder._sub_classes

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.

You don't have to patch the builder of dsv2. You can just remove 'kimi_k2' in it's condition.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I've removed kimi_k2 from DeepseekV2ModelConfigBuilder.

@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

@grimoire

Copy link
Copy Markdown
Collaborator

Please fix lint

@qescccczmr
qescccczmr requested a review from lvhan028 August 24, 2026 02:55
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