Add ORTGenAI Engine support

Author: RJKeevil

backends/generative.go

@@ -1,5 +1,7 @@
 package backends
 
+import "context"
+
 // GuidanceType specifies the constrained-generation strategy.
 type GuidanceType string
 
@@ -19,8 +21,8 @@ type Guidance struct {
 }
 
 type SequenceDelta struct {
-	Token string
-	Index int
+	Token    string
+	Sequence int
 }
 
 // Message represents a single message in a conversation.
@@ -30,3 +32,19 @@ type Message struct {
 	Content   string   `json:"content"`
 	ImageURLs []string `json:"image_urls,omitempty"` // File paths or data URIs for multimodal support
 }
+
+// GenerativeModel abstracts either a generative session or engine.
+type GenerativeModel interface {
+	Generate(ctx context.Context, inputs [][]Message, tools []string, options *GenerativeOptions) (chan SequenceDelta, chan error, error)
+	GetStatistics() PipelineStatistics
+	Destroy() error
+}
+
+// GenerativeOptions contains settings for text generation.
+type GenerativeOptions struct {
+	MaxLength   int
+	Temperature *float64
+	TopP        *float64
+	Seed        *int
+	Guidance    *Guidance
+}

backends/model.go

@@ -42,7 +42,20 @@ func LoadModel(path string, onnxFilename string, options *options.Options, isGen
 		IsGenerative: isGenerative,
 	}
 
-	if !isGenerative {
+	if isGenerative {
+		// creation of the session. Only one output (either token or sentence embedding).
+		if options.Backend != "ORT" {
+			return nil, fmt.Errorf("generative models are only supported with ORT backend currently")
+		}
+		if onnxFilename != "" {
+			return nil, fmt.Errorf("onnx filename should not be provided for generative models as we currently rely on genai_config for the onnx backend")
+		}
+
+		err := createORTGenerativeSession(model, options)
+		if err != nil {
+			return nil, err
+		}
+	} else {
 		err := loadModelConfig(model)
 		if err != nil {
 			return nil, err
@@ -55,18 +68,6 @@ func LoadModel(path string, onnxFilename string, options *options.Options, isGen
 		if tkErr != nil {
 			return nil, tkErr
 		}
-	} else {
-		// creation of the session. Only one output (either token or sentence embedding).
-		if options.Backend != "ORT" {
-			return nil, fmt.Errorf("generative models are only supported with ORT backend currently")
-		}
-		if onnxFilename != "" {
-			return nil, fmt.Errorf("onnx filename should not be provided for generative models as we currently rely on genai_config for the onnx backend")
-		}
-		err := createORTGenerativeSession(model, options)
-		if err != nil {
-			return nil, err
-		}
 	}
 
 	model.Destroy = func() error {

backends/model_gomlx.go

@@ -45,7 +45,7 @@ func createGoMLXModelBackend(model *Model, options *options.Options) error {
 		modelParsed, err = parser.ParseReader(model.OnnxReader)
 	} else {
 		modelParsed, err = parser.ParseFile(fileutil.PathJoinSafe(model.Path, model.OnnxPath))
-		if err != nil {
+		if err != nil && modelParsed != nil {
 			modelParsed = modelParsed.WithBaseDir(model.Path)
 		}
 	}
@@ -97,7 +97,7 @@ func createGoMLXModelBackend(model *Model, options *options.Options) error {
 		return errors.Join(contextErr, modelParsed.Close())
 	}
 
-	maxCache, batchBuckets, sequenceBuckets := getCacheAndBucketSizes(options, model, config)
+	maxCache, batchBuckets, sequenceBuckets := getCacheAndBucketSizes(options, config)
 	exec.SetMaxCache(maxCache)
 
 	model.GoMLXModel = &GoMLXModel{
@@ -121,7 +121,7 @@ func createGoMLXModelBackend(model *Model, options *options.Options) error {
 	return err
 }
 
-func getCacheAndBucketSizes(options *options.Options, model *Model, backend string) (int, []int, []int) {
+func getCacheAndBucketSizes(options *options.Options, backend string) (int, []int, []int) {
 	bucketsSpecified := false
 	// Use configured buckets or fall back to defaults.
 	batchBuckets := defaultBatchBuckets

backends/model_ort.go

@@ -11,6 +11,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 
 	ort "github.com/yalue/onnxruntime_go"
 	"golang.org/x/sync/errgroup"
@@ -23,11 +24,14 @@ import (
 type ORTModel struct {
 	Session           *ort.DynamicAdvancedSession
 	GenerativeSession *ortgenai.Session
+	GenerativeEngine  *ortgenai.Engine
 	SessionOptions    *ort.SessionOptions
 	Options           *options.OrtOptions
 	Destroy           func() error
 }
 
+var generativeBackendMutex = sync.Mutex{}
+
 func mapORTOptions(options *options.Options) ([]string, map[string]map[string]string, error) {
 	if options == nil || options.ORTOptions == nil {
 		return []string{}, map[string]map[string]string{}, nil // let default EPs be used
@@ -87,11 +91,54 @@ func mapORTOptions(options *options.Options) ([]string, map[string]map[string]st
 }
 
 func createORTGenerativeSession(model *Model, options *options.Options) error {
-
 	if strings.HasPrefix(model.Path, "s3:") {
 		return errors.New("ORT Gen AI does not support S3 paths. Please download the model to a local directory and try again")
 	}
 
+	err := initialiseORTGenAI(options)
+	if err != nil {
+		return err
+	}
+
+	providers, providerOptions, err := mapORTOptions(options)
+	if err != nil {
+		return fmt.Errorf("error mapping ORT options for generative session: %w", err)
+	}
+
+	if options.ORTOptions.UseEngine {
+		ortGenAiEngine, err := ortgenai.CreateEngineWithOptions(model.Path, providers, providerOptions)
+		if err != nil {
+			return fmt.Errorf("error creating ortgenai engine: %w", err)
+		}
+		model.ORTModel = &ORTModel{
+			GenerativeEngine: ortGenAiEngine,
+			Options:          options.ORTOptions,
+			Destroy: func() error {
+				ortGenAiEngine.Destroy()
+				return nil
+			},
+		}
+	} else {
+		ortGenAiSession, err := ortgenai.CreateSessionWithOptions(model.Path, providers, providerOptions)
+		if err != nil {
+			return fmt.Errorf("error creating ortgenai session: %w", err)
+		}
+		model.ORTModel = &ORTModel{
+			GenerativeSession: ortGenAiSession,
+			Options:           options.ORTOptions,
+			Destroy: func() error {
+				ortGenAiSession.Destroy()
+				return nil
+			},
+		}
+	}
+	return nil
+}
+
+func initialiseORTGenAI(options *options.Options) error {
+	generativeBackendMutex.Lock()
+	defer generativeBackendMutex.Unlock()
+
 	if !ortgenai.IsInitialized() {
 		if options == nil || options.ORTOptions == nil {
 			return fmt.Errorf("ORT options must be provided to initialize ortgenai")
@@ -124,22 +171,6 @@ func createORTGenerativeSession(model *Model, options *options.Options) error {
 			return fmt.Errorf("error initializing the ort genai environment: %w", err)
 		}
 	}
-	providers, providerOptions, err := mapORTOptions(options)
-	if err != nil {
-		return fmt.Errorf("error mapping ORT options for generative session: %w", err)
-	}
-	ortGenAiSession, err := ortgenai.CreateSessionWithOptions(model.Path, providers, providerOptions)
-	if err != nil {
-		return fmt.Errorf("error creating ortgenai session: %w", err)
-	}
-	model.ORTModel = &ORTModel{
-		GenerativeSession: ortGenAiSession,
-		Options:           options.ORTOptions,
-		Destroy: func() error {
-			ortGenAiSession.Destroy()
-			return nil
-		},
-	}
 	return nil
 }
 
@@ -156,8 +187,9 @@ func runGenerativeORTSessionOnBatch(ctx context.Context, batch *PipelineBatch, p
 	}
 
 	session := p.Model.ORTModel.GenerativeSession
-	if session == nil {
-		return nil, nil, errors.New("ORT generative session is not initialized")
+	engine := p.Model.ORTModel.GenerativeEngine
+	if session == nil && engine == nil {
+		return nil, nil, errors.New("ORT generative session/engine is not initialized")
 	}
 
 	inputs, ok := batch.InputValues.([][]ortgenai.Message)
@@ -183,6 +215,10 @@ func runGenerativeORTSessionOnBatch(ctx context.Context, batch *PipelineBatch, p
 	generateCtx, cancel := context.WithCancel(ctx)
 
 	if batch.Images != nil {
+		if session == nil {
+			cancel()
+			return nil, nil, errors.New("multimodal generation requires a session, but only engine is initialized")
+		}
 		images, ok := batch.Images.(*ortgenai.Images)
 		if !ok {
 			cancel()
@@ -194,7 +230,11 @@ func runGenerativeORTSessionOnBatch(ctx context.Context, batch *PipelineBatch, p
 			return nil, nil, fmt.Errorf("error during multimodal generation start: %w", err)
 		}
 	} else {
-		ortTokenStream, ortErrorStream, err = session.Generate(generateCtx, inputs, tools, &ortgenai.GenerationOptions{MaxLength: maxLength, Temperature: temperature, TopP: topP, Seed: seed, Guidance: ortGuidance})
+		if engine != nil {
+			ortTokenStream, ortErrorStream, err = engine.Generate(generateCtx, inputs, tools, &ortgenai.GenerationOptions{MaxLength: maxLength, Temperature: temperature, TopP: topP, Seed: seed, Guidance: ortGuidance})
+		} else {
+			ortTokenStream, ortErrorStream, err = session.Generate(generateCtx, inputs, tools, &ortgenai.GenerationOptions{MaxLength: maxLength, Temperature: temperature, TopP: topP, Seed: seed, Guidance: ortGuidance})
+		}
 		if err != nil {
 			cancel()
 			return nil, nil, fmt.Errorf("error during generation start: %w", err)
@@ -237,14 +277,14 @@ func runGenerativeORTSessionOnBatch(ctx context.Context, batch *PipelineBatch, p
 				if !ok {
 					return
 				}
-				idx := tokenDelta.Sequence
-				if completeSequences[idx] {
+				sequence := tokenDelta.Sequence
+				if completeSequences[sequence] {
 					// Already complete; ignore further tokens for this sequence.
 					continue
 				}
 				if tokenDelta.EOSReached {
 					// EOS terminates sequence; no token content to forward.
-					completeSequences[idx] = true
+					completeSequences[sequence] = true
 					completedCount++
 					if completedCount == totalSequences {
 						cancel()
@@ -254,26 +294,26 @@ func runGenerativeORTSessionOnBatch(ctx context.Context, batch *PipelineBatch, p
 				}
 
 				// Forward token immediately.
-				tokenStream <- SequenceDelta{Token: tokenDelta.Token, Index: idx}
+				tokenStream <- SequenceDelta{Token: tokenDelta.Token, Sequence: sequence}
 
 				// Detect stop sequences using a bounded rolling window.
 				// Window size is maxStopLen + len(current token) to catch:
 				//  - stops spanning the boundary between previous tail and this token
 				//  - stops fully contained inside a single (possibly long) token
 				if len(filteredStops) > 0 && maxStopLen > 0 {
-					tail := tails[idx] + tokenDelta.Token
+					tail := tails[sequence] + tokenDelta.Token
 					keep := maxStopLen + len(tokenDelta.Token)
 					if keep < maxStopLen {
 						keep = maxStopLen
 					}
 					if len(tail) > keep {
 						tail = tail[len(tail)-keep:]
 					}
-					tails[idx] = tail
+					tails[sequence] = tail
 
 					for _, s := range filteredStops {
 						if strings.Contains(tail, s) {
-							completeSequences[idx] = true
+							completeSequences[sequence] = true
 							completedCount++
 							if completedCount == totalSequences {
 								cancel()

backends/model_ort_disabled.go

@@ -11,7 +11,8 @@ import (
 
 type ORTModel struct {
 	Destroy           func() error
-	GenerativeSession disabledGenerativeSession // placeholder when ORT disabled
+	GenerativeSession *disabledGenerativeSession // placeholder when ORT disabled
+	GenerativeEngine  *disabledGenerativeEngine  // placeholder when ORT disabled
 }
 
 func createORTModelBackend(_ *Model, _ *options.Options) error {
@@ -46,12 +47,19 @@ func CreateMessagesORT(_ *PipelineBatch, _ any, _ string) error {
 	return errors.New("ORT is not enabled")
 }
 
-type disabledGenerativeSession struct{}
+type (
+	disabledGenerativeSession struct{}
+	disabledGenerativeEngine  struct{}
+)
 
 func (*disabledGenerativeSession) GetStatistics() disabledStatistics {
 	return disabledStatistics{}
 }
 
+func (*disabledGenerativeEngine) GetStatistics() disabledStatistics {
+	return disabledStatistics{}
+}
+
 type disabledStatistics struct {
 	AvgPrefillSeconds              float64
 	TokensPerSecond                float64

docker-bake.hcl

@@ -5,7 +5,7 @@ variable "GO_VERSION"                { default = "1.26.2" }
 variable "GOTESTSUM_VERSION"         { default = "1.13.0" }
 variable "GOPJRT_VERSION"            { default = "0.98.0" }
 variable "ONNXRUNTIME_VERSION"       { default = "1.24.4" }
-variable "ONNXRUNTIME_GENAI_VERSION" { default = "0.12.2" }
+variable "ONNXRUNTIME_GENAI_VERSION" { default = "0.13.1" }
 variable "JAX_CUDA_VERSION"          { default = "0.9.1" }
 
 target "base" {

go.mod

@@ -1,13 +1,13 @@
 module github.com/knights-analytics/hugot
 
-go 1.25.0
+go 1.26.0
 
 require (
 	github.com/daulet/tokenizers v1.27.0
 	github.com/gomlx/go-huggingface v0.3.5
 	github.com/gomlx/gomlx v0.27.3
 	github.com/gomlx/onnx-gomlx v0.4.2
-	github.com/knights-analytics/ortgenai v0.2.0
+	github.com/knights-analytics/ortgenai v0.3.0
 	github.com/stretchr/testify v1.11.1
 	github.com/viant/afs v1.30.0
 	github.com/yalue/onnxruntime_go v1.27.0

go.sum

@@ -58,8 +58,8 @@ github.com/janpfeifer/go-benchmarks v0.1.1 h1:gLLy07/JrOKSnMWeUxSnjTdhkglgmrNR2I
 github.com/janpfeifer/go-benchmarks v0.1.1/go.mod h1:5AagXCOUzevvmYFQalcgoa4oWPyH1IkZNckolGWfiSM=
 github.com/janpfeifer/must v0.2.0 h1:yWy1CE5gtk1i2ICBvqAcMMXrCMqil9CJPkc7x81fRdQ=
 github.com/janpfeifer/must v0.2.0/go.mod h1:S6c5Yg/YSMR43cJw4zhIq7HFMci90a7kPY9XA4c8UIs=
-github.com/knights-analytics/ortgenai v0.2.0 h1:WOZAHxbvlHswydkDrBJi/XnqdzGT1kErUA4hnzoCajA=
-github.com/knights-analytics/ortgenai v0.2.0/go.mod h1:NsxP23iC77IP6q+gRTC+v79v3hyv7fKH1iSa7D+DoVk=
+github.com/knights-analytics/ortgenai v0.3.0 h1:ruHiWzxGnKVowlDQ4zKA3dafLLmFxus6R1vqA0rp6Ss=
+github.com/knights-analytics/ortgenai v0.3.0/go.mod h1:lSbQsRP5wY5NS+4W5CUGhdxjTzERQkR7WprAFxrBSt4=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=

hugot.go

@@ -192,10 +192,9 @@ func NewPipeline[T backends.Pipeline](s *Session, pipelineConfig backends.Pipeli
 	defer pipelineLock.Unlock()
 
 	_, getError := GetPipeline[T](s, pipelineConfig.Name)
-	var notFoundError *pipelineNotFoundError
 	if getError == nil {
 		return pipeline, fmt.Errorf("pipeline %s has already been initialised", pipelineConfig.Name)
-	} else if !errors.As(getError, &notFoundError) {
+	} else if _, ok := errors.AsType[*pipelineNotFoundError](getError); !ok {
 		return pipeline, getError
 	}
 

hugot_ort_test.go

@@ -331,9 +331,11 @@ func TestTextGenerationPipelineORTCuda(t *testing.T) {
 	if os.Getenv("CI") != "" {
 		t.SkipNow()
 	}
-	session, err := NewORTSession(t.Context(), options.WithCuda(map[string]string{
-		"device_id": "0",
-	}))
+	session, err := NewORTSession(t.Context(),
+		options.WithCuda(map[string]string{
+			"device_id": "0",
+		}),
+		options.WithUseEngine())
 	checkT(t, err)
 	defer func(session *Session) {
 		destroyErr := session.Destroy()