Skip to content

feat(kimi): route a Kimi K3.1 checkpoint through the Text path - #797

Draft
kalebbroo wants to merge 1 commit into
mainfrom
feat/kimi-k31-text-route
Draft

kalebbroo wants to merge 1 commit into
mainfrom
feat/kimi-k31-text-route

Conversation

@kalebbroo

Copy link
Copy Markdown
Contributor

Summary

Wires Kimi K3.1 (model_type kimi_k31) into the engine's Text path as a host reference family, so a Kimi checkpoint directory loads through HfTextDirectoryLoader.LoadHostReference the way ERNIE-5.5 and DeepSeek-V4.1 do.

  • HfTextDirectoryLoader: kimi_k31 is a supported host reference family, loaded as KimiK31TextModel on the backend the caller passes. RequireSupported refuses, by name, a Kimi directory with no shard index, no chat_template.jinja, or no tokenizer.json (the Moonshot repo ships tiktoken.model and tokenization_kimi.py instead).
  • TextService.LoadHfDirectory: the Kimi branch picks its backend from the requested device (CUDA or Vulkan when asked, CPU otherwise). A device that cannot be created throws; the load never falls back to the CPU. The DeepSeek and ERNIE paths are unchanged, and they still force the CPU.
  • KimiK31TextModel and KimiK31GenerationModel (new, src/HartsyInference.LLM/Kimi/) adapt KimiK31Checkpoint to the shared pipeline. They reuse HostReplaySequenceState and do not batch decode, so Kimi requests take the pipeline route, which gates the device they name.
  • KimiK31Model: ForwardToken is split into ForwardTokenHidden and ProjectHidden. The decoder's math is unchanged, and the existing Kimi tests pass.
  • KimiK31Refusals (new): image input on a Kimi slot is refused by name (NotSupportedException). Otherwise the image would reach RunText and be dropped silently. Multi-GPU placement (+ composite keys, two or more shard devices, tensor-parallel degree above 1) is refused by name. This is the placement hook for the sharding work still in progress.
  • Vision: the checkpoint loads text-only. The vision tower and mm_projector.* are never read. Image requests are refused.
  • Catalog row kimi-k3.1-pro (moonshotai/Kimi-K3.1-Pro): Structural, CliDrivable = false, no Assets. Revision pinned to d9e99b0792a5e93b78ea9bea62a68167e36ddd9e (read from the HF API sha). License text summarized in the row's notes from the repo's LICENSE file ("Kimi K3.1-Pro License").

Tests (CPU lane)

  • dotnet test on HartsyInference.Engine.Tests, CPU filter: 314 passed, 0 failed, 1 skipped (FfmpegFrameSourceTests, resource-guarded).
  • dotnet test on HartsyInference.LLM.Tests, CPU filter: 178 passed, 0 failed.
  • New Engine tests: KimiK31TextPathTests (route load on CPU, finite logits from prefill and ProjectLogits, deterministic pipeline generation, rollback and replay, vision checkpoint loads text-only and refuses image input by name, multi-GPU refusals by name, single-device requests pass, missing template and missing tokenizer.json refused by name) and KimiK31CatalogTests.
  • Synthetic checkpoint: KimiSyntheticCheckpoint (Engine test helper) reuses the ERNIE tokenizer and template files; vocab is 259 to cover the merged id 258.

Not run

  • No GPU tests. KimiK31ModelGpuTests (GpuIntegration) was not run. The CUDA and Vulkan paths through KimiK31TextModel are wired but unverified: host-resident weights on a device backend, the device gate, and the output tensors are untested.
  • No real-weight run. The Moonshot repo has no tokenizer.json, so the real checkpoint cannot load through the route until one is supplied. The loader widens dense weights to F32 at load (about 114 GB of bf16 in the research doc), and there is no host memory guard yet. TextService logs a warning about this.
  • No TextService-level test. The route's device choice and refusals sit in TextService, which needs the whole engine. The refusals are unit tested through KimiK31Refusals.
  • The HF revision is pinned from the API sha. The HF tree was not checked for file-level hashes.

Notes for reviewers

  • KimiK31Model.cs and TextService.cs are shared with other work. The KimiK31Model change is additive, in the ForwardToken region. TextService also gets a small refactor: the slot-disposal code moves from ReplaceWithHostBackend into DisposeSlotBackends, with no behavior change for the existing paths.
  • Changelog entry under the existing ## Unreleased heading, per the author's instruction. VersionSuffix is untouched.

- HfTextDirectoryLoader supports kimi_k31 as a host reference family, loading KimiK31TextModel
  on the backend the request names (CUDA or Vulkan when asked, CPU otherwise). No CPU fallback.
- KimiK31TextModel and KimiK31GenerationModel adapt KimiK31Checkpoint to the shared pipeline.
  KimiK31Model splits ForwardToken into ForwardTokenHidden and ProjectHidden.
- Image input and multi-GPU placement are refused by name (KimiK31Refusals).
- The vision tower is not loaded; the checkpoint loads text-only.
- kimi-k3.1-pro catalog row: Structural, not CLI-drivable, pinned to d9e99b07.
- CPU-lane tests cover the route, logits, replay, refusals and the catalog row.

This branch has not been deployed

No deployments
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