Apple MLX: Dense prepare and Sequential

github.com/openfluke/welvet/nativemlxdarwin arm64

Why it exists

Show prepare-once Dense execute and a small Sequential tower on the same MLX session.

What it is

SyncToMLX a Dense with PrepareMLX, then SyncToMLX a two-Dense Sequential and forward once.

Prepare vs online

Dense can PrepareMLX(batch) outside the timed loop so repeated Forward calls stay execute-only. Sequential rebuilds a welvet.mlx Sequential from supported children (Dense / Softmax / RMSNorm / LayerNorm today) and host-ping-pongs child MLX evals.

These shapes are API demos: not performance claims. Live large-shape matrices live in w2a go_v_mlx (§113).

Requirements: chapter 110.

Go example

examples/112-apple-mlx-forward/main.go

Run:cd welvet/examples/112-apple-mlx-forward && source ../env.sh && go run .

// Dense PrepareMLX + Sequential SyncToMLX forward on Apple MLX.
package main

import (
	"fmt"
	"log"
	"runtime"

	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/dense"
	"github.com/openfluke/welvet/layers/sequential"
	"github.com/openfluke/welvet/nativemlx"
	"github.com/openfluke/welvet/quant"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	if runtime.GOOS != "darwin" || runtime.GOARCH != "arm64" {
		fmt.Println("SKIP: Apple MLX needs Apple silicon macOS and cgo")
		return nil
	}
	if !nativemlx.Available() {
		fmt.Println("SKIP: need libwelvet_mlx.dylib + Metal GPU")
		return nil
	}

	if err := demoDensePrepare(); err != nil {
		return err
	}
	return demoSequential()
}

func demoDensePrepare() error {
	init := make([]float32, 4*4)
	for i := range init {
		init[i] = float32(i+1) * 0.05
	}
	l, err := dense.NewConfigured(4, 4, core.ActivationLinear, core.DTypeFloat32, quant.FormatNone, init)
	if err != nil {
		return err
	}
	if err := l.SyncToMLX(); err != nil {
		return err
	}
	if err := l.PrepareMLX(1); err != nil {
		return err
	}
	l.Exec.Backend = core.BackendMLX
	x := core.NewTensor[float32](1, 4)
	for i := range x.Data {
		x.Data[i] = 0.1 * float32(i+1)
	}
	_, post, err := dense.Forward(l, x)
	if err != nil {
		return err
	}
	fmt.Printf("dense prepare+exec out=%v\n", post.Data)
	return nil
}

func demoSequential() error {
	// Two identity Dense children (Depth=2, Dim=3).
	packed := append(identity3(), identity3()...)
	seq, err := sequential.NewConfigured(sequential.Config{Dim: 3, Depth: 2}, core.DTypeFloat32, quant.FormatNone, packed)
	if err != nil {
		return err
	}
	if err := seq.SyncToMLX(); err != nil {
		return err
	}
	seq.Exec.Backend = core.BackendMLX
	x := &core.Tensor[float32]{Shape: []int{1, 3}, Data: []float32{1, 2, 3}}
	_, post, err := sequential.Forward(seq, x)
	if err != nil {
		return err
	}
	fmt.Printf("sequential mlx out=%v synced=%v\n", post.Data, seq.MLXSynced())
	return nil
}

func identity3() []float32 {
	return []float32{
		1, 0, 0,
		0, 1, 0,
		0, 0, 1,
	}
}

Output

exit 0 · last run via go run .

dense prepare+exec out=[0.15 0.35 0.5500001 0.75]
sequential mlx out=[1 2 3] synced=true

Validate (harness)

cd welvet/examples/112-apple-mlx-forward && source ../env.sh && go run .