Skip to content

feat: add native Kimi-K3 inference support - #93

Open
earlvanze wants to merge 5 commits into
FlashML-org:mainfrom
earlvanze:feat/kimi-k3-support
Open

feat: add native Kimi-K3 inference support#93
earlvanze wants to merge 5 commits into
FlashML-org:mainfrom
earlvanze:feat/kimi-k3-support

Conversation

@earlvanze

@earlvanze earlvanze commented Aug 23, 2026

Copy link
Copy Markdown

Closes #92.

What changed

  • adds native Kimi-K3 config parsing, model registration, KDA/MLA attention,
    latent-MoE routing, SiTU, and checkpoint weight mapping
  • supports the official full configuration and the public eight-layer
    inference-optimization/Kimi-K3-0.40B-MXFP4 development checkpoint
  • supports TP=1 with MXFP4 routed experts through the offload-family backends
  • makes hybrid KDA/MLA KV allocation layer-aware
  • bounds CUDA graph batch sizes against the available MoE cache slots
  • makes the offload copy plan size from cache capacity instead of expert count
  • loads custom Hugging Face tokenizers with trust_remote_code=True and adds
    the required tiktoken dependency
  • selects naive request caching for Kimi until linear-attention state snapshots
    can support radix reuse safely

Validation

  • 61 focused model, weight-loader, KV-pool, MoE, graph, and tokenizer tests
    passed while CUDA was healthy
  • 22/22 config, loader, graph-batch, and tokenizer tests passed again after the
    benchmark host lost GPU access
  • exact checkpoint audit loaded 204 resident tensors with matching shapes and
    all seven routed-expert banks from revision
    ef3902c318fb8e13c3507e26055656e687fdfe38
  • OpenAI-compatible streaming generation ran end to end on an RTX 3090 through
    32 simultaneous requests

Raw metrics, methodology, host/checkpoint metadata, and the disclosed host
failure are preserved in the
benchmark artifact.

The benchmark uses the 0.40B development checkpoint only as a functional test.
Its throughput is not presented as a proxy for the full 2.8T model. During a
sustained repetition, the live-swapped NVIDIA driver lost the passed-through
GPU; completed samples are retained, the interrupted sample is not reported,
and the failure is documented in the artifact.

Scope

  • text generation only
  • tensor parallel size 1
  • Kimi routed experts must be MXFP4
  • hybrid radix reuse is intentionally disabled for Kimi pending KDA state
    snapshot support

This contribution was directed by the operator and implemented/tested with
Codex assistance. The implementation was manually checked against the official
configuration and checkpoint keys.

Full-checkpoint qualification evidence

The full moonshotai/Kimi-K3 checkpoint passed an end-to-end OpenAI-compatible qualification on 4× NVIDIA H200. The versioned record includes the exact server arguments, 21m30s cold start to first /v1/models, warmup and streaming measurements, and a 1/2/4/8/16-request concurrency sweep with 31/31 successful requests.

The branch was rebased on upstream main (bd372b6). The Kimi, generic offload, GPT-OSS, and graph batch-size regression suite passes: 59 tests.

@earlvanze
earlvanze force-pushed the feat/kimi-k3-support branch from 0cb565f to 39320b5 Compare August 24, 2026 02:38
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.

Add native Kimi-K3 architecture support

1 participant