Module cpu_forward 

One-token K3 CPU forward: AttnRes + mixer + MLP, then logits.

Structs

AttnResStream

Matches HF KimiDecoderLayer._forward_attn_residual: mix the incoming prefix with already-archived blocks, then at layer_idx % block_size == 0 archive that incoming prefix and reset the intra-block sum. Layer 0 therefore archives the embedding as its own source, not embed + mixer.

K3LayerCtx

Per-layer geometry the GPU wrapper and forward_token share.

Functions

forward_one_layer

One decoder layer: AttnRes + (KDA|MLA) + MLP. Mutates this layer’s LayerCache and the token’s AttnResStream.

forward_one_layer_with_cores

Same as forward_one_layer, with replaceable KDA / MLA / routed-expert cores.

forward_one_layer_with_kda_decode

Same as forward_one_layer, with a replaceable KDA conv+recurrent core.

forward_one_layer_with_mla_decode

Same as forward_one_layer, with a replaceable MLA rope+SDPA+gate core.

forward_token

Embed one token and run every decoder layer at pos. Updates the hybrid cache. AttnRes is per-token (layer depth), not carried across the sequence.

logits

LM head over a hidden vector. W is [vocab, hidden].

Type Aliases

DenseMlpCore

Optional resident dense/shared MLP, with local TP intermediate width.

HiddenReduce

After row-parallel o_proj / dense MLP down. None at TP=1.