XingfenD
c57a9bf9e6
- MultiHeadAttention类添加KV cache机制 - TransformerBlock支持use_cache参数 - SimpleTransformer新增prefill和decode方法 - 添加InferenceEngine推理引擎类 - 更新example.py演示推理过程 - 更新README.md文档
1.4 KiB
1.4 KiB
Simple Transformer
A minimal Transformer implementation using NumPy with KV Cache support for efficient inference.
Files
transformer.py- Core Transformer model implementation with KV Cacheexample.py- Usage examples (training & inference)
Usage
python3 example.py
Model Architecture
The implementation includes:
- Multi-Head Attention: Scaled dot-product attention with multiple heads
- Feed-Forward Network: Two-layer fully connected network with ReLU activation
- Layer Normalization: Applied after each sub-layer
- Positional Encoding: Sinusoidal position embeddings
- KV Cache: Efficient inference with prefill/decode separation
Inference: Prefill vs Decode
Prefill Stage
- Process the complete input prompt in parallel
- Initialize KV cache for all layers
- Generate the first token
- Computation: O(n²)
Decode Stage
- Process one token at a time
- Reuse cached Key and Value tensors
- Only compute new Query
- Computation: O(n) per step
Model Parameters
Default configuration:
- Vocabulary size: 1000
- Model dimension: 512
- Number of heads: 8
- Number of layers: 6
- Feed-forward dimension: 2048
- Max sequence length: 100
Total parameters: ~19.4M
Classes
SimpleTransformer- Transformer model with prefill/decode methodsInferenceEngine- Manages inference process with KV cache