XingfenD
c57a9bf9e6
- MultiHeadAttention类添加KV cache机制 - TransformerBlock支持use_cache参数 - SimpleTransformer新增prefill和decode方法 - 添加InferenceEngine推理引擎类 - 更新example.py演示推理过程 - 更新README.md文档
55 lines
1.4 KiB
Markdown
55 lines
1.4 KiB
Markdown
# Simple Transformer
|
|
|
|
A minimal Transformer implementation using NumPy with **KV Cache** support for efficient inference.
|
|
|
|
## Files
|
|
|
|
- `transformer.py` - Core Transformer model implementation with KV Cache
|
|
- `example.py` - Usage examples (training & inference)
|
|
|
|
## Usage
|
|
|
|
```bash
|
|
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 methods
|
|
- `InferenceEngine` - Manages inference process with KV cache |