Add Model

Add Model — Adapt New LLM Architectures to xinfer

Phase 0: Gather Required Information

Before starting, the agent must collect the inputs below. If any are missing, ask the user explicitly.

Resolution order (try top to bottom, stop at the first that succeeds):

1. Local model path (preferred) — If the user provides a local directory containing safetensors, read config.json and inspect weight tensors directly from disk. If the user provides a .gguf file, extract metadata and tensor info from the GGUF header (GGUF is self-contained; there is no separate config.json). No further input is needed in either case.
2. HuggingFace model ID — Look for the model in the local HuggingFace Hub cache (~/.cache/huggingface/hub/). If not cached, fetch config.json from the HuggingFace model repo. If that also fails, fall back to step 3.
3. Manual input — Ask the user to provide:
   • config.json contents (for safetensors models) or GGUF metadata (for GGUF models).
   • Weight tensor info (names, shapes, dtypes). The user can obtain this by clicking a weight file in the HuggingFace model repo.

Additionally, ask the user if they can provide the Python reference implementation (modeling_<arch>.py from HuggingFace Transformers). This is not strictly required, but significantly improves accuracy — it clarifies the exact forward pass, attention variants, MoE routing, activation functions, and normalization order that config fields alone cannot fully describe.

Extracting info from local weights

For safetensors (Python required):

# scripts/inspect_weights.py
import json, sys, struct
path = sys.argv[1]
with open(path, "rb") as f:
    n = struct.unpack("<Q", f.read(8))[0]
    header = json.loads(f.read(n))
for k, v in sorted(header.items()):
    if k != "__metadata__":
        print(f"{k}\t{v.get('shape')}\t{v.get('dtype')}")

For GGUF, use the gguf_helper CLI or Python gguf library to list tensors and metadata keys.

Phase 1: Analyze the Architecture

Read and compare:

1. config.json — identify: architectures, hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, head_dim, intermediate_size, hidden_act, vocab_size, rope_theta, partial_rotary_factor, sliding_window, tie_word_embeddings.
2. MoE fields (if any): num_experts, num_experts_per_tok / top_k_experts, moe_intermediate_size, norm_topk_prob, routed_scaling_factor.
3. Special features: attention_k_eq_v, layer_types array, layer_scalar, final_logit_softcapping, attn_logit_softcapping, dual head_dim, dual RoPE, per-expert scaling, shared experts, etc.
4. Python reference — understand the forward pass, especially:
   • Attention: GQA/MQA, head dim, RoPE variant, sliding window, QK norm, V norm
   • MLP: standard SiLU-gate, GeluTanh, custom SwiGLU, etc.
   • MoE: router structure, expert computation, parallel dense+MoE, renormalization
   • Layer structure: pre-norm vs post-norm, residual connections, layer scalar

Identify which existing model in src/models/ is the closest match. Use it as a starting template.

Architecture classification

Phase 2: Check if New Kernels Are Needed

If the model requires operators not in attention.rs or src/models/layers/:

1. Clone attention.rs to a local sibling directory (if not already present):

   cd .. && git clone https://github.com/guoqingbao/attention.rs.git

   Then switch to the commit that this project relying on.
2. Point xinfer/Cargo.toml to local attention.rs:

   # In [dependencies], change the git URL to:
   attention-rs = { path = "../attention.rs", ... }

3. Add both CUDA and Metal kernel implementations:
   • CUDA kernel: attention.rs/src/kernels/src/<name>.cu
   • Metal kernel: attention.rs/src/metal-kernels/src/<name>.metal
   • FFI bindings: attention.rs/src/kernels/src/ffi.rs
   • Rust wrapper: attention.rs/src/<name>.rs
   • Update attention.rs/src/lib.rs: add pub mod <name>;
   • Update attention.rs/src/kernels/build.rs: add rerun-if-changed for .cu
   • Update attention.rs/src/metal-kernels/build.rs: add to METAL_SOURCES
4. For CUDA BF16 kernels (Ampere+ only), guard with #ifndef NO_BF16_KERNEL and provide F16 fallback or dummy stubs for older GPUs.
5. Verify compilation: cargo check --features metal (macOS) or cargo check --features cuda (Linux).

Phase 3: Implement the Model

3a. Create model file

Create src/models/<arch>.rs. Follow the established pattern from the closest existing model:

Key struct pattern:

pub struct <Arch>DecoderLayer {
    self_attn: Attention,
    mlp: MLP,
    moe: Option<<Arch>MoE>,      // if MoE
    input_layernorm: NormX,
    post_attention_layernorm: NormX,
    // ... additional norms for hybrid MoE
}

pub struct <Arch>ForCausalLM {
    embed_tokens: candle_nn::Embedding,
    layers: Vec<<Arch>DecoderLayer>,
    norm: NormX,
    lm_head: ReplicatedLinear,
    // ...
}

Required public methods on <Arch>ForCausalLM:

pub fn new(vb: &VarBuilderX, comm: Rc<Comm>, config: &Config, dtype: DType,
           is_rope_i: bool, device: &Device,
           progress_reporter: Arc<RwLock<Box<dyn ProgressLike>>>) -> Result<Self>

pub fn forward(&self, input_ids: &Tensor, positions: &Tensor,
               kv_caches: Option<&Vec<(Tensor, Tensor)>>,
               input_metadata: &InputMetadata,
               _embeded_inputs: bool) -> Result<Tensor>

pub fn forward_embedding(&self, ..same args..) -> Result<Tensor>

pub fn get_vocab_size(&self) -> usize

The 5th _embeded_inputs: bool argument is required by the model_call! macro.

3b. Weight loading paths

The model must support two VarBuilder types via vb.is_qvar_builder():

For norm and lm_head:

3c. MoE construction

For models with MoE, the Gemma4MoE / FusedMoe selection pattern:

let moe = if is_qvar_builder {
    FusedMoeGGUF::new(config, vb.clone(), comm.clone(), dtype)?
} else if quant_config == "mxfp4" {
    FusedMoeMxfp4::new(config, vb.pp("mlp"), comm.clone(), dtype)?
} else if config.quant.is_some() {
    FusedMoeISQ::new(config, vb.pp("mlp"), comm.clone(), dtype)?
} else {
    FusedMoe::new(config, vb.pp("mlp"), comm.clone(), dtype)?
};

Important: If the model's router gate is NOT at the standard mlp.gate path (e.g. Gemma4 uses router.proj), use FusedMoe::new_with_gate(config, gate_vb, experts_vb, ...) and FusedMoeISQ::new_with_gate(...).

For GGUF models with packed ffn_gate_up_exps (instead of separate ffn_gate_exps + ffn_up_exps), FusedMoeGGUF::new() auto-detects and handles this.

3d. Packed expert layout

If the model uses packed gate_up_proj with a non-standard layout, add the architecture name to the layout resolvers in src/models/layers/moe.rs:

• resolve_packed_gate_up_layout() — InterPacked if shape is [experts, 2*intermediate, hidden]
• resolve_packed_down_layout() — HiddenInter if shape is [experts, hidden, intermediate]

Phase 4: Register the Model

4a. src/models/mod.rs

pub mod <arch>;

4b. src/utils/config.rs

Add variant to ModelType enum:

pub enum ModelType {
    // ...
    <Arch>,
}

4c. src/utils/mod.rs — get_arch_rope function

Add mappings in order:

1. GGUF canonical_arch: "<gguf_arch>" => "<HFArchitectureName>".to_string()
2. Architecture to ModelType + chat template: Add match arm for HF architecture names
3. Rope type selection: ("<arch_lower>", false) in the rope map
4. HF config parsing (multimodal wrapper): If the model wraps text_config, add handler to extract nested config, MoE config, and rope_parameters
5. GGUF extra_config_json: If the model has special GGUF metadata (sliding_window_pattern, layer_types, etc.), add extraction block
6. GGUF MoE config: Add mod_cfg construction from GGUF expert metadata
7. hidden_act override: If the model uses a non-Silu activation, override after config construction
8. require_model_penalty(): Add architecture names
9. Multimodal GGUF warning: Add if applicable