Module kimi_k3 

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.

The CPU reference math is re-exported from metrale_model_weights::kimi_k3_host.

Owner: model-arch, Kimi K3. Invariants: none beyond the types.

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

2026-09-25: K3BoundLayer, the Kimi K3 decoder layer the engine binds, and K3HostShared, the state its layers share.

dense_cuda

2026-09-25: CUDA SiTU-GLU MLP for Kimi K3 (a dense layer’s MLP or a LatentMoE layer’s shared expert) on the resident BF16/FP32 weights, with FP32 activations.

device_cache

2026-09-25: DeviceHybridCache: device copies of every layer’s Kimi K3 mixer state, built from a host HybridCache.

kda

2026-09-25: Re-exports the Kimi K3 KDA CPU reference, metrale_model_weights::kimi_k3_host::kda.

kda_cuda

2026-09-25: Host launch of the Kimi K3 KDA decode kernels (PTX module MODULE) for one token.

latent_moe

2026-09-25: Re-exports the Kimi K3 LatentMoE CPU reference (latent down projection, sigmoid top-k router, SiTU-GLU experts, up projection), metrale_model_weights::kimi_k3_host::latent_moe.

mla

2026-09-25: Re-exports the Kimi K3 gated NoPE MLA CPU reference, metrale_model_weights::kimi_k3_host::mla.

mla_cuda

2026-09-25: Host launch of the Kimi K3 MLA decode kernels (PTX module MODULE) for one token.

moe_cuda

2026-09-25: Host launch of the Kimi K3 packed LatentMoE experts on the MXFP4 E8M0 pointer-table GEMM, PTRTABLE_E8M0 in PTX module MODULE.

state

2026-09-25: K3CpuFallbackState, one sequence’s Kimi K3 mixer state: the host LayerCache and the device copies the CUDA mixers use.

tp

2026-09-25: The Kimi K3 tensor-parallel plan (metrale_model_weights::kimi_k3_host::tp::tensor_plan) in TpShardKind terms.

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.