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 `avarok_core::kimi_k3` so Mac unit tests compile without
6//! spark-storage. CUDA launch: [`kda_cuda`].
7
8pub mod dense_cuda;
9pub mod device_cache;
10pub mod kda;
11pub mod kda_cuda;
12pub mod latent_moe;
13pub mod mla;
14pub mod mla_cuda;
15pub mod moe_cuda;
16pub mod tp;
17
18pub use avarok_core::kimi_k3::{
19    AttnResHub, HybridCache, K3Graph, KDA_L2_EPS, KdaConfig, KdaState, LatentMoeConfig, LayerCache,
20    MixerKind, MlaConfig, MlaKv, MlpKind, attnres_blend, attnres_mix, attnres_softmax_mix,
21    cuda_kda_enabled, cuda_mla_enabled, gated_mla_attend, kda_decode_token, kda_from,
22    latent_moe_forward, mix_routed_experts, mla_decode_token, mla_from, moe_from, sigmoid_topk,
23    situ_glu, situ_glu_vec, softcap,
24};
25pub use device_cache::{DeviceHybridCache, DeviceLayerCache};
26pub use kda_cuda::{
27    K3KdaDecodeKernels, KdaDeviceState, launch_k3_kda_decode_token,
28    launch_k3_kda_decode_token_on_device,
29};
30pub use mla_cuda::{
31    K3MlaDecodeKernels, MlaDeviceKv, launch_k3_mla_decode_token,
32    launch_k3_mla_decode_token_on_device,
33};
34pub use moe_cuda::{K3MoeGemmKernels, launch_k3_latent_moe_experts};
35pub use tp::{supports_tp, tensor_plan};
36
37pub mod bound;
38mod host_decode;
39pub mod state;
40pub use avarok_core::kimi_k3::{Ablation, K3CpuModel, K3LayerSpec, greedy_decode};
41pub use state::K3CpuFallbackState;
42#[cfg(test)]
43mod host_decode_kda;
44#[cfg(test)]
45mod host_decode_mla;
spark_model/kimi_k3/mod.rs

spark_model/kimi_k3/
mod.rs
