Implementation:Ollama Ollama XModels GLM4 MoE Lite

Knowledge Sources	Ollama
Domains	Model Architecture, MLX Runtime
Last Updated	2025-02-15 00:00 GMT

Overview

Implements the GLM4-MoE-Lite transformer model architecture for the MLX runner, supporting Multi-head Latent Attention (MLA) and Mixture of Experts (MoE).

Description

Registers under "Glm4MoeLiteForCausalLM" and "GLM4MoeLite" in the model registry. The architecture consists of alternating dense and MoE blocks. MLAAttention implements multi-latent attention with compressed KV projections, RoPE position encoding, and scaled dot-product attention. MoE uses a gated routing mechanism (MoEGate) with top-k expert selection, optional shared experts, and norm-based probability scaling. Expert weights are stacked for efficient batch processing using GatherQMM.

Usage

Automatically loaded when a model with the "Glm4MoeLiteForCausalLM" architecture is detected. This is the first model architecture implemented for the MLX runner.

Code Reference

Source Location

• Repository: Ollama
• File: x/models/glm4_moe_lite/glm4_moe_lite.go
• Lines: 1-788

Signature

type Config struct {
    HiddenSize            int32   `json:"hidden_size"`
    NumHiddenLayers       int32   `json:"num_hidden_layers"`
    NumAttentionHeads     int32   `json:"num_attention_heads"`
    QLoraRank             int32   `json:"q_lora_rank"`
    KVLoraRank            int32   `json:"kv_lora_rank"`
    NRoutedExperts        int32   `json:"n_routed_experts"`
    NumExpertsPerTok      int32   `json:"num_experts_per_tok"`
    // ... more fields
}

type MLAAttention struct { ... }
func (a *MLAAttention) Forward(x *mlx.Array, c cache.Cache, B, L int32, cfg *Config) *mlx.Array

type MoEGate struct { ... }
func (g *MoEGate) Forward(x *mlx.Array, cfg *Config) (*mlx.Array, *mlx.Array)

type Model struct { ... }
func (m *Model) Forward(inputs *mlx.Array, cache []cache.Cache) *mlx.Array
func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error

Import

import "github.com/ollama/ollama/x/models/glm4_moe_lite"

I/O Contract

Inputs

Name	Type	Required	Description
inputs	*mlx.Array	Yes	Token IDs tensor [batch, seq_len]
cache	[]cache.Cache	Yes	KV caches per layer

Outputs

Name	Type	Description
output	*mlx.Array	Hidden state tensor [batch, seq_len, hidden_size]

Usage Examples

// Registration happens in init()
func init() {
    base.Register("Glm4MoeLiteForCausalLM", newModel)
}

// Model is instantiated via the base registry
model, err := base.New(root)
if err != nil {
    return err
}

// Forward pass
output := model.Forward(tokenTensor, caches)
logits := model.Unembed(output)

Related Pages

• Principle:Ollama_Ollama_ModelArchitecture