# BitLinear Project Structure

Complete directory tree and file descriptions.

```
BitLinear/
│
├── README.md                      # Project overview and quick start
├── LICENSE                        # MIT License
├── setup.py                       # Build system with torch.utils.cpp_extension
├── pyproject.toml                 # Tool configurations (pytest, black, mypy)
├── requirements.txt               # Core dependencies
├── requirements-dev.txt           # Development dependencies
├── .gitignore                     # Git ignore rules
├── IMPLEMENTATION_GUIDE.md        # Step-by-step implementation roadmap
│
├── bitlinear/                     # Main package
│   ├── __init__.py               # Package exports
│   ├── layers.py                 # BitLinear and MultiTernaryLinear modules
│   ├── functional.py             # Core functional implementations
│   ├── quantization.py           # Ternary quantization utilities
│   ├── packing.py                # Base-3 packing for memory efficiency
│   │
│   └── cpp/                      # C++/CUDA extensions
│       ├── bitlinear.cpp         # PyBind11 bindings and CPU implementation
│       └── bitlinear_kernel.cu   # CUDA kernel implementations
│
├── tests/                         # Test suite
│   ├── __init__.py
│   ├── test_functional.py        # Tests for functional API
│   ├── test_layers.py            # Tests for layer modules
│   └── test_quantization.py     # Tests for quantization and packing
│
└── examples/                      # Usage examples
    ├── basic_usage.py            # Simple usage demonstration
    └── transformer_example.py    # Transformer integration example
```

## File Descriptions

### Root Level

- **README.md**: Project overview, installation instructions, quick start guide, and citations
- **LICENSE**: MIT License for open-source distribution
- **setup.py**: Build configuration using PyTorch's cpp_extension, handles CPU/CUDA builds
- **pyproject.toml**: Configuration for pytest, black, mypy, and coverage
- **requirements.txt**: Core runtime dependencies (torch, numpy)
- **requirements-dev.txt**: Development tools (pytest, black, flake8, mypy)
- **.gitignore**: Ignores Python cache, build artifacts, CUDA objects
- **IMPLEMENTATION_GUIDE.md**: Detailed implementation roadmap with phases and best practices

### bitlinear/ (Main Package)

#### Python Modules

- **`__init__.py`**: Package initialization, exports main classes and functions
- **`layers.py`**: nn.Module implementations
  - `BitLinear`: Drop-in replacement for nn.Linear with ternary weights
  - `MultiTernaryLinear`: Sum of k ternary components
  - `convert_linear_to_bitlinear()`: Recursive model conversion utility

- **`functional.py`**: Core functional implementations
  - `bitlinear_python()`: Pure PyTorch ternary matmul with scaling
  - `greedy_ternary_decomposition()`: Iterative residual quantization
  - `multi_ternary_linear_python()`: Multi-component forward pass
  - `activation_quant()`: Activation quantization for full BitNet

- **`quantization.py`**: Quantization utilities
  - `absmax_scale()`: Compute absmax scaling factors
  - `ternary_quantize()`: Quantize to {-1, 0, +1}
  - `weight_to_ternary()`: Full quantization pipeline
  - `quantize_activations_absmax()`: 8-bit activation quantization
  - `dequantize_scale()`: Reverse quantization

- **`packing.py`**: Memory optimization
  - `pack_ternary_base3()`: Pack 5 ternary values per byte
  - `unpack_ternary_base3()`: Unpack base-3 encoded weights
  - `compute_compression_ratio()`: Calculate compression statistics
  - `estimate_memory_savings()`: Memory estimation utilities

#### C++/CUDA Extensions

- **`cpp/bitlinear.cpp`**: C++ interface
  - PyBind11 module definition
  - CPU implementations: `bitlinear_cpu_forward()`, `multi_ternary_cpu_forward()`
  - Device dispatcher (routes to CPU or CUDA)
  - Packing utilities in C++

- **`cpp/bitlinear_kernel.cu`**: CUDA kernels
  - `bitlinear_forward_kernel()`: Optimized ternary matmul kernel
  - `multi_ternary_forward_kernel()`: Fused multi-component kernel
  - Kernel launchers with error handling
  - TODO: Tensor Core optimization

### tests/

Comprehensive test suite using pytest:

- **`test_functional.py`**: Tests for functional API
  - Shape correctness
  - Numerical correctness vs. nn.Linear
  - Greedy decomposition quality
  - Multi-ternary equivalence

- **`test_layers.py`**: Tests for layer modules
  - Initialization and parameter counts
  - Forward pass shapes
  - Compatibility with nn.Linear
  - Conversion utilities
  - Gradient flow (QAT)
  - Integration with Transformer blocks

- **`test_quantization.py`**: Tests for quantization
  - Absmax scaling (global and per-channel)
  - Ternary quantization values and thresholds
  - Reconstruction quality
  - Base-3 packing roundtrip
  - Compression ratios
  - Memory estimation

### examples/

Demonstration scripts:

- **`basic_usage.py`**: Minimal example showing basic API
  - Creating BitLinear layers
  - Forward pass
  - Conversion from nn.Linear

- **`transformer_example.py`**: Realistic Transformer example
  - Complete Transformer block implementation
  - Conversion to BitLinear
  - Output comparison
  - Memory savings calculation

## Key Design Patterns

### 1. Progressive Enhancement
- Python baseline → C++ CPU → CUDA GPU
- Each layer fully functional before adding next

### 2. Drop-in Compatibility
- Same interface as nn.Linear
- Same initialization arguments
- Same forward signature
- Works with existing PyTorch features

### 3. Modular Testing
- Unit tests for each component
- Integration tests for full pipelines
- Performance benchmarks separate

### 4. Extensive Documentation
- Docstrings explain mathematical operations
- TODO comments mark implementation points
- References to papers for algorithms
- Type hints for clarity

## Build Targets

### CPU-only (Development)
```bash
pip install -e .
```

### With CUDA (Production)
```bash
CUDA_HOME=/usr/local/cuda pip install -e .
```

### Testing
```bash
pip install -e ".[dev]"
pytest tests/ -v
```

## What's NOT Implemented Yet

All files are **stubs with TODOs**:
- ✅ Structure is complete
- ✅ Interfaces are defined
- ✅ Documentation is written
- ❌ Logic is NOT implemented (by design)
- ❌ Tests will skip/fail until implementation

## Next Steps

Follow IMPLEMENTATION_GUIDE.md:
1. Start with `quantization.py` (absmax_scale, ternary_quantize)
2. Move to `functional.py` (bitlinear_python)
3. Implement `layers.py` (BitLinear module)
4. Test with examples
5. Add C++/CUDA if needed

## Design Philosophy

**Correctness > Speed > Memory**
1. First make it work (Python)
2. Then make it fast (C++/CUDA)
3. Then make it efficient (packing)

Every component is:
- Well-documented
- Testable
- Modular
- Extensible