axonml-llm

Overview

axonml-llm provides nine large-language-model architectures for the AxonML framework, all implemented in pure Rust on top of axonml-tensor and axonml-autograd. Shared infrastructure includes multi-head / causal self-attention with a KV cache, a FlashAttention kernel, RoPE, RMSNorm, a HuggingFace weight loader (safetensors), a state-dict name-mapping helper, a pretrained-weights hub with on-disk caching, a configurable text-generation sampler, and an HF-style tokenizer.

Supported architectures

Preset configurations

Shared building blocks

• Attention (attention) — MultiHeadSelfAttention, CausalSelfAttention, scaled_dot_product_attention, plus FlashAttention + FlashAttentionConfig. KVCache and per-layer LayerKVCache for incremental decoding.
• Transformer (transformer) — TransformerBlock, TransformerEncoder, TransformerDecoder with configurable depth/width/heads/activation and pre- or post-norm.
• Embeddings (embedding) — TokenEmbedding, PositionalEmbedding (sinusoidal), GPT2Embedding, BertEmbedding.
• Generation (generation) — GenerationConfig with greedy(), sampling(temp), top_k_sampling(k, temp), nucleus_sampling(p, temp), beam_search(beams), plus builder methods with_max_tokens, with_eos_token, with_repetition_penalty. TextGenerator drives the next-token logic.
• Weight loading — HFLoader (safetensors), load_llama_from_hf, load_mistral_from_hf, and generic LoadStateDict trait with LoadResult + map_hf_to_axonml / map_axonml_to_hf name mappers.
• Hub — PretrainedLLM registry (llm_registry(), list_models()), download_weights(name, force) with an on-disk cache under $XDG_CACHE_HOME/axonml/hub/llm.
• Tokenizer — HFTokenizer + SpecialTokens (HuggingFace-compatible tokenizer.json).

Modules

Usage

Add the crate to your Cargo.toml:

[dependencies]
axonml-llm = "0.6.1"

GPT-2 text generation

use axonml_llm::{GPT2LMHead, GPT2Config};
use axonml_tensor::Tensor;

let config = GPT2Config::small();
let model = GPT2LMHead::new(&config);

let input_ids = Tensor::from_vec(vec![50256u32, 1, 2, 3], &[1, 4]).unwrap();
let output  = model.generate(&input_ids, 50, 0.8, Some(50));
let greedy  = model.generate_greedy(&input_ids, 50);

BERT sequence classification

use axonml_llm::{BertForSequenceClassification, BertConfig};
use axonml_tensor::Tensor;

let config = BertConfig::base();
let model  = BertForSequenceClassification::new(&config, /*num_labels=*/ 2);

let input_ids = Tensor::from_vec(vec![101u32, 2054, 2003, 1996, 102], &[1, 5]).unwrap();
let logits = model.forward_classification(&input_ids);

BERT masked language modeling

use axonml_llm::{BertForMaskedLM, BertConfig};
use axonml_tensor::Tensor;

let model = BertForMaskedLM::new(&BertConfig::base());
let input_ids = Tensor::from_vec(vec![101u32, 2054, 103, 1996, 102], &[1, 5]).unwrap();
let logits = model.forward_mlm(&input_ids); // [batch, seq, vocab]

LLaMA / Mistral / Phi

use axonml_llm::{LLaMAForCausalLM, LLaMAConfig, MistralForCausalLM, MistralConfig, PhiForCausalLM, PhiConfig};

let llama   = LLaMAForCausalLM::new(&LLaMAConfig::llama2_7b());
let mistral = MistralForCausalLM::new(&MistralConfig::mistral_7b());
let phi     = PhiForCausalLM::new(&PhiConfig::phi2());

SSM (Mamba-style) causal LM

use axonml_llm::{SSMForCausalLM, SSMConfig};

let cfg = SSMConfig::from_d_model(/*d_model=*/ 512, /*vocab_size=*/ 32000);
let ssm = SSMForCausalLM::new(&cfg);

Hydra hybrid SSM + windowed attention

use axonml_llm::{HydraModel, HydraConfig};

let hydra = HydraModel::new(&HydraConfig::base()); // 768d, 24 layers, window 256

Chimera sparse MoE + differential attention

use axonml_llm::{ChimeraModel, ChimeraConfig};

let cfg = ChimeraConfig::default_2b(); // 8 experts per layer, top-2 active
let model = ChimeraModel::new(&cfg);
// Returns (logits, lb_loss) when used via forward_with_loss for training.

Trident 1.58-bit ternary LM

use axonml_llm::{TridentModel, TridentConfig};

let cfg = TridentConfig::default_150m(); // ternary weights + RoPE + GQA + ReLU² + SubLN
let model = TridentModel::new(&cfg);

Custom transformer encoder

use axonml_llm::TransformerEncoder;
use axonml_autograd::Variable;
use axonml_tensor::Tensor;

let encoder = TransformerEncoder::new(
    6,       // num_layers
    512,     // hidden_size
    8,       // num_heads
    2048,    // intermediate_size
    0.1,     // dropout
    1e-12,   // layer_norm_eps
    "gelu",  // activation
    false,   // pre_norm
);

let input  = Variable::new(Tensor::randn(&[2, 128, 512]), false);
let output = encoder.forward(&input);

Generation configuration

use axonml_llm::{GenerationConfig, TextGenerator};

let config = GenerationConfig::nucleus_sampling(0.95, 0.8)
    .with_max_tokens(100)
    .with_repetition_penalty(1.2)
    .with_eos_token(50256);

let generator = TextGenerator::new(config);
let next_token = generator.get_next_token(&logits, &generated_so_far);

Loading HuggingFace weights

use axonml_llm::{load_llama_from_hf, LLaMAConfig};

let mut model = axonml_llm::LLaMAForCausalLM::new(&LLaMAConfig::llama2_7b());
load_llama_from_hf(&mut model, "path/to/model.safetensors")?;

Pretrained hub

use axonml_llm::{llm_registry, download_llm_weights};

for m in llm_registry().values() {
    println!("{}  ({} params)", m.name, m.num_parameters);
}

let path = download_llm_weights("bert-base-uncased", /*force=*/ false)?;

Configuration reference

BERT

GPT-2

LLaMA / Mistral / Phi

Hybrid / SSM / MoE / Ternary

Generation strategies

Tests

cargo test -p axonml-llm
cargo test -p axonml-llm -- --nocapture

License

Licensed under either of:

• MIT License
• Apache License, Version 2.0

at your option.

Last updated: 2026-04-16 (v0.6.1)