core/go-mlx
8
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions

feat(metal): add inference metrics (timing, throughput, memory)

Browse source 
Instrument Generate, Classify, and BatchGenerate with:
- Prefill/decode timing (separate phases)
- Token counts (prompt + generated)
- Throughput (tok/s for each phase)
- Peak and active GPU memory via Metal allocator

Wire through metalAdapter.Metrics() to go-inference interface.
Test validates all fields populated after generation.

Gemma3-1B 4-bit on M3 Ultra: prefill 246 tok/s, decode 82 tok/s,
peak 6.2 GB GPU memory.

Co-Authored-By: Virgil <virgil@lethean.io>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Snider 2026-02-19 23:34:40 +00:00
parent 197326bd22
commit a44e9f5789
4 changed files with 162 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
internal/metal/batch.go
View file

@ -7,6 +7,7 @@ import (
"fmt"
"math"
"sort"
"time"
)
// ClassifyResult holds the output for a single prompt in batch classification.
@ -24,16 +25,22 @@ type BatchResult struct {
// Classify runs batched prefill-only inference. Each prompt gets a single
// forward pass and the token at the last position is sampled.
func (m *Model) Classify(ctx context.Context, prompts []string, cfg GenerateConfig, returnLogits bool) ([]ClassifyResult, error) {
m.lastMetrics = Metrics{}
if len(prompts) == 0 {
return nil, nil
}
totalStart := time.Now()
ResetPeakMemory()
// Tokenise all prompts.
encoded := make([][]int32, len(prompts))
lengths := make([]int, len(prompts))
totalPromptTokens := 0
for i, p := range prompts {
encoded[i] = m.tokenizer.Encode(p)
lengths[i] = len(encoded[i])
totalPromptTokens += lengths[i]
}
// Sort by length descending for minimal padding. Track original indices.
@ -111,21 +118,40 @@ func (m *Model) Classify(ctx context.Context, prompts []string, cfg GenerateConf
results[origIdx] = sortedResults[si]
}
totalDur := time.Since(totalStart)
m.lastMetrics = Metrics{
PromptTokens: totalPromptTokens,
GeneratedTokens: int(N), // One token sampled per prompt
PrefillDuration: totalDur,
TotalDuration: totalDur,
PeakMemoryBytes: GetPeakMemory(),
ActiveMemoryBytes: GetActiveMemory(),
}
if totalDur > 0 {
m.lastMetrics.PrefillTokensPerSec = float64(totalPromptTokens) / totalDur.Seconds()
}
return results, nil
}
// BatchGenerate runs batched autoregressive generation.
func (m *Model) BatchGenerate(ctx context.Context, prompts []string, cfg GenerateConfig) ([]BatchResult, error) {
m.lastMetrics = Metrics{}
if len(prompts) == 0 {
return nil, nil
}
totalStart := time.Now()
ResetPeakMemory()
// Tokenise all prompts.
encoded := make([][]int32, len(prompts))
lengths := make([]int, len(prompts))
totalPromptTokens := 0
for i, p := range prompts {
encoded[i] = m.tokenizer.Encode(p)
lengths[i] = len(encoded[i])
totalPromptTokens += lengths[i]
}
// Sort by length descending.
@ -156,6 +182,7 @@ func (m *Model) BatchGenerate(ctx context.Context, prompts []string, cfg Generat
}
// Prefill with mask.
prefillStart := time.Now()
mask := buildBatchMask(N, L, sortedLengths)
tokens := FromValues(padded, int(N), int(L))
caches := m.newCachesN(int(N))
@ -163,6 +190,7 @@ func (m *Model) BatchGenerate(ctx context.Context, prompts []string, cfg Generat
if err := Eval(logits); err != nil {
return nil, fmt.Errorf("batch prefill: %w", err)
}
prefillDur := time.Since(prefillStart)
sampler := newSampler(cfg.Temperature, cfg.TopP, 0, cfg.TopK)
eosID := m.tokenizer.EOSToken()
@ -258,14 +286,34 @@ func (m *Model) BatchGenerate(ctx context.Context, prompts []string, cfg Generat
// Unsort results back to original order.
sortedResults := make([]BatchResult, N)
totalGenerated := 0
for si := range states {
sortedResults[si] = BatchResult{Tokens: states[si].tokens}
totalGenerated += len(states[si].tokens)
}
results := make([]BatchResult, N)
for si, origIdx := range indices {
results[origIdx] = sortedResults[si]
}
totalDur := time.Since(totalStart)
decodeDur := totalDur - prefillDur
m.lastMetrics = Metrics{
PromptTokens: totalPromptTokens,
GeneratedTokens: totalGenerated,
PrefillDuration: prefillDur,
DecodeDuration: decodeDur,
TotalDuration: totalDur,
PeakMemoryBytes: GetPeakMemory(),
ActiveMemoryBytes: GetActiveMemory(),
}
if prefillDur > 0 {
m.lastMetrics.PrefillTokensPerSec = float64(totalPromptTokens) / prefillDur.Seconds()
}
if decodeDur > 0 {
m.lastMetrics.DecodeTokensPerSec = float64(totalGenerated) / decodeDur.Seconds()
}
return results, nil
}

48
internal/metal/generate.go
View file

@ -6,6 +6,7 @@ import (
"context"
"fmt"
"iter"
"time"
)
// Token represents a single generated token.
@ -30,6 +31,19 @@ type GenerateConfig struct {
RepeatPenalty float32
}
// Metrics holds performance metrics from the last inference operation.
type Metrics struct {
PromptTokens int
GeneratedTokens int
PrefillDuration time.Duration
DecodeDuration time.Duration
TotalDuration time.Duration
PrefillTokensPerSec float64
DecodeTokensPerSec float64
PeakMemoryBytes uint64
ActiveMemoryBytes uint64
}
// Model wraps a loaded transformer model for text generation.
type Model struct {
model InternalModel
@ -37,6 +51,7 @@ type Model struct {
modelType string
contextLen int // 0 = unbounded (model default)
lastErr error
lastMetrics Metrics
}
// ModelType returns the architecture identifier (e.g. "gemma3", "qwen3").
@ -45,6 +60,9 @@ func (m *Model) ModelType() string { return m.modelType }
// Err returns the error from the last Generate/Chat call, if any.
func (m *Model) Err() error { return m.lastErr }
// LastMetrics returns performance metrics from the last inference operation.
func (m *Model) LastMetrics() Metrics { return m.lastMetrics }
// Close releases all model weight arrays and associated resources.
// After Close, the Model must not be used for generation.
func (m *Model) Close() error {
@ -75,13 +93,41 @@ func (m *Model) Chat(ctx context.Context, messages []ChatMessage, cfg GenerateCo
// Metal memory promptly rather than waiting for GC finalisers.
func (m *Model) Generate(ctx context.Context, prompt string, cfg GenerateConfig) iter.Seq[Token] {
m.lastErr = nil
m.lastMetrics = Metrics{}
return func(yield func(Token) bool) {
totalStart := time.Now()
ResetPeakMemory()
tokens := m.tokenizer.Encode(prompt)
promptLen := len(tokens)
caches := m.newCaches()
sampler := newSampler(cfg.Temperature, cfg.TopP, 0, cfg.TopK)
var genCount int
var prefillDur time.Duration
defer func() {
decodeDur := time.Since(totalStart) - prefillDur
totalDur := time.Since(totalStart)
m.lastMetrics = Metrics{
PromptTokens: promptLen,
GeneratedTokens: genCount,
PrefillDuration: prefillDur,
DecodeDuration: decodeDur,
TotalDuration: totalDur,
PeakMemoryBytes: GetPeakMemory(),
ActiveMemoryBytes: GetActiveMemory(),
}
if prefillDur > 0 {
m.lastMetrics.PrefillTokensPerSec = float64(promptLen) / prefillDur.Seconds()
}
if decodeDur > 0 {
m.lastMetrics.DecodeTokensPerSec = float64(genCount) / decodeDur.Seconds()
}
}()
// Prefill: process entire prompt
prefillStart := time.Now()
input := FromValues(tokens, len(tokens))
input = Reshape(input, 1, int32(len(tokens)))
logits := m.model.Forward(input, caches)
@ -89,6 +135,7 @@ func (m *Model) Generate(ctx context.Context, prompt string, cfg GenerateConfig)
m.lastErr = fmt.Errorf("prefill: %w", err)
return
}
prefillDur = time.Since(prefillStart)
// Track generated token IDs for repeat penalty.
var history []int32
@ -129,6 +176,7 @@ func (m *Model) Generate(ctx context.Context, prompt string, cfg GenerateConfig)
}
}
genCount++
text := m.tokenizer.DecodeToken(id)
if !yield(Token{ID: id, Text: text}) {
return

51
mlx_test.go
View file

@ -443,6 +443,57 @@ func TestLlama_Inference(t *testing.T) {
}
}
// --- Metrics tests ---
// TestGenerate_Metrics validates that metrics are populated after generation.
func TestGenerate_Metrics(t *testing.T) {
modelPath := gemma3ModelPath(t)
m, err := inference.LoadModel(modelPath)
if err != nil {
t.Fatalf("LoadModel: %v", err)
}
defer func() { m.Close(); mlx.ClearCache() }()
ctx := context.Background()
var count int
for range m.Generate(ctx, "Hello world", inference.WithMaxTokens(8)) {
count++
}
if err := m.Err(); err != nil {
t.Fatalf("Generate error: %v", err)
}
met := m.Metrics()
if met.PromptTokens == 0 {
t.Error("PromptTokens should be > 0")
}
if met.GeneratedTokens != count {
t.Errorf("GeneratedTokens = %d, want %d", met.GeneratedTokens, count)
}
if met.PrefillDuration == 0 {
t.Error("PrefillDuration should be > 0")
}
if met.TotalDuration == 0 {
t.Error("TotalDuration should be > 0")
}
if met.PrefillTokensPerSec == 0 {
t.Error("PrefillTokensPerSec should be > 0")
}
if met.DecodeTokensPerSec == 0 {
t.Error("DecodeTokensPerSec should be > 0")
}
if met.PeakMemoryBytes == 0 {
t.Error("PeakMemoryBytes should be > 0")
}
t.Logf("Metrics: prompt=%d tokens, generated=%d tokens", met.PromptTokens, met.GeneratedTokens)
t.Logf(" Prefill: %s (%.0f tok/s)", met.PrefillDuration, met.PrefillTokensPerSec)
t.Logf(" Decode: %s (%.1f tok/s)", met.DecodeDuration, met.DecodeTokensPerSec)
t.Logf(" Total: %s", met.TotalDuration)
t.Logf(" Peak memory: %.1f MB", float64(met.PeakMemoryBytes)/(1024*1024))
}
// --- Batch Inference tests (Gemma3-1B) ---
// TestClassify_Batch validates batched prefill-only classification.

15
register_metal.go
View file

@ -164,6 +164,21 @@ func (a *metalAdapter) BatchGenerate(ctx context.Context, prompts []string, opts
return out, nil
}
func (a *metalAdapter) Metrics() inference.GenerateMetrics {
m := a.m.LastMetrics()
return inference.GenerateMetrics{
PromptTokens: m.PromptTokens,
GeneratedTokens: m.GeneratedTokens,
PrefillDuration: m.PrefillDuration,
DecodeDuration: m.DecodeDuration,
TotalDuration: m.TotalDuration,
PrefillTokensPerSec: m.PrefillTokensPerSec,
DecodeTokensPerSec: m.DecodeTokensPerSec,
PeakMemoryBytes: m.PeakMemoryBytes,
ActiveMemoryBytes: m.ActiveMemoryBytes,
}
}
func (a *metalAdapter) ModelType() string { return a.m.ModelType() }
func (a *metalAdapter) Err() error { return a.m.Err() }
func (a *metalAdapter) Close() error { return a.m.Close() }