Module kimi_k3 

Kimi K3 host graph. This slice is KDA only.

Math lives in avarok_core::kimi_k3 so Mac unit tests compile without spark-storage. CUDA launch: kda_cuda.

Re-exports

pub use device_cache::DeviceHybridCache;

pub use device_cache::DeviceLayerCache;

pub use kda_cuda::K3KdaDecodeKernels;

pub use kda_cuda::KdaDeviceState;

pub use kda_cuda::launch_k3_kda_decode_token;

pub use kda_cuda::launch_k3_kda_decode_token_on_device;

pub use mla_cuda::K3MlaDecodeKernels;

pub use mla_cuda::MlaDeviceKv;

pub use mla_cuda::launch_k3_mla_decode_token;

pub use mla_cuda::launch_k3_mla_decode_token_on_device;

pub use moe_cuda::K3MoeGemmKernels;

pub use moe_cuda::launch_k3_latent_moe_experts;

pub use tp::supports_tp;

pub use tp::tensor_plan;

pub use state::K3CpuFallbackState;

Modules

bound

BF16/FP32 twin layer bind. Decode copies hidden D2H, runs mixer+MLP+AttnRes, copies H2D. LinearAttention / KDA conv+recurrent uses CUDA kda_decode with per-sequence resident recurrence unless K3_CUDA_KDA=0. FullAttention / MLA uses CUDA mla_decode unless K3_CUDA_MLA=0. Packed LatentMoE experts launch DSV4 moe_w4a16_grouped_gemm_ptrtable_e8m0. Router / down / up / shared / SiTU stay on the host.

dense_cuda

Opt-in K3 dense/shared MLP using original resident weights and FP32 I/O. Only input and final output cross the host boundary; no weight duplication.

device_cache

GPU-resident hybrid cache. Host HybridCache bytes are prefix snapshot only. Decode must not D2H/H2D conv, recurrent, or the growing MLA KV.

kda

K3 KDA CPU ref (twin 8×32 conv-4; prod 96×128 bound −5).

kda_cuda

Host launch for K3 CUDA KDA decode (kda_decode PTX module).

latent_moe

Latent down / sigmoid-noaux_tc router / SiTU experts / up CPU ref.

mla

Gated NoPE MLA CPU ref. Do not reuse qwen3_attention blindly. CUDA: super::mla_cuda (default; K3_CUDA_MLA=0 CPU).

mla_cuda

Host launch for K3 CUDA gated-NoPE MLA decode (mla_decode PTX module).

moe_cuda

Host launch for K3 packed LatentMoE experts (moe_w4a16 E8M0 ptrtable).

state

Per-sequence K3 state: resident CUDA KDA recurrence and MLA KV.

tp

K3 device-binding adapter for the shared tensor-parallel storage plan. Runtime loading partitions packed and dense weights before GPU allocation.

Structs

Ablation

Graph mutants. Mix=0 / force-expert-0 are C1; skip-layer is C2; o_proj TP is C7.

AttnResHub

Stream map keyed by hidden/residual pointer bits (DevicePtr.0). Layer 0 inserts; last layer removes on success. Any Err drops the entry (umbrella 940bd4eeb).

HybridCache

Per-sequence hybrid cache. Slot identity is the layer index; a prefix hit that writes KDA state into the wrong slot is the C4 mutant.

K3CpuModel

CPU twin (or tiny stand-in) ready for greedy decode.

K3Graph

K3LayerSpec

KdaConfig

KDA geometry. Tiny dims are legal for CPU tests; production is 128/4.

KdaState

Per-sequence KDA state. Both buffers are FP32, read-modify-write.

LatentMoeConfig

MlaConfig

MlaKv

One MLA layer’s host KV (unpaged CPU stand-in).

Enums

LayerCache

MixerKind

MlpKind

Constants

KDA_L2_EPS

FLA use_qk_l2norm_in_kernel eps. CUDA k3_kda_recurrent_step_f32 uses the same.

Functions

attnres_blend

Test / ablation lever: mix=0 returns skip; mix=1 returns mixed.

attnres_mix

Apply AttnRes with an explicit mix lever. mix=0 is identity skip.

attnres_softmax_mix

Softmax mixture over residual sources. sources[0] is conventionally the skip (current partial block). query is [hidden] (*_res_proj).

cuda_kda_enabled

LinearAttention BoundLayer uses CUDA kda_decode unless K3_CUDA_KDA=0. Projections, AttnRes, and MLP stay on the host either way.

cuda_mla_enabled

FullAttention BoundLayer uses CUDA mla_decode unless K3_CUDA_MLA=0. Same polarity as K3_CUDA_KDA. Projections stay on the host either way.

gated_mla_attend

Gated NoPE attend: optionally skip RoPE, SDPA, optional output gate.

greedy_decode

Greedy decode. Returns prompt || new_tokens (length prompt.len() + max_new).

kda_decode_token

One decode token: conv update then recurrent step.

kda_from

Map parsed ModelConfig onto KDA geometry.

latent_moe_forward

Routed latent path + optional shared expert (identity-scale for tests).

mix_routed_experts

Mix selected SiTU-GLU experts into latent. Empty w1 means packed-only (no host dequant) — CUDA grouped GEMM must have run instead.

mla_decode_token

One decode token: optional RoPE, append K/V, SDPA, optional output gate. Projections stay in mla_mixer. CUDA k3_mla_* matches this order.

mla_from

moe_from

sigmoid_topk

noaux_tc: sigmoid(logits) for mix weights; scores + bias only ranks.

situ_glu

One SiTU-GLU coordinate. Product is bounded by |β1 * β2|.

situ_glu_vec

Elementwise SiTU-GLU over paired gate/up vectors.

softcap

Smooth cap. beta must be finite and non-zero.