Experimental Native Vulkan Backend for Kimi K3 — Looking for Moonshot Engineering Feedback

#187
by netcat420 - opened

Hi Moonshot team,

I've been working on an experimental native Rust + Vulkan training/inference backend for Kimi K3, and I'd really appreciate it if any of the K3 engineers could take a look at the implementation and see whether this direction could be useful to you.

Repo: https://github.com/necat101/Hierarchos-Native

I'm reaching out now in particular because the K3 implementation has reached a point where outside technical validation would be extremely valuable. I may also have to slow or temporarily pause development in the near future due to significant financial, housing, and medical circumstances that could require me to relocate, so I wanted to get the work in front of the people most familiar with K3 while I still have the ability to actively iterate on it.

The goal of the project is to execute modern Transformer architectures directly through Vulkan compute rather than depending on CUDA, PyTorch, Triton, or another vendor-specific training runtime.

Over the last few days I've specifically implemented a native Kimi K3 / KimiLinear path instead of mapping K3 onto K2.5, DeepSeek, Qwen, or another related architecture.

The current K3 text implementation includes:

  • hybrid KDA + MLA layer scheduling
  • independent Q/K/V short convolutions
  • recurrent KDA state
  • KDA output gating
  • MLA output gating
  • SiTU
  • Stable LatentMoE
  • Attention Residuals / AttnRes
  • native forward + backward propagation
  • AdamW training
  • cached token-by-token generation
  • SafeTensors loading/export
  • Rust-native Hugging Face model handling

I've also been implementing this against the released Kimi K3 modeling semantics rather than merely trying to make the graph execute.

Current numerical validation results:

  • isolated KDA: 5.59e-8 max absolute logit error
  • isolated MLA: 6.15e-8
  • mixed production-style KDA/MLA graph: 5.22e-8
  • two full AdamW steps: 5.96e-8 max absolute parameter divergence across 18,450 compared values

My acceptance target has been a maximum divergence of 2e-7 against the reference implementation.

I'm particularly interested in whether this Vulkan direction could eventually be useful to Moonshot as an additional hardware path.

Since Vulkan isn't inherently tied to CUDA, the longer-term goal is to make architectures like K3 usable across compute-capable hardware from NVIDIA, AMD, Intel, and potentially other vendors with suitable Vulkan support.

Given the hardware-access constraints faced by AI developers in China, I'm especially curious whether a vendor-neutral backend like this could have practical value for Moonshot — whether for experimentation, fine-tuning, research hardware, development machines, inference, or eventually larger-scale workloads on non-NVIDIA systems.

I want to be careful not to overstate the current state of the project: this is not yet a claim that the full 2.8T K3 checkpoint runs efficiently in production.

At present, my validated contract is specifically the Kimi K3 language/text backbone. MoonViT / the vision projector are not implemented yet, and the released compressed-tensors MXFP4-packed checkpoint representation is currently rejected rather than silently approximated. The native path presently expects an unquantized FP32/BF16-compatible KimiLinear package.

The numerical results above are also correctness/parity tests using smaller deterministic configurations, not throughput benchmarks for the complete 2.8T model.

That's one of the reasons I'm posting here.

If anyone from the Kimi K3 engineering team is willing to inspect the architecture implementation, parity harnesses, or Vulkan execution path, I'd be extremely interested in feedback on:

  • whether I've reproduced the K3/KimiLinear semantics correctly
  • edge cases I may have missed in KDA, AttnRes, or Stable LatentMoE
  • what would be required to support the official MXFP4 representation correctly
  • whether testing the backend on hardware available to Moonshot would be worthwhile
  • whether a Vulkan-native execution path could have value in reducing dependence on a single GPU software ecosystem

I'm very willing to implement corrections or compatibility changes if your engineers identify discrepancies.

I also want to mention the personal situation above only because it may affect the project's development pace. I've been funding and developing this independently, and upcoming housing/relocation pressures combined with medical and financial issues may make it difficult for me to sustain the same amount of development time indefinitely. The code and validation work are already public, though, and I would much rather have knowledgeable engineers evaluate what exists now than allow the project to quietly stall without anyone from the K3 team ever seeing it.

If this direction looks useful to Moonshot, even just as an experimental hardware backend, hearing that would also help me understand which parts are most important to preserve and prioritize if my available development time becomes limited.

Thanks for open-sourcing K3 and its custom modeling code. It made this kind of independent implementation and numerical validation possible.

GitHub:
https://github.com/necat101/Hierarchos-Native

Sign up or log in to comment