aha/docs/concepts.md

Architecture & Design

This document provides an in-depth look at the architecture and design principles behind AHA.

Overview

AHA (High-performance AI inference engine) is a Rust-based library built on the Candle framework. It provides a unified interface for running multiple state-of-the-art AI models locally, without requiring API keys or cloud services.

Key Characteristics

• Local-First: All inference runs on your machine
• Multi-Modal: Support for text, vision, audio, OCR, and ASR models
• Cross-Platform: Linux, macOS, and Windows support
• GPU-Accelerated: Optional CUDA and Metal support
• Memory-Safe: Built with Rust for safety and performance
• OpenAI-Compatible: Easy integration with existing tools

Architecture Principles

1. Local-First Design

AHA is designed to run entirely on your local machine:

• No cloud dependencies: All models are downloaded and run locally
• Privacy-preserving: Your data never leaves your machine
• No API keys required: Once downloaded, models work indefinitely
• Offline capable: Models work without internet connection after download

2. Unified Model Interface

All models implement a common GenerateModel trait, providing:

• Consistent API across different model types
• Easy model switching without code changes
• Streaming response support for real-time outputs
• Standardized error handling

3. Cross-Platform Support

AHA abstracts platform differences:

• Device abstraction: Automatic CPU/GPU detection and selection
• Precision handling: Dynamic F32/F16/BF16 selection based on hardware
• Path management: Consistent model storage across platforms

Core Components

┌─────────────────────────────────────────────────────────────┐
│                         CLI Layer                           │
│  (main.rs - Command parsing, model download, service mgmt)  │
└─────────────────────────────────────────────────────────────┘
                              │
┌─────────────────────────────────────────────────────────────┐
│                        HTTP API Layer                        │
│  (api.rs - OpenAI-compatible endpoints, streaming, auth)     │
└─────────────────────────────────────────────────────────────┘
                              │
┌─────────────────────────────────────────────────────────────┐
│                     Model Abstraction Layer                  │
│           (GenerateModel trait - unified interface)          │
└─────────────────────────────────────────────────────────────┘
                              │
        ┌─────────────────────┼─────────────────────┐
        │                     │                     │
┌───────▼────────┐  ┌────────▼─────────┐  ┌───────▼────────┐
│  Text Models   │  │  Vision Models   │  │  Audio Models  │
│  - Qwen3       │  │  - Qwen2.5VL     │  │  - VoxCPM      │
│  - MiniCPM4    │  │  - Qwen3VL       │  │  - VoxCPM1.5   │
└────────────────┘  └──────────────────┘  └────────────────┘
        │                     │                     │
┌───────▼────────┐  ┌────────▼─────────┐  ┌───────▼────────┐
│   OCR Models   │  │   ASR Models     │  │  Image Models  │
│  - DeepSeek    │  │  - GLM-ASR       │  │  - RMBG2.0     │
│  - Hunyuan     │  │  - Fun-ASR       │  │                │
│  - PaddleOCR   │  │  - Qwen3-ASR     │  │                │
└────────────────┘  └──────────────────┘  └────────────────┘
                              │
┌─────────────────────────────────────────────────────────────┐
│                      Utility Modules                         │
│  - tokenizer: Tokenization utilities                         │
│  - chat_template: Chat format handling                       │
│  - position_embed: Positional embeddings                     │
│  - utils: Common utilities (audio, image, download)          │
└─────────────────────────────────────────────────────────────┘
                              │
┌─────────────────────────────────────────────────────────────┐
│                   Candle ML Framework                        │
│  (Tensor operations, model loading, device management)       │
└─────────────────────────────────────────────────────────────┘

CLI Layer (src/main.rs)

The CLI layer provides command-line interface functionality:

• Command parsing: Uses clap for argument parsing
• Model management: Automatic download and caching
• Service control: Start/stop HTTP server
• Direct inference: Run models without server

Available Commands:

• cli - Download model and start service (default)
• serv - Start service with existing model
• download - Download model only
• run - Direct model inference
• list - List supported models

HTTP API Layer (src/api.rs)

The HTTP API layer provides REST endpoints:

• OpenAI-compatible: Matches OpenAI API format
• Streaming support: Real-time response generation
• Multi-modal: Handles text, images, and audio
• Thread-safe: Uses RwLock for concurrent requests

Endpoints:

• POST /chat/completions - Chat and text generation
• POST /images/remove_background - Image background removal
• POST /audio/speech - Text-to-speech synthesis

Model Abstraction Layer

All models implement the GenerateModel trait:

pub trait GenerateModel {
    
    // Generate response
    fn generate(&mut self, prompt: &str, params: GenerationParams) -> Result<String>;

    // Generate with streaming
    fn generate_stream(&mut self, prompt: &str, params: GenerationParams)
        -> Result<Box<dyn Iterator<Item = Result<String>>>>;
}

This provides:

• Polymorphism: Treat different models uniformly
• Extensibility: Easy to add new models
• Type safety: Compile-time guarantees

Utility Modules

Tokenizer (src/tokenizer/)

• Loads tokenizers from model configurations
• Handles special tokens
• Manages vocabulary

Chat Template (src/chat_template/)

• Formats chat messages into model prompts
• Supports multiple chat formats (ChatML, etc.)
• Handles system messages and role tags

Position Embeddings (src/position_embed/)

• Implements positional encoding for transformers
• Supports RoPE (Rotary Position Embedding)
• Handles M-RoPE for multimodal models

Utils (src/utils/)

• audio_utils.rs - Audio processing (WAV, MP3)
• image_utils.rs - Image processing (resize, encode/decode)
• tensor_utils.rs - Tensor utility methods
• mod.rs - Common utilities and constants

Design Patterns

1. Trait-Based Abstraction

The GenerateModel trait provides a unified interface:

// All models implement this trait
impl GenerateModel for Qwen3VL { /* ... */ }
impl GenerateModel for VoxCPM { /* ... */ }
impl GenerateModel for DeepSeekOCR { /* ... */ }

// Usage is model-agnostic
let mut model: Box<dyn GenerateModel> = load_model(model_type)?;
let result = model.generate(prompt, params)?;

2. Factory Pattern

Model loading uses a factory function:

pub fn load_model(
    model_type: &str,
    model_path: &str,
    device: &Device,
) -> Result<Box<dyn GenerateModel>> {
    match model_type {
        "qwen3vl-2b" => Ok(Box::new(qwen3vl::generate::Qwen3VLGenerate::init(...)?)),
        "voxcpm1.5" => Ok(Box::new(voxcpm::generate::VoxCPMGenerate::init(...)?)),
        // ... other models
        _ => Err(anyhow!("Unsupported model: {}", model_type)),
    }
}

3. Command Pattern

CLI subcommands encapsulate different operations:

match command {
    Commands::Cli { model, port, address } => { /* download and serve */ }
    Commands::Serv { model, weight_path, port } => { /* serve only */ }
    Commands::Download { model, save_dir } => { /* download only */ }
    Commands::Run { model, input, weight_path } => { /* direct inference */ }
    Commands::List => { /* list models */ }
}

Model Organization

Each model follows a consistent structure:

src/models/{model_name}/
├── config.rs       # Model configuration and generation parameters
├── model.rs        # Core model architecture (layers, attention)
├── generate.rs     # Inference logic (implements GenerateModel trait)
├── processor.rs    # Model-specific processing (for complex models)
└── mod.rs          # Module declaration and exports

Example: Qwen3VL

src/models/qwen3vl/
├── config.rs       # Qwen3VLConfig, GenerationConfig
├── model.rs        # Qwen3VL transformer layers, attention mechanisms
├── generate.rs     # Qwen3VLGenerate implementation
├── processor.rs    # Image and text processing for multimodal input
└── mod.rs          # Exports public API

Performance Optimizations

GPU Acceleration

AHA supports GPU acceleration through:

• CUDA: For NVIDIA GPUs (Linux, Windows)
• Metal: For Apple Silicon (macOS)

Enable with:

cargo build --features cuda    # NVIDIA GPUs
cargo build --features metal   # Apple Silicon

Flash Attention

Flash Attention optimizes long-sequence processing:

• Reduces memory usage
• Improves inference speed
• Especially beneficial for vision models

Enable with:

cargo build --features cuda,flash-attn

Memory-Mapped Tensors

Models use memory-mapped files for:

• Faster loading times
• Reduced memory footprint
• Concurrent model loading

Precision Optimization

Dynamic precision selection based on hardware:

• F32: Maximum accuracy (CPU-only)
• F16: Balanced performance (GPU)
• BF16: Best for modern GPUs

Security Considerations

Local-Only Processing

• No external API calls after model download
• No telemetry or data collection
• Data remains entirely on the local system

Memory Safety

• Rust's ownership system prevents memory leaks
• No buffer overflows or use-after-free bugs
• Thread-safe concurrent operations

Input Validation

• File size limits (5MB strings, 100MB files)
• Path validation to prevent directory traversal
• Type-safe request handling

Data Flow

Request Flow

┌─────────┐
│ Client  │
└────┬────┘
     │ HTTP Request
     ▼
┌──────────────────────────────────────────────────────────┐
│  Rocket HTTP Server                                      │
│  - Route request to endpoint                             │
│  - Parse request body                                    │
│  - Extract parameters                                    │
└────────────┬─────────────────────────────────────────────┘
             │
             ▼
┌──────────────────────────────────────────────────────────┐
│  API Handler (api.rs)                                    │
│  - Acquire model lock                                    │
│  - Prepare input (tokenize, process images/audio)        │
│  - Call model.generate() or generate_stream()           │
└────────────┬─────────────────────────────────────────────┘
             │
             ▼
┌──────────────────────────────────────────────────────────┐
│  Model Implementation (models/{model}/generate.rs)       │
│  - Load weights from memory-mapped files                 │
│  - Run forward pass through Candle tensors              │
│  - Decode output tokens                                  │
└────────────┬─────────────────────────────────────────────┘
             │
             ▼
┌──────────────────────────────────────────────────────────┐
│  Candle Framework                                        │
│  - Execute on CPU or GPU device                          │
│  - Manage tensor operations                              │
└────────────┬─────────────────────────────────────────────┘
             │
             ▼
┌──────────────────────────────────────────────────────────┐
│  Response Generation                                     │
│  - Format response (JSON / streaming)                    │
│  - Return to client                                      │
└──────────────────────────────────────────────────────────┘

Model Loading Flow

User specifies model
        │
        ▼
Check if --weight-path provided
        │
    ┌───┴───┐
    │       │
   Yes      No
    │       │
    ▼       ▼
Use local    Download from ModelScope
path         │
    │        ▼
    │    Save to ~/.aha/{model}/
    │        │
    └───┬────┘
        ▼
Load model weights into memory
        │
        ▼
Initialize model (init())
        │
        ▼
Ready for inference

Extension Points

Adding a New Model

1. Create model directory under src/models/
2. Implement GenerateModel trait
3. Add model to factory function in mod.rs
4. Add CLI mapping in main.rs
5. Add test case in tests/

Custom Processing

Models can override default processing:

• Custom tokenization
• Special input/output formats
• Model-specific optimizations

See Also

• Installation Guide - Setup and installation
• Getting Started - Quick start guide
• API Reference - REST API documentation
• Development - Contributing guide