bitmamba 0.1.0

BitMamba: 1.58-bit Mamba language model with infinite context window - includes OpenAI-compatible API server
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

# BitMamba

A 1.58-bit Mamba language model with **infinite context window**, implemented in Rust.

[![Crates.io](https://img.shields.io/crates/v/bitmamba.svg)](https://crates.io/crates/bitmamba)
[![Documentation](https://docs.rs/bitmamba/badge.svg)](https://docs.rs/bitmamba)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

## Features

- **Infinite Context Window** - Mamba's SSM maintains fixed-size state regardless of sequence length
- **1.58-bit Weights** - BitNet-style quantization for efficient inference  
- **CPU Inference** - No GPU required
- **OpenAI-Compatible API** - Drop-in replacement for OpenAI API, works with Cline, Continue, etc.
- **Streaming Support** - Server-Sent Events for real-time token generation

## Installation

```bash
cargo install bitmamba
```

Or build from source:

```bash
git clone https://github.com/rileyseaburg/bitmamba
cd bitmamba
cargo build --release
```

## Usage

### CLI

```bash
# Run inference directly
bitmamba
```

### OpenAI-Compatible Server

```bash
# Start the server
bitmamba-server
```

The server runs at `http://localhost:8000` with these endpoints:

| Endpoint | Method | Description |
|----------|--------|-------------|
| `/v1/models` | GET | List available models |
| `/v1/chat/completions` | POST | Chat completions (streaming supported) |
| `/v1/completions` | POST | Text completions |
| `/health` | GET | Health check |

### Configure with Cline/Continue

```json
{
  "apiProvider": "openai-compatible",
  "baseUrl": "http://localhost:8000/v1",
  "model": "bitmamba-student"
}
```

### As a Library

```rust
use bitmamba::{BitMambaStudent, load_model, load_tokenizer};

fn main() -> anyhow::Result<()> {
    let (model, tokenizer) = bitmamba::load()?;
    
    let prompt = "def fibonacci(n):";
    let tokens = tokenizer.encode(prompt, true)?;
    let output = model.generate(tokens.get_ids(), 50, 0.7)?;
    
    println!("{}", tokenizer.decode(&output, true)?);
    Ok(())
}
```

## Model

The default model is [`rileyseaburg/bitmamba-student`](https://huggingface.co/rileyseaburg/bitmamba-student) on Hugging Face, a 278M parameter BitMamba model distilled from Qwen2.5-Coder-1.5B.

### Architecture

- **Hidden Size**: 768
- **Layers**: 12 BitMamba blocks
- **State Size**: 16 (SSM state dimension)
- **Expand Factor**: 2
- **Vocab Size**: 151,665 (Qwen tokenizer)

### BitMamba Block

```
Input -> RMSNorm -> BitLinear (in_proj) -> Conv1d -> SiLU -> SSM Scan -> Gate -> BitLinear (out_proj) -> Residual
```

The **SSM Scan** is the key component that enables infinite context:

```rust
// Fixed-size state, O(1) memory per token
h = dA * h + dB * x  // State update
y = h @ C + D * x     // Output
```

## Performance

| Metric | Value |
|--------|-------|
| Parameters | 278M |
| Memory (inference) | ~1.1 GB |
| Context Window | Unlimited |
| Quantization | 1.58-bit weights |

## Citation

If you use BitMamba in your research, please cite:

```bibtex
@software{bitmamba2024,
  author = {Seaburg, Riley},
  title = {BitMamba: 1.58-bit Mamba with Infinite Context},
  year = {2024},
  url = {https://github.com/rileyseaburg/bitmamba}
}
```

## Related Work

- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://arxiv.org/abs/2312.00752)
- [BitNet: Scaling 1-bit Transformers for Large Language Models](https://arxiv.org/abs/2310.11453)
- [The Era of 1-bit LLMs](https://arxiv.org/abs/2402.17764)

## License

MIT License - see [LICENSE](LICENSE) for details.