Introduction to Mixture of Experts (MoE) Models in LLM

#learnai #oxlo #ai

We're going to build a lightweight MoE-style router that classifies incoming questions and dispatches them to specialized expert prompts, backed by Oxlo.ai's actual MoE and dense models. This gives you hands-on intuition for how Mixture of Experts architectures work at the application layer, without waiting weeks to train a sparse gate from scratch. If you handle varied workloads, routing to the right model is often cheaper and faster than sending everything to a single giant endpoint.

What you'll need

Python 3.10 or newer
The OpenAI SDK: pip install openai
An Oxlo.ai API key from https://portal.oxlo.ai

Step 1: Create the Oxlo.ai client

Set up the OpenAI-compatible client pointing at Oxlo.ai. I keep my key in an environment variable, but you can paste it directly for local testing.

import os
from openai import OpenAI

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key=os.getenv("OXLO_API_KEY", "YOUR_OXLO_API_KEY")
)

Step 2: Build the task gate

The gate is a tiny classifier prompt that returns a single tag: code, math, creative, or general. We run this against a fast dense model so the routing decision itself stays cheap.

GATE_PROMPT = """You are a routing gate. Read the user query and reply with exactly one tag from this list: code, math, creative, general. No explanation, just the tag."""

def classify_task(user_message: str) -> str:
    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": GATE_PROMPT},
            {"role": "user", "content": user_message},
        ],
        max_tokens=10,
    )
    tag = response.choices[0].message.content.strip().lower()
    return tag if tag in {"code", "math", "creative", "general"} else "general"

Step 3: Define the expert prompts

Each expert is just a system prompt tuned for a domain. This mirrors the expert networks inside a true MoE transformer. Here is the full set.

EXPERTS = {
    "code": """You are an expert software engineer. Write clean, commented code. Explain tradeoffs. Prefer Python unless asked otherwise.""",
    "math": """You are a mathematician. Show your reasoning step by step. Use LaTeX formatting for equations. Double-check arithmetic.""",
    "creative": """You are a creative writer. Be vivid but concise. Avoid clichés.""",
    "general": """You are a helpful assistant. Keep answers factual and brief."""
}

Step 4: Wire the router to Oxlo.ai

Now we connect the gate to the experts. For routine tasks, I use a dense model like Qwen 3 32B. The function returns the chosen tag, the model name, and the answer so you can debug routing decisions.

def ask_expert(user_message: str) -> dict:
    tag = classify_task(user_message)
    system_prompt = EXPERTS[tag]
    
    response = client.chat.completions.create(
        model="qwen-3-32b",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    )
    
    return {
        "tag": tag,
        "model": "qwen-3-32b",
        "answer": response.choices[0].message.content,
    }

Step 5: Add the MoE backbone for heavy reasoning

For hard math or deep reasoning, we swap in Oxlo.ai's DeepSeek V4 Flash. It is a sparse Mixture-of-Experts model with a 1M context window and near state-of-the-art open-source reasoning. Because Oxlo.ai uses flat per-request pricing instead of per-token pricing, sending a long multi-shot prompt to a massive MoE does not scale your cost with input length. See https://oxlo.ai/pricing for current plan details. That makes it practical to invoke the heavy model exactly when the gate decides it is needed.

def ask_expert(user_message: str, force_moe: bool = False) -> dict:
    tag = classify_task(user_message)
    system_prompt = EXPERTS[tag]
    
    # Route heavy reasoning to the MoE model
    model = "deepseek-v4-flash" if (force_moe or tag == "math") else "qwen-3-32b"
    
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    )
    
    return {
        "tag": tag,
        "model": model,
        "answer": response.choices[0].message.content,
    }

Step 6: Run it

Test the router with a routine coding question and a hard math problem. The gate should route coding to Qwen 3 32B and math to DeepSeek V4 Flash.

if __name__ == "__main__":
    # Routine coding task
    r1 = ask_expert("Write a Python function that computes Fibonacci numbers with memoization.")
    print(f"Tag: {r1['tag']} | Model: {r1['model']}")
    print(r1['answer'][:400] + "\n")
    
    # Hard reasoning task
    r2 = ask_expert(
        "Prove that the sum of the first n odd positive integers is n^2. Show every step.",
        force_moe=True
    )
    print(f"Tag: {r2['tag']} | Model: {r2['model']}")
    print(r2['answer'][:600])

Example output:

Tag: code | Model: qwen-3-32b
def fibonacci(n, memo=None):
    if memo is None:
        memo = {}
    if n in memo:
        return memo[n]
    if n <= 1:
        return n
    memo[n] = fibonacci(n - 1, memo) + fibonacci(n - 2, memo)
    return memo[n]

Tag: math | Model: deepseek-v4-flash
Proof by induction:
Base case (n = 1): The first odd positive integer is 1, and 1^2 = 1. So the statement holds.
Inductive step: Assume the sum of the first k odd integers is k^2. The (k+1)th odd integer is 2k+1. Adding this to both sides gives k^2 + 2k + 1 = (k+1)^2. Thus, by induction, the formula holds for all positive integers n.

Wrap-up and next steps

Replace the prompt-based gate with a lightweight classifier fine-tuned on your own logs, or add streaming by setting stream=True in the chat completion calls. If you want to go deeper, swap the dense backbone for GLM 5 or DeepSeek R1 671B MoE on Oxlo.ai for long-horizon agentic tasks.

DEV Community