avarok_core/kimi_k3/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Kimi K3 host graph. This slice is **KDA only**.
4//!
5//! K3-DECISION: KDA is a new backend. Do not copy GDN / Mamba-2 /
6//! Qwen3-Next / glm5next_kda kernels into `kernels/gb10/kimi-k3/`.
7//! CUDA: `kda_decode.cu` (default; `K3_CUDA_KDA=0` CPU escape).
8
9pub mod attnres;
10pub mod binding_memory;
11pub mod cache;
12pub mod expert_backend;
13pub mod kda;
14pub mod latent_moe;
15pub mod layer;
16pub mod mla;
17pub mod situ;
18pub mod tp;
19pub mod weights;
20pub mod weights_geometry;
21
22pub use attnres::{AttnResHub, attnres_blend, attnres_mix, attnres_softmax_mix};
23pub use cache::{HybridCache, LayerCache, MlaKv};
24pub use expert_backend::{ExpertBackendKind, MmapExpertStore, PrefetchPlanner, kind_from_env};
25pub use kda::{KDA_L2_EPS, KdaConfig, KdaState, cuda_kda_enabled, kda_decode_token, kda_from};
26pub use latent_moe::{
27    LatentMoeConfig, latent_moe_forward, mix_routed_experts, moe_from, sigmoid_topk,
28};
29pub use layer::{K3Graph, K3LayerSpec, MixerKind, MlpKind};
30pub use mla::{MlaConfig, cuda_mla_enabled, gated_mla_attend, mla_decode_token, mla_from};
31pub use situ::{situ_glu, situ_glu_vec, softcap};
32
33mod cpu_bind;
34pub mod cpu_forward;
35mod cpu_load;
36pub mod cpu_weights;
37pub mod greedy;
38pub mod mxfp4;
39pub mod ops;
40pub use cpu_bind::{assemble_layer, text_key};
41pub use cpu_forward::{
42    AttnResStream, HiddenReduce, K3LayerCtx, forward_one_layer, forward_one_layer_with_cores,
43    forward_one_layer_with_kda_decode, forward_one_layer_with_mla_decode, forward_token,
44};
45pub use cpu_weights::{Ablation, K3CpuLayer, K3CpuModel};
46pub use greedy::greedy_decode;
47
48#[cfg(test)]
49mod c1;
50
51#[cfg(test)]
52mod c2;
53
54#[cfg(test)]
55mod c3;
56
57#[cfg(test)]
58mod c4;
59
60#[cfg(test)]
61mod c5;
62
63#[cfg(test)]
64mod c6;
65
66#[cfg(test)]
67mod c7;
68
69#[cfg(test)]
70mod host_fallback;
71
72#[cfg(test)]
73mod kda_cuda_parity;
74
75#[cfg(test)]
76mod mla_cuda_parity;