Near-term roadmap

[jan-wassenberg]

We're sharing a roadmap of ideas for improving and speeding up Gemma. If you'd like to join in and help us get there faster, please reach out so we can coordinate :)

Threading

• (jan-wassenberg) detect: total #logical, per-logical: package, chiplet, core, smt
• (jan-wassenberg) detect: CPU name, L2D/L3 size
• (Z.A.) CCX-aware pinning - ready, awaiting Highway 1.2 release
• (jan-wassenberg) more efficient ThreadPool
• command line arg to disable pinning
• detect NUMA

[x] Dot product

• Add _mm*_dpbf16_ps to HWY_AVX3_SPR and HWY_AVX3_ZEN4 targets, plus define HWY_NATIVE_DOT_BF16 in set_macros-inl.h
• Faster SFP decode via table lookup
• Add new NEON_* target that uses vbfdot for ReorderWidenMulAccumulate
• If !defined(HWY_NATIVE_DOT_BF16) || !HWY_NATIVE_DOT_BF16, decompress bf16->f32 to temp array before MatVec (idea by Samuel, thank you!) - in Factor out deinterleaving of bf16 vectors for MatVecs. #166
• Apply even/odd trick to SFP

Matmul

• (pculliton) implement basic matmul and test. Not using BLAS because we want to fuse matmul and decompression.
• (pculliton) 4x4 unrolled and vectorized matmul
• (szabadka, B.B.) Update Prefill to use matmul (activation @ weights) instead of MatVec. Almost there.
• Fused decompression inside matmul
• Support offsets within the matrix, required by some call sites
• (jan-wassenberg) Decompress weights to bf16 when native
• (jan-wassenberg) Cache-aware tiling/packing
• (jan-wassenberg) NUMA aware
• (jan-wassenberg) 64-bit precision
• (B.B.) Larger batch size
• (A.V.) Avoid allocations for decompression

Compression

• (pculliton, A.R.) Eval infrastructure
• (A.R.) Arbiter model for eval
• (Ray) add metadata to tensors, remove RawWeights
• add TOC to BlobStore
• decide whether NUQ is launchable

Optimizations

• Replace attention matVec with matmul - requires reshaping a matrix
• Convert f32 activations to bf16 beforehand if HWY_NATIVE_DOT_BF16
• Integrate wraparound support into matmul
• Fuse softmax and sampling
• Vectorize RoPE
• Faster/more accurate hwy/contrib/math functions by updating the polynomials
• Vectorize RMSNorm
• (A.R.?, ...) Smaller KVCache: bf16, possibly reorder for better locality

Usability

• warn if unknown arguments given. std::map of known arg names?
• multiple .cc files to speed up builds
• Actionable error codes as return values: kLoadFailed, kSeqTooShort
• move eval/test files to tests/
• Ctrl+C signal handler to ensure profiler results are printed without requiring %q input
• add --prompt flag to run.cc
• random prompt generation for debug_prompt.cc

File format

• store ModelInfo in weights BlobStore
• store tensor info in BlobStore
• store tokenizer in BlobStore

New models

• (Daniel) Support PaliGemma
• Split Model into ModelFamily and ModelSize

[x] General infra

• (pculliton) Python wrapper
• (pculliton, ...) Improved CI: run on Kaggle infra
• AuxOut to hold timing info instead of printing in GenerateImpl.
• Sampling struct holds rng and temperature, to reduce length of args
• (P. C.) use new HWY_EXPORT_T to simplify dispatch - ready, awaiting Highway 1.2 release

[jan-wassenberg]

Making good progress :)

[MathiasSchindler]

Is Paligemma part of the scope of gemma.cpp?

[jan-wassenberg]

Let's discuss in #185 :)