Liquid AI LFM2.5-2.6B — The 2.6B Parameter Edge Model That Defies Scaling Laws: Architecture Revolution and Deployment Practice

"""
NAS Simulation: Evolutionary Architecture Search
Finding the optimal conv/attention ratio for LFM2.5-2.6B
"""
import numpy as np
from dataclasses import dataclass, field
from typing import List, Optional, Tuple
import random
import math

@dataclass
class Architecture:
    """Individual architecture encoding"""
    layer_types: List[str]  # 'conv' or 'gqa'
    conv_kernel: int
    kv_heads: int
    ffn_scale: float
    hidden_dim: int = 2048
    total_layers: int = 30
    
    def kv_cache_size(self, seq_len: int = 131072) -> float:
        """Estimate KV cache size in MB"""
        gqa_count = sum(1 for t in self.layer_types if t == 'gqa')
        if gqa_count == 0:
            return 0.0
        head_dim = self.hidden_dim // 32
        # 2 (K+V) * kv_heads * head_dim * seq_len * 2 bytes (FP16)
        bytes_per_layer = 2 * self.kv_heads * head_dim * seq_len * 2
        return (bytes_per_layer * gqa_count) / (1024 * 1024)
    
    def compute_cost(self) -> float:
        """Relative compute cost estimation"""
        conv_cost = sum(1 for t in self.layer_types if t == 'conv') * self.conv_kernel * 0.3
        gqa_cost = sum(1 for t in self.layer_types if t == 'gqa') * 1.0
        ffn_cost = self.total_layers * self.ffn_scale * 0.4
        return conv_cost + gqa_cost + ffn_cost
    
    def num_gqa(self) -> int:
        return sum(1 for t in self.layer_types if t == 'gqa')
    
    def num_conv(self) -> int:
        return self.total_layers - self.num_gqa()

def random_architecture() -> Architecture:
    """Generate random architecture"""
    total = 30
    num_gqa = random.randint(4, 12)
    positions = sorted(random.sample(range(total), num_gqa))
    types = ['gqa' if i in positions else 'conv' for i in range(total)]
    
    return Architecture(
        layer_types=types,
        conv_kernel=random.choice([3, 5, 7]),
        kv_heads=random.choice([8, 16]),
        ffn_scale=random.choice([3.0, 3.5, 4.0, 5.25])
    )

def crossover(a1: Architecture, a2: Architecture) -> Architecture:
    """Crossover two architectures"""
    child_types = [
        random.choice([a1.layer_types[i], a2.layer_types[i]])
        for i in range(len(a1.layer_types))
    ]
    return Architecture(
        layer_types=child_types,
        conv_kernel=random.choice([a1.conv_kernel, a2.conv_kernel]),
        kv_heads=random.choice([a1.kv_heads, a2.kv_heads]),
        ffn_scale=random.choice([a1.ffn_scale, a2.ffn_scale])
    )

def mutate(arch: Architecture) -> Architecture:
    """Mutate architecture"""
    new_types = arch.layer_types.copy()
    
    # Swap a conv and a gqa
    if random.random() < 0.3:
        conv_indices = [i for i, t in enumerate(new_types) if t == 'conv']
        gqa_indices = [i for i, t in enumerate(new_types) if t == 'gqa']
        if conv_indices and gqa_indices:
            ci = random.choice(conv_indices)
            gi = random.choice(gqa_indices)
            new_types[ci], new_types[gi] = new_types[gi], new_types[ci]
    
    return Architecture(
        layer_types=new_types,
        conv_kernel=arch.conv_kernel if random.random() > 0.2 else random.choice([3, 5, 7]),
        kv_heads=arch.kv_heads if random.random() > 0.2 else random.choice([8, 16]),
        ffn_scale=arch.ffn_scale if random.random() > 0.2 else random.choice([3.0, 3.5, 4.0, 5.25])
    )

def fitness(arch: Architecture, target_kv_mb: float = 500.0) -> float:
    """Fitness function: balance capability vs resource constraints"""
    kv = arch.kv_cache_size()
    if kv > target_kv_mb * 1.5:
        return -float('inf')
    
    gqa_count = arch.num_gqa()
    conv_count = arch.num_conv()
    
    # GQA provides global reasoning capability
    score = gqa_count * 1.5 + conv_count * 0.8 - arch.compute_cost() * 0.3
    
    # Prefer evenly distributed GQA layers
    gqa_positions = [i for i, t in enumerate(arch.layer_types) if t == 'gqa']
    if gqa_positions:
        spread = (max(gqa_positions) - min(gqa_positions)) / len(arch.layer_types)
        score += spread * 2.0
    
    return score

def evolutionary_search(
    pop_size: int = 50,
    generations: int = 100,
    elite_ratio: float = 0.2
) -> List[Architecture]:
    """Evolutionary NAS search"""
    population = [random_architecture() for _ in range(pop_size)]
    best_archs = []
    
    for gen in range(generations):
        scored = [(arch, fitness(arch)) for arch in population]
        scored.sort(key=lambda x: x[1], reverse=True)
        
        best_archs.append(scored[0][0])
        
        elite_count = int(pop_size * elite_ratio)
        elites = [arch for arch, _ in scored[:elite_count]]
        
        next_gen = elites.copy()
        while len(next_gen) < pop_size:
            parent = random.choice(elites)
            if random.random() < 0.7 and len(elites) > 1:
                parent2 = random.choice(elites)
                child = crossover(parent, parent2)
            else:
                child = parent
            
            if random.random() < 0.4:
                child = mutate(child)
            
            next_gen.append(child)
        
        population = next_gen
        
        if (gen + 1) % 20 == 0:
            best = scored[0][0]
            print(f"Gen {gen+1}: fitness={scored[0][1]:.2f}, "
                  f"GQA={best.num_gqa()}, Conv={best.num_conv()}, "
                  f"KV cache={best.kv_cache_size():.0f}MB")
    
    return best_archs

if __name__ == "__main__":
    print("Neural Architecture Search Simulation")
    print("=" * 50)
    print(f"Search space: 30 layers, kernel=[3,5,7], KV heads=[8,16]")
    print()
    
    results = evolutionary_search(pop_size=60, generations=80)
    
    final = results[-1]
    print(f"\nBest Architecture Found:")
    print(f"  Total layers: {final.total_layers}")
    print(f"  ConvBlocks: {final.num_conv()}")
    print(f"  GQA layers: {final.num_gqa()}")
    print(f"  Conv kernel: {final.conv_kernel}")
    print(f"  KV heads: {final.kv_heads}")
    print(f"  FFN scale: {final.ffn_scale}")
    print(f"  KV cache (128K): {final.kv_cache_size():.0f}MB")
    print(f"  Layer pattern: {''.join('C' if t == 'conv' else 'A' for t in final.layer_types)}")
    print(f"  (C=ConvBlock, A=GQA Attention)")

"""
MOPD (Multi-Domain On-Policy Distillation) Implementation
Core algorithm for distilling multiple domain experts into one student model
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Dict, Optional
from dataclasses import dataclass

@dataclass
class MOPDConfig:
    """MOPD configuration"""
    vocab_size: int = 128000
    hidden_dim: int = 2048
    num_teachers: int = 4
    kl_weight: float = 0.5
    ce_weight: float = 1.0
    temperature: float = 2.0
    domain_weights: List[float] = None

class MOPDLoss(nn.Module):
    """
    Multi-Domain On-Policy Distillation Loss
    
    L = α * L_CE + β * T² * Σ(domain_weight * KL(P_teacher || P_student))
    """
    
    def __init__(self, config: MOPDConfig):
        super().__init__()
        self.config = config
        if config.domain_weights is None:
            self.config.domain_weights = [1.0, 1.0, 1.0, 1.0]
    
    def forward(
        self,
        student_logits: torch.Tensor,
        teacher_logits_list: List[torch.Tensor],
        labels: torch.Tensor,
        domain_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        """
        Compute MOPD loss
        
        Args:
            student_logits: [batch, seq_len, vocab]
            teacher_logits_list: list of 4 teacher logits
            labels: target token ids [batch, seq_len]
            domain_ids: per-sample domain labels [batch]
            attention_mask: [batch, seq_len]
        """
        batch_size, seq_len, vocab_size = student_logits.shape
        
        if attention_mask is None:
            attention_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
        
        # 1. Cross-entropy loss (standard language modeling)
        ce_loss = F.cross_entropy(
            student_logits.view(-1, vocab_size),
            labels.view(-1),
            reduction='none'
        ).view(batch_size, seq_len)
        ce_loss = (ce_loss * attention_mask).sum() / attention_mask.sum()
        
        # 2. KL divergence loss (distillation)
        student_log_probs = F.log_softmax(
            student_logits / self.config.temperature, dim=-1
        )
        
        kl_loss = 0.0
        for domain_idx in range(self.config.num_teachers):
            domain_mask = (domain_ids == domain_idx)
            if domain_mask.sum() == 0:
                continue
            
            teacher_probs = F.softmax(
                    teacher_logits_list[domain_idx] / self.config.temperature, 
                    dim=-1
            )
            
            domain_kl = F.kl_div(
                student_log_probs[domain_mask],
                teacher_probs[domain_mask],
                reduction='sum',
                log_target=False
            )
            kl_loss += self.config.domain_weights[domain_idx] * domain_kl
        
        kl_loss = kl_loss / attention_mask.sum()
        
        # 3. Combined loss with temperature scaling
        total_loss = (self.config.ce_weight * ce_loss + 
                     self.config.kl_weight * kl_loss * 
                     (self.config.temperature ** 2))
        
        return {
            'total_loss': total_loss,
            'ce_loss': ce_loss,
            'kl_loss': kl_loss,
        }

class GRPOTrainer:
    """
    GRPO (Group Relative Policy Optimization) for Agentic RL
    
    The key idea: use a group of trajectories to compute advantages
    relative to the group mean, avoiding the need for a separate value model.
    """
    
    def __init__(
        self,
        clip_epsilon: float = 0.2,
        kl_coeff: float = 0.01,
        group_size: int = 8
    ):
        self.clip_epsilon = clip_epsilon
        self.kl_coeff = kl_coeff
        self.group_size = group_size
    
    def compute_advantage(self, rewards: List[float]) -> torch.Tensor:
        """
        Compute group-relative advantages
        
        A_i = (R_i - mean(R)) / std(R)
        """
        rewards_t = torch.tensor(rewards)
        mean = rewards_t.mean()
        std = rewards_t.std() + 1e-8
        return (rewards_t - mean) / std
    
    def compute_loss(
        self,
        old_log_probs: torch.Tensor,
        new_log_probs: torch.Tensor,
        advantages: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute GRPO loss
        
        L_GRPO = -E[min(r * A, clip(r, 1-ε, 1+ε) * A)] + β * KL
        where r = exp(new_log_prob - old_log_prob)
        """
        ratios = torch.exp(new_log_probs - old_log_probs)
        
        surr1 = ratios * advantages
        surr2 = torch.clamp(ratios, 
                           1.0 - self.clip_epsilon, 
                           1.0 + self.clip_epsilon) * advantages
        
        policy_loss = -torch.min(surr1, surr2).mean()
        kl_penalty = (old_log_probs - new_log_probs).mean()
        
        return policy_loss + self.kl_coeff * kl_penalty


def demo_mopd():
    """Demonstrate MOPD training pipeline"""
    config = MOPDConfig()
    loss_fn = MOPDLoss(config)
    
    print("MOPD Training Pipeline Demo")
    print("=" * 60)
    print(f"  Vocabulary: {config.vocab_size:,}")
    print(f"  Teachers: {config.num_teachers} (math, code, tool, reasoning)")
    print(f"  Temperature: {config.temperature}")
    print(f"  KL weight: {config.kl_weight}, CE weight: {config.ce_weight}")
    print()
    
    # Simulate training data
    batch_size, seq_len = 4, 512
    dummy_student = torch.randn(batch_size, seq_len, config.vocab_size)
    dummy_teachers = [
        torch.randn(batch_size, seq_len, config.vocab_size)
        for _ in range(config.num_teachers)
    ]
    dummy_labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    dummy_domains = torch.randint(0, config.num_teachers, (batch_size,))
    
    losses = loss_fn(dummy_student, dummy_teachers, dummy_labels, dummy_domains)
    
    print(f"Loss components:")
    print(f"  Total: {losses['total_loss']:.4f}")
    print(f"  CE: {losses['ce_loss']:.4f}")
    print(f"  KL: {losses['kl_loss']:.4f}")
    
    # GRPO demo
    grpo = GRPOTrainer(group_size=8)
    rewards = [0.2, 0.5, 0.8, 1.0, 0.3, 0.6, 0.9, 0.4]
    advantages = grpo.compute_advantage(rewards)
    
    print(f"\nGRPO Demo (group_size=8):")
    print(f"  Rewards: {[f'{r:.2f}' for r in rewards]}")
    print(f"  Advantages: {[f'{a:.2f}' for a in advantages.tolist()]}")
    
    old_log_probs = torch.randn(8) * 0.1
    new_log_probs = old_log_probs + torch.randn(8) * 0.05
    loss = grpo.compute_loss(old_log_probs, new_log_probs, advantages)
    print(f"  GRPO loss: {loss:.4f}")

if __name__ == "__main__":
    demo_mopd()

// engine.go - Edge inference engine with GQA-optimized KV cache
package main

import (
	"fmt"
	"math"
	"sync"
	"time"
)

// ============================================================
// Core Data Types
// ============================================================

// ModelConfig holds LFM2.5-2.6B architecture parameters
type ModelConfig struct {
	HiddenDim  int
	NumLayers  int
	NumQHeads  int
	NumKVHeads int
	ConvKernel int
	VocabSize  int
	MaxSeqLen  int
	FFNScale   float64
	RoPETheta  float64
	LayerTypes []string
}

// LFM25Config returns the default LFM2.5-2.6B configuration
func LFM25Config() ModelConfig {
	// 30 layers: 22 ConvBlocks + 8 GQA
	// GQA positions: every 2-3 ConvBlocks, insert one GQA
	layerTypes := make([]string, 30)
	gqaPos := map[int]bool{
		2: true, 5: true, 8: true, 11: true,
		14: true, 18: true, 22: true, 26: true,
	}
	for i := 0; i < 30; i++ {
		if gqaPos[i] {
			layerTypes[i] = "gqa"
		} else {
			layerTypes[i] = "conv"
		}
	}
	return ModelConfig{
		HiddenDim:  2048,
		NumLayers:  30,
		NumQHeads:  32,
		NumKVHeads: 8,
		ConvKernel: 3,
		VocabSize:  128000,
		MaxSeqLen:  131072,
		FFNScale:   5.25,
		RoPETheta:  10000000.0,
		LayerTypes: layerTypes,
	}
}

// ============================================================
// GQA-Optimized KV Cache
// ============================================================

// KVCache implements GQA-aware key-value cache with memory tracking
type KVCache struct {
	keys      [][][]float32 // [layer][head][pos*headDim]
	values    [][][]float32
	headDim   int
	mu        sync.RWMutex
}

// NewKVCache creates a new KV cache
func NewKVCache(config ModelConfig) *KVCache {
	numGQA := 0
	for _, t := range config.LayerTypes {
		if t == "gqa" {
			numGQA++
		}
	}

	headDim := config.HiddenDim / config.NumQHeads

	kc := &KVCache{
		keys:    make([][][]float32, numGQA),
		values:  make([][][]float32, numGQA),
		headDim: headDim,
	}

	for l := 0; l < numGQA; l++ {
		kc.keys[l] = make([][]float32, config.NumKVHeads)
		kc.values[l] = make([][]float32, config.NumKVHeads)
		for h := 0; h < config.NumKVHeads; h++ {
			// Pre-allocate with reasonable capacity
			kc.keys[l][h] = make([]float32, 0, config.MaxSeqLen*headDim)
			kc.values[l][h] = make([]float32, 0, config.MaxSeqLen*headDim)
		}
	}

	return kc
}

// Append adds new K/V pairs to the cache
func (kc *KVCache) Append(layerIdx, head int, k, v []float32) {
	kc.mu.Lock()
	defer kc.mu.Unlock()
	kc.keys[layerIdx][head] = append(kc.keys[layerIdx][head], k...)
	kc.values[layerIdx][head] = append(kc.values[layerIdx][head], v...)
}

// Get retrieves K/V from cache for the last `numTokens` positions
func (kc *KVCache) Get(layerIdx, head, numTokens int) ([]float32, []float32) {
	kc.mu.RLock()
	defer kc.mu.RUnlock()

	hd := kc.headDim
	keys := kc.keys[layerIdx][head]
	vals := kc.values[layerIdx][head]

	start := len(keys) - numTokens*hd
	if start < 0 {
		start = 0
	}

	return keys[start:], vals[start:]
}

// MemoryUsageMB returns approximate memory usage in MB
func (kc *KVCache) MemoryUsageMB() float64 {
	kc.mu.RLock()
	defer kc.mu.RUnlock()

	var totalBytes int64
	for l := range kc.keys {
		for h := range kc.keys[l] {
			totalBytes += int64(len(kc.keys[l][h])) * 4
			totalBytes += int64(len(kc.values[l][h])) * 4
		}
	}
	return float64(totalBytes) / (1024 * 1024)
}

// ============================================================
# RoPE (Rotary Position Embedding)
// ============================================================

// RoPECache pre-computes sin/cos values for all positions up to maxSeqLen
type RoPECache struct {
	sin [][]float32
	cos [][]float32
}

// NewRoPECache pre-computes RoPE values
func NewRoPECache(config ModelConfig) *RoPECache {
	headDim := config.HiddenDim / config.NumQHeads
	rope := &RoPECache{
		sin: make([][]float32, config.MaxSeqLen),
		cos: make([][]float32, config.MaxSeqLen),
	}

	for pos := 0; pos < config.MaxSeqLen; pos++ {
		rope.sin[pos] = make([]float32, headDim)
		rope.cos[pos] = make([]float32, headDim)
		for d := 0; d < headDim; d += 2 {
			theta := math.Pow(config.RoPETheta, -float64(d)/float64(headDim))
			angle := float64(pos) * theta
			sin := float32(math.Sin(angle))
			cos := float32(math.Cos(angle))
			rope.sin[pos][d] = sin
			rope.cos[pos][d] = cos
			if d+1 < headDim {
				rope.sin[pos][d+1] = sin
				rope.cos[pos][d+1] = cos
			}
		}
	}
	return rope
}

// Apply applies RoPE to query or key at a given position
func (rope *RoPECache) Apply(x []float32, pos int, headDim int) []float32 {
	result := make([]float32, len(x))
	copy(result, x)

	for h := 0; h < len(x)/headDim; h++ {
		offset := h * headDim
		for d := 0; d < headDim; d += 2 {
			i := offset + d
			j := offset + d + 1
			if j >= len(x) {
				break
			}
			sin := rope.sin[pos][d]
			cos := rope.cos[pos][d]
			result[i] = x[i]*cos - x[j]*sin
			result[j] = x[i]*sin + x[j]*cos
		}
	}
	return result
}

// ============================================================
// GQA Attention Layer
// ============================================================

// GQALayer implements Grouped Query Attention
type GQALayer struct {
	numQHeads  int
	numKVHeads int
	headDim    int
	groupSize  int
}

// NewGQALayer creates a GQA attention layer
func NewGQALayer(numQHeads, numKVHeads, headDim int) *GQALayer {
	return &GQALayer{
		numQHeads:  numQHeads,
		numKVHeads: numKVHeads,
		headDim:    headDim,
		groupSize:  numQHeads / numKVHeads,
	}
}

// Forward runs GQA attention forward pass
func (gqa *GQALayer) Forward(
	x []float32,
	pos int,
	kvCache *KVCache,
	gqaIdx int,
	rope *RoPECache,
) []float32 {
	hd := gqa.headDim
	numKV := gqa.numKVHeads

	// Simulate QKV projections (in practice, these are learned weights)
	q := make([]float32, gqa.numQHeads*hd)
	k := make([]float32, numKV*hd)
	v := make([]float32, numKV*hd)
	for i := range q {
		q[i] = x[i%len(x)] * 0.1
	}
	for i := range k {
		k[i] = x[i%len(x)] * 0.1
		v[i] = x[i%len(x)] * 0.1
	}

	// Apply RoPE
	q = rope.Apply(q, pos, hd)
	k = rope.Apply(k, pos, hd)

	// Update KV cache
	for h := 0; h < numKV; h++ {
		offset := h * hd
		kvCache.Append(gqaIdx, h, k[offset:offset+hd], v[offset:offset+hd])
	}

	// GQA: each KV head serves groupSize Q heads
	output := make([]float32, gqa.numQHeads*hd)
	scale := float32(math.Sqrt(float64(hd)))

	for qh := 0; qh < gqa.numQHeads; qh++ {
		kvh := qh / gqa.groupSize
		qOffset := qh * hd

		// Retrieve KV from cache
		cachedK, cachedV := kvCache.Get(gqaIdx, kvh, pos+1)
		numCached := len(cachedK) / hd

		// Compute attention over all cached positions
		var attnSum float32
		outVec := make([]float32, hd)

		for cp := 0; cp < numCached; cp++ {
			// Score = Q * K^T / sqrt(d)
			var score float32
			for d := 0; d < hd; d++ {
				score += q[qOffset+d] * cachedK[cp*hd+d]
			}
			score /= scale
			attn := float32(math.Exp(float64(score)))

			for d := 0; d < hd; d++ {
				outVec[d] += attn * cachedV[cp*hd+d]
			}
			attnSum += attn
		}

		// Normalize
		if attnSum > 1e-10 {
			for d := 0; d < hd; d++ {
				output[qOffset+d] = outVec[d] / attnSum
			}
		}
	}

	return output
}

// ============================================================
// ConvBlock
// ============================================================

// ConvBlock implements a double-gated short convolution block
type ConvBlock struct {
	hiddenDim int
	kernel    int
}

// NewConvBlock creates a convolution block
func NewConvBlock(hiddenDim, kernel int) *ConvBlock {
	return &ConvBlock{hiddenDim: hiddenDim, kernel: kernel}
}

// Forward runs causal convolution forward pass
func (cb *ConvBlock) Forward(x []float32, stateCache [][]float32, step int) []float32 {
	output := make([]float32, len(x))
	copy(output, x)

	for d := 0; d < cb.hiddenDim; d++ {
		var sum float32
		for k := 0; k < cb.kernel; k++ {
			pos := step - k
			if pos < 0 {
				continue
			}
			// Gate: closer positions have higher weight
			gate := float32(1.0 / float64(k+1))
			sum += stateCache[pos][d] * gate
		}
		output[d] = sum * 0.1
	}

	return output
}

// ============================================================
// SwiGLU FFN
// ============================================================

// SwiGLUFFN implements SwiGLU-activated feed-forward network
type SwiGLUFFN struct {
	hiddenDim    int
	intermediate int
}

func swish(x float32) float32 {
	return x / (1 + float32(math.Exp(float64(-x))))
}

func (ffn *SwiGLUFFN) Forward(x []float32) []float32 {
	output := make([]float32, len(x))
	for i, v := range x {
		gate := swish(v)
		// SwiGLU: output = swish(xW1) * (xV1) * W2 (simplified)
		output[i] = gate * v * 0.5
	}
	return output
}

// ============================================================
// Inference Engine
// ============================================================

// InferenceEngine is the complete edge inference runtime
type InferenceEngine struct {
	config  ModelConfig
	kvCache *KVCache
	rope    *RoPECache
	gqa     *GQALayer
	conv    *ConvBlock
	ffn     *SwiGLUFFN
}

// NewInferenceEngine creates a new inference engine
func NewInferenceEngine(config ModelConfig) *InferenceEngine {
	headDim := config.HiddenDim / config.NumQHeads
	intermediate := int(float64(config.HiddenDim) * config.FFNScale)

	return &InferenceEngine{
		config:  config,
		kvCache: NewKVCache(config),
		rope:    NewRoPECache(config),
		gqa:     NewGQALayer(config.NumQHeads, config.NumKVHeads, headDim),
		conv:    NewConvBlock(config.HiddenDim, config.ConvKernel),
		ffn:     NewSwiGLUFFN(config.HiddenDim, intermediate),
	}
}

// Generate generates tokens autoregressively
func (e *InferenceEngine) Generate(prompt []int, maxTokens int) []int {
	output := make([]int, 0, maxTokens)
	hidden := make([]float32, e.config.HiddenDim)

	// State cache for convolution layers
	convCache := make([][]float32, e.config.MaxSeqLen)
	for i := range convCache {
		convCache[i] = make([]float32, e.config.HiddenDim)
	}

	gqaIdx := 0

	for step := 0; step < maxTokens; step++ {
		// Embedding lookup (simplified)
		for i := range hidden {
			hidden[i] = float32(step%100) * 0.01
		}

		// Layer-by-layer inference
		for layer := 0; layer < e.config.NumLayers; layer++ {
			if e.config.LayerTypes[layer] == "conv" {
				hidden = e.conv.Forward(hidden, convCache, step)
				copy(convCache[step], hidden)
			} else {
				hidden = e.gqa.Forward(hidden, step, e.kvCache, gqaIdx, e.rope)
				gqaIdx++
			}
			// FFN after each layer
			hidden = e.ffn.Forward(hidden)
		}

		// Sample next token (simplified)
		nextToken := step % 1000
		output = append(output, nextToken)
	}

	return output
}

// ============================================================
// Quantization Utilities
// ============================================================

// Q4Weight represents a 4-bit quantized weight group
type Q4Weight struct {
	Scale float32
	Min   float32
	Data  []int8 // Only lower 4 bits used
}

// QuantizeQ4KM performs Q4_K_M quantization (group-wise)
func QuantizeQ4KM(weights []float32, groupSize int) []Q4Weight {
	numGroups := (len(weights) + groupSize - 1) / groupSize
	result := make([]Q4Weight, numGroups)

	for g := 0; g < numGroups; g++ {
		start := g * groupSize
		end := start + groupSize
		if end > len(weights) {
			end = len(weights)
		}
		group := weights[start:end]

		// Find min/max
		minVal := float32(math.Inf(1))
		maxVal := float32(math.Inf(-1))
		for _, w := range group {
			if w < minVal {
				minVal = w
			}
			if w > maxVal {
				maxVal = w
			}
		}

		scale := (maxVal - minVal) / 15.0
		if scale < 1e-10 {
			scale = 1.0
		}

		data := make([]int8, len(group))
		for i, w := range group {
			q := int8((w - minVal) / scale)
			if q < 0 {
				q = 0
			}
			if q > 15 {
				q = 15
			}
			data[i] = q
		}

		result[g] = Q4Weight{Scale: scale, Min: minVal, Data: data}
	}

	return result
}

// DequantizeQ4KM reconstructs weights from Q4_K_M format
func DequantizeQ4KM(qweights []Q4Weight) []float32 {
	totalLen := 0
	for _, qw := range qweights {
		totalLen += len(qw.Data)
	}

	result := make([]float32, totalLen)
	idx := 0
	for _, qw := range qweights {
		for _, qi := range qw.Data {
			result[idx] = qw.Min + float32(qi)*qw.Scale
			idx++
		}
	}
	return result
}

// ============================================================
// Benchmark Runner
// ============================================================

func runBenchmark() {
	config := LFM25Config()
	engine := NewInferenceEngine(config)

	// Count layer types
	gqaCount, convCount := 0, 0
	for _, t := range config.LayerTypes {
		if t == "gqa" {
			gqaCount++
		} else {
			convCount++
		}
	}

	fmt.Println("LFM2.5-2.6B Edge Inference Engine Benchmark")
	fmt.Println("=" * 60)
	fmt.Printf("Model Configuration:\n")
	fmt.Printf("  Layers: %d (Conv=%d, GQA=%d)\n", config.NumLayers, convCount, gqaCount)
	fmt.Printf("  Hidden: %d | Heads: %dQ/%dKV\n", config.HiddenDim, config.NumQHeads, config.NumKVHeads)
	fmt.Printf("  Vocab: %d | Context: %d\n", config.VocabSize, config.MaxSeqLen)

	// Speed test
	fmt.Println("\nInference Speed Test:")
	prompt := make([]int, 128)
	for i := range prompt {
		prompt[i] = i % 1000
	}

	numRuns := 5
	var totalTime float64
	var totalTokens int

	for run := 0; run < numRuns; run++ {
		start := time.Now()
		output := engine.Generate(prompt, 256)
		elapsed := time.Since(start).Seconds()
		totalTime += elapsed
		totalTokens += len(output)
	}

	avgSpeed := float64(totalTokens) / totalTime
	fmt.Printf("  Average speed: %.1f tok/s\n", avgSpeed)

	// Memory analysis
	fmt.Println("\nMemory Analysis:")
	kvMem := engine.kvCache.MemoryUsageMB()
	fmt.Printf("  KV cache: %.1f MB\n", kvMem)
	fmt.Printf("  Model weights (BF16): ~2500 MB\n")
	fmt.Printf("  Total (BF16): ~%.1f MB\n", 2500.0+kvMem)

	// Quantization test
	fmt.Println("\nQuantization Test:")
	numWeights := 100000
	weights := make([]float32, numWeights)
	for i := range weights {
		weights[i] = float32(i) * 0.0001
	}

	origSize := numWeights * 4
	qweights := QuantizeQ4KM(weights, 32)
	quantSize := 0
	for _, qw := range qweights {
		quantSize += 4 + 4 + len(qw.Data) // scale(float32) + min(float32) + data
	}

	fmt.Printf("  Original: %.2f MB\n", float64(origSize)/(1024*1024))
	fmt.Printf("  Q4_K_M:   %.2f MB\n", float64(quantSize)/(1024*1024))
	fmt.Printf("  Ratio:    %.1f:1\n", float64(origSize)/float64(quantSize))

	// Deploy summary
	fmt.Println("\nDeployment Summary:")
	fmt.Printf("  Q4 quantized model: ~%.0f MB\n", 2500.0/4.0)
	fmt.Printf("  KV cache (128K):    ~%.0f MB\n", kvMem)
	fmt.Printf("  Total (Q4):         ~%.0f MB\n", 2500.0/4.0+kvMem)
	fmt.Printf("  Memory budget:      2500 MB ✅\n")
}

func main() {
	runBenchmark()
}

"""
Quantization-Aware Deployment Toolchain
Supports: dynamic quantization, KV cache optimization, memory profiling
"""
import numpy as np
from typing import Tuple, Dict, List, Optional
from dataclasses import dataclass, field
from enum import Enum
import time
import math

class QuantMethod(Enum):
    Q4_0 = "q4_0"
    Q4_K_M = "q4_k_m"
    Q8_0 = "q8_0"
    BF16 = "bf16"

@dataclass
class DeployConfig:
    """Deployment configuration"""
    method: QuantMethod = QuantMethod.Q4_K_M
    group_size: int = 32
    kv_cache_bits: int = 8
    memory_budget_mb: float = 2500.0
    use_sym: bool = False

class BlockQuantizer:
    """Block-wise quantizer for efficient edge deployment"""
    
    @staticmethod
    def quantize_q4_km(weights: np.ndarray, group_size: int = 32) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Q4_K_M quantization: per-group scale and min
        
        Returns:
            quantized: packed 4-bit data
            scales: per-group scale factors
            mins: per-group minimum values
        """
        n = len(weights)
        num_groups = (n + group_size - 1) // group_size
        
        quantized = np.zeros(num_groups * group_size // 2, dtype=np.uint8)
        scales = np.zeros(num_groups, dtype=np.float16)
        mins = np.zeros(num_groups, dtype=np.float16)
        
        for g in range(num_groups):
            start = g * group_size
            end = min(start + group_size, n)
            group = weights[start:end]
            
            w_min = group.min()
            w_max = group.max()
            scale = (w_max - w_min) / 15.0 if w_max != w_min else 1.0
            
            scales[g] = scale
            mins[g] = w_min
            
            for i in range(0, len(group), 2):
                v0 = int((group[i] - w_min) / scale)
                v1 = int((group[i + 1] - w_min) / scale) if i + 1 < len(group) else 0
                v0 = max(0, min(15, v0))
                v1 = max(0, min(15, v1))
                quantized[g * group_size // 2 + i // 2] = (v0 & 0x0F) | ((v1 & 0x0F) << 4)
        
        return quantized, scales, mins
    
    @staticmethod
    def dequantize_q4_km(quantized: np.ndarray, scales: np.ndarray, 
                         mins: np.ndarray, group_size: int = 32,
                         total_len: Optional[int] = None) -> np.ndarray:
        """Dequantize Q4_K_M weights"""
        num_groups = len(scales)
        if total_len is None:
            total_len = num_groups * group_size
        
        result = np.zeros(total_len, dtype=np.float32)
        
        for g in range(num_groups):
            start = g * group_size
            end = min(start + group_size, total_len)
            scale = scales[g]
            w_min = mins[g]
            
            for i in range(start, end):
                byte_idx = g * group_size // 2 + i // 2
                if i % 2 == 0:
                    qi = quantized[byte_idx] & 0x0F
                else:
                    qi = (quantized[byte_idx] >> 4) & 0x0F
                result[i] = w_min + float(qi) * scale
        
        return result

class KVCacheOptimizer:
    """KV cache quantization and optimization"""
    
    def __init__(self, bits: int = 8):
        self.bits = bits
        self.max_val = 2 ** (bits - 1) - 1
    
    def quantize(self, tensor: np.ndarray) -> Tuple[np.ndarray, float]:
        """Per-tensor quantization"""
        abs_max = np.max(np.abs(tensor))
        if abs_max < 1e-10:
            abs_max = 1.0
        scale = self.max_val / abs_max
        quantized = np.clip(np.round(tensor * scale), 
                           -self.max_val - 1, self.max_val).astype(np.int8)
        return quantized, scale
    
    def dequantize(self, quantized: np.ndarray, scale: float) -> np.ndarray:
        """Dequantize KV cache tensor"""
        return quantized.astype(np.float32) / scale
    
    @staticmethod
    def estimate_size(num_layers: int, num_kv_heads: int, head_dim: int,
                      seq_len: int, bits: int) -> float:
        """Estimate KV cache size in MB"""
        bytes_per_elem = bits / 8
        total_bytes = (2 * num_layers * num_kv_heads * head_dim * 
                      seq_len * bytes_per_elem)
        return total_bytes / (1024 * 1024)

class MemoryProfiler:
    """Memory usage profiler for edge deployment"""
    
    def __init__(self, config: DeployConfig):
        self.config = config
    
    def profile(self, num_params: float, num_gqa: int, num_kv_heads: int,
                head_dim: int, seq_len: int) -> Dict[str, float]:
        """Profile memory usage for given configuration"""
        bits = {"bf16": 16, "q8_0": 8, "q4_k_m": 4, "q4_0": 4}[self.config.method.value]
        
        model_mem = num_params * (bits / 8) / (1024 ** 3)
        kv_mem = KVCacheOptimizer.estimate_size(
            num_gqa, num_kv_heads, head_dim, seq_len, self.config.kv_cache_bits
        ) / 1024  # GB
        
        total = model_mem + kv_mem
        within_budget = total * 1024 <= self.config.memory_budget_mb
        
        return {
            "model_gb": model_mem,
            "kv_cache_gb": kv_mem,
            "total_gb": total,
            "total_mb": total * 1024,
            "within_budget": within_budget,
            "budget_mb": self.config.memory_budget_mb
        }


def analyze_deployment():
    """Full deployment analysis for LFM2.5-2.6B"""
    config = DeployConfig()
    profiler = MemoryProfiler(config)
    
    # LFM2.5-2.6B architecture
    arch = {
        "num_params": 2.69e9,
        "num_layers": 30,
        "num_gqa": 8,
        "hidden_dim": 2048,
        "num_kv_heads": 8,
        "head_dim": 64,
        "max_seq_len": 131072,
    }
    
    print("=" * 70)
    print("LFM2.5-2.6B Edge Deployment Analysis")
    print("=" * 70)
    
    print(f"\n📐 Architecture:")
    print(f"  Parameters: {arch['num_params']/1e9:.2f}B")
    print(f"  Layers: {arch['num_layers']} (22 Conv + 8 GQA)")
    print(f"  Hidden: {arch['hidden_dim']}")
    print(f"  KV heads: {arch['num_kv_heads']}")
    print(f"  Max context: {arch['max_seq_len']:,}")
    
    print(f"\n📊 Quantization Comparison:")
    for name, bits in [("BF16", 16), ("Q8_0", 8), ("Q4_K_M", 4), ("Q4_0", 4)]:
        config.method = QuantMethod(name.lower())
        p = profiler.profile(
            arch["num_params"], arch["num_gqa"], 
            arch["num_kv_heads"], arch["head_dim"],
            arch["max_seq_len"]
        )
        flag = "✅" if p["within_budget"] else "❌"
        print(f"  {name:>8}: model={p['model_gb']*1024:>7.0f}MB, "
              f"KV={p['kv_cache_gb']*1024:>7.0f}MB, "
              f"total={p['total_mb']:>7.0f}MB {flag}")
    
    print(f"\n📏 Context Length Scaling (Q4_K_M):")
    for seq_len in [4096, 8192, 16384, 32768, 65536, 131072]:
        config.method = QuantMethod.Q4_K_M
        p = profiler.profile(
            arch["num_params"], arch["num_gqa"],
            arch["num_kv_heads"], arch["head_dim"],
            seq_len
        )
        flag = "✅" if p["within_budget"] else "❌"
        print(f"  {seq_len:>6,}: model={p['model_gb']*1024:>6.0f}MB + "
              f"KV={p['kv_cache_gb']*1024:>6.0f}MB = {p['total_mb']:>6.0f}MB {flag}")
    
    # Quantization accuracy simulation
    print(f"\n🔧 Quantization Accuracy:")
    np.random.seed(42)
    test_weights = np.random.randn(10000).astype(np.float32) * 0.1
    
    q_data, scales, mins = BlockQuantizer.quantize_q4_km(test_weights, 32)
    deq = BlockQuantizer.dequantize_q4_km(q_data, scales, mins, 32, len(test_weights))
    
    mse = np.mean((test_weights - deq) ** 2)
    peak_snr = 20 * math.log10(np.max(np.abs(test_weights)) / math.sqrt(mse))
    
    print(f"  Q4_K_M MSE: {mse:.8f}")
    print(f"  Peak SNR: {peak_snr:.1f} dB")
    print(f"  Compression: {len(test_weights)*4/(len(q_data)+len(scales)*2+len(mins)*2):.1f}:1")
    
    print(f"\n💡 Recommendation:")
    print(f"  Smartphone: Q4_K_M + 8-bit KV cache → 32K-64K context @ ~30 tok/s")
    print(f"  Laptop:     Q4_K_M/Q8_0 + 8-bit KV → 128K context @ 113-220 tok/s")
    print(f"  Server:     BF16/FP8 → 128K @ ~15K tok/s on H100")

if __name__ == "__main__":
    analyze_deployment()

// benchmark_analysis.go - Cross-model benchmark comparison
package main

import (
	"fmt"
	"math"
	"sort"
)

// BenchmarkData holds all benchmark results
type BenchmarkData struct {
	Name     string
	Category string
	Models   map[string]float64
}

// LoadBenchmarks returns the complete benchmark dataset
func LoadBenchmarks() []BenchmarkData {
	return []BenchmarkData{
		{"AA Omniscience", "stem", map[string]float64{
			"LFM2.5-2.6B": -29.50, "Gemma4-E2B": -74.47,
			"Gemma4-E4B": -49.03, "Qwen3.5-4B": -54.30, "Qwen3.5-9B": -50.43,
		}},
		{"AIME25", "stem", map[string]float64{
			"LFM2.5-2.6B": 51.87, "Gemma4-E2B": 26.33,
			"Gemma4-E4B": 34.27, "Qwen3.5-4B": 49.33, "Qwen3.5-9B": 56.07,
		}},
		{"LiveCodeBenchv6", "coding", map[string]float64{
			"LFM2.5-2.6B": 59.41, "Gemma4-E2B": 54.92,
			"Gemma4-E4B": 63.77, "Qwen3.5-4B": 60.85, "Qwen3.5-9B": 69.86,
		}},
		{"IFBench", "instruction", map[string]float64{
			"LFM2.5-2.6B": 59.17, "Gemma4-E2B": 34.08,
			"Gemma4-E4B": 39.24, "Qwen3.5-4B": 48.40, "Qwen3.5-9B": 56.47,
		}},
		{"Multi-IF", "instruction", map[string]float64{
			"LFM2.5-2.6B": 80.07, "Gemma4-E2B": 69.44,
			"Gemma4-E4B": 77.35, "Qwen3.5-4B": 55.67, "Qwen3.5-9B": 62.55,
		}},
		{"IFStruct", "instruction", map[string]float64{
			"LFM2.5-2.6B": 85.49, "Gemma4-E2B": 64.85,
			"Gemma4-E4B": 76.65, "Qwen3.5-4B": 36.25, "Qwen3.5-9B": 78.50,
		}},
		{"BFCLv4", "tool", map[string]float64{
			"LFM2.5-2.6B": 56.88, "Gemma4-E2B": 36.98,
			"Gemma4-E4B": 46.39, "Qwen3.5-4B": 50.56, "Qwen3.5-9B": 60.13,
		}},
		{"ToolSandbox", "tool", map[string]float64{
			"LFM2.5-2.6B": 77.83, "Gemma4-E2B": 52.40,
			"Gemma4-E4B": 65.00, "Qwen3.5-4B": 75.55, "Qwen3.5-9B": 76.44,
		}},
		{"τ³-Bench Banking", "agent", map[string]float64{
			"LFM2.5-2.6B": 5.67, "Gemma4-E2B": 3.35,
			"Gemma4-E4B": 4.12, "Qwen3.5-4B": 5.45, "Qwen3.5-9B": 5.15,
		}},
		{"Claw-Eval avg", "agent", map[string]float64{
			"LFM2.5-2.6B": 62.85, "Gemma4-E2B": 53.14,
			"Gemma4-E4B": 58.02, "Qwen3.5-4B": 62.28, "Qwen3.5-9B": 66.53,
		}},
		{"PinchBench", "agent", map[string]float64{
			"LFM2.5-2.6B": 68.22, "Gemma4-E2B": 44.24,
			"Gemma4-E4B": 55.09, "Qwen3.5-4B": 71.26, "Qwen3.5-9B": 71.45,
		}},
		{"BrowseComp+", "agent", map[string]float64{
			"LFM2.5-2.6B": 26.89, "Gemma4-E2B": 8.31,
			"Gemma4-E4B": 15.90, "Qwen3.5-4B": 24.46, "Qwen3.5-9B": 27.23,
		}},
	}
}

// ModelInfo describes a model's architecture
type ModelInfo struct {
	Name       string
	ParamsB    float64
	Layers     int
	HiddenDim  int
	ContextLen int
}

func getModels() []ModelInfo {
	return []ModelInfo{
		{"LFM2.5-2.6B", 2.69, 30, 2048, 131072},
		{"Gemma4-E2B", 5.1, 26, 2048, 128000},
		{"Gemma4-E4B", 8.0, 32, 2560, 128000},
		{"Qwen3.5-4B", 4.7, 32, 2560, 262144},
		{"Qwen3.5-9B", 9.7, 36, 3584, 262144},
	}
}

// EfficiencyRank stores per-model efficiency metrics
type EfficiencyRank struct {
	ModelName     string
	ParamsB       float64
	AvgEfficiency float64
	NormScore     float64
	WinCount      int
}

func computeEfficiency(data []BenchmarkData) []EfficiencyRank {
	models := []string{"LFM2.5-2.6B", "Gemma4-E2B", "Gemma4-E4B", "Qwen3.5-4B", "Qwen3.5-9B"}
	params := map[string]float64{
		"LFM2.5-2.6B": 2.69, "Gemma4-E2B": 5.1,
		"Gemma4-E4B": 8.0, "Qwen3.5-4B": 4.7, "Qwen3.5-9B": 9.7,
	}

	rankMap := make(map[string]*EfficiencyRank)
	for _, m := range models {
		rankMap[m] = &EfficiencyRank{ModelName: m, ParamsB: params[m]}
	}

	for _, b := range data {
		// Find best efficiency in this benchmark
		bestEff := math.Inf(-1)
		effMap := make(map[string]float64)

		for model, score := range b.Models {
			eff := score / params[model]
			effMap[model] = eff
			if eff > bestEff {
				bestEff = eff
			}
		}

		for model, eff := range effMap {
			r := rankMap[model]
			r.AvgEfficiency += eff
			r.NormScore += (eff / bestEff) * 100
			if eff == bestEff {
				r.WinCount++
			}
		}
	}

	numBench := len(data)
	var result []EfficiencyRank
	for _, r := range rankMap {
		r.AvgEfficiency /= float64(numBench)
		r.NormScore /= float64(numBench)
		result = append(result, *r)
	}

	sort.Slice(result, func(i, j int) bool {
		return result[i].NormScore > result[j].NormScore
	})

	return result
}

func main() {
	data := LoadBenchmarks()
	models := getModels()

	fmt.Println("LFM2.5-2.6B Benchmark Deep Analysis")
	fmt.Println("=" * 70)

	// Model comparison
	fmt.Println("\n📋 Model Comparison:")
	fmt.Printf("%-16s %8s %8s %10s %10s\n", "Model", "Params(B)", "Layers", "Hidden", "Context")
	fmt.Println("-" * 55)
	for _, m := range models {
		fmt.Printf("%-16s %8.1f %8d %10d %10d\n",
			m.Name, m.ParamsB, m.Layers, m.HiddenDim, m.ContextLen)
	}

	// Efficiency ranking
	fmt.Println("\n⚡ Parameter Efficiency Ranking:")
	rankings := computeEfficiency(data)
	fmt.Printf("%-16s %10s %15s %15s %10s\n", "Model", "Params(B)", "Avg Eff", "Norm Score", "Wins")
	fmt.Println("-" * 70)
	for _, r := range rankings {
		fmt.Printf("%-16s %10.1f %15.2f %14.1f%% %10d\n",
			r.ModelName, r.ParamsB, r.AvgEfficiency, r.NormScore, r.WinCount)
	}

	// Efficiency ratios
	fmt.Println("\n📊 Efficiency Ratio (LFM2.5-2.6B vs Competitors):")
	fmt.Printf("%-20s %15s %15s\n", "Benchmark", "vs Gemma4-E4B", "vs Qwen3.5-9B")
	fmt.Println("-" * 55)

	lfmP := 2.69
	gemmaP := 8.0
	qwenP := 9.7

	keyBenchmarks := []string{"IFBench", "IFStruct", "BFCLv4", "ToolSandbox", "AIME25", "Multi-IF"}

	// Build lookup
	benchMap := make(map[string]map[string]float64)
	for _, b := range data {
		benchMap[b.Name] = b.Models
	}

	for _, name := range keyBenchmarks {
		if scores, ok := benchMap[name]; ok {
			lfmEff := scores["LFM2.5-2.6B"] / lfmP
			gemmaEff := scores["Gemma4-E4B"] / gemmaP
			qwenEff := scores["Qwen3.5-9B"] / qwenP

			fmt.Printf("%-20s %13.1fx %13.1fx\n", name, lfmEff/gemmaEff, lfmEff/qwenEff)
		}
	}

	// Summary
	fmt.Println("\n" + "=" * 70)
	fmt.Println("Core Conclusions")
	fmt.Println("=" * 70)
	fmt.Println(`
1. LFM2.5-2.6B leads ALL instruction-following benchmarks (IFBench 59.17, IFStruct 85.49)
2. Tool use matches Qwen3.5-9B, ToolSandbox even surpasses it (77.83 vs 76.44)
3. Agentic tasks beat Gemma 4 entire series, on par with Qwen3.5 series
4. Parameter efficiency: 3-4x Gemma 4-E4B, 5-6x Qwen3.5-9B
5. Code generation is the only weakness → use larger models for coding
6. <2.5GB memory + 220 tok/s → unmatched for edge deployment`)
}

"""
Complete on-device agent inference loop
Features: tool calling, multi-step reasoning, context management
"""
import json
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
import re
import time

@dataclass
class Tool:
    """Tool definition"""
    name: str
    description: str
    parameters: Dict[str, Any]
    function: Callable
    
    def schema(self) -> Dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": self.parameters,
                    "required": list(self.parameters.keys())
                }
            }
        }

@dataclass
class Message:
    """Conversation message"""
    role: str
    content: str
    tool_calls: Optional[List[Dict]] = None
    tool_call_id: Optional[str] = None

@dataclass
class AgentState:
    """Agent execution state"""
    messages: List[Message] = field(default_factory=list)
    max_turns: int = 10
    current_turn: int = 0

class ToolCallParser:
    """Parse tool calls from model output"""
    
    @staticmethod
    def parse(text: str) -> List[Dict]:
        """Extract tool calls from text"""
        calls = []
        
        # JSON format
        for match in re.findall(r'```json\s*(\[.*?\])\s*```', text, re.DOTALL):
            try:
                items = json.loads(match)
                if isinstance(items, list):
                    calls.extend(items)
                else:
                    calls.append(items)
            except json.JSONDecodeError:
                pass
        
        # Function call format
        for match in re.findall(r'<function=(\w+)>(.*?)</function>', text, re.DOTALL):
            name, args_str = match
            try:
                args = json.loads(args_str)
                calls.append({"name": name, "arguments": args})
            except json.JSONDecodeError:
                pass
        
        return calls

class OnDeviceAgent:
    """
    On-device agent inference engine
    Runs LFM2.5-2.6B agent loop locally
    """
    
    def __init__(self, tools: List[Tool]):
        self.tools = {t.name: t for t in tools}
        self.parser = ToolCallParser()
        self.state = AgentState()
        self.stats = {"turns": 0, "tool_calls": 0, "total_time": 0.0}
    
    def build_system_prompt(self) -> str:
        """Build system prompt with tool definitions"""
        tools_desc = "\n\n".join([
            f"## {t.name}\n{t.description}\n"
            f"Parameters: {json.dumps(t.parameters, indent=2)}"
            for t in self.tools.values()
        ])
        
        return f"""You are an on-device AI Agent running on LFM2.5-2.6B.
You can use these tools:

{tools_desc}

Tool call format:
```json
[{{"name": "tool_name", "arguments": {{"param": "value"}}}}]