Skip to content

Repository files navigation

SmolVLM on RK3588 NPU

High-performance implementation of SmolVLM-256M optimized for the Rockchip RK3588 NPU. Features novel quantization patterns and memory management techniques to enable stable FP16/INT8 hybrid inference on edge devices.

Key Features

Sandwich Quantization

The RK3588 NPU suffers from underflow when processing raw FP16 activations in deep layers. We wrap sensitive layers in a "Sandwich" of scalers:

  • InputScaler: Multiplies input by 10x before entering the NPU block.
  • OutputDescaler: Divides output by 10x after leaving the NPU block.

NanoTiled Attention

Standard self-attention requires transposing large matrices ($1024 \times 1024$), which triggers "Illegal Transpose" errors on the NPU. We decompose attention into small $32 \times 32$ tiles for safe transposes and block-by-block matmul.

Hybrid Split Architecture

The vision encoder is split into 24 separate shards (attention + MLP for each of 12 layers), distributed round-robin across NPU cores 0, 1, and 2.

Repository Structure

├── src/
│   ├── rknn_patterns/          # Reusable NPU-safe patterns (sandwich, tiling, blocks)
│   ├── smolvlm_convert/        # x86 conversion pipeline (torch → ONNX → RKNN)
│   ├── smolvlm_infer/          # On-device shard runner & orchestrator
│   ├── rkllm_bindings/          # ctypes wrapper for bundled librkllmrt.so
│   └── npu_backend/             # Low-level NPU backend (experimental)
│
├── scripts/
│   ├── convert.py               # Host-side RKNN conversion
│   ├── run_inference.py         # On-device inference entry point
│   ├── validate.py              # Layer-wise cosine similarity validation
│   ├── sweep_hybrid_fallback.py # Automated compile → push → validate sweep
│   ├── export_vanilla_baseline.py
│   ├── validate_vanilla_baseline.py
│   ├── test_zero_copy_scheduler.py
│   ├── debug/                   # NPU ioctl debugging scripts
│   ├── poc/                     # Proof-of-concept explorations (replay, DRM, mmap)
│   └── experiments/             # Ablation, benchmarking, and analysis scripts
│
├── docs/
│   ├── SBC_SSH_BENCH_GUIDE.md   # SSH benchmarking workflow for RK3588
│   ├── papers/                  # Workshop submissions, reviews, camera-ready plans
│   ├── research/                # SGTF spec, cross-platform research, surveys
│   ├── technical/               # Implementation notes, zero-copy reports
│   ├── reference/               # Vendor documentation (RKNN operator list)
│   └── archive/                 # Historical logs, baseline results, meeting notes
│
├── data/                        # Example input images
├── sbc_results_runs/            # Experiment result logs (CSV, text summaries)
├── external/                    # Third-party tools & kernel source (not tracked)
├── rk3588_open_compiler_project/ # Related research subproject (separate)
├── overleaf_clone/              # LaTeX paper source (separate git repo)
│
├── pyproject.toml
├── requirements_convert.txt     # Host (x86) dependencies
├── requirements_infer.txt       # SBC (aarch64) dependencies
├── build_rkllm_wheel.sh         # Build script for RKLLM wheel
└── MANIFEST.in

Getting Started

1. Conversion (x86 Host)

Requires Python 3.10 for rknn-toolkit2.

pip install -r requirements_convert.txt
python scripts/convert.py

This generates .rknn shards in smolvlm_subshards/.

2. Inference (RK3588)

Transfer smolvlm_subshards/ and src/ to the board.

pip install rkllm_bindings-0.1.0-py3-none-any.whl
pip install -r requirements_infer.txt
python scripts/run_inference.py

Validation

Layer-wise cosine similarity check against the CPU (PyTorch) baseline:

python scripts/validate.py --layers 0-11

Common flags:

  • --layers 0-11 or subsets (--layers 0 1 2)
  • --ablate-sandwich — disable sandwich scaling
  • --cumulative — error accumulation instead of teacher forcing
  • --model-path <shard_dir> — custom shard directory

Vanilla Baseline Validation

# Host export:
python scripts/export_vanilla_baseline.py --layer <L> --mode full|mlp

# SBC validation:
python scripts/validate_vanilla_baseline.py --layer <L> --vanilla-prefix <prefix>

SBC Benchmarking over SSH

For the full workflow (SSH access, per-layer validation, automated sweeps):

See docs/SBC_SSH_BENCH_GUIDE.md.

Documentation

Directory Contents
docs/papers/ Workshop submissions, review responses, camera-ready plans
docs/research/ SGTF spec, cross-platform research, literature surveys
docs/technical/ Implementation notes, zero-copy orchestration reports
docs/archive/ Historical baselines, failure analyses, experiment data
docs/reference/ Vendor reference documentation

About

No description, website, or topics provided.

Resources

Stars

50 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages