metrale_model_arch/kimi_k3/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Kimi K3 in the engine: the bound decoder layer, its per-sequence state, and the CUDA launches for the KDA and MLA mixers, the packed LatentMoE experts and the dense MLP.
4//!
5//! The CPU reference math is re-exported from `metrale_model_weights::kimi_k3_host`.
6//!
7//! Owner: model-arch, Kimi K3.
8//! Invariants: none beyond the types.
9
10pub mod dense_cuda;
11pub mod device_cache;
12pub mod kda;
13pub mod kda_cuda;
14pub mod latent_moe;
15pub mod mla;
16pub mod mla_cuda;
17pub mod moe_cuda;
18pub mod tp;
19
20pub use device_cache::{DeviceHybridCache, DeviceLayerCache};
21pub use kda_cuda::{
22    K3KdaDecodeKernels, KdaDeviceState, launch_k3_kda_decode_token,
23    launch_k3_kda_decode_token_on_device,
24};
25pub use metrale_model_weights::kimi_k3_host::{
26    AttnResHub, HybridCache, K3Graph, KDA_L2_EPS, KdaConfig, KdaState, LatentMoeConfig, LayerCache,
27    MixerKind, MlaConfig, MlaKv, MlpKind, attnres_blend, attnres_mix, attnres_softmax_mix,
28    cuda_kda_enabled, cuda_mla_enabled, gated_mla_attend, kda_decode_token, kda_from,
29    latent_moe_forward, mix_routed_experts, mla_decode_token, mla_from, moe_from, sigmoid_topk,
30    situ_glu, situ_glu_vec, softcap,
31};
32pub use mla_cuda::{
33    K3MlaDecodeKernels, MlaDeviceKv, launch_k3_mla_decode_token,
34    launch_k3_mla_decode_token_on_device,
35};
36pub use moe_cuda::{K3MoeGemmKernels, launch_k3_latent_moe_experts};
37pub use tp::{supports_tp, tensor_plan};
38
39pub mod bound;
40mod host_decode;
41pub mod state;
42pub use metrale_model_weights::kimi_k3_host::{Ablation, K3CpuModel, K3LayerSpec, greedy_decode};
43pub use state::K3CpuFallbackState;
44#[cfg(test)]
45mod host_decode_kda;
46#[cfg(test)]
47mod host_decode_mla;
metrale_model_arch/kimi_k3/mod.rs

metrale_model_arch/kimi_k3/
mod.rs
