spark_model/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//! Math lives in `atlas_core::kimi_k3` so Mac unit tests compile without
6//! spark-storage. CUDA launch: [`kda_cuda`].
7
8pub mod device_cache;
9pub mod kda;
10pub mod kda_cuda;
11pub mod latent_moe;
12pub mod mla;
13pub mod mla_cuda;
14pub mod moe_cuda;
15pub mod tp;
16
17pub use atlas_core::kimi_k3::{
18    AttnResHub, HybridCache, K3Graph, KDA_L2_EPS, KdaConfig, KdaState, LatentMoeConfig, LayerCache,
19    MixerKind, MlaConfig, MlaKv, MlpKind, attnres_blend, attnres_mix, attnres_softmax_mix,
20    cuda_kda_enabled, cuda_mla_enabled, gated_mla_attend, kda_decode_token, kda_from,
21    latent_moe_forward, mix_routed_experts, mla_decode_token, mla_from, moe_from, sigmoid_topk,
22    situ_glu, situ_glu_vec, softcap,
23};
24pub use device_cache::{DeviceHybridCache, DeviceLayerCache};
25pub use kda_cuda::{
26    K3KdaDecodeKernels, KdaDeviceState, launch_k3_kda_decode_token,
27    launch_k3_kda_decode_token_on_device,
28};
29pub use mla_cuda::{
30    K3MlaDecodeKernels, MlaDeviceKv, launch_k3_mla_decode_token,
31    launch_k3_mla_decode_token_on_device,
32};
33pub use moe_cuda::{K3MoeGemmKernels, launch_k3_latent_moe_experts};
34pub use tp::{supports_tp, tensor_plan};