How moonshotai/nano-kpu works
nano-kpu is a hybrid-attention MoE inference chip, with a functional-simulation and synthesis-measurement flow. Its principal workflow is a host RUN command plus a memory image: the chip fetches token and weight data, runs token-by-token transformer inference, and writes logits, argmax results, and status to external memory before returning DONE. The model stages include KDA, MLA, attention-residual mixing, and sigmoid-routed top-2 MoE with a shared expert. The evaluation flow compares chip output with the float32 reference and reports correctness, throughput, area, and timing. Some RTL wiring is inferred from documented capabilities and module names; unsampled source limits finer-grained call relationships.
Main entry points: msh_chip_top.v Chip controller · evaluate.py Evaluation runner · msh_mem.v Memory interface · msh_vec.v Vector operations
Components 21
Interface and memory
- Chip controller
rtl/msh_chip_top.v - Fetch engine
rtl/msh_fetch.v - Memory interface
rtl/msh_mem.v - Packed image
harness/memmap.py - Memory macros
harness/macros
Inference datapath
- Token sequencer
rtl/msh_seq.v - Vector operations
rtl/msh_vec.v - Quantized GEMV
rtl/msh_gemv.v - KDA attention
rtl/msh_kda.v - MLA attention
rtl/msh_mla.v - MoE routing
rtl/msh_deq.v - Logit generation
rtl/msh_lg.v - Nonlinearity tables
rtl/roms
Model and evaluation
- Fixed-point model
rtl/selfmodel/fxmodel.py - Float reference
reference/model.py - Weight generation
reference/gen_weights.py - Evaluation runner
harness/evaluate.py - Correctness metrics
harness/scoring.py
Synthesis measurement
- Area and timing
harness
Other components
- Host command
- DRAM model
harness/tb/tb_main.cpp
Connections 30
- Host command → Chip controller: issues RUN
- Chip controller → Fetch engine: dispatches
- Fetch engine → Memory interface: requests data
- Memory interface → DRAM model: reads and writes
- DRAM model → Memory interface: returns reads
- Packed image → DRAM model: loads image
- Memory interface → Memory macros: uses macros
- Token sequencer → Fetch engine: requests tokens
- Chip controller → Token sequencer: controls decode
- Chip controller → KDA attention: dispatches layer
- Chip controller → MLA attention: dispatches layer
- Chip controller → MoE routing: dispatches experts
- KDA attention → Quantized GEMV: computes projections
- MLA attention → Quantized GEMV: computes projections
- MoE routing → Quantized GEMV: computes experts
- KDA attention → Vector operations: uses operations
- MLA attention → Vector operations: uses operations
- MoE routing → Vector operations: uses operations
- Vector operations → Nonlinearity tables: looks up functions
- Chip controller → Logit generation: produces output
- Logit generation → Memory interface: writes results
- Evaluation runner → Weight generation: generates tensors
- Evaluation runner → Packed image: packs inputs
- Evaluation runner → Float reference: computes golden output
- Evaluation runner → Chip controller: runs simulation
- Evaluation runner → Correctness metrics: compares outputs
- Evaluation runner → Area and timing: measures design
- Fixed-point model → Nonlinearity tables: loads tables
- Float reference → Weight generation: consumes tensors
- Correctness metrics → Float reference: compares against